"""Independent CPU tests of the new diagnostic oracle, not GPU evidence.

AST tests consume the real pinned AJ/C/R source. Tiny tensor/helper doubles
exercise switch/installation code without claiming a fullgraph simulation,
executing any CUDA kernel, or restoring state as an equivalence oracle.
"""
import ast
import copy
import hashlib
from pathlib import Path
import sys
import tempfile
from types import SimpleNamespace
import unittest
from unittest.mock import patch
import uuid

import numpy as np

import boundary_oracle as b


ROOT = Path(__file__).resolve().parent.parent
AJ = ROOT / 'apl-timing-lab-v1' / 'gpu_timing.py'


def sha(value):
    return hashlib.sha256(value).hexdigest()


def bits(value):
    return np.asarray(value).tobytes()


def fixture(tick=41):
    """Raw global-phase rings; the live oracle must NOT rotate spike slots."""
    v = np.array([-.051, -.047], np.float64)
    g = np.array([-0., .00013], np.float64)
    ref, counts = np.array([22, 3], np.int32), np.array([5, 2], np.int32)
    ring = np.zeros((19, 2), np.uint8)
    ring[(tick - 1) % 19, 0] = 1
    ring[(tick - 18) % 19, 1] = 1
    # Distinct samples across every row expose off-by-one and 18/19 confusion.
    apl_ring = np.arange(36, dtype=np.float32).reshape(18, 2) / np.float32(16)
    apl_ring[0, 0] = np.float32(-0.)
    apl_state = np.array([.5, -0.], np.float32)
    parent = dict(steps=tick, global_next_tick=tick, fresh=True, manifest_sha256='a' * 64,
        state={key: value.copy() for key, value in zip(('v', 'g', 'ref', 'counts'), (v, g, ref, counts))},
        terminal=dict(spike_ring=ring.copy(), apl_state=apl_state.copy(),
            apl_pending=np.concatenate((apl_ring[tick % 18:], apl_ring[:tick % 18])), next_tick=np.asarray(tick, np.int64)))
    return parent, [v, g, ref, counts, ring, apl_ring, apl_state]


class Tensor:
    """NumPy-only tensor double, including assignment accounting."""
    def __init__(self, value):
        self.a = np.array(value, copy=True)
        self.assignments = []

    def cpu(self): return self
    def numpy(self): return self.a
    def numel(self): return self.a.size

    def __getitem__(self, key):
        return Tensor(self.a[key.a if isinstance(key, Tensor) else key])

    def __setitem__(self, key, value):
        index = key.a if isinstance(key, Tensor) else key
        self.assignments.append(np.array(index, copy=True))
        self.a[index] = value.a if isinstance(value, Tensor) else value

    def __add__(self, other): return Tensor(self.a + (other.a if isinstance(other, Tensor) else other))
    def __mul__(self, other): return Tensor(self.a * other)


class CanonicalEdgeIdentity:
    """No actual fullgraph arrays: only the private test's admission double."""
    def __len__(self): return 25582938
    def __getitem__(self, key): return np.asarray(key, dtype=np.int64).copy()


class TinyHelper:
    @staticmethod
    def map_installation(base, edges, apl_mask, plan):
        nz = plan['plastic_delta_mv'] != np.float32(0)
        rows = plan['edge_id'][nz]
        # Independent frozen bridge expression; do not invoke oracle helper.
        expected = np.add(base[rows], np.multiply(plan['plastic_delta_mv'][nz].astype(np.float64), np.float64(.001)))
        return dict(rows=rows, expected_si=expected, missing_candidate_count=0)

    @staticmethod
    def expected_hashes(base, apl_mask, mapping):
        result = base.copy()
        result[mapping['rows']] = mapping['expected_si']
        return sha(bits(base)), sha(bits(result))

    @staticmethod
    def bit_equal(a, c):
        return a.dtype == c.dtype and a.shape == c.shape and a.tobytes() == c.tobytes()


class BoundaryOracleTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Use the actual pinned dependency, with no Torch/CUDA import.
        cls.c = b.frozen()
        cls.source = AJ.read_bytes()

    def test_real_AJ_hook_rollback_and_entire_original_functions_identical(self):
        source_before = sha(self.source)
        derived, audit = b.derive_oracle(self.c, self.source)
        old_reference, original_audit = self.c.R.derive_ast(self.source, 'reference')
        loops = [n for n in ast.walk(derived) if isinstance(n, ast.For) and isinstance(n.target, ast.Name) and n.target.id == 'tick']
        self.assertEqual(len(loops), 1)
        self.assertEqual(ast.unparse(loops[0].body[0]), b.HOOK)
        hookless = copy.deepcopy(derived)
        loop = next(n for n in ast.walk(hookless) if isinstance(n, ast.For) and isinstance(n.target, ast.Name) and n.target.id == 'tick')
        del loop.body[0]
        self.assertEqual(self.c.R._dump(hookless), self.c.R._dump(old_reference))
        # Independently roll back the inherited R edits as well. This reaches
        # the ACTUAL frozen AJ functions, not merely the adapter's own receipt.
        loader, run = hookless.body
        admission = [n for n in ast.walk(loader) if isinstance(n, ast.Compare) and ast.unparse(n) == "0 < steps <= spec['max_steps']"]
        self.assertEqual(len(admission), 1)
        admission[0].comparators[-1] = ast.Constant(value=1000)
        body = next(n for n in run.body if isinstance(n, ast.Try)).body
        for text in (self.c.R._INSTALL, self.c.R._FINISH):
            indices = [i for i, n in enumerate(body) if ast.unparse(n) == text]
            self.assertEqual(len(indices), 1)
            del body[indices[0]]
        originals = ast.parse(self.source)
        original_module = ast.Module(body=[next(n for n in originals.body if isinstance(n, ast.FunctionDef) and n.name == name)
            for name in ('load_inputs', 'run_gpu')], type_ignores=[])
        self.assertEqual(self.c.R._dump(hookless), self.c.R._dump(original_module))
        self.assertEqual(audit['base'], original_audit)
        self.assertTrue(audit['full_AST_rollback_identical'])
        self.assertFalse(audit['restore_called'])
        self.assertFalse(audit['production_backend'])
        self.assertEqual(sha(self.source), source_before)
        compile(derived, '<actual-AJ-oracle-CPU-compile-only>', 'exec')

    def test_oracle_has_one_hook_before_all_frozen_tick_operations_and_no_restore(self):
        tree, _ = b.derive_oracle(self.c, self.source)
        loop = next(n for n in ast.walk(tree) if isinstance(n, ast.For) and isinstance(n.target, ast.Name) and n.target.id == 'tick')
        self.assertEqual(ast.unparse(loop.body[1]), 'requests.zero_()')
        calls = [ast.unparse(n.func) for n in ast.walk(tree) if isinstance(n, ast.Call)]
        self.assertEqual(calls.count('_oracle.switch'), 1)
        self.assertFalse(any('restore' in name for name in calls))
        self.assertFalse(any(isinstance(n, ast.Name) and n.id == '_continuation' for n in ast.walk(tree)))
        oracle_class = b.runtime_class(self.c)
        self.assertIs(oracle_class.__bases__[0], self.c.R._Runtime)
        self.assertFalse(issubclass(oracle_class, self.c.ContinuedRuntime))
        self.assertFalse(hasattr(oracle_class, 'restore'))
        self.assertIn('due = ring[(tick - 18) % 19]', ast.unparse(tree))
        self.assertIn('slot = tick % apl_delay', ast.unparse(tree))

    def test_real_source_pin_rejects_even_whitespace_drift(self):
        for changed in (self.source + b'\n', self.source.replace(b'tick >= 18', b'tick >= 17')):
            with self.assertRaisesRegex(ValueError, 'AJ source pin'):
                b.derive_oracle(self.c, changed)

    def test_frozen_C_pin_and_existing_module_binding_refuse_before_execution(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / 'tiny.py'
            path.write_text('VALUE = 17\n', encoding='utf-8')
            name = 'boundary_oracle_cpu_fixture_' + uuid.uuid4().hex
            with self.assertRaisesRegex(ValueError, 'Unverified'):
                b.load(name, path, '0' * 64)
            self.assertNotIn(name, sys.modules)
            try:
                module = b.load(name, path, b.digest(path))
                self.assertEqual(module.VALUE, 17)
                with self.assertRaisesRegex(ValueError, 'Unverified'):
                    b.load(name, path, b.digest(path))
                self.assertIs(sys.modules[name], module)
            finally:
                sys.modules.pop(name, None)  # Only this test-created tiny module.
            frozen_path = Path(tmp) / 'runtime_continuation.py'
            frozen_path.write_text('raise AssertionError("must not execute")\n', encoding='utf-8')
            with patch.object(b, 'LAB', Path(tmp)), self.assertRaisesRegex(ValueError, 'Unverified'):
                b.frozen()

    def test_expected_bridge_is_original_base_plus_widened_after_delta_not_old_effective(self):
        base = np.array([.000275, .000550, -.00011, .0], np.float64)
        mask = np.array([0, 1], np.int64)
        before = np.array([-.00275, -.013125], np.float32)
        after = np.array([-.0055, -.02125], np.float32)
        old_effective = base[mask] + before.astype(np.float64) * np.float64(.001)
        expected = np.add(base[mask], np.multiply(after.astype(np.float64), np.float64(.001)))
        actual = b.expected_candidate_values(base, mask, after)
        self.assertEqual(bits(actual), bits(expected))
        self.assertNotEqual(bits(actual), bits(old_effective + after.astype(np.float64) * np.float64(.001)))
        self.assertNotEqual(bits(actual), bits(base[mask] + (after * np.float32(.001)).astype(np.float64)))
        zero = np.array([0., -0.], np.float32)
        self.assertEqual(bits(b.expected_candidate_values(base, mask, zero)), bits(base[mask]))
        self.assertEqual(bits(base), bits(np.array([.000275, .000550, -.00011, .0], np.float64)))

    def test_candidate_bridge_rejects_dtype_bounds_order_sign_and_nonfinite(self):
        base, mask, delta = np.array([.01, .02], np.float64), np.array([0, 1], np.int64), np.array([-.001, -.002], np.float32)
        cases = [(base.astype(np.float32), mask, delta), (base, mask.astype(np.int32), delta),
            (base, mask, delta.astype(np.float64)), (base, mask[::-1].copy(), delta),
            (base, np.array([0, 0], np.int64), delta), (base, np.array([-1, 1], np.int64), delta),
            (base, np.array([0, 2], np.int64), delta), (base, mask, np.array([np.nan, 0.], np.float32)),
            (base, mask, np.array([np.inf, 0.], np.float32)), (base, mask, np.array([.01, 0.], np.float32)),
            (base, mask, np.array([-20., 0.], np.float32)), (np.array([np.inf, .02], np.float64), mask, delta)]
        for args in cases:
            with self.subTest(args=[(a.dtype.str, a.tolist()) for a in args]), self.assertRaises(ValueError):
                b.expected_candidate_values(*args)

    def test_live_raw_boundary_checks_18_and_19_phases_without_rotating_spike_ring(self):
        for tick in (1, 5, 17, 18, 19, 20, 35, 36, 37, 38, 41, 1200):
            parent, arrays = fixture(tick)
            before = [bits(value) for value in arrays]
            actual, checks = b.boundary_arrays(parent, tick, *arrays)
            self.assertTrue(all(checks.values()))
            self.assertEqual(bits(actual['spike_ring']), bits(arrays[4]))
            order = np.array([(tick + i) % 18 for i in range(18)])
            self.assertEqual(bits(actual['apl_pending']), bits(arrays[5][order]))
            self.assertEqual(before, [bits(value) for value in arrays])
            if tick % 19:
                wrong = list(arrays); wrong[4] = self.c.rotate_spike_ring(arrays[4], tick)
                with self.subTest(tick=tick), self.assertRaisesRegex(ValueError, 'spike_ring'):
                    b.boundary_arrays(parent, tick, *wrong)

    def test_boundary_rejects_rawbit_dtype_queue_clock_or_counts_corruption(self):
        for variant in ('negative_zero_g', 'negative_zero_apl', 'g_dtype', 'v_shape', 'counts', 'spike', 'apl_phase', 'clock', 'continued', 'booltick'):
            parent, arrays = fixture(41)
            tick = 41
            if variant == 'negative_zero_g': arrays[1][0] = 0.
            if variant == 'negative_zero_apl': arrays[6][1] = np.float32(0.)
            if variant == 'g_dtype': arrays[1] = arrays[1].astype(np.float32)
            if variant == 'v_shape': arrays[0] = arrays[0][:-1]
            if variant == 'counts': arrays[3][0] += 1
            if variant == 'spike': arrays[4][2, 0] ^= 1
            if variant == 'apl_phase': arrays[5] = np.roll(arrays[5], 1, axis=0)
            if variant == 'clock': parent['global_next_tick'] += 1
            if variant == 'continued': parent['fresh'] = False
            if variant == 'booltick': tick = True
            with self.subTest(variant=variant), self.assertRaises(ValueError): b.boundary_arrays(parent, tick, *arrays)

    def runtime_fixture(self, out):
        parent, arrays = fixture(41)
        base = np.array([.000275, .000550, -.00011, .0], np.float64)
        mask = np.array([0, 1], np.int64)
        old = dict(edge_id=mask, plastic_delta_mv=np.array([-.00275, -.013125], np.float32))
        new = dict(edge_id=mask, plastic_delta_mv=np.array([-.0055, -.02125], np.float32))
        helper = TinyHelper()
        apl_mask = SimpleNamespace(sum=lambda: 5845)
        prior = helper.map_installation(base, None, apl_mask, old)
        following = helper.map_installation(base, None, apl_mask, new)
        before_sha = helper.expected_hashes(base, apl_mask, prior)[1]
        after_sha = helper.expected_hashes(base, apl_mask, following)[1]
        parent['summary'] = {'effective_weights_raw_sha256': before_sha}
        saved = {}
        # Only these diagnostic I/O and GPU-observer dependencies are doubles;
        # the actual Oracle install/switch implementation is executed below.
        r = SimpleNamespace(_Runtime=self.c.R._Runtime, FiniteObserver=lambda *_: object(),
            _copy_input=lambda *_: None, digest=lambda p: sha(str(p).encode()),
            _save=lambda p, value: saved.__setitem__(Path(p).name, value))
        cls = b.runtime_class(SimpleNamespace(R=r))
        boundary = dict(parent=parent, after_plan=new,
            value=dict(expected_effective_weight_before_sha256=before_sha, expected_effective_weight_after_sha256=after_sha),
            paths={key: out / key for key in ('before_overlay', 'before_manifest', 'after_overlay', 'after_manifest')})
        obj = cls('reference', {'max_steps': 60, 'steps': 60}, helper, old, {}, {},
            SimpleNamespace(int64=np.int64, float64=np.float64), out / 'old.npz', out / 'old.json', boundary=boundary)
        weights = Tensor(base)
        graph = dict(nodes=range(166700), edge_id=CanonicalEdgeIdentity(), weights=base)
        upload = lambda value, dtype: Tensor(np.asarray(value, dtype=dtype))
        args = [Tensor(value) for value in arrays[:-1]] + [SimpleNamespace(state=Tensor(arrays[-1]))]
        return obj, weights, graph, apl_mask, upload, args, saved, base, new

    def test_actual_install_then_one_switch_replaces_base_plus_after_and_never_state_or_queues(self):
        with tempfile.TemporaryDirectory() as tmp:
            out = Path(tmp)
            obj, weights, graph, mask, upload, args, saved, base, new = self.runtime_fixture(out)
            obj.install(weights, graph, mask, upload, out, {}, 'CPU-double')
            before = weights.a.copy()
            state_before = [bits(x.a) for x in args[:-1]] + [bits(args[-1].state.a)]
            obj.switch(41, weights, *args, out)
            expected = base.copy()
            expected[new['edge_id']] = np.add(base[new['edge_id']], np.multiply(new['plastic_delta_mv'].astype(np.float64), np.float64(.001)))
            self.assertEqual(bits(weights.a), bits(expected))
            wrong = before.copy(); wrong[new['edge_id']] += new['plastic_delta_mv'].astype(np.float64) * np.float64(.001)
            self.assertNotEqual(bits(weights.a), bits(wrong))
            self.assertEqual(bits(weights.a[2:]), bits(before[2:]))
            self.assertEqual(obj.switches, 1)
            self.assertEqual(obj.installations, 1)
            self.assertEqual(state_before, [bits(x.a) for x in args[:-1]] + [bits(args[-1].state.a)])
            self.assertFalse(saved['boundary-comparison.json']['restore_called'])
            self.assertFalse(saved['boundary-comparison.json']['queues_replaced'])
            self.assertEqual(len(weights.assignments), 2)  # One old-overlay install, one replacement.
            with np.load(out / 'boundary-snapshot.npz', allow_pickle=False) as recorded:
                self.assertEqual(bits(recorded['g']), state_before[1])
            with self.assertRaisesRegex(ValueError, 'repeated'): obj.switch(41, weights, *args, out)
            with self.assertRaisesRegex(ValueError, 'one intact'): obj.install(weights, graph, mask, upload, out, {}, 'CPU-double')

    def test_switch_refuses_wrong_time_or_preexisting_weight_state_drift(self):
        for variant in ('before_install', 'wrong_tick', 'weight_drift', 'state_drift', 'after_hash'):
            with tempfile.TemporaryDirectory() as tmp:
                out = Path(tmp)
                obj, weights, graph, mask, upload, args, _, _, _ = self.runtime_fixture(out)
                obj.install(weights, graph, mask, upload, out, {}, 'CPU-double')
                tick = 41
                if variant == 'before_install': obj.installations = 0
                if variant == 'wrong_tick': tick = 42
                if variant == 'weight_drift': weights.a[3] = .0001
                if variant == 'state_drift': args[1].a[0] = 0.
                if variant == 'after_hash': obj.after_sha = 'f' * 64
                with self.subTest(variant=variant), self.assertRaises(ValueError): obj.switch(tick, weights, *args, out)
                self.assertEqual(obj.switches, 0)
                self.assertFalse((out / 'boundary-snapshot.npz').exists())

    def test_finish_requires_the_boundary_and_summary_discloses_diagnostic_scope(self):
        cls = b.runtime_class(self.c)
        obj = object.__new__(cls)
        obj.switches = 0
        with self.assertRaisesRegex(ValueError, 'did not execute'): obj.finish()
        obj.mode = 'reference'; obj.spec = dict(max_steps=100, steps=100)
        obj.result = {'first_failure_tick': -1}; obj.boundary_tick = 41; obj.switches = 1
        obj.after_sha = 'a' * 64; obj.before_sha = 'b' * 64
        result = obj.summary({'complete': True, 'effective_weights_raw_sha256': 'b' * 64, 'learning': False})
        self.assertTrue(result['within_trial_weight_updates'])
        self.assertFalse(result['production_backend'])
        self.assertFalse(result['restored_from_checkpoint'])
        self.assertFalse(result['learning_pass_claimed'])
        self.assertEqual(result['effective_weights_raw_sha256'], 'a' * 64)
        self.assertEqual(result['initial_effective_weights_raw_sha256'], 'b' * 64)


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