# geney/folding/delta_mfe.py — delta-MFE over overlapping windows (mutation-level score).
"""Compute mean delta-MFE over 39nt overlapping windows covering the mutation position.

Includes empirical null distribution from exhaustive local substitution landscape.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Optional

import numpy as np

from .vienna import fold, fold_full

__all__ = [
    "compute_delta_mfe",
    "compute_folding_profile",
    "compute_null_distribution",
    "WindowResult",
    "FoldingProfile",
    "NullDistribution",
]


def _windows_covering_position(seq_len: int, position_0based: int, window_size: int):
    """
    Yield (start, end) for each window of length window_size that contains position_0based.
    Windows are [start, start+window_size); position is contained iff start <= position_0based < start+window_size.
    So start in [position_0based - window_size + 1, position_0based], clamped to [0, seq_len - window_size].
    """
    if window_size <= 0 or position_0based < 0 or position_0based >= seq_len:
        return
    lo = max(0, position_0based - window_size + 1)
    hi = min(position_0based, seq_len - window_size)
    for s in range(lo, hi + 1):
        yield s, s + window_size


def compute_delta_mfe(
    ref_seq: str,
    var_seq: str,
    position_0based: int,
    window_size: int = 39,
    fold_fn=None,
) -> float:
    """
    Mean over overlapping windows of (MFE(var_window) - MFE(ref_window)).
    Only windows that contain the mutation position are used.

    Args:
        ref_seq: Reference (e.g. mRNA) sequence.
        var_seq: Variant sequence (same length; mutation applied at position_0based).
        position_0based: 0-based index of the mutated position in both sequences.
        window_size: Window length in nt (default 39).
        fold_fn: Optional callable(seq) -> MFE; default is Vienna fold().

    Returns:
        Mean delta-MFE (positive = variant less stable / higher MFE).
    """
    if fold_fn is None:
        fold_fn = fold
    n = len(ref_seq)
    if n != len(var_seq):
        raise ValueError("ref_seq and var_seq must have the same length")
    if position_0based < 0 or position_0based >= n:
        raise ValueError(
            f"position_0based must be in [0, {n}), got {position_0based}"
        )
    if window_size > n:
        raise ValueError(
            f"window_size ({window_size}) cannot exceed sequence length ({n})"
        )

    deltas = []
    for start, end in _windows_covering_position(n, position_0based, window_size):
        ref_win = ref_seq[start:end]
        var_win = var_seq[start:end]
        mfe_ref = fold_fn(ref_win)
        mfe_var = fold_fn(var_win)
        deltas.append(mfe_var - mfe_ref)

    if not deltas:
        return 0.0
    return sum(deltas) / len(deltas)


# ---------------------------------------------------------------------------
# Rich per-window output
# ---------------------------------------------------------------------------

@dataclass
class WindowResult:
    """Per-window folding detail for a single overlapping window."""
    start: int
    end: int
    ref_seq: str
    var_seq: str
    ref_structure: str
    var_structure: str
    ref_mfe: float
    var_mfe: float
    delta_mfe: float
    mutation_offset: int  # position of mutation within this window


@dataclass
class NullDistribution:
    """Empirical null distribution from exhaustive local substitution landscape.

    Generated by computing delta-MFE for every possible single-nucleotide
    substitution at every position in the folding window. ViennaRNA is fast
    enough (~0.1ms per 39nt fold) that this is practical: a 39nt window
    produces ~117 substitutions (39 positions × 3 alternatives each).

    Attributes:
        deltas: All null delta-MFE values (sorted).
        mean: Mean of null distribution.
        std: Standard deviation of null distribution.
        n_samples: Number of substitutions sampled.
        empirical_p: Fraction of null deltas with |delta| >= observed |delta|.
        z_score: (|observed| - mean(|null|)) / std(|null|).
        percentile: Percentile of observed |delta| in null distribution.
    """
    deltas: np.ndarray
    mean: float
    std: float
    n_samples: int
    empirical_p: float
    z_score: float
    percentile: float


@dataclass
class FoldingProfile:
    """Aggregated folding result for a single mutation position."""
    position: int          # 0-based mRNA position
    ref_base: str
    var_base: str
    delta_mfe: float       # mean over windows
    score: float           # min(10.0, |delta_mfe|)  [legacy; see enriched score below]
    windows: list[WindowResult] = field(default_factory=list)
    # Null distribution enrichment (populated when compute_with_null=True)
    null: Optional[NullDistribution] = None


def compute_null_distribution(
    ref_seq: str,
    position_0based: int,
    window_size: int = 39,
    observed_abs_delta: float = 0.0,
    radius: Optional[int] = None,
) -> NullDistribution:
    """Build empirical null by exhaustive single-nucleotide substitution landscape.

    For every position in the local window around the mutation, try all 3
    alternative bases and compute the mean delta-MFE across overlapping windows.
    This gives us the distribution of delta-MFE values expected by chance for
    *any* single-nucleotide change in this local sequence context.

    ViennaRNA folds a 39nt sequence in ~0.05ms, so exhaustive enumeration of
    ~117 substitutions (39 pos × 3 alts) takes ~25ms — negligible.

    Args:
        ref_seq: Reference mRNA sequence.
        position_0based: 0-based index of the actual mutation (center of window).
        window_size: Folding window size.
        observed_abs_delta: |delta_mfe| of the real mutation (for p-value).
        radius: How far from the mutation to sample. Default = window_size
            (covers the full folding context). Use smaller values for speed.

    Returns:
        NullDistribution with empirical p-value and z-score.
    """
    n = len(ref_seq)
    if radius is None:
        radius = window_size

    bases = "ACGT"
    ref_upper = ref_seq.upper()
    null_deltas: list[float] = []

    # Sample positions within radius of the mutation that are covered by
    # at least one folding window
    lo = max(0, position_0based - radius)
    hi = min(n, position_0based + radius + 1)

    for pos in range(lo, hi):
        ref_base = ref_upper[pos]
        for alt in bases:
            if alt == ref_base:
                continue
            # Build variant sequence with this substitution
            var_seq = ref_upper[:pos] + alt + ref_upper[pos + 1:]

            # Compute mean delta-MFE across overlapping windows at this position
            win_deltas = []
            for start, end in _windows_covering_position(n, pos, window_size):
                ref_win = ref_upper[start:end]
                var_win = var_seq[start:end]
                mfe_ref = fold(ref_win)
                mfe_var = fold(var_win)
                win_deltas.append(mfe_var - mfe_ref)

            if win_deltas:
                null_deltas.append(sum(win_deltas) / len(win_deltas))

    null_arr = np.array(null_deltas, dtype=np.float64)
    abs_null = np.abs(null_arr)

    n_samples = len(null_arr)
    null_mean = float(abs_null.mean()) if n_samples > 0 else 0.0
    null_std = float(abs_null.std()) if n_samples > 0 else 1.0

    # Empirical p-value: fraction of null with |delta| >= observed
    if n_samples > 0:
        empirical_p = float(np.sum(abs_null >= observed_abs_delta)) / n_samples
        # Continuity correction for small samples
        empirical_p = max(empirical_p, 1.0 / (n_samples + 1))
    else:
        empirical_p = 1.0

    # Z-score
    if null_std > 0:
        z_score = (observed_abs_delta - null_mean) / null_std
    else:
        z_score = 0.0

    # Percentile
    if n_samples > 0:
        percentile = float(np.sum(abs_null <= observed_abs_delta)) / n_samples * 100
    else:
        percentile = 50.0

    return NullDistribution(
        deltas=null_arr,
        mean=round(null_mean, 4),
        std=round(null_std, 4),
        n_samples=n_samples,
        empirical_p=round(empirical_p, 6),
        z_score=round(z_score, 4),
        percentile=round(percentile, 2),
    )


def compute_folding_profile(
    ref_seq: str,
    var_seq: str,
    position_0based: int,
    window_size: int = 39,
    compute_null: bool = True,
    null_radius: Optional[int] = None,
) -> FoldingProfile:
    """Compute rich folding profile with per-window structures and MFEs.

    Like compute_delta_mfe but returns full detail: dot-bracket structures,
    per-window sequences, individual MFEs, and a 0-10 normalized score.

    When ``compute_null=True`` (default), also generates an empirical null
    distribution from exhaustive local substitution landscape and provides
    an empirical p-value. This lets downstream scoring distinguish real
    structural effects from thermal noise.

    Args:
        ref_seq: Reference mRNA sequence.
        var_seq: Variant mRNA sequence (same length; mutation at position_0based).
        position_0based: 0-based index of the mutated position.
        window_size: Window length in nt (default 39).
        compute_null: If True (default), compute null distribution for p-value.
        null_radius: Radius around mutation for null sampling (default=window_size).

    Returns:
        FoldingProfile with per-window WindowResult entries, aggregated score,
        and optional NullDistribution.
    """
    n = len(ref_seq)
    if n != len(var_seq):
        raise ValueError("ref_seq and var_seq must have the same length")
    if position_0based < 0 or position_0based >= n:
        raise ValueError(
            f"position_0based must be in [0, {n}), got {position_0based}"
        )
    if window_size > n:
        raise ValueError(
            f"window_size ({window_size}) cannot exceed sequence length ({n})"
        )

    # Cache ref-window folds — only ref windows covering this position
    ref_cache: dict[int, tuple[str, float]] = {}
    for start, end in _windows_covering_position(n, position_0based, window_size):
        ref_win = ref_seq[start:end]
        ref_cache[start] = fold_full(ref_win)

    # Fold variant windows and build results
    window_results: list[WindowResult] = []
    deltas: list[float] = []
    for start, end in _windows_covering_position(n, position_0based, window_size):
        var_win = var_seq[start:end]
        var_structure, var_mfe = fold_full(var_win)
        ref_structure, ref_mfe = ref_cache[start]
        delta = var_mfe - ref_mfe
        deltas.append(delta)
        window_results.append(WindowResult(
            start=start,
            end=end,
            ref_seq=ref_seq[start:end],
            var_seq=var_win,
            ref_structure=ref_structure,
            var_structure=var_structure,
            ref_mfe=ref_mfe,
            var_mfe=var_mfe,
            delta_mfe=delta,
            mutation_offset=position_0based - start,
        ))

    mean_delta = sum(deltas) / len(deltas) if deltas else 0.0
    abs_delta = abs(mean_delta)

    # Compute null distribution if requested
    null_dist = None
    if compute_null and n >= window_size:
        null_dist = compute_null_distribution(
            ref_seq, position_0based, window_size,
            observed_abs_delta=abs_delta,
            radius=null_radius,
        )

    # Score: nonlinear transform of |ΔMFE|, compressing thermal noise (<2 kcal/mol).
    # k=0.5 tuned so |ΔMFE|=2→2.6, |ΔMFE|=5→7.1, |ΔMFE|=10→9.9.
    # The empirical p-value from the null distribution is reported as a separate
    # annotation (not blended into the score) so downstream consumers can apply
    # their own significance thresholds.
    k = 0.5
    enriched_score = 10.0 * (1.0 - math.exp(-k * abs_delta))

    return FoldingProfile(
        position=position_0based,
        ref_base=ref_seq[position_0based],
        var_base=var_seq[position_0based],
        delta_mfe=mean_delta,
        score=min(10.0, enriched_score),
        windows=window_results,
        null=null_dist,
    )
