# geney/pipelines/_mutation.py — Unified single-mutation pipeline.
from __future__ import annotations

import logging
import time

from ..results import (
    SplicingResult,
    ProteinResult,
    TISResult,
    FoldingResult,
    TranscriptionResult,
    MiRNAResult,
    MutationResult,
)
from ..splicing import SpliceSimulator, TranscriptLibrary

from ._common import (
    SPLICING_CONTEXT_BP,
    SPLICE_MAX_DISTANCE,
    _prepare_mutation,
    _build_splice_site_index,
    _splicing_base_report,
    _oncosplice_dataframe,
    _prepare_tis,
    _tis_completed_dict,
    _run_folding_single,
    _transcription_result_from_dict,
)

logger = logging.getLogger(__name__)


def mutation_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
    splicing_engine: str = "ensemble",
    context_bp: int = 1000,
    run_splicing: bool = True,
    run_all_transcripts: bool = False,
    run_tis: bool = True,
    run_folding: bool = True,
    run_transcription: bool = True,
    run_mirna: bool = False,
    mirna_seeds: dict[str, str] | None = None,
    mirna_full_prediction: bool = True,
    tis_proximity_window: int = 500,
    tis_probability_threshold: float = 0.1,
    folding_window_size: int = 39,
    folding_compute_null: bool = False,
    gene_load_timeout: float | None = None,
    max_isoforms: int | None = None,
    min_total_prob: float = 0.5,
    min_delta_magnitude: float = 0.2,
) -> 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.
        max_isoforms: Max number of isoforms to score (splicing); *None* = no limit.
        min_total_prob: Min combined donor/acceptor prob for viable isoform paths (default 0.5).
        min_delta_magnitude: Min |Δprob| to count a splice site gain or loss (default 0.2).
    """
    _t_total = time.perf_counter()
    timing: dict[str, float] = {}
    errors: dict[str, str] = {}

    # -- Prepare (parse → gene → transcript → mutate → context) ---------------
    _t0 = time.perf_counter()
    p = _prepare_mutation(mut_id, organism, transcript_id, context_bp, gene_load_timeout)
    timing["prepare"] = round(time.perf_counter() - _t0, 4)

    if p["status"] != "ok":
        timing["total"] = round(time.perf_counter() - _t_total, 4)
        errors["prepare"] = p["status"]
        return MutationResult(
            splicing=SplicingResult(status=p["status"],
                                   reason=f"Prepare failed: {p['status']}"),
            timing=timing,
            errors=errors,
        )

    m = p["event"]
    central_pos = p["central_pos"]
    gene = p["gene"]
    selected = p["selected"]
    reference_transcript = p["ref_transcript"]
    mutated_transcript = p["mut_transcript"]
    mutation_context = p["mutation_context"]

    # -- 1. Splicing / Oncosplice ---------------------------------------------
    splicing_result = SplicingResult()
    _shared_splicing_preds = None  # Saved for multi-transcript reuse
    multi_transcript_result = None
    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(
                upstream=SPLICING_CONTEXT_BP,
                downstream=SPLICING_CONTEXT_BP,
            )
            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")

            # Save for multi-transcript reuse (zero cost — just a reference)
            _shared_splicing_preds = splicing_preds

            # 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=SPLICE_MAX_DISTANCE,
                min_total_prob=min_total_prob,
                min_delta_magnitude=min_delta_magnitude,
                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)

    # -- 1a. Multi-transcript splicing (reuses shared SpliceAI predictions) ---
    if run_all_transcripts and _shared_splicing_preds is not None:
        _t0 = time.perf_counter()
        try:
            from ._multi_transcript import analyze_all_transcripts
            multi_transcript_result = analyze_all_transcripts(
                gene=gene,
                mutation_event=m,
                splicing_preds=_shared_splicing_preds,
                primary_transcript=reference_transcript,
                primary_splicing_result=splicing_result,
                splicing_engine=splicing_engine,
                central_pos=central_pos,
                max_isoforms=max_isoforms,
                min_total_prob=min_total_prob,
                min_delta_magnitude=min_delta_magnitude,
            )
        except Exception as e:
            logger.debug("Multi-transcript analysis failed for %s: %s", mut_id, e)
            errors["multi_transcript"] = str(e)
        timing["multi_transcript"] = round(time.perf_counter() - _t0, 4)

    # -- 1b. Protein-level analysis (direct ORF translation) -------------------
    protein_result = None
    _t0 = time.perf_counter()
    try:
        ref_prot = reference_transcript.protein or ""
        var_prot = getattr(mutated_transcript, "protein", "") or ""
        if ref_prot:
            from ..protein import analyze_protein_change
            from ._common import _get_conservation_vector
            cons = _get_conservation_vector(reference_transcript)
            protein_result = ProteinResult(
                **{k: v for k, v in analyze_protein_change(
                    ref_prot, var_prot, cons,
                    mut_id=mut_id, gene=m.gene,
                    transcript_id=reference_transcript.transcript_id,
                ).__dict__.items() if k in ProteinResult.__dataclass_fields__}
            )
    except Exception as e:
        protein_result = ProteinResult(status="error", reason=str(e),
                                       mut_id=mut_id, gene=m.gene)
        errors["protein"] = str(e)
    timing["protein"] = round(time.perf_counter() - _t0, 4)

    # -- 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,
                                                    compute_null=folding_compute_null)
        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)

    # -- 5. miRNA binding site disruption ----------------------------------------
    mirna_result = None
    if run_mirna:
        # Auto-load bundled miRNA panel if no seeds provided
        if not mirna_seeds:
            from ..mirna.reference import get_mirna_panel
            mirna_seeds = get_mirna_panel()
    if run_mirna and mirna_seeds:
        _t0 = time.perf_counter()
        try:
            from ..mirna import analyze_mirna_disruption

            ref_mrna_seq = getattr(reference_transcript.mature_mrna, "seq", "")
            var_mrna_obj = getattr(mutated_transcript, "mature_mrna", None)
            var_mrna_seq = getattr(var_mrna_obj, "seq", "") if var_mrna_obj else ""

            if ref_mrna_seq and var_mrna_seq:
                cds_len = len(reference_transcript.protein or "") * 3
                utr3_start = cds_len if cds_len > 0 else None

                result = analyze_mirna_disruption(
                    ref_mrna=ref_mrna_seq,
                    var_mrna=var_mrna_seq,
                    mirna_seeds=mirna_seeds,
                    utr3_start=utr3_start,
                    mut_id=mut_id,
                    gene=m.gene,
                )

                # ── Full miBSIM prediction (on changed sites only) ──────
                # Only runs when mirna_full_prediction=True AND sites actually changed.
                # Skips entirely when no gained/lost/modified sites detected.
                mibsim_repression_ref = None
                mibsim_repression_var = None
                has_changes = result.n_gained + result.n_lost + result.n_modified > 0
                if mirna_full_prediction and has_changes:
                    try:
                        from ..mirna.features import compute_site_features
                        from ..mirna.mibsim import predict_transcript_repression

                        utr5 = getattr(reference_transcript, "utr5_seq",
                                getattr(reference_transcript, "utr5", "")) or ""
                        orf_seq = getattr(reference_transcript, "cds_seq",
                                  getattr(reference_transcript, "cds", "")) or ""
                        utr3_seq = getattr(reference_transcript, "utr3_seq",
                                   getattr(reference_transcript, "utr3", "")) or ""

                        if not orf_seq and ref_mrna_seq:
                            prot_len = len(reference_transcript.protein or "")
                            orf_len = prot_len * 3
                            utr5 = utr5 or ref_mrna_seq[:max(0, len(ref_mrna_seq) - orf_len)]
                            orf_seq = ref_mrna_seq[len(utr5):len(utr5) + orf_len]
                            utr3_seq = utr3_seq or ref_mrna_seq[len(utr5) + orf_len:]

                        if orf_seq:
                            cons_vec = _get_conservation_vector(reference_transcript)
                            has_real_cons = cons_vec is not None and not (
                                hasattr(cons_vec, 'size') and cons_vec.size > 0
                                and __import__('numpy').allclose(cons_vec, 1.0))

                            def _build_site_inputs(sites):
                                inputs = []
                                for site in sites:
                                    mir_seq = mirna_seeds.get(site.mirna_id, "")
                                    if not mir_seq:
                                        continue
                                    feats = compute_site_features(
                                        utr5=utr5, orf=orf_seq, utr3=utr3_seq,
                                        mirna=mir_seq, rna_start=site.position,
                                        seed_length=site.end - site.position,
                                        seed_type=site.seed_type,
                                        region="UTR3" if site.region == "utr3" else "ORF",
                                        with_conservation=has_real_cons, with_thermo=False)
                                    inputs.append({"features": feats, "position": site.position,
                                                   "region": site.region, "seed_type": site.seed_type})
                                return inputs

                            if result.ref_sites:
                                ref_inputs = _build_site_inputs(result.ref_sites)
                                if ref_inputs:
                                    mibsim_repression_ref = predict_transcript_repression(ref_inputs)["total_repression"]
                            if result.var_sites:
                                var_inputs = _build_site_inputs(result.var_sites)
                                if var_inputs:
                                    mibsim_repression_var = predict_transcript_repression(var_inputs)["total_repression"]

                    except Exception as e:
                        logger.debug("miBSIM full prediction failed for %s: %s", mut_id, e)

                mirna_result = MiRNAResult(
                    status=result.status,
                    score=result.score,
                    n_sites_ref=result.n_sites_ref,
                    n_sites_var=result.n_sites_var,
                    n_gained=result.n_gained,
                    n_lost=result.n_lost,
                    n_modified=result.n_modified,
                    n_unchanged=result.n_unchanged,
                    total_repression_ref=mibsim_repression_ref if mibsim_repression_ref is not None else result.total_repression_ref,
                    total_repression_var=mibsim_repression_var if mibsim_repression_var is not None else result.total_repression_var,
                    repression_delta=round(
                        (mibsim_repression_var or result.total_repression_var)
                        - (mibsim_repression_ref or result.total_repression_ref), 4),
                    top_affected_mirnas=result.top_affected_mirnas,
                    site_changes=result.site_changes,
                    n_mirnas_scanned=result.n_mirnas_scanned,
                    mut_id=mut_id,
                    gene=m.gene,
                )
            else:
                mirna_result = MiRNAResult(
                    status="skipped", reason="no_mature_mrna",
                    mut_id=mut_id, gene=m.gene)
        except Exception as e:
            mirna_result = MiRNAResult(
                status="error", reason=str(e), mut_id=mut_id, gene=m.gene)
            errors["mirna"] = str(e)
        timing["mirna"] = round(time.perf_counter() - _t0, 4)

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

    return MutationResult(
        splicing=splicing_result,
        multi_transcript=multi_transcript_result,
        protein=protein_result,
        tis=tis_result,
        folding=folding_result,
        transcription=transcription_result,
        mirna=mirna_result,
        transcript=mutated_transcript,
        mutation_context=mutation_context,
        timing=timing,
        errors=errors,
    )
