#!/usr/bin/env python3
"""
dna_codec.py — Hachimoji DNA Encoding/Decoding Layer

Translates between binary data and Hachimoji 8-base DNA sequences.
Uses the SilverSight Hachimoji alphabet (A, T, G, C, B, S, P, Z)
at 3 bits per base — 50% denser than standard 2-bit DNA encoding.

Encoding table (from HachimojiBase.lean):
    000 → A (Φ)    001 → T (Λ)    010 → G (Ρ)    011 → C (Κ)
    100 → B (Ω)    101 → S (Σ)    110 → P (Π)    111 → Z (Ζ)

Melting temperature contributions (nearest-neighbor simplified):
    G/C pairs: 3 bonds → high Tm contribution
    A/T pairs: 2 bonds → low Tm contribution
    B/S, P/Z: synthetic pairs, assigned Tm values for encoding purposes
"""

from __future__ import annotations

from typing import List, Tuple

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

# 3-bit → base
BITS_TO_BASE = {
    0b000: "A",
    0b001: "T",
    0b010: "G",
    0b011: "C",
    0b100: "B",
    0b101: "S",
    0b110: "P",
    0b111: "Z",
}

# base → 3-bit
BASE_TO_BITS = {v: k for k, v in BITS_TO_BASE.items()}

# All 8 bases in order
HACHIMOJI_BASES = list("ATGCBSPZ")

# Greek ↔ Latin (from HachimojiBase.lean)
LATIN_TO_GREEK = {
    "A": "Φ", "T": "Λ", "G": "Ρ", "C": "Κ",
    "B": "Ω", "S": "Σ", "P": "Π", "Z": "Ζ",
}
GREEK_TO_LATIN = {v: k for k, v in LATIN_TO_GREEK.items()}

# ============================================================
# §2  MELTING TEMPERATURE MODEL
# ============================================================

# Per-base Tm contribution (°C, simplified nearest-neighbor model)
# Standard bases: G/C = 3 bonds, A/T = 2 bonds
# Synthetic bases: assigned for sorting purposes
TM_CONTRIBUTION = {
    "A": 2.0,   # 2 hydrogen bonds
    "T": 2.0,   # 2 hydrogen bonds
    "G": 3.0,   # 3 hydrogen bonds (canonical Watson-Crick)
    "C": 3.0,   # 3 hydrogen bonds (canonical Watson-Crick)
    "B": 2.5,   # synthetic, intermediate
    "S": 2.5,   # synthetic, intermediate
    "P": 3.5,   # synthetic, higher affinity (hachimoji pair)
    "Z": 3.5,   # synthetic, higher affinity (hachimoji pair)
}

# GC-equivalent weight for sorting (normalized)
# Maps each base to a "GC-equivalent" score for sorting purposes
GC_EQUIVALENT = {
    "A": 0.0,
    "T": 0.0,
    "G": 1.0,
    "C": 1.0,
    "B": 0.5,
    "S": 0.5,
    "P": 1.5,
    "Z": 1.5,
}


def melting_temperature(sequence: str) -> float:
    """Compute simplified melting temperature for a Hachimoji DNA sequence.

    Uses a per-base contribution model (simplified from nearest-neighbor).
    For sorting purposes — relative order matters more than absolute °C.

    Args:
        sequence: Hachimoji DNA string (characters from ATGCBSPZ)

    Returns:
        Melting temperature in °C (simplified)
    """
    if not sequence:
        return 0.0
    base_tm = sum(TM_CONTRIBUTION.get(b, 0.0) for b in sequence)
    # Wallace rule approximation: Tm ≈ sum of per-base contributions
    return base_tm


