#!/usr/bin/env python3
"""
dna_qubo_sort.py — Naive QUBO-DNA Encoding via Tm Sorting (Approach A)

Encodes QUBO solutions as DNA sequences where melting temperature (Tm)
is approximately proportional to QUBO energy. The key idea:

  x_i = 0 → low-Tm base (A, Tm=2.0)
  x_i = 1 → high-Tm base (P, Tm=3.5)

With 3 bases per variable for majority-vote decoding redundancy.

For diagonal QUBOs with positive entries, Tm sort ≈ energy sort.
For general QUBOs (with off-diagonal or negative terms), correlation
is degraded — this is expected and documented.

This is the "naive" approach. For better results, see dna_qubo_nn.py
which uses nearest-neighbor stacking thermodynamics.
"""

from __future__ import annotations

import random
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

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


# ============================================================
# §1  QUBO-DNA CANDIDATE
# ============================================================

@dataclass
class QuboDnaCandidate:
    """A QUBO solution encoded as a DNA sequence with Tm metadata."""
    x: List[int]
    sequence: str
    energy: float
    tm: float
    gc: float

    def to_dict(self) -> dict:
        return {
            "x": self.x,
            "sequence": self.sequence,
            "energy": round(self.energy, 6),
            "tm": round(self.tm, 4),
            "gc": round(self.gc, 4),
        }


# ============================================================
# §2  ENCODING
# ============================================================

def encode_qubo_candidate(
    x: List[int],
    Q: List[List[float]],
    bits_per_var: int = 3,
) -> QuboDnaCandidate:
    """Encode a QUBO solution as a DNA sequence.

    Uses the simple encoding: 0 → A (low Tm), 1 → P (high Tm).
    Tm is monotonically proportional to the number of 1s in x.

    Args:
        x: binary solution vector
        Q: QUBO matrix (for energy computation)
        bits_per_var: bases per variable (default 3)

    Returns:
        QuboDnaCandidate with all metadata
    """
    sequence = encode_binary_vector(x, bits_per_var)
    energy = qubo_energy(x, Q)
    tm = melting_temperature(sequence)
    gc = raw_gc_content(sequence)

    return QuboDnaCandidate(
        x=x,
        sequence=sequence,
        energy=energy,
        tm=tm,
        gc=gc,
    )


# ============================================================
# §3  CANDIDATE GENERATION
# ============================================================

def generate_all_candidates(
    Q: List[List[float]],
    n_vars: int,
    bits_per_var: int = 3,
) -> List[QuboDnaCandidate]:
    """Generate all QUBO candidates by brute force.

    Args:
        Q: QUBO matrix
        n_vars: number of variables (must be <= 12 for brute force)
        bits_per_var: bases per variable

    Returns:
        List of all QuboDnaCandidate (2^n_vars items)
    """
    candidates = []
    for i in range(2**n_vars):
        x = [(i >> j) & 1 for j in range(n_vars)]
        candidates.append(encode_qubo_candidate(x, Q, bits_per_var))
    return candidates


def generate_random_candidates(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 1000,
    seed: int = 42,
    bits_per_var: int = 3,
) -> List[QuboDnaCandidate]:
    """Generate random QUBO candidates.

    Args:
        Q: QUBO matrix
        n_vars: number of variables
        n_samples: number of random samples
        seed: RNG seed
        bits_per_var: bases per variable

    Returns:
        List of QuboDnaCandidate
    """
    rng = random.Random(seed)
    seen = set()
    candidates = []
    for _ in range(n_samples):
        x = tuple(rng.randint(0, 1) for _ in range(n_vars))
        if x in seen:
            continue
        seen.add(x)
        candidates.append(encode_qubo_candidate(list(x), Q, bits_per_var))
    return candidates


# ============================================================
# §4  SORTING
# ============================================================

def sort_by_energy(candidates: List[QuboDnaCandidate]) -> List[QuboDnaCandidate]:
    """Sort candidates by QUBO energy (ascending). Ground truth."""
    return sorted(candidates, key=lambda c: c.energy)


def sort_by_tm(candidates: List[QuboDnaCandidate]) -> List[QuboDnaCandidate]:
    """Sort candidates by melting temperature (ascending)."""
    return sorted(candidates, key=lambda c: c.tm)


def sort_by_gc(candidates: List[QuboDnaCandidate]) -> List[QuboDnaCandidate]:
    """Sort candidates by GC content (ascending)."""
    return sorted(candidates, key=lambda c: c.gc)


