#!/usr/bin/env python3
"""
dna_qubo_sort.py — QUBO Energy Minimization via DNA Sorting

The "backdoor energy sort": encode QUBO candidate solutions as Hachimoji
DNA sequences where melting temperature (Tm) correlates with QUBO energy.
Then sort by Tm — the sorted order IS the energy ranking.

Theory:
    QUBO energy E(x) = x^T Q x
    For a diagonal-dominant QUBO with non-negative entries:
        E(x) ≈ Σ_i Q_ii · x_i + cross terms

    If we encode x_i=0 → A/T (low Tm) and x_i=1 → G/C (high Tm),
    then Tm(sequence) ∝ Σ_i (Tm contribution for x_i) ∝ E(x).

    Sorting by Tm (ascending) = sorting by energy (ascending).
    Top of sorted list = minimum energy = optimal solution.

Adleman's insight (1994): let molecular physics do the sorting.
    - Generate all 2^n candidate sequences (massively parallel)
    - Sort by physical property (gel electrophoresis, affinity)
    - Read the result

Modern insight: samtools sort processes billions of reads.
    The infrastructure already exists.
"""

from __future__ import annotations

import hashlib
import json
import math
import random
import struct
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

from dna_codec import (
    HACHIMOJI_BASES,
    decode_binary_vector,
    encode_binary_vector,
    gc_content,
    melting_temperature,
    qubo_energy,
    sequence_stats,
)
from q16_canonical import float_to_q16, q16_to_float, q16_to_bytes, q16_from_bytes


# ============================================================
# §1  QUBO → DNA ENCODING
# ============================================================

@dataclass
class QuboDnaCandidate:
    """A single QUBO solution encoded as a DNA sequence."""
    x: List[int]          # binary solution vector
    sequence: str         # Hachimoji DNA encoding
    energy: float         # QUBO energy E(x)
    energy_q16: int       # Q16_16 fixed-point energy
    tm: float             # melting temperature
    gc: float             # GC-equivalent content

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


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

    The encoding maps:
        x_i = 0 → 'A' repeated bits_per_var times (low Tm)
        x_i = 1 → 'P' repeated bits_per_var times (high Tm)

    This ensures Tm(sequence) is monotonically related to
    the number of 1s in x, which drives energy for non-negative Q.

    Args:
        x: binary solution vector
        Q: QUBO matrix
        bits_per_var: bases per variable

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

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


# ============================================================
# §2  BRUTE-FORCE GENERATION (small QUBOs)
# ============================================================

def generate_all_candidates(
    Q: List[List[float]],
    n_vars: int,
    bits_per_var: int = 3,
    max_candidates: int = 4096,
) -> List[QuboDnaCandidate]:
    """Generate all 2^n candidate solutions and encode as DNA.

    Only feasible for small n (n ≤ 12 gives 4096 candidates).
    For larger QUBOs, use generate_random_candidates().

    Args:
        Q: QUBO matrix
        n_vars: number of binary variables
        bits_per_var: bases per variable
        max_candidates: safety cap

    Returns:
        List of QuboDnaCandidate, one per solution
    """
    if n_vars > 12:
        raise ValueError(
            f"n_vars={n_vars} too large for brute force (2^{n_vars} = {2**n_vars}). "
            f"Use generate_random_candidates() instead."
        )

    candidates = []
    for i in range(min(2**n_vars, max_candidates)):
        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 = 10000,
    seed: int = 42,
    bits_per_var: int = 3,
) -> List[QuboDnaCandidate]:
    """Generate random candidate solutions for large QUBOs.

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

    Returns:
        List of QuboDnaCandidate
    """
    rng = random.Random(seed)
    candidates = []
    seen = set()

    for _ in range(n_samples):
        # Generate random binary vector
        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


# ============================================================
# §3  DNA SORTING (the backdoor)
# ============================================================

def sort_by_tm(candidates: List[QuboDnaCandidate]) -> List[QuboDnaCandidate]:
    """Sort candidates by melting temperature (ascending).

    Low Tm = low GC content = fewer 1s in x = lower energy (for non-negative Q).
    This is the "backdoor" — sorting by a physical property IS energy ranking.

    Args:
        candidates: list of QuboDnaCandidate

    Returns:
        Sorted list (lowest Tm first = lowest energy first)
    """
    return sorted(candidates, key=lambda c: c.tm)


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


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


def sort_by_energy_q16(candidates: List[QuboDnaCandidate]) -> List[QuboDnaCandidate]:
    """Sort candidates by Q16_16 fixed-point energy (ascending)."""
    return sorted(candidates, key=lambda c: c.energy_q16)


# ============================================================
# §4  SAMTOOLS COMPATIBLE OUTPUT
# ============================================================

