# geney/enrichment.py — Lightweight enrichment utilities.
"""
Pure-function enrichments that add interpretive value to pipeline outputs
without requiring additional model inference. All functions are O(n) or
better and add negligible runtime.
"""
from __future__ import annotations

import re
from typing import Optional


# ============================================================================
# NMD prediction
# ============================================================================

def predict_nmd_eligibility(
    variant_protein: str,
    reference_protein: str,
    exon_boundaries: list[int] | None = None,
    ptc_threshold_nt: int = 55,
) -> dict:
    """Predict nonsense-mediated decay eligibility for a variant isoform.

    NMD is triggered when a premature termination codon (PTC) is located
    >50-55 nt upstream of the last exon-exon junction (EEJ). This is a
    simplified rule-of-thumb; actual NMD escape can depend on many factors.

    Note on units: This function operates on PROTEIN sequences (amino acids),
    but the NMD rule is defined in NUCLEOTIDES (>50-55 nt from last EEJ).
    The conversion uses integer division: ptc_threshold_aa = ptc_threshold_nt // 3.
    exon_boundaries should be provided in amino acid positions (not nucleotide).

    Args:
        variant_protein: Variant protein sequence (may contain '*' for stop).
        reference_protein: Reference protein sequence.
        exon_boundaries: List of amino acid positions at exon-exon junctions.
            These are in AA coordinates, not nucleotide. If None, uses a
            length-based heuristic (PTC near C-terminus → NMD escape).
        ptc_threshold_nt: Nucleotide distance from last EEJ for NMD (default 55).
            Converted to AA internally via integer division by 3.

    Returns:
        Dict with keys: has_ptc, ptc_position, truncation_fraction,
        nmd_eligible, nmd_escape, reason.
    """
    ref_len = len(reference_protein.rstrip("*"))
    var_len = len(variant_protein.rstrip("*"))

    # Check for premature termination
    # A stop in the variant that's earlier than the reference stop
    has_ptc = var_len < ref_len * 0.95  # >5% truncation
    ptc_position = var_len if has_ptc else None
    truncation_fraction = 1.0 - (var_len / ref_len) if ref_len > 0 else 0.0

    if not has_ptc:
        return {
            "has_ptc": False,
            "ptc_position": None,
            "truncation_fraction": round(truncation_fraction, 4),
            "nmd_eligible": False,
            "nmd_escape": False,
            "reason": "no_premature_termination",
        }

    # NMD eligibility: PTC must be >55nt (~18 AA) upstream of last EEJ
    ptc_threshold_aa = ptc_threshold_nt // 3

    if exon_boundaries and len(exon_boundaries) > 1:
        last_eej = max(exon_boundaries)
        distance_to_last_eej = last_eej - var_len
        nmd_eligible = distance_to_last_eej > ptc_threshold_aa
    else:
        # Heuristic: if truncation is in the last ~20 AA, likely NMD escape
        distance_from_end = ref_len - var_len
        nmd_eligible = distance_from_end > ptc_threshold_aa

    # Known NMD escape scenarios
    nmd_escape = False
    escape_reason = None

    # Last-exon PTCs escape NMD
    if exon_boundaries and var_len > max(exon_boundaries):
        nmd_escape = True
        escape_reason = "ptc_in_last_exon"

    # Very early PTCs (<200 nt from start) can trigger reinitiation
    if var_len < 67:  # ~200 nt / 3
        nmd_escape = True
        escape_reason = "ptc_near_start_possible_reinitiation"

    return {
        "has_ptc": True,
        "ptc_position": ptc_position,
        "truncation_fraction": round(truncation_fraction, 4),
        "nmd_eligible": nmd_eligible and not nmd_escape,
        "nmd_escape": nmd_escape,
        "reason": escape_reason or ("nmd_eligible" if nmd_eligible else "ptc_near_end_nmd_escape"),
    }


# ============================================================================
# Kozak context scoring
# ============================================================================

# Kozak consensus: (gcc)gccRccAUGG where R = A or G
# Position weights relative to AUG (positions -6 to +4):
#   Critical: -3 (must be A or G) and +4 (must be G)
#   Important: -6 to -1 context
KOZAK_WEIGHTS = {
    -6: {"G": 0.23, "C": 0.18, "A": 0.29, "T": 0.30},
    -5: {"G": 0.19, "C": 0.39, "A": 0.23, "T": 0.19},
    -4: {"G": 0.23, "C": 0.42, "A": 0.19, "T": 0.16},
    -3: {"G": 0.47, "C": 0.05, "A": 0.53, "T": 0.0},  # Critical: purine
    -2: {"G": 0.15, "C": 0.55, "A": 0.15, "T": 0.15},
    -1: {"G": 0.15, "C": 0.55, "A": 0.15, "T": 0.15},
    # +1,+2,+3 = ATG (fixed)
    +4: {"G": 0.55, "C": 0.15, "A": 0.15, "T": 0.15},  # Critical: G
}


