#!/usr/bin/env python3
"""
test_dna_lut.py — Tests for symbology LUT adaptation
"""

import os
import sys
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python"))

from dna_lut import (
    LUT,
    BASES,
    BASE_TO_INDEX,
    INDEX_TO_BASE,
    build_direct_lut,
    build_monotone_lut,
    build_positional_lut,
    demo_qubo,
    dna_to_int,
    int_to_dna,
    qubo_energy,
    sequence_length_for_solutions,
    verify_all_encodings,
)


class TestIntDnaRoundtrip(unittest.TestCase):
    """§1: Integer ↔ DNA conversion."""

    def test_roundtrip(self):
        for val in [0, 1, 42, 255, 511]:
            seq = int_to_dna(val, 3)
            back = dna_to_int(seq)
            self.assertEqual(back, val, f"Failed for {val}: {seq}")

    def test_length(self):
        """Sequence length is correct."""
        self.assertEqual(len(int_to_dna(0, 5)), 5)
        self.assertEqual(len(int_to_dna(999, 4)), 4)

    def test_ordering(self):
        """Lexicographic order = integer order with ABCGPSTZ bases."""
        seqs = [int_to_dna(i, 3) for i in range(100)]
        self.assertEqual(seqs, sorted(seqs))

    def test_sequence_length_for_solutions(self):
        """Minimum length calculation."""
        self.assertEqual(sequence_length_for_solutions(1), 1)
        self.assertEqual(sequence_length_for_solutions(8), 1)
        self.assertEqual(sequence_length_for_solutions(9), 2)
        self.assertEqual(sequence_length_for_solutions(64), 2)
        self.assertEqual(sequence_length_for_solutions(65), 3)


class TestQuboEnergy(unittest.TestCase):
    """§2: QUBO energy computation."""

    def test_identity(self):
        Q = [[1.0, 0], [0, 2.0]]
        self.assertAlmostEqual(qubo_energy([0, 0], Q), 0.0)
        self.assertAlmostEqual(qubo_energy([1, 0], Q), 1.0)
        self.assertAlmostEqual(qubo_energy([0, 1], Q), 2.0)
        self.assertAlmostEqual(qubo_energy([1, 1], Q), 3.0)

    def test_negative(self):
        Q = [[-1.0, -0.5], [-0.5, -1.0]]
        self.assertAlmostEqual(qubo_energy([0, 0], Q), 0.0)
        self.assertAlmostEqual(qubo_energy([1, 1], Q), -3.0)


class TestDirectLUT(unittest.TestCase):
    """§3: Direct encoding LUT."""

    def test_size(self):
        Q = [[1.0, 0], [0, 2.0]]
        lut = build_direct_lut(Q, 2)
        self.assertEqual(lut.size(), 4)

    def test_lookup(self):
        Q = [[3.0, 0], [0, 2.0]]
        lut = build_direct_lut(Q, 2)
        # x=[0,0] → seq="AA", energy=0
        x, e = lut.lookup("AA")
        self.assertEqual(x, [0, 0])
        self.assertAlmostEqual(e, 0.0)

    def test_not_monotone(self):
        """Direct encoding is generally not monotone."""
        Q = [[3.0, 0, 0], [0, 2.0, 0], [0, 0, 1.0]]
        lut = build_direct_lut(Q, 3)
        mono, _ = lut.verify_monotone()
        # For this specific Q, "AAA" (E=0) is first in both orderings
        # but the rest won't match
        self.assertIsInstance(mono, bool)


