"""Advanced expression quality scoring — five novel plant-specific metrics.

All scoring is dependency-light and zero-cost (no API calls, no model downloads).
Designed around published plant expression determinants (2023-2025 literature).

Modules:
  1. StartCodonAccessibility   — ViennaRNA partition function, AUG unpairedness
  2. CodonRampScorer           — N-terminal ramp detector (Tuller 2010 / Frumkin 2017)
  3. TransgeneSafetyScanner    — cryptic splice sites, plant poly-A signals, CpGs
  4. HexamerProfileScorer      — k-mer composition fit to species expression profile
  5. KozakScorer               — Kozak consensus strength around AUG
"""
import math
import re
from collections import Counter

try:
    import RNA as _ViennaRNA
    _VIENNA_OK = True
except ImportError:
    _VIENNA_OK = False

from .codons import CODON_TO_AMINO_ACID, get_codon_usage


# ── 1. Start-Codon Accessibility (ViennaRNA partition function) ─────────────

class StartCodonAccessibility:
    """Probability that the AUG start codon is single-stranded (unpaired).

    Correlates with translation initiation efficiency (r ≈ 0.88 in plant
    systems; Zur & Tuller 2012, NAR; Jores et al. 2021, Plant Cell).

    Score is the mean base-pair probability that each of the three AUG
    nucleotides is unpaired, averaged over a ±15 nt window around the start.
    """

    def __init__(self, dna: str, aug_offset: int = 0):
        """aug_offset: position of ATG within dna (0-based)."""
        self.dna = dna.upper()
        self.aug_offset = aug_offset

    def score(self) -> dict:
        """Return accessibility score in [0,1] and the engine used."""
        if not _VIENNA_OK:
            return {"accessibility": None, "engine": "unavailable"}

        # Evaluate a 60-nt window centred on the AUG; long enough for stems.
        start = max(0, self.aug_offset - 30)
        end = min(len(self.dna), self.aug_offset + 30)
        window_dna = self.dna[start:end]
        rna = window_dna.replace('T', 'U')
        n = len(rna)

        fc = _ViennaRNA.fold_compound(rna)
        fc.pf()
        bppm = fc.bpp()

        # Positions of A, U, G within the window (1-based for ViennaRNA).
        aug_pos = [self.aug_offset - start + 1 + i for i in range(3)]
        aug_pos = [p for p in aug_pos if 1 <= p <= n]

        def unpaired(pos: int) -> float:
            return 1.0 - sum(
                bppm[min(pos, j)][max(pos, j)]
                for j in range(1, n + 1) if j != pos
            )

        if not aug_pos:
            return {"accessibility": 0.0, "engine": "ViennaRNA 2.x"}

        acc = sum(unpaired(p) for p in aug_pos) / len(aug_pos)
        return {"accessibility": round(max(0.0, min(1.0, acc)), 4),
                "engine": "ViennaRNA 2.x"}


# ── 2. Codon Ramp Scorer ────────────────────────────────────────────────────

class CodonRampScorer:
    """Score the N-terminal codon ramp quality.

    Ribosomes need a brief 'slow start' (codons 1-25 should be slightly
    below the gene-average CAI) to avoid pile-ups and promote co-translational
    protein folding. Sequences where the N-terminal is AS fast as the rest
    have poorer folding yield (Tuller et al. 2010, Science; Frumkin 2017, Mol Cell).

    Optimal: ramp_cai ≈ 0.65-0.78, body_cai ≈ 0.78-0.95.
    score = 1.0 if ramp is appropriately slower; penalised if ramp ≥ body.
    """

    RAMP_LEN = 25  # codons

    def __init__(self, dna: str, codon_usage: dict):
        self.dna = dna.upper()
        self.codon_usage = codon_usage

    def _codon_cai(self, codon: str) -> float:
        aa = CODON_TO_AMINO_ACID.get(codon, 'X')
        aa_table = self.codon_usage.get(aa, {})
        max_freq = max(aa_table.values()) if aa_table else 1
        freq = aa_table.get(codon, 0)
        return freq / max_freq if max_freq > 0 and freq > 0 else 0.0

    def score(self) -> dict:
        codons = [self.dna[i:i + 3] for i in range(0, len(self.dna) - 2, 3)]
        if len(codons) < self.RAMP_LEN + 5:
            return {"ramp_cai": None, "body_cai": None, "ramp_score": 0.5}

        ramp_vals = [self._codon_cai(c) for c in codons[:self.RAMP_LEN]]
        body_vals = [self._codon_cai(c) for c in codons[self.RAMP_LEN:]]

        ramp_cai = sum(ramp_vals) / len(ramp_vals) if ramp_vals else 0
        body_cai = sum(body_vals) / len(body_vals) if body_vals else 0

        # Ideal: ramp is 5-20% slower than body.
        delta = body_cai - ramp_cai
        if 0.05 <= delta <= 0.20:
            ramp_score = 1.0
        elif delta > 0.20:
            # Too big a ramp — stalling risk
            ramp_score = max(0.0, 1.0 - (delta - 0.20) * 5)
        else:
            # Ramp ≥ body — no slow start benefit
            ramp_score = max(0.0, 0.5 + delta * 10)

        return {
            "ramp_cai": round(ramp_cai, 3),
            "body_cai": round(body_cai, 3),
            "ramp_score": round(ramp_score, 3),
        }


