import random
import unittest
import torch
from lightnest import rv32
from lightnest.learned_patch import (Example, LearnedPatch, MAX_WINDOW, PatchNetwork, Predictor, TRAIN_REGS, TEST_REGS, TRAIN_TEMPLATES, TEST_TEMPLATES, batch, encode, independent_expected, make_example, validate_candidate)


class ForcedNetwork(torch.nn.Module):
    def __init__(self, offset, new):
        super().__init__()
        self.offset, self.new = offset, new

    def forward(self, command, binary):
        scores = torch.full((1, 257), -30.)
        scores[0, self.offset] = 30.
        bits = torch.full((1, 256, 8), -30.)
        if self.offset < 256:
            for i in range(8):
                bits[0, self.offset, i] = 30. if self.new & (1 << i) else -30.
        return scores, bits


class LearnedPatchTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torch.set_num_threads(2)

    def test_dataset_oracle_and_semantics(self):
        rng = random.Random(1)
        for split in ('train', 'combined'):
            for _ in range(60):
                e = make_example(rng, split)
                self.assertLessEqual(len(e.binary), 256)
                diffs = [i for i, (a, b) in enumerate(zip(e.binary, e.target)) if a != b]
                self.assertEqual(len(diffs), int(e.present))
                if not e.present:
                    self.assertEqual(e.pointer, MAX_WINDOW)
                for x in ([3, 7], [-27, 12], [0x7fffffff, 1]):
                    expected = independent_expected(e, x)
                    self.assertEqual(rv32.simulate(e.target, x), expected)
                if e.present:
                    patch = LearnedPatch(e.pointer, e.binary[e.pointer], e.replacement)
                    self.assertTrue(validate_candidate(e.binary, patch, [([3, 7], rv32.simulate(e.target, [3, 7]))]))

    def test_disjoint_holdouts(self):
        self.assertFalse(set(TRAIN_REGS) & set(TEST_REGS))
        self.assertFalse(set(TRAIN_TEMPLATES) & set(TEST_TEMPLATES))
        rng = random.Random(2)
        for _ in range(100):
            train, test = make_example(rng), make_example(rng, 'combined')
            self.assertLessEqual(train.prefix_words, 16)
            self.assertGreaterEqual(test.prefix_words, 22)
            self.assertTrue(set(test.registers) <= set(TEST_REGS))
            self.assertIn(test.template, TEST_TEMPLATES)

    def test_proposal_follows_neural_output_not_labels(self):
        # A deliberately nonsensical instruction still returns forced neural bytes.
        predictor = Predictor(ForcedNetwork(5, 123))
        patch = predictor.propose('arbitrary unknown text', bytes(range(32)), base_offset=2**40)
        self.assertEqual((patch.relative_offset, patch.old_byte, patch.new_byte), (5, 5, 123))
        self.assertEqual(patch.absolute_offset, 2**40 + 5)
        self.assertIn('unvalidated', predictor.last_status)

    def test_none_and_reject(self):
        self.assertIsNone(Predictor(ForcedNetwork(256, 0)).propose('change add to sub', b'1234'))
        rng = random.Random(3)
        e = make_example(rng)
        bad = LearnedPatch(0, e.binary[0], 255)
        self.assertFalse(validate_candidate(e.binary, bad, [([7, 4], rv32.simulate(e.target, [7, 4]))]))
        self.assertFalse(validate_candidate(e.binary, bad, []))
        self.assertFalse(validate_candidate(e.binary, None, [([7, 4], 3)]))

    def test_resources_and_bounds(self):
        model = PatchNetwork()
        self.assertLess(sum(p.numel() for p in model.parameters()), 1_000_000)
        e = make_example(random.Random(4))
        c, b, _, _ = batch([e])
        logits, bits = model(c, b)
        self.assertEqual(tuple(logits.shape), (1, 257))
        self.assertEqual(tuple(bits.shape), (1, 256, 8))
        for command, binary in [('a'*97, b'1'), ('a', b'1'*257), ('', b'1'), ('a', b'')]:
            with self.assertRaises(ValueError):
                encode(command, binary)
        with self.assertRaises(ValueError):
            Predictor(model).propose('a', b'1', -1)


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


def test_nonfinite_policy_fails_closed(tmp_path):
    import pytest
    import torch
    from lightnest.learned_patch import PatchNetwork, Predictor, load_predictor
    for threshold in (float('nan'), float('inf'), 0., 1.1):
        with pytest.raises(ValueError):
            Predictor(PatchNetwork(), threshold)
    class Broken(torch.nn.Module):
        def forward(self, command, binary):
            return torch.full((1,257),float('nan')), torch.zeros((1,256,8))
    with pytest.raises(ValueError,match='nonfinite'):
        Predictor(Broken()).propose('change add to sub',b'\x13\0\0\0')
    state=PatchNetwork().state_dict()
    state['pointer.bias'][0]=float('nan')
    path=tmp_path/'bad.pt';torch.save({'state_dict':state,'threshold':.55},path)
    with pytest.raises(ValueError,match='tensor'):
        load_predictor(path)