# geney/mirna/features.py — miBSIM 152-feature computation engine (Python port).
"""
Complete Python port of the miBSIM feature computation pipeline.

Computes 152 features per miRNA binding site across 5 categories:
  1. Sequence (nucleotide composition, AU content, 3' pairing, PUM/RBP sites)
  2. Thermodynamic (ΔG duplex/binding/open, structural accessibility, SPS, MRE counts)
  3. Conservation (phastCons/phyloP for seed/site/flank regions)
  4. Codon usage (CUB, tAI, TDR, charge, slow AA scores)
  5. Biochemical (Kd, occupancy — delegated to mibsim.predict_occupancy)

Ported from MATLAB: calc_gen_features.m and all sub-functions.
"""
from __future__ import annotations

import json
import math
import re
import subprocess
import tempfile
from pathlib import Path
from typing import Optional

import numpy as np
from scipy.stats import binom

# ============================================================================
# Watson-Crick pairing
# ============================================================================

_WC_PAIRS = {("A", "U"), ("U", "A"), ("G", "C"), ("C", "G"),
             ("A", "T"), ("T", "A")}
_WOBBLE_PAIRS = _WC_PAIRS | {("G", "A"), ("A", "G"), ("U", "C"), ("C", "U")}

_RC_MAP = str.maketrans("ACGTU", "UGCAA")


def _rc(seq: str) -> str:
    """Reverse complement for RNA."""
    return seq.upper().translate(_RC_MAP)[::-1]


def is_watson_crick(s1: str, s2: str, wobble: bool = False) -> list[bool]:
    """Check Watson-Crick pairing between two antiparallel RNA strands.

    s2 is reverse-complemented, then compared positionally to s1.
    If s1[i] == RC(s2)[i], the bases form a Watson-Crick pair.
    With wobble=True, G-U pairs are also accepted.
    """
    s1 = s1.upper().replace("T", "U")
    s2_rc = _rc(s2).replace("T", "U")

    if not wobble:
        return [a == b for a, b in zip(s1, s2_rc)]
    else:
        _WOBBLE_OK = {("G", "U"), ("U", "G")}
        return [a == b or (a, b) in _WOBBLE_OK for a, b in zip(s1, s2_rc)]


def _count_wc(s1: str, s2: str) -> int:
    """Count Watson-Crick pairs between two RNA strands."""
    return sum(is_watson_crick(s1, s2))


# ============================================================================
# AU params loading
# ============================================================================

_AU_PARAMS: dict | None = None
_AU_PATH = Path(__file__).parent / "au_params.json"

_SEED_TYPE_TO_AU_KEY = {
    "6mer": "six", "7mer-A1": "sevenA1", "7mer-m8": "sevenm8", "8mer": "eight",
}


def _load_au_params() -> dict:
    global _AU_PARAMS
    if _AU_PARAMS is None:
        with open(_AU_PATH) as f:
            _AU_PARAMS = json.load(f)
    return _AU_PARAMS


# ============================================================================
# Sequence features
# ============================================================================

def calc_au_content(rna: str, rna_start: int, seed_length: int, seed_type: str) -> float:
    """Weighted AU content flanking the seed site."""
    au_key = _SEED_TYPE_TO_AU_KEY.get(seed_type, "six")
    params = _load_au_params()[au_key]
    window = int(params["window_length"])
    up_shift = int(params["up_shift"])
    down_shift = int(params["down_shift"])
    weights_up = np.array(params["weights_up"], dtype=float)
    weights_down = np.array(params["weights_down"], dtype=float)

    seq = rna.upper().replace("T", "U")
    n = len(seq)

    # Upstream window
    up_end = rna_start + seed_length - up_shift
    up_start = max(0, up_end - window)
    seq_up = seq[up_start:up_end]

    # Downstream window
    down_start = rna_start + seed_length + down_shift
    down_end = min(n, down_start + window)
    seq_down = seq[down_start:down_end]

    # Count weighted AU positions
    wu = weights_up[:len(seq_up)]
    wd = weights_down[:len(seq_down)]

    au_up = sum(1.0 / wu[i] for i, c in enumerate(seq_up) if c in "AU")
    au_down = sum(1.0 / wd[i] for i, c in enumerate(seq_down) if c in "AU")

    total_weight = sum(1.0 / wu) + sum(1.0 / wd)
    return (au_up + au_down) / total_weight if total_weight > 0 else 0.0


def calc_utr_pos(rna_start: int, region_start: int, region_end: int,
                 mirna_len: int, seed_length: int) -> tuple[float, float, float]:
    """Log10 distances from site to region termini."""
    mir_end = rna_start + seed_length
    mir_start = mir_end - mirna_len + 1

    d_start = abs(mir_start - region_start + 1)
    d_end = abs(region_end - mir_end + 1)

    dist_start = math.log10(d_start) if d_start > 0 else 0.0
    dist_end = math.log10(d_end) if d_end > 0 else 0.0
    min_dist = min(dist_start, dist_end)

    return min_dist, dist_start, dist_end


