#!/usr/bin/env python3
"""
test_dna_codec.py — Tests for the Hachimoji DNA encoding/decoding layer

Test categories:
    1. Base encoding roundtrip
    2. Binary vector encoding
    3. QUBO energy computation
    4. Tm sort verification (the critical test)
    5. SAM/FASTA export
"""

import json
import math
import os
import random
import sys
import unittest

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

from dna_codec import (
    BASE_TO_BITS,
    BITS_TO_BASE,
    HACHIMOJI_BASES,
    decode_binary_vector,
    decode_dna_to_bytes,
    dna_to_greek,
    encode_binary_vector,
    encode_bytes_to_dna,
    gc_content,
    greek_to_dna,
    melting_temperature,
    qubo_energy,
    raw_gc_content,
    sequence_stats,
)

from dna_qubo_sort import (
    QuboDnaCandidate,
    candidates_to_fasta,
    candidates_to_sam,
    demo_max_cut,
    demo_number_partition,
    encode_qubo_candidate,
    generate_all_candidates,
    generate_random_candidates,
    run_qubo_dna_sort,
    sort_by_energy,
    sort_by_gc,
    sort_by_tm,
    spearman_rank_correlation,
    verify_tm_sort,
)


class TestBaseEncoding(unittest.TestCase):
    """§1: Base encoding roundtrip."""

    def test_all_bases_present(self):
        """All 8 Hachimoji bases are defined."""
        self.assertEqual(len(HACHIMOJI_BASES), 8)
        self.assertEqual(len(BITS_TO_BASE), 8)
        self.assertEqual(len(BASE_TO_BITS), 8)

    def test_bijection(self):
        """BITS_TO_BASE and BASE_TO_BITS are inverses."""
        for bits, base in BITS_TO_BASE.items():
            self.assertEqual(BASE_TO_BITS[base], bits)
        for base, bits in BASE_TO_BITS.items():
            self.assertEqual(BITS_TO_BASE[bits], base)

    def test_3bit_coverage(self):
        """All 3-bit values (0-7) map to a base."""
        for i in range(8):
            self.assertIn(i, BITS_TO_BASE)
            self.assertEqual(len(BITS_TO_BASE[i]), 1)  # single character

    def test_greek_roundtrip(self):
        """Latin → Greek → Latin roundtrip."""
        for base in HACHIMOJI_BASES:
            greek = dna_to_greek(base)
            back = greek_to_dna(greek)
            self.assertEqual(back, base, f"Failed for {base} → {greek} → {back}")


class TestByteEncoding(unittest.TestCase):
    """§2: Byte → DNA encoding roundtrip."""

    def test_single_byte(self):
        """Single byte encodes and decodes."""
        for b in range(256):
            data = bytes([b])
            dna = encode_bytes_to_dna(data)
            decoded = decode_dna_to_bytes(dna)
            self.assertEqual(decoded[:1], data, f"Failed for byte {b}")

    def test_empty(self):
        """Empty input produces empty output."""
        dna = encode_bytes_to_dna(b"")
        self.assertEqual(dna, "")
        decoded = decode_dna_to_bytes("")
        self.assertEqual(decoded, b"")

    def test_known_sequence(self):
        """Known byte → DNA mapping (ABCGPSTZ ordering, ASCII-sorted)."""
        # 0x00 = 00000000 → 000 000 00(pad 0) → A A A (3 bases)
        dna = encode_bytes_to_dna(b"\x00")
        self.assertEqual(dna, "AAA")

        # 0xFF = 11111111 → 111 111 11(pad 0) → Z Z T (3 bases)
        # With ABCGPSTZ: 111=7→Z, 111=7→Z, 110=6→T
        dna = encode_bytes_to_dna(b"\xff")
        self.assertEqual(dna[0], "Z")  # 111 = 7
        self.assertEqual(dna[1], "Z")  # 111 = 7
        self.assertEqual(dna[2], "T")  # 110 = 6 (padded with 0)

        # 0b11010001 = 110 100 01(pad 0) = 110 100 010 → T P C
        # With ABCGPSTZ: 110=6→T, 100=4→P, 010=2→C
        dna = encode_bytes_to_dna(b"\xd1")
        self.assertEqual(dna[0], "T")  # 110 = 6
        self.assertEqual(dna[1], "P")  # 100 = 4
        self.assertEqual(dna[2], "C")  # 010 = 2

    def test_random_roundtrip(self):
        """Random bytes roundtrip through DNA encoding."""
        rng = random.Random(42)
        for _ in range(100):
            length = rng.randint(1, 50)
            data = bytes(rng.randint(0, 255) for _ in range(length))
            dna = encode_bytes_to_dna(data)
            decoded = decode_dna_to_bytes(dna)
            # May have padding zeros at the end
            self.assertEqual(decoded[:length], data)

    def test_invalid_base_raises(self):
        """Invalid base character raises ValueError."""
        with self.assertRaises(ValueError):
            decode_dna_to_bytes("AXG")