class TestMonotoneLUT(unittest.TestCase):
    """§4: Monotone encoding — THE FIX."""

    def test_is_monotone(self):
        """Monotone encoding: lexicographic sort = energy sort."""
        Q = [[3.0, 0, 0], [0, 2.0, 0], [0, 0, 1.0]]
        lut = build_monotone_lut(Q, 3)
        mono, corr = lut.verify_monotone()
        self.assertTrue(mono, "Monotone encoding should be monotone")
        self.assertAlmostEqual(corr, 1.0)

    def test_monotone_ising(self):
        """Monotone encoding works for Ising chain (negative Q)."""
        Q = [[-1.0, -0.5, 0], [-0.5, -1.0, -0.5], [0, -0.5, -1.0]]
        lut = build_monotone_lut(Q, 3)
        mono, corr = lut.verify_monotone()
        self.assertTrue(mono)
        self.assertAlmostEqual(corr, 1.0)

    def test_monotone_random(self):
        """Monotone encoding works for random QUBO."""
        Q = demo_qubo(6, seed=42, style="random")
        lut = build_monotone_lut(Q, 6)
        mono, corr = lut.verify_monotone()
        self.assertTrue(mono)
        self.assertAlmostEqual(corr, 1.0)

    def test_first_lex_is_min_energy(self):
        """First sequence in lex order has minimum energy."""
        Q = demo_qubo(6, seed=42, style="ising")
        lut = build_monotone_lut(Q, 6)
        by_seq = lut.sort_by_sequence()
        by_energy = lut.sort_by_energy()
        self.assertEqual(by_seq[0][0], by_energy[0][0])
        self.assertAlmostEqual(by_seq[0][2], by_energy[0][2])

    def test_all_energies_sorted(self):
        """All energies are in ascending order by sequence."""
        Q = demo_qubo(8, seed=42, style="banded")
        lut = build_monotone_lut(Q, 8)
        by_seq = lut.sort_by_sequence()
        energies = [e for _, _, e in by_seq]
        for i in range(len(energies) - 1):
            self.assertLessEqual(energies[i], energies[i + 1])

    def test_json_roundtrip(self):
        """Monotone LUT survives JSON serialization."""
        Q = [[1.0, 0], [0, 2.0]]
        lut = build_monotone_lut(Q, 2)
        j = lut.to_json()
        lut2 = LUT.from_json(j)
        self.assertEqual(lut.size(), lut2.size())
        for s in lut.entries:
            self.assertAlmostEqual(lut.energy(s), lut2.energy(s))


class TestPositionalLUT(unittest.TestCase):
    """§5: Position-dependent encoding."""

    def test_encode_decode(self):
        """Positional encoding roundtrips."""
        Q = [[3.0, 0, 0], [0, -2.0, 0], [0, 0, 1.0]]
        lut = build_positional_lut(Q, 3)
        for s, (x, e) in lut.entries.items():
            # Verify the sequence decodes back to the solution
            # (we can't easily decode positional without the map,
            #  but we can verify the entry exists)
            self.assertIsInstance(x, list)
            self.assertIsInstance(e, float)

    def test_ising_sign_fix(self):
        """Positional encoding fixes sign inversion for Ising."""
        Q = [[-1.0, -0.5, 0], [-0.5, -1.0, -0.5], [0, -0.5, -1.0]]

        lut_direct = build_direct_lut(Q, 3)
        _, corr_direct = lut_direct.verify_monotone()

        lut_pos = build_positional_lut(Q, 3)
        _, corr_pos = lut_pos.verify_monotone()

        # Positional should have better correlation than direct
        self.assertGreater(corr_pos, corr_direct - 0.1)


class TestVerifyAll(unittest.TestCase):
    """§6: Comparative verification."""

    def test_monotone_always_perfect(self):
        """Monotone encoding always has correlation = 1.0."""
        for style in ["diagonal", "banded", "ising", "random"]:
            Q = demo_qubo(6, seed=42, style=style)
            results = verify_all_encodings(Q, 6)
            self.assertTrue(
                results["monotone"]["is_monotone"],
                f"Monotone failed for {style}"
            )
            self.assertAlmostEqual(
                results["monotone"]["rank_correlation"], 1.0,
                msg=f"Monotone correlation not 1.0 for {style}"
            )

    def test_direct_not_always_monotone(self):
        """Direct encoding is not always monotone."""
        Q = demo_qubo(6, seed=42, style="ising")
        results = verify_all_encodings(Q, 6)
        self.assertFalse(
            results["direct"]["is_monotone"],
            "Direct encoding should NOT be monotone for Ising"
        )


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