# geney/visualization/_tis.py — TIS card data for frontend rendering.
#
# Produces a JSON-serializable dict with:
#   - mRNA sequences (ref/var windows around 5' end)
#   - ATG sites with ref/var probabilities (thresholded)
#   - Canonical start status and metrics
from __future__ import annotations

from ..results import TISResult

# Default thresholds for site inclusion.
# A site is included if: canonical, OR prob >= MIN_PROB, OR |delta| >= MIN_DELTA.
_DEFAULT_MIN_PROB = 0.1
_DEFAULT_MIN_DELTA = 0.01


def tis_card(tis: TISResult | None,
             min_prob: float = _DEFAULT_MIN_PROB,
             min_delta: float = _DEFAULT_MIN_DELTA) -> dict:
    """Build TIS card data for frontend rendering.

    Args:
        tis: TISResult from pipeline.
        min_prob: Minimum probability in either ref or var to include a site.
        min_delta: Minimum |probability change| to include a site (even if
            below min_prob). The canonical site is always included.

    Returns:
        JSON-serializable dict with mRNA track data and ATG site probabilities.
    """
    if tis is None:
        return {"type": "tis", "score": 0.0, "status": "not_run"}

    if tis.status != "completed":
        return {
            "type": "tis",
            "score": 0.0,
            "status": tis.status,
            "reason": tis.reason,
        }

    # Filter ATG sites by threshold
    sites = []
    for site in tis.atg_sites:
        ref_p = float(site.get("ref_probability", 0))
        var_p = float(site.get("var_probability", 0))
        delta = var_p - ref_p
        is_canonical = site.get("location") == "canonical"

        # Always include canonical; otherwise require meaningful probability
        # AND a nonzero change, OR a large change regardless of probability.
        if not is_canonical:
            has_prob = max(ref_p, var_p) >= min_prob
            has_delta = abs(delta) >= min_delta
            if not (has_prob and has_delta) and abs(delta) < min_delta * 5:
                continue

        sites.append({
            "position": int(site["position"]),
            "location": site.get("location", "unknown"),
            "in_frame": bool(site.get("in_frame", False)),
            "ref_prob": round(ref_p, 4),
            "var_prob": round(var_p, 4),
            "delta": round(delta, 4),
            "ref_usage": round(float(site.get("ref_usage", 0)), 4),
            "var_usage": round(float(site.get("var_usage", 0)), 4),
        })

    # Sort: canonical first, then by |delta| descending
    sites.sort(key=lambda s: (s["location"] != "canonical", -abs(s["delta"])))

    return {
        "type": "tis",
        "score": round(min(10.0, tis.score * 10), 1),
        "status": "completed",

        # mRNA info
        "mrna_length": tis.ref_mrna_length,
        "var_mrna_length": tis.var_mrna_length,
        "ref_mrna_window": tis.ref_mrna_window or None,
        "var_mrna_window": tis.var_mrna_window or None,

        # Canonical metrics
        "canonical_intact": tis.initiation_intact,
        "canonical_prob_ref": round(tis.canonical_prob_ref, 4),
        "canonical_prob_var": round(tis.canonical_prob_var, 4),
        "canonical_usage_ref": round(tis.canonical_usage_ref, 4),
        "canonical_usage_var": round(tis.canonical_usage_var, 4),

        # Mutation context
        "mutation_positions": tis.mutation_mrna_positions,

        # Alternatives
        "num_alternatives_gained": tis.num_alternative_tis,
        "top_alternative_position": tis.top_alternative_position,
        "top_alternative_prob": round(tis.top_alternative_prob, 4),
        "top_alternative_location": tis.top_alternative_location,
        "top_alternative_in_frame": tis.top_alternative_in_frame,

        # Thresholded ATG sites
        "sites": sites,
    }
