#!/usr/bin/env python3
"""
test_dna_nn.py — Tests for nearest-neighbor QUBO-DNA encoding

Tests the critical property: NN Tm sort = energy sort for banded QUBOs.
"""

import math
import os
import random
import sys
import unittest

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

from dna_qubo_nn import (
    NNQuboCandidate,
    apply_permutation,
    base_to_value,
    choose_base,
    classify_qubo,
    decode_nn_sequence,
    demo_banded_qubo,
    demo_full_qubo,
    demo_ising_chain,
    encode_nn_candidate,
    generate_nn_candidates,
    inverse_permutation,
    nn_stack_dg,
    nn_tm_estimate,
    reorder_qubo_for_bandwidth,
    sort_by_energy,
    sort_by_tm_proxy,
    spearman_rank_correlation,
    verify_nn_sort,
)


class TestBaseChoice(unittest.TestCase):
    """§1: Base selection encodes value + diagonal sign."""

    def test_zero_positive_diagonal(self):
        """x=0, Q_ii >= 0 → A"""
        self.assertEqual(choose_base(0, 3.0), "A")
        self.assertEqual(choose_base(0, 0.0), "A")

    def test_zero_negative_diagonal(self):
        """x=0, Q_ii < 0 → T"""
        self.assertEqual(choose_base(0, -1.0), "T")

    def test_one_positive_diagonal(self):
        """x=1, Q_ii >= 0 → G"""
        self.assertEqual(choose_base(1, 3.0), "G")

    def test_one_negative_diagonal(self):
        """x=1, Q_ii < 0 → C"""
        self.assertEqual(choose_base(1, -1.0), "C")

    def test_base_to_value(self):
        """Low-Tm bases → 0, high-Tm bases → 1."""
        self.assertEqual(base_to_value("A"), 0)
        self.assertEqual(base_to_value("T"), 0)
        self.assertEqual(base_to_value("G"), 1)
        self.assertEqual(base_to_value("C"), 1)


class TestNNThermodynamics(unittest.TestCase):
    """§2: Nearest-neighbor thermodynamic model."""

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

    def test_stacking_energies_finite(self):
        """All dinucleotide stacking energies are finite."""
        for b1 in "ATGC":
            for b2 in "ATGC":
                dg = nn_stack_dg(b1, b2)
                self.assertTrue(math.isfinite(dg))

    def test_gc_gc_most_stable(self):
        """GC/GC stacking is among the most stable."""
        gc_gc = nn_stack_dg("G", "C")
        at_at = nn_stack_dg("A", "T")
        self.assertLess(gc_gc, at_at)  # more negative = more stable

    def test_empty_sequence(self):
        """Empty sequence has zero Tm."""
        self.assertEqual(nn_tm_estimate(""), 0.0)

    def test_single_base(self):
        """Single base has per-base Tm only."""
        tm_g = nn_tm_estimate("G")
        tm_a = nn_tm_estimate("A")
        self.assertGreater(tm_g, tm_a)


class TestEncodeDecode(unittest.TestCase):
    """§3: NN encoding and decoding."""

    def test_roundtrip(self):
        """Encode → decode roundtrip preserves binary vector."""
        Q = [[3.0, 0, 0], [0, 2.0, 0], [0, 0, 1.0]]
        for x in [[0, 0, 0], [1, 1, 1], [0, 1, 0], [1, 0, 1]]:
            cand = encode_nn_candidate(x, Q)
            decoded = decode_nn_sequence(cand.sequence)
            self.assertEqual(decoded, x, f"Failed for x={x}")

    def test_all_zeros(self):
        """All-zero vector → all low-Tm bases."""
        n = 5
        Q = [[2.0] * n for _ in range(n)]
        x = [0] * n
        cand = encode_nn_candidate(x, Q)
        for base in cand.sequence:
            self.assertIn(base, "AT")

    def test_all_ones(self):
        """All-one vector → all high-Tm bases."""
        n = 5
        Q = [[2.0] * n for _ in range(n)]
        x = [1] * n
        cand = encode_nn_candidate(x, Q)
        for base in cand.sequence:
            self.assertIn(base, "GC")

    def test_negative_diagonal_uses_different_bases(self):
        """Negative Q_ii uses T/C instead of A/G."""
        Q_pos = [[3.0]]
        Q_neg = [[-3.0]]
        self.assertEqual(encode_nn_candidate([0], Q_pos).sequence, "A")
        self.assertEqual(encode_nn_candidate([0], Q_neg).sequence, "T")
        self.assertEqual(encode_nn_candidate([1], Q_pos).sequence, "G")
        self.assertEqual(encode_nn_candidate([1], Q_neg).sequence, "C")


class TestTmSortBanded(unittest.TestCase):
    """§4: THE CRITICAL TEST — banded QUBOs."""

    def test_diagonal_perfect(self):
        """Diagonal QUBO: Tm sort = energy sort exactly."""
        n = 6
        Q = [[3.0 if i == j else 0.0 for j in range(n)] for i in range(n)]
        result = verify_nn_sort(Q, n)
        self.assertTrue(result.tm_sort_matches_energy)
        self.assertGreater(result.rank_correlation, 0.9)

    def test_banded_high_correlation(self):
        """Banded QUBO: Tm sort has high rank correlation."""
        Q = demo_banded_qubo(8, seed=42)
        result = verify_nn_sort(Q, 8)
        self.assertGreater(result.rank_correlation, 0.35)
        self.assertGreater(result.top_k_agreement, 0)

    def test_ising_chain_perfect(self):
        """1D Ising chain: NN encoding captures all interactions."""
        Q = demo_ising_chain(8, h=1.0, J=0.5)
        result = verify_nn_sort(Q, 8)
        # All interactions are adjacent → NN captures everything
        self.assertGreater(result.rank_correlation, 0.9)

    def test_ising_chain_optimal_match(self):
        """Ising chain: Tm-optimal is energy-optimal."""
        Q = demo_ising_chain(6, h=1.0, J=0.5)
        result = verify_nn_sort(Q, 6)
        self.assertTrue(result.tm_sort_matches_energy)


class TestTmSortFull(unittest.TestCase):
    """§5: Full QUBO — tests the limits."""

    def test_full_qubo_degraded(self):
        """Full QUBO: NN encoding loses non-adjacent interactions."""
        Q = demo_full_qubo(8, seed=42)
        result = verify_nn_sort(Q, 8)
        # Should be worse than banded but still positive
        self.assertGreater(result.rank_correlation, 0.0)

    def test_full_qubo_still_finds_decent(self):
        """Full QUBO: Tm-optimal is at least in the top-k by energy."""
        Q = demo_full_qubo(8, seed=42)
        result = verify_nn_sort(Q, 8, top_k=10)
        self.assertGreater(result.top_k_agreement, 0)


class TestBandwidthReordering(unittest.TestCase):
    """§6: QUBO reordering for bandwidth optimization."""

    def test_identity_permutation(self):
        """Already-banded QUBO keeps similar ordering."""
        Q = demo_banded_qubo(6, seed=42)
        Q_reordered, perm = reorder_qubo_for_bandwidth(Q, 6)
        # Should be similar (not necessarily identical)
        self.assertEqual(len(perm), 6)

    def test_permutation_inverse(self):
        """Permutation and inverse are correct."""
        perm = [2, 0, 1, 3]
        inv = inverse_permutation(perm)
        for i in range(4):
            self.assertEqual(perm[inv[i]], i)

    def test_apply_permutation(self):
        """Apply permutation correctly reorders vector."""
        x = [10, 20, 30, 40]
        perm = [2, 0, 3, 1]
        result = apply_permutation(x, perm)
        self.assertEqual(result, [30, 10, 40, 20])

    def test_reordering_preserves_energy(self):
        """Reordering doesn't change QUBO energies."""
        Q = demo_full_qubo(6, seed=42)
        Q_reord, perm = reorder_qubo_for_bandwidth(Q, 6)
        inv = inverse_permutation(perm)

        # Energy of a solution should be the same
        x_orig = [1, 0, 1, 0, 1, 0]
        x_reord = apply_permutation(x_orig, perm)

        e_orig = sum(Q[i][j] * x_orig[i] * x_orig[j] for i in range(6) for j in range(6))
        e_reord = sum(Q_reord[i][j] * x_reord[i] * x_reord[j] for i in range(6) for j in range(6))

        self.assertAlmostEqual(e_orig, e_reord, places=10)


class TestQuboClassification(unittest.TestCase):
    """§7: QUBO type classification."""

    def test_diagonal(self):
        """Diagonal matrix classified as 'diagonal'."""
        Q = [[3.0, 0], [0, 2.0]]
        self.assertEqual(classify_qubo(Q, 2), "diagonal")

    def test_banded(self):
        """Tridiagonal matrix classified as 'banded'."""
        Q = [[3.0, 1.0, 0], [1.0, 2.0, 0.5], [0, 0.5, 1.0]]
        self.assertEqual(classify_qubo(Q, 3), "banded")

    def test_full(self):
        """Full matrix classified as 'full'."""
        Q = [[3.0, 1.0, 0.5], [1.0, 2.0, 0.3], [0.5, 0.3, 1.0]]
        self.assertEqual(classify_qubo(Q, 3), "full")


class TestSpearman(unittest.TestCase):
    """§8: Rank correlation."""

    def test_perfect_correlation(self):
        """Identical lists → correlation = 1.0."""
        vals = [1.0, 2.0, 3.0, 4.0, 5.0]
        self.assertAlmostEqual(spearman_rank_correlation(vals, vals), 1.0)

    def test_inverse_correlation(self):
        """Reversed list → correlation = -1.0 (approximately)."""
        vals = [1.0, 2.0, 3.0, 4.0, 5.0]
        rev = list(reversed(vals))
        corr = spearman_rank_correlation(vals, rev)
        self.assertLess(corr, -0.9)

    def test_empty(self):
        """Empty lists → 0.0."""
        self.assertEqual(spearman_rank_correlation([], []), 0.0)


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