"""Tests for the hyperdimensional full adder.

Two of these are not ordinary arithmetic tests:

* ``TestGate.test_keys_are_not_parity`` is a regression test for a real bug.
  XORing shared zero/one atoms collapsed the 8-row truth table onto 2 keys,
  because ``1^1^1 == 1^0^0``.  The fix was one atom pair per gate slot.
* ``TestNoDecode`` patches out ``hv_to_int`` inside the interpreter and runs
  ADD/SUB, which is the only way to show the operands never leave the vector
  domain: if any decode remained, the program would fault.
"""

import random
import sys
import unittest
from pathlib import Path
from unittest import mock

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from hdc.adder import AdderError, HyperAdder, wrap
from hdc.assembler import assemble
from hdc.hypervec import (
    MAX_INT_BITS,
    bind,
    hv_bits_of_int,
    hv_from_bits,
    hv_to_bits,
    hv_to_int,
    int_to_hv,
    sim,
)
from hdc.interpreter import Machine
from hdc.memory import HDCStorage

D = 2000


def build(source):
    m = Machine(storage=HDCStorage(d=D), d=D)
    m.load(assemble(source, d=D))
    return m


class TestGate(unittest.TestCase):
    def setUp(self):
        self.adder = HyperAdder(d=D)

    def test_table_has_eight_correct_rows(self):
        self.assertEqual(len(self.adder.rows), 8)
        for row in self.adder.rows:
            total = row.a ^ row.b ^ row.c
            carry = 1 if (row.a + row.b + row.c) >= 2 else 0
            self.assertEqual((row.total, row.carry), (total, carry), row.render())

    def test_rows_are_stored_in_hdc_storage(self):
        for row in self.adder.rows:
            self.assertIn(row.cell, self.adder.storage)
            self.assertNotEqual(self.adder.storage.read(row.cell), 0)

    def test_keys_are_not_parity(self):
        """Regression: a shared atom pair reduces every key to its parity."""
        keys = [row.key for row in self.adder.rows]
        self.assertEqual(len(set(keys)), 8, "truth-table keys collided")
        worst = min(
            sim(a, b, D) for i, a in enumerate(keys) for j, b in enumerate(keys) if i != j
        )
        self.assertLess(worst, 0.15, "row keys are not near-orthogonal")

    def test_full_adder_all_eight_inputs(self):
        for a in (0, 1):
            for b in (0, 1):
                for c in (0, 1):
                    total, carry, confidence = self.adder.full_adder(a, b, c)
                    self.assertEqual(total, a ^ b ^ c, (a, b, c))
                    self.assertEqual(carry, 1 if a + b + c >= 2 else 0, (a, b, c))
                    self.assertGreater(confidence, 0.3)

    def test_separation_reports_real_numbers(self):
        report = self.adder.separation()
        self.assertEqual(report["rows"], 8)
        self.assertEqual(report["table_cells"], 8)
        self.assertLess(report["worst_cross_similarity"], 0.15)


class TestRipple(unittest.TestCase):
    def setUp(self):
        self.adder = HyperAdder(d=D)

    def test_single_bit_carry_chain(self):
        bits, stages = self.adder.add_bits([1, 1], [1, 0], 0)
        # 1 + 3 = 4 -> 100
        self.assertEqual(bits[:3], [0, 0, 1])
        self.assertEqual([s.carry_out for s in stages[:3]], [1, 1, 0])

    def test_carry_propagates_across_many_bits(self):
        bits, _ = self.adder.add_bits(hv_bits_of_int(0x7FFFFFFF), hv_bits_of_int(1), 0)
        self.assertEqual(hv_to_int(hv_from_bits(bits, D), D), wrap(0x7FFFFFFF + 1))

    def test_add_matches_python(self):
        rng = random.Random(7)
        for _ in range(25):
            a = rng.randrange(-(1 << 31), 1 << 31)
            b = rng.randrange(-(1 << 31), 1 << 31)
            self.assertEqual(self.adder.decode(self.adder.add(a, b)), wrap(a + b), (a, b))

    def test_sub_matches_python(self):
        rng = random.Random(8)
        for _ in range(25):
            a = rng.randrange(-(1 << 31), 1 << 31)
            b = rng.randrange(-(1 << 31), 1 << 31)
            self.assertEqual(self.adder.decode(self.adder.sub(a, b)), wrap(a - b), (a, b))

    def test_known_values(self):
        cases = [(0, 0, 0), (1, 1, 2), (7, 35, 42), (-7, 35, 28), (5, -11, -6),
                 (-1, 1, 0), (-2, -3, -5), (65535, 1, 65536)]
        for a, b, expected in cases:
            self.assertEqual(self.adder.decode(self.adder.add(a, b)), expected, (a, b))

    def test_hv_to_hv_layer(self):
        a, b = 1234, -5678
        out = self.adder.add_hv(int_to_hv(a, D), int_to_hv(b, D))
        self.assertEqual(self.adder.decode(out), wrap(a + b))
        self.assertEqual(out, self.adder.add(a, b), "HV path differs from int path")

    def test_sub_hv_layer(self):
        a, b = 99, 1234
        out = self.adder.sub_hv(int_to_hv(a, D), int_to_hv(b, D))
        self.assertEqual(self.adder.decode(out), wrap(a - b))

    def test_bit_extraction_matches_codec(self):
        for n in (0, 1, -1, 42, -42, 0x55555555, -0x55555556, (1 << 31) - 1):
            self.assertEqual(
                hv_to_bits(int_to_hv(n, D), D), hv_bits_of_int(n), "value %d" % n
            )


