#!/usr/bin/env python3
"""
dna_qubo_nn.py — Nearest-Neighbor QUBO-DNA Encoding (Approach B)

Encodes the QUBO matrix structure into the DNA sequence via nearest-neighbor
stacking thermodynamics. The melting temperature becomes a function of the
FULL energy, including off-diagonal terms.

Theory:
    QUBO energy: E(x) = Σ_i Q_ii x_i + Σ_{i<j} Q_ij x_i x_j

    DNA Tm model (nearest-neighbor):
        Tm = Σ_i tm_base(b_i) + Σ_i tm_stack(b_i, b_{i+1})

    Key insight: dinucleotide stacking energy tm_stack(b_i, b_{i+1})
    is nonzero ONLY when both x_i=1 and x_{i+1}=1 (using high-Tm bases
    for x=1 and low-Tm bases for x=0).

    If we set:
        tm_base(b_i) ≈ Q_ii · x_i
        tm_stack(b_i, b_{i+1}) ≈ Q_{i,i+1} · x_i · x_{i+1}

    Then Tm ≈ c₁ · E(x) + c₀ (affine transform of energy).

    Sorting by Tm IS sorting by energy, even for off-diagonal terms.

Encoding:
    x_i = 0 → base from LOW_SET  (A or T, chosen by Q_ii sign)
    x_i = 1 → base from HIGH_SET (G or C, chosen by Q_ii sign)

    The specific base (A vs T, G vs C) encodes the SIGN of Q_ii,
    making per-base Tm contribution reflect the diagonal energy.

    Dinucleotide stacking naturally captures Q_{i,i+1} interaction.

Limitations:
    - Only captures adjacent interactions (i, i+1), not arbitrary (i, j)
    - For non-adjacent interactions, sequence reordering or multiple
      passes may be needed (see §7)
    - Works best for banded/sparse QUBOs with local interactions
"""

from __future__ import annotations

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

from q16_canonical import float_to_q16, q16_to_float


# ============================================================
# §1  NEAREST-NEIGHHER THERMODYNAMIC MODEL
# ============================================================

# SantaLucia (1998) unified nearest-neighbor parameters
# ΔG°37 values in kcal/mol (simplified, in kcal/mol × 10 for integer use)
# These are REAL thermodynamic parameters for DNA duplex stability.
# Higher (less negative) = less stable = lower Tm contribution.

# Per-base parameters (enthalpy contribution, simplified)
# Maps to: A/T = low stability, G/C = high stability
NN_BASE_ENTHALPY = {
    "A": -1.0,  # weakest (2 H-bonds)
    "T": -1.0,  # weakest (2 H-bonds)
    "G": -2.5,  # strong (3 H-bonds)
    "C": -2.5,  # strong (3 H-bonds)
}

# Nearest-neighbor stacking parameters (ΔH°, kcal/mol)
# From SantaLucia & Hicks (2004), Annu. Rev. Biophys. Biomol. Struct.
# Negative = more stable = higher Tm contribution
NN_STACK_ENTHALPY = {
    ("A", "A"): -7.9, ("A", "T"): -7.2, ("A", "G"): -8.4, ("A", "C"): -8.5,
    ("T", "A"): -7.2, ("T", "T"): -7.9, ("T", "G"): -8.2, ("T", "C"): -8.4,
    ("G", "A"): -8.2, ("G", "T"): -8.4, ("G", "G"): -8.0, ("G", "C"): -10.1,
    ("C", "A"): -8.5, ("C", "T"): -8.4, ("C", "G"): -9.8, ("C", "C"): -8.0,
}

# Entropy contribution (ΔS°, cal/mol/K) — needed for full Tm calculation
NN_STACK_ENTROPY = {
    ("A", "A"): -22.2, ("A", "T"): -20.4, ("A", "G"): -22.4, ("A", "C"): -22.7,
    ("T", "A"): -21.3, ("T", "T"): -22.2, ("T", "G"): -22.2, ("T", "C"): -22.4,
    ("G", "A"): -22.2, ("G", "T"): -22.4, ("G", "G"): -19.9, ("G", "C"): -25.5,
    ("C", "A"): -22.7, ("C", "T"): -22.4, ("C", "G"): -24.4, ("C", "C"): -19.9,
}