class TestBinaryVectorEncoding(unittest.TestCase):
    """§3: Binary vector encoding."""

    def test_all_zeros(self):
        """All-zero vector → all A's (3 bases per var by default)."""
        x = [0, 0, 0, 0]
        seq = encode_binary_vector(x)
        self.assertEqual(seq, "A" * 12)  # 4 vars × 3 bases = 12

    def test_all_ones(self):
        """All-one vector → all P's (3 bases per var by default)."""
        x = [1, 1, 1, 1]
        seq = encode_binary_vector(x)
        self.assertEqual(seq, "P" * 12)  # 4 vars × 3 bases = 12

    def test_mixed(self):
        """Mixed vector encodes correctly (3 bases per var)."""
        x = [0, 1, 0, 1]
        seq = encode_binary_vector(x)
        self.assertEqual(seq, "AAAPPPAAAPPP")  # 0→AAA, 1→PPP, 0→AAA, 1→PPP

    def test_single_base_per_var(self):
        """bases_per_var=1 gives one base per variable."""
        x = [0, 1, 0, 1]
        seq = encode_binary_vector(x, bases_per_var=1)
        self.assertEqual(seq, "APAP")

    def test_roundtrip(self):
        """Binary vector roundtrip."""
        for _ in range(50):
            rng = random.Random(42)
            x = [rng.randint(0, 1) for _ in range(20)]
            seq = encode_binary_vector(x)
            decoded = decode_binary_vector(seq)
            self.assertEqual(decoded, x)


class TestMeltingTemperature(unittest.TestCase):
    """§4: Melting temperature model."""

    def test_empty(self):
        """Empty sequence has Tm = 0."""
        self.assertEqual(melting_temperature(""), 0.0)

    def test_all_a_low_tm(self):
        """All-A sequence has low Tm."""
        tm_a = melting_temperature("A" * 10)
        tm_g = melting_temperature("G" * 10)
        self.assertLess(tm_a, tm_g)

    def test_gc_richer_higher_tm(self):
        """GC-rich sequences have higher Tm than AT-rich."""
        at_rich = "AATT" * 10
        gc_rich = "GGCC" * 10
        self.assertLess(melting_temperature(at_rich), melting_temperature(gc_rich))

    def test_synthetic_bases_intermediate(self):
        """Synthetic bases have intermediate Tm values."""
        tm_a = TM_CONTRIBUTION["A"]
        tm_g = TM_CONTRIBUTION["G"]
        tm_p = TM_CONTRIBUTION["P"]
        self.assertLess(tm_a, tm_p)
        self.assertLess(tm_p, tm_g + 1)  # P is between A and G

    def test_gc_content_range(self):
        """GC content is in [0, 1]."""
        for seq in ["AAAA", "GGGG", "ATGC", "PPZZ", ""]:
            gc = gc_content(seq)
            self.assertGreaterEqual(gc, 0.0)
            self.assertLessEqual(gc, 1.0)

    def test_tm_monotone_with_ones(self):
        """Tm is monotonically related to number of 1s in binary vector."""
        # For encode_binary_vector: 0→A (Tm=2), 1→P (Tm=3.5)
        # More 1s = higher Tm
        for n_ones in range(10):
            x = [1] * n_ones + [0] * (10 - n_ones)
            seq = encode_binary_vector(x)
            tm = melting_temperature(seq)
            if n_ones > 0:
                x_prev = [1] * (n_ones - 1) + [0] * (11 - n_ones)
                seq_prev = encode_binary_vector(x_prev)
                tm_prev = melting_temperature(seq_prev)
                self.assertGreaterEqual(tm, tm_prev)


