"""Shape compute_vadug's raw trace into display + audit data.

trace entries look like {word, role, coeff, v, a, d, u, g} (per-word deltas);
trace['structures'] is a list of StructureMatch (has .name or str()).
"""
from __future__ import annotations

_INERT = {None, "", "GAS", "CONNECTOR"}


def _struct_names(structures) -> list[str]:
    out = []
    for s in structures or []:
        out.append(getattr(s, "pattern", None) or getattr(s, "name", None) or str(s))
    return out


def shape_trace(trace: dict) -> dict:
    """Shape raw trace into display and audit data.

    Args:
        trace: raw trace dict from compute_vadug with 'trace' list and 'structures'

    Returns:
        Shaped trace dict with:
        - words: list of per-word entries with role and dimension deltas
        - structures: list of matched pattern names
        - contributors: top 5 words by absolute impact
        - unknown_tokens: alpha words with inert roles and zero movement
        - suspected_gap: bool, True if very few words fired
    """
    entries = trace.get("trace", []) or []
    words = []
    unknown = []

    for e in entries:
        role = e.get("role")
        dv = int(e.get("v", 0))
        da = int(e.get("a", 0))
        dd = int(e.get("d", 0))
        du = int(e.get("u", 0))
        dg = int(e.get("g", 0))

        words.append({
            "word": e.get("word", ""),
            "role": role or "GAS",
            "dv": dv,
            "da": da,
            "dd": dd,
            "du": du,
            "dg": dg,
        })

        # Check if word moved
        moved = any(abs(x) > 0 for x in [dv, da, dd, du, dg])

        # Flag as unknown: alpha word, inert role, no movement
        if role in _INERT and not moved and e.get("word", "").isalpha():
            unknown.append(e.get("word", ""))

    # Top 5 contributors by total absolute impact
    contributors = sorted(
        words,
        key=lambda w: abs(w["dv"]) + abs(w["da"]) + abs(w["dd"]) +
                     abs(w["du"]) + abs(w["dg"]),
        reverse=True
    )[:5]

    # Flag suspected gaps: few words actually fired
    fired = [w for w in words if w["role"] not in _INERT]
    suspected_gap = len(words) >= 3 and len(fired) <= max(1, len(words) // 5)

    return {
        "words": words,
        "structures": _struct_names(trace.get("structures")),
        "contributors": contributors,
        "unknown_tokens": unknown,
        "suspected_gap": suspected_gap,
    }