def candidates_to_sam(
    candidates: List[QuboDnaCandidate],
    qubo_name: str = "QUBO",
) -> str:
    """Export candidates as a SAM file for samtools sort.

    Each candidate becomes a SAM read with:
        - QNAME: energy ranking
        - SEQ: Hachimoji DNA sequence
        - QUAL: quality scores based on Tm
        - TAG: energy, gc content

    This produces a valid SAM file that can be sorted with:
        samtools sort -t TM -o sorted.bam input.sam

    Args:
        candidates: list of QuboDnaCandidate
        qubo_name: name for the QUBO problem

    Returns:
        SAM format string
    """
    lines = []
    # SAM header
    lines.append("@HD\tVN:1.6\tSO:unsorted")
    lines.append(f"@PG\tID:dna_qubo_sort\tPN:dna_qubo_sort\tVN:0.1")
    lines.append(f"@CO\tQUBO problem: {qubo_name}, {len(candidates)} candidates")

    for i, cand in enumerate(candidates):
        qname = f"{qubo_name}_{i:06d}"
        flag = 0
        rname = "*"
        pos = 0
        mapq = 0
        cigar = "*"
        rnext = "*"
        pnext = 0
        tlen = len(cand.sequence)
        seq = cand.sequence
        qual = "~" * len(seq)  # placeholder quality

        # Custom tags
        tags = [
            f"TM:f:{cand.tm:.4f}",
            f"GC:f:{cand.gc:.4f}",
            f"EN:f:{cand.energy:.6f}",
            f"EQ:i:{cand.energy_q16}",
            f"XV:Z:{','.join(str(v) for v in cand.x)}",
        ]

        line = "\t".join(
            str(x) for x in [qname, flag, rname, pos, mapq, cigar, rnext, pnext, tlen, seq, qual]
        ) + "\t" + "\t".join(tags)
        lines.append(line)

    return "\n".join(lines) + "\n"


def candidates_to_fasta(candidates: List[QuboDnaCandidate]) -> str:
    """Export candidates as FASTA (sequence only, for external tools)."""
    lines = []
    for i, cand in enumerate(candidates):
        lines.append(f">candidate_{i:06d} energy={cand.energy:.6f} tm={cand.tm:.4f}")
        # Wrap at 80 chars
        seq = cand.sequence
        for j in range(0, len(seq), 80):
            lines.append(seq[j : j + 80])
    return "\n".join(lines) + "\n"


# ============================================================
# §5  VERIFICATION: does Tm sorting = energy sorting?
# ============================================================

@dataclass
class SortVerification:
    """Results of verifying that Tm sort matches energy sort."""
    n_candidates: int
    n_vars: int
    tm_sort_matches_energy: bool
    rank_correlation: float       # Spearman-like correlation
    top_k_agreement: int          # number of top-k that match
    top_k: int
    min_energy_candidate: QuboDnaCandidate
    min_tm_candidate: QuboDnaCandidate
    energy_sort: List[QuboDnaCandidate]
    tm_sort: List[QuboDnaCandidate]

    def to_dict(self) -> dict:
        return {
            "n_candidates": self.n_candidates,
            "n_vars": self.n_vars,
            "tm_sort_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_candidate.to_dict(),
            "min_tm": self.min_tm_candidate.to_dict(),
        }


def spearman_rank_correlation(
    values_a: List[float],
    values_b: List[float],
) -> float:
    """Compute Spearman rank correlation between two value lists."""
    n = len(values_a)
    if n < 2:
        return 0.0

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

    rank_a = rank(values_a)
    rank_b = rank(values_b)

    # Spearman formula: 1 - (6 * Σd²) / (n * (n² - 1))
    d_sq_sum = sum((ra - rb) ** 2 for ra, rb in zip(rank_a, rank_b))
    denom = n * (n * n - 1)
    if denom == 0:
        return 0.0
    return 1.0 - (6.0 * d_sq_sum) / denom


def verify_tm_sort(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
    top_k: int = 5,
) -> SortVerification:
    """Verify that sorting by Tm gives the same order as sorting by energy.

    This is THE critical test. If Tm sort ≠ energy sort, the backdoor doesn't work.

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

    Returns:
        SortVerification with detailed results
    """
    if n_samples == 0:
        candidates = generate_all_candidates(Q, n_vars)
    else:
        candidates = generate_random_candidates(Q, n_vars, n_samples, seed)

    energy_sorted = sort_by_energy(candidates)
    tm_sorted = sort_by_tm(candidates)

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

    # Top-k agreement
    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 SortVerification(
        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_candidate=energy_sorted[0],
        min_tm_candidate=tm_sorted[0],
        energy_sort=energy_sorted,
        tm_sort=tm_sorted,
    )


