"""Pure synthetic credit tests; no brain execution or historical dynamic read."""
import ast
import hashlib
import math
from pathlib import Path
import types
import unittest

import numpy as np

from credit import CATEGORY_CODES, N_GLOBAL, pair_identity, replay_credit


SOURCE_PINS = {
    'phase4o_event_replay.py': 'b169937bab495c77d8ee6c1013877e7500e7354140d9fcc90595d984238ba44e',
    'phase4l.py': 'e97d58d8c1f44920d0e953f8f5831ddc86fa0198784e86460e315de4116ac553',
}


def frozen_tree(filename):
    """Only static source bytes are read; never import the old script or its main."""
    here = Path(__file__).absolute().parent
    candidates = [here.parent / 'models/male-cns-lif' / filename,
                  here.parent / filename,
                  here.parent / ('runtime-credit-v1' if filename.startswith('phase4o') else 'olfactory-grounding-v1') / filename]
    path = next((path for path in candidates if path.is_file()), None)
    if path is None:
        raise unittest.SkipTest('Frozen arithmetic source not installed at declared static paths: ' + filename)
    source = path.read_bytes()
    if hashlib.sha256(source).hexdigest() != SOURCE_PINS[filename]:
        raise AssertionError('Frozen arithmetic source SHA mismatch: ' + str(path))
    return ast.parse(source, filename=str(path)), path


def assert_bits(test, first, second):
    test.assertEqual(first.dtype, second.dtype)
    test.assertEqual(first.shape, second.shape)
    np.testing.assert_array_equal(first.view(np.uint8), second.view(np.uint8))


def events(rows):
    if not rows:
        return np.empty(0, np.int32), np.empty(0, np.int32)
    array = np.asarray(rows, np.int32)
    return array[:, 0].copy(), array[:, 1].copy()


