# geney/outcome_tree.py — Sequence outcome tree: probability-weighted isoform × TIS paths.
"""
Builds a probability tree of all possible sequence outcomes from a mutation,
starting from the pre-mRNA and generating all splicing isoforms, then
predicting TIS usage per isoform to produce protein outcomes.

Each level's weights sum to 1. Leaf probability = product down the path.

    pre_mrna (weight=1.0)
    ├── mRNA_isoform_1 (weight=0.65)
    │   ├── protein_canonical (weight=0.90)
    │   └── protein_upstream  (weight=0.10)
    └── mRNA_reference (weight=0.35)
        └── protein_canonical (weight=1.0)

Usage::

    from geney import build_outcome_tree

    tree = build_outcome_tree("KRAS:12:25227343:G:T")
    print(tree.n_protein_outcomes, "protein outcomes")
    for m in tree.mrna_outcomes:
        print(f"  mRNA {m.isoform_id} weight={m.weight}")
        for p in m.proteins:
            print(f"    protein {p.tis_location} len={len(p.sequence)} abs={p.absolute_weight}")
"""
from __future__ import annotations

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

from Bio.Seq import Seq

logger = logging.getLogger(__name__)


# ============================================================================
# Data structures
# ============================================================================

@dataclass
class ProteinOutcome:
    """A leaf node: one possible protein product."""
    sequence: str
    tis_position: int              # 0-based mRNA position of start codon
    tis_location: str              # "canonical" / "upstream" / "downstream"
    in_frame: bool
    weight: float                  # TIS usage fraction (sums to 1 within parent mRNA)
    absolute_weight: float         # product of all weights from root to this leaf

    def to_dict(self) -> dict:
        return {
            "sequence": self.sequence,
            "length": len(self.sequence),
            "tis_position": self.tis_position,
            "tis_location": self.tis_location,
            "in_frame": self.in_frame,
            "weight": round(self.weight, 4),
            "absolute_weight": round(self.absolute_weight, 6),
        }


@dataclass
class MRNAOutcome:
    """A branch node: one possible mature mRNA isoform."""
    sequence: str
    isoform_id: str
    splicing_summary: str          # "reference", "exon_skipping", etc.
    weight: float                  # isoform prevalence (sums to 1 across siblings)
    proteins: list[ProteinOutcome] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "sequence_length": len(self.sequence),
            "isoform_id": self.isoform_id,
            "splicing_summary": self.splicing_summary,
            "weight": round(self.weight, 4),
            "proteins": [p.to_dict() for p in self.proteins],
        }


@dataclass
class OutcomeTree:
    """Root: the full probability tree of sequence outcomes."""
    mut_id: str
    gene: str
    pre_mrna: str                  # scoped pre-mRNA with mutation applied
    reference_mrna: str
    reference_protein: str
    mrna_outcomes: list[MRNAOutcome] = field(default_factory=list)

    @property
    def n_mrna_outcomes(self) -> int:
        return len(self.mrna_outcomes)

    @property
    def n_protein_outcomes(self) -> int:
        return sum(len(m.proteins) for m in self.mrna_outcomes)

    @property
    def top_outcome(self) -> Optional[ProteinOutcome]:
        best = None
        for m in self.mrna_outcomes:
            for p in m.proteins:
                if best is None or p.absolute_weight > best.absolute_weight:
                    best = p
        return best

    def to_dict(self) -> dict:
        top = self.top_outcome
        return {
            "mut_id": self.mut_id,
            "gene": self.gene,
            "pre_mrna_length": len(self.pre_mrna),
            "reference_mrna_length": len(self.reference_mrna),
            "reference_protein_length": len(self.reference_protein),
            "n_mrna_outcomes": self.n_mrna_outcomes,
            "n_protein_outcomes": self.n_protein_outcomes,
            "top_outcome": top.to_dict() if top else None,
            "mrna_outcomes": [m.to_dict() for m in self.mrna_outcomes],
        }


# ============================================================================
# Helpers
# ============================================================================

def _translate_from_position(mrna: str, start_pos: int) -> str:
    """Translate mRNA from a given start position to the first stop codon."""
    coding = mrna[start_pos:]
    if len(coding) < 3:
        return ""
    try:
        return str(Seq(coding).translate(to_stop=True))
    except Exception:
        return ""


