"""
Layer 1 (v5) - Pose -> structured defensive features, with caching.

Identity is now by MOTION CONTINUITY, not color: we keep two tracks (the two
fighters) and assign each frame's detections to whichever track was nearest a
moment ago. A fighter can't teleport, so this survives them swapping sides and
needs no distinguishing colors. Color (when the two fighters read clearly
different) is used only to *label* the tracks "Red"/"Blue"; otherwise they're
"Fighter A"/"Fighter B".

Public API:
  extract_features(video_path)      -> features dict (runs YOLO live)
  cache_poses(video_path, npz_path) -> saves raw per-frame keypoints (slow step)
  features_from_cache(npz_path)     -> features dict (instant; re-derives)

Every threshold is a STARTING GUESS - calibrate with eval.py.
"""

import os
import warnings
import numpy as np
import cv2
from ultralytics import YOLO

warnings.filterwarnings("ignore")

# ---- Tunables --------------------------------------------------------------
MODEL            = "yolo11n-pose.pt"
VID_STRIDE       = 1      # 1 = every frame (best quality); raise to speed up long clips on CPU
KP_CONF_MIN      = 0.30
MIN_SH_PX        = 20.0
TRACK_GATE_PX    = 160.0  # a track only grabs detections within this many px of its last position
SIZE_RATIO_MIN   = 0.55   # candidate size vs track size band — rejects the much-smaller
SIZE_RATIO_MAX   = 1.9    #   bystander a track would otherwise drift onto
# NOTE: a torso-color gate was tried here (hard and reacquisition-only) and
# REMOVED: gym lighting makes per-frame color too noisy — it blocked the real
# fighter more than bystanders (coverage 19.7% -> ~10%). Identity through
# occlusion is SAM2's job, not a color heuristic's.
MIN_COLOR_FRAMES = 10     # frames of torso color needed before naming a track by color
GUARD_UP         = 0.80
GUARD_DOWN       = 1.60
IN_RANGE         = 2.50
TAG_DIST         = 0.60
TAG_REFRACTORY_S = 0.30
PUNCH_SPEED_SW_S = 3.5    # wrist speed (shoulder-widths/sec) above which a hand is "throwing"
EXCHANGE_PAD_S   = 0.40   # merge punches this close together into one exchange
EXCHANGE_MIN_S   = 0.15   # an exchange must last at least this long
FLURRY_SPEED_SW_S = 5.0   # combined both-fighter arm activity for a "flurry" window
FLURRY_MIN_S      = 0.4   # a flurry must sustain at least this long
FLURRY_HEAVY_SW_S = 9.0   # peak combined activity above this = "heavy" flurry
DROP_MIN_S       = 0.50   # hysteresis: seconds below guard before an episode opens
CLEAR_MIN_S      = 0.30   # hysteresis: seconds back up before an episode closes
PRESSURE_MARGIN  = 5.0    # pct-point gap between advancing/retreating to call a stance
NET_RATE_SW10S   = 0.7    # net drift toward/away from opponent (shoulder-widths per 10s)
                          #   that alone decides pressuring/backing up — a coach's eye
                          #   reads GROUND GIVEN, not step counts
PAN_DOMINANCE    = 0.50   # common (camera) motion > this x relative motion = pan-dominated
BUCKET_S         = 30
MIN_PAIR_FRAMES  = 10
# ---------------------------------------------------------------------------

NOSE, L_SH, R_SH, L_WR, R_WR, L_HIP, R_HIP = 0, 5, 6, 9, 10, 11, 12
L_KNEE, R_KNEE = 13, 14
L_ANK, R_ANK = 15, 16


# ============================ pose + tracking (slow) ========================

def _torso_bgr(frame, kp):
    """Median BGR over the torso box, or None. Median (not mean) so ring ropes /
    background bleeding into the box don't tint the read."""
    pts = kp[[L_SH, R_SH, L_HIP, R_HIP]]
    pts = pts[~np.isnan(pts).any(axis=1)]
    if len(pts) < 2:
        return None
    h, w = frame.shape[:2]
    x0, y0 = pts.min(0)
    x1, y1 = pts.max(0)
    x0, x1 = int(max(0, x0)), int(min(w, x1))
    y0, y1 = int(max(0, y0)), int(min(h, y1))
    if x1 <= x0 or y1 <= y0:
        return None
    roi = frame[y0:y1, x0:x1].reshape(-1, 3)
    return np.median(roi, axis=0)


def _color_name(bgr):
    """Classify a BGR color into a small palette name, or None if ambiguous.
    Replaces the old (R-B) 'redness' rule, which scored YELLOW gear the same as
    RED (yellow = high R, low B) and then force-named the other fighter 'Blue'
    even when nobody wore blue."""
    if bgr is None or np.isnan(np.asarray(bgr, dtype=float)).any():
        return None
    px = np.uint8([[np.asarray(bgr, dtype=np.uint8)]])
    h, s, v = cv2.cvtColor(px, cv2.COLOR_BGR2HSV)[0, 0]
    if v < 50:
        return "black"
    if s < 45:
        return "white" if v > 175 else "gray"
    if h < 8 or h >= 170:
        return "red"
    if h < 22:
        return "orange"
    if h < 38:
        return "yellow"
    if h < 85:
        return "green"
    if h < 128:
        return "blue"
    return "purple"


def _candidates(r):
    """All people in a frame above the size floor: list of (keypoints, hip_center, shoulder_px)."""
    cands = []
    if r.keypoints is None or r.boxes is None or len(r.boxes) == 0:
        return cands
    xy = r.keypoints.xy.cpu().numpy()
    cf = (r.keypoints.conf.cpu().numpy()
          if r.keypoints.conf is not None else np.ones(xy.shape[:2]))
    xy[cf < KP_CONF_MIN] = np.nan
    shs = np.linalg.norm(xy[:, L_SH] - xy[:, R_SH], axis=1)
    for i in range(len(xy)):
        if not (shs[i] >= MIN_SH_PX):
            continue
        kp = xy[i]
        hc = np.nanmean(kp[[L_HIP, R_HIP]], axis=0)
        if np.isnan(hc).any():
            hc = kp[NOSE]
        if np.isnan(hc).any():
            continue
        cands.append((kp, hc, float(shs[i])))
    return cands


def _d(p, q):
    return float(np.linalg.norm(p - q))


