# oncosplice/pipelines.py
from __future__ import annotations

from collections import namedtuple
from datetime import datetime
import threading
from typing import Optional
import numpy as np
import pandas as pd

from seqmat import Gene

from .splicing import SpliceSimulator, TranscriptLibrary
from .variants import MutationalEvent
from .oncosplice import Oncosplice
from .utils import select_transcript, _is_protein_coding, _position_in_transcript


MutationResult = namedtuple(
    "MutationResult", ["splicing", "tis", "folding", "transcription", "transcript", "tss_context"]
)


def _load_gene_with_timeout(gene_name: str, organism: str, timeout_seconds: Optional[float]) -> Optional[Gene]:
    """Load gene via Gene.from_file, optionally with a timeout to avoid hanging on slow I/O."""
    result: list = []
    exc: list = []

    def _run():
        try:
            result.append(Gene.from_file(gene_name, organism=organism))
        except Exception as e:
            exc.append(e)

    if timeout_seconds is not None and timeout_seconds > 0:
        thread = threading.Thread(target=_run, daemon=True)
        thread.start()
        thread.join(timeout=timeout_seconds)
        if thread.is_alive():
            raise TimeoutError(
                f"Loading gene '{gene_name}' from seqmat did not complete within {timeout_seconds}s. "
                "Your seqmat data path may be on a slow or network drive. Try increasing the timeout "
                "or moving data to a local disk."
            )
        if exc:
            raise exc[0]
        return result[0] if result else None
    return Gene.from_file(gene_name, organism=organism)


def _get_conservation_vector(transcript) -> np.ndarray:
    """
    Get conservation vector for a transcript, with fallback to ones.

    Args:
        transcript: Transcript object that may have cons_vector attribute

    Returns:
        numpy array of conservation scores (ones if not available)
    """
    cons_vector = getattr(transcript, 'cons_vector', None)
    if cons_vector is not None:
        return cons_vector

    # Fallback: create vector of ones matching protein length
    protein_length = len(getattr(transcript, 'protein', ''))
    if protein_length == 0:
        protein_length = 1  # Minimum length to avoid empty array
    return np.ones(protein_length)


def _extract_gained_alternatives(result) -> list[dict]:
    """Extract alternative TIS sites that gained significant probability.

    Filters the full analysis DataFrame for non-canonical sites whose variant
    probability exceeds the reference, sorted by probability change (largest
    gain first).
    """
    df = result.analysis_df
    if df is None or df.empty:
        return []
    alts = df[
        (df["location"] != "canonical")
        & (df["prob_change"] > 0)
        & (df["var_probability"] > 0)
    ].sort_values("prob_change", ascending=False)
    return alts.to_dict(orient="records")