def calc_3p_pairing(rna: str, rna_start: int, mirna: str,
                    seed_length: int, mir_start: int = 1) -> float:
    """3' complementary pairing score (miRNA positions 13-16 to target)."""
    seq = rna.upper().replace("T", "U")
    mir = mirna.upper().replace("T", "U")
    n = len(seq)

    seed_end = mir_start + seed_length - 1
    rna_nt1 = rna_start + seed_length - 1 + mir_start - 1

    # Pad if needed
    if rna_nt1 >= n:
        seq = seq + "A" * (rna_nt1 - n + 20)

    min_offset = seed_end - 13 + 1
    max_offset = len(mir) - 16
    best = 0.0

    for offset in range(min_offset, max_offset + 1):
        pos = rna_nt1 - 15 - offset
        if pos < 0:
            continue
        # Core 4-nt pairing at miRNA 13-16
        rna_window = seq[pos:pos + 4]
        mir_window = mir[12:16]
        if len(rna_window) < 4:
            continue
        score = _count_wc(mir_window, rna_window)
        score -= max(0, (abs(offset) - 2) / 2)

        # Extend upstream from position 13
        i = 1
        while 12 - i >= seed_end and pos + 4 + i - 1 < len(seq):
            if not is_watson_crick(mir[12 - i], seq[pos + 4 + i - 1])[0]:
                break
            score += 0.5
            i += 1

        # Extend downstream from position 16
        i = 1
        while 15 + i < len(mir) and pos - i >= 0:
            if not is_watson_crick(mir[15 + i], seq[pos - i])[0]:
                break
            score += 0.5
            i += 1

        best = max(best, score)

    return best


def calc_3p_pairing_sw(rna: str, rna_start: int, mirna: str, seed_type: str) -> float:
    """Sliding window average 3' pairing score."""
    seq = rna.upper().replace("T", "U")
    mir = mirna.upper().replace("T", "U")
    window_len = 7

    if seed_type in ("8mer", "7mer-m8"):
        mir_3p_start = 8  # 0-indexed
    else:
        mir_3p_start = 7

    mir_3p_len = len(mir) - mir_3p_start
    rna_5p_end = rna_start
    rna_5p_start = max(0, rna_5p_end - mir_3p_len)
    rna_5p = seq[rna_5p_start:rna_5p_end]

    if len(rna_5p) < window_len:
        return 0.0

    scores = []
    for i in range(len(rna_5p) - window_len + 1):
        mir_slice = mir[mir_3p_start + i:mir_3p_start + i + window_len]
        rna_slice = rna_5p[i:i + window_len]
        if len(mir_slice) == window_len:
            scores.append(_count_wc(mir_slice, rna_slice))

    return float(np.mean(scores)) if scores else 0.0


_PUM_MOTIFS = [
    r"[ACGU]GUA[ACGU]AU[AU]",  # PUM1-like
    r"[CGU][CU]C[CG][AU]C[CG]C",  # PUM2-like
]


def calc_pum(rna: str, rna_start: int, seed_length: int,
             orf_start: int, orf_end: int) -> tuple[list[float], list[float], list[float]]:
    """Count Pumilio binding motifs in ORF, UTR3, and neighborhood."""
    seq = rna.upper().replace("T", "U")
    n = len(seq)

    orf_pum = [0.0, 0.0]
    utr3_pum = [0.0, 0.0]
    nei_pum = [0.0, 0.0]

    for i, motif in enumerate(_PUM_MOTIFS):
        positions = [m.start() for m in re.finditer(motif, seq)]
        # Filter out seed region ± 8 nt
        positions = [p for p in positions if p <= rna_start - 8 or p > rna_start + seed_length - 1]

        orf_hits = [p for p in positions if orf_start <= p <= orf_end - 8]
        if orf_hits:
            orf_pum[i] = 0.5 * (len(orf_hits) - 1) + 1

        utr3_hits = [p for p in positions if orf_end < p <= n - 8]
        if utr3_hits:
            utr3_pum[i] = 0.5 * (len(utr3_hits) - 1) + 1

        nei_hits = [p for p in positions
                    if rna_start - 400 <= p <= rna_start + seed_length - 1 + 400]
        nei_pum[i] = len(nei_hits)

    return orf_pum, utr3_pum, nei_pum


# 20 RBP binding motifs from TargetScan/miBSIM
RBP_SITES = [
    "UAUUU", "UAUUUAU", "AUUUAU", "GUGAAG", "AUUUAUU",
    "UAUUUA", "UAUUUU", "AGCCA", "AGAGAA", "UUUAU",
    "UUUUU", "UUUUAAA", "UUUUUU", "UUUUUUU", "UUUAAA",
    "UUUAAAA", "UUUUG", "UGUUU", "GAAAA", "CAAAA",
]


# ============================================================================
# Conservation features
# ============================================================================