def _associate(cands, lastA, lastB):
    """Gated nearest-neighbour tracking. Each track state is (position, size,
    torso color); a track grabs its closest detection within TRACK_GATE_PX that
    is also consistent in size (SIZE_RATIO band) and torso color (COLOR_GATE).
    Position gate: far-away bystanders can't steal a track. Size gate: nearby
    much-smaller ones can't either (a ringside spectator is much smaller on
    screen than the fighter whose track he'd inherit).
    Returns (kpA, kpB, newLastA, newLastB)."""
    if not cands:
        return None, None, lastA, lastB
    if lastA is None or lastB is None:                 # seed with the two largest
        if len(cands) >= 2:
            top = sorted(range(len(cands)), key=lambda i: -cands[i][2])[:2]
            top.sort(key=lambda i: cands[i][1][0])     # left -> A, right -> B
            iA, iB = top
            return (cands[iA][0], cands[iB][0],
                    (cands[iA][1], cands[iA][2]), (cands[iB][1], cands[iB][2]))
        return None, None, lastA, lastB

    def nearest(last, exclude=-1):
        pos, sz = last
        best, bd = -1, TRACK_GATE_PX
        for i, c in enumerate(cands):
            if i == exclude:
                continue
            if not (SIZE_RATIO_MIN * sz < c[2] < SIZE_RATIO_MAX * sz):
                continue
            d = _d(c[1], pos)
            if d < bd:
                bd, best = d, i
        return best

    iA, iB = nearest(lastA), nearest(lastB)
    if iA != -1 and iA == iB:                          # both want the same person
        if _d(cands[iA][1], lastA[0]) <= _d(cands[iA][1], lastB[0]):
            iB = nearest(lastB, exclude=iA)
        else:
            iA = nearest(lastA, exclude=iB)

    def updated(last, i):
        if i == -1:
            return last
        # smooth the size so a fighter closing on the camera stays in-band
        return (cands[i][1], 0.7 * last[1] + 0.3 * cands[i][2])

    kpA = cands[iA][0] if iA != -1 else None
    kpB = cands[iB][0] if iB != -1 else None
    return kpA, kpB, updated(lastA, iA), updated(lastB, iB)


def _label_tracks(colsA, colsB):
    """Name tracks by their ACTUAL dominant torso color ('Yellow fighter' vs
    'Red fighter'), but only when the two colors clearly differ; otherwise the
    neutral 'Fighter A/B' (A = starts on the left)."""
    def track_color(cols):
        cols = [c for c in cols if c is not None]
        if len(cols) < MIN_COLOR_FRAMES:
            return None
        return _color_name(np.median(np.stack(cols), axis=0))
    ca, cb = track_color(colsA), track_color(colsB)
    if ca and cb and ca != cb and "gray" not in (ca, cb):
        return f"{ca.capitalize()} fighter", f"{cb.capitalize()} fighter"
    return ("Fighter A", "Fighter B")


def _seed_frame(all_cands):
    """Pick the frame to seed identity from: where the two largest people are
    BOTH big (the actual fighters, mid-action) — not merely the first frame
    with 2 detections, which happily seeds on a coach/bystander when a fighter
    is off-screen at t=0 and then follows the wrong person for the whole clip.
    Score = size of the SECOND-largest person (both must be big), with a small
    proximity bonus (fighters engage; bystanders hug the edges)."""
    best, best_score = -1, -1.0
    for i, cands in enumerate(all_cands):
        if len(cands) < 2:
            continue
        two = sorted(cands, key=lambda c: -c[2])[:2]
        pair_dist = _d(two[0][1], two[1][1])
        score = two[1][2] - 0.02 * pair_dist
        if score > best_score:
            best, best_score = i, score
    return best


def _track_span(all_cands, poses, cols, start, step, lastA, lastB):
    """Gated NN association from `start` moving by `step` (+1 fwd / -1 back),
    continuing from known track positions; writes poses + per-track colors."""
    n = len(all_cands)
    i = start
    while 0 <= i < n:
        cands = all_cands[i]
        kpA, kpB, lastA, lastB = _associate(cands, lastA, lastB)
        if kpA is not None:
            poses[i, 0] = kpA
            cols[0].append(_cand_color(cands, kpA))
        if kpB is not None:
            poses[i, 1] = kpB
            cols[1].append(_cand_color(cands, kpB))
        i += step


def _cand_color(cands, kp):
    for c in cands:
        if c[0] is kp:
            return c[3]
    return None


def _collect_poses(video_path, vid_stride=1):
    """Return (poses (n,2,17,2): per frame [trackA, trackB], labels tuple).

    Two passes: (1) YOLO over the whole clip collecting every person candidate
    (+ torso color); (2) seed identity at the BEST frame (see _seed_frame) and
    associate outward in both directions. Backward tracking recovers the start
    of the clip that frame-0 seeding used to get wrong."""
    model = YOLO(MODEL)
    all_cands = []
    for r in model.predict(source=video_path, stream=True, classes=[0], verbose=False, vid_stride=vid_stride):
        all_cands.append([(kp, hc, sh, _torso_bgr(r.orig_img, kp))
                          for kp, hc, sh in _candidates(r)])
    n = len(all_cands)
    poses = np.full((n, 2, 17, 2), np.nan)
    if n == 0:
        return np.zeros((0, 2, 17, 2)), ("Fighter A", "Fighter B")

    seed = _seed_frame(all_cands)
    if seed < 0:                                  # never saw two people at once
        return poses, ("Fighter A", "Fighter B")

    # assign identity at the seed frame: two largest people, left -> A
    scands = all_cands[seed]
    two = sorted(range(len(scands)), key=lambda j: -scands[j][2])[:2]
    two.sort(key=lambda j: scands[j][1][0])
    iA, iB = two
    poses[seed, 0], poses[seed, 1] = scands[iA][0], scands[iB][0]
    cols = ([scands[iA][3]], [scands[iB][3]])
    lastA = (scands[iA][1], scands[iA][2])            # track state: (position, size)
    lastB = (scands[iB][1], scands[iB][2])

    _track_span(all_cands, poses, cols, seed + 1, +1, lastA, lastB)  # -> end
    _track_span(all_cands, poses, cols, seed - 1, -1, lastA, lastB)  # -> start
    return poses, _label_tracks(cols[0], cols[1])


def _fps(video_path):
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    cap.release()
    return fps


# ============================ feature math (fast) ===========================

def _episodes(mask, min_len):
    mask = np.nan_to_num(mask).astype(bool)
    count, run = 0, 0
    for v in mask:
        run = run + 1 if v else 0
        if run == min_len:
            count += 1
    return count


def _hysteresis_intervals(signal, hi, lo, enter_frames, exit_frames):
    """Sustained excursions of `signal` above `hi`, debounced by hysteresis.
    Returns a list of (start_frame, end_frame) intervals; len() = episode count.

    An episode OPENS only after the signal stays above `hi` for `enter_frames`
    in a row, and CLOSES only after it drops below `lo` for `exit_frames` in a
    row. NaN frames are 'no evidence': they hold the current state and don't
    advance either counter, so a one-frame dropout or a flicker between the
    thresholds can't split one real lapse into many. This replaces the old
    run-length count that booked a fresh episode on every brief blip (e.g. the
    spurious '33 guard-dropped episodes' on a 180-frame read)."""
    intervals = []
    in_episode = False
    above = below = 0
    start = last = 0
    for i, v in enumerate(signal):
        if np.isnan(v):
            continue
        last = i
        if not in_episode:
            above = above + 1 if v > hi else 0
            if above >= enter_frames:
                start = i - enter_frames + 1
                in_episode = True
                below = 0
        else:
            below = below + 1 if v < lo else 0
            if below >= exit_frames:
                intervals.append((start, i - exit_frames + 1))
                in_episode = False
                above = 0
    if in_episode:
        intervals.append((start, last))
    return intervals


def _mmss(seconds):
    seconds = max(0, int(round(seconds)))
    return f"{seconds // 60}:{seconds % 60:02d}"