class ReplayTests(unittest.TestCase):
    def setUp(self):
        self.kc = np.array([3, 9, 12], np.int32)

    def test_gate_900_includes_same_tick_901_is_final_only_no_18_shift(self):
        ticks, nodes = events([(882, 12), (900, 3), (901, 9)])
        result = replay_credit(ticks, nodes, self.kc)
        self.assertEqual(result['trace_gate'][0], np.float32(1))
        self.assertEqual(result['trace_gate'][1], np.float32(0))
        decay = np.float32(math.exp(-.1 / 100))
        expected = np.float32(1)
        for _ in range(18):
            expected = np.float32(expected * decay)
        self.assertEqual(result['trace_gate'][2].view(np.uint32), expected.view(np.uint32))
        self.assertTrue(np.all(result['trace_final'] > 0))
        np.testing.assert_array_equal(result['spike_count_gate'], [1, 0, 1])
        np.testing.assert_array_equal(result['spike_count_total'], [1, 1, 1])
        self.assertEqual(result['event_shift_steps'], 0)
        self.assertTrue(result['gate_is_virtual_readout_not_DAN_activity'])

    def test_spike_cap_resets_recency_not_accumulated_count(self):
        ticks, nodes = events([(0, 3), (7, 3), (7, 9)])
        result = replay_credit(ticks, nodes, self.kc, steps=15, gate_tick=7)
        np.testing.assert_array_equal(result['trace_gate'], np.array([1, 1, 0], np.float32))
        self.assertEqual(result['trace_final'][0].view(np.uint32), result['trace_final'][1].view(np.uint32))
        np.testing.assert_array_equal(result['spike_count_total'], [2, 1, 0])

    def test_tick_zero_gate_and_last_tick_have_no_extra_decay(self):
        ticks, nodes = events([(0, 3), (2, 9)])
        result = replay_credit(ticks, nodes, self.kc, steps=3, gate_tick=0)
        self.assertEqual(result['trace_gate'][0], np.float32(1))
        self.assertEqual(result['trace_final'][1], np.float32(1))
        self.assertEqual(result['spike_count_gate'].tolist(), [1, 0, 0])

    def test_empty_events_are_positive_zero_with_exact_dtypes(self):
        ticks, nodes = events([])
        result = replay_credit(ticks, nodes, self.kc, steps=2, gate_tick=1)
        for key in ('trace_gate', 'trace_final'):
            self.assertEqual(result[key].dtype, np.dtype('float32'))
            self.assertTrue(np.all(result[key].view(np.uint32) == 0))
        for key in ('spike_count_gate', 'spike_count_total'):
            self.assertEqual(result[key].dtype, np.dtype('int32'))
            self.assertFalse(result[key].any())

    def test_non_kc_events_ignored_only_after_full_validation(self):
        ticks, nodes = events([(0, 2), (0, 3), (1, 4)])
        for value in (ticks, nodes, self.kc):
            value.flags.writeable = False
        before = [value.tobytes() for value in (ticks, nodes, self.kc)]
        result = replay_credit(ticks, nodes, self.kc, steps=2, gate_tick=1)
        self.assertEqual(result['total_input_events'], 3)
        self.assertEqual(result['total_KC_events'], 1)
        self.assertEqual(before, [value.tobytes() for value in (ticks, nodes, self.kc)])
        result['trace_gate'][0] = 0
        self.assertGreater(result['trace_final'][0], 0)

    def test_invalid_event_order_duplicate_ranges_and_dtype(self):
        bad_rows = [[(-1, 3)], [(1000, 3)], [(0, -1)], [(0, N_GLOBAL)],
                    [(1, 3), (0, 3)], [(0, 9), (0, 3)], [(0, 3), (0, 3)], [(0, 2), (0, 2)]]
        for rows in bad_rows:
            with self.subTest(rows=rows), self.assertRaises(ValueError):
                replay_credit(*events(rows), self.kc)
        for dtype in (np.int64, np.float32, np.float64):
            with self.subTest(dtype=dtype), self.assertRaises(ValueError):
                replay_credit(np.array([], dtype), np.array([], dtype), self.kc)
        with self.assertRaises(ValueError):
            replay_credit(np.array([0], np.int32), np.array([], np.int32), self.kc)
        with self.assertRaises(ValueError):
            replay_credit(np.array([np.nan]), np.array([3], np.int32), self.kc)

    def test_invalid_kc_identities_and_frozen_parameters(self):
        empty = np.empty(0, np.int32)
        for kc in (empty, np.array([9, 3], np.int32), np.array([3, 3], np.int32),
                   np.array([-1], np.int32), np.array([N_GLOBAL], np.int32), self.kc.astype(np.int64)):
            with self.subTest(kc=kc.tolist()), self.assertRaises(ValueError):
                replay_credit(empty, empty, kc)
        for keyword, value in (('steps', 0), ('steps', True), ('gate_tick', -1), ('gate_tick', 1000),
                               ('gate_tick', 900.), ('dt_ms', .2), ('tau_ms', 20), ('tau_ms', np.inf),
                               ('cap', .5), ('cap', True)):
            with self.subTest(keyword=keyword, value=value), self.assertRaises(ValueError):
                replay_credit(empty, empty, self.kc, **{keyword: value})

    def test_independent_frozen_numpy_loop_ast_matches_all_trace_bits(self):
        tree, path = frozen_tree('phase4o_event_replay.py')
        loops = [node for node in ast.walk(tree) if isinstance(node, ast.For) and isinstance(node.target, ast.Name)
                 and node.target.id == 'step' and any(isinstance(child, ast.AugAssign) and
                 isinstance(child.target, ast.Name) and child.target.id == 'trace' for child in node.body)]
        self.assertEqual(len(loops), 1)
        self.assertEqual(len(loops[0].body), 4)
        code = compile(ast.fix_missing_locations(ast.Module(body=[loops[0]], type_ignores=[])), str(path), 'exec')
        ticks, nodes = events([(0, 3), (0, 12), (1, 9), (22, 3), (899, 9), (900, 3), (901, 12), (999, 3)])
        result = replay_credit(ticks, nodes, self.kc)
        lookup = {int(node): index for index, node in enumerate(self.kc)}
        for gate, field in ((900, 'trace_gate'), (999, 'trace_final')):
            by_step = {}
            for tick, node in zip(ticks, nodes):
                by_step.setdefault(int(tick), []).append(lookup[int(node)])
            namespace = dict(np=np, trace=np.zeros(len(self.kc), np.float32),
                decay=np.float32(math.exp(-.1 / 100)), by_step=by_step, step_gate=gate)
            exec(code, namespace)
            assert_bits(self, result[field], namespace['trace'])

    def test_independent_frozen_torch_cpu_statement_matches_when_available(self):
        try:
            import torch
        except (ImportError, OSError) as error:
            self.skipTest('Torch CPU unavailable; comparison explicitly not performed: ' + str(error))
        tree, path = frozen_tree('phase4l.py')
        expressions = [node for node in ast.walk(tree) if isinstance(node, ast.Expr) and
            'trace.mul_(trace_decay).add_' in ast.unparse(node)]
        self.assertEqual(len(expressions), 1)
        code = compile(ast.fix_missing_locations(ast.Module(body=expressions, type_ignores=[])), str(path), 'exec')
        ticks, nodes = events([(0, 3), (22, 3), (899, 9), (900, 3), (901, 12), (999, 9)])
        result = replay_credit(ticks, nodes, self.kc)
        trace = torch.zeros(len(self.kc), dtype=torch.float32, device='cpu')
        namespace = dict(trace=trace, trace_decay=math.exp(-.1 / 100), torch=torch,
            self=types.SimpleNamespace(kc_gpu=torch.tensor(self.kc.astype(np.int64), device='cpu')))
        gate = None
        for tick in range(1000):
            spikes = torch.zeros(13, dtype=torch.bool, device='cpu')
            active = nodes[ticks == tick]
            if len(active):
                spikes[torch.tensor(active.astype(np.int64), device='cpu')] = True
            namespace['spikes'] = spikes
            exec(code, namespace)
            if tick == 900:
                gate = trace.numpy().copy()
        assert_bits(self, result['trace_gate'], gate)
        assert_bits(self, result['trace_final'], trace.numpy())