def oncosplice_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    splicing_engine: str = "spliceai",
    organism: str = "hg38",
    gene_load_timeout: Optional[float] = None,
) -> pd.DataFrame:
    """
    Run the full oncosplice pipeline for a mutation.

    Returns DataFrame with all viable isoforms and their oncosplice scores.
    Returns empty DataFrame if no suitable transcript contains the mutation.

    Args:
        mut_id: Mutation ID (GENE:CHR:POS:REF:ALT, use | for compound events).
        transcript_id: Optional transcript ID.
        splicing_engine: Splicing engine (default "spliceai").
        organism: Genome build (default "hg38").
        gene_load_timeout: If set, raise TimeoutError when loading gene data
            from seqmat takes longer than this many seconds. Use when seqmat
            data is on a slow/network drive to fail fast instead of hanging.
    """
    m = MutationalEvent(mut_id)
    assert m.compatible(), "Mutations in event are incompatible"

    gene = _load_gene_with_timeout(m.gene, organism, gene_load_timeout)
    if gene is None:
        return pd.DataFrame()
    central_pos = m.central_position

    # Smart transcript selection
    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        # No transcript contains this mutation - return empty DataFrame
        return pd.DataFrame()

    reference_transcript = (
        selected
        .generate_pre_mrna()
        .generate_mature_mrna()
        .generate_protein()
    )

    tl = TranscriptLibrary(reference_transcript, m)

    tl.predict_splicing(central_pos, engine=splicing_engine, inplace=True)
    splicing_results = tl.get_event_columns("event")

    ss = SpliceSimulator(
        splicing_results, tl.event, feature="event", max_distance=100_000_000
    )

    base_report = pd.Series({
        "mut_id": mut_id,
        "gene": m.gene,
        "transcript_id": reference_transcript.transcript_id,
        "primary_transcript": getattr(reference_transcript, 'primary_transcript', False),
        "splicing_engine": splicing_engine,
        "central_position": central_pos,
        "mutation_count": len(m.positions),
        "time_of_execution": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    })

    # Get conservation vector with fallback to ones
    cons_vector = _get_conservation_vector(reference_transcript)

    ss_metadata = ss.report(central_pos)
    rows = []
    for variant_transcript, isoform_metadata in ss.get_viable_transcripts(metadata=True):
        onco = Oncosplice(
            reference_transcript.protein,
            variant_transcript.protein,
            cons_vector,
        )
        rows.append(
            pd.concat([
                base_report,
                ss_metadata,
                isoform_metadata,
                pd.Series({
                    "reference_mrna": reference_transcript.mature_mrna.seq,
                    "variant_mrna": variant_transcript.mature_mrna.seq,
                    "conservation_vector": cons_vector,
                }),
                onco.get_analysis_series(),
            ])
        )

    return pd.DataFrame(rows)


def oncosplice_top_isoform(
    mut_id: str,
    transcript_id: str | None = None,
    splicing_engine: str = "spliceai",
    organism: str = "hg38",
) -> pd.Series | None:
    """
    Get the most likely non-reference isoform for a mutation.

    Returns Series with full oncosplice analysis, or None if no missplicing detected.
    """
    df = oncosplice_pipeline(mut_id, transcript_id, splicing_engine, organism)

    if df.empty:
        return None

    variants = df[df["summary"] != "-"]

    if variants.empty:
        return None

    return variants.iloc[0]


def max_splicing_delta(
    mut_id: str,
    transcript_id: str | None = None,
    splicing_engine: str = "spliceai",
    organism: str = "hg38",
) -> float | None:
    """
    Get the maximum splice site probability change for a mutation.

    Returns None if no suitable transcript contains the mutation.
    """
    m = MutationalEvent(mut_id)
    assert m.compatible(), "Mutations in event are incompatible"

    gene = Gene.from_file(m.gene, organism=organism)
    central_pos = m.central_position

    # Smart transcript selection
    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return None

    reference_transcript = (
        selected
        .generate_pre_mrna()
        .generate_mature_mrna()
        .generate_protein()
    )

    tl = TranscriptLibrary(reference_transcript, m)
    splicing_results = tl.predict_splicing(
        central_pos, engine=splicing_engine, inplace=True
    ).get_event_columns("event")

    ss = SpliceSimulator(
        splicing_results, tl.event, feature="event", max_distance=100_000_000
    )

    return ss.max_splicing_delta("event_prob")


# Keep old name for backwards compatibility
oncosplice_pipeline_single_transcript = oncosplice_pipeline


# ============================================================================
# TIS Analysis Pipelines
# ============================================================================