def calc_conservation(phastcons100: np.ndarray | None, phastcons20: np.ndarray | None,
                      phylops100: np.ndarray | None, phylops20: np.ndarray | None,
                      rna_start: int, seed_length: int) -> dict[str, float]:
    """Extract phastCons/phyloP scores for seed, site, and flanking regions.

    Feature naming matches miBSIM model convention:
      phastcons_seed100, phastcons_site100, phastcons_flank100,
      phylops_site100, phylops_flank100, phylops100_pos1..6, etc.
    """
    out: dict[str, float] = {}

    # (metric_base, alignment_suffix, vector)
    configs = [
        ("phastcons", "100", phastcons100),
        ("phastcons", "20", phastcons20),
        ("phylops", "100", phylops100),
        ("phylops", "20", phylops20),
    ]

    for metric, aln, vec in configs:
        if vec is None or len(vec) == 0:
            out[f"{metric}_seed{aln}"] = 0.0
            out[f"{metric}_site{aln}"] = 0.0
            out[f"{metric}_flank{aln}"] = 0.0
            if metric == "phylops":
                for p in range(1, 7):
                    out[f"{metric}{aln}_pos{p}"] = 0.0
            continue

        n = len(vec)
        seed_start = rna_start + seed_length - 6
        seed_end = rna_start + seed_length
        site_start = rna_start
        site_end = rna_start + seed_length

        seed_idx = list(range(max(0, seed_start), min(n, seed_end)))
        site_idx = list(range(max(0, site_start), min(n, site_end)))
        flank_start = max(0, rna_start - 40)
        flank_end = min(n, rna_start + seed_length + 40)
        flank_idx = list(range(flank_start, rna_start)) + list(range(site_end, flank_end))

        out[f"{metric}_seed{aln}"] = float(np.mean(vec[seed_idx])) if seed_idx else 0.0
        out[f"{metric}_site{aln}"] = float(np.mean(vec[site_idx])) if site_idx else 0.0
        out[f"{metric}_flank{aln}"] = float(np.mean(vec[flank_idx])) if flank_idx else 0.0

        # Per-position phyloP for seed (reversed to match MATLAB convention)
        if metric == "phylops":
            seed_vals = [float(vec[i]) if i < n else 0.0
                         for i in range(seed_end - 1, max(0, seed_start) - 1, -1)]
            for p, val in enumerate(seed_vals[:6], 1):
                out[f"{metric}{aln}_pos{p}"] = val
            # Pad if fewer than 6 positions
            for p in range(len(seed_vals) + 1, 7):
                out[f"{metric}{aln}_pos{p}"] = 0.0

    return out


def calc_prob_binom(rna: str, rna_start: int, seed_length: int) -> float:
    """Binomial probability of observing the target site by chance (Markov order-1)."""
    seq = rna.lower().replace("t", "u")
    target = seq[rna_start:rna_start + seed_length]
    if len(target) < 2:
        return 1.0

    # Build Markov order-1 transition matrix
    bases = "acgu"
    counts = np.zeros((4, 4))
    b2i = {b: i for i, b in enumerate(bases)}

    for i in range(len(seq) - 1):
        a, b = seq[i], seq[i + 1]
        if a in b2i and b in b2i:
            counts[b2i[a], b2i[b]] += 1

    # Normalize rows
    row_sums = counts.sum(axis=1, keepdims=True)
    row_sums[row_sums == 0] = 1
    trans = counts / row_sums

    # Probability of target sequence under Markov model
    p = 1.0
    for i in range(len(target) - 1):
        a, b = target[i], target[i + 1]
        if a in b2i and b in b2i:
            p *= trans[b2i[a], b2i[b]]
        else:
            p *= 0.25

    # Binomial test
    k = seq.count(target)
    n_trials = max(1, len(seq) - seed_length + 1)
    return float(1.0 - binom.cdf(k, n_trials, max(p, 1e-20)))


# ============================================================================
# Thermodynamic features (ViennaRNA wrappers)
# ============================================================================

def _vienna_available() -> bool:
    """Check if ViennaRNA Python API is available."""
    try:
        import RNA  # noqa: F401
        return True
    except ImportError:
        return False


def _run_vienna_cmd(cmd: str, input_text: str) -> str:
    """Run a ViennaRNA command-line tool with stdin input."""
    try:
        result = subprocess.run(
            cmd.split(), input=input_text, capture_output=True, text=True, timeout=30)
        return result.stdout
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return ""


def calc_dg_duplex_binding(rna: str, mirna: str, mir_start: int, rna_pos1: int,
                           rna_start: int, seed_type: str, seed_length: int
                           ) -> tuple[float, float, float, float]:
    """ΔG duplex and binding energy for miRNA-RNA complex via RNAcofold.

    Returns (dg_duplex, dg_duplex_seed, dg_binding, dg_binding_seed).
    All values in kcal/mol. Returns (0,0,0,0) if ViennaRNA unavailable.
    """
    seq = rna.upper().replace("T", "U")
    mir = mirna.upper().replace("T", "U")

    # Determine pairing
    if seed_type in ("8mer", "7mer-A1"):
        pair_rna = seq[rna_start:rna_start + seed_length - 1]
        pair_mir = mir[mir_start:mir_start + seed_length - 1]
    else:
        pair_rna = seq[rna_start:rna_start + seed_length]
        pair_mir = mir[mir_start:mir_start + seed_length]

    pairing = is_watson_crick(pair_rna, pair_mir)

    # Build constraint
    site_const = "".join("(" if p else "." for p in pairing)
    seed_const = "".join(")" if p else "." for p in reversed(pairing))

    len_no_const = len(mir) - 8 if seed_type in ("8mer", "7mer-m8") else len(mir) - 7

    # Target window
    start = max(0, rna_pos1 - len(mir) + 1)
    end = min(len(seq), rna_pos1 + 1)
    target = seq[start:end]
    if len(target) < len(mir):
        target = target + "A" * (len(mir) - len(target))

    # Constraint string
    pre_dots = "." * min(len_no_const, max(0, rna_start - start))
    constraint = pre_dots + site_const + ".&." + seed_const + "." * len_no_const

    # Full duplex with constraints
    inp = f"{target}&{mir}\n{constraint}\n@\n"
    out = _run_vienna_cmd("RNAcofold --constraint --partfunc", inp)
    vals = re.findall(r"-?\d+\.\d+", out)
    dg_duplex = float(vals[0]) if len(vals) >= 1 else 0.0
    dg_binding = float(vals[-1]) if len(vals) >= 2 else 0.0

    # Seed-only (no constraints)
    target_seed = seq[rna_start:rna_start + seed_length]
    if seed_type in ("8mer", "7mer-A1"):
        mir_seed = mir[:seed_length]
    else:
        mir_seed = mir[1:seed_length + 1]
    inp2 = f"{target_seed}&{mir_seed}\n@\n"
    out2 = _run_vienna_cmd("RNAcofold --partfunc", inp2)
    vals2 = re.findall(r"-?\d+\.\d+", out2)
    dg_duplex_seed = float(vals2[0]) if len(vals2) >= 1 else 0.0
    dg_binding_seed = float(vals2[-1]) if len(vals2) >= 2 else 0.0

    return dg_duplex, dg_duplex_seed, dg_binding, dg_binding_seed