def gc_content(sequence: str) -> float:
    """Compute GC-equivalent content for a Hachimoji sequence.

    Standard GC content counts G+C / total. For Hachimoji, we use
    the GC-equivalent weight: G/C = 1.0, B/S = 0.5, P/Z = 1.5, A/T = 0.0.

    Args:
        sequence: Hachimoji DNA string

    Returns:
        GC-equivalent content in [0, 1] (normalized)
    """
    if not sequence:
        return 0.0
    total = sum(GC_EQUIVALENT.get(b, 0.0) for b in sequence)
    # Normalize: max possible per base is 1.5 (P/Z), so divide by 1.5 * len
    max_possible = 1.5 * len(sequence)
    return total / max_possible if max_possible > 0 else 0.0


def raw_gc_content(sequence: str) -> float:
    """Standard GC content (G+C only, ignoring synthetic bases)."""
    if not sequence:
        return 0.0
    gc = sum(1 for b in sequence if b in ("G", "C"))
    return gc / len(sequence)


# ============================================================
# §3  ENCODING: bytes → Hachimoji DNA
# ============================================================

def bytes_to_bits(data: bytes) -> List[int]:
    """Convert bytes to a list of bits (MSB first per byte)."""
    bits = []
    for byte in data:
        for i in range(7, -1, -1):
            bits.append((byte >> i) & 1)
    return bits


def bits_to_bytes(bits: List[int]) -> bytes:
    """Convert a list of bits back to bytes (MSB first per byte)."""
    # Pad to multiple of 8
    padded = bits + [0] * ((8 - len(bits) % 8) % 8)
    result = bytearray()
    for i in range(0, len(padded), 8):
        byte = 0
        for j in range(8):
            byte = (byte << 1) | padded[i + j]
        result.append(byte)
    return bytes(result)


def encode_bytes_to_dna(data: bytes) -> str:
    """Encode arbitrary bytes to a Hachimoji DNA sequence.

    Chunks bits into groups of 3, maps each to a base.
    Pads with zeros if needed (last group may be short).

    Args:
        data: arbitrary bytes

    Returns:
        Hachimoji DNA string
    """
    bits = bytes_to_bits(data)
    # Pad to multiple of 3
    while len(bits) % 3 != 0:
        bits.append(0)

    sequence = []
    for i in range(0, len(bits), 3):
        chunk = (bits[i] << 2) | (bits[i + 1] << 1) | bits[i + 2]
        sequence.append(BITS_TO_BASE[chunk])

    return "".join(sequence)


def decode_dna_to_bytes(sequence: str) -> bytes:
    """Decode a Hachimoji DNA sequence back to bytes.

    Args:
        sequence: Hachimoji DNA string

    Returns:
        Decoded bytes (may include padding zeros at the end)
    """
    bits = []
    for base in sequence:
        if base not in BASE_TO_BITS:
            raise ValueError(f"Invalid Hachimoji base: {base!r}")
        chunk = BASE_TO_BITS[base]
        bits.append((chunk >> 2) & 1)
        bits.append((chunk >> 1) & 1)
        bits.append(chunk & 1)

    return bits_to_bytes(bits)


# ============================================================
# §4  BINARY VECTOR ENCODING: x ∈ {0,1}^n → DNA
# ============================================================

def encode_binary_vector(x: List[int], bits_per_var: int = 3) -> str:
    """Encode a binary vector as a Hachimoji DNA sequence.

    Each variable gets `bits_per_var` bases (default 3 = one Hachimoji base).
    This ensures every variable has a dedicated position in the sequence.

    Args:
        x: binary vector (list of 0s and 1s)
        bits_per_var: number of bases per variable (default 3)

    Returns:
        Hachimoji DNA string of length len(x) * bits_per_var
    """
    sequence = []
    for val in x:
        # Encode each variable as a base (0 → A, 1 → P for max Tm separation)
        # Using A (low Tm) for 0 and P (high Tm) for 1
        # This makes the Tm directly proportional to the number of 1s
        for _ in range(bits_per_var):
            sequence.append("A" if val == 0 else "P")
    return "".join(sequence)