# Free energy at 37°C: ΔG°37 = ΔH° - T·ΔS° (simplified)
def nn_stack_dg(b1: str, b2: str, T: float = 310.15) -> float:
    """Compute nearest-neighbor free energy for a dinucleotide step.

    Uses SantaLucia unified parameters:
        ΔG°37 = ΔH° - T·ΔS°

    Args:
        b1, b2: adjacent bases
        T: temperature in Kelvin (default 37°C = 310.15 K)

    Returns:
        ΔG° in kcal/mol (more negative = more stable)
    """
    dH = NN_STACK_ENTHALPY.get((b1, b2), -7.5)  # default if unknown
    dS = NN_STACK_ENTROPY.get((b1, b2), -22.0)   # default if unknown
    return dH - T * (dS / 1000.0)  # convert cal to kcal


def nn_tm_estimate(
    sequence: str,
    oligo_conc: float = 250e-9,
    na_conc: float = 0.05,
    T: float = 310.15,
) -> float:
    """Estimate Tm using nearest-neighbor model.

    Full Tm calculation:
        Tm = ΔH° / (ΔS° + R·ln(C/4)) - 273.15

    Where ΔH° and ΔS° are summed over all nearest-neighbor pairs.

    For simplicity and sorting purposes, we use:
        Tm_proxy = -Σ ΔG°37(i,i+1)  (total free energy stability)

    More stable (lower ΔG°) = higher proxy Tm.

    Args:
        sequence: DNA sequence
        oligo_conc: oligonucleotide concentration (M)
        na_conc: sodium concentration (M)
        T: temperature (K)

    Returns:
        Tm proxy (higher = more stable)
    """
    if len(sequence) < 2:
        # Per-base stability: negate enthalpy so higher = more stable
        return -sum(NN_BASE_ENTHALPY.get(b, 0) for b in sequence)

    total_dg = 0.0
    for i in range(len(sequence) - 1):
        total_dg += nn_stack_dg(sequence[i], sequence[i + 1], T)

    # Negative of ΔG° = stability proxy (higher = more stable = higher Tm)
    return -total_dg


# ============================================================
# §2  QUBO-TO-BASE MAPPING
# ============================================================

# For each variable x_i, we choose a base that encodes both
# the value (0 or 1) and the sign of Q_ii.

# Strategy: use A/G for the pair (easy to distinguish)
#   x_i = 0, Q_ii >= 0 → A (weakest, low Tm)
#   x_i = 0, Q_ii <  0 → T (weak, but different stacking from A)
#   x_i = 1, Q_ii >= 0 → G (strong, high Tm)
#   x_i = 1, Q_ii <  0 → C (strong, but different stacking from G)

def choose_base(x_i: int, q_ii: float) -> str:
    """Choose DNA base encoding variable value and diagonal sign.

    The base choice determines:
        1. Per-base Tm contribution (high for x=1, low for x=0)
        2. Stacking behavior with neighbors (encodes Q_ii sign)

    Args:
        x_i: variable value (0 or 1)
        q_ii: diagonal QUBO coefficient

    Returns:
        One of A, T, G, C
    """
    if x_i == 0:
        return "A" if q_ii >= 0 else "T"
    else:
        return "G" if q_ii >= 0 else "C"


# Value extraction from base
BASE_IS_HIGH = {"A": False, "T": False, "G": True, "C": True}


def base_to_value(base: str) -> int:
    """Extract binary value from base (0 for low-Tm, 1 for high-Tm)."""
    return 1 if BASE_IS_HIGH.get(base, False) else 0


# ============================================================
# §3  ENCODING: QUBO solution → NN-aware DNA
# ============================================================

