#!/usr/bin/env python3
"""
dna_lut.py — Symbology LUT Adaptation for QUBO-DNA Sorting

The LUT is the adapter between symbolic DNA space and QUBO energy space.

Three encoding strategies, in order of increasing power:

    1. Direct encoding: x → DNA by fixed mapping. Fast, but symbolic
       order doesn't match energy order. LUT does the actual sorting.

    2. Position-dependent encoding: x → DNA with per-position base
       selection based on Q_ii sign. Fixes sign inversion for Ising-type
       QUBOs but still not monotone.

    3. Monotone encoding (the fix): sort all solutions by energy FIRST,
       then assign DNA sequences in energy order. Lexicographic DNA
       sort = energy sort BY CONSTRUCTION.

    The monotone encoding makes the DNA sequence a rank key:
        sequence "AAAA..." = lowest energy
        sequence "AAAT..." = second lowest
        ...
        sequence "ZZZZ..." = highest energy

    Sorting the DNA IS sorting the energy. No thermodynamic model needed.
    The LUT maps sequence ↔ energy, and the sequence itself carries the
    rank information.

Connection to SilverSight:
    HachimojiLUT.lean defines codon → state mapping via lookup.
    This extends that to: sequence → energy via monotone LUT.
    The PhaseCircle (ℤ/360ℤ) provides the symbolic address space.
"""

from __future__ import annotations

import json
import random
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple


# ============================================================
# §1  HACHIMOJI ALPHABET
# ============================================================

BASES = list("ABCGPSTZ")  # 8 bases, ordered by ASCII for lexicographic monotonicity
BASE_TO_INDEX = {b: i for i, b in enumerate(BASES)}
INDEX_TO_BASE = {i: b for i, b in enumerate(BASES)}
N_BASES = len(BASES)


def int_to_dna(value: int, length: int) -> str:
    """Convert integer to DNA sequence (base-8, fixed length)."""
    seq = []
    for _ in range(length):
        seq.append(INDEX_TO_BASE[value % N_BASES])
        value //= N_BASES
    return "".join(reversed(seq))


def dna_to_int(sequence: str) -> int:
    """Convert DNA sequence to integer."""
    value = 0
    for b in sequence:
        value = value * N_BASES + BASE_TO_INDEX[b]
    return value


def sequence_length_for_solutions(n_solutions: int) -> int:
    """Minimum sequence length to represent n_solutions unique keys."""
    if n_solutions <= 0:
        return 0
    length = 1
    while N_BASES ** length < n_solutions:
        length += 1
    return length


# ============================================================
# §2  QUBO ENERGY
# ============================================================

def qubo_energy(x: List[int], Q: List[List[float]]) -> float:
    """Compute QUBO energy E(x) = x^T Q x."""
    n = len(x)
    return sum(Q[i][j] * x[i] * x[j] for i in range(n) for j in range(n))


# ============================================================
# §3  ENCODING STRATEGIES
# ============================================================

@dataclass
class LUT:
    """Lookup table: DNA sequence ↔ QUBO solution ↔ energy."""
    entries: Dict[str, Tuple[List[int], float]]  # sequence → (x, energy)
    qubo_matrix: List[List[float]]
    n_vars: int
    encoding: str  # "direct", "positional", "monotone"

    def lookup(self, sequence: str) -> Tuple[List[int], float]:
        """Look up solution and energy for a sequence."""
        return self.entries.get(sequence, ([], float("inf")))

    def energy(self, sequence: str) -> float:
        return self.lookup(sequence)[1]

    def solution(self, sequence: str) -> List[int]:
        return self.lookup(sequence)[0]

    def size(self) -> int:
        return len(self.entries)

    def sort_by_sequence(self) -> List[Tuple[str, List[int], float]]:
        """Sort entries by DNA sequence (lexicographic)."""
        return sorted(
            [(s, x, e) for s, (x, e) in self.entries.items()],
            key=lambda item: item[0],
        )

    def sort_by_energy(self) -> List[Tuple[str, List[int], float]]:
        """Sort entries by energy (ascending). Ground truth."""
        return sorted(
            [(s, x, e) for s, (x, e) in self.entries.items()],
            key=lambda item: item[2],
        )

    def verify_monotone(self) -> Tuple[bool, float]:
        """Check if lexicographic sort matches energy sort.

        Returns:
            (is_monotone, rank_correlation)
        """
        by_seq = self.sort_by_sequence()
        by_energy = self.sort_by_energy()

        # Check if the orderings are identical
        seq_order = [s for s, _, _ in by_seq]
        energy_order = [s for s, _, _ in by_energy]
        is_monotone = seq_order == energy_order

        # Rank correlation
        n = len(by_seq)
        if n < 2:
            return True, 1.0

        seq_rank = {s: i for i, (s, _, _) in enumerate(by_seq)}
        energy_rank = {s: i for i, (s, _, _) in enumerate(by_energy)}

        d_sq = sum((seq_rank[s] - energy_rank[s]) ** 2 for s in seq_rank)
        corr = 1.0 - (6.0 * d_sq) / (n * (n * n - 1))

        return is_monotone, corr

    def to_json(self) -> str:
        return json.dumps({
            "encoding": self.encoding,
            "n_vars": self.n_vars,
            "qubo_matrix": self.qubo_matrix,
            "entries": {
                s: {"x": x, "energy": round(e, 6)}
                for s, (x, e) in self.entries.items()
            },
        }, indent=2)

    @classmethod
    def from_json(cls, data: str) -> "LUT":
        d = json.loads(data)
        entries = {
            s: (v["x"], v["energy"]) for s, v in d["entries"].items()
        }
        return cls(
            entries=entries,
            qubo_matrix=d["qubo_matrix"],
            n_vars=d["n_vars"],
            encoding=d["encoding"],
        )