def decode_binary_vector(sequence: str, bits_per_var: int = 3) -> List[int]:
    """Decode a Hachimoji DNA sequence back to a binary vector.

    Args:
        sequence: Hachimoji DNA string
        bits_per_var: bases per variable (must match encoding)

    Returns:
        Binary vector
    """
    if len(sequence) % bits_per_var != 0:
        raise ValueError(
            f"Sequence length {len(sequence)} not divisible by bits_per_var={bits_per_var}"
        )

    result = []
    for i in range(0, len(sequence), bits_per_var):
        chunk = sequence[i : i + bits_per_var]
        # Majority vote: if most bases are P → 1, else → 0
        p_count = chunk.count("P")
        result.append(1 if p_count > bits_per_var / 2 else 0)

    return result


# ============================================================
# §5  QUBO ENCODING: energy → DNA sequence structure
# ============================================================

def qubo_energy(x: List[int], Q: List[List[float]]) -> float:
    """Compute QUBO energy E(x) = x^T Q x.

    Args:
        x: binary vector
        Q: QUBO matrix (n×n, upper triangular or symmetric)

    Returns:
        Energy value
    """
    n = len(x)
    energy = 0.0
    for i in range(n):
        for j in range(n):
            energy += Q[i][j] * x[i] * x[j]
    return energy


def encode_qubo_solution(
    x: List[int],
    Q: List[List[float]],
    bits_per_var: int = 3,
) -> Tuple[str, float]:
    """Encode a QUBO solution as a DNA sequence with energy metadata.

    The sequence encodes the binary vector. The energy is stored as
    a header (first 8 bases encode Q16_16 energy value).

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

    Returns:
        (dna_sequence, energy)
    """
    from q16_canonical import float_to_q16, q16_to_bytes, q16_from_bytes

    energy = qubo_energy(x, Q)

    # Encode energy as Q16_16 header (8 bases = 24 bits = 3 bytes)
    energy_q16 = float_to_q16(energy)
    energy_bytes = q16_to_bytes(energy_q16)  # 4 bytes
    energy_dna = encode_bytes_to_dna(energy_bytes)  # ~11 bases

    # Encode solution vector
    solution_dna = encode_binary_vector(x, bits_per_var)

    return energy_dna + "|" + solution_dna, energy


def decode_qubo_solution(
    sequence: str,
    n_vars: int,
    bits_per_var: int = 3,
) -> Tuple[List[int], float]:
    """Decode a QUBO solution from DNA sequence.

    Args:
        sequence: DNA sequence with energy header
        n_vars: number of variables in the QUBO
        bits_per_var: bases per variable

    Returns:
        (binary_vector, energy)
    """
    from q16_canonical import q16_from_bytes, q16_to_float

    # Split header and body
    parts = sequence.split("|")
    if len(parts) != 2:
        raise ValueError("Invalid QUBO DNA format: expected energy|solution")

    energy_dna, solution_dna = parts

    # Decode energy header
    energy_bytes = decode_dna_to_bytes(energy_dna)
    energy_q16 = q16_from_bytes(energy_bytes[:4])  # first 4 bytes
    energy = q16_to_float(energy_q16)

    # Decode solution vector
    x = decode_binary_vector(solution_dna, bits_per_var)

    return x[:n_vars], energy


# ============================================================
# §6  CONVENIENCE
# ============================================================

def dna_to_greek(sequence: str) -> str:
    """Convert Latin Hachimoji DNA to Greek notation."""
    return "".join(LATIN_TO_GREEK.get(b, b) for b in sequence)


def greek_to_dna(sequence: str) -> str:
    """Convert Greek Hachimoji notation to Latin DNA."""
    return "".join(GREEK_TO_LATIN.get(g, g) for g in sequence)


def sequence_stats(sequence: str) -> dict:
    """Compute useful statistics for a Hachimoji DNA sequence."""
    return {
        "length": len(sequence),
        "gc_content": round(gc_content(sequence), 4),
        "raw_gc_content": round(raw_gc_content(sequence), 4),
        "melting_temperature": round(melting_temperature(sequence), 2),
        "base_counts": {b: sequence.count(b) for b in HACHIMOJI_BASES if b in sequence},
        "greek": dna_to_greek(sequence),
    }