def _wrist_speed(w, s, fps):
    """Per-frame wrist speed in shoulder-widths/sec (NaN where unknown)."""
    d = np.linalg.norm(np.diff(w, axis=0), axis=1) / s * fps
    return np.concatenate([[np.nan], d])


def _dilate(mask, k):
    """Grow True regions by k frames each side (merges near-adjacent punches)."""
    if k <= 0:
        return mask
    out = mask.copy()
    for i in np.where(mask)[0]:
        out[max(0, i - k):i + k + 1] = True
    return out


def _runs(mask, min_len):
    """(start, end) frame intervals of True runs lasting at least min_len."""
    out, start = [], None
    for i, v in enumerate(mask):
        if v and start is None:
            start = i
        elif not v and start is not None:
            if i - start >= min_len:
                out.append((start, i - 1))
            start = None
    if start is not None and len(mask) - start >= min_len:
        out.append((start, len(mask) - 1))
    return out


def _events(mask, refractory):
    mask = np.nan_to_num(mask).astype(bool)
    frames, last = [], -10 ** 9
    for i, v in enumerate(mask):
        if v and i - last > refractory:
            frames.append(i)
            last = i
    return np.array(frames, dtype=int)


def _interp1(a):
    """Linear-fill NaNs in a 1-D signal."""
    a = a.astype(float).copy()
    idx = np.arange(len(a)); good = ~np.isnan(a)
    if good.sum() >= 2:
        a[~good] = np.interp(idx[~good], idx[good], a[good])
    elif good.sum() == 1:
        a[~good] = a[good][0]
    else:
        a[:] = 0.0
    return a


def _side_arrays(P):
    return {
        "nose": P[:, NOSE],
        "shm": np.nanmean(P[:, [L_SH, R_SH]], axis=1),
        "hip": np.nanmean(P[:, [L_HIP, R_HIP]], axis=1),
        "wl": P[:, L_WR],
        "wr": P[:, R_WR],
        "hl": P[:, L_HIP],
        "hr": P[:, R_HIP],
        "kl": P[:, L_KNEE],
        "kr": P[:, R_KNEE],
        "al": P[:, L_ANK],
        "ar": P[:, R_ANK],
        "sh": np.linalg.norm(P[:, L_SH] - P[:, R_SH], axis=1),
    }