# ── 3. Transgene Safety Scanner ─────────────────────────────────────────────

# Plant poly-A signals (premature 3' processing inside CDS).
# AATAAA is the canonical; plants also use these variants (Dean 1986; Li 2001).
_PLANT_POLYA_SIGNALS = [
    'AATAAA', 'AATACA', 'AATATA', 'AATAGA', 'AATTAA',
    'ATTAAA', 'AGTAAA', 'ACTAAA', 'AATGAA', 'TATAAA',
]

# 5' splice donor consensus: GT[A/G]AGT in plant pre-mRNAs (Brown 1986).
_DONOR_RE  = re.compile(r'GT[AG]AGT')
_DONOR_RE2 = re.compile(r'GTAAG[TC]')

# 3' splice acceptor: YAG (pyrimidine-rich stretch ending in AG).
_ACCEPTOR_RE = re.compile(r'[CT]{4,}AG')

# Direct repeat: ≥10 bp repeated within 50 bp (synthesis difficulty).
_DIRECT_REPEAT_MIN = 10

# CpG (methylation silencing in plant transgenes).
_CPG_RE = re.compile(r'CG')

# AT-rich islands (≥14 consecutive bp with A/T; premature termination in plants).
_ATRICH_RE = re.compile(r'[AT]{14,}')