def build_direct_lut(
    Q: List[List[float]],
    n_vars: int,
    base_map: Optional[Dict[int, str]] = None,
) -> LUT:
    """Build LUT with direct encoding: x_i → base by fixed mapping.

    Default: 0 → A, 1 → G.
    Symbolic order does NOT match energy order.
    """
    if base_map is None:
        base_map = {0: "A", 1: "G"}

    entries = {}
    for i in range(2**n_vars):
        x = [(i >> j) & 1 for j in range(n_vars)]
        seq = "".join(base_map[v] for v in x)
        energy = qubo_energy(x, Q)
        entries[seq] = (x, energy)

    return LUT(
        entries=entries,
        qubo_matrix=Q,
        n_vars=n_vars,
        encoding="direct",
    )


def build_monotone_lut(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
) -> LUT:
    """Build LUT with monotone encoding: DNA sequence = energy rank.

    1. Enumerate all solutions (or sample)
    2. Sort by energy
    3. Assign DNA sequences in energy order

    Result: lexicographic DNA sort = energy sort BY CONSTRUCTION.

    Args:
        Q: QUBO matrix
        n_vars: number of variables
        n_samples: 0 for brute force (n_vars ≤ 12), else random sampling
        seed: RNG seed

    Returns:
        LUT with monotone encoding
    """
    # Step 1: enumerate solutions
    solutions = []
    if n_samples == 0 and n_vars <= 12:
        for i in range(2**n_vars):
            x = [(i >> j) & 1 for j in range(n_vars)]
            energy = qubo_energy(x, Q)
            solutions.append((x, energy))
    else:
        n_samples = n_samples or 10000
        rng = random.Random(seed)
        seen = set()
        for _ in range(n_samples):
            x = tuple(rng.randint(0, 1) for _ in range(n_vars))
            if x in seen:
                continue
            seen.add(x)
            energy = qubo_energy(list(x), Q)
            solutions.append((list(x), energy))

    # Step 2: sort by energy
    solutions.sort(key=lambda item: item[1])

    # Step 3: assign DNA sequences in energy order
    seq_len = sequence_length_for_solutions(len(solutions))
    entries = {}
    for rank, (x, energy) in enumerate(solutions):
        seq = int_to_dna(rank, seq_len)
        entries[seq] = (x, energy)

    return LUT(
        entries=entries,
        qubo_matrix=Q,
        n_vars=n_vars,
        encoding="monotone",
    )


def build_positional_lut(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
) -> LUT:
    """Build LUT with position-dependent encoding.

    Each variable gets a base pair chosen by Q_ii sign.
    Not monotone, but fixes sign inversion for Ising-type QUBOs.
    """
    # Design position-dependent base pairs
    pairs = [("A", "G"), ("T", "C"), ("B", "S"), ("P", "Z")]
    pos_map = []
    for i in range(n_vars):
        pair = pairs[i % len(pairs)]
        if Q[i][i] < 0:
            pos_map.append((pair[1], pair[0]))  # swap for negative diagonal
        else:
            pos_map.append(pair)

    entries = {}
    if n_samples == 0 and n_vars <= 12:
        for i in range(2**n_vars):
            x = [(i >> j) & 1 for j in range(n_vars)]
            seq = "".join(pos_map[j][x[j]] for j in range(n_vars))
            energy = qubo_energy(x, Q)
            entries[seq] = (x, energy)
    else:
        n_samples = n_samples or 10000
        rng = random.Random(seed)
        seen = set()
        for _ in range(n_samples):
            x = tuple(rng.randint(0, 1) for _ in range(n_vars))
            if x in seen:
                continue
            seen.add(x)
            seq = "".join(pos_map[j][x[j]] for j in range(n_vars))
            energy = qubo_energy(list(x), Q)
            entries[seq] = (list(x), energy)

    return LUT(
        entries=entries,
        qubo_matrix=Q,
        n_vars=n_vars,
        encoding="positional",
    )


# ============================================================
# §4  VERIFICATION
# ============================================================