# Need TM_CONTRIBUTION for the test
from dna_codec import TM_CONTRIBUTION


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

    def test_zero_matrix(self):
        """Zero matrix → zero energy for any x."""
        Q = [[0.0] * 4 for _ in range(4)]
        self.assertEqual(qubo_energy([1, 0, 1, 0], Q), 0.0)

    def test_identity_matrix(self):
        """Identity matrix → energy = number of 1s."""
        n = 5
        Q = [[1.0 if i == j else 0.0 for j in range(n)] for i in range(n)]
        for k in range(n + 1):
            x = [1] * k + [0] * (n - k)
            self.assertAlmostEqual(qubo_energy(x, Q), float(k))

    def test_known_qubo(self):
        """Known QUBO with hand-computed energy."""
        # Q = [[1, -1], [-1, 1]]
        # x = [0, 0] → E = 0
        # x = [1, 0] → E = 1
        # x = [0, 1] → E = 1
        # x = [1, 1] → E = 0
        Q = [[1.0, -1.0], [-1.0, 1.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), 1.0)
        self.assertAlmostEqual(qubo_energy([1, 1], Q), 0.0)


class TestTmSortVerification(unittest.TestCase):
    """§6: THE CRITICAL TEST — does Tm sorting = energy sorting?"""

    def test_diagonal_qubo_perfect_correlation(self):
        """For diagonal Q with positive entries, Tm sort = energy sort exactly.

        This is the ideal case: E(x) = Σ Q_ii * x_i, and Tm ∝ Σ x_i.
        When all Q_ii are equal, Tm sort is identical to energy sort.
        """
        n = 6
        Q = [[3.0 if i == j else 0.0 for j in range(n)] for i in range(n)]
        result = verify_tm_sort(Q, n)
        self.assertTrue(result.tm_sort_matches_energy)
        self.assertGreater(result.rank_correlation, 0.99)

    def test_uniform_qubo_rank_correlation(self):
        """For uniform Q, rank correlation should be very high."""
        n = 8
        Q = [[2.0 if i == j else 0.5 for j in range(n)] for i in range(n)]
        result = verify_tm_sort(Q, n)
        self.assertGreater(result.rank_correlation, 0.8)

    def test_max_cut_sort(self):
        """Max-Cut QUBO: Tm sort should find a low-energy solution."""
        Q = demo_max_cut(6, seed=42)
        result = verify_tm_sort(Q, 6)
        # Tm-optimal should be in the top-k by energy
        self.assertGreater(result.top_k_agreement, 0)

    def test_negative_entries_degrade_correlation(self):
        """Negative Q entries break the Tm-energy correlation.

        This is expected: Tm encoding assumes non-negative contribution per 1.
        Negative Q entries mean more 1s can LOWER energy, inverting the sort.
        """
        n = 4
        Q = [[-3.0 if i == j else 0.0 for j in range(n)] for i in range(n)]
        result = verify_tm_sort(Q, n)
        # Correlation should be negative (inverted!)
        self.assertLess(result.rank_correlation, 0.5)


class TestSamFastaExport(unittest.TestCase):
    """§7: SAM and FASTA export."""

    def test_sam_format(self):
        """SAM output is valid-ish format."""
        x = [1, 0, 1]
        Q = [[1.0, 0, 0], [0, 1.0, 0], [0, 0, 1.0]]
        cand = encode_qubo_candidate(x, Q)
        sam = candidates_to_sam([cand], "test")

        lines = sam.strip().split("\n")
        # Header lines start with @
        self.assertTrue(lines[0].startswith("@HD"))
        self.assertTrue(lines[1].startswith("@PG"))
        # Data line (3rd line after @CO comment)
        data_line = [l for l in lines if not l.startswith("@")][0]
        self.assertIn("TM:f:", data_line)
        self.assertIn("EN:f:", data_line)

    def test_fasta_format(self):
        """FASTA output starts with >."""
        x = [1, 0, 1]
        Q = [[1.0, 0, 0], [0, 1.0, 0], [0, 0, 1.0]]
        cand = encode_qubo_candidate(x, Q)
        fasta = candidates_to_fasta([cand])

        lines = fasta.strip().split("\n")
        self.assertTrue(lines[0].startswith(">"))
        # Second line is sequence
        self.assertTrue(all(c in HACHIMOJI_BASES for c in lines[1]))


class TestSequenceStats(unittest.TestCase):
    """§8: Sequence statistics."""

    def test_stats_keys(self):
        """sequence_stats returns expected keys."""
        stats = sequence_stats("ATGCBSPZ")
        self.assertIn("length", stats)
        self.assertIn("gc_content", stats)
        self.assertIn("melting_temperature", stats)
        self.assertIn("base_counts", stats)
        self.assertIn("greek", stats)

    def test_stats_length(self):
        """Length is correct."""
        self.assertEqual(sequence_stats("ATGC")["length"], 4)
        self.assertEqual(sequence_stats("")["length"], 0)


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