def calc_dg_open(rna: str, rna_pos1: int, mirna_length: int) -> float:
    """ΔG open: energy cost to keep target site accessible for RISC."""
    seq = rna.upper().replace("T", "U")
    risc_win = (70 - mirna_length) // 2
    start = max(0, rna_pos1 - mirna_length - risc_win)
    end = min(len(seq), rna_pos1 + risc_win + 1)
    region = seq[start:end]
    if len(region) < end - start:
        region += "A" * (end - start - len(region))

    # Unconstrained
    out1 = _run_vienna_cmd("RNAfold --partfunc", f"{region}\n@\n")
    vals1 = re.findall(r"-?\d+\.\d+", out1)
    dg_free = float(vals1[1]) if len(vals1) >= 2 else 0.0

    # Constrained (all unpaired)
    constraint = "x" * len(region)
    out2 = _run_vienna_cmd("RNAfold -C --partfunc", f"{region}\n{constraint}\n@\n")
    vals2 = re.findall(r"-?\d+\.\d+", out2)
    dg_const = float(vals2[1]) if len(vals2) >= 2 else 0.0

    return dg_const - dg_free


def calc_sa(rna: str, rna_start: int, seed_type: str, region_end: int) -> float:
    """Structural accessibility via RNAplfold (log10 unpaired probability)."""
    seq = rna.upper().replace("T", "U")

    if seed_type in ("8mer", "7mer-m8", "offset-6mer"):
        seven_pos = rna_start + 1
    elif seed_type == "offset-7mer":
        seven_pos = rna_start + 2
    elif seed_type == "6mer-A1":
        seven_pos = rna_start - 1
    else:
        seven_pos = rna_start

    target_line = seven_pos + 6 + 2  # 0-indexed + header offset

    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
        f.write(f"{seq}\n@\n")
        fname = f.name

    try:
        subprocess.run(
            f"RNAplfold -L 40 -W 80 -u 14".split(),
            input=seq + "\n@\n", capture_output=True, text=True, timeout=60)

        # Read plfold_lunp output
        lunp_path = Path("plfold_lunp")
        if lunp_path.exists():
            lines = lunp_path.read_text().splitlines()
            if target_line < len(lines):
                vals = re.findall(r"[\d.]+e?[+-]?\d*", lines[target_line])
                if vals:
                    prob = float(vals[-1])
                    return math.log10(prob) if prob > 0 else -10.0
    except Exception:
        pass

    return 0.0


def calc_sps(rna: str, mirna: str, rna_start: int, seed_type: str, seed_length: int) -> float:
    """Seed pairing stability via RNAup."""
    seq = rna.upper().replace("T", "U")
    mir = mirna.upper().replace("T", "U")

    target_seed = seq[rna_start:rna_start + seed_length]
    if seed_type in ("8mer", "7mer-A1"):
        mir_seed = mir[:seed_length]
    else:
        mir_seed = mir[1:seed_length + 1]

    out = _run_vienna_cmd("RNAup -o", f"{target_seed}&{mir_seed}\n@\n")
    vals = re.findall(r"-?\d+\.\d+", out)
    return float(vals[0]) if vals else 0.0


def calc_mre(region_seq: str, mirna: str, max_energies: list[float],
             mismatch: int = 2) -> list[int]:
    """Count MRE sites at various energy thresholds via RNAduplex."""
    seq = region_seq.upper().replace("T", "U")
    mir = mirna.upper().replace("T", "U")

    out = _run_vienna_cmd("RNAduplex -e 100", f"{seq}\n{mir}\n@\n")
    if not out:
        return [0] * len(max_energies)

    energies = []
    for line in out.splitlines():
        if "&" not in line:
            continue
        energy_match = re.findall(r"\(\s*(-?\d+\.\d+)\)", line)
        if not energy_match:
            continue
        energy = float(energy_match[0])

        # Check seed pairing quality
        pos_matches = re.findall(r"(\d+),(\d+)", line)
        if len(pos_matches) >= 2:
            start_pos = int(pos_matches[1][0])
            if start_pos <= 8:
                pairing_str = line.split()[0] if line.split() else ""
                dots_in_seed = pairing_str[1:8].count(".") if len(pairing_str) > 7 else 99
                if dots_in_seed >= mismatch:
                    energy = 999  # Exclude
        energies.append(energy)

    diff = abs(max_energies[1] - max_energies[0]) if len(max_energies) > 1 else 5
    counts = []
    for m in max_energies:
        counts.append(sum(1 for e in energies if e <= m and e > m - diff))
    return counts