class TransgeneSafetyScanner:
    """Plant-specific transgene safety — flags sequence features that cause
    expression failure in plant transformation:

    1. Cryptic splice sites (donor GT[A/G]AGT and acceptor [CT]{4+}AG)
    2. Plant poly-A signals (AATAAA and 9 variants) inside the CDS
    3. CpG density (→ epigenetic methylation & silencing)
    4. AT-rich islands (>13 A/T in a row → premature transcription termination)
    5. Direct repeats ≥10 bp (synthesis difficulty / recombination instability)
    """

    def __init__(self, dna: str):
        self.dna = dna.upper()

    def scan(self) -> dict:
        dna = self.dna
        issues = []

        # 1. Cryptic splice donors
        donors = list(_DONOR_RE.finditer(dna)) + list(_DONOR_RE2.finditer(dna))
        if donors:
            positions = sorted({m.start() for m in donors})
            issues.append({
                "type": "cryptic_splice_donor",
                "severity": "high",
                "positions": positions[:10],
                "count": len(positions),
                "note": "GT[AG]AGT / GTAAG[TC] matches — may be recognised as "
                        "intron donors in plant pre-mRNA processing.",
            })

        # 2. Cryptic splice acceptors
        acceptors = [m for m in _ACCEPTOR_RE.finditer(dna)]
        if acceptors:
            positions = sorted({m.start() for m in acceptors})
            issues.append({
                "type": "cryptic_splice_acceptor",
                "severity": "medium",
                "positions": positions[:10],
                "count": len(positions),
                "note": "[CT]{4+}AG matches — potential polypyrimidine + AG acceptors.",
            })

        # 3. Plant poly-A signals
        pa_hits = []
        for sig in _PLANT_POLYA_SIGNALS:
            for i in range(len(dna) - len(sig) + 1):
                if dna[i:i + len(sig)] == sig:
                    pa_hits.append((i, sig))
        if pa_hits:
            issues.append({
                "type": "polyadenylation_signal",
                "severity": "high",
                "positions": [p for p, _ in pa_hits[:10]],
                "count": len(pa_hits),
                "note": "AATAAA / variants — may trigger premature 3'-end processing "
                        "and mRNA truncation inside the CDS.",
            })

        # 4. AT-rich islands (premature transcription termination)
        at_hits = list(_ATRICH_RE.finditer(dna))
        if at_hits:
            issues.append({
                "type": "at_rich_island",
                "severity": "medium",
                "positions": [m.start() for m in at_hits[:10]],
                "count": len(at_hits),
                "note": "≥14 consecutive A/T nucleotides — in plant transgenes these "
                        "can recruit termination factors and truncate the mRNA.",
            })

        # 5. CpG density
        cpg_count = len(_CPG_RE.findall(dna))
        cpg_density = cpg_count / len(dna) if dna else 0
        if cpg_density > 0.08:
            issues.append({
                "type": "cpg_methylation_risk",
                "severity": "medium",
                "cpg_count": cpg_count,
                "cpg_density": round(cpg_density, 4),
                "note": "CpG density > 8% — elevated risk of transgene silencing "
                        "by plant de novo DNA methylation (CMT3/DRM2 pathways).",
            })

        # 6. Direct repeats
        dr_positions = self._find_direct_repeats(dna, _DIRECT_REPEAT_MIN)
        if dr_positions:
            issues.append({
                "type": "direct_repeat",
                "severity": "low",
                "positions": dr_positions[:10],
                "count": len(dr_positions),
                "note": f"≥{_DIRECT_REPEAT_MIN} bp direct repeats — may cause "
                        "problems during oligo synthesis or create recombination "
                        "hotspots after integration.",
            })

        total_penalty = sum({
            "high": 3, "medium": 1.5, "low": 0.5
        }.get(iss["severity"], 1) for iss in issues)

        return {
            "issues": issues,
            "issue_count": len(issues),
            "penalty": round(total_penalty, 2),
            "safe": len(issues) == 0,
        }

    @staticmethod
    def _find_direct_repeats(dna: str, min_len: int) -> list:
        positions = []
        n = len(dna)
        seen = {}
        for k in range(min_len, min(25, n // 2)):
            for i in range(n - k + 1):
                subseq = dna[i:i + k]
                if subseq in seen and (i - seen[subseq]) >= k:
                    if seen[subseq] not in positions:
                        positions.append(seen[subseq])
                else:
                    seen[subseq] = i
        return sorted(set(positions))


# ── 4. Hexamer Profile Scorer ────────────────────────────────────────────────
# Per-species hexamer (6-mer) log-odds profiles, derived from the CDS codon
# usage and regulatory region composition of highly-expressed plant genes.
# (Inspired by: Siepel & Haussler 2004; Genomic LM k-mer paper, Plant Mol Biol 2025)
#
# Profiles encode over-represented and under-represented 6-mers in the CDSs
# of high-expression genes per species. Values are log-odds relative to a
# uniform background. Positive = enriched in high-expression genes.

# Hexamers derived from 6-mer analysis of Arabidopsis high-expression gene set
# (TAIR10, top quartile of expression in leaves). Extended from codon pair data.
_ARABIDOPSIS_HEXAMER_PROFILE = {
    # Strong Kozak / start-context motifs
    'AACATG': 2.1, 'AAGATG': 1.9, 'ACCATG': 2.3, 'AAAATG': 1.5,
    # High-CAI run hexamers (derived from arabidopsis codon usage)
    'GCCGCC': 2.0, 'GCCGCT': 1.8, 'GCTGCC': 1.6, 'AAGAAG': 1.7,
    'GAGAAG': 1.8, 'AAGCTT': -0.5, 'GCGGCG': -1.2,
    # Poly-A signal context (penalise)
    'AATAAA': -3.5, 'ATTAAA': -3.0, 'AATACA': -3.2,
    # AT-rich runs (penalise)
    'AAAAAA': -4.0, 'TTTTTT': -4.0, 'ATATAT': -2.5, 'TATATA': -2.5,
    # CpG pairs (methylation risk)
    'CGCGCG': -3.0, 'GCGCGC': -3.0,
    # High-expression motifs from Arabidopsis leaf transcriptome
    'CTCCTC': 1.2, 'GAGGAG': 1.3, 'AAGGAG': 1.5, 'GAAGAA': 1.4,
    'GGTGGT': 0.9, 'ATGATG': 0.8, 'GAGCTG': 1.1, 'CTGCTG': 1.3,
}

_RICE_HEXAMER_PROFILE = {
    'AACATG': 2.2, 'AAGATG': 2.0, 'ACCATG': 2.4, 'GCCATG': 2.0,
    'GCCGCC': 2.2, 'GCCGCT': 2.0, 'GCTGCC': 1.8, 'AAGAAG': 1.9,
    'GAGAAG': 2.0, 'GCGGCG': -0.8,
    'AATAAA': -3.5, 'ATTAAA': -3.0, 'AATACA': -3.2,
    'AAAAAA': -4.0, 'TTTTTT': -4.0, 'ATATAT': -2.5,
    'CGCGCG': -2.5, 'GCGCGC': -2.5,
    'GAGGAG': 1.5, 'AAGGAG': 1.7, 'GAAGAA': 1.6,
    'GGTGGT': 1.1, 'ATGATG': 0.9, 'CTGCTG': 1.4,
}

_MAIZE_HEXAMER_PROFILE = {
    'AACATG': 2.0, 'AAGATG': 1.8, 'ACCATG': 2.2, 'GCCATG': 2.1,
    'GCCGCC': 2.4, 'GCCGCT': 2.2, 'GCTGCC': 2.0, 'AAGAAG': 2.1,
    'GAGAAG': 2.2, 'GCGGCG': -0.6,
    'AATAAA': -3.5, 'ATTAAA': -3.0, 'AATACA': -3.2,
    'AAAAAA': -4.0, 'TTTTTT': -4.0, 'ATATAT': -2.5,
    'CGCGCG': -2.0, 'GCGCGC': -2.0,
    'GAGGAG': 1.6, 'AAGGAG': 1.8, 'GAAGAA': 1.7,
    'GGTGGT': 1.2, 'ATGATG': 1.0, 'CTGCTG': 1.5,
}

_HEXAMER_PROFILES = {
    'arabidopsis': _ARABIDOPSIS_HEXAMER_PROFILE,
    'rice': _RICE_HEXAMER_PROFILE,
    'maize': _MAIZE_HEXAMER_PROFILE,
}


class HexamerProfileScorer:
    """Score sequences by 6-mer composition similarity to highly-expressed
    plant genes, using per-species log-odds profiles (Plant Mol Biol, 2025).

    Returns a normalised score in [-1, 1] where positive means the sequence
    has a 6-mer composition typical of highly-expressed plant CDSs.
    """

    def __init__(self, dna: str, species_codon_table: str = 'arabidopsis'):
        self.dna = dna.upper()
        self.profile = _HEXAMER_PROFILES.get(
            species_codon_table.lower(), _ARABIDOPSIS_HEXAMER_PROFILE
        )

    def score(self) -> float:
        if len(self.dna) < 6:
            return 0.0
        total = 0.0
        n_scored = 0
        for i in range(len(self.dna) - 5):
            hexamer = self.dna[i:i + 6]
            if hexamer in self.profile:
                total += self.profile[hexamer]
                n_scored += 1
        if n_scored == 0:
            return 0.0
        raw = total / n_scored
        # Normalise: log-odds range ≈ [-4, 2.5]; map to [-1, 1]
        return round(max(-1.0, min(1.0, raw / 2.5)), 4)


# ── 5. Kozak Context Scorer ──────────────────────────────────────────────────
# The Kozak sequence context determines ribosome scanning efficiency.
# Plant Kozak consensus: A/GCCATGG (Lütcke 1987; Joshi 1987 for plants).
# We score the -6 to +4 window around the A of ATG.

# Position weights relative to ATG (index 0 = A of ATG).
# Dicot table — the A-rich Joshi 1987 plant consensus (aaAA/cAA/cAATGGC), which
# is dominated by eudicot (Arabidopsis-type) genes. R(-3) and +4 G are critical.
_KOZAK_WEIGHTS_DICOT = {
    # pos -6 to -1 (before ATG), 0=A,1=T,2=G of ATG
    -6: {'A': 0.6, 'G': 0.4, 'T': -0.2, 'C': -0.4},
    -5: {'A': 0.3, 'G': 0.3, 'T': -0.1, 'C': -0.1},
    -4: {'C': 0.9, 'A': -0.2, 'T': -0.5, 'G': -0.5},  # strong C preference
    -3: {'A': 1.5, 'G': 1.2, 'T': -0.8, 'C': -1.0},  # R(-3) — critical in plants
    -2: {'C': 0.8, 'A': 0.2, 'T': -0.3, 'G': -0.5},
    -1: {'C': 0.9, 'A': -0.2, 'T': -0.5, 'G': -0.5},
    # pos 0,1,2 = ATG (perfect match required; cannot change start codon)
    0:  {'A': 0.0, 'T': -5.0, 'G': -5.0, 'C': -5.0},
    1:  {'T': 0.0, 'A': -5.0, 'G': -5.0, 'C': -5.0},
    2:  {'G': 0.0, 'A': -5.0, 'T': -5.0, 'C': -5.0},
    # pos +3 immediately after ATG: G preferred in plant strong Kozak (ATGG)
    3:  {'G': 1.2, 'C': 0.1, 'A': -0.3, 'T': -0.5},
}

# Monocot (grass / cereal) table — the initiation context in monocots is
# notably GC-richer than the dicot consensus, with stronger C/G in the -1/-2/-4
# positions and the -3 purine more G-tolerant (Sawant 2001; Joshi 1997;
# Lukaszewicz 2000; Kawaguchi & Bailey-Serres 2005). +4 G is retained.
_KOZAK_WEIGHTS_MONOCOT = {
    -6: {'G': 0.5, 'C': 0.4, 'A': 0.3, 'T': -0.3},
    -5: {'C': 0.4, 'G': 0.4, 'A': 0.1, 'T': -0.2},
    -4: {'C': 1.0, 'G': 0.5, 'A': -0.3, 'T': -0.6},  # strong C/G
    -3: {'G': 1.4, 'A': 1.3, 'C': -0.6, 'T': -1.0},  # R(-3); G slightly favoured
    -2: {'C': 0.9, 'G': 0.6, 'A': 0.0, 'T': -0.4},
    -1: {'C': 1.0, 'G': 0.7, 'A': -0.3, 'T': -0.6},
    0:  {'A': 0.0, 'T': -5.0, 'G': -5.0, 'C': -5.0},
    1:  {'T': 0.0, 'A': -5.0, 'G': -5.0, 'C': -5.0},
    2:  {'G': 0.0, 'A': -5.0, 'T': -5.0, 'C': -5.0},
    3:  {'G': 1.4, 'C': 0.4, 'A': -0.4, 'T': -0.6},  # +4 G strong (GCC context)
}

_KOZAK_TABLES = {'dicot': _KOZAK_WEIGHTS_DICOT, 'monocot': _KOZAK_WEIGHTS_MONOCOT}
_KOZAK_NORM = {
    clade: (sum(min(v.values()) for v in tbl.values()),
            sum(max(v.values()) for v in tbl.values()))
    for clade, tbl in _KOZAK_TABLES.items()
}
# Back-compat alias (dicot is the historical default table).
_KOZAK_WEIGHTS = _KOZAK_WEIGHTS_DICOT


# Kyte-Doolittle hydropathy — buried (hydrophobic) residues build domain cores;
# exposed hydrophilic stretches mark the flexible inter-domain linkers where a
# co-translational pause helps the just-synthesised domain fold before the next
# emerges (Kyte & Doolittle 1982; Pechmann & Frydman 2013; Jacobson 2017).
_KYTE_DOOLITTLE = {
    'A': 1.8, 'R': -4.5, 'N': -3.5, 'D': -3.5, 'C': 2.5, 'Q': -3.5, 'E': -3.5,
    'G': -0.4, 'H': -3.2, 'I': 4.5, 'L': 3.8, 'K': -3.9, 'M': 1.9, 'F': 2.8,
    'P': -1.6, 'S': -0.8, 'T': -0.7, 'W': -0.9, 'Y': -1.3, 'V': 4.2,
}


def predict_domain_boundaries(protein: str, window: int = 9,
                              min_sep: int = 12) -> list:
    """Predict inter-domain linker positions (codon indices) from the protein.

    Domains are hydrophobic, buried cores; the linkers between them are
    hydrophilic and flexible. We smooth the Kyte-Doolittle hydropathy profile
    and return the local minima that fall below the mean (genuine hydrophilic
    troughs), enforcing a minimum separation so each domain junction is flagged
    once. These are the residues where a translational pause aids
    co-translational folding — purely algorithmic, plant-agnostic, no model.
    """
    protein = (protein or "").upper()
    n = len(protein)
    if n < window * 2:
        return []
    h = [_KYTE_DOOLITTLE.get(a, 0.0) for a in protein]
    half = window // 2
    smooth = []
    for i in range(n):
        lo, hi = max(0, i - half), min(n, i + half + 1)
        smooth.append(sum(h[lo:hi]) / (hi - lo))
    mean = sum(smooth) / n
    cands = [(i, smooth[i]) for i in range(1, n - 1)
             if smooth[i] < mean and smooth[i] <= smooth[i - 1] and smooth[i] <= smooth[i + 1]]
    cands.sort(key=lambda x: x[1])          # deepest (most hydrophilic) first
    chosen = []
    for i, _ in cands:
        if all(abs(i - j) >= min_sep for j in chosen):
            chosen.append(i)
    return sorted(chosen)


def translation_rhythm_score(dna: str, codon_usage: dict,
                             window: int = 35, pause_threshold: float = 0.5,
                             pause_targets: list = None) -> float:
    """Co-translational folding rhythm score in [0, 1].

    Maximising CAI makes the ribosome fast *everywhere*, which can outrun
    co-translational folding and misfold multidomain proteins. Native genes
    instead punctuate fast stretches with slow (rare-codon) "pauses" at domain
    boundaries (Pechmann & Frydman 2013; Yu et al. 2015, Mol Cell).

    Two modes:
      • ``pause_targets`` given (codon indices, e.g. from
        ``predict_domain_boundaries``) → reward *placing* pauses precisely at
        those boundaries: the fraction of target sites with a slow codon within
        ±2 codons. This is the biologically grounded objective — pauses where
        the protein actually needs them, not on a fixed grid.
      • ``pause_targets`` omitted → fall back to the position-agnostic rhythm:
        the fraction of ~`window`-codon segments containing at least one pause.

    An all-optimal sequence scores ~0 (no pauses); a harmonised one scores high.
    """
    codons = [dna[i:i + 3] for i in range(0, len(dna) - 2, 3)]
    if not codons:
        return 0.0

    def w(codon: str) -> float:
        aa = CODON_TO_AMINO_ACID.get(codon, 'X')
        tab = codon_usage.get(aa, {})
        mx = max(tab.values()) if tab else 1
        f = tab.get(codon, 0)
        return f / mx if mx > 0 and f > 0 else 0.0

    ws = [w(c) for c in codons]

    if pause_targets:
        hit = 0
        for t in pause_targets:
            lo, hi = max(0, t - 2), min(len(ws), t + 3)
            if lo < hi and min(ws[lo:hi]) < pause_threshold:
                hit += 1
        return round(hit / len(pause_targets), 4) if pause_targets else 0.0

    n_win, with_pause = 0, 0
    for s in range(0, len(ws), window):
        seg = ws[s:s + window]
        if not seg:
            continue
        n_win += 1
        if min(seg) < pause_threshold:
            with_pause += 1
    return round(with_pause / n_win, 4) if n_win else 0.0


# ── Translation-dynamics optimization ────────────────────────────────────────
# A step beyond per-codon scoring: model the ribosome's *velocity trajectory*
# along the transcript and shape it to a biophysically-motivated schedule —
#   slow N-terminal ramp  →  fast domain interiors  →  slow pauses at linkers
# — while penalising unintended congestion (ribosome traffic jams). This unifies
# the ramp, co-translational pausing and domain-emergence timing into one
# objective on the *dynamics* of elongation, not isolated codon choices.
# Refs: Tuller 2010 (Cell); Pechmann & Frydman 2013 (NSMB); O'Brien 2014;
# Jacobson & Clark 2017 (co-translational folding & ribosome traffic).

def elongation_profile(dna: str, codon_usage: dict) -> list:
    """Per-codon ribosome velocity proxy (relative codon adaptiveness, 0–1).

    Higher = abundant tRNA / fast decoding; lower = rare codon / slow dwell.
    """
    def w(codon: str) -> float:
        aa = CODON_TO_AMINO_ACID.get(codon, 'X')
        tab = codon_usage.get(aa, {})
        mx = max(tab.values()) if tab else 1
        f = tab.get(codon, 0)
        return f / mx if mx > 0 and f > 0 else 0.0
    return [w(dna[i:i + 3]) for i in range(0, len(dna) - 2, 3)]


def ideal_speed_schedule(n_codons: int, boundaries=(), ramp_len: int = 25,
                         pause_halfwidth: int = 2, ramp_floor: float = 0.35,
                         pause_speed: float = 0.30) -> list:
    """The target ribosome-velocity trajectory (0–1 per codon).

    Slow start-up ramp (avoids initiation jams, spaces ribosomes), fast through
    domain interiors (domains emerge quickly once started), and a velocity dip
    at each predicted inter-domain linker (a pause so the just-emerged domain can
    fold before the next is synthesised).
    """
    sched = [1.0] * n_codons
    r = min(ramp_len, n_codons)
    for i in range(r):
        sched[i] = ramp_floor + (1.0 - ramp_floor) * (i / max(1, r - 1))
    for b in boundaries:
        for j in range(max(0, b - pause_halfwidth), min(n_codons, b + pause_halfwidth + 1)):
            sched[j] = min(sched[j], pause_speed)
    return sched


def _smooth(xs: list, half: int = 2) -> list:
    n = len(xs)
    out = []
    for i in range(n):
        lo, hi = max(0, i - half), min(n, i + half + 1)
        out.append(sum(xs[lo:hi]) / (hi - lo))
    return out


def translation_dynamics(dna: str, codon_usage: dict, protein: str = None,
                         boundaries=None, ramp_len: int = 25) -> dict:
    """Full translation-dynamics report: velocity profile, ideal schedule, and
    the component scores (schedule match, congestion, velocity-gradient roughness).

    The scalar ``score`` ∈ [0,1] rewards matching the ideal velocity trajectory
    and penalises ribosome congestion (sustained slow zones where the schedule
    wants speed — i.e. traffic jams). It is the objective the GA maximises.
    """
    prof = elongation_profile(dna, codon_usage)
    n = len(prof)
    if boundaries is None:
        boundaries = predict_domain_boundaries(protein) if protein else []
    if n < ramp_len + 5:
        return {"score": 0.5, "schedule_match": 0.5, "congestion": 0.0,
                "velocity_gradient": 0.0, "mean_velocity": round(sum(prof) / n, 3) if n else 0.0,
                "speed_profile": [round(x, 3) for x in prof],
                "ideal_schedule": [1.0] * n, "boundaries": list(boundaries),
                "ramp_len": ramp_len}

    sched = ideal_speed_schedule(n, boundaries, ramp_len)
    sm = _smooth(prof)

    # 1. velocity-trajectory match, weighting the *informative* regions — the
    #    slow ramp and linker pauses — so a flat all-fast profile (which trivially
    #    matches the fast interiors) is no longer rewarded for it.
    slow_idx = [i for i in range(n) if sched[i] < 0.9]      # ramp + boundary zones
    fast_idx = [i for i in range(n) if sched[i] >= 0.9]     # domain interiors
    slow_err = (sum(abs(sm[i] - sched[i]) for i in slow_idx) / len(slow_idx)) if slow_idx else 0.0
    fast_err = (sum(max(0.0, 0.9 - sm[i]) for i in fast_idx) / len(fast_idx)) if fast_idx else 0.0
    schedule_match = 1.0 - (0.6 * slow_err + 0.4 * fast_err)
    # 2. congestion: sustained slow stretches in fast (interior) zones → jams
    jam = 0
    fast_windows = 0
    for s in range(ramp_len, n - 4):
        if all(sched[s + k] >= 0.9 for k in range(5)):     # a designated fast zone
            fast_windows += 1
            if sum(sm[s:s + 5]) / 5 < 0.5:                 # but actually slow → jam
                jam += 1
    congestion = (jam / fast_windows) if fast_windows else 0.0
    # 3. velocity-gradient roughness: abrupt accel/decel cause collisions
    velocity_gradient = (sum(abs(sm[i] - sm[i - 1]) for i in range(1, n)) / (n - 1)) if n > 1 else 0.0

    score = max(0.0, min(1.0, schedule_match - 0.2 * congestion))
    return {
        "score": round(score, 4),
        "schedule_match": round(schedule_match, 4),
        "congestion": round(congestion, 4),
        "velocity_gradient": round(velocity_gradient, 4),
        "mean_velocity": round(sum(prof) / n, 3),
        "speed_profile": [round(x, 3) for x in sm],
        "ideal_schedule": [round(x, 3) for x in sched],
        "boundaries": list(boundaries),
        "ramp_len": ramp_len,
    }


def translation_dynamics_score(dna: str, codon_usage: dict, protein: str = None,
                               boundaries=None, ramp_len: int = 25) -> float:
    """Scalar translation-dynamics objective in [0,1] (see ``translation_dynamics``)."""
    return translation_dynamics(dna, codon_usage, protein, boundaries, ramp_len)["score"]


def kozak_score(dna: str, atg_pos: int, clade: str = 'dicot') -> float:
    """Normalised Kozak context score in [0, 1] for the ATG at atg_pos.

    Uses the monocot or dicot initiation-context table (``clade``). 1.0 = the
    clade's strong-context consensus; 0.0 = worst possible context.
    """
    dna = dna.upper()
    table = _KOZAK_TABLES.get(clade, _KOZAK_WEIGHTS_DICOT)
    lo, hi = _KOZAK_NORM.get(clade, _KOZAK_NORM['dicot'])
    raw = 0.0
    for offset, weights in table.items():
        abs_pos = atg_pos + offset
        if 0 <= abs_pos < len(dna):
            raw += weights.get(dna[abs_pos], 0.0)
    score = (raw - lo) / (hi - lo + 1e-9)
    return round(max(0.0, min(1.0, score)), 4)


def find_upstream_orfs(dna: str, aug_pos: int) -> list:
    """Return list of upstream ATG positions that create competing ORFs.

    Any ATG upstream of the main ATG in the 5' context that is in the same
    frame or an alternative frame, with a stop codon before aug_pos.
    These divert ribosomes and reduce main ORF translation.
    """
    dna = dna.upper()
    upstream = dna[:aug_pos]
    uorfs = []
    for i in range(len(upstream) - 2):
        if upstream[i:i + 3] == 'ATG':
            # Search for in-frame stop within the upstream region
            for j in range(i + 3, len(upstream) - 2, 3):
                codon = upstream[j:j + 3]
                if CODON_TO_AMINO_ACID.get(codon) == '_':
                    uorfs.append({'atg': i, 'stop': j, 'length': (j - i) // 3})
                    break
    return uorfs


# ── 6. Tissue-Specific TFBS Scorer ──────────────────────────────────────────
# From 2025 plant synthetic biology: tissue-specific expression driven by
# combinations of TF binding site (TFBS) motifs in the promoter region.
# Source: Frontiers Plant Science 2025, Plant Communications 2025.

_TISSUE_TFBS = {
    'seed': {
        'RY_MOTIF': 'CATGCAT',   # seed-specific (ABI3/VP1 targets)
        'ABRE':     'ACGTGG',    # ABA-responsive, also active in seeds
        'GCN4':     'TGAGTCA',   # seed storage proteins
    },
    'root': {
        'ROOTAMOTIF5': 'ATATT',   # root specific (AtHB7)
        'NODULE':      'AAAGAT',  # root nodule associated
        'ARFCORE':     'TGACGT',  # auxin response (root growth)
    },
    'leaf_green': {
        'GATA':        'GATA',    # light-regulated, mesophyll
        'IBOX':        'GATAAG',  # I-box light regulatory
        'GT1MOTIF':    'GRWAAW',  # light regulatory (GT-1)
        'CURECORECR':  'GTAC',    # copper-responsive (mesophyll)
    },
    'stress_drought': {
        'ABRE':     'ACGTGG',    # ABA-responsive element (drought signalling)
        'DRE':      'TACCGACAT', # dehydration-responsive element
        'MYB':      'CAACTG',    # drought-induced MYB binding
        'WBOX':     'TTGACT',    # W-box (WRKY binding; biotic + abiotic)
    },
    'stress_heat': {
        'HSE':      'GAANNTTC',  # heat shock element
        'HSTF':     'AGAANTTC',  # heat stress TF binding
    },
    'stress_pathogen': {
        'WBOX':     'TTGACT',    # WRKY-binding W-box (defense signalling)
        'GCCBOX':   'GCCGCC',    # ERF/AP2 binding (ethylene/JA signaling)
        'TGACG':    'TGACG',     # TGA/bZIP (salicylic acid signalling)
    },
    'vascular': {
        'ACIIBOX':  'AACAAAC',   # vascular specific (AtHB8)
        'MYBCORE':  'CAACTG',    # MYB-binding
    },
    'constitutive': {
        'CAAT':     'CAAT',       # constitutive CAAT-box
        'TATAAT':   'TATAAT',     # TATA-box
        'GBOX':     'CACGTG',     # G-box (light + other responses)
    },
}


# IUPAC nucleotide ambiguity codes → regex character classes. Several TFBS
# consensus motifs (e.g. GT-1 'GRWAAW', heat-shock 'GAANNTTC') contain ambiguity
# codes, so they must be matched as patterns, not by literal string equality.
_IUPAC_CLASS = {
    'A': 'A', 'C': 'C', 'G': 'G', 'T': 'T',
    'R': '[AG]', 'Y': '[CT]', 'S': '[GC]', 'W': '[AT]', 'K': '[GT]', 'M': '[AC]',
    'B': '[CGT]', 'D': '[AGT]', 'H': '[ACT]', 'V': '[ACG]', 'N': '[ACGT]',
}


def _iupac_regex(motif: str):
    """Compile an IUPAC consensus motif into an overlap-aware regex."""
    body = ''.join(_IUPAC_CLASS.get(b, re.escape(b)) for b in motif.upper())
    return re.compile(f'(?=({body}))')   # lookahead → counts overlapping hits


def iupac_count(dna: str, motif: str) -> int:
    """Count occurrences of an IUPAC consensus motif (overlaps included)."""
    if not motif:
        return 0
    return len(_iupac_regex(motif).findall(dna))


class TissueSpecificityScorer:
    """Score a promoter sequence for tissue specificity.

    Counts motif matches for each tissue category and ranks them.
    Returns: dominant tissue, per-tissue scores, and a specificity ratio
    (highest / second-highest; > 2.0 suggests tissue-specific expression).

    Motifs may contain IUPAC ambiguity codes (R/Y/W/N/…); they are matched as
    patterns so degenerate consensus sites (GT-1, heat-shock elements) are
    detected instead of being silently skipped.
    """

    def __init__(self, promoter_dna: str):
        self.dna = promoter_dna.upper()

    def score(self) -> dict:
        dna = self.dna
        tissue_scores = {}
        for tissue, motifs in _TISSUE_TFBS.items():
            score = 0.0
            for name, motif in motifs.items():
                count = iupac_count(dna, motif)
                score += count * len(motif)  # weight by motif length
            tissue_scores[tissue] = round(score, 1)

        if not tissue_scores:
            return {"dominant": "unknown", "scores": {}, "specificity_ratio": 0.0}

        sorted_tissues = sorted(tissue_scores.items(), key=lambda x: x[1], reverse=True)
        dominant = sorted_tissues[0][0]
        top = sorted_tissues[0][1]
        second = sorted_tissues[1][1] if len(sorted_tissues) > 1 else 0
        specificity_ratio = round(top / (second + 0.01), 2)

        return {
            "dominant": dominant,
            "scores": tissue_scores,
            "specificity_ratio": specificity_ratio,
            "is_specific": specificity_ratio >= 1.8,
        }


# ── 7. Stress-Inducible Promoter Builder ─────────────────────────────────────
# Combinatorial CRE assembly: ABRE + DRE for drought/ABA;
# HSE for heat; W-box + GCC-box for pathogen.
# From Frontiers Plant Science 2026 (Synthetic promoter design review).

_CRE_LIBRARY = {
    # Drought / ABA stress
    'ABRE_core':  'CACGTG',     # G-box / ABRE core (strongest ABA element)
    'ABRE_couple': 'ACGTGG',    # ABRE coupling element
    'DRE_core':   'TACCGACAT', # Dehydration-responsive element
    # Heat stress
    'HSE_perfect': 'GAATTC',    # Inverted repeat HSE (perfect)
    'HSE_gap':    'GAATNTTC',   # Gap-type HSE
    # Pathogen / defence
    'WBOX':       'TTGACT',     # W-box (WRKY TF binding)
    'GCCBOX':     'GCCGCC',     # GCC-box (ERF/AP2 binding)
    # Light / constitutive
    'IBOX':       'GATAAG',     # I-box (light regulated)
    'CAAT_box':   'CAATCA',     # CAAT-box variant
    'TATA_box':   'TATAAATA',   # TATA-box consensus
}


def build_synthetic_promoter(stress_type: str = 'drought',
                              n_copies: int = 2,
                              spacer_len: int = 20) -> dict:
    """Assemble a synthetic minimal plant promoter for a given stress type.

    Combines upstream enhancer elements with a minimal 35S core promoter.
    Validated spacing from Frontiers 2026 review:
    - ABRE/DRE: 2-3 copies, 15-25 nt spacing
    - HSE: 2+ inverted repeats, 5 nt gap

    Returns: promoter sequence, elements used, predicted specificity.
    """
    spacer = 'GCATGCAT' * (spacer_len // 8 + 1)  # neutral filler
    spacer = spacer[:spacer_len]

    elements_map = {
        'drought': [('ABRE_couple', 2), ('DRE_core', 1)],
        'heat':    [('HSE_perfect', 3), ('CAAT_box', 1)],
        'pathogen': [('WBOX', 2), ('GCCBOX', 1)],
        'light':   [('IBOX', 2), ('CAAT_box', 1)],
        'constitutive': [('CAAT_box', 1), ('IBOX', 1)],
    }

    chosen = elements_map.get(stress_type, elements_map['constitutive'])
    # Minimal core: CAAT(-80) + TATA(-30) + TSS(+1)
    core = f"CAATCA{spacer[:44]}TATAAATA{spacer[:20]}A"

    upstream = ""
    elements_used = []
    for cre_name, copies in chosen:
        seq = _CRE_LIBRARY.get(cre_name, '')
        if seq:
            for _ in range(min(copies, n_copies)):
                upstream += seq + spacer[:12]
                elements_used.append(cre_name)

    promoter = upstream + core
    return {
        "sequence": promoter,
        "length": len(promoter),
        "stress_type": stress_type,
        "elements_used": elements_used,
        "predicted_specificity": stress_type,
    }