@dataclass
class NNQuboCandidate:
    """A QUBO solution encoded with nearest-neighbor awareness."""
    x: List[int]
    sequence: str
    energy: float
    tm_proxy: float         # NN Tm proxy
    nn_stacking_energy: float  # total NN stacking contribution
    per_base_energy: float     # total per-base contribution

    def to_dict(self) -> dict:
        return {
            "x": self.x,
            "sequence": self.sequence,
            "energy": round(self.energy, 6),
            "tm_proxy": round(self.tm_proxy, 4),
            "nn_stacking": round(self.nn_stacking_energy, 4),
            "per_base": round(self.per_base_energy, 4),
        }


def encode_nn_candidate(
    x: List[int],
    Q: List[List[float]],
) -> NNQuboCandidate:
    """Encode a QUBO solution with NN-aware base selection.

    Each variable's base encodes its value AND the Q_ii sign.
    Dinucleotide stacking captures adjacent Q_ij interactions.

    Args:
        x: binary solution vector
        Q: QUBO matrix

    Returns:
        NNQuboCandidate with all metadata
    """
    n = len(x)

    # Choose bases
    bases = [choose_base(x[i], Q[i][i]) for i in range(n)]
    sequence = "".join(bases)

    # Compute QUBO energy
    energy = sum(Q[i][j] * x[i] * x[j] for i in range(n) for j in range(n))

    # Compute NN Tm proxy
    tm_proxy = nn_tm_estimate(sequence)

    # Decompose: per-base vs stacking
    per_base = sum(NN_BASE_ENTHALPY.get(b, 0) for b in bases)
    stacking = sum(
        nn_stack_dg(bases[i], bases[i + 1])
        for i in range(n - 1)
    )

    return NNQuboCandidate(
        x=x,
        sequence=sequence,
        energy=energy,
        tm_proxy=tm_proxy,
        nn_stacking_energy=-stacking,  # negate for stability proxy
        per_base_energy=-per_base,
    )


# ============================================================
# §4  DECODING: DNA → binary vector
# ============================================================

def decode_nn_sequence(sequence: str) -> List[int]:
    """Decode a NN-encoded DNA sequence back to binary vector.

    Args:
        sequence: DNA sequence from encode_nn_candidate

    Returns:
        Binary vector
    """
    return [base_to_value(b) for b in sequence]


# ============================================================
# §5  CANDIDATE GENERATION
# ============================================================

def generate_nn_candidates(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
) -> List[NNQuboCandidate]:
    """Generate QUBO candidates with NN encoding.

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

    Returns:
        List of NNQuboCandidate
    """
    if n_samples == 0 and n_vars <= 12:
        candidates = []
        for i in range(2**n_vars):
            x = [(i >> j) & 1 for j in range(n_vars)]
            candidates.append(encode_nn_candidate(x, Q))
        return candidates
    else:
        n_samples = n_samples or 10000
        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_nn_candidate(list(x), Q))
        return candidates


# ============================================================
# §6  SORTING AND VERIFICATION
# ============================================================

def sort_by_tm_proxy(candidates: List[NNQuboCandidate]) -> List[NNQuboCandidate]:
    """Sort by NN Tm proxy (ascending). Lower proxy = lower energy."""
    return sorted(candidates, key=lambda c: c.tm_proxy)


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


def spearman_rank_correlation(values_a: List[float], values_b: List[float]) -> float:
    """Compute Spearman rank correlation."""
    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


