# geney/pipelines.py
from __future__ import annotations

from datetime import datetime
import threading
import time
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
from .results import (
    SplicingResult,
    TISResult, TISSite,
    FoldingResult, FoldingWindow, FoldingMutation,
    TranscriptionResult,
    MutationResult,
)


# ============================================================================
# Internal helpers
# ============================================================================

_gene_cache: dict[tuple[str, str], Gene] = {}


def _build_splice_site_index(gene) -> dict[frozenset, str]:
    """Build a lookup from splice-site signature to transcript ID.

    Each key is a frozenset of (donor, acceptor) tuples; the value is the
    transcript_id of the annotated transcript that uses that exact splice
    pattern.  O(T) construction, O(1) lookup per isoform.
    """
    index: dict[frozenset, str] = {}
    transcripts = getattr(gene, 'transcripts', None) or getattr(gene, '_transcripts', {})
    if hasattr(transcripts, 'values'):
        transcript_list = list(transcripts.values())
    elif hasattr(transcripts, '__iter__'):
        transcript_list = list(transcripts)
    else:
        return index

    for t in transcript_list:
        tid = getattr(t, 'transcript_id', None)
        donors = tuple(sorted(getattr(t, 'donors', [])))
        acceptors = tuple(sorted(getattr(t, 'acceptors', [])))
        if tid and (donors or acceptors):
            index[frozenset(donors + acceptors)] = tid
    return index


def _match_known_transcript(variant_transcript, splice_index: dict[frozenset, str]) -> str | None:
    """Return the transcript ID if the variant's splice sites match a known transcript."""
    if not splice_index:
        return None
    donors = tuple(sorted(getattr(variant_transcript, 'donors', [])))
    acceptors = tuple(sorted(getattr(variant_transcript, 'acceptors', [])))
    return splice_index.get(frozenset(donors + acceptors))

_GENE_CACHE_MAX = 32


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.

    Caches up to _GENE_CACHE_MAX genes in memory to avoid repeated EFS reads
    for consecutive mutations in the same gene.
    """
    key = (gene_name, organism)
    if key in _gene_cache:
        return _gene_cache[key]

    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."
            )
        if exc:
            raise exc[0]
        gene = result[0] if result else None
    else:
        gene = Gene.from_file(gene_name, organism=organism)

    if gene is not None:
        if len(_gene_cache) >= _GENE_CACHE_MAX:
            _gene_cache.pop(next(iter(_gene_cache)))
        _gene_cache[key] = gene
    return gene


def _get_conservation_vector(transcript) -> np.ndarray:
    """Get conservation vector for a transcript, with fallback to ones."""
    cons_vector = getattr(transcript, 'cons_vector', None)
    if cons_vector is not None:
        return cons_vector
    protein_length = len(getattr(transcript, 'protein', ''))
    if protein_length == 0:
        protein_length = 1
    return np.ones(protein_length)


def _extract_gained_alternatives(result) -> list[dict]:
    """Extract alternative TIS sites that gained significant probability."""
    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 _splicing_base_report(mut_id, m, ref_transcript, splicing_engine, central_pos):
    """Build the shared base report Series for splicing results."""
    return pd.Series({
        "mut_id": mut_id,
        "gene": m.gene,
        "transcript_id": ref_transcript.transcript_id,
        "primary_transcript": getattr(ref_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"),
    })


def _oncosplice_dataframe(ref_transcript, ss, ss_metadata, base_report, max_isoforms=None, splice_index=None):
    """Run Oncosplice on all viable transcripts, return SplicingResult."""
    cons_vector = _get_conservation_vector(ref_transcript)
    rows = []
    for variant_transcript, isoform_metadata in ss.get_viable_transcripts(metadata=True, max_isoforms=max_isoforms):
        onco = Oncosplice(
            ref_transcript.protein,
            variant_transcript.protein,
            cons_vector,
        )
        matched_tid = _match_known_transcript(variant_transcript, splice_index) if splice_index else None
        row = pd.concat([
            isoform_metadata,
            pd.Series({
                "variant_mrna": variant_transcript.mature_mrna.seq,
                "known_transcript_id": matched_tid,
            }),
            onco.get_analysis_series(),
        ])
        rows.append(row)

    isoforms_df = pd.DataFrame(rows)

    # Drop columns that are stored once on SplicingResult, not per-isoform
    for col in ("reference_protein", "reference_length"):
        if col in isoforms_df.columns:
            isoforms_df = isoforms_df.drop(columns=col)

    missplicing = float(ss_metadata.get("missplicing", 0.0) or 0.0)

    max_oncosplice = 0.0
    if not isoforms_df.empty and "oncosplice_score" in isoforms_df.columns:
        max_oncosplice = float(isoforms_df["oncosplice_score"].max())

    return SplicingResult(
        score=missplicing,
        missplicing=missplicing,
        max_oncosplice_score=max_oncosplice,
        isoforms=isoforms_df,
        reference_protein=ref_transcript.protein,
        reference_mrna=ref_transcript.mature_mrna.seq,
        conservation_vector=cons_vector,
        mut_id=base_report.get("mut_id", ""),
        gene=base_report.get("gene", ""),
        transcript_id=base_report.get("transcript_id"),
        splicing_engine=base_report.get("splicing_engine", ""),
        central_position=int(base_report.get("central_position", 0)),
        donor_df=ss.donor_df,
        acceptor_df=ss.acceptor_df,
        splice_events=ss_metadata,
    )


def _mrna_mutation_positions(mrna, mutation_args):
    """Map genomic mutation positions to 0-based mRNA positions."""
    if not hasattr(mrna, "index"):
        return []
    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)}
    return [pos_to_mrna[pos] for pos, _, _ in mutation_args if pos in pos_to_mrna]


def _tis_completed_dict(result, mut_id, gene, transcript_id,
                        ref_mrna_seq, var_mrna_seq, mrna_positions):
    """Build a TISResult from a TISShiftResult."""
    from .tis import get_tis_summary

    atg_sites = result.analysis_df.to_dict("records") if not result.analysis_df.empty else []

    # Build canonical TISSite
    canonical_row = next((s for s in atg_sites if s.get("location") == "canonical"), None)
    canonical_site = None
    if canonical_row:
        canonical_site = TISSite(
            position=canonical_row["position"],
            location="canonical",
            in_frame=True,
            ref_probability=canonical_row["ref_probability"],
            var_probability=canonical_row["var_probability"],
            prob_change=canonical_row["prob_change"],
            ref_usage=canonical_row.get("ref_usage", 0.0),
            var_usage=canonical_row.get("var_usage", 0.0),
        )

    return TISResult(
        status="completed",
        score=result.tis_shift_score,
        canonical=canonical_site,
        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,
        initiation_intact=result.initiation_intact,
        gained_alternative_tis=(gained := _extract_gained_alternatives(result)),
        num_alternative_tis=len(gained),
        top_alternative_position=gained[0]["position"] if gained else None,
        top_alternative_prob=gained[0]["var_probability"] if gained else 0.0,
        top_alternative_location=gained[0]["location"] if gained else None,
        top_alternative_in_frame=gained[0]["in_frame"] if gained else None,
        atg_sites=atg_sites,
        ref_mrna_length=len(ref_mrna_seq),
        var_mrna_length=len(var_mrna_seq),
        ref_mrna_window=ref_mrna_seq[:600],
        var_mrna_window=var_mrna_seq[:600],
        mutation_mrna_positions=mrna_positions,
        mut_id=mut_id,
        gene=gene,
        transcript_id=transcript_id,
        summary=get_tis_summary(result),
        tis_result=result,
    )


def _prepare_tis(ref_transcript, m, mut_id, proximity_window=500,
                 skip_if_not_near_start=True):
    """Prepare TIS inputs. Returns (ref_mrna, var_mrna, positions) or a TISResult for skip/error."""
    from .tis.helpers import (
        _mutation_in_mature_mrna,
        _mutation_near_start_codon,
        _apply_mutation_to_mrna,
        _get_mature_mrna_seq,
    )

    tid = getattr(ref_transcript, "transcript_id", None)

    if not _mutation_in_mature_mrna(ref_transcript, m.positions):
        return TISResult(status="skipped", score=0.0, reason="mutation_not_in_mature_mrna",
                         mut_id=mut_id, gene=m.gene, transcript_id=tid)

    if skip_if_not_near_start:
        if not _mutation_near_start_codon(ref_transcript, m.positions, proximity_window):
            return TISResult(status="skipped", score=0.0, reason="mutation_far_from_start_codon",
                             mut_id=mut_id, gene=m.gene, transcript_id=tid,
                             proximity_window=proximity_window)

    ref_mrna_seq = _get_mature_mrna_seq(ref_transcript)
    if not ref_mrna_seq:
        return TISResult(status="error", score=0.0, reason="no_mature_mrna", mut_id=mut_id)

    mrna = ref_transcript.mature_mrna
    var_mrna_seq = (_apply_mutation_to_mrna(ref_mrna_seq, mrna.index, m.mutation_args())
                    if hasattr(mrna, "index") else ref_mrna_seq)

    positions = _mrna_mutation_positions(mrna, m.mutation_args())
    return ref_mrna_seq, var_mrna_seq, positions


def _transcription_result_from_dict(raw: dict, mut_id: str, gene: str) -> TranscriptionResult:
    """Convert a raw transcription prediction dict to a TranscriptionResult."""
    return TranscriptionResult(
        status=raw.get("status", "completed"),
        score=raw.get("regulatory_score", 0.0),
        tissue_fold_changes=raw.get("tissue_fold_changes", {}),
        tissue_predicted_levels=raw.get("tissue_predicted_levels", {}),
        max_lfc_tissue=raw.get("max_lfc_tissue", ""),
        max_lfc_value=raw.get("max_lfc_value", 0.0),
        max_abs_lfc=raw.get("max_abs_lfc", 0.0),
        num_affected_tissues=raw.get("num_affected_tissues", 0),
        mut_id=mut_id,
        gene=gene,
    )


def _run_folding_single(ref_transcript, m, mut_id, window_size):
    """Run folding analysis for a single mutation. Returns FoldingResult."""
    from .folding.delta_mfe import compute_folding_profile as _cfp
    from .folding.vienna import fold_available
    from .folding.reference import complement

    tid = getattr(ref_transcript, "transcript_id", None)

    if not fold_available():
        return FoldingResult(status="error", reason="ViennaRNA not installed", mut_id=mut_id)

    mrna = ref_transcript.mature_mrna
    if mrna is None or not hasattr(mrna, "seq") or not hasattr(mrna, "index"):
        return FoldingResult(status="error", reason="no_mature_mrna", mut_id=mut_id)

    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 = []
    central_window = None
    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, coding_alt = ref.upper(), 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 = _cfp(ref_seq, var_seq, mrna_pos, window_size)

            # Build FoldingWindow from the best-centered window
            fw = None
            if profile.windows:
                best = min(profile.windows,
                           key=lambda w: abs(w.mutation_offset - len(w.ref_seq) // 2))
                fw = FoldingWindow(
                    ref_sequence=best.ref_seq,
                    var_sequence=best.var_seq,
                    ref_structure=best.ref_structure,
                    var_structure=best.var_structure,
                    ref_mfe=best.ref_mfe,
                    var_mfe=best.var_mfe,
                    delta_mfe=best.var_mfe - best.ref_mfe,
                    mutation_offset=best.mutation_offset,
                )
                if central_window is None:
                    central_window = fw

            # Build legacy-compatible dict entry
            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}
            if fw:
                entry.update({
                    "ref_window": fw.ref_sequence,
                    "var_window": fw.var_sequence,
                    "ref_structure": fw.ref_structure,
                    "var_structure": fw.var_structure,
                    "mfe_ref": fw.ref_mfe,
                    "mfe_var": fw.var_mfe,
                    "mutation_offset": fw.mutation_offset,
                })
            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)})

    first_delta = per_mut[0]["delta_mfe"] if per_mut else None
    first_score = next((e.get("score", 0.0) for e in per_mut if e.get("status") == "ok"), 0.0)

    return FoldingResult(
        status="completed",
        score=first_score or 0.0,
        delta_mfe=first_delta,
        mutations=per_mut,
        central_window=central_window,
        mut_id=mut_id,
        gene=m.gene,
        transcript_id=tid,
        window_size=window_size,
    )


# ============================================================================
# Splicing Pipelines
# ============================================================================

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

    Returns SplicingResult with all viable isoforms and their oncosplice scores.
    Returns empty SplicingResult if no suitable transcript contains the mutation.
    """
    m = MutationalEvent(mut_id)
    if not m.compatible():
        raise ValueError(f"Mutations in event are incompatible: {mut_id}")

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

    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return SplicingResult()

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

    # Scoped pre-mRNA for splicing prediction
    ref_for_splicing = selected.clone()
    ref_for_splicing.generate_pre_mrna(
        region_start=central_pos - 7500,
        region_end=central_pos + 7500,
    )
    tl = TranscriptLibrary(ref_for_splicing, m, individual_mutations=False)
    tl.predict_splicing(central_pos, engine=splicing_engine, inplace=True)
    splicing_results = tl.get_event_columns("event")

    # Clear scoped pre-mRNA so isoform clones use direct-exon assembly
    tl.event._pre_mrna = None

    ss = SpliceSimulator(
        splicing_results, tl.event, feature="event", max_distance=100_000_000,
        mutations=list(m.mutation_args()),
    )

    base_report = _splicing_base_report(mut_id, m, reference_transcript, splicing_engine, central_pos)
    ss_metadata = ss.report(central_pos)
    splice_index = _build_splice_site_index(gene)
    return _oncosplice_dataframe(reference_transcript, ss, ss_metadata, base_report, max_isoforms=max_isoforms, splice_index=splice_index)


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."""
    result = oncosplice_pipeline(mut_id, transcript_id, splicing_engine, organism)
    if result.empty:
        return None
    variants = result.isoforms[result.isoforms["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."""
    m = MutationalEvent(mut_id)
    if not m.compatible():
        raise ValueError(f"Mutations in event are incompatible: {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 None

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

    # Scoped pre-mRNA for splicing prediction
    ref_for_splicing = selected.clone()
    ref_for_splicing.generate_pre_mrna(
        region_start=central_pos - 7500,
        region_end=central_pos + 7500,
    )
    tl = TranscriptLibrary(ref_for_splicing, m, individual_mutations=False)
    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")


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

    Applies the mutation directly to the mature mRNA (without splicing changes)
    and compares TIS probabilities between reference and variant.
    """
    from .tis import compute_tis_shift_metrics

    m = MutationalEvent(mut_id)
    if not m.compatible():
        raise ValueError(f"Mutations in event are incompatible: {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 None

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

    tis_inputs = _prepare_tis(reference_transcript, m, mut_id,
                              proximity_window, skip_if_not_near_start)
    if isinstance(tis_inputs, TISResult):
        return tis_inputs

    ref_mrna_seq, var_mrna_seq, mrna_positions = tis_inputs

    try:
        result = compute_tis_shift_metrics(
            ref_mrna=ref_mrna_seq,
            var_mrna=var_mrna_seq,
            canonical_start=None,
            probability_threshold=probability_threshold,
            use_prior=use_prior,
        )
        return _tis_completed_dict(
            result, mut_id, m.gene, reference_transcript.transcript_id,
            ref_mrna_seq, var_mrna_seq, mrna_positions,
        )
    except Exception as e:
        return TISResult(status="error", score=0.0, reason=str(e),
                         mut_id=mut_id, gene=m.gene)


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.

    Fast pre-screening that doesn't run the full TITER model.
    """
    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_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)

    affects_start = False
    try:
        mrna = reference_transcript.mature_mrna
        if hasattr(mrna, 'seq') and hasattr(mrna, 'index'):
            seq = mrna.seq.upper()
            atg_pos = seq.find('ATG')
            if atg_pos != -1 and atg_pos < len(mrna.index):
                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

    return {
        'should_analyze': mutation_in_mrna and (near_start or affects_start),
        '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",
    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,
    max_isoforms: int | 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
    subsystem.  Returns a :class:`MutationResult` namedtuple.

    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"``).
        context_bp: Base pairs on each side of mutation 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.
    """
    from .utils import _apply_mutation_safe

    _t_total = time.perf_counter()
    timing: dict[str, float] = {}
    errors: dict[str, str] = {}

    # -- Parse ----------------------------------------------------------------
    _t0 = time.perf_counter()
    m = MutationalEvent(mut_id)
    if not m.compatible():
        raise ValueError(f"Mutations in event are incompatible: {mut_id}")
    central_pos = m.central_position
    timing["parse"] = round(time.perf_counter() - _t0, 4)

    # -- Gene load ------------------------------------------------------------
    _t0 = time.perf_counter()
    gene = _load_gene_with_timeout(m.gene, organism, gene_load_timeout)
    timing["gene_load"] = round(time.perf_counter() - _t0, 4)
    if gene is None:
        timing["total"] = round(time.perf_counter() - _t_total, 4)
        return MutationResult(timing=timing)

    # -- Transcript prep ------------------------------------------------------
    _t0 = time.perf_counter()
    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        timing["transcript_prep"] = round(time.perf_counter() - _t0, 4)
        timing["total"] = round(time.perf_counter() - _t_total, 4)
        return MutationResult(timing=timing)

    reference_transcript = (
        selected
        .generate_mature_mrna()
        .generate_protein()
    )
    timing["transcript_prep"] = round(time.perf_counter() - _t0, 4)

    # -- Mutate transcript ----------------------------------------------------
    _t0 = time.perf_counter()
    mutated_transcript = selected.clone().generate_mature_mrna()
    for pos, ref, alt in m.mutation_args():
        try:
            _apply_mutation_safe(mutated_transcript.mature_mrna, pos, ref, alt)
        except Exception:
            pass

    # Scoped pre-mRNA read for mutation context (includes intronic sequence)
    ctx_transcript = selected.clone()
    ctx_transcript.generate_pre_mrna(
        region_start=central_pos - context_bp,
        region_end=central_pos + context_bp,
    )
    for pos, ref, alt in m.mutation_args():
        try:
            _apply_mutation_safe(ctx_transcript.pre_mrna, pos, ref, alt)
        except Exception:
            pass
    mutation_context = ctx_transcript.pre_mrna

    mutated_transcript.generate_protein()
    timing["mutate_transcript"] = round(time.perf_counter() - _t0, 4)

    # -- 1. Splicing / Oncosplice ---------------------------------------------
    splicing_result = SplicingResult()
    if run_splicing:
        _t0 = time.perf_counter()
        try:
            # Use scoped pre-mRNA read for splicing prediction — only fetch ~15kb
            ref_for_splicing = selected.clone()
            ref_for_splicing.generate_pre_mrna(
                region_start=central_pos - 7500,
                region_end=central_pos + 7500,
            )
            tl = TranscriptLibrary(ref_for_splicing, m, individual_mutations=False)
            tl.predict_splicing(central_pos, engine=splicing_engine, inplace=True)
            splicing_preds = tl.get_event_columns("event")

            # Clear scoped pre-mRNA so isoform clones use direct-exon assembly
            tl.event._pre_mrna = None

            ss = SpliceSimulator(
                splicing_preds, tl.event, feature="event", max_distance=100_000_000,
                mutations=list(m.mutation_args()),
            )

            base_report = _splicing_base_report(mut_id, m, reference_transcript, splicing_engine, central_pos)
            ss_metadata = ss.report(central_pos)
            splice_index = _build_splice_site_index(gene)
            splicing_result = _oncosplice_dataframe(reference_transcript, ss, ss_metadata, base_report, max_isoforms=max_isoforms, splice_index=splice_index)
        except Exception as e:
            splicing_result = SplicingResult()
            errors["splicing"] = str(e)
        timing["spliceai"] = round(time.perf_counter() - _t0, 4)

    # -- 1b. Oncosplice scoring (included in spliceai timing above) -----------
    timing["oncosplice"] = 0.0  # oncosplice runs inside the splicing block

    # -- 2. TIS ---------------------------------------------------------------
    tis_result = None
    if run_tis:
        _t0 = time.perf_counter()
        try:
            from .tis import compute_tis_shift_metrics

            tis_inputs = _prepare_tis(reference_transcript, m, mut_id, tis_proximity_window)
            if isinstance(tis_inputs, TISResult):
                tis_result = tis_inputs
            else:
                ref_mrna_seq, var_mrna_seq, mrna_positions = tis_inputs
                result = compute_tis_shift_metrics(
                    ref_mrna=ref_mrna_seq,
                    var_mrna=var_mrna_seq,
                    canonical_start=None,
                    probability_threshold=tis_probability_threshold,
                )
                tis_result = _tis_completed_dict(
                    result, mut_id, m.gene, reference_transcript.transcript_id,
                    ref_mrna_seq, var_mrna_seq, mrna_positions,
                )
        except Exception as e:
            tis_result = TISResult(status="error", score=0.0, reason=str(e),
                                   mut_id=mut_id, gene=m.gene)
            errors["tis"] = str(e)
        timing["tis"] = round(time.perf_counter() - _t0, 4)

    # -- 3. mRNA Folding ------------------------------------------------------
    folding_result = None
    if run_folding:
        _t0 = time.perf_counter()
        try:
            folding_result = _run_folding_single(reference_transcript, m, mut_id, folding_window_size)
        except Exception as e:
            folding_result = FoldingResult(status="error", reason=str(e),
                                           mut_id=mut_id, gene=m.gene)
            errors["folding"] = str(e)
        timing["folding"] = round(time.perf_counter() - _t0, 4)

    # -- 4. Transcription (Borzoi) --------------------------------------------
    transcription_result = None
    if run_transcription:
        _t0 = time.perf_counter()
        try:
            from .transcription.prediction import predict_transcription_effect
            raw = predict_transcription_effect(
                chrom=str(m.chrom), pos=central_pos,
                mutations=m.mutation_args(), organism=organism,
            )
            transcription_result = _transcription_result_from_dict(raw, mut_id, m.gene)
        except Exception as e:
            transcription_result = TranscriptionResult(
                status="error", reason=str(e), mut_id=mut_id, gene=m.gene)
            errors["borzoi"] = str(e)
        timing["borzoi"] = round(time.perf_counter() - _t0, 4)

    timing["total"] = round(time.perf_counter() - _t_total, 4)

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


# ============================================================================
# Batch-optimized Mutation Pipeline
# ============================================================================

def mutation_pipeline_batch(
    mut_ids: list[str],
    transcript_id: str | None = None,
    organism: str = "hg38",
    splicing_engine: str = "spliceai",
    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,
    max_gpu_batch: int | None = None,
    n_folding_workers: int = 4,
    n_io_workers: int = 8,
    max_isoforms: int | None = None,
) -> list[MutationResult]:
    """Batch-optimized mutation pipeline.

    Returns identical results to calling :func:`mutation_pipeline` N times,
    but internally batches GPU inference across mutations for higher throughput.

    Execution phases (CPU and GPU work overlapped where possible):
    1. PREPARE — Parse mutations, deduplicate gene loads, build transcripts.
    2. Launch FOLDING (CPU) + BORZOI SEQ EXTRACTION (I/O) in background.
    3. SPLICING — Batched GPU prediction, per-mutation Oncosplice assembly.
    4. TIS — Cross-mutation mega-batch via ``compute_tis_shift_metrics_batch``.
    5. Collect FOLDING results (already complete from background).
    6. TRANSCRIPTION — GPU inference on pre-extracted sequences.
    7. ASSEMBLE — Build MutationResult per mutation.
    """
    from concurrent.futures import ThreadPoolExecutor
    from .utils import _apply_mutation_safe

    if max_gpu_batch is None:
        max_gpu_batch = 8

    n = len(mut_ids)
    if n == 0:
        return []

    # ================================================================
    # Phase 1: PREPARE
    # ================================================================

    events = []
    for mid in mut_ids:
        m = MutationalEvent(mid)
        if not m.compatible():
            raise ValueError(f"Mutations in event are incompatible: {mid}")
        events.append(m)

    gene_names = list({m.gene for m in events})
    gene_cache: dict = {}

    def _load_gene(name):
        return name, _load_gene_with_timeout(name, organism, gene_load_timeout)

    with ThreadPoolExecutor(max_workers=min(n_io_workers, len(gene_names))) as pool:
        for name, gene in pool.map(_load_gene, gene_names):
            gene_cache[name] = gene

    prep: list[dict] = []
    for mid, m in zip(mut_ids, events):
        gene = gene_cache.get(m.gene)
        central_pos = m.central_position

        if gene is None:
            prep.append({"status": "no_gene", "event": m, "mut_id": mid})
            continue

        selected = select_transcript(gene, central_pos, transcript_id)
        if selected is None:
            prep.append({"status": "no_transcript", "event": m, "mut_id": mid})
            continue

        ref_transcript = (
            selected
            .generate_mature_mrna()
            .generate_protein()
        )

        mut_transcript = selected.clone().generate_mature_mrna()
        for pos, ref, alt in m.mutation_args():
            try:
                _apply_mutation_safe(mut_transcript.mature_mrna, pos, ref, alt)
            except Exception:
                pass

        # Scoped pre-mRNA read for mutation context (intronic context)
        ctx_transcript = selected.clone()
        ctx_transcript.generate_pre_mrna(
            region_start=central_pos - context_bp,
            region_end=central_pos + context_bp,
        )
        for pos, ref, alt in m.mutation_args():
            try:
                _apply_mutation_safe(ctx_transcript.pre_mrna, pos, ref, alt)
            except Exception:
                pass
        mutation_context = ctx_transcript.pre_mrna

        mut_transcript.generate_protein()

        prep.append({
            "status": "ok",
            "event": m,
            "mut_id": mid,
            "central_pos": central_pos,
            "ref_transcript": ref_transcript,
            "mut_transcript": mut_transcript,
            "mutation_context": mutation_context,
        })

    # ================================================================
    # Phase 2+4: Launch FOLDING (CPU) and BORZOI SEQ EXTRACTION (I/O)
    #            in background, then run GPU phases
    # ================================================================

    folding_results: list[FoldingResult | None] = [None] * n
    folding_future = None

    if run_folding:
        fold_args = [(i, p) for i, p in enumerate(prep) if p["status"] == "ok"]
        if fold_args:
            def _fold_one(args):
                i, p = args
                try:
                    return i, _run_folding_single(
                        p["ref_transcript"], p["event"], p["mut_id"], folding_window_size
                    )
                except Exception as e:
                    return i, FoldingResult(status="error", reason=str(e), mut_id=p["mut_id"])

            # Launch folding in background — CPU work overlaps with GPU phases
            from concurrent.futures import Future
            fold_pool = ThreadPoolExecutor(max_workers=min(n_folding_workers, len(fold_args)))
            folding_future = [fold_pool.submit(_fold_one, args) for args in fold_args]

    # Pre-extract Borzoi sequences in background (I/O bound, overlaps with GPU)
    transcription_results: list[TranscriptionResult | None] = [None] * n
    trans_eligible: list[int] = []
    trans_seq_pairs = None

    if run_transcription:
        from .transcription.prediction import _extract_sequences as _borzoi_extract

        trans_specs_raw: list[dict] = []
        for i, p in enumerate(prep):
            if p["status"] != "ok":
                continue
            m = p["event"]
            trans_eligible.append(i)
            trans_specs_raw.append({
                "chrom": str(m.chrom),
                "pos": p["central_pos"],
                "mutations": m.mutation_args(),
                "organism": organism,
            })

        if trans_specs_raw:
            def _extract_one(spec):
                return _borzoi_extract(
                    spec["chrom"], spec["pos"], spec["mutations"],
                    spec.get("organism", "hg38"),
                )

            # Start sequence extraction in background (disk I/O)
            io_pool = ThreadPoolExecutor(max_workers=min(n_io_workers, len(trans_specs_raw)))
            trans_seq_future = io_pool.map(_extract_one, trans_specs_raw)

    # ================================================================
    # Phase 2: SPLICING (GPU)
    # ================================================================

    splicing_results: list[SplicingResult] = [SplicingResult() for _ in range(n)]

    if run_splicing:
        from .splicing.engines import (
            _prepare_splice_input,
            _assemble_splice_df,
            batch_run_splicing_engine,
            adjoin_splicing_outcomes as _adjoin,
        )

        splice_jobs: list[dict] = []
        tl_map: dict[int, TranscriptLibrary] = {}

        for i, p in enumerate(prep):
            if p["status"] != "ok":
                continue
            try:
                m = p["event"]
                # Create a transcript with scoped pre-mRNA for splicing prediction
                ref_for_splicing = p["ref_transcript"].clone()
                ref_for_splicing.generate_pre_mrna(
                    region_start=p["central_pos"] - 7500,
                    region_end=p["central_pos"] + 7500,
                )
                tl = TranscriptLibrary(ref_for_splicing, m, individual_mutations=False)
                tl_map[i] = tl
                for key, t in tl:
                    sp = _prepare_splice_input(t.pre_mrna, p["central_pos"])
                    splice_jobs.append({"mut_idx": i, "tl_key": key, "prep": sp})
            except Exception:
                splicing_results[i] = SplicingResult()

        if splice_jobs:
            all_seqs = [j["prep"]["padded_seq"] for j in splice_jobs]
            try:
                all_preds = batch_run_splicing_engine(all_seqs, engine=splicing_engine)
            except Exception:
                all_preds = []
                for seq in all_seqs:
                    try:
                        from .splicing.engines import run_splicing_engine as _rse
                        all_preds.append(_rse(seq=seq, engine=splicing_engine))
                    except Exception:
                        all_preds.append(([], []))

            from collections import defaultdict
            per_mut_splicing: dict[int, dict[str, pd.DataFrame]] = defaultdict(dict)

            for job, (donor_probs, acceptor_probs) in zip(splice_jobs, all_preds):
                try:
                    df = _assemble_splice_df(donor_probs, acceptor_probs, job["prep"])
                    per_mut_splicing[job["mut_idx"]][job["tl_key"]] = df
                except Exception:
                    pass

            splice_index_cache: dict[str, dict] = {}
            for i, preds_dict in per_mut_splicing.items():
                if i not in tl_map:
                    continue
                tl = tl_map[i]
                p = prep[i]
                m = p["event"]
                ref_transcript = p["ref_transcript"]
                central_pos = p["central_pos"]

                try:
                    tl.splicing_results = _adjoin(preds_dict, tl.ref)
                    splicing_preds = tl.get_event_columns("event")

                    # Clear scoped pre-mRNA so isoform clones use direct-exon assembly
                    tl.event._pre_mrna = None

                    ss = SpliceSimulator(
                        splicing_preds, tl.event, feature="event", max_distance=100_000_000,
                        mutations=list(m.mutation_args()),
                    )

                    base_report = _splicing_base_report(p["mut_id"], m, ref_transcript, splicing_engine, central_pos)
                    ss_metadata = ss.report(central_pos)
                    if m.gene not in splice_index_cache:
                        gene = gene_cache.get(m.gene)
                        splice_index_cache[m.gene] = _build_splice_site_index(gene) if gene else {}
                    splicing_results[i] = _oncosplice_dataframe(ref_transcript, ss, ss_metadata, base_report, max_isoforms=max_isoforms, splice_index=splice_index_cache[m.gene])
                except Exception:
                    splicing_results[i] = SplicingResult()

    # ================================================================
    # Phase 3: TIS (cross-mutation mega-batch, GPU)
    # ================================================================

    tis_results: list[TISResult | None] = [None] * n

    if run_tis:
        from .tis.metrics import compute_tis_shift_metrics_batch

        tis_eligible: list[int] = []
        tis_pairs: list[tuple[str, str]] = []
        tis_meta: list[dict] = []

        for i, p in enumerate(prep):
            if p["status"] != "ok":
                continue

            tis_inputs = _prepare_tis(p["ref_transcript"], p["event"], p["mut_id"], tis_proximity_window)
            if isinstance(tis_inputs, TISResult):
                tis_results[i] = tis_inputs
                continue

            ref_mrna_seq, var_mrna_seq, mrna_positions = tis_inputs
            tis_eligible.append(i)
            tis_pairs.append((ref_mrna_seq, var_mrna_seq))
            tis_meta.append({
                "mut_id": p["mut_id"],
                "gene": p["event"].gene,
                "transcript_id": p["ref_transcript"].transcript_id,
                "ref_mrna_seq": ref_mrna_seq,
                "var_mrna_seq": var_mrna_seq,
                "mrna_positions": mrna_positions,
            })

        if tis_pairs:
            try:
                batch_results = compute_tis_shift_metrics_batch(
                    tis_pairs, probability_threshold=tis_probability_threshold,
                )
                for j, (idx, result) in enumerate(zip(tis_eligible, batch_results)):
                    meta = tis_meta[j]
                    tis_results[idx] = _tis_completed_dict(
                        result, meta["mut_id"], meta["gene"], meta["transcript_id"],
                        meta["ref_mrna_seq"], meta["var_mrna_seq"], meta["mrna_positions"],
                    )
            except Exception as e:
                for j, idx in enumerate(tis_eligible):
                    meta = tis_meta[j]
                    tis_results[idx] = TISResult(
                        status="error", score=0.0, reason=str(e),
                        mut_id=meta["mut_id"], gene=meta["gene"])

    # ================================================================
    # Collect FOLDING results (should be done by now — ran during GPU phases)
    # ================================================================

    if folding_future is not None:
        for fut in folding_future:
            idx, result = fut.result()
            folding_results[idx] = result
        fold_pool.shutdown(wait=False)

    # ================================================================
    # Phase 5: TRANSCRIPTION (GPU — sequences pre-extracted during GPU phases)
    # ================================================================

    if run_transcription and trans_eligible:
        # Wait for background sequence extraction to finish
        try:
            trans_seq_pairs = list(trans_seq_future)
            io_pool.shutdown(wait=False)
        except Exception:
            trans_seq_pairs = None

    if run_transcription and trans_eligible and trans_seq_pairs:
        from .transcription.model import _load_borzoi
        from .transcription.prediction import (
            _predict_tracks_batch,
            _compute_tissue_changes,
            _build_gtex_mapping,
            _normalize_score,
        )

        try:
            model, device, dtype, use_autocast = _load_borzoi()
            gtex_map = _build_gtex_mapping()

            # Flatten all WT+mut sequences for batched inference
            all_seqs: list[str] = []
            for wt, mut in trans_seq_pairs:
                all_seqs.append(wt)
                all_seqs.append(mut)

            # Chunked GPU forward passes
            all_tracks: list[np.ndarray] = []
            for i in range(0, len(all_seqs), max_gpu_batch):
                chunk = all_seqs[i : i + max_gpu_batch]
                all_tracks.extend(
                    _predict_tracks_batch(model, chunk, device, dtype, use_autocast)
                )

            # Compute per-mutation tissue changes
            for j, idx in enumerate(trans_eligible):
                p = prep[idx]
                wt_tracks = all_tracks[j * 2]
                mut_tracks = all_tracks[j * 2 + 1]
                tissue_data = _compute_tissue_changes(wt_tracks, mut_tracks, gtex_map)
                score = _normalize_score(tissue_data["max_abs_lfc"])
                tissue_data["regulatory_score"] = round(score, 4)
                tissue_data["status"] = "completed"
                transcription_results[idx] = _transcription_result_from_dict(
                    tissue_data, p["mut_id"], p["event"].gene)
        except Exception as e:
            for idx in trans_eligible:
                p = prep[idx]
                transcription_results[idx] = TranscriptionResult(
                    status="error", reason=str(e),
                    mut_id=p["mut_id"], gene=p["event"].gene)

    # ================================================================
    # Phase 6: ASSEMBLE
    # ================================================================

    results: list[MutationResult] = []
    for i in range(n):
        p = prep[i]
        if p["status"] != "ok":
            results.append(MutationResult())
        else:
            results.append(MutationResult(
                splicing=splicing_results[i],
                tis=tis_results[i],
                folding=folding_results[i],
                transcription=transcription_results[i],
                transcript=p["mut_transcript"],
                mutation_context=p["mutation_context"],
            ))

    return results


# ============================================================================
# mRNA Folding Pipeline (re-export)
# ============================================================================
from .folding.pipeline import folding_pipeline, folding_query_batch  # noqa: E402
