# geney/pipelines/_multi_transcript.py — Multi-transcript splicing from shared SpliceAI.
"""
Reuses a single SpliceAI forward pass to analyze splicing impact across ALL
protein-coding transcripts of a gene. The expensive neural network inference
runs once; only the cheap graph enumeration + Oncosplice alignment repeats
per transcript (~50-100ms each).

This catches transcript-specific effects invisible to single-transcript
analysis: a mutation might be benign on the primary/MANE transcript but
devastating on an alternative isoform expressed in a specific tissue.
"""
from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import Optional

import pandas as pd

import numpy as np

from ..results import SplicingResult
from ..splicing import SpliceSimulator
from ..utils import get_all_transcripts, _apply_mutation_safe

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

logger = logging.getLogger(__name__)


@dataclass
class TranscriptSplicingResult:
    """Splicing analysis result for a single protein-coding transcript."""
    transcript_id: str
    is_primary: bool
    is_protein_coding: bool = True
    splicing_result: SplicingResult = field(default_factory=SplicingResult)
    n_isoforms: int = 0
    max_oncosplice_score: float = 0.0
    missplicing: float = 0.0
    has_impact: bool = False
    timing_seconds: float = 0.0


@dataclass
class NonCodingSpliceCheck:
    """Lightweight splice impact check for a non-protein-coding transcript.

    No isoform enumeration or protein alignment — just checks whether
    the SpliceAI-predicted splice site changes overlap with this
    transcript's annotated donors/acceptors.
    """
    transcript_id: str
    biotype: str                            # e.g. "lncRNA", "processed_transcript"
    has_impact: bool = False
    n_affected_donors: int = 0
    n_affected_acceptors: int = 0
    affected_donor_positions: list[int] = field(default_factory=list)
    affected_acceptor_positions: list[int] = field(default_factory=list)
    max_donor_delta: float = 0.0            # Largest |Δprob| at an annotated donor
    max_acceptor_delta: float = 0.0         # Largest |Δprob| at an annotated acceptor


@dataclass
class MultiTranscriptResult:
    """Aggregate splicing analysis across all transcripts of a gene.

    Attributes:
        primary_result: SplicingResult for the primary/MANE transcript.
        per_transcript: Per-transcript results (protein-coding only).
        non_coding_checks: Lightweight splice checks for non-coding transcripts.
        n_transcripts_analyzed: Total transcripts analyzed (coding + non-coding).
        n_transcripts_affected: How many show any splicing impact.
        n_non_coding_affected: How many non-coding transcripts are affected.
        worst_transcript_id: Transcript with the highest Oncosplice score.
        worst_oncosplice_score: Highest score across protein-coding transcripts.
        transcript_scores: Summary {transcript_id: max_oncosplice_score}.
        timing: Total wall-clock time for multi-transcript analysis.
    """
    primary_result: SplicingResult
    per_transcript: list[TranscriptSplicingResult] = field(default_factory=list)
    non_coding_checks: list[NonCodingSpliceCheck] = field(default_factory=list)
    n_transcripts_analyzed: int = 0
    n_transcripts_affected: int = 0
    n_non_coding_affected: int = 0
    worst_transcript_id: str = ""
    worst_oncosplice_score: float = 0.0
    transcript_scores: dict[str, float] = field(default_factory=dict)
    shared_splice_events: Optional[pd.Series] = None
    timing: float = 0.0


