import random
import unittest
from lightnest.learned_patch import MAX_WINDOW
from lightnest.train_patch_rl import execution_reward, passes
from lightnest.train_robust_patch_rl import FAMILIES, TRAIN_COMMANDS, TEST_COMMANDS, make_task

class RobustTests(unittest.TestCase):
    def test_templates_disjoint(self):
        self.assertFalse(set(sum(TRAIN_COMMANDS.values(),())) & set(sum(TEST_COMMANDS.values(),())))

    def test_semantic_oracles_and_none(self):
        rng=random.Random(123)
        for held in (False,True):
            for family in FAMILIES:
                e,cc=make_task(rng,family,held)
                negative=family in ('already_correct','negation','keep','negation_dead')
                self.assertEqual(passes(e.binary,cc),negative)
                self.assertEqual(execution_reward(e.binary,MAX_WINDOW,0,cc),1 if negative else -1)
                if negative:
                    self.assertEqual(execution_reward(e.binary,0,e.binary[0],cc),-1)

    def test_distractor_edit_not_rewarded(self):
        e,cc=make_task(random.Random(51),'dead_before')
        # Exhaustively evaluate single-byte add/sub toggles for this test only.
        successes=[]
        for i,b in enumerate(e.binary):
            if execution_reward(e.binary,i,b^64,cc)>0: successes.append(i)
        self.assertEqual(len(successes),1)

if __name__=='__main__': unittest.main()