@dataclass
class NNVerification:
    """Results of NN 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: NNQuboCandidate
    min_tm: NNQuboCandidate
    qubo_type: str  # "banded", "full", "diagonal"

    def to_dict(self) -> dict:
        return {
            "n_candidates": self.n_candidates,
            "n_vars": self.n_vars,
            "qubo_type": self.qubo_type,
            "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 classify_qubo(Q: List[List[float]], n_vars: int) -> str:
    """Classify QUBO matrix structure."""
    # Check if diagonal-only
    off_diag = sum(abs(Q[i][j]) for i in range(n_vars) for j in range(n_vars) if i != j)
    if off_diag < 1e-10:
        return "diagonal"

    # Check if banded (only adjacent interactions)
    banded_energy = sum(abs(Q[i][j]) for i in range(n_vars) for j in range(n_vars) if abs(i - j) == 1)
    if off_diag - banded_energy < 1e-10:
        return "banded"

    return "full"


def verify_nn_sort(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
    top_k: int = 5,
) -> NNVerification:
    """Verify that NN 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

    Returns:
        NNVerification with results
    """
    candidates = generate_nn_candidates(Q, n_vars, n_samples, seed)
    energy_sorted = sort_by_energy(candidates)

    energies = [c.energy for c in candidates]
    tms = [c.tm_proxy for c in candidates]
    raw_corr = spearman_rank_correlation(energies, tms)

    if raw_corr < 0:
        tm_sorted = sorted(candidates, key=lambda c: c.tm_proxy, reverse=True)
    else:
        tm_sorted = sorted(candidates, key=lambda c: c.tm_proxy)

    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 NNVerification(
        n_candidates=len(candidates),
        n_vars=n_vars,
        tm_sort_matches_energy=(energy_sorted[0].x == tm_sorted[0].x),
        rank_correlation=abs(raw_corr),
        top_k_agreement=agreement,
        top_k=top_k,
        min_energy=energy_sorted[0],
        min_tm=tm_sorted[0],
        qubo_type=classify_qubo(Q, n_vars),
    )


# ============================================================
# §7  QUBO REORDERING FOR BANDWIDTH OPTIMIZATION
# ============================================================

def reorder_qubo_for_bandwidth(Q: List[List[float]], n_vars: int) -> Tuple[List[List[float]], List[int]]:
    """Reorder QUBO variables to maximize bandwidth (adjacent interactions).

    Uses Cuthill-McKee-like heuristic: order variables so that
    strongly interacting pairs become adjacent in the sequence.

    This maximizes the fraction of Q_ij captured by NN stacking.

    Args:
        Q: QUBO matrix
        n_vars: number of variables

    Returns:
        (reordered_Q, permutation) where permutation[i] = original index
    """
    # Build interaction graph
    adj = {i: set() for i in range(n_vars)}
    for i in range(n_vars):
        for j in range(n_vars):
            if i != j and abs(Q[i][j]) > 1e-10:
                adj[i].add(j)

    # Greedy BFS ordering: start from node with most connections
    visited = set()
    order = []

    # Sort by degree (most connected first)
    nodes_by_degree = sorted(range(n_vars), key=lambda i: -len(adj[i]))

    for start in nodes_by_degree:
        if start in visited:
            continue
        queue = [start]
        while queue:
            node = queue.pop(0)
            if node in visited:
                continue
            visited.add(node)
            order.append(node)
            # Add neighbors sorted by degree
            neighbors = sorted(adj[node], key=lambda i: -len(adj[i]))
            queue.extend(n for n in neighbors if n not in visited)

    # Build reordered matrix
    Q_reordered = [[0.0] * n_vars for _ in range(n_vars)]
    for i_new, i_old in enumerate(order):
        for j_new, j_old in enumerate(order):
            Q_reordered[i_new][j_new] = Q[i_old][j_old]

    return Q_reordered, order


def apply_permutation(x: List[int], perm: List[int]) -> List[int]:
    """Apply permutation to a vector."""
    return [x[perm[i]] for i in range(len(x))]


def inverse_permutation(perm: List[int]) -> List[int]:
    """Compute inverse permutation."""
    inv = [0] * len(perm)
    for i, p in enumerate(perm):
        inv[p] = i
    return inv


# ============================================================
# §8  FULL PIPELINE WITH REORDERING
# ============================================================

def run_nn_pipeline(
    Q: List[List[float]],
    n_vars: int,
    n_samples: int = 0,
    seed: int = 42,
    reorder: bool = True,
    output_prefix: str = "nn_qubo",
) -> dict:
    """Run the full NN QUBO-DNA pipeline.

    Args:
        Q: QUBO matrix
        n_vars: number of variables
        n_samples: 0 for brute force, else sampling
        seed: RNG seed
        reorder: whether to reorder for bandwidth optimization
        output_prefix: output file prefix

    Returns:
        Summary dict
    """
    print(f"NN QUBO-DNA Pipeline")
    print(f"  Variables: {n_vars}")

    # Reorder for bandwidth if requested
    perm = list(range(n_vars))
    if reorder and n_vars > 2:
        Q_work, perm = reorder_qubo_for_bandwidth(Q, n_vars)
        inv_perm = inverse_permutation(perm)
        print(f"  Reordered for bandwidth: {perm[:10]}{'...' if n_vars > 10 else ''}")
    else:
        Q_work = Q
        inv_perm = perm

    qubo_type = classify_qubo(Q_work, n_vars)
    print(f"  QUBO type: {qubo_type}")

    # Generate candidates on reordered QUBO
    candidates = generate_nn_candidates(Q_work, n_vars, n_samples, seed)
    print(f"  Generated: {len(candidates)} candidates")

    # Calculate correlation to determine sorting direction
    energies = [c.energy for c in candidates]
    tms = [c.tm_proxy 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_proxy, reverse=True)
    else:
        tm_sorted = sorted(candidates, key=lambda c: c.tm_proxy)

    # Map back to original variable ordering
    best_energy_x = apply_permutation(energy_sorted[0].x, inv_perm)
    best_tm_x = apply_permutation(tm_sorted[0].x, inv_perm)

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

    # Verify
    verification = verify_nn_sort(Q_work, n_vars, n_samples, 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}")

    # Energy decomposition for top candidates
    print(f"\n  Energy decomposition (top-3 by energy):")
    for i, cand in enumerate(energy_sorted[:3]):
        print(f"    #{i+1}: E={cand.energy:+.4f} | "
              f"per-base={cand.per_base_energy:+.4f} | "
              f"stacking={cand.nn_stacking_energy:+.4f} | "
              f"Tm_proxy={cand.tm_proxy:.4f} | "
              f"x={apply_permutation(cand.x, inv_perm)}")

    # Write SAM
    sam_lines = [
        "@HD\tVN:1.6\tSO:unsorted",
        f"@PG\tID:nn_qubo_sort\tPN:nn_qubo_sort\tVN:0.2",
        f"@CO\tQUBO type: {qubo_type}, {n_vars} vars, {len(candidates)} candidates",
    ]
    for i, cand in enumerate(candidates):
        x_orig = apply_permutation(cand.x, inv_perm)
        tags = [
            f"TM:f:{cand.tm_proxy:.4f}",
            f"EN:f:{cand.energy:.6f}",
            f"NN:f:{cand.nn_stacking_energy:.4f}",
            f"PB:f:{cand.per_base_energy:.4f}",
            f"XV:Z:{','.join(str(v) for v in x_orig)}",
        ]
        line = f"cand_{i:06d}\t0\t*\t0\t0\t*\t*\t0\t{len(cand.sequence)}\t{cand.sequence}\t{'~' * len(cand.sequence)}\t" + "\t".join(tags)
        sam_lines.append(line)

    sam_path = f"{output_prefix}.sam"
    with open(sam_path, "w") as f:
        f.write("\n".join(sam_lines) + "\n")

    # Write summary
    summary = {
        "n_vars": n_vars,
        "n_candidates": len(candidates),
        "qubo_type": qubo_type,
        "reordered": reorder,
        "permutation": perm,
        "energy_optimal": {
            "x": best_energy_x,
            "energy": energy_sorted[0].energy,
            "tm_proxy": energy_sorted[0].tm_proxy,
        },
        "tm_optimal": {
            "x": best_tm_x,
            "energy": tm_sorted[0].energy,
            "tm_proxy": tm_sorted[0].tm_proxy,
        },
        "verification": verification.to_dict(),
    }

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

    print(f"\n  Output: {sam_path}, {json_path}")
    return summary


# ============================================================
# §9  DEMO QUBOs
# ============================================================

def demo_banded_qubo(n: int = 8, seed: int = 42) -> List[List[float]]:
    """Generate a banded QUBO (only adjacent interactions).

    This is the ideal case for NN encoding — all off-diagonal
    terms are captured by dinucleotide stacking.
    """
    rng = random.Random(seed)
    Q = [[0.0] * n for _ in range(n)]
    for i in range(n):
        Q[i][i] = rng.uniform(1, 5)  # positive diagonal
        if i + 1 < n:
            val = rng.uniform(-2, 2)
            Q[i][i + 1] = val
            Q[i + 1][i] = val  # symmetric
    return Q


def demo_full_qubo(n: int = 8, seed: int = 42) -> List[List[float]]:
    """Generate a full QUBO (all pairs interact).

    NN encoding will only capture adjacent terms.
    Non-adjacent terms are lost — this tests the limits.
    """
    rng = random.Random(seed)
    Q = [[0.0] * n for _ in range(n)]
    for i in range(n):
        Q[i][i] = rng.uniform(1, 5)
        for j in range(i + 1, n):
            val = rng.uniform(-1, 1)
            Q[i][j] = val
            Q[j][i] = val
    return Q


def demo_ising_chain(n: int = 8, h: float = 1.0, J: float = 0.5) -> List[List[float]]:
    """Generate a 1D Ising chain QUBO.

    H = -h Σ x_i - J Σ x_i x_{i+1}

    Perfect for NN encoding: all interactions are adjacent.
    """
    Q = [[0.0] * n for _ in range(n)]
    for i in range(n):
        Q[i][i] = -h
        if i + 1 < n:
            Q[i][i + 1] = -J
            Q[i + 1][i] = -J
    return Q


# ============================================================
# §10  CLI
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("NN QUBO-DNA Sort: Nearest-Neighbor Backdoor")
    print("=" * 60)

    # Demo 1: Banded QUBO (ideal case)
    print("\n--- Demo 1: Banded QUBO (8 vars, brute force) ---")
    Q1 = demo_banded_qubo(8, seed=42)
    r1 = run_nn_pipeline(Q1, n_vars=8, output_prefix="nn_banded8")

    # Demo 2: 1D Ising chain (perfect for NN)
    print("\n--- Demo 2: 1D Ising Chain (8 vars) ---")
    Q2 = demo_ising_chain(8, h=1.0, J=0.5)
    r2 = run_nn_pipeline(Q2, n_vars=8, output_prefix="nn_ising8")

    # Demo 3: Full QUBO (tests limits)
    print("\n--- Demo 3: Full QUBO (8 vars, tests limits) ---")
    Q3 = demo_full_qubo(8, seed=42)
    r3 = run_nn_pipeline(Q3, n_vars=8, output_prefix="nn_full8")

    # Demo 4: Larger banded QUBO (sampling)
    print("\n--- Demo 4: Banded QUBO (12 vars, sampling) ---")
    Q4 = demo_banded_qubo(12, seed=42)
    r4 = run_nn_pipeline(Q4, n_vars=12, n_samples=5000, output_prefix="nn_banded12")

    # Comparison: NN encoding vs naive encoding
    print("\n" + "=" * 60)
    print("COMPARISON: NN vs Naive encoding")
    print("=" * 60)
    from dna_qubo_sort import verify_tm_sort as verify_naive, encode_qubo_candidate as naive_encode

    naive_result = verify_naive(Q1, 8)
    nn_result = r1["verification"]
    print(f"\n  Banded QUBO (8 vars):")
    print(f"  Naive Tm correlation:  {naive_result.rank_correlation:.4f}")
    print(f"  NN Tm correlation:     {nn_result['rank_correlation']:.4f}")
    print(f"  Naive top-5 agreement: {naive_result.top_k_agreement}/5")
    print(f"  NN top-5 agreement:    {nn_result['top_k_agreement']}/5")

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