# geney/folding/precompute.py — Step 2 precomputation: delta-MFE table generation.
"""Build the full delta-MFE table for every position × alt in a transcript.

Folding windows are cut from the **mature mRNA** (5'UTR + CDS + 3'UTR) so
that windows near CDS boundaries extend into UTRs, matching the biology.
Codon classification (syn / non-syn) still uses the **CDS-only** sequence.

The table is used by :mod:`geney.folding.query` for fast, Vienna-free scoring.
"""

from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor
from typing import Callable

import pandas as pd

from .codons import classify_substitution
from .delta_mfe import _windows_covering_position
from .reference import find_cds_offset
from .vienna import fold

__all__ = [
    "precompute_ref_window_mfes",
    "precompute_transcript_deltas",
    "precompute_batch",
]

_BASES = ("A", "C", "G", "T")


# ---------------------------------------------------------------------------
# Reference-window MFE cache (one Vienna call per window, reused for all alts)
# ---------------------------------------------------------------------------

def precompute_ref_window_mfes(
    seq: str,
    window_size: int = 39,
    fold_fn: Callable[[str], float] | None = None,
) -> list[float]:
    """Compute MFE for every overlapping window in *seq*.

    ``result[i]`` = MFE of ``seq[i : i + window_size]``.

    Args:
        seq: DNA/RNA sequence.
        window_size: Window length (default 39).
        fold_fn: Custom fold callable; defaults to :func:`vienna.fold`.
    """
    if fold_fn is None:
        fold_fn = fold
    seq = seq.upper()
    n = len(seq)
    if window_size > n:
        return []
    return [fold_fn(seq[i : i + window_size]) for i in range(n - window_size + 1)]


# ---------------------------------------------------------------------------
# Per-transcript delta table
# ---------------------------------------------------------------------------

def precompute_transcript_deltas(
    mrna_seq: str,
    cds_seq: str,
    transcript_id: str = "",
    window_size: int = 39,
    fold_fn: Callable[[str], float] | None = None,
    ref_mfes: list[float] | None = None,
    cds_offset: int | None = None,
) -> pd.DataFrame:
    """Generate the delta-MFE table for one transcript.

    For every CDS position and each of 3 possible alternative bases:

    * **delta_mfe** — mean over overlapping windows of
      ``MFE(variant_window) − MFE(reference_window)``, where windows are
      cut from the mature mRNA (so CDS-boundary windows extend into UTRs).
    * **functional_class** — ``'synonymous'`` or ``'non_synonymous'``
      (determined from the CDS sequence).

    Args:
        mrna_seq: Full mature mRNA sequence (5'UTR + CDS + 3'UTR) used for
            folding windows.
        cds_seq: CDS-only DNA sequence (in-frame, no UTR) used for codon
            classification.
        transcript_id: Label stored in each output row.
        window_size: Folding window length (default 39).
        fold_fn: Custom fold callable; defaults to :func:`vienna.fold`.
        ref_mfes: Pre-computed reference window MFEs over the **mRNA**
            (from :func:`precompute_ref_window_mfes`).  Computed internally
            if *None*.
        cds_offset: 0-based start of CDS within *mrna_seq* (= 5'UTR
            length).  Auto-detected via :func:`find_cds_offset` if *None*.

    Returns:
        DataFrame with columns ``transcript_id``, ``cds_position_1based``,
        ``ref_base``, ``alt_base``, ``delta_mfe``, ``functional_class``.
        Three rows per CDS position.
    """
    if fold_fn is None:
        fold_fn = fold

    mrna = mrna_seq.upper()
    cds = cds_seq.upper().strip("N")  # Ensembl CDS may have leading/trailing N's
    n_mrna = len(mrna)
    n_cds = len(cds)

    if cds_offset is None:
        cds_offset = find_cds_offset(mrna, cds)

    if ref_mfes is None:
        ref_mfes = precompute_ref_window_mfes(mrna, window_size, fold_fn)

    rows: list[dict] = []
    for cds_pos in range(n_cds):
        mrna_pos = cds_offset + cds_pos
        ref_base = cds[cds_pos]

        for alt_base in _BASES:
            if alt_base == ref_base:
                continue

            # Synonymous / non-synonymous (from CDS sequence)
            try:
                func_class = classify_substitution(cds, cds_pos, alt_base)
            except ValueError:
                func_class = "non_synonymous"

            # Delta-MFE: average over all mRNA windows covering this position
            deltas: list[float] = []
            for start, end in _windows_covering_position(n_mrna, mrna_pos, window_size):
                ref_win_mfe = ref_mfes[start]
                var_win = mrna[start:mrna_pos] + alt_base + mrna[mrna_pos + 1 : end]
                var_win_mfe = fold_fn(var_win)
                deltas.append(var_win_mfe - ref_win_mfe)

            delta_mfe = sum(deltas) / len(deltas) if deltas else 0.0

            rows.append(
                {
                    "transcript_id": transcript_id,
                    "cds_position_1based": cds_pos + 1,
                    "ref_base": ref_base,
                    "alt_base": alt_base,
                    "delta_mfe": delta_mfe,
                    "functional_class": func_class,
                }
            )

    return pd.DataFrame(rows)


