import random
import unittest
import torch
from dataclasses import replace
from lightnest.learned_patch import MAX_WINDOW, PatchNetwork, make_example
from lightnest.train_patch_rl import cases_for, execution_reward, sample_actions, oracle


class PatchRLTests(unittest.TestCase):
    def test_execution_reward_not_teacher_fields(self):
        rng = random.Random(9)
        e = make_example(rng)
        while not e.present:
            e = make_example(rng)
        cases = cases_for(e.target_operation, rng)
        self.assertEqual(execution_reward(e.binary, e.pointer, e.replacement, cases), 1)
        self.assertEqual(execution_reward(e.binary, MAX_WINDOW, 0, cases), -1)
        poisoned = replace(e, pointer=-777, replacement=-99, target=b'bad')
        self.assertEqual(execution_reward(poisoned.binary, e.pointer, e.replacement, cases), 1)

    def test_negative_requires_none(self):
        rng = random.Random(9)
        e = make_example(rng)
        while e.present:
            e = make_example(rng)
        cases = cases_for(e.target_operation, rng)
        self.assertEqual(execution_reward(e.binary, MAX_WINDOW, 0, cases), 1)
        self.assertEqual(execution_reward(e.binary, 0, e.binary[0], cases), -1)
        with self.assertRaises(ValueError):
            execution_reward(e.binary, MAX_WINDOW, 0, [])

    def test_sampling_is_differentiable(self):
        torch.set_num_threads(2)
        model = PatchNetwork()
        es = [make_example(random.Random(i)) for i in range(3)]
        pos, byte, logp = sample_actions(model, es)
        self.assertEqual(pos.shape, (3,4))
        self.assertTrue(bool(((byte >= 0) & (byte <= 255)).all()))
        (-logp.mean()).backward()
        self.assertIsNotNone(model.pointer.weight.grad)
        self.assertIsNotNone(model.replacement.weight.grad)

    def test_signed_composed_oracle(self):
        self.assertEqual(oracle('add', [2147483647, 1]), -2147483648)
        self.assertEqual(oracle('sub', [5, 2], 7), 10)

class PatchRLUpdateTests(unittest.TestCase):
    def test_ties_skip_adam_momentum(self):
        from lightnest.train_patch_rl import reinforce_step
        model = torch.nn.Linear(1,1)
        opt = torch.optim.Adam(model.parameters(), lr=.01)
        x = torch.ones(1,4,1)
        lp = model(x).squeeze(-1)
        # First create a nonzero momentum state using nonidentical log-probs.
        lp = lp * torch.tensor([[1.,2.,3.,4.]])
        stepped, _, groups = reinforce_step(opt, model, lp, torch.tensor([[1.,-1.,1.,-1.]]))
        self.assertTrue(stepped); self.assertEqual(groups,1)
        before = {k:v.clone() for k,v in model.state_dict().items()}
        step_before = opt.state[model.weight]['step'].clone()
        stepped, loss, groups = reinforce_step(opt, model, model(x).squeeze(-1), torch.ones(1,4))
        self.assertFalse(stepped); self.assertEqual(loss,0); self.assertEqual(groups,0)
        self.assertTrue(all(torch.equal(v,model.state_dict()[k]) for k,v in before.items()))
        self.assertTrue(torch.equal(step_before,opt.state[model.weight]['step']))

    def test_train_templates_disjoint_from_eval_templates(self):
        from lightnest.train_patch_rl import RL_TEMPLATES
        from lightnest.learned_patch import TRAIN_TEMPLATES, TEST_TEMPLATES
        self.assertFalse(set(RL_TEMPLATES) & set(TRAIN_TEMPLATES + TEST_TEMPLATES))