# ============================================================================
# Codon features (using standard human codon tables)
# ============================================================================

# Standard genetic code
_CODON_TABLE = {
    "TTT": "F", "TTC": "F", "TTA": "L", "TTG": "L", "CTT": "L", "CTC": "L",
    "CTA": "L", "CTG": "L", "ATT": "I", "ATC": "I", "ATA": "I", "ATG": "M",
    "GTT": "V", "GTC": "V", "GTA": "V", "GTG": "V", "TCT": "S", "TCC": "S",
    "TCA": "S", "TCG": "S", "CCT": "P", "CCC": "P", "CCA": "P", "CCG": "P",
    "ACT": "T", "ACC": "T", "ACA": "T", "ACG": "T", "GCT": "A", "GCC": "A",
    "GCA": "A", "GCG": "A", "TAT": "Y", "TAC": "Y", "TAA": "*", "TAG": "*",
    "CAT": "H", "CAC": "H", "CAA": "Q", "CAG": "Q", "AAT": "N", "AAC": "N",
    "AAA": "K", "AAG": "K", "GAT": "D", "GAC": "D", "GAA": "E", "GAG": "E",
    "TGT": "C", "TGC": "C", "TGA": "*", "TGG": "W", "CGT": "R", "CGC": "R",
    "CGA": "R", "CGG": "R", "AGT": "S", "AGC": "S", "AGA": "R", "AGG": "R",
    "GGT": "G", "GGC": "G", "GGA": "G", "GGG": "G",
}

# Human codon usage frequencies (per thousand, from Kazusa)
_HUMAN_CUF = {
    "TTT": 17.6, "TTC": 20.3, "TTA": 7.7, "TTG": 12.9, "CTT": 13.2, "CTC": 19.6,
    "CTA": 7.2, "CTG": 39.6, "ATT": 16.0, "ATC": 20.8, "ATA": 7.5, "ATG": 22.0,
    "GTT": 11.0, "GTC": 14.5, "GTA": 7.1, "GTG": 28.1, "TCT": 15.2, "TCC": 17.7,
    "TCA": 12.2, "TCG": 4.4, "CCT": 17.5, "CCC": 19.8, "CCA": 16.9, "CCG": 6.9,
    "ACT": 13.1, "ACC": 18.9, "ACA": 15.1, "ACG": 6.1, "GCT": 18.4, "GCC": 27.7,
    "GCA": 15.8, "GCG": 7.4, "TAT": 12.2, "TAC": 15.3, "CAT": 10.9, "CAC": 15.1,
    "CAA": 12.3, "CAG": 34.2, "AAT": 17.0, "AAC": 19.1, "AAA": 24.4, "AAG": 31.9,
    "GAT": 21.8, "GAC": 25.1, "GAA": 29.0, "GAG": 39.6, "TGT": 10.6, "TGC": 12.6,
    "TGG": 13.2, "CGT": 4.5, "CGC": 10.4, "CGA": 6.2, "CGG": 11.4, "AGT": 12.1,
    "AGC": 19.5, "AGA": 12.2, "AGG": 12.0, "GGT": 10.8, "GGC": 22.2, "GGA": 16.5,
    "GGG": 16.5, "TAA": 1.0, "TAG": 0.8, "TGA": 1.6,
}

# Human tAI weights (from dos Reis et al. 2004, optimized)
_HUMAN_TAI = {
    "TTT": 0.296, "TTC": 1.000, "TTA": 0.068, "TTG": 0.359, "CTT": 0.185,
    "CTC": 0.576, "CTA": 0.068, "CTG": 1.000, "ATT": 0.407, "ATC": 1.000,
    "ATA": 0.185, "ATG": 1.000, "GTT": 0.296, "GTC": 0.576, "GTA": 0.068,
    "GTG": 1.000, "TCT": 0.407, "TCC": 0.576, "TCA": 0.185, "TCG": 0.068,
    "CCT": 0.407, "CCC": 0.576, "CCA": 0.359, "CCG": 0.068, "ACT": 0.407,
    "ACC": 1.000, "ACA": 0.359, "ACG": 0.068, "GCT": 0.576, "GCC": 1.000,
    "GCA": 0.359, "GCG": 0.068, "TAT": 0.296, "TAC": 1.000, "CAT": 0.296,
    "CAC": 1.000, "CAA": 0.359, "CAG": 1.000, "AAT": 0.296, "AAC": 1.000,
    "AAA": 0.407, "AAG": 1.000, "GAT": 0.407, "GAC": 1.000, "GAA": 0.407,
    "GAG": 1.000, "TGT": 0.185, "TGC": 1.000, "TGG": 1.000, "CGT": 0.068,
    "CGC": 0.576, "CGA": 0.068, "CGG": 0.068, "AGT": 0.185, "AGC": 0.576,
    "AGA": 0.185, "AGG": 0.068, "GGT": 0.296, "GGC": 1.000, "GGA": 0.296,
    "GGG": 0.068,
}


def _translate(dna_seq: str) -> str:
    """Translate DNA to protein."""
    return "".join(_CODON_TABLE.get(dna_seq[i:i+3], "X")
                   for i in range(0, len(dna_seq) - 2, 3))


