# geney/visualization/_folding.py — Folding card data for frontend rendering.
#
# Produces a JSON-serializable dict with:
#   - Ref/var sequences and dot-bracket structures
#   - Base-pair arcs (parsed from dot-bracket) for arc diagram rendering
#   - MFE values and delta
#   - Mutation position within the window
#
# The frontend draws: arc diagram (arcs above a sequence strip, ref vs var).
from __future__ import annotations

from ..results import FoldingResult, FoldingWindow


def _parse_dot_bracket(structure: str) -> list[dict]:
    """Parse dot-bracket notation into base pair arcs.

    Returns list of {"i": open_pos, "j": close_pos} dicts, sorted by i.
    O(n) time, O(n) space.
    """
    stack: list[int] = []
    pairs: list[dict] = []
    for i, c in enumerate(structure):
        if c == '(':
            stack.append(i)
        elif c == ')':
            if stack:
                j = stack.pop()
                pairs.append({"i": j, "j": i})
    pairs.sort(key=lambda p: p["i"])
    return pairs


def _classify_arcs(ref_arcs: list[dict], var_arcs: list[dict]) -> tuple[list[dict], list[dict]]:
    """Annotate arcs with change status: 'kept', 'lost', 'gained'."""
    ref_set = {(a["i"], a["j"]) for a in ref_arcs}
    var_set = {(a["i"], a["j"]) for a in var_arcs}

    for arc in ref_arcs:
        key = (arc["i"], arc["j"])
        arc["status"] = "kept" if key in var_set else "lost"

    for arc in var_arcs:
        key = (arc["i"], arc["j"])
        arc["status"] = "kept" if key in ref_set else "gained"

    return ref_arcs, var_arcs


def _window_data(window: FoldingWindow) -> dict:
    """Extract arc diagram data from a FoldingWindow."""
    ref_arcs = _parse_dot_bracket(window.ref_structure)
    var_arcs = _parse_dot_bracket(window.var_structure)
    _classify_arcs(ref_arcs, var_arcs)

    return {
        "ref_sequence": window.ref_sequence,
        "var_sequence": window.var_sequence,
        "ref_structure": window.ref_structure,
        "var_structure": window.var_structure,
        "ref_arcs": ref_arcs,
        "var_arcs": var_arcs,
        "ref_mfe": round(window.ref_mfe, 2),
        "var_mfe": round(window.var_mfe, 2),
        "delta_mfe": round(window.delta_mfe, 2),
        "mutation_offset": window.mutation_offset,
        "sequence_length": len(window.ref_sequence),
    }


def folding_card(folding: FoldingResult | None) -> dict:
    """Build folding card data for frontend rendering.

    Returns:
        JSON-serializable dict with arc diagram data and MFE comparison.
    """
    if folding is None:
        return {"type": "folding", "score": 0.0, "status": "not_run"}

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

    card: dict = {
        "type": "folding",
        "score": round(min(10.0, folding.score), 1),
        "status": "completed",
        "delta_mfe": round(folding.delta_mfe, 2) if folding.delta_mfe is not None else None,
        "window_size": folding.window_size,
    }

    # Arc diagram data from central window
    if folding.central_window is not None:
        card["window"] = _window_data(folding.central_window)
    else:
        card["window"] = None

    # Per-mutation summary (positions and status, no heavy data)
    mutations = []
    for m in folding.mutations:
        mutations.append({
            "genomic_position": m.get("genomic_position"),
            "mrna_position": m.get("mrna_position"),
            "status": m.get("status"),
            "delta_mfe": round(m["delta_mfe"], 2) if m.get("delta_mfe") is not None else None,
        })
    card["mutations"] = mutations

    return card