def tis_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
    proximity_window: int = 500,
    probability_threshold: float = 0.1,
    use_prior: bool = False,
    skip_if_not_near_start: bool = True,
) -> dict | None:
    """
    Run TIS (Translation Initiation Site) analysis for a mutation.

    This analyzes whether a mutation affects translation initiation by
    applying the mutation directly to the mature mRNA (WITHOUT splicing changes)
    and comparing TIS probabilities between reference and variant.

    Efficiency optimizations:
    - Skips analysis if mutation is not in mature mRNA
    - Skips analysis if mutation is far from start codon (configurable)

    Args:
        mut_id: Mutation ID (e.g., "GENE:CHR:POS:REF:ALT")
        transcript_id: Optional specific transcript ID
        organism: Genome build (default "hg38")
        proximity_window: Max distance from 5' end to run analysis (default 500nt)
        probability_threshold: Min probability to consider a TIS (default 0.1)
        use_prior: Whether to use codon prior in scoring (default False)
        skip_if_not_near_start: If True, skip analysis for mutations far from start

    Returns:
        Dictionary with TIS analysis results, or None if:
        - No suitable transcript found
        - Mutation not in mature mRNA
        - Mutation too far from start codon (if skip_if_not_near_start=True)
    """
    from .tis import compute_tis_shift_metrics, get_tis_summary, TISShiftResult
    from .tis.helpers import (
        _mutation_in_mature_mrna,
        _mutation_near_start_codon,
        _apply_mutation_to_mrna,
        _get_mature_mrna_seq,
    )

    m = MutationalEvent(mut_id)
    assert m.compatible(), "Mutations in event are incompatible"

    gene = Gene.from_file(m.gene, organism=organism)
    central_pos = m.central_position

    # Select transcript
    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return None

    # Generate reference transcript with mature mRNA
    reference_transcript = (
        selected
        .generate_pre_mrna()
        .generate_mature_mrna()
        .generate_protein()
    )

    # Check if mutation is in mature mRNA
    if not _mutation_in_mature_mrna(reference_transcript, m.positions):
        return {
            'status': 'skipped',
            'reason': 'mutation_not_in_mature_mrna',
            'mut_id': mut_id,
            'gene': m.gene,
            'transcript_id': reference_transcript.transcript_id,
        }

    # Check if mutation is near start codon
    if skip_if_not_near_start:
        if not _mutation_near_start_codon(reference_transcript, m.positions, proximity_window):
            return {
                'status': 'skipped',
                'reason': 'mutation_far_from_start_codon',
                'mut_id': mut_id,
                'gene': m.gene,
                'transcript_id': reference_transcript.transcript_id,
                'proximity_window': proximity_window,
            }

    # Get reference mature mRNA
    ref_mrna_seq = _get_mature_mrna_seq(reference_transcript)
    if not ref_mrna_seq:
        return {
            'status': 'error',
            'reason': 'no_mature_mrna',
            'mut_id': mut_id,
        }

    # Apply mutations directly to mature mRNA (NO splicing changes)
    mrna = reference_transcript.mature_mrna
    if hasattr(mrna, 'index'):
        var_mrna_seq = _apply_mutation_to_mrna(
            ref_mrna_seq,
            mrna.index,
            m.mutation_args()
        )
    else:
        # Fallback: can't apply mutations without index mapping
        var_mrna_seq = ref_mrna_seq  # Will show no change

    # Find mutation position(s) in mRNA coordinates for visualization
    mutation_mrna_positions = []
    if hasattr(mrna, 'index'):
        idx_list = mrna.index.tolist() if hasattr(mrna.index, 'tolist') else list(mrna.index)
        pos_to_mrna = {pos: i for i, pos in enumerate(idx_list)}
        for pos, ref, alt in m.mutation_args():
            if pos in pos_to_mrna:
                mutation_mrna_positions.append(pos_to_mrna[pos])

    # Run TIS analysis
    try:
        result = compute_tis_shift_metrics(
            ref_mrna=ref_mrna_seq,
            var_mrna=var_mrna_seq,
            canonical_start=None,  # Auto-detect
            probability_threshold=probability_threshold,
            use_prior=use_prior,
        )

        summary = get_tis_summary(result)

        return {
            'status': 'completed',
            'mut_id': mut_id,
            'gene': m.gene,
            'transcript_id': reference_transcript.transcript_id,
            'ref_mrna_length': len(ref_mrna_seq),
            'var_mrna_length': len(var_mrna_seq),
            # Key metrics
            'canonical_prob_ref': result.canonical_ref_prob,
            'canonical_prob_var': result.canonical_var_prob,
            'canonical_prob_change': result.canonical_prob_change,
            'canonical_usage_ref': result.canonical_ref_usage,
            'canonical_usage_var': result.canonical_var_usage,
            'canonical_usage_change': result.canonical_usage_change,
            'tis_shift_score': result.tis_shift_score,
            'initiation_intact': result.initiation_intact,
            'num_alternative_tis': result.num_alternative_tis,
            'top_alternative_position': result.top_alternative_position,
            'top_alternative_prob': result.top_alternative_prob,
            'top_alternative_location': result.top_alternative_location,
            'top_alternative_in_frame': result.top_alternative_in_frame,
            'gained_alternative_tis': _extract_gained_alternatives(result),
            'atg_sites': result.analysis_df.to_dict('records') if not result.analysis_df.empty else [],
            # Full result object for detailed analysis
            'tis_result': result,
            'summary': summary,
            # Sequence context for frontend visualization
            'ref_mrna_window': ref_mrna_seq[:600],
            'var_mrna_window': var_mrna_seq[:600],
            'mutation_mrna_positions': mutation_mrna_positions,
        }
    except Exception as e:
        return {
            'status': 'error',
            'reason': str(e),
            'mut_id': mut_id,
            'gene': m.gene,
        }