def _geomean(values: list[float]) -> float:
    """Geometric mean of positive values."""
    vals = [v for v in values if v > 0]
    if not vals:
        return 0.0
    return float(np.exp(np.mean(np.log(vals))))


def _relative_codon_freq(codon: str) -> float:
    """Relative codon frequency within its synonymous family."""
    aa = _CODON_TABLE.get(codon, "X")
    if aa in ("*", "X"):
        return 0.0
    synonyms = [c for c, a in _CODON_TABLE.items() if a == aa]
    freqs = [_HUMAN_CUF.get(c, 0.0) for c in synonyms]
    total = sum(freqs)
    return _HUMAN_CUF.get(codon, 0.0) / total if total > 0 else 0.0


def calc_codon_features(rna: str, rna_start: int, seed_length: int,
                        orf_start: int, orf_end: int, region: str
                        ) -> dict[str, float]:
    """Codon usage, tAI, charge, and slow AA scores for window and whole ORF."""
    seq = rna.upper().replace("U", "T")
    result = {}

    regions_to_calc = ["orf"]
    if region == "ORF":
        regions_to_calc.append("window")
    else:
        # Non-ORF sites get sentinel 999 for window metrics
        for prefix in ("cub_global", "cub_local", "cub_local_CAI", "tAI_score",
                       "TDR_score", "charge_score", "slow_score"):
            result[f"{prefix}_win"] = 999.0

    for calc_type in regions_to_calc:
        if calc_type == "orf":
            win_start = orf_start
            win_end = orf_end
        else:
            window = 25
            start_rf = (rna_start - orf_start) % 3
            rna_end = rna_start + seed_length - 1
            end_rf = (rna_end - orf_start) % 3
            win_start = max(orf_start, rna_start - 3 * window - start_rf)
            win_end = min(orf_end, rna_end + 2 - end_rf + 3 * window)

        codons_str = seq[win_start:win_end]
        # Trim to multiple of 3
        codons_str = codons_str[:len(codons_str) - len(codons_str) % 3]
        if not codons_str:
            suffix = "orf" if calc_type == "orf" else "win"
            for prefix in ("cub_global", "cub_local", "cub_local_CAI", "tAI_score",
                           "TDR_score", "charge_score", "slow_score"):
                result[f"{prefix}_{suffix}"] = 0.0
            continue

        codons = [codons_str[i:i+3] for i in range(0, len(codons_str), 3)]
        protein = _translate(codons_str)

        # Slow AAs (D, P)
        slow_count = sum(1 for aa in protein if aa in "DP")
        n_aa = len(protein)
        if slow_count > 0:
            slow_score = -1 * slow_count / n_aa + (n_aa - slow_count) / slow_count
        else:
            slow_score = 0.0

        # Charged AAs
        pos_count = sum(1 for aa in protein if aa in "RKH")
        neg_count = sum(1 for aa in protein if aa in "ED")
        charge_score = (pos_count - neg_count) / n_aa if n_aa > 0 else 0.0

        # Per-codon scores
        cub_global = []
        cub_local = []
        tai_scores = []

        for codon in codons:
            if codon in _HUMAN_CUF:
                cub_global.append(_HUMAN_CUF[codon])
            rel = _relative_codon_freq(codon)
            if rel > 0:
                cub_local.append(rel)
            if codon in _HUMAN_TAI:
                tai_scores.append(_HUMAN_TAI[codon])

        suffix = "orf" if calc_type == "orf" else "win"
        result[f"cub_global_{suffix}"] = _geomean(cub_global)
        result[f"cub_local_{suffix}"] = _geomean(cub_local)
        result[f"cub_local_CAI_{suffix}"] = _geomean(cub_local)  # CAI ≈ local CUB
        result[f"tAI_score_{suffix}"] = _geomean(tai_scores)
        result[f"TDR_score_{suffix}"] = _geomean(tai_scores)  # TDR ≈ tAI approximation
        result[f"charge_score_{suffix}"] = charge_score
        result[f"slow_score_{suffix}"] = slow_score

    return result


# ============================================================================
# Master feature computation
# ============================================================================

