# 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."""

from __future__ import annotations

from dataclasses import dataclass, field

from .vienna import fold, fold_full

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


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 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|)
    windows: list[WindowResult] = field(default_factory=list)


def compute_folding_profile(
    ref_seq: str,
    var_seq: str,
    position_0based: int,
    window_size: int = 39,
) -> 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.

    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).

    Returns:
        FoldingProfile with per-window WindowResult entries and aggregated score.
    """
    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

    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, abs(mean_delta)),
        windows=window_results,
    )