def spearman_rank_correlation(values_a: List[float], values_b: List[float]) -> float:
    """Compute Spearman rank correlation.

    Returns 1.0 for perfect positive correlation, -1.0 for perfect negative.
    """
    n = len(values_a)
    if n < 2:
        return 0.0

    def rank(vals):
        sorted_idx = sorted(range(n), key=lambda i: vals[i])
        ranks = [0] * n
        for r, idx in enumerate(sorted_idx):
            ranks[idx] = r
        return ranks

    ra = rank(values_a)
    rb = rank(values_b)
    d_sq = sum((a - b) ** 2 for a, b in zip(ra, rb))
    denom = n * (n * n - 1)
    return 1.0 - (6.0 * d_sq) / denom if denom > 0 else 0.0


# ============================================================
# §5  Tm SORT VERIFICATION
# ============================================================

@dataclass
class TmVerification:
    """Results of Tm sort verification."""
    n_candidates: int
    n_vars: int
    tm_sort_matches_energy: bool
    rank_correlation: float
    top_k_agreement: int
    top_k: int
    min_energy: QuboDnaCandidate
    min_tm: QuboDnaCandidate

    def to_dict(self) -> dict:
        return {
            "n_candidates": self.n_candidates,
            "n_vars": self.n_vars,
            "tm_matches_energy": self.tm_sort_matches_energy,
            "rank_correlation": round(self.rank_correlation, 4),
            "top_k_agreement": self.top_k_agreement,
            "top_k": self.top_k,
            "min_energy": self.min_energy.to_dict(),
            "min_tm": self.min_tm.to_dict(),
        }


def verify_tm_sort(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
    top_k: int = 5,
    bits_per_var: int = 3,
) -> TmVerification:
    """Verify that Tm sort matches energy sort.

    Args:
        Q: QUBO matrix
        n_vars: number of variables
        n_samples: 0 for brute force, else random sampling
        seed: RNG seed
        top_k: number of top candidates to compare
        bits_per_var: bases per variable

    Returns:
        TmVerification with results
    """
    if n_samples == 0 and n_vars <= 12:
        candidates = generate_all_candidates(Q, n_vars, bits_per_var)
    else:
        n_samples = n_samples or 10000
        candidates = generate_random_candidates(Q, n_vars, n_samples, seed, bits_per_var)

    energies = [c.energy for c in candidates]
    tms = [c.tm for c in candidates]
    corr = spearman_rank_correlation(energies, tms)

    energy_sorted = sort_by_energy(candidates)
    if corr < 0:
        tm_sorted = sorted(candidates, key=lambda c: c.tm, reverse=True)
    else:
        tm_sorted = sorted(candidates, key=lambda c: c.tm)

    top_k = min(top_k, len(candidates))
    energy_top = set(tuple(c.x) for c in energy_sorted[:top_k])
    tm_top = set(tuple(c.x) for c in tm_sorted[:top_k])
    agreement = len(energy_top & tm_top)

    return TmVerification(
        n_candidates=len(candidates),
        n_vars=n_vars,
        tm_sort_matches_energy=(energy_sorted[0].x == tm_sorted[0].x),
        rank_correlation=corr,
        top_k_agreement=agreement,
        top_k=top_k,
        min_energy=energy_sorted[0],
        min_tm=tm_sorted[0],
    )


# ============================================================
# §6  SAM / FASTA EXPORT
# ============================================================

def candidates_to_sam(
    candidates: List[QuboDnaCandidate],
    name: str = "qubo_sort",
) -> str:
    """Export candidates to SAM format.

    SAM is a tab-delimited format used in bioinformatics.
    We use optional tags for energy, Tm, and solution vector.

    Args:
        candidates: list of QuboDnaCandidate
        name: sample name

    Returns:
        SAM-formatted string
    """
    lines = [
        "@HD\tVN:1.6\tSO:unsorted",
        f"@PG\tID:{name}\tPN:{name}\tVN:0.1",
    ]
    for i, cand in enumerate(candidates):
        tags = [
            f"TM:f:{cand.tm:.4f}",
            f"EN:f:{cand.energy:.6f}",
            f"GC:f:{cand.gc:.4f}",
            f"XV:Z:{','.join(str(v) for v in cand.x)}",
        ]
        line = (
            f"cand_{i:06d}\t0\t*\t0\t0\t*\t*\t0\t{len(cand.sequence)}\t"
            f"{cand.sequence}\t{'~' * len(cand.sequence)}\t"
            + "\t".join(tags)
        )
        lines.append(line)
    return "\n".join(lines) + "\n"


def candidates_to_fasta(
    candidates: List[QuboDnaCandidate],
    prefix: str = "qubo",
) -> str:
    """Export candidates to FASTA format.

    FASTA is a simple format with header lines (>) followed by sequence.

    Args:
        candidates: list of QuboDnaCandidate
        prefix: prefix for sequence IDs

    Returns:
        FASTA-formatted string
    """
    lines = []
    for i, cand in enumerate(candidates):
        header = f">{prefix}_cand_{i:06d} energy={cand.energy:.6f} tm={cand.tm:.4f}"
        lines.append(header)
        lines.append(cand.sequence)
    return "\n".join(lines) + "\n"