# ---------------------------------------------------------------------------
# Batch (multi-transcript) precomputation with optional parallelism
# ---------------------------------------------------------------------------

def _precompute_one(args: tuple[str, str, str, int]) -> pd.DataFrame:
    """Worker target for :func:`precompute_batch` — must be top-level for pickling."""
    transcript_id, mrna_seq, cds_seq, window_size = args
    return precompute_transcript_deltas(mrna_seq, cds_seq, transcript_id, window_size)


def precompute_batch(
    mrna_sequences: dict[str, str],
    cds_sequences: dict[str, str],
    transcript_ids: list[str] | None = None,
    window_size: int = 39,
    n_workers: int = 1,
    output_path: str | None = None,
) -> pd.DataFrame:
    """Precompute delta-MFE tables for multiple transcripts.

    Only transcripts present in **both** *mrna_sequences* and
    *cds_sequences* are processed.

    Args:
        mrna_sequences: ``{transcript_id: mrna_dna_sequence}`` — from
            the cDNA FASTA (``cdna.all.fa.gz``).
        cds_sequences: ``{transcript_id: cds_dna_sequence}`` — from
            the CDS FASTA (``cds.all.fa.gz``).
        transcript_ids: Subset to process (default: intersection of both
            dicts).
        window_size: Folding window length (default 39).
        n_workers: Parallel workers (default 1 = serial).
            Each worker runs in a separate process so ViennaRNA is safe
            across the GIL boundary.
        output_path: If provided, save the combined table here
            (``.parquet`` or CSV).

    Returns:
        Combined DataFrame with delta-MFE rows for all processed transcripts.
    """
    common_ids = set(mrna_sequences.keys()) & set(cds_sequences.keys())

    if transcript_ids is None:
        transcript_ids = sorted(common_ids)
    else:
        transcript_ids = [tid for tid in transcript_ids if tid in common_ids]

    tasks = [
        (tid, mrna_sequences[tid], cds_sequences[tid], window_size)
        for tid in transcript_ids
    ]

    if not tasks:
        return pd.DataFrame()

    if n_workers > 1 and len(tasks) > 1:
        with ProcessPoolExecutor(max_workers=n_workers) as pool:
            frames = list(pool.map(_precompute_one, tasks))
    else:
        frames = [_precompute_one(t) for t in tasks]

    result = pd.concat(frames, ignore_index=True)

    if output_path:
        if str(output_path).endswith(".parquet"):
            result.to_parquet(output_path, index=False)
        else:
            result.to_csv(output_path, index=False)

    return result