def _run_tis_for_mrna(ref_mrna: str, var_mrna: str, probability_threshold: float):
    """Run TIS on a single mRNA pair. Returns list of (position, location, in_frame, usage)."""
    from .tis.metrics import compute_tis_shift_metrics

    result = compute_tis_shift_metrics(
        ref_mrna=ref_mrna,
        var_mrna=var_mrna,
        canonical_start=None,
        probability_threshold=probability_threshold,
    )

    sites = []
    if result.analysis_df is not None and not result.analysis_df.empty:
        for _, row in result.analysis_df.iterrows():
            usage = row.get("var_usage", 0.0)
            if usage > 0:
                sites.append({
                    "position": int(row["position"]),
                    "location": row.get("location", "unknown"),
                    "in_frame": bool(row.get("in_frame", False)),
                    "usage": float(usage),
                })

    return sites


# ============================================================================
# Builder
# ============================================================================

def build_outcome_tree(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
    splicing_engine: str = "ensemble",
    gene_load_timeout: float | None = None,
    max_isoforms: int | None = None,
    min_total_prob: float = 0.5,
    min_delta_magnitude: float = 0.2,
    tis_probability_threshold: float = 0.1,
    min_isoform_prevalence: float = 0.01,
    min_tis_usage: float = 0.01,
    max_proteins_per_mrna: int = 5,
) -> OutcomeTree:
    """Build a probability-weighted sequence outcome tree from scratch.

    Loads the gene, runs splicing prediction, generates all viable isoforms,
    runs TIS per isoform, and translates proteins per start site.

    Args:
        mut_id: Mutation ID (GENE:CHR:POS:REF:ALT, or pipe-separated compound).
        transcript_id: Specific transcript, or auto-select.
        organism: Genome build.
        splicing_engine: "spliceai" or "pangolin".
        gene_load_timeout: Max seconds for gene loading.
        max_isoforms: Cap on splice isoforms from SpliceSimulator.
        min_total_prob: Min combined splice-site probability for a path.
        min_delta_magnitude: Min |delta_prob| to count a splice gain/loss.
        tis_probability_threshold: Min ContextScore for a TIS site.
        min_isoform_prevalence: Prune isoforms below this prevalence.
        min_tis_usage: Prune TIS sites below this usage fraction.
        max_proteins_per_mrna: Cap TIS sites per isoform.

    Returns:
        OutcomeTree with mRNA isoforms and protein leaves, weights summing to 1.
    """
    from .variants import MutationalEvent
    from .utils import select_transcript, _apply_mutation_safe
    from .splicing import SpliceSimulator, TranscriptLibrary
    from .pipelines._common import (
        SPLICING_CONTEXT_BP,
        SPLICE_MAX_DISTANCE,
        _load_gene_with_timeout,
    )

    # -- 1. Parse mutation and load gene --------------------------------------
    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:
        raise ValueError(f"Gene not found: {m.gene}")

    central_pos = m.central_position
    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        raise ValueError(f"No transcript covers position {central_pos}")

    # -- 2. Generate reference sequences --------------------------------------
    reference_transcript = selected.generate_mature_mrna().generate_protein()
    ref_mrna = reference_transcript.mature_mrna.seq
    ref_protein = reference_transcript.protein or ""

    # -- 3. Build scoped pre-mRNA and run splicing ----------------------------
    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()),
    )

    # Get pre-mRNA with mutation applied (for the tree root)
    pre_mrna_transcript = selected.clone()
    pre_mrna_transcript.generate_pre_mrna(
        upstream=SPLICING_CONTEXT_BP,
        downstream=SPLICING_CONTEXT_BP,
    )
    for pos, ref, alt in m.mutation_args():
        try:
            _apply_mutation_safe(pre_mrna_transcript.pre_mrna, pos, ref, alt)
        except Exception:
            pass
    pre_mrna_seq = pre_mrna_transcript.pre_mrna.seq if pre_mrna_transcript.pre_mrna else ""

    # -- 4. Collect all viable isoforms from SpliceSimulator ------------------
    isoform_data = []
    for variant_transcript, isoform_metadata in ss.get_viable_transcripts(
        metadata=True, max_isoforms=max_isoforms
    ):
        mrna_obj = variant_transcript.mature_mrna
        if mrna_obj is None or not hasattr(mrna_obj, "seq"):
            continue

        mrna_seq = mrna_obj.seq
        protein_seq = variant_transcript.protein or ""
        prevalence = float(isoform_metadata.get("isoform_prevalence", 0.0))
        isoform_id = str(isoform_metadata.get("isoform_id", ""))
        summary = str(isoform_metadata.get("summary", ""))

        if mrna_seq and prevalence > 0:
            isoform_data.append({
                "mrna": mrna_seq,
                "protein": protein_seq,
                "prevalence": prevalence,
                "isoform_id": isoform_id,
                "summary": summary,
            })

    # Add reference isoform if not already represented
    total_prev = sum(d["prevalence"] for d in isoform_data)
    ref_already = any(d["mrna"] == ref_mrna for d in isoform_data)
    if not ref_already and ref_mrna:
        ref_prev = max(0.0, 1.0 - total_prev)
        if ref_prev > 0:
            isoform_data.append({
                "mrna": ref_mrna,
                "protein": ref_protein,
                "prevalence": ref_prev,
                "isoform_id": "reference",
                "summary": "reference",
            })

    # Normalize prevalences to sum to 1
    total_prev = sum(d["prevalence"] for d in isoform_data)
    if total_prev > 0:
        for d in isoform_data:
            d["prevalence"] /= total_prev

    # Filter and re-normalize
    isoform_data = [d for d in isoform_data if d["prevalence"] >= min_isoform_prevalence]
    total_prev = sum(d["prevalence"] for d in isoform_data)
    if total_prev > 0:
        for d in isoform_data:
            d["prevalence"] /= total_prev

    isoform_data.sort(key=lambda d: -d["prevalence"])

    # -- 5. Run TIS per isoform and build protein leaves ----------------------
    mrna_outcomes = []

    for iso in isoform_data:
        iso_mrna = iso["mrna"]
        iso_weight = iso["prevalence"]
        iso_protein = iso["protein"]

        proteins = []
        try:
            tis_sites = _run_tis_for_mrna(ref_mrna, iso_mrna, tis_probability_threshold)

            # Filter and cap
            tis_sites = [s for s in tis_sites if s["usage"] >= min_tis_usage]
            tis_sites.sort(key=lambda s: -s["usage"])
            tis_sites = tis_sites[:max_proteins_per_mrna]

            # Re-normalize usages to sum to 1
            total_usage = sum(s["usage"] for s in tis_sites)
            if total_usage > 0:
                for s in tis_sites:
                    s["usage"] /= total_usage

            for site in tis_sites:
                pos = site["position"]
                location = site["location"]
                in_frame = site["in_frame"]
                usage = site["usage"]

                # Translate protein from this start position
                if location == "canonical" and iso_protein:
                    protein_seq = iso_protein
                else:
                    protein_seq = _translate_from_position(iso_mrna, pos)

                proteins.append(ProteinOutcome(
                    sequence=protein_seq,
                    tis_position=pos,
                    tis_location=location,
                    in_frame=in_frame,
                    weight=round(usage, 6),
                    absolute_weight=round(iso_weight * usage, 6),
                ))

        except Exception as e:
            logger.debug("TIS failed for isoform %s: %s", iso["isoform_id"], e)

        # Fallback: if TIS produced nothing, use the splice-derived protein
        if not proteins:
            proteins.append(ProteinOutcome(
                sequence=iso_protein or _translate_from_position(iso_mrna, 0),
                tis_position=0,
                tis_location="canonical",
                in_frame=True,
                weight=1.0,
                absolute_weight=round(iso_weight, 6),
            ))

        mrna_outcomes.append(MRNAOutcome(
            sequence=iso_mrna,
            isoform_id=iso["isoform_id"],
            splicing_summary=iso["summary"],
            weight=round(iso_weight, 4),
            proteins=proteins,
        ))

    return OutcomeTree(
        mut_id=mut_id,
        gene=m.gene,
        pre_mrna=pre_mrna_seq,
        reference_mrna=ref_mrna,
        reference_protein=ref_protein,
        mrna_outcomes=mrna_outcomes,
    )