class IdentityTests(unittest.TestCase):
    def setUp(self):
        self.a, self.b = np.zeros((5, 4), np.float32), np.zeros((5, 4), np.float32)
        self.a[:3, 0] = 1  # A specific
        self.a[:3, 1] = .1; self.b[:4, 1] = .2  # shared
        self.b[:3, 2] = .3  # B specific
        self.a[:2, 3] = 1; self.b[:2, 3] = 1  # no majority
        self.source = np.array([0, 1, 1, 2, 3], np.int32)

    def test_three_of_five_categories_and_full_set_jaccard(self):
        result = pair_identity(self.a, self.b, self.source)
        np.testing.assert_array_equal(result['kc_category'], np.array([1, 3, 2, 0], np.int8))
        np.testing.assert_array_equal(result['edge_category'], np.array([1, 3, 3, 2, 0], np.int8))
        self.assertEqual(result['summary']['KC']['active_set_jaccard'], 1 / 3)
        self.assertEqual(result['summary']['edges']['active_set_jaccard'], .5)
        self.assertEqual(result['summary']['KC']['category_counts'], {name: 1 for name in CATEGORY_CODES})
        self.assertTrue(result['summary']['edge_counts_are_source_degree_weighted'])

    def test_tiny_positive_no_epsilon_and_no_weight_update_outputs(self):
        self.a[:3, 0] = np.nextafter(np.float32(0), np.float32(1))
        result = pair_identity(self.a, self.b, self.source)
        self.assertTrue(result['kc_active_a'][0])
        self.assertFalse(result['summary']['weights_modified'])
        self.assertFalse(result['summary']['weight_delta_generated'])
        self.assertFalse(any('delta' in key or 'weight' in key for key in result if key != 'summary'))

    def test_empty_union_and_empty_edge_mapping_explicit(self):
        result = pair_identity(np.zeros_like(self.a), np.zeros_like(self.b), np.empty(0, np.int32))
        self.assertTrue(result['summary']['KC']['empty_union'])
        self.assertEqual(result['summary']['KC']['active_set_jaccard'], 1.)
        self.assertTrue(result['summary']['edges']['empty_union'])
        self.assertEqual(result['edge_category'].dtype, np.dtype('int8'))

    def test_readonly_inputs_and_returned_mapping_owns_memory(self):
        for value in (self.a, self.b, self.source):
            value.flags.writeable = False
        before = [value.tobytes() for value in (self.a, self.b, self.source)]
        result = pair_identity(self.a, self.b, self.source)
        result['edge_category'][0] = 0
        self.assertEqual(result['kc_category'][0], 1)
        self.assertEqual(before, [value.tobytes() for value in (self.a, self.b, self.source)])

    def test_reject_wrong_dtype_shape_nonfinite_negative_cap_and_source(self):
        for bad in (self.a.astype(np.float64), self.a[:4], np.zeros((5, 0), np.float32)):
            with self.subTest(shape=bad.shape), self.assertRaises(ValueError):
                pair_identity(bad, self.b, self.source)
        for value in (np.nan, np.inf, -.1, 1.01):
            bad = self.a.copy(); bad[0, 0] = value
            with self.subTest(value=value), self.assertRaises(ValueError):
                pair_identity(bad, self.b, self.source)
        with self.assertRaises(ValueError):
            pair_identity(self.a, self.b[:, :3], self.source)
        for source in (np.array([-1], np.int32), np.array([4], np.int32), self.source.astype(np.int64)):
            with self.subTest(source=source.tolist()), self.assertRaises(ValueError):
                pair_identity(self.a, self.b, source)


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