def combined_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    splicing_engine: str = "spliceai",
    organism: str = "hg38",
    run_tis_analysis: bool = True,
    tis_proximity_window: int = 500,
    tis_probability_threshold: float = 0.1,
    tis_use_prior: bool = False,
) -> pd.DataFrame:
    """Run combined splicing + TIS analysis pipeline.

    .. deprecated::
        Use :func:`mutation_pipeline` instead, which loads the gene once and
        also supports mRNA-folding analysis.

    Returns a DataFrame with splicing isoform rows and TIS columns merged in,
    preserving the original return format for backwards compatibility.
    """
    import warnings
    warnings.warn(
        "combined_pipeline is deprecated — use mutation_pipeline() instead.",
        DeprecationWarning,
        stacklevel=2,
    )

    r = mutation_pipeline(
        mut_id=mut_id,
        transcript_id=transcript_id,
        organism=organism,
        splicing_engine=splicing_engine,
        run_splicing=True,
        run_tis=run_tis_analysis,
        run_folding=False,
        run_transcription=False,
        tis_proximity_window=tis_proximity_window,
        tis_probability_threshold=tis_probability_threshold,
    )

    df = r.splicing
    if df.empty:
        return df

    tis_result = r.tis
    if run_tis_analysis:
        if tis_result and tis_result.get("status") == "completed":
            df["tis_analyzed"] = True
            df["tis_canonical_prob_change"] = tis_result["canonical_prob_change"]
            df["tis_canonical_usage_change"] = tis_result["canonical_usage_change"]
            df["tis_shift_score"] = tis_result["tis_shift_score"]
            df["tis_initiation_intact"] = tis_result["initiation_intact"]
            df["tis_num_alternatives"] = tis_result["num_alternative_tis"]
            df["tis_top_alt_position"] = tis_result["top_alternative_position"]
            df["tis_top_alt_prob"] = tis_result["top_alternative_prob"]
            df["tis_top_alt_location"] = tis_result["top_alternative_location"]
        elif tis_result and tis_result.get("status") == "skipped":
            df["tis_analyzed"] = False
            df["tis_skip_reason"] = tis_result.get("reason", "unknown")
        else:
            df["tis_analyzed"] = False
            df["tis_skip_reason"] = "error" if tis_result else "no_result"
    else:
        df["tis_analyzed"] = False
        df["tis_skip_reason"] = "disabled"

    return df