def _knee_angle(hip, knee, ank):
    """Per-frame joint angle AT the knee in degrees (180 = straight leg).
    2D-projected, so it under-reads bend when the leg points at the camera —
    treat downstream labels as tendencies, not measurements."""
    v1, v2 = hip - knee, ank - knee
    denom = np.linalg.norm(v1, axis=1) * np.linalg.norm(v2, axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        cos = np.sum(v1 * v2, axis=1) / np.where(denom > 1e-6, denom, np.nan)
    return np.degrees(np.arccos(np.clip(cos, -1.0, 1.0)))


def _analyze(me, opp, scale, fps):
    s = scale
    gl = (me["wl"][:, 1] - me["nose"][:, 1]) / s
    gr = (me["wr"][:, 1] - me["nose"][:, 1]) / s
    guard_min = np.nanmin(np.stack([gl, gr]), axis=0)
    valid = int(np.sum(~np.isnan(guard_min)))
    hands_up = guard_min < GUARD_UP
    dropped = guard_min > GUARD_DOWN

    head_off = (me["nose"][:, 0] - me["hip"][:, 0]) / s
    head_std = float(np.nanstd(head_off))
    move_cat = "low" if head_std < 0.10 else "medium" if head_std < 0.20 else "high"

    # gap normalized by the fighters' shoulder width IN THAT FRAME (larger of
    # the two: robust when one blades sideways). The clip-median ruler misread
    # perspective on moving cameras: fighters near the lens have a bigger pixel
    # gap at the same real distance, so genuine in-range moments read "long".
    sh_pair = np.nanmax(np.stack([me["sh"], opp["sh"]]), axis=0)
    sh_pair = np.where(np.isnan(sh_pair) | (sh_pair < MIN_SH_PX), s, sh_pair)
    gap = np.abs(me["hip"][:, 0] - opp["hip"][:, 0]) / sh_pair
    in_range = gap < IN_RANGE
    danger = in_range & dropped

    d = np.nanmin(np.stack([
        np.linalg.norm(opp["wl"] - me["nose"], axis=1) / sh_pair,
        np.linalg.norm(opp["wr"] - me["nose"], axis=1) / sh_pair,
    ]), axis=0)
    tag_frames = _events(d < TAG_DIST, int(TAG_REFRACTORY_S * fps))

    # pressure / retreat: does this fighter walk the opponent down, or give
    # ground? Two honesty rules learned from coach feedback:
    # (1) only frame pairs where BOTH fighters were actually measured count —
    #     interpolating across tracking gaps used to fabricate motion;
    # (2) a handheld pan moves both fighters across the frame together, and a
    #     single camera cannot separate that from real footwork — when the two
    #     fighters' motions correlate like a pan, say "unclear" instead of
    #     printing a confident wrong label.
    my_raw = me["hip"][:, 0] / s
    op_raw = opp["hip"][:, 0] / s
    ok_pair = (~np.isnan(my_raw[:-1]) & ~np.isnan(my_raw[1:])
               & ~np.isnan(op_raw[:-1]) & ~np.isnan(op_raw[1:]))
    dmy = np.diff(my_raw)[ok_pair]
    dop = np.diff(op_raw)[ok_pair]
    adv = dmy * np.sign(op_raw[:-1] - my_raw[:-1])[ok_pair]
    net10 = 0.0
    if ok_pair.sum() < 10:
        pct_adv = pct_ret = 0.0
        stance = "unknown (tracking too thin)"
    else:
        pct_adv = round(100 * (adv > 0.03).mean(), 1)
        pct_ret = round(100 * (adv < -0.03).mean(), 1)
        # NET ground gained/given decides the label (coach feedback: a fighter
        # giving ground in fast bursts while shuffling forward slowly counted
        # as 'holding ground' by step-counting — the eye reads net ground, so
        # net wins; percentages stay as supporting detail).
        net10 = float(adv.sum() * fps * 10 / max(len(adv), 1))
        if net10 < -NET_RATE_SW10S or pct_ret - pct_adv > PRESSURE_MARGIN:
            stance = "backing up"
        elif net10 > NET_RATE_SW10S or pct_adv - pct_ret > PRESSURE_MARGIN:
            stance = "pressuring"
        else:
            stance = "holding ground"
        # decompose motion: common component (both drift together = camera pan)
        # vs relative (fighters actually closing/opening = real footwork).
        # Correlation can't catch a smooth pan (zero variance), this can.
        common = float(np.median(np.abs(dmy + dop) / 2))
        relative = float(np.median(np.abs(dmy - dop)))
        if common > PAN_DOMINANCE * relative + 0.015 and common > 0.02:
            stance = "unclear (camera moves with the fighters)"

    # range profile (robust: from the distance distribution, not velocity)
    gv = gap[~np.isnan(gap)]
    close = int(round(100 * np.mean(gv < 1.8))) if len(gv) else 0
    longr = int(round(100 * np.mean(gv > 3.2))) if len(gv) else 0
    mid = max(0, 100 - close - longr)

    def pct(mask):
        return round(100 * np.nan_to_num(mask).astype(bool).sum() / max(valid, 1), 1)

    drop_iv = _hysteresis_intervals(
        guard_min, GUARD_DOWN, GUARD_UP,
        max(2, int(DROP_MIN_S * fps)), max(2, int(CLEAR_MIN_S * fps)))

    # exchanges: in punching range while at least one of MY hands moves at punch
    # speed. (Velocity metrics were a dead end at ~20% coverage; at SAM2-level
    # coverage they're usable — the speed threshold is still an uncalibrated guess.)
    spd = np.fmax(_wrist_speed(me["wl"], s, fps), _wrist_speed(me["wr"], s, fps))
    throwing = _dilate(np.nan_to_num(spd) > PUNCH_SPEED_SW_S, int(EXCHANGE_PAD_S * fps))
    exch_mask = np.nan_to_num(in_range).astype(bool) & throwing
    exch_iv = _runs(exch_mask, max(1, int(EXCHANGE_MIN_S * fps)))
    # in-range stretches as timeline events (pairwise: same for both fighters).
    # Same min length as exchanges so an exchange can never sit outside a
    # displayed in-range band.
    range_iv = _runs(np.nan_to_num(in_range).astype(bool), max(1, int(EXCHANGE_MIN_S * fps)))

    # stance: which foot leads toward the opponent (left lead = orthodox), and
    # the MATCHUP over time: same lead foot = closed stance, opposite = open.
    face = np.sign(_interp1(opp["hip"][:, 0]) - _interp1(me["hip"][:, 0]))
    my_valid = ~np.isnan(me["al"][:, 0]) & ~np.isnan(me["ar"][:, 0])
    my_left_lead = (me["al"][:, 0] - me["ar"][:, 0]) * face > 0
    if my_valid.sum() >= 10:
        orth = 100 * my_left_lead[my_valid].mean()
        stance_str = f"orthodox {orth:.0f}% / southpaw {100 - orth:.0f}%"
    else:
        stance_str = "unknown (feet rarely visible)"
    opp_valid = ~np.isnan(opp["al"][:, 0]) & ~np.isnan(opp["ar"][:, 0])
    opp_left_lead = (opp["al"][:, 0] - opp["ar"][:, 0]) * (-face) > 0
    both_st = my_valid & opp_valid
    open_iv, closed_iv = [], []
    if both_st.sum() >= 10:
        closed = 100 * (my_left_lead[both_st] == opp_left_lead[both_st]).mean()
        matchup_str = f"closed {closed:.0f}% / open {100 - closed:.0f}%"
        # matchup stretches as timeline events (min 0.5s so one noisy ankle
        # frame can't flicker the band)
        same = my_left_lead == opp_left_lead
        open_iv = _runs(both_st & ~same, max(2, int(0.5 * fps)))
        closed_iv = _runs(both_st & same, max(2, int(0.5 * fps)))
    else:
        matchup_str = "unknown (feet rarely visible)"

    # leg drive: does the fighter BEND THE KNEES at punch release (sitting down
    # on shots) or arm-punch upright? Knee angle from hip-knee-ankle, compared
    # between release windows (wrist at punch speed ±0.2s) and the baseline.
    # 2D projection + uncalibrated thresholds — labels are coarse on purpose.
    knee = np.nanmean(np.stack([
        _knee_angle(me["hl"], me["kl"], me["al"]),
        _knee_angle(me["hr"], me["kr"], me["ar"]),
    ]), axis=0)
    knee_ok = ~np.isnan(knee)
    release = _dilate(np.nan_to_num(spd) > PUNCH_SPEED_SW_S, max(1, int(0.2 * fps)))
    if knee_ok.sum() >= 30:
        base_deg = float(np.nanmedian(knee[knee_ok]))
        load = knee_ok & release
        if load.sum() >= 8:
            load_deg = float(np.nanmedian(knee[load]))
            dip = base_deg - load_deg
            lbl = ("sits down on punches" if dip >= 6 else
                   "some knee bend into shots" if dip >= 3 else
                   "upright when punching, little visible leg drive")
            leg_drive = (f"{lbl} (knees {load_deg:.0f}\N{DEGREE SIGN} at release "
                         f"vs {base_deg:.0f}\N{DEGREE SIGN} baseline)")
        else:
            leg_drive = (f"baseline knee bend {base_deg:.0f}\N{DEGREE SIGN} "
                         "(too few tracked punches to judge loading)")
    else:
        leg_drive = "unknown (legs rarely visible)"

    # weight placement: where the hips sit between the feet, signed toward the
    # opponent (positive = over the front foot). Only frames with a real side-on
    # base (ankle spread > 0.25 shoulder-widths) count — a bladed or facing-
    # camera stance makes the projection meaningless.
    mid_ank = (me["al"][:, 0] + me["ar"][:, 0]) / 2
    base_w = np.abs(me["al"][:, 0] - me["ar"][:, 0])
    with np.errstate(invalid="ignore", divide="ignore"):
        shift = (me["hip"][:, 0] - mid_ank) / np.where(base_w > 1e-6, base_w, np.nan) * face
    w_ok = ~np.isnan(shift) & (base_w > 0.25 * s)
    if w_ok.sum() >= 30:
        sv = shift[w_ok]
        front = 100 * float(np.mean(sv > 0.12))
        rear = 100 * float(np.mean(sv < -0.12))
        centered = max(0.0, 100 - front - rear)
        lbl = ("rear-loaded" if rear - front > 15 else
               "front-heavy" if front - rear > 15 else "balanced")
        weight_side = (f"{lbl} (weight over front foot {front:.0f}% / "
                       f"centered {centered:.0f}% / rear foot {rear:.0f}%)")
        tw = w_ok & release
        if tw.sum() >= 8:
            ft = 100 * float(np.mean(shift[tw] > 0.12))
            if ft - front >= 12:
                weight_side += f"; drives onto the front foot when throwing ({ft:.0f}%)"
            elif front - ft >= 12:
                weight_side += f"; stays OFF the front foot when throwing ({ft:.0f}%)"
    else:
        weight_side = "unknown (feet rarely visible or stance too bladed to read)"

    # ---- habit metrics (all 2D; thresholds are uncalibrated coarse labels) ----
    my_thr = np.nan_to_num(spd) > PUNCH_SPEED_SW_S
    opp_spd = np.fmax(_wrist_speed(opp["wl"], s, fps), _wrist_speed(opp["wr"], s, fps))
    opp_thr = np.nan_to_num(opp_spd) > PUNCH_SPEED_SW_S
    exch_b = np.nan_to_num(exch_mask).astype(bool)

    # guard recovery: after each throwing burst, seconds until a hand is back
    # above guard height ("leaves the jab hanging" is this number)
    rec_times, left_out = [], 0
    for _a, _b in _runs(my_thr, 1):
        rec, t, saw_valid = None, _b + 1, False
        cap = min(len(guard_min), t + int(3 * fps))
        while t < cap:
            g = guard_min[t]
            if not np.isnan(g):
                saw_valid = True
                if g < GUARD_UP:
                    rec = (t - _b - 1) / fps
                    break
            t += 1
        if rec is not None:
            rec_times.append(rec)
        elif saw_valid:
            left_out += 1        # measured frames existed and the guard stayed low
        # all-NaN window = tracking gap, no evidence either way -> not counted
    if len(rec_times) + left_out >= 3:
        med_r = float(np.median(rec_times + [3.0] * left_out))   # failures count
        lbl = ("snaps the guard straight back" if med_r <= 0.3 else
               "adequate recovery" if med_r <= 0.7 else
               "leaves the hands out after punching")
        guard_recovery = (f"{lbl} (median {med_r:.1f}s back to guard"
                          + (f", slowest {max(rec_times):.1f}s" if rec_times else "")
                          + (f"; {left_out}x stayed low 3s+" if left_out else "") + ")")
    else:
        guard_recovery = "unknown (too few tracked punches)"

    # context of each sustained guard drop: punching, under fire, or off the ball
    ctx = {"while punching": 0, "under fire": 0, "off the ball": 0}
    for _a, _b in drop_iv:
        seg = slice(_a, _b + 1)
        own, theirs = float(np.mean(my_thr[seg])), float(np.mean(opp_thr[seg]))
        if theirs > 0.15 and theirs >= own:
            ctx["under fire"] += 1
        elif own > 0.15:
            ctx["while punching"] += 1
        else:
            ctx["off the ball"] += 1
    guard_drop_context = (", ".join(f"{v} {k}" for k, v in ctx.items() if v)
                          or "no sustained drops")

    # squaring up: torso width vs own height, exchanges vs baseline (self-
    # normalized, so no absolute threshold needed)
    body_h = np.nanmax(np.stack([me["al"][:, 1], me["ar"][:, 1]]), axis=0) - me["nose"][:, 1]
    with np.errstate(invalid="ignore", divide="ignore"):
        sq = me["sh"] / np.where(body_h > 1e-6, body_h, np.nan)
    sq_ok = ~np.isnan(sq)
    if sq_ok.sum() >= 30 and (sq_ok & exch_b).sum() >= 8:
        qr = float(np.nanmedian(sq[sq_ok & exch_b])) / max(float(np.nanmedian(sq[sq_ok])), 1e-6)
        lbl = ("squares up in exchanges" if qr > 1.15 else
               "stays bladed in exchanges" if qr < 0.9 else "keeps shape in exchanges")
        squaring_up = f"{lbl} (torso {100 * (qr - 1):+.0f}% wider than baseline when trading)"
    elif sq_ok.sum() >= 30:
        squaring_up = "unknown (no measured exchanges to compare)"
    else:
        squaring_up = "unknown (body rarely fully visible)"

    # chin: nose height above the shoulder line, exchanges vs baseline
    chin = (me["shm"][:, 1] - me["nose"][:, 1]) / s
    c_ok = ~np.isnan(chin)
    if c_ok.sum() >= 30:
        base_c = float(np.nanmedian(chin[c_ok]))
        if (c_ok & exch_b).sum() >= 8:
            ex_c = float(np.nanmedian(chin[c_ok & exch_b]))
            lbl = ("lifts the chin in exchanges" if ex_c - base_c > 0.08 else
                   "tucks tighter in exchanges" if ex_c - base_c < -0.08 else
                   "keeps the chin steady")
            chin_position = (f"{lbl} (head {base_c:.2f} SW above the shoulder line; "
                             f"{ex_c:.2f} when trading)")
        else:
            chin_position = f"head {base_c:.2f} SW above the shoulder line"
    else:
        chin_position = "unknown"

    # base width (uses base_w from the weight block)
    wv = (base_w / s)[~np.isnan(base_w)]
    if len(wv) >= 30:
        med_w = float(np.median(wv))
        lbl = "narrow base" if med_w < 0.65 else "wide base" if med_w > 1.5 else "normal base"
        stance_width = (f"{lbl} (median {med_w:.1f} SW between the feet; "
                        f"narrow <0.6 on {100 * float(np.mean(wv < 0.6)):.0f}% of frames)")
    else:
        stance_width = "unknown (feet rarely visible)"

    # crossed feet: lead/rear ankle order flips while moving laterally fast
    # (speed gate keeps stationary stance switches out — mostly)
    hipv = np.abs(np.concatenate([[np.nan], np.diff(me["hip"][:, 0])])) / s * fps
    crossed, prev_i, last_hit = [], None, -10 ** 9
    ax, bx = me["al"][:, 0], me["ar"][:, 0]
    for i in range(len(ax)):
        if np.isnan(ax[i]) or np.isnan(bx[i]) or abs(ax[i] - bx[i]) < 0.15 * s:
            continue
        if (prev_i is not None and np.sign(ax[i] - bx[i]) != np.sign(ax[prev_i] - bx[prev_i])
                and i - prev_i <= int(0.7 * fps) and i - last_hit > fps):
            v = np.nanmax(hipv[max(0, prev_i):i + 1])
            if v and v > 0.8:
                crossed.append(i)
                last_hit = i
        prev_i = i
    crossed_feet = (f"{len(crossed)}x while moving ({', '.join(_mmss(i / fps) for i in crossed)})"
                    if crossed else "none seen")

    # reaction to incoming: head moved laterally around the opponent's bursts
    head_rel = (me["nose"][:, 0] - me["hip"][:, 0]) / s
    n_ep = n_moved = 0
    for _a, _b in _runs(opp_thr, 1):
        w0 = max(0, _a - int(0.15 * fps))
        w1 = min(len(head_rel), _b + int(0.4 * fps) + 1)
        seg = head_rel[w0:w1]
        seg = seg[~np.isnan(seg)]
        if len(seg) >= 3:
            n_ep += 1
            n_moved += float(np.max(seg) - np.min(seg)) > 0.25
    reacts_to_punches = (f"moves the head on {100 * n_moved / n_ep:.0f}% of the "
                         f"opponent's punches (n={n_ep})"
                         if n_ep >= 4 else "unknown (too few opponent punches tracked)")

    # round trend: first vs last 30s bucket (needs a 60s+ clip)
    bframes = int(BUCKET_S * fps)
    nb = int(np.ceil(len(guard_min) / max(bframes, 1)))
    round_trend = "clip too short for a trend (<60s)"
    if nb >= 2 and valid >= 60:
        per = []
        for k in range(nb):
            seg = slice(k * bframes, (k + 1) * bframes)
            g = guard_min[seg]
            g = g[~np.isnan(g)]
            per.append((100 * float(np.mean(g < GUARD_UP)) if len(g) >= 10 else None,
                        100 * float(np.mean(my_thr[seg]))))
        firsts = [p for p in per if p[0] is not None]
        if len(firsts) >= 2:
            (h0, t0), (h1, t1) = firsts[0], firsts[-1]
            fading = (h0 - h1 > 10) or (t0 > 2 and t1 < 0.6 * t0)
            round_trend = (f"{'fading' if fading else 'steady'} across the round "
                           f"(hands up {h0:.0f}%→{h1:.0f}%, throwing {t0:.0f}%→{t1:.0f}%)")

    metrics = {
        "frames_analyzed": valid,
        "hands_up_pct": pct(hands_up),
        "guard_dropped_episodes": len(drop_iv),
        "guard_dropped_at": [f"{_mmss(s / fps)}-{_mmss(e / fps)}" for s, e in drop_iv],
        "head_movement": f"{move_cat} (std={round(head_std, 3)})",
        "avg_distance_shoulderwidths": round(float(np.nanmean(gap)), 2),
        "pct_in_range_with_hands_down": pct(danger),
        "pressure": f"{stance} (net {net10:+.1f} SW/10s; advancing {pct_adv}% / retreating {pct_ret}%)",
        "range_profile": f"close {close}% / mid {mid}% / long {longr}%",
        "stance": stance_str,
        "stance_matchup": matchup_str,
        "leg_drive": leg_drive,
        "weight_side": weight_side,
        "guard_recovery": guard_recovery,
        "guard_drop_context": guard_drop_context,
        "squaring_up": squaring_up,
        "chin_position": chin_position,
        "stance_width": stance_width,
        "crossed_feet": crossed_feet,
        "reacts_to_punches": reacts_to_punches,
        "round_trend": round_trend,
        "open_stance_at": [f"{_mmss(a / fps)}-{_mmss(b / fps)}" for a, b in open_iv],
        "closed_stance_at": [f"{_mmss(a / fps)}-{_mmss(b / fps)}" for a, b in closed_iv],
        "in_range_at": [f"{_mmss(a / fps)}-{_mmss(b / fps)}" for a, b in range_iv],
        "in_range_throwing_pct": pct(exch_mask),
        "exchanges_at": [f"{_mmss(a / fps)}-{_mmss(b / fps)}" for a, b in exch_iv],
        "head_contact_events_approx_EXPERIMENTAL": int(len(tag_frames)),
        "head_contact_times_approx_EXPERIMENTAL": [_mmss(i / fps) for i in tag_frames],
    }
    # strided-frame indices for render_overlay / the dashboard strip
    overlay = {"guard_intervals": drop_iv, "exchange_intervals": exch_iv,
               "range_intervals": range_iv, "open_intervals": open_iv,
               "closed_intervals": closed_iv}
    return metrics, overlay


def _exchange_activity(sA, sB, scale, fps):
    """Graded EXCHANGE INTENSITY over time, and the flurry windows it peaks in.

    This is the robust alternative to counting individual punches: instead of
    committing to 'a punch happened at instant t' (fragile), it integrates two
    reliable, camera-robust signals over a window — the two fighters are IN
    RANGE, and their arms are ACTIVE (wrist speed relative to their own hips,
    which cancels camera pan and footwork). High + sustained = a flurry. Even
    when per-punch calls are wrong, the hot windows are real exchanges.

    Returns (flurries [(start, end, 'light'|'heavy')], intensity series)."""
    sh_pair = np.nanmax(np.stack([sA["sh"], sB["sh"]]), axis=0)
    sh_pair = np.where(np.isnan(sh_pair) | (sh_pair < MIN_SH_PX), scale, sh_pair)
    gap = np.abs(sA["hip"][:, 0] - sB["hip"][:, 0]) / sh_pair
    in_range = gap < IN_RANGE                      # NaN gap -> False (not in range)

    def arm(s):
        aL = _wrist_speed(s["wl"] - s["hip"], scale, fps)
        aR = _wrist_speed(s["wr"] - s["hip"], scale, fps)
        return np.fmax(np.nan_to_num(aL), np.nan_to_num(aR))

    intensity = np.where(in_range, arm(sA) + arm(sB), 0.0)
    k = max(1, int(0.4 * fps))                     # ~0.4s smoothing
    smooth = np.convolve(intensity, np.ones(k) / k, mode="same")
    flurries = []
    for a, b in _runs(smooth > FLURRY_SPEED_SW_S, max(2, int(FLURRY_MIN_S * fps))):
        peak = float(np.max(smooth[a:b + 1])) if b >= a else 0.0
        flurries.append((a, b, "heavy" if peak > FLURRY_HEAVY_SW_S else "light"))
    return flurries, smooth


def _confidence(frames_analyzed, two_fighter_pct):
    """How much to trust this fighter's read, from how much of the clip we
    actually measured. Surfaced to the user AND the LLM so a thin read gets
    hedged instead of stated as fact. Thresholds are uncalibrated guesses."""
    if frames_analyzed >= 120 and two_fighter_pct >= 40:
        level = "high"
    elif frames_analyzed >= 50 and two_fighter_pct >= 20:
        level = "medium"
    else:
        level = "low"
    return (f"{level} — read from {frames_analyzed} frames; "
            f"both fighters were in view for {two_fighter_pct}% of the clip")


def _features_from_poses(poses, fps, labels=("Fighter A", "Fighter B"), mask_overlap=None):
    """Returns (features_dict, overlay_info) or (None, None).

    mask_overlap: optional per-frame SAM2 mask overlap ratio (intersection /
    smaller mask). When given, it upgrades the clinch read — mask overlap
    catches tie-ups that keypoint distance misreads."""
    if len(poses) == 0:
        return None, None
    A, B = poses[:, 0], poses[:, 1]
    shsA = np.linalg.norm(A[:, L_SH] - A[:, R_SH], axis=1)
    shsB = np.linalg.norm(B[:, L_SH] - B[:, R_SH], axis=1)
    both = int(np.sum(~np.isnan(shsA) & ~np.isnan(shsB)))
    all_sh = np.concatenate([shsA, shsB])
    all_sh = all_sh[~np.isnan(all_sh)]
    if both < MIN_PAIR_FRAMES or len(all_sh) == 0:
        return None, None
    scale = float(np.median(all_sh))
    cov = round(100 * both / max(len(poses), 1), 1)
    sA, sB = _side_arrays(A), _side_arrays(B)
    fa, ov_a = _analyze(sA, sB, scale, fps)
    fb, ov_b = _analyze(sB, sA, scale, fps)
    # Per-fighter confidence (non-underscore key so it survives build_prompt and
    # reaches the coach, unlike _meta which is stripped before the LLM sees it).
    fa["confidence"] = _confidence(fa["frames_analyzed"], cov)
    fb["confidence"] = _confidence(fb["frames_analyzed"], cov)
    # exchange intensity / flurries — pairwise, so the same list on both fighters
    flurries, _intensity = _exchange_activity(sA, sB, scale, fps)
    flur_str = [f"{_mmss(a / fps)}-{_mmss(b / fps)} ({lvl})" for a, b, lvl in flurries]
    fa["flurries_at"] = fb["flurries_at"] = flur_str
    ov_a["flurry_intervals"] = ov_b["flurry_intervals"] = flurries
    # Per-fighter NOT-TRACKED intervals (>=0.5s of missing pose: occluded, off
    # frame, or too close to the lens for keypoints). Timeline honesty layer —
    # shows WHERE the metrics are blind, like the overlay's "not tracked" flag.
    # Lives in _meta so it never reaches the LLM (coverage already hedges it).
    # frame-accurate event intervals for the dashboard timeline (seconds, 2dp).
    # The *_at strings are m:ss for the LLM/report — second-quantized, which
    # chopped continuous timeline bands into 1s blocks. The dashboard reads
    # these instead; the strings stay for the coach.
    def _sec(ivs):
        return [[round(a / fps, 2), round((b + 1) / fps, 2)] for a, b in ivs]
    timeline = {}
    for lab, ov in ((labels[0], ov_a), (labels[1], ov_b)):
        timeline[lab] = {
            "guard": _sec(ov["guard_intervals"]),
            "range": _sec(ov["range_intervals"]),
            "exch": _sec(ov["exchange_intervals"]),
            "open": _sec(ov["open_intervals"]),
            "closed": _sec(ov["closed_intervals"]),
        }
    timeline["flurries"] = [[round(a / fps, 2), round((b + 1) / fps, 2), lvl]
                            for a, b, lvl in flurries]

    # clinch / tie-up: bodies close AND arms quiet (a close+busy moment is an
    # exchange, not a clinch). SAM2 mask overlap, when available, also counts —
    # overlapping silhouettes are the direct evidence of a tie-up.
    sh_pair = np.nanmax(np.stack([sA["sh"], sB["sh"]]), axis=0)
    sh_pair = np.where(np.isnan(sh_pair) | (sh_pair < MIN_SH_PX), scale, sh_pair)
    pair_gap = np.abs(sA["hip"][:, 0] - sB["hip"][:, 0]) / sh_pair
    arm_all = (np.nan_to_num(np.fmax(_wrist_speed(sA["wl"], scale, fps),
                                     _wrist_speed(sA["wr"], scale, fps)))
               + np.nan_to_num(np.fmax(_wrist_speed(sB["wl"], scale, fps),
                                       _wrist_speed(sB["wr"], scale, fps))))
    tie = np.nan_to_num(pair_gap < 1.1) & (arm_all < 3.0)
    if mask_overlap is not None and len(mask_overlap) == len(tie):
        tie = tie | (np.nan_to_num(mask_overlap) > 0.30)
    clinch_iv = _runs(tie.astype(bool), max(2, int(0.6 * fps)))
    clinch_frames = sum(b - a + 1 for a, b in clinch_iv)
    clinch_str = (f"{100 * clinch_frames / max(len(poses), 1):.0f}% of the clip tied up ("
                  + ", ".join(f"{_mmss(a / fps)}-{_mmss(b / fps)}" for a, b in clinch_iv) + ")"
                  if clinch_iv else "no clinch time detected")
    fa["clinch"] = fb["clinch"] = clinch_str
    timeline["clinch"] = [[round(a / fps, 2), round((b + 1) / fps, 2)] for a, b in clinch_iv]

    # measured size hint: same-frame shoulder-width ratio (A/B). Same-frame, so
    # camera distance mostly cancels; still 2D, so it grounds (not replaces) the
    # VLM's qualitative size read.
    pair_mask = ~np.isnan(shsA) & ~np.isnan(shsB)
    with np.errstate(invalid="ignore", divide="ignore"):
        sw_ratio = round(float(np.nanmedian(shsA[pair_mask] / shsB[pair_mask])), 2) \
            if pair_mask.sum() >= 10 else None

    # height ratio (A/B): nose-to-lowest-ankle vertical extent, same frame only.
    # Shoulder width reads BREADTH — a lanky fighter can out-height a stockier
    # one while reading "similar" on shoulders; this is the tallness signal.
    # Median over frames dampens crouching; still 2D and depth-skewed.
    def _body_h(P):
        ank_y = np.nanmax(np.stack([P[:, L_ANK, 1], P[:, R_ANK, 1]]), axis=0)
        return ank_y - P[:, NOSE, 1]
    hA, hB = _body_h(A), _body_h(B)
    hmask = ~np.isnan(hA) & ~np.isnan(hB) & (hA > 0) & (hB > 0)
    with np.errstate(invalid="ignore", divide="ignore"):
        height_ratio = round(float(np.nanmedian(hA[hmask] / hB[hmask])), 2) \
            if hmask.sum() >= 8 else None
    gap_min = max(2, int(0.5 * fps))
    # a fighter counts as NOT TRACKED when the shoulders are missing (no scale
    # ruler) OR the whole pose is too sparse to trust (<6 confident keypoints —
    # mostly out of frame, fully absent, or too close to the lens). Shoulder-NaN
    # alone missed the "half a leg visible" case.
    kp_count = [np.sum(~np.isnan(poses[:, i, :, 0]), axis=1) for i in range(2)]
    gone = [np.isnan(shsA) | (kp_count[0] < 6),
            np.isnan(shsB) | (kp_count[1] < 6)]
    gap_runs = {labels[0]: _runs(gone[0], gap_min),
                labels[1]: _runs(gone[1], gap_min)}
    not_tracked = {lab: [f"{_mmss(a / fps)}-{_mmss(b / fps)}" for a, b in runs]
                   for lab, runs in gap_runs.items()}
    for lab, runs in gap_runs.items():
        timeline[lab]["gaps"] = _sec(runs)
    # BOTH fighters gone at once = the camera lost the action entirely — shown
    # on the shared Round lane so total blind spots are unmissable
    timeline["both_gone"] = _sec(_runs(gone[0] & gone[1], gap_min))
    feats = {
        "_meta": {
            "frames": int(len(poses)),
            "frames_with_two_fighters": both,
            "two_fighter_pct": cov,
            "scale_px": round(scale, 1),
            "not_tracked_at": not_tracked,
            "shoulder_width_ratio": sw_ratio,
            "height_ratio": height_ratio,
            "timeline": timeline,
        },
        labels[0]: fa,
        labels[1]: fb,
    }
    return feats, {labels[0]: ov_a, labels[1]: ov_b}


# ============================ public API ====================================

def extract_features_full(video_path, target_frames=350):
    """extract_features + everything render_overlay needs. Returns None on
    failure, else a dict: features, poses, stride, fps (original), labels,
    overlay (per-fighter event intervals in strided-frame indices)."""
    fps = _fps(video_path) or 30.0
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    stride = max(VID_STRIDE, total // target_frames) if total > target_frames else VID_STRIDE
    poses, labels = _collect_poses(video_path, vid_stride=stride)
    feats, overlay = _features_from_poses(poses, fps / stride, labels)
    if not feats:
        return None
    return {"features": feats, "poses": poses, "stride": stride,
            "fps": fps, "labels": labels, "overlay": overlay}


def extract_features(video_path, target_frames=350):
    """Adaptively sample frames so any clip finishes in bounded time on CPU:
    short clips run every frame (full quality); long clips are subsampled to
    ~target_frames. Defensive metrics are aggregates, so a bounded sample is fine."""
    full = extract_features_full(video_path, target_frames)
    return full["features"] if full else None


def cache_poses(video_path, npz_path):
    poses, labels = _collect_poses(video_path)
    np.savez_compressed(npz_path, poses=poses, fps=_fps(video_path), labels=np.array(labels))
    return npz_path


def features_from_cache(npz_path):
    d = np.load(npz_path, allow_pickle=False)
    labels = tuple(str(x) for x in d["labels"]) if "labels" in d.files else ("Fighter A", "Fighter B")
    feats, _ = _features_from_poses(d["poses"], float(d["fps"]), labels)
    return feats


# ============================ tracking overlay ==============================

_LIMBS = [(5, 7), (7, 9), (6, 8), (8, 10), (5, 6), (5, 11), (6, 12),
          (11, 12), (11, 13), (13, 15), (12, 14), (14, 16), (0, 5), (0, 6)]
_OVERLAY_MAX_W = 960


_DRAW_COLORS = {   # BGR, chosen to stay visible over ring/canvas footage
    "red": (60, 60, 230), "orange": (0, 140, 255), "yellow": (40, 220, 240),
    "green": (80, 200, 80), "blue": (230, 140, 40), "purple": (200, 100, 200),
    "white": (240, 240, 240), "black": (80, 80, 80), "gray": (160, 160, 160),
}


def _track_color(label):
    """BGR color matched to the track's display name."""
    n = label.lower()
    for name, bgr in _DRAW_COLORS.items():
        if name in n:
            return bgr
    return (0, 165, 255) if n.endswith("a") else (255, 220, 0)


def _draw_fighter(img, kp, color, label, guard_down):
    """Draw one fighter; returns True if any keypoint was visible this frame."""
    pts = [None if np.isnan(p).any() else (int(p[0]), int(p[1])) for p in kp]
    for a, b in _LIMBS:
        if pts[a] and pts[b]:
            cv2.line(img, pts[a], pts[b], color, 2, cv2.LINE_AA)
    for p in pts:
        if p:
            cv2.circle(img, p, 3, color, -1, cv2.LINE_AA)
    vis = [p for p in pts if p]
    if not vis:
        return False
    top = min(vis, key=lambda p: p[1])
    anchor = (max(5, top[0] - 40), max(20, top[1] - 12))
    cv2.putText(img, label, anchor, cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 3, cv2.LINE_AA)
    cv2.putText(img, label, anchor, cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 1, cv2.LINE_AA)
    if guard_down:
        fl = (anchor[0], anchor[1] + 18)
        cv2.putText(img, "GUARD DOWN", fl, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 3, cv2.LINE_AA)
        cv2.putText(img, "GUARD DOWN", fl, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
    return True


def render_overlay(video_path, full, out_path):
    """Write a tracking-overlay video: per-fighter skeletons color-matched to
    their labels, a running clock (matches guard_dropped_at timestamps), a
    GUARD DOWN flag during drop episodes, and an explicit 'not tracked' marker
    when a fighter has no pose that frame (so tracking gaps are visible instead
    of looking like a broken feature). Pure CPU; uses the poses already
    extracted (no second model pass). Re-encodes to H.264 via ffmpeg so browsers
    can play it inline.

    Returns (path_or_None, note): `note` explains any failure or warns about
    thin tracking; every stage is also print-logged for the Space logs."""
    import shutil
    import subprocess
    import tempfile

    poses, stride, fps = full["poses"], full["stride"], full["fps"]
    labels, overlay = full["labels"], full["overlay"]
    eff_fps = max(1.0, fps / stride)

    # per-fighter guard-down membership over strided frames
    n = len(poses)
    down = {}
    for lab in labels:
        m = np.zeros(n, dtype=bool)
        for s, e in overlay.get(lab, {}).get("guard_intervals", []):
            m[max(0, s):min(n, e + 1)] = True
        down[lab] = m

    cap = cv2.VideoCapture(video_path)
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    if not w or not h:
        cap.release()
        print(f"[overlay] could not open/read video ({w}x{h}): {video_path}", flush=True)
        return None, "Couldn't re-read the uploaded video to draw on it."
    sc = min(1.0, _OVERLAY_MAX_W / w)
    ow, oh = int(w * sc) // 2 * 2, int(h * sc) // 2 * 2   # even dims for H.264
    print(f"[overlay] source {w}x{h} -> {ow}x{oh}, stride={stride}, eff_fps={eff_fps:.1f}", flush=True)

    tmp = tempfile.mktemp(suffix=".mp4")
    vw = cv2.VideoWriter(tmp, cv2.VideoWriter_fourcc(*"mp4v"), eff_fps, (ow, oh))
    written = 0
    idx = 0
    seen = {lab: 0 for lab in labels}
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        if idx % stride == 0:
            si = idx // stride
            if si < n:
                if (frame.shape[1], frame.shape[0]) != (ow, oh):
                    frame = cv2.resize(frame, (ow, oh))
                missing = []
                for t, lab in enumerate(labels):
                    drawn = _draw_fighter(frame, poses[si, t] * sc, _track_color(lab),
                                          lab, bool(down[lab][si]))
                    if drawn:
                        seen[lab] += 1
                    else:
                        missing.append(lab)
                for j, lab in enumerate(missing):
                    cv2.putText(frame, f"{lab}: not tracked", (ow - 230, 22 + 18 * j),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 3, cv2.LINE_AA)
                    cv2.putText(frame, f"{lab}: not tracked", (ow - 230, 22 + 18 * j),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 200), 1, cv2.LINE_AA)
                clock = _mmss(si / eff_fps)
                cv2.putText(frame, clock, (10, 28), cv2.FONT_HERSHEY_SIMPLEX,
                            0.8, (255, 255, 255), 3, cv2.LINE_AA)
                cv2.putText(frame, clock, (10, 28), cv2.FONT_HERSHEY_SIMPLEX,
                            0.8, (30, 30, 30), 1, cv2.LINE_AA)
                vw.write(frame)
                written += 1
        idx += 1
    cap.release()
    vw.release()
    print(f"[overlay] wrote {written} frames; visibility: "
          + ", ".join(f"{lab} {100 * c // max(written, 1)}%" for lab, c in seen.items()), flush=True)
    if not written:
        return None, "No frames could be rendered from this clip."

    # warn when tracking was thin — the overlay will show mostly bare video
    thin = [f"{lab} visible in only {100 * c // written}% of frames"
            for lab, c in seen.items() if c < 0.3 * written]
    note = ("⚠️ Thin tracking: " + "; ".join(thin) + " — stretches without a skeleton "
            "mean the model lost that fighter (see 'not tracked' marker).") if thin else ""

    # mp4v doesn't play in browsers; H.264 via ffmpeg is required for inline playback
    if shutil.which("ffmpeg"):
        r = subprocess.run(
            ["ffmpeg", "-y", "-loglevel", "error", "-i", tmp, "-c:v", "libx264",
             "-pix_fmt", "yuv420p", "-movflags", "+faststart", out_path],
            capture_output=True)
        if r.returncode == 0 and os.path.exists(out_path):
            os.remove(tmp)
            print(f"[overlay] h264 ok -> {out_path}", flush=True)
            return out_path, note
        print(f"[overlay] ffmpeg re-encode failed: {r.stderr.decode(errors='replace')[:300]}", flush=True)
        os.remove(tmp)
        return None, "Overlay rendered but couldn't be converted to a browser-playable format."
    print("[overlay] ffmpeg not found; skipping (mp4v won't play in browsers)", flush=True)
    os.remove(tmp)
    return None, "Overlay disabled: ffmpeg isn't available on this server."