# ============================================================
# §7  DEMO QUBOs
# ============================================================

def demo_max_cut(n: int = 6, seed: int = 42) -> List[List[float]]:
    """Generate a Max-Cut QUBO.

    Max-Cut: E(x) = Σ_{(i,j) in E} (x_i + x_j - 2*x_i*x_j)
    Equivalent to: Q_ii = degree(i), Q_ij = -2 for edges
    """
    rng = random.Random(seed)
    Q = [[0.0] * n for _ in range(n)]

    # Create a random graph (Erdos-Renyi with p=0.5)
    for i in range(n):
        for j in range(i + 1, n):
            if rng.random() < 0.5:
                Q[i][i] += 1.0
                Q[j][j] += 1.0
                Q[i][j] = -2.0
                Q[j][i] = -2.0

    return Q


def demo_number_partition(seed: int = 42) -> List[List[float]]:
    """Generate a small Number Partitioning QUBO.

    NPP: partition numbers into two sets with equal sum.
    Q_ii = a_i^2, Q_ij = 2*a_i*a_j
    """
    rng = random.Random(seed)
    numbers = [rng.randint(1, 20) for _ in range(6)]
    n = len(numbers)
    Q = [[0.0] * n for _ in range(n)]

    for i in range(n):
        Q[i][i] = numbers[i] ** 2
        for j in range(i + 1, n):
            Q[i][j] = 2.0 * numbers[i] * numbers[j]
            Q[j][i] = 2.0 * numbers[i] * numbers[j]

    return Q


# ============================================================
# §8  FULL PIPELINE
# ============================================================

def run_qubo_dna_sort(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
    bits_per_var: int = 3,
) -> dict:
    """Run the full QUBO-DNA sort pipeline.

    Args:
        Q: QUBO matrix
        n_vars: number of variables
        n_samples: 0 for brute force, else sampling
        seed: RNG seed
        bits_per_var: bases per variable

    Returns:
        Summary dict with verification results
    """
    print(f"QUBO-DNA Sort Pipeline")
    print(f"  Variables: {n_vars}")

    if n_samples == 0 and n_vars <= 12:
        candidates = generate_all_candidates(Q, n_vars, bits_per_var)
        print(f"  Brute force: {len(candidates)} candidates")
    else:
        n_samples = n_samples or 10000
        candidates = generate_random_candidates(Q, n_vars, n_samples, seed, bits_per_var)
        print(f"  Random sampling: {len(candidates)} candidates")

    # Calculate correlation to determine sorting direction
    energies = [c.energy for c in candidates]
    tms = [c.tm for c in candidates]
    corr = spearman_rank_correlation(energies, tms)

    # Sort
    energy_sorted = sort_by_energy(candidates)
    if corr < 0:
        tm_sorted = sorted(candidates, key=lambda c: c.tm, reverse=True)
    else:
        tm_sorted = sorted(candidates, key=lambda c: c.tm)

    print(f"\n  Min energy: {energy_sorted[0].energy:.4f}, x={energy_sorted[0].x}")
    print(f"  Min Tm:     {tm_sorted[0].tm:.4f}, x={tm_sorted[0].x}")

    # Verify
    verification = verify_tm_sort(Q, n_vars, n_samples, seed, bits_per_var)
    print(f"\n  Tm sort = Energy sort: {verification.tm_sort_matches_energy}")
    print(f"  Rank correlation: {verification.rank_correlation:.4f}")
    print(f"  Top-{verification.top_k} agreement: {verification.top_k_agreement}/{verification.top_k}")

    return {
        "n_vars": n_vars,
        "n_candidates": len(candidates),
        "energy_optimal": energy_sorted[0].to_dict(),
        "tm_optimal": tm_sorted[0].to_dict(),
        "verification": verification.to_dict(),
    }


# ============================================================
# §9  CLI
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("QUBO-DNA Sort: Naive Tm Encoding")
    print("=" * 60)

    # Demo 1: Diagonal QUBO (ideal case)
    print("\n--- Demo 1: Diagonal QUBO (6 vars) ---")
    n = 6
    Q = [[3.0 if i == j else 0.0 for j in range(n)] for i in range(n)]
    r1 = run_qubo_dna_sort(Q, n)

    # Demo 2: Max-Cut
    print("\n--- Demo 2: Max-Cut (6 vars) ---")
    Q2 = demo_max_cut(6, seed=42)
    r2 = run_qubo_dna_sort(Q2, 6)

    # Demo 3: Ising (negative diagonal)
    print("\n--- Demo 3: Ising Chain (6 vars) ---")
    Q3 = [[-1.0 if i == j else (-0.5 if abs(i-j) == 1 else 0.0) for j in range(6)] for i in range(6)]
    r3 = run_qubo_dna_sort(Q3, 6)

    print("\n" + "=" * 60)
    print("DONE")
    print("=" * 60)