def tis_quick_check(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
) -> dict:
    """
    Quick check if a mutation is likely to affect TIS.

    This is a fast pre-screening that doesn't run the full TITER model.
    Use this to filter mutations before running full TIS analysis.

    Args:
        mut_id: Mutation ID
        transcript_id: Optional transcript ID
        organism: Genome build

    Returns:
        Dictionary with:
        - 'should_analyze': bool, whether full TIS analysis is warranted
        - 'mutation_in_mrna': bool
        - 'near_start_codon': bool
        - 'affects_start_codon': bool, if mutation directly hits ATG
    """
    from .tis.helpers import _mutation_in_mature_mrna, _mutation_near_start_codon

    m = MutationalEvent(mut_id)
    gene = Gene.from_file(m.gene, organism=organism)
    central_pos = m.central_position

    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return {
            'should_analyze': False,
            'reason': 'no_transcript',
            'mutation_in_mrna': False,
            'near_start_codon': False,
            'affects_start_codon': False,
        }

    reference_transcript = (
        selected
        .generate_pre_mrna()
        .generate_mature_mrna()
    )

    mutation_in_mrna = _mutation_in_mature_mrna(reference_transcript, m.positions)
    near_start = _mutation_near_start_codon(reference_transcript, m.positions, window=500)

    # Check if mutation directly affects start codon
    affects_start = False
    try:
        mrna = reference_transcript.mature_mrna
        if hasattr(mrna, 'seq') and hasattr(mrna, 'index'):
            # Find first ATG
            seq = mrna.seq.upper()
            atg_pos = seq.find('ATG')
            if atg_pos != -1 and atg_pos < len(mrna.index):
                # Get genomic positions of ATG
                atg_genomic = set(mrna.index[atg_pos:atg_pos+3])
                affects_start = any(pos in atg_genomic for pos in m.positions)
    except Exception:
        pass

    should_analyze = mutation_in_mrna and (near_start or affects_start)

    return {
        'should_analyze': should_analyze,
        'mutation_in_mrna': mutation_in_mrna,
        'near_start_codon': near_start,
        'affects_start_codon': affects_start,
        'gene': m.gene,
        'transcript_id': getattr(reference_transcript, 'transcript_id', None),
    }


# ============================================================================
# Unified Mutation Pipeline
# ============================================================================