def score_kozak_context(mrna: str, atg_position: int) -> dict:
    """Score the Kozak consensus context around an ATG codon.

    Args:
        mrna: mRNA sequence.
        atg_position: 0-based position of the A in ATG.

    Returns:
        Dict with kozak_score (0-1), kozak_strength ("strong"/"moderate"/"weak"),
        critical_minus3 (bool), critical_plus4 (bool), context_seq.
    """
    seq = mrna.upper().replace("U", "T")
    n = len(seq)

    # Extract context sequence
    ctx_start = max(0, atg_position - 6)
    ctx_end = min(n, atg_position + 7)  # AUG + 4
    context_seq = seq[ctx_start:ctx_end]

    score = 0.0
    max_score = 0.0

    for offset, weights in KOZAK_WEIGHTS.items():
        # Offsets relative to A in ATG (0-indexed):
        #   -6..-1 = upstream bases
        #   +4 = first base after G of ATG (0-indexed: atg_position + 3)
        pos = atg_position + offset if offset < 0 else atg_position + offset - 1
        if 0 <= pos < n:
            base = seq[pos]
            score += weights.get(base, 0.0)
            max_score += max(weights.values())
        else:
            max_score += max(weights.values())

    kozak_score = score / max_score if max_score > 0 else 0.0

    # Critical positions
    minus3_pos = atg_position - 3
    plus4_pos = atg_position + 3  # G after ATG
    critical_minus3 = (0 <= minus3_pos < n and seq[minus3_pos] in "AG")
    critical_plus4 = (0 <= plus4_pos < n and seq[plus4_pos] == "G")

    # Strength classification
    if critical_minus3 and critical_plus4:
        strength = "strong"
    elif critical_minus3 or critical_plus4:
        strength = "moderate"
    else:
        strength = "weak"

    return {
        "kozak_score": round(kozak_score, 4),
        "kozak_strength": strength,
        "critical_minus3": critical_minus3,
        "critical_plus4": critical_plus4,
        "context_seq": context_seq,
    }


# ============================================================================
# Base-pair distance (structural dissimilarity from dot-bracket)
# ============================================================================

def _parse_pairs(structure: str) -> set[tuple[int, int]]:
    """Extract base pairs from dot-bracket notation as a set of (i, j) tuples."""
    stack: list[int] = []
    pairs: set[tuple[int, int]] = set()
    for i, c in enumerate(structure):
        if c == "(":
            stack.append(i)
        elif c == ")":
            if stack:
                j = stack.pop()
                pairs.add((j, i))
    return pairs


def base_pair_distance(ref_structure: str, var_structure: str) -> dict:
    """Compute structural dissimilarity between two dot-bracket structures.

    Args:
        ref_structure: Reference dot-bracket notation.
        var_structure: Variant dot-bracket notation.

    Returns:
        Dict with bp_distance (int), ref_bp_count, var_bp_count,
        shared_bp_count, jaccard_similarity.
    """
    ref_pairs = _parse_pairs(ref_structure)
    var_pairs = _parse_pairs(var_structure)

    shared = ref_pairs & var_pairs
    union = ref_pairs | var_pairs
    distance = len(ref_pairs - var_pairs) + len(var_pairs - ref_pairs)
    jaccard = len(shared) / len(union) if union else 1.0

    return {
        "bp_distance": distance,
        "ref_bp_count": len(ref_pairs),
        "var_bp_count": len(var_pairs),
        "shared_bp_count": len(shared),
        "jaccard_similarity": round(jaccard, 4),
    }


# ============================================================================
# Structural motif detection
# ============================================================================

# G-quadruplex: 4 runs of 2+ Gs separated by 1-7 nt loops
_GQUAD_PATTERN = re.compile(
    r"G{2,}[ACGTU]{1,7}G{2,}[ACGTU]{1,7}G{2,}[ACGTU]{1,7}G{2,}",
    re.IGNORECASE,
)

# Simple hairpin: complementary regions flanking a small loop
# This is a lightweight heuristic, not a thermodynamic prediction
_HAIRPIN_MIN_STEM = 4
_HAIRPIN_MAX_LOOP = 8


def detect_structural_motifs(sequence: str) -> dict:
    """Detect known RNA structural motifs in a sequence.

    Args:
        sequence: RNA or DNA sequence.

    Returns:
        Dict with g_quadruplex_count, g_quadruplex_positions,
        has_g_quadruplex (bool).
    """
    seq = sequence.upper().replace("U", "T")
    gquad_matches = list(_GQUAD_PATTERN.finditer(seq))

    return {
        "g_quadruplex_count": len(gquad_matches),
        "g_quadruplex_positions": [(m.start(), m.end()) for m in gquad_matches],
        "has_g_quadruplex": len(gquad_matches) > 0,
    }


# ============================================================================
# Variant classification hint
# ============================================================================