class TestNoDecode(unittest.TestCase):
    """Prove ADD/SUB execute without materialising an integer operand."""

    def test_add_and_sub_without_hv_to_int(self):
        m = build(
            """
            ldi R0, 7
            ldi R1, 35
            add R2, R0, R1
            sub R3, R0, R1
            halt
            """
        )
        with mock.patch("hdc.interpreter.hv_to_int", side_effect=AssertionError("decoded")):
            result = m.run()
        self.assertEqual(result["status"], "halted")
        self.assertEqual(m.adder.decode(m.regs[2]), 42)
        self.assertEqual(m.adder.decode(m.regs[3]), -28)

    def test_trace_reports_stage_count(self):
        m = build("ldi R0, 2\n ldi R1, 3\n add R2, R0, R1\n halt")
        m.run()
        notes = [e.note for e in m.trace if e.op == "ADD"]
        self.assertEqual(len(notes), 1)
        self.assertIn("stages", notes[0])
        self.assertIn("radix 8", notes[0])
        self.assertEqual(len(m.adder.last_stages), 4)


class TestMachineIntegration(unittest.TestCase):
    def test_program_arithmetic_through_circuit(self):
        m = build(
            """
            ldi R0, 21
            ldi R1, 2
            mul R2, R0, R1
            add R3, R2, R0
            sub R4, R3, R1
            out R3
            out R4
            halt
            """
        )
        m.run()
        self.assertEqual(m.output, ["63", "61"])

    def test_adder_stages_are_counted(self):
        m = build("ldi R0, 1\n ldi R1, 2\n add R2, R0, R1\n sub R3, R1, R0\n halt")
        m.run()
        per_op = (m.adder.width + m.adder.radix - 1) // m.adder.radix
        self.assertEqual(m.adder.stages_executed, 2 * per_op)

    def test_wraparound_is_hardware_like(self):
        m = build("ldi R0, 2147483647\n ldi R1, 1\n add R2, R0, R1\n halt")
        m.run()
        self.assertEqual(m.adder.decode(m.regs[2]), -(1 << 31))


class TestSliceGates(unittest.TestCase):
    """Wide gates: fewer stages, bigger tables, lazy materialisation."""

    def test_stage_counts_per_radix(self):
        for radix, stages in ((1, 32), (2, 16), (4, 8), (8, 4)):
            adder = HyperAdder(d=D, radix=radix)
            bits, _ = adder.add_bits(hv_bits_of_int(1), hv_bits_of_int(1), 0)
            self.assertEqual(len(adder.last_stages), stages, "radix %d" % radix)
            self.assertEqual(adder.decode(hv_from_bits(bits, D)), 2)

    def test_radix8_arithmetic(self):
        adder = HyperAdder(d=D, radix=8)
        report = adder.verify(20)
        self.assertEqual(report["add_correct"], 20)
        self.assertEqual(report["sub_correct"], 20)
        self.assertEqual(report["hv_to_hv_correct"], 20)

    def test_all_radices_agree(self):
        cases = [(7, 35), (-7, 35), (-1, 1), ((1 << 31) - 1, 1), (-12345, -54321)]
        for radix in (1, 2, 4, 8):
            adder = HyperAdder(d=D, radix=radix)
            for a, b in cases:
                self.assertEqual(
                    adder.decode(adder.add(a, b)), wrap(a + b), (radix, a, b)
                )

    def test_byte_table_is_lazy(self):
        adder = HyperAdder(d=D, radix=8)
        self.assertEqual(len(adder.gate.rows), 0)
        adder.add(1, 1)
        # 131 072 rows exist in principle; four additions touch only a few.
        self.assertEqual(adder.gate.row_count, 1 << 17)
        self.assertLess(len(adder.gate.rows), 16)
        self.assertFalse(adder.gate.eager)

    def test_rows_materialise_into_storage(self):
        adder = HyperAdder(d=D, radix=8)
        adder.add(300, 400)
        self.assertEqual(len(adder.gate.storage), len(adder.gate.rows))
        for row in adder.gate.rows.values():
            self.assertIn(row.cell, adder.gate.storage)

    def test_gate_row_reads_back_as_built(self):
        # evaluate() asserts internally that the decoded fields equal the row;
        # this drives every combination through it for a 2-bit gate.
        adder = HyperAdder(d=D, radix=2)
        gate = adder.gate
        for c in (0, 1):
            for b in range(4):
                for a in range(4):
                    total, carry, confidence = gate.evaluate(a, b, c)
                    self.assertEqual((total, carry), ((a + b + c) & 3, (a + b + c) >> 2))
                    self.assertGreater(confidence, 0.1)

    def test_similarity_fallback_resolves_nearest_row(self):
        adder = HyperAdder(d=D, radix=8)
        adder.add(1, 1)
        row, confidence = adder.gate.resolve_similar(0xFF, 0xFF, 1)
        self.assertIsNotNone(row)
        self.assertLessEqual(confidence, 1.0)

    def test_machine_default_is_four_stages(self):
        m = build("ldi R0, 1\n ldi R1, 2\n add R2, R0, R1\n halt")
        self.assertEqual(m.adder.radix, 8)
        m.run()
        self.assertEqual(len(m.adder.last_stages), 4)
        self.assertEqual(m.adder.decode(m.regs[2]), 3)


class TestCodecIsTwosComplement(unittest.TestCase):
    def test_store_slots(self):
        self.assertEqual(MAX_INT_BITS, 32)

    def test_negative_values_are_dense(self):
        # -1 sets every bit, which is why the codec threshold had to be tuned
        # for 32-atom bundles rather than the old 8-member assumption.
        self.assertEqual(sum(hv_bits_of_int(-1)), 32)

    def test_range_enforced(self):
        with self.assertRaises(ValueError):
            int_to_hv(1 << 31, D)
        with self.assertRaises(ValueError):
            int_to_hv(-(1 << 31) - 1, D)


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