def mutation_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
    splicing_engine: str = "spliceai",
    tss_context_bp: int = 1000,
    run_splicing: bool = True,
    run_tis: bool = True,
    run_folding: bool = True,
    run_transcription: bool = True,
    tis_proximity_window: int = 500,
    tis_probability_threshold: float = 0.1,
    folding_window_size: int = 39,
    gene_load_timeout: float | None = None,
) -> MutationResult:
    """Run splicing, TIS, mRNA-folding, and transcription analyses for a single mutation.

    Loads the gene and selects a transcript **once**, then dispatches to each
    analysis system.  Transcription (Borzoi) runs independently using raw
    genomic sequence via ``SeqMat.from_fasta``.  Returns a
    :class:`MutationResult` namedtuple carrying the results together with the
    mutated transcript and a SeqMat clone of the region around the
    transcription start site.

    Args:
        mut_id: Mutation ID (``GENE:CHR:POS:REF:ALT``; pipe-separated for compound).
        transcript_id: Optional specific transcript ID.
        organism: Genome build (default ``"hg38"``).
        splicing_engine: Splicing engine (default ``"spliceai"``).
        tss_context_bp: Base pairs on each side of TSS to extract (default 1000).
        run_splicing: Whether to run splicing/oncosplice analysis.
        run_tis: Whether to run TIS analysis.
        run_folding: Whether to run mRNA folding analysis.
        run_transcription: Whether to run Borzoi transcription-effect analysis.
        tis_proximity_window: Max distance from 5' end for TIS analysis (nt).
        tis_probability_threshold: Min probability to consider a TIS.
        folding_window_size: Local folding window for delta-MFE (nt).
        gene_load_timeout: Timeout for gene loading (seconds); *None* = no limit.

    Returns:
        :class:`MutationResult` with fields *splicing* (DataFrame), *tis* (dict
        or None), *folding* (dict or None), *transcription* (dict or None),
        *transcript* (mutated Transcript), and *tss_context* (SeqMat clone
        around the TSS).
    """
    from .utils import _apply_mutation_safe

    # -- parse mutation -------------------------------------------------------
    m = MutationalEvent(mut_id)
    assert m.compatible(), "Mutations in event are incompatible"
    central_pos = m.central_position

    # -- load gene (single load) ----------------------------------------------
    gene = _load_gene_with_timeout(m.gene, organism, gene_load_timeout)
    if gene is None:
        return MutationResult(pd.DataFrame(), None, None, None, None, None)

    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return MutationResult(pd.DataFrame(), None, None, None, None, None)

    # -- build reference transcript (unmutated) for oncosplice ----------------
    reference_transcript = (
        selected
        .generate_pre_mrna()
        .generate_mature_mrna()
        .generate_protein()
    )

    # -- build mutated transcript + genomic context around mutation -----------
    mutated_transcript = selected.clone().generate_pre_mrna()
    for pos, ref, alt in m.mutation_args():
        try:
            _apply_mutation_safe(mutated_transcript.pre_mrna, pos, ref, alt)
        except Exception:
            pass  # mutation may be outside this transcript's range

    tss_context = mutated_transcript.pre_mrna.clone(
        central_pos - tss_context_bp, central_pos + tss_context_bp
    )
    mutated_transcript.generate_mature_mrna().generate_protein()

    # -- 1. Splicing / Oncosplice ---------------------------------------------
    splicing_result = pd.DataFrame()
    if run_splicing:
        try:
            tl = TranscriptLibrary(reference_transcript, m)
            tl.predict_splicing(central_pos, engine=splicing_engine, inplace=True)
            splicing_results = tl.get_event_columns("event")

            ss = SpliceSimulator(
                splicing_results, tl.event, feature="event", max_distance=100_000_000
            )

            base_report = pd.Series({
                "mut_id": mut_id,
                "gene": m.gene,
                "transcript_id": reference_transcript.transcript_id,
                "primary_transcript": getattr(reference_transcript, "primary_transcript", False),
                "splicing_engine": splicing_engine,
                "central_position": central_pos,
                "mutation_count": len(m.positions),
                "time_of_execution": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            })

            cons_vector = _get_conservation_vector(reference_transcript)
            ss_metadata = ss.report(central_pos)

            rows = []
            for variant_transcript, isoform_metadata in ss.get_viable_transcripts(metadata=True):
                onco = Oncosplice(
                    reference_transcript.protein,
                    variant_transcript.protein,
                    cons_vector,
                )
                rows.append(
                    pd.concat([
                        base_report,
                        ss_metadata,
                        isoform_metadata,
                        pd.Series({
                            "reference_mrna": reference_transcript.mature_mrna.seq,
                            "variant_mrna": variant_transcript.mature_mrna.seq,
                            "conservation_vector": cons_vector,
                        }),
                        onco.get_analysis_series(),
                    ])
                )

            splicing_result = pd.DataFrame(rows)
        except Exception:
            splicing_result = pd.DataFrame()

    # -- 2. TIS ---------------------------------------------------------------
    tis_result = None
    if run_tis:
        try:
            from .tis import compute_tis_shift_metrics, get_tis_summary
            from .tis.helpers import (
                _mutation_in_mature_mrna,
                _mutation_near_start_codon,
                _apply_mutation_to_mrna,
                _get_mature_mrna_seq,
            )

            if not _mutation_in_mature_mrna(reference_transcript, m.positions):
                tis_result = {
                    "status": "skipped",
                    "reason": "mutation_not_in_mature_mrna",
                    "mut_id": mut_id,
                    "gene": m.gene,
                    "transcript_id": reference_transcript.transcript_id,
                }
            elif not _mutation_near_start_codon(reference_transcript, m.positions, tis_proximity_window):
                tis_result = {
                    "status": "skipped",
                    "reason": "mutation_far_from_start_codon",
                    "mut_id": mut_id,
                    "gene": m.gene,
                    "transcript_id": reference_transcript.transcript_id,
                    "proximity_window": tis_proximity_window,
                }
            else:
                ref_mrna_seq = _get_mature_mrna_seq(reference_transcript)
                if not ref_mrna_seq:
                    tis_result = {"status": "error", "reason": "no_mature_mrna", "mut_id": mut_id}
                else:
                    mrna = reference_transcript.mature_mrna
                    if hasattr(mrna, "index"):
                        var_mrna_seq = _apply_mutation_to_mrna(
                            ref_mrna_seq, mrna.index, m.mutation_args()
                        )
                    else:
                        var_mrna_seq = ref_mrna_seq

                    mutation_mrna_positions = []
                    if hasattr(mrna, "index"):
                        idx_list = mrna.index.tolist() if hasattr(mrna.index, "tolist") else list(mrna.index)
                        pos_to_mrna = {pos: i for i, pos in enumerate(idx_list)}
                        for pos, ref, alt in m.mutation_args():
                            if pos in pos_to_mrna:
                                mutation_mrna_positions.append(pos_to_mrna[pos])

                    result = compute_tis_shift_metrics(
                        ref_mrna=ref_mrna_seq,
                        var_mrna=var_mrna_seq,
                        canonical_start=None,
                        probability_threshold=tis_probability_threshold,
                    )
                    summary = get_tis_summary(result)

                    tis_result = {
                        "status": "completed",
                        "mut_id": mut_id,
                        "gene": m.gene,
                        "transcript_id": reference_transcript.transcript_id,
                        "ref_mrna_length": len(ref_mrna_seq),
                        "var_mrna_length": len(var_mrna_seq),
                        "canonical_prob_ref": result.canonical_ref_prob,
                        "canonical_prob_var": result.canonical_var_prob,
                        "canonical_prob_change": result.canonical_prob_change,
                        "canonical_usage_ref": result.canonical_ref_usage,
                        "canonical_usage_var": result.canonical_var_usage,
                        "canonical_usage_change": result.canonical_usage_change,
                        "tis_shift_score": result.tis_shift_score,
                        "initiation_intact": result.initiation_intact,
                        "num_alternative_tis": result.num_alternative_tis,
                        "top_alternative_position": result.top_alternative_position,
                        "top_alternative_prob": result.top_alternative_prob,
                        "top_alternative_location": result.top_alternative_location,
                        "top_alternative_in_frame": result.top_alternative_in_frame,
                        "gained_alternative_tis": _extract_gained_alternatives(result),
                        "atg_sites": result.analysis_df.to_dict("records") if not result.analysis_df.empty else [],
                        "tis_result": result,
                        "summary": summary,
                        "ref_mrna_window": ref_mrna_seq[:600],
                        "var_mrna_window": var_mrna_seq[:600],
                        "mutation_mrna_positions": mutation_mrna_positions,
                    }
        except Exception as e:
            tis_result = {
                "status": "error",
                "reason": str(e),
                "mut_id": mut_id,
                "gene": m.gene,
            }

    # -- 3. mRNA Folding ------------------------------------------------------
    folding_result = None
    if run_folding:
        try:
            from .folding.delta_mfe import compute_folding_profile as _compute_folding_profile
            from .folding.pipeline import _central_window_dict
            from .folding.vienna import fold_available
            from .folding.reference import complement

            if not fold_available():
                folding_result = {
                    "status": "error",
                    "reason": "ViennaRNA not installed (pip install ViennaRNA)",
                    "mut_id": mut_id,
                }
            else:
                mrna = reference_transcript.mature_mrna
                if mrna is None or not hasattr(mrna, "seq") or not hasattr(mrna, "index"):
                    folding_result = {"status": "error", "reason": "no_mature_mrna", "mut_id": mut_id}
                else:
                    ref_seq = mrna.seq.upper()
                    idx_list = mrna.index.tolist() if hasattr(mrna.index, "tolist") else list(mrna.index)
                    pos_to_mrna = {gpos: i for i, gpos in enumerate(idx_list)}

                    per_mut: list[dict] = []
                    for pos, ref, alt in m.mutation_args():
                        mrna_pos = pos_to_mrna.get(pos)
                        if mrna_pos is None:
                            per_mut.append({
                                "genomic_position": pos, "ref": ref, "alt": alt,
                                "status": "not_in_mrna", "delta_mfe": None,
                            })
                            continue

                        actual_ref = ref_seq[mrna_pos]
                        coding_ref = ref.upper()
                        coding_alt = alt.upper()

                        if actual_ref != coding_ref:
                            coding_ref = complement(coding_ref)
                            coding_alt = complement(coding_alt)
                            if actual_ref != coding_ref:
                                per_mut.append({
                                    "genomic_position": pos, "ref": ref, "alt": alt,
                                    "status": "ref_mismatch", "delta_mfe": None,
                                    "expected_ref": actual_ref,
                                })
                                continue

                        var_seq = ref_seq[:mrna_pos] + coding_alt + ref_seq[mrna_pos + 1:]
                        try:
                            profile = _compute_folding_profile(ref_seq, var_seq, mrna_pos, folding_window_size)
                            entry = {
                                "genomic_position": pos, "mrna_position": mrna_pos,
                                "ref": ref, "alt": alt, "coding_alt": coding_alt,
                                "status": "ok", "delta_mfe": profile.delta_mfe,
                                "score": profile.score,
                            }
                            entry.update(_central_window_dict(profile))
                            per_mut.append(entry)
                        except Exception as e:
                            per_mut.append({
                                "genomic_position": pos, "ref": ref, "alt": alt,
                                "status": "fold_error", "delta_mfe": None, "error": str(e),
                            })

                    folding_result = {
                        "status": "completed",
                        "mut_id": mut_id,
                        "gene": m.gene,
                        "transcript_id": getattr(reference_transcript, "transcript_id", None),
                        "window_size": folding_window_size,
                        "mutations": per_mut,
                        "delta_mfe": per_mut[0]["delta_mfe"] if per_mut else None,
                    }
        except Exception as e:
            folding_result = {
                "status": "error",
                "reason": str(e),
                "mut_id": mut_id,
                "gene": m.gene,
            }

    # -- 4. Transcription (Borzoi) — runs independently, own seq extraction ---
    transcription_result = None
    if run_transcription:
        try:
            from .transcription.prediction import predict_transcription_effect
            # Use first mutation's coordinates for Borzoi
            positions = m.mutation_args()
            pos, ref, alt = positions[0]
            transcription_result = predict_transcription_effect(
                chrom=str(m.chrom), pos=pos,
                ref=ref, alt=alt, organism=organism,
            )
            transcription_result["mut_id"] = mut_id
            transcription_result["gene"] = m.gene
        except Exception as e:
            transcription_result = {
                "status": "error",
                "reason": str(e),
                "mut_id": mut_id,
                "gene": m.gene,
            }

    return MutationResult(
        splicing=splicing_result,
        tis=tis_result,
        folding=folding_result,
        transcription=transcription_result,
        transcript=mutated_transcript,
        tss_context=tss_context,
    )


# ============================================================================
# mRNA Folding Pipeline
# ============================================================================
# Re-export from the folding subpackage for discoverability alongside the
# other pipeline entry-points.
from .folding.pipeline import folding_pipeline, folding_query_batch  # noqa: E402