def _check_non_coding_splice_impact(
    transcript,
    splicing_preds: pd.DataFrame,
    feature: str = "event",
    delta_threshold: float = 0.2,
) -> NonCodingSpliceCheck:
    """Check if SpliceAI-predicted changes affect a non-coding transcript's splice sites.

    This is the lightweight path: no isoform enumeration, no protein alignment.
    Just checks whether any annotated donor/acceptor in this transcript has a
    significant probability change in the shared SpliceAI predictions.

    Cost: ~0.1ms per transcript (set intersection + DataFrame lookup).
    """
    tid = getattr(transcript, 'transcript_id', 'unknown')
    biotype = getattr(transcript, 'biotype', 'non_coding') or 'non_coding'
    donors = set(getattr(transcript, 'donors', []) or [])
    acceptors = set(getattr(transcript, 'acceptors', []) or [])

    feature_col = f"{feature}_prob"
    affected_donors = []
    affected_acceptors = []
    max_d_delta = 0.0
    max_a_delta = 0.0

    # Check donors
    if hasattr(splicing_preds, 'donors'):
        donor_df = splicing_preds.donors
        if feature_col in donor_df.columns and 'ref_prob' in donor_df.columns:
            for pos in donors:
                if pos in donor_df.index:
                    row = donor_df.loc[pos]
                    delta = abs(float(row.get(feature_col, 0)) - float(row.get('ref_prob', 0)))
                    if delta >= delta_threshold:
                        affected_donors.append(pos)
                        max_d_delta = max(max_d_delta, delta)

    # Check acceptors
    if hasattr(splicing_preds, 'acceptors'):
        acceptor_df = splicing_preds.acceptors
        if feature_col in acceptor_df.columns and 'ref_prob' in acceptor_df.columns:
            for pos in acceptors:
                if pos in acceptor_df.index:
                    row = acceptor_df.loc[pos]
                    delta = abs(float(row.get(feature_col, 0)) - float(row.get('ref_prob', 0)))
                    if delta >= delta_threshold:
                        affected_acceptors.append(pos)
                        max_a_delta = max(max_a_delta, delta)

    has_impact = len(affected_donors) > 0 or len(affected_acceptors) > 0

    return NonCodingSpliceCheck(
        transcript_id=tid,
        biotype=biotype,
        has_impact=has_impact,
        n_affected_donors=len(affected_donors),
        n_affected_acceptors=len(affected_acceptors),
        affected_donor_positions=sorted(affected_donors),
        affected_acceptor_positions=sorted(affected_acceptors),
        max_donor_delta=round(max_d_delta, 4),
        max_acceptor_delta=round(max_a_delta, 4),
    )


