# 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,
)


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 - SPLICING_CONTEXT_BP,
        region_end=central_pos + 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,
        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 - SPLICING_CONTEXT_BP,
        region_end=central_pos + 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
    )

    return ss.max_splicing_delta("event_prob")