def compute_site_features(
    utr5: str,
    orf: str,
    utr3: str,
    mirna: str,
    rna_start: int,
    seed_length: int,
    seed_type: str,
    region: str,
    phastcons100: np.ndarray | None = None,
    phastcons20: np.ndarray | None = None,
    phylops100: np.ndarray | None = None,
    phylops20: np.ndarray | None = None,
    with_conservation: bool = True,
    with_thermo: bool = True,
) -> dict[str, float]:
    """Compute the full 152-feature vector for a single miRNA binding site.

    This is the Python equivalent of MATLAB's calc_gen_features.m.

    Args:
        utr5: 5'UTR sequence.
        orf: ORF/CDS sequence.
        utr3: 3'UTR sequence.
        mirna: Mature miRNA sequence.
        rna_start: 0-based position of the seed match start in the full mRNA.
        seed_length: Seed length (6, 7, or 8).
        seed_type: "6mer", "7mer-A1", "7mer-m8", or "8mer".
        region: "UTR5", "ORF", or "UTR3".
        phastcons100/20: PhastCons conservation vectors.
        phylops100/20: PhyloP conservation vectors.
        with_conservation: Include conservation + codon features.
        with_thermo: Include thermodynamic features (requires ViennaRNA).

    Returns:
        Dict mapping feature names to float values (152 features).
    """
    rna = (utr5 + orf + utr3).upper().replace("T", "U")
    mir = mirna.upper().replace("T", "U")
    n = len(rna)
    orf_start = len(utr5)
    orf_end = len(utr5) + len(orf)

    mir_start = 1 if seed_type in ("6mer", "6mer-A1", "7mer-A1", "7mer-m8", "8mer") else 2
    if seed_type in ("8mer", "7mer-A1"):
        rna_pos1 = rna_start + seed_length - 1
    else:
        rna_pos1 = rna_start + seed_length

    # Region boundaries
    if region == "UTR5":
        region_start, region_end = 0, len(utr5)
        curr_region = utr5.upper().replace("T", "U")
    elif region == "ORF":
        region_start, region_end = orf_start, orf_end
        curr_region = orf.upper().replace("T", "U")
    else:
        region_start, region_end = orf_end, n
        curr_region = utr3.upper().replace("T", "U")

    f: dict[str, float] = {}

    # ── Conservation ──
    f["prob_binom"] = calc_prob_binom(rna, rna_start, seed_length)
    f["prob_exact"] = f["prob_binom"]  # Approximation (exact requires spatt binary)

    if with_conservation and any(v is not None for v in (phastcons100, phastcons20, phylops100, phylops20)):
        cons = calc_conservation(phastcons100, phastcons20, phylops100, phylops20, rna_start, seed_length)
        f.update(cons)
    else:
        # Zero-fill conservation features with correct naming
        for metric in ("phastcons", "phylops"):
            for aln in ("100", "20"):
                f[f"{metric}_seed{aln}"] = 0.0
                f[f"{metric}_site{aln}"] = 0.0
                f[f"{metric}_flank{aln}"] = 0.0
        for aln in ("100", "20"):
            for p in range(1, 7):
                f[f"phylops{aln}_pos{p}"] = 0.0

    # ── Sequence features ──
    dist, dist_start, dist_end = calc_utr_pos(rna_start, region_start, region_end, len(mir), seed_length)
    f["distance_score"] = dist
    f["dist_start"] = dist_start
    f["dist_end"] = dist_end

    f["au_content"] = calc_au_content(rna, rna_start, seed_length, seed_type)

    # miRNA nucleotide identities at positions 1 and 8
    for base in "ACG":
        f[f"miR_pos1_{base}"] = 1.0 if len(mir) > 0 and mir[0] == base else 0.0
        f[f"miR_pos8_{base}"] = 1.0 if len(mir) > 7 and mir[7] == base else 0.0

    # RNA nucleotide identities at positions 1, 8, 9
    rna_p1 = rna[rna_pos1] if rna_pos1 < n else ""
    rna_p8 = rna[rna_pos1 - 7] if rna_pos1 - 7 >= 0 else ""
    rna_p9 = rna[rna_pos1 - 8] if rna_pos1 - 8 >= 0 else ""
    for base in "ACG":
        f[f"RNA_pos1_{base}"] = 1.0 if rna_p1 == base else 0.0
        f[f"RNA_pos8_{base}"] = 1.0 if rna_p8 == base else 0.0
        f[f"RNA_pos9_{base}"] = 1.0 if rna_p9 == base else 0.0

    f["dist_start_rel"] = round(10 ** dist_start) / max(1, len(curr_region))
    for base in "ACG":
        f[f"region_{base}"] = curr_region.count(base) / max(1, len(curr_region))

    seed_seq = rna[rna_start:rna_start + seed_length]
    for base in "ACG":
        f[f"site_{base}"] = seed_seq.count(base) / max(1, seed_length)

    # Region lengths and AU content
    for name, seq_r in [("UTR5", utr5), ("ORF", orf), ("UTR3", utr3)]:
        sr = seq_r.upper().replace("T", "U")
        if sr:
            f[f"{name}_len"] = math.log10(len(sr))
            f[f"{name}_AU"] = 1.0 - (sr.count("G") + sr.count("C")) / len(sr)
        else:
            f[f"{name}_len"] = 0.0
            f[f"{name}_AU"] = 0.0

    # TargetScan 7 seed scores (approximated as 0 when no TS7 data)
    f["TA"] = 0.0
    f["TA_hela"] = 0.0
    f["SPS_6mer"] = 0.0
    f["SPS_7mer_m8"] = 0.0
    f["mean_SPS"] = 0.0

    # Longest pairing run in full duplex
    mir_end = max(0, rna_pos1 - len(mir) + 1)
    duplex_rna = rna[mir_end:mir_end + len(mir)]
    if len(duplex_rna) == len(mir):
        wc = is_watson_crick(duplex_rna, mir)
        runs = re.findall(r"1+", "".join("1" if p else "0" for p in wc))
        f["longest_pairing"] = max(len(r) for r in runs) if runs else 0.0
    else:
        f["longest_pairing"] = 0.0

    # Paired miR end (last 7 nt)
    end_rna = rna[mir_end:mir_end + 7]
    end_mir = mir[-7:] if len(mir) >= 7 else mir
    f["paired_miR_end"] = float(_count_wc(end_rna, end_mir)) if len(end_rna) == len(end_mir) else 0.0

    # 3' pairing scores
    f["pairing3p_sw"] = calc_3p_pairing_sw(rna, rna_start, mir, seed_type)
    f["pairing3p"] = calc_3p_pairing(rna, rna_start, mir, seed_length, mir_start)

    # miR cognate region composition
    mir_cog = rna[mir_end:rna_pos1 + 1]
    cog_len = max(1, len(mir_cog))
    for base in "ACG":
        f[f"miR_cog_{base}"] = mir_cog.count(base) / cog_len

    # Upstream/downstream 50nt flanking composition
    up_start = max(0, mir_end - 50)
    cog_up = rna[up_start:max(0, mir_end)]
    up_len = max(1, len(cog_up))
    for base in "ACG":
        f[f"miR_cog_up_{base}"] = cog_up.count(base) / up_len

    down_end = min(n, rna_pos1 + 51)
    cog_down = rna[min(n, rna_pos1 + 1):down_end]
    dn_len = max(1, len(cog_down))
    for base in "ACG":
        f[f"miR_cog_down_{base}"] = cog_down.count(base) / dn_len

    # PUM sites
    orf_pum, utr3_pum, nei_pum = calc_pum(rna, rna_start, seed_length, orf_start, orf_end)
    f["ORF_pum1"], f["ORF_pum2"] = orf_pum
    f["UTR3_pum1"], f["UTR3_pum2"] = utr3_pum
    f["nei_pum1"], f["nei_pum2"] = nei_pum

    # RBP sites in 3'UTR
    utr3_u = utr3.upper().replace("T", "U")
    for j, motif in enumerate(RBP_SITES, 1):
        f[f"RBP{j}"] = float(utr3_u.count(motif))

    # Non-canonical seed site counts in 3'UTR
    for seed_name in ("6mer_A1", "offset_7mer", "offset_6mer", "centered",
                      "CDNST1", "CDNST2", "CDNST3", "CDNST4"):
        f[f"utr3_{seed_name}"] = 0.0  # Would need full non-canonical seed scanner

    # 3-mer seed complement counts
    seed_3mer = _rc(mir[1:4])
    for name, seq_r in [("utr5", utr5), ("orf", orf), ("utr3", utr3)]:
        sr = seq_r.upper().replace("U", "T")
        seed_3_dna = seed_3mer.replace("U", "T")
        f[f"{name}_3mer"] = float(sr.count(seed_3_dna)) if sr else 0.0

    # ── Thermodynamic features ──
    if with_thermo:
        try:
            dg_d, dg_ds, dg_b, dg_bs = calc_dg_duplex_binding(
                rna, mir, mir_start, rna_pos1, rna_start, seed_type, seed_length)
            f["dg_duplex"] = dg_d
            f["dg_duplex_seed"] = dg_ds
            f["dg_binding"] = dg_b
            f["dg_binding_seed"] = dg_bs
        except Exception:
            f["dg_duplex"] = f["dg_duplex_seed"] = f["dg_binding"] = f["dg_binding_seed"] = 0.0

        try:
            f["dg_open"] = calc_dg_open(rna, rna_pos1, len(mir))
        except Exception:
            f["dg_open"] = 0.0

        try:
            f["SPS"] = calc_sps(rna, mir, rna_start, seed_type, seed_length)
        except Exception:
            f["SPS"] = 0.0

        try:
            sa = calc_sa(rna, rna_start, seed_type, region_end)
            f["SA"] = sa if not math.isnan(sa) else 999.0
        except Exception:
            f["SA"] = 999.0

        max_energies = [0, -5, -10, -15, -20, -25, -30]
        orf_u = orf.upper().replace("T", "U")
        utr3_seq = utr3.upper().replace("T", "U")

        try:
            orf_mres = calc_mre(orf_u, mir, max_energies, mismatch=2)
        except Exception:
            orf_mres = [0] * len(max_energies)

        try:
            utr_mres = calc_mre(utr3_seq, mir, max_energies, mismatch=2)
        except Exception:
            utr_mres = [0] * len(max_energies)

        for i, e in enumerate(max_energies):
            f[f"MRE_therm_ORF_m2_{abs(e)}"] = float(orf_mres[i])
            f[f"MRE_therm_UTR_m2_{abs(e)}"] = float(utr_mres[i])
    else:
        # Zero-fill thermo features
        for key in ("dg_duplex", "dg_duplex_seed", "dg_binding", "dg_binding_seed",
                    "dg_open", "SPS", "SA"):
            f[key] = 0.0
        for prefix in ("MRE_therm_ORF_m2", "MRE_therm_UTR_m2"):
            for e in (0, 5, 10, 15, 20, 25, 30):
                f[f"{prefix}_{e}"] = 0.0

    # ── Codon features ──
    if with_conservation:
        orf_dna = orf.upper().replace("U", "T")
        rna_dna = rna.replace("U", "T")
        codon_feats = calc_codon_features(rna_dna, rna_start, seed_length, orf_start, orf_end, region)
        f.update(codon_feats)
    else:
        for prefix in ("cub_global", "cub_local", "cub_local_CAI", "tAI_score",
                       "TDR_score", "charge_score", "slow_score"):
            f[f"{prefix}_win"] = 0.0
            f[f"{prefix}_orf"] = 0.0

    # ── Biochemical features (placeholders — filled by mibsim.predict_occupancy) ──
    for key in ("kds_val", "logKd", "occ", "occ_init", "log_occ", "log_occ_init"):
        f.setdefault(key, 0.0)

    return f