def analyze_all_transcripts(
    gene,
    mutation_event,
    splicing_preds: pd.DataFrame,
    primary_transcript,
    primary_splicing_result: SplicingResult,
    splicing_engine: str = "ensemble",
    central_pos: int = 0,
    max_isoforms: int | None = None,
    min_total_prob: float = 0.5,
    min_delta_magnitude: float = 0.2,
    max_transcripts: int = 20,
) -> MultiTranscriptResult:
    """Apply shared SpliceAI predictions to all protein-coding transcripts.

    The core optimization: SpliceAI runs once, producing a position-indexed
    DataFrame of donor/acceptor probabilities. This function reuses that
    DataFrame for every transcript in the gene, only recomputing the cheap
    parts (graph traversal + protein alignment).

    Cost: ~50-100ms per additional transcript (vs ~5-10s for SpliceAI).

    Args:
        gene: seqmat Gene object.
        mutation_event: MutationalEvent with mutation details.
        splicing_preds: Shared SpliceAI predictions DataFrame
            (position-indexed, from TranscriptLibrary.get_event_columns).
        primary_transcript: The already-analyzed primary transcript.
        primary_splicing_result: The already-computed SplicingResult for primary.
        splicing_engine: Engine name for metadata.
        central_pos: Central mutation position.
        max_isoforms: Max isoforms per transcript.
        min_total_prob: Min combined prob for viable paths.
        min_delta_magnitude: Min |Δprob| for splice site changes.
        max_transcripts: Safety limit on number of transcripts to analyze.

    Returns:
        MultiTranscriptResult with per-transcript and aggregate results.
    """
    t_start = time.perf_counter()

    m = mutation_event
    primary_tid = getattr(primary_transcript, "transcript_id", "")

    # Get ALL transcripts — protein-coding for full analysis, non-coding for splice check
    coding_transcripts, non_coding_transcripts = get_all_transcripts(gene, position=central_pos)

    if len(coding_transcripts) > max_transcripts:
        logger.info("Gene %s has %d coding transcripts; limiting to %d",
                     m.gene, len(coding_transcripts), max_transcripts)
        coding_transcripts = coding_transcripts[:max_transcripts]

    splice_index = _build_splice_site_index(gene)
    per_transcript: list[TranscriptSplicingResult] = []
    non_coding_checks: list[NonCodingSpliceCheck] = []
    transcript_scores: dict[str, float] = {}
    worst_tid = primary_tid
    worst_score = primary_splicing_result.max_oncosplice_score

    # Add the primary transcript result (already computed, no extra work)
    per_transcript.append(TranscriptSplicingResult(
        transcript_id=primary_tid,
        is_primary=True,
        splicing_result=primary_splicing_result,
        n_isoforms=len(primary_splicing_result.isoforms),
        max_oncosplice_score=primary_splicing_result.max_oncosplice_score,
        missplicing=primary_splicing_result.missplicing,
        has_impact=not primary_splicing_result.empty,
    ))
    transcript_scores[primary_tid] = primary_splicing_result.max_oncosplice_score

    # Analyze each additional protein-coding transcript
    for transcript in coding_transcripts:
        tid = getattr(transcript, "transcript_id", "")
        if tid == primary_tid:
            continue  # Already have this one

        t0 = time.perf_counter()
        try:
            # Build reference for this transcript
            ref_t = transcript.generate_mature_mrna().generate_protein()

            # Create mutated version for this transcript
            mut_t = transcript.clone().generate_mature_mrna()
            for pos, ref, alt in m.mutation_args():
                try:
                    _apply_mutation_safe(mut_t.mature_mrna, pos, ref, alt)
                except Exception:
                    pass

            # Clear pre-mRNA so isoform clones use direct-exon assembly
            # (matches the behavior in _mutation.py)
            mut_t._pre_mrna = None

            # Build SpliceSimulator for THIS transcript using the SHARED predictions
            ss = SpliceSimulator(
                splicing_preds,
                mut_t,
                feature="event",
                max_distance=SPLICE_MAX_DISTANCE,
                min_total_prob=min_total_prob,
                min_delta_magnitude=min_delta_magnitude,
                mutations=list(m.mutation_args()),
            )

            # Build metadata and run Oncosplice
            mut_id = f"{m.gene}:{m.chrom}:{central_pos}:{':'.join(str(x) for x in m.mutation_args()[0][1:])}"
            base_report = _splicing_base_report(mut_id, m, ref_t, splicing_engine, central_pos)
            ss_metadata = ss.report(central_pos)
            splicing_result = _oncosplice_dataframe(
                ref_t, ss, ss_metadata, base_report,
                max_isoforms=max_isoforms, splice_index=splice_index,
            )

            elapsed = time.perf_counter() - t0
            onco_score = splicing_result.max_oncosplice_score
            has_impact = not splicing_result.empty

            per_transcript.append(TranscriptSplicingResult(
                transcript_id=tid,
                is_primary=False,
                splicing_result=splicing_result,
                n_isoforms=len(splicing_result.isoforms),
                max_oncosplice_score=onco_score,
                missplicing=splicing_result.missplicing,
                has_impact=has_impact,
                timing_seconds=round(elapsed, 4),
            ))
            transcript_scores[tid] = onco_score

            if onco_score > worst_score:
                worst_score = onco_score
                worst_tid = tid

        except Exception as e:
            logger.debug("Multi-transcript analysis failed for %s/%s: %s", m.gene, tid, e)
            per_transcript.append(TranscriptSplicingResult(
                transcript_id=tid,
                is_primary=False,
                splicing_result=SplicingResult(),
                timing_seconds=round(time.perf_counter() - t0, 4),
            ))
            transcript_scores[tid] = 0.0

    # ── Lightweight non-coding transcript checks ──
    # ~0.1ms each: just check if any annotated splice site has a significant delta
    for nc_transcript in non_coding_transcripts:
        try:
            check = _check_non_coding_splice_impact(
                nc_transcript, splicing_preds,
                feature="event", delta_threshold=min_delta_magnitude,
            )
            non_coding_checks.append(check)
        except Exception as e:
            nc_tid = getattr(nc_transcript, 'transcript_id', 'unknown')
            logger.debug("Non-coding splice check failed for %s/%s: %s", m.gene, nc_tid, e)

    total_time = time.perf_counter() - t_start
    n_coding_affected = sum(1 for r in per_transcript if r.has_impact)
    n_nc_affected = sum(1 for c in non_coding_checks if c.has_impact)

    return MultiTranscriptResult(
        primary_result=primary_splicing_result,
        per_transcript=per_transcript,
        non_coding_checks=non_coding_checks,
        n_transcripts_analyzed=len(per_transcript) + len(non_coding_checks),
        n_transcripts_affected=n_coding_affected + n_nc_affected,
        n_non_coding_affected=n_nc_affected,
        worst_transcript_id=worst_tid,
        worst_oncosplice_score=worst_score,
        transcript_scores=transcript_scores,
        timing=round(total_time, 4),
    )