# ============================================================
# §6  DEMO QUBO PROBLEMS
# ============================================================

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

    Max-Cut: partition vertices into two sets to maximize edges between them.
    QUBO formulation: minimize -Σ_{(i,j)∈E} x_i(1-x_j) + x_j(1-x_i)

    Args:
        n: number of vertices
        seed: RNG seed

    Returns:
        QUBO matrix (n×n)
    """
    rng = random.Random(seed)
    Q = [[0.0] * n for _ in range(n)]

    # Random graph with ~50% edge probability
    for i in range(n):
        for j in range(i + 1, n):
            if rng.random() < 0.5:
                # Edge exists: reward for having endpoints in different sets
                Q[i][j] = -1.0
                Q[j][i] = -1.0
                Q[i][i] += 1.0
                Q[j][j] += 1.0

    return Q


def demo_number_partition(numbers: List[int]) -> List[List[float]]:
    """Generate a Number Partitioning QUBO.

    Given numbers, partition into two sets with equal sum.
    QUBO: minimize (Σ s_i · a_i)² where s_i ∈ {-1, +1}

    Args:
        numbers: list of integers to partition

    Returns:
        QUBO matrix
    """
    n = len(numbers)
    Q = [[0.0] * n for _ in range(n)]

    for i in range(n):
        for j in range(n):
            Q[i][j] = float(numbers[i] * numbers[j])

    return Q


# ============================================================
# §7  FULL PIPELINE
# ============================================================

def run_qubo_dna_sort(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
    output_prefix: str = "qubo_dna",
) -> dict:
    """Run the full QUBO → DNA → Sort → Verify pipeline.

    Args:
        Q: QUBO matrix
        n_vars: number of variables
        n_samples: 0 for brute force, >0 for sampling
        seed: RNG seed
        output_prefix: prefix for output files

    Returns:
        Summary dict
    """
    print(f"QUBO-DNA Sort Pipeline")
    print(f"  Variables: {n_vars}")
    print(f"  Matrix size: {len(Q)}×{len(Q[0])}")

    # Generate candidates
    if n_samples == 0 and n_vars <= 12:
        candidates = generate_all_candidates(Q, n_vars)
        print(f"  Generated: {len(candidates)} (brute force)")
    else:
        n_samples = n_samples or 10000
        candidates = generate_random_candidates(Q, n_vars, n_samples, seed)
        print(f"  Generated: {len(candidates)} (random sampling)")

    # Sort by different methods
    energy_sorted = sort_by_energy(candidates)
    tm_sorted = sort_by_tm(candidates)
    gc_sorted = sort_by_gc(candidates)

    print(f"\n  Energy-optimal: x={energy_sorted[0].x}, E={energy_sorted[0].energy:.6f}")
    print(f"  Tm-optimal:     x={tm_sorted[0].x}, E={tm_sorted[0].energy:.6f}")
    print(f"  GC-optimal:     x={gc_sorted[0].x}, E={gc_sorted[0].energy:.6f}")

    # Verify Tm sort
    verification = verify_tm_sort(Q, n_vars, n_samples if n_samples else 0, seed)
    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}")

    # Export SAM
    sam_content = candidates_to_sam(candidates, output_prefix)
    sam_path = f"{output_prefix}.sam"
    with open(sam_path, "w") as f:
        f.write(sam_content)
    print(f"\n  SAM output: {sam_path} ({len(candidates)} reads)")

    # Export FASTA
    fasta_content = candidates_to_fasta(candidates)
    fasta_path = f"{output_prefix}.fasta"
    with open(fasta_path, "w") as f:
        f.write(fasta_content)
    print(f"  FASTA output: {fasta_path}")

    # Summary
    summary = {
        "n_vars": n_vars,
        "n_candidates": len(candidates),
        "energy_optimal": energy_sorted[0].to_dict(),
        "tm_optimal": tm_sorted[0].to_dict(),
        "gc_optimal": gc_sorted[0].to_dict(),
        "verification": verification.to_dict(),
        "files": {
            "sam": sam_path,
            "fasta": fasta_path,
        },
    }

    # Write summary JSON
    json_path = f"{output_prefix}_summary.json"
    with open(json_path, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"  Summary: {json_path}")

    return summary


# ============================================================
# §8  CLI
# ============================================================

if __name__ == "__main__":
    import sys

    print("=" * 60)
    print("QUBO-DNA Sort: Backdoor Energy Minimization")
    print("=" * 60)

    # Demo 1: Small Max-Cut
    print("\n--- Demo 1: Max-Cut (6 vertices, brute force) ---")
    Q1 = demo_max_cut(6, seed=42)
    result1 = run_qubo_dna_sort(Q1, n_vars=6, output_prefix="qubo_dna_maxcut6")

    # Demo 2: Number Partition
    print("\n--- Demo 2: Number Partition ---")
    numbers = [3, 7, 1, 5, 9, 2, 8, 4]
    Q2 = demo_number_partition(numbers)
    result2 = run_qubo_dna_sort(Q2, n_vars=len(numbers), output_prefix="qubo_dna_partition8")

    # Demo 3: Larger random QUBO
    print("\n--- Demo 3: Random QUBO (10 vars, sampling) ---")
    rng = random.Random(42)
    n3 = 10
    Q3 = [[rng.uniform(-2, 2) if i != j else rng.uniform(0, 5) for j in range(n3)] for i in range(n3)]
    result3 = run_qubo_dna_sort(Q3, n_vars=n3, n_samples=5000, output_prefix="qubo_dna_random10")

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