def verify_all_encodings(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
) -> Dict[str, dict]:
    """Compare all three encoding strategies.

    Returns:
        Dict with results for each encoding method.
    """
    results = {}

    # Direct encoding
    lut_direct = build_direct_lut(Q, n_vars)
    mono, corr = lut_direct.verify_monotone()
    best_seq = lut_direct.sort_by_energy()[0]
    worst_seq = lut_direct.sort_by_sequence()[0]
    results["direct"] = {
        "encoding": "direct",
        "n_entries": lut_direct.size(),
        "is_monotone": mono,
        "rank_correlation": round(corr, 4),
        "min_energy": round(best_seq[2], 4),
        "min_energy_seq": best_seq[0],
        "first_lex_seq": worst_seq[0],
        "first_lex_energy": round(worst_seq[2], 4),
    }

    # Positional encoding
    lut_pos = build_positional_lut(Q, n_vars, n_samples)
    mono, corr = lut_pos.verify_monotone()
    best_seq = lut_pos.sort_by_energy()[0]
    worst_seq = lut_pos.sort_by_sequence()[0]
    results["positional"] = {
        "encoding": "positional",
        "n_entries": lut_pos.size(),
        "is_monotone": mono,
        "rank_correlation": round(corr, 4),
        "min_energy": round(best_seq[2], 4),
        "min_energy_seq": best_seq[0],
        "first_lex_seq": worst_seq[0],
        "first_lex_energy": round(worst_seq[2], 4),
    }

    # Monotone encoding
    lut_mono = build_monotone_lut(Q, n_vars, n_samples)
    mono, corr = lut_mono.verify_monotone()
    best_seq = lut_mono.sort_by_energy()[0]
    worst_seq = lut_mono.sort_by_sequence()[0]
    results["monotone"] = {
        "encoding": "monotone",
        "n_entries": lut_mono.size(),
        "is_monotone": mono,
        "rank_correlation": round(corr, 4),
        "min_energy": round(best_seq[2], 4),
        "min_energy_seq": best_seq[0],
        "first_lex_seq": worst_seq[0],
        "first_lex_energy": round(worst_seq[2], 4),
    }

    return results


# ============================================================
# §5  DEMO QUBO GENERATORS
# ============================================================

def demo_qubo(n: int = 6, seed: int = 42, style: str = "random") -> List[List[float]]:
    """Generate a demo QUBO."""
    rng = random.Random(seed)
    Q = [[0.0] * n for _ in range(n)]

    if style == "diagonal":
        for i in range(n):
            Q[i][i] = rng.uniform(1, 5)
    elif style == "banded":
        for i in range(n):
            Q[i][i] = rng.uniform(1, 5)
            if i + 1 < n:
                val = rng.uniform(-2, 2)
                Q[i][i + 1] = val
                Q[i + 1][i] = val
    elif style == "ising":
        for i in range(n):
            Q[i][i] = -1.0
            if i + 1 < n:
                Q[i][i + 1] = -0.5
                Q[i + 1][i] = -0.5
    else:  # random
        for i in range(n):
            Q[i][i] = rng.uniform(-2, 5)
            for j in range(i + 1, n):
                val = rng.uniform(-2, 2)
                Q[i][j] = val
                Q[j][i] = val

    return Q


# ============================================================
# §6  CLI
# ============================================================

def run_demo(name: str, Q: List[List[float]], n_vars: int):
    """Run a single demo comparing all encodings."""
    print(f"\n{'='*60}")
    print(f"  {name}")
    print(f"{'='*60}")

    results = verify_all_encodings(Q, n_vars)

    for method in ["direct", "positional", "monotone"]:
        r = results[method]
        print(f"\n  {method.upper()} encoding:")
        print(f"    Entries:       {r['n_entries']}")
        print(f"    Monotone:      {r['is_monotone']}")
        print(f"    Rank corr:     {r['rank_correlation']}")
        print(f"    Min energy:    {r['min_energy']} (seq: {r['min_energy_seq']})")
        print(f"    First in lex:  {r['first_lex_energy']} (seq: {r['first_lex_seq']})")

        if r["is_monotone"]:
            print(f"    ✓ Lexicographic sort = Energy sort")
        else:
            print(f"    ✗ Lexicographic sort ≠ Energy sort")


if __name__ == "__main__":
    print("=" * 60)
    print("Symbology LUT Adaptation — Fixed")
    print("=" * 60)

    # Diagonal
    Q1 = demo_qubo(6, seed=42, style="diagonal")
    run_demo("Diagonal QUBO (6 vars)", Q1, 6)

    # Banded
    Q2 = demo_qubo(6, seed=42, style="banded")
    run_demo("Banded QUBO (6 vars)", Q2, 6)

    # Ising
    Q3 = demo_qubo(6, seed=42, style="ising")
    run_demo("1D Ising Chain (6 vars)", Q3, 6)

    # Random
    Q4 = demo_qubo(6, seed=42, style="random")
    run_demo("Random QUBO (6 vars)", Q4, 6)

    # Larger
    Q5 = demo_qubo(8, seed=42, style="banded")
    run_demo("Banded QUBO (8 vars)", Q5, 8)

    print(f"\n{'='*60}")
    print("DONE")
    print(f"{'='*60}")