# geney/pipelines/_splicing.py — Splicing-specific pipelines.
from __future__ import annotations

from typing import Optional

import pandas as pd
from seqmat import Gene

from ..splicing import SpliceSimulator, TranscriptLibrary
from ..variants import MutationalEvent
from ..results import SplicingResult
from ..utils import select_transcript

from ._common import (
    SPLICING_CONTEXT_BP,
    SPLICE_MAX_DISTANCE,
    _load_gene_with_timeout,
    _build_splice_site_index,
    _splicing_base_report,
    _oncosplice_dataframe,
    _check_branch_point,
)


def oncosplice_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    splicing_engine: str = "ensemble",
    organism: str = "hg38",
    gene_load_timeout: Optional[float] = None,
    max_isoforms: int | None = None,
    min_total_prob: float = 0.5,
    min_delta_magnitude: float = 0.2,
) -> 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.

    Args:
        min_total_prob: Minimum combined donor/acceptor probability for an isoform
            path to be considered viable (default 0.5). Lower values include more
            isoforms; higher values restrict to higher-probability paths.
        min_delta_magnitude: Minimum magnitude of probability change to count a
            gain or loss (default 0.2). Both gains (e.g. +0.25) and losses
            (e.g. -0.25) are included when |Δprob| >= this value.
    """
    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(status="no_gene", reason=f"Gene not found: {m.gene}")
    central_pos = m.central_position

    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return SplicingResult(status="no_transcript",
                              reason=f"No transcript covers position {central_pos}")

    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(
        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_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=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)
    result = _oncosplice_dataframe(reference_transcript, ss, ss_metadata, base_report, max_isoforms=max_isoforms, splice_index=splice_index)

    # Branch point check
    bp_disrupted, bp_info = _check_branch_point(
        reference_transcript, m.positions, ref_for_splicing.pre_mrna)
    result.branch_point_disrupted = bp_disrupted
    result.branch_point_info = bp_info

    return result


def oncosplice_top_isoform(
    mut_id: str,
    transcript_id: str | None = None,
    splicing_engine: str = "ensemble",
    organism: str = "hg38",
    min_total_prob: float = 0.5,
    min_delta_magnitude: float = 0.2,
) -> pd.Series | None:
    """Get the most likely non-reference isoform for a mutation."""
    result = oncosplice_pipeline(
        mut_id, transcript_id, splicing_engine, organism,
        min_total_prob=min_total_prob, min_delta_magnitude=min_delta_magnitude,
    )
    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 = "ensemble",
    organism: str = "hg38",
    min_total_prob: float = 0.5,
    min_delta_magnitude: float = 0.2,
) -> 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(
        upstream=SPLICING_CONTEXT_BP,
        downstream=SPLICING_CONTEXT_BP,
    )
    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=SPLICE_MAX_DISTANCE,
        min_total_prob=min_total_prob,
        min_delta_magnitude=min_delta_magnitude,
    )

    return ss.max_splicing_delta("event_prob")
