# geney/visualization/_splicing.py — Splicing card data for frontend rendering.
#
# Produces a JSON-serializable dict with:
#   - Reference/variant sequences (protein, mRNA)
#   - Exon map (reference exon boundaries, disruption flags)
#   - Splice site probability changes (significant sites only)
#   - Per-isoform details (aligned proteins, scores, prevalence)
from __future__ import annotations

from typing import Optional

import numpy as np
import pandas as pd

from ..results import SplicingResult


def _safe_float(v) -> Optional[float]:
    if v is None or (isinstance(v, float) and np.isnan(v)):
        return None
    return round(float(v), 4)


def _extract_splice_sites(df: Optional[pd.DataFrame], site_type: str,
                          min_delta: float = 0.01) -> list[dict]:
    """Extract splice sites with significant probability changes."""
    if df is None or df.empty:
        return []

    sites = []
    if isinstance(df.columns, pd.MultiIndex):
        prefix = "donors" if site_type == "donor" else "acceptors"
        ref_col, var_col, ann_col = (prefix, "ref_prob"), (prefix, "event_prob"), (prefix, "annotated")
    else:
        ref_col, var_col, ann_col = "ref_prob", "event_prob", "annotated"

    has_ref = ref_col in df.columns
    has_var = var_col in df.columns
    has_ann = ann_col in df.columns
    if not (has_ref and has_var):
        return []

    for pos in df.index:
        ref_p = float(df.loc[pos, ref_col])
        var_p = float(df.loc[pos, var_col])
        delta = var_p - ref_p
        if abs(delta) < min_delta:
            continue

        site = {
            "position": int(pos),
            "type": site_type,
            "ref_prob": round(ref_p, 4),
            "var_prob": round(var_p, 4),
            "delta": round(delta, 4),
        }
        if has_ann:
            site["annotated"] = bool(df.loc[pos, ann_col])
        sites.append(site)

    sites.sort(key=lambda s: abs(s["delta"]), reverse=True)
    return sites


def _extract_exons(transcript) -> list[dict]:
    """Extract exon boundaries from transcript object."""
    exons_raw = getattr(transcript, "exons", None)
    if exons_raw is None:
        return []
    try:
        if len(exons_raw) == 0:
            return []
    except TypeError:
        return []

    exons = []
    for i, ex in enumerate(exons_raw):
        try:
            start, end = int(ex[0]), int(ex[1])
        except (TypeError, IndexError):
            continue
        exons.append({
            "number": i + 1,
            "start": min(start, end),
            "end": max(start, end),
        })

    exons.sort(key=lambda e: e["start"])
    return exons


def _flag_disrupted_exons(exons: list[dict], splice_sites: list[dict],
                          threshold: float = 0.1) -> None:
    """Mark exons whose flanking splice sites are disrupted (in-place)."""
    disrupted_positions = {
        s["position"] for s in splice_sites if abs(s["delta"]) >= threshold
    }
    for exon in exons:
        margin = 5
        exon["disrupted"] = any(
            (exon["start"] - margin) <= pos <= (exon["end"] + margin)
            for pos in disrupted_positions
        )


def _isoform_rows(isoforms_df: pd.DataFrame) -> list[dict]:
    """Extract per-isoform data as a list of dicts."""
    if isoforms_df.empty:
        return []

    rows = []
    for _, row in isoforms_df.iterrows():
        iso: dict = {}

        # Splicing event info
        for col in ("summary", "isoform_id", "known_transcript_id"):
            if col in row.index:
                val = row[col]
                iso[col] = str(val) if val is not None and str(val) != "nan" else None

        # Scores
        for col in ("oncosplice_score", "percentile", "isoform_prevalence"):
            if col in row.index:
                iso[col] = _safe_float(row[col])

        # Protein data
        for col in ("variant_protein", "aligned_reference_protein", "aligned_variant_protein"):
            if col in row.index:
                val = row[col]
                iso[col] = str(val) if isinstance(val, str) and val else None

        if "variant_protein" in row.index:
            vp = row["variant_protein"]
            iso["variant_protein_length"] = len(vp) if isinstance(vp, str) else None

        # Variant mRNA
        if "variant_mrna" in row.index:
            val = row["variant_mrna"]
            iso["variant_mrna"] = str(val) if isinstance(val, str) and val else None

        # Structural changes
        for col in ("number_of_deletions", "number_of_insertions", "modified_positions_count",
                     "variant_length"):
            if col in row.index:
                val = row[col]
                iso[col] = int(val) if val is not None and not (isinstance(val, float) and np.isnan(val)) else None

        # Splice event details
        for col in ("pes", "pir", "es", "ne", "ir"):
            if col in row.index:
                val = row[col]
                iso[col] = str(val) if val and str(val).strip() else None

        rows.append(iso)

    # Sort by oncosplice_score descending
    rows.sort(key=lambda r: r.get("oncosplice_score") or 0, reverse=True)
    return rows


def splicing_card(splicing: SplicingResult, transcript=None) -> dict:
    """Build splicing card data for frontend rendering.

    Args:
        splicing: SplicingResult from pipeline.
        transcript: Optional transcript object (for exon coordinates).

    Returns:
        JSON-serializable dict with all data the frontend needs.
    """
    card: dict = {
        "type": "splicing",
        "score": round(min(10.0, splicing.max_oncosplice_score
                           if splicing.max_oncosplice_score > 0
                           else abs(splicing.missplicing) * 10), 1),
        "status": "completed" if not splicing.empty else "no_data",
        "missplicing": round(splicing.missplicing, 4),
        "oncosplice_score": round(splicing.max_oncosplice_score, 4),
        "mutation_position": splicing.central_position or None,
        "gene": splicing.gene or None,
        "transcript_id": splicing.transcript_id,

        # Reference sequences
        "reference_protein": splicing.reference_protein or None,
        "reference_protein_length": len(splicing.reference_protein) if splicing.reference_protein else 0,
        "reference_mrna": splicing.reference_mrna or None,

        # Per-residue conservation (list of floats, same length as reference_protein)
        "conservation_vector": (
            splicing.conservation_vector.tolist()
            if splicing.conservation_vector is not None
            else None
        ),

        # Isoform count
        "num_isoforms": len(splicing.isoforms),
    }

    # Splice site changes
    donors = _extract_splice_sites(splicing.donor_df, "donor")
    acceptors = _extract_splice_sites(splicing.acceptor_df, "acceptor")
    card["splice_sites"] = donors + acceptors
    card["splice_sites"].sort(key=lambda s: abs(s["delta"]), reverse=True)

    # Exon map
    exons = _extract_exons(transcript) if transcript is not None else []
    if exons:
        _flag_disrupted_exons(exons, card["splice_sites"])
    card["exons"] = exons

    # Per-isoform details (aligned proteins, scores, variant sequences)
    card["isoforms"] = _isoform_rows(splicing.isoforms)

    return card