def classify_variant(
    splicing_score: float,
    tis_score: float,
    folding_score: float,
    transcription_score: float,
    splicing_confidence: str = "high",
    has_ptc: bool = False,
    nmd_eligible: bool = False,
    truncation_fraction: float = 0.0,
    n_concordant: int = 0,
    predominant_direction: str = "unchanged",
) -> dict:
    """Produce a variant classification hint from multi-mechanism scores.

    This is NOT a clinical classifier — it's a computational suggestion
    to guide further investigation. Classifications:
    - "likely_lof": Strong evidence for loss of function
    - "possible_lof": Moderate evidence for loss of function
    - "possible_gof": Evidence suggesting gain of function
    - "regulatory_effect": Primarily transcriptional impact
    - "translational_effect": Primarily TIS/folding impact
    - "vus": Variant of uncertain significance

    Returns:
        Dict with classification, evidence_summary, contributing_mechanisms.
    """
    evidence = []
    contributors = []

    # Loss of function evidence
    lof_score = 0.0

    if has_ptc and nmd_eligible:
        lof_score += 3.0
        evidence.append("premature_stop_with_nmd")
        contributors.append("splicing")

    if splicing_score > 5.0 and splicing_confidence != "low":
        lof_score += 2.0
        evidence.append("high_splicing_disruption")
        contributors.append("splicing")

    if truncation_fraction > 0.3:
        lof_score += 1.5
        evidence.append(f"protein_truncation_{truncation_fraction:.0%}")
        contributors.append("splicing")

    if tis_score > 5.0:
        lof_score += 1.5
        evidence.append("translation_initiation_disrupted")
        contributors.append("tis")

    if transcription_score > 5.0 and predominant_direction == "downregulated":
        lof_score += 1.0
        evidence.append("predicted_downregulation")
        contributors.append("transcription")

    # Concordance bonus
    if n_concordant >= 3:
        lof_score += 1.0
        evidence.append("multi_mechanism_concordance")

    # Classification
    if lof_score >= 5.0:
        classification = "likely_lof"
    elif lof_score >= 3.0:
        classification = "possible_lof"
    elif transcription_score > 5.0 and predominant_direction == "upregulated":
        classification = "possible_gof"
        evidence.append("predicted_upregulation")
        contributors.append("transcription")
    elif transcription_score > 4.0 and splicing_score < 2.0:
        classification = "regulatory_effect"
        contributors.append("transcription")
    elif (tis_score > 3.0 or folding_score > 3.0) and splicing_score < 2.0:
        classification = "translational_effect"
        if tis_score > 3.0:
            contributors.append("tis")
        if folding_score > 3.0:
            contributors.append("folding")
    else:
        classification = "vus"

    return {
        "classification": classification,
        "evidence_summary": evidence,
        "contributing_mechanisms": list(set(contributors)),
        "lof_score": round(lof_score, 1),
    }


# ============================================================================
# Standalone missense scoring via Oncosplice
# ============================================================================

def score_missense_standalone(
    reference_protein: str,
    variant_protein: str,
    conservation_vector,
    window_length: int = 13,
) -> dict | None:
    """Score a missense mutation's conservation impact without requiring splicing.

    This runs the Oncosplice alignment + conservation scoring on the
    direct ref→var protein change (typically a single AA substitution).
    Gives a conservation-weighted functional impact even when there's
    no splice site disruption.

    Args:
        reference_protein: Reference protein sequence.
        variant_protein: Direct-translation variant protein sequence.
        conservation_vector: Per-residue conservation array.
        window_length: Smoothing window (default 13).

    Returns:
        Dict with missense_oncosplice_score, missense_percentile,
        n_mismatches, mismatch_positions, conservation_at_change.
        Returns None if proteins are identical.
    """
    import numpy as np

    ref = reference_protein.rstrip("*")
    var = variant_protein.rstrip("*")

    if ref == var:
        return None

    # Find mismatched positions (for same-length proteins, these are missense)
    min_len = min(len(ref), len(var))
    mismatches = [i for i in range(min_len) if ref[i] != var[i]]

    if not mismatches and len(ref) == len(var):
        return None

    # Run Oncosplice on the direct protein change
    from .oncosplice import Oncosplice

    onco = Oncosplice(ref, var, conservation_vector, window_length=window_length)

    # Extract conservation at mismatch positions
    cons = conservation_vector
    if hasattr(cons, '__len__'):
        cons_at_change = [
            float(cons[i]) if i < len(cons) else 0.0
            for i in mismatches
        ]
    else:
        cons_at_change = []

    return {
        "missense_oncosplice_score": onco.score,
        "missense_percentile": onco.percentile,
        "n_mismatches": len(mismatches),
        "mismatch_positions": mismatches,
        "conservation_at_change": cons_at_change,
        "mean_conservation_at_change": (
            sum(cons_at_change) / len(cons_at_change)
            if cons_at_change else 0.0
        ),
    }
