"""
SAM2 two-fighter tracker — Space GPU path (occlusion-robust identity + crop-pose).

Validated in Colab: two-fighter coverage 23% (pose-only) -> 58% (this) on hard
footage. SAM2 segments + tracks each fighter through the clip; pose runs on each
fighter's mask CROP (isolated fighter = easier detection) and the result feeds
the UNCHANGED feature extractor + overlay renderer.

Space-specific differences from the Colab original:
- Frames are extracted to JPEGs with cv2 and fed to SAM2's init_state as a
  directory (SAM2's mp4 loader needs decord, which has no py3.13 wheels).
- Identity seeds at the BEST frame (both fighters big, near — same idea as the
  pose path's _seed_frame), then propagates FORWARD AND BACKWARD; the Colab
  version seeded at frame 0 and inherited the bystander bug.
- fp16 autocast on pre-Ampere GPUs (T4 has no bfloat16), bf16 on Ampere+.
- Long clips are strided to ~TARGET_FRAMES like the pose path.
"""

import os
import shutil
import tempfile

import numpy as np
import cv2
import torch

import pose_features as pf

SAM2_ID = "facebook/sam2-hiera-base-plus"
SAM2_POSE_MODEL = "yolo11m-pose.pt"   # bigger model is fine: it only sees small crops
                                      # (distinct name: POSE_MODEL is the env knob for
                                      # pose_features' full-frame model)
DET_MODEL = "yolo11n.pt"
KP_CONF = 0.30
# Frame BUDGET, not a frame RATE -- so effective fps is TARGET_FRAMES divided
# by clip length, and a long clip is measured more coarsely than a short one.
# At 300 a 30s clip was sampled at 10 fps and a 120s clip at 2.5 fps, which is
# below the floor where a punch can be detected at all: a hand travels further
# between two samples than the whole extension the detector looks for. 750
# holds every clip up to the 60s limit at >=12.5 fps.
#
# This is a GPU-cost decision as much as an accuracy one -- it is ~2.5x the
# frames of the old budget. Lower it and clips get coarser, do not get faster
# per frame. See docs/research/accuracy-for-combat-sports-athletes.md section 1.
TARGET_FRAMES = int(os.environ.get("SAM2_TARGET_FRAMES", "750"))
SEED_DET_STRIDE = 10             # run the person detector on every Nth sampled frame

# Per-process caches — constructing YOLO/SAM2 per request costs seconds each,
# and _tracks_collapsed retries were re-paying it up to 3x per clip.
_models = None
_predictor = None


def _get_models():
    # Imports stay lazy (heavy CUDA init) but must live HERE: at call time this
    # runs at module scope, where a function-local import in the caller is not
    # visible.
    global _models
    if _models is None:
        from ultralytics import YOLO
        _models = (YOLO(DET_MODEL), YOLO(SAM2_POSE_MODEL))
    return _models


def _get_predictor():
    global _predictor
    if _predictor is None:
        from sam2.sam2_video_predictor import SAM2VideoPredictor
        _predictor = SAM2VideoPredictor.from_pretrained(SAM2_ID)
    return _predictor


def _extract_jpeg_frames(video_path, tmpdir):
    """Sample the clip to ~TARGET_FRAMES and write JPEGs SAM2 can read.
    Returns (frame_paths, stride, fps, (w, h))."""
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    stride = max(1, total // TARGET_FRAMES) if total > TARGET_FRAMES else 1
    paths, idx, si = [], 0, 0
    size = None
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        if idx % stride == 0:
            p = os.path.join(tmpdir, f"{si:05d}.jpg")
            cv2.imwrite(p, frame, [cv2.IMWRITE_JPEG_QUALITY, 92])
            paths.append(p)
            size = (frame.shape[1], frame.shape[0])
            si += 1
        idx += 1
    cap.release()
    return paths, stride, fps, size


def _iou(a, b):
    ix0, iy0 = max(a[0], b[0]), max(a[1], b[1])
    ix1, iy1 = min(a[2], b[2]), min(a[3], b[3])
    iw, ih = max(0.0, ix1 - ix0), max(0.0, iy1 - iy0)
    inter = iw * ih
    ua = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter / ua if ua > 0 else 0.0

MAX_SEED_IOU = 0.25       # the two seed boxes must be DIFFERENT people:
MIN_CENTER_SEP = 0.6      #   low overlap AND centers > this many box-widths apart


def _valid_pair(bi, bj):
    """True if two person boxes are plausibly two DIFFERENT people. YOLO can
    emit two boxes for one person (part + whole); seeding SAM2 with such a pair
    tracks the same fighter twice — 'Fighter A and B are the same guy'."""
    if _iou(bi, bj) > MAX_SEED_IOU:
        return False
    wi, wj = bi[2] - bi[0], bj[2] - bj[0]
    ci = ((bi[0] + bi[2]) / 2, (bi[1] + bi[3]) / 2)
    cj = ((bj[0] + bj[2]) / 2, (bj[1] + bj[3]) / 2)
    sep = float(np.hypot(ci[0] - cj[0], ci[1] - cj[1]))
    return sep > MIN_CENTER_SEP * (wi + wj) / 2


def _seed_candidates(det, frame_paths, k=3):
    """Ranked [(seed_idx, two boxes left->right), ...] of frames where two
    genuinely DISTINCT large people are visible. Candidates are spaced apart in
    time so a retry after a collapsed pair tries a different moment, not the
    neighbouring frame with the same problem."""
    scored = []
    for i in range(0, len(frame_paths), SEED_DET_STRIDE):
        frame = cv2.imread(frame_paths[i])
        r = det.predict(frame, classes=[0], verbose=False,
                        imgsz=640, conf=0.25, **pf.predict_kwargs())[0]
        if r.boxes is None or len(r.boxes) < 2:
            continue
        b = r.boxes.xyxy.cpu().numpy()
        areas = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1])
        order = np.argsort(-areas)[:4]                   # consider the 4 largest
        best = None
        for x in range(len(order)):
            for y in range(x + 1, len(order)):
                bi, bj = b[order[x]], b[order[y]]
                if not _valid_pair(bi, bj):
                    continue
                score = float(np.sqrt(min(areas[order[x]], areas[order[y]])))
                if best is None or score > best[0]:
                    best = (score, bi, bj)
        if best is not None:
            two = np.stack([best[1], best[2]])
            scored.append((best[0], i, two[np.argsort(two[:, 0])]))
    scored.sort(key=lambda t: -t[0])
    picked = []
    for s, i, boxes in scored:
        if all(abs(i - j) >= SEED_DET_STRIDE * 3 for _, j, _ in picked):
            picked.append((s, i, boxes))
        if len(picked) >= k:
            break
    return [(i, boxes) for _, i, boxes in picked]


def _tracks_collapsed(poses):
    """True if the two tracks are really the same person: among frames where
    both are present, their hips stay within a fraction of a shoulder-width."""
    hipA = np.nanmean(poses[:, 0, [pf.L_HIP, pf.R_HIP]], axis=1)
    hipB = np.nanmean(poses[:, 1, [pf.L_HIP, pf.R_HIP]], axis=1)
    both = ~np.isnan(hipA).any(1) & ~np.isnan(hipB).any(1)
    if both.sum() < 10:
        return True                                       # too little signal to trust
    gaps = np.linalg.norm(hipA[both] - hipB[both], axis=1)
    sh = np.concatenate([
        np.linalg.norm(poses[:, 0, pf.L_SH] - poses[:, 0, pf.R_SH], axis=1),
        np.linalg.norm(poses[:, 1, pf.L_SH] - poses[:, 1, pf.R_SH], axis=1)])
    scale = np.nanmedian(sh)
    return bool(np.median(gaps) < 0.6 * scale)


def _autocast_dtype():
    if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
        return torch.bfloat16
    return torch.float16                                  # T4 (sm_75) has no bf16


def _pose_in_mask(frame, mask, pose):
    """Pose on the crop around a fighter's mask -> full-frame keypoints or None."""
    ys, xs = np.where(mask)
    if len(xs) < 50:
        return None
    x0, x1, y0, y1 = int(xs.min()), int(xs.max()), int(ys.min()), int(ys.max())
    pad = int(0.12 * max(x1 - x0, y1 - y0))
    h, w = frame.shape[:2]
    x0, y0 = max(0, x0 - pad), max(0, y0 - pad)
    x1, y1 = min(w, x1 + pad), min(h, y1 + pad)
    crop = frame[y0:y1, x0:x1]
    if crop.size == 0:
        return None
    r = pose.predict(crop, classes=[0], verbose=False,
                     imgsz=640, conf=0.25, **pf.predict_kwargs())[0]
    if r.keypoints is None or r.boxes is None or len(r.boxes) == 0:
        return None
    # pick the detection ON THIS FIGHTER'S MASK, not the largest in the crop:
    # at close range both fighters appear in both crops, and largest-wins gave
    # BOTH tracks the same person's pose exactly during exchanges
    sub = mask[y0:y1, x0:x1]
    best_i, best_ov = -1, 0.0
    for j, (bx0, by0, bx1, by1) in enumerate(r.boxes.xyxy.cpu().numpy().astype(int)):
        bx0, by0 = max(0, bx0), max(0, by0)
        bx1, by1 = min(sub.shape[1], bx1), min(sub.shape[0], by1)
        if bx1 <= bx0 or by1 <= by0:
            continue
        ov = float(sub[by0:by1, bx0:bx1].mean())    # share of box on the mask
        if ov > best_ov:
            best_ov, best_i = ov, j
    if best_i < 0 or best_ov < 0.15:
        return None                                  # nobody actually on the mask
    i = best_i
    kp = r.keypoints.xy[i].cpu().numpy()
    cf = r.keypoints.conf[i].cpu().numpy()
    kp[cf < KP_CONF] = np.nan
    kp[:, 0] += x0
    kp[:, 1] += y0
    return kp


def sam2_tracks(video_path, progress_cb=None):
    """RawTracks from SAM2 mask propagation + crop pose.

    Naming, left/right ordering, features and metadata are the Tracker seam's
    shared tail (trackers.finalize) — this only produces tracks. Raises on
    failure; the chain logs it and falls back.

    SAM2 has no idea which fighter the operator marked, so `seed_a_track` is
    left None: on a tapped clip the tail recovers the subject from the taps,
    and says so honestly when it can't.

    progress_cb(frac_0_to_1, desc) — optional; drives the app's loading bar
    through the three long stages (seed scan / SAM2 propagation / crop pose)."""
    from trackers import RawTracks, TAG_SAM2

    cb = progress_cb or (lambda f, d: None)
    tmpdir = tempfile.mkdtemp(prefix="sam2_frames_")
    try:
        cb(0.01, "sampling frames from the clip")
        frame_paths, stride, fps, size = _extract_jpeg_frames(video_path, tmpdir)
        n = len(frame_paths)
        if n < pf.MIN_PAIR_FRAMES:
            raise RuntimeError(f"clip too short after sampling ({n} frames)")
        print(f"[sam2] {n} frames (stride {stride}) at {size}", flush=True)

        cb(0.04, "finding the two fighters")
        det, pose = _get_models()
        candidates = _seed_candidates(det, frame_paths)
        if not candidates:
            raise RuntimeError("no frame with two distinct people found to seed SAM2")

        predictor = _get_predictor()
        poses = cols = None
        for attempt, (seed, boxes) in enumerate(candidates, 1):
            print(f"[sam2] attempt {attempt}: seeding at frame {seed}/{n}", flush=True)
            retry = f" (retry {attempt})" if attempt > 1 else ""
            masks = {}
            done = 0
            with torch.inference_mode(), torch.autocast("cuda", dtype=_autocast_dtype()):
                # Offload frames + inference state to CPU RAM: init_state
                # otherwise loads every sampled frame onto the GPU at once —
                # ~480 full-res frames is a real OOM risk on a 16 GB T4.
                # Per-frame transfer cost is minor next to propagation.
                state = predictor.init_state(
                    tmpdir,
                    offload_video_to_cpu=True,
                    offload_state_to_cpu=True,
                )
                for oid, box in enumerate(boxes):
                    predictor.add_new_points_or_box(state, frame_idx=seed, obj_id=oid, box=box)
                for fidx, obj_ids, logits in predictor.propagate_in_video(state):
                    masks[fidx] = {int(o): (logits[i, 0] > 0).cpu().numpy()
                                   for i, o in enumerate(obj_ids)}
                    done += 1
                    if done % 10 == 0:
                        cb(0.10 + 0.40 * min(done / max(n, 1), 1.0),
                           f"tracking both fighters{retry} — frame {done}/{n}")
                for fidx, obj_ids, logits in predictor.propagate_in_video(state, reverse=True):
                    masks.setdefault(fidx, {}).update(
                        {int(o): (logits[i, 0] > 0).cpu().numpy()
                         for i, o in enumerate(obj_ids)})
                    done += 1
                    if done % 10 == 0:
                        cb(0.10 + 0.40 * min(done / max(n, 1), 1.0),
                           f"tracking both fighters{retry} — frame {done}/{n}")
            print(f"[sam2] propagation done ({len(masks)} frames with masks)", flush=True)

            poses = np.full((n, 2, 17, 2), np.nan)
            overlap = np.full(n, np.nan)     # per-frame mask overlap -> clinch metric
            cols = ([], [])
            for si, p in enumerate(frame_paths):
                if si % 10 == 0:
                    cb(0.50 + 0.50 * si / max(n, 1),
                       f"reading each fighter's pose{retry} — frame {si}/{n}")
                fm = masks.get(si)
                if not fm:
                    continue
                frame = cv2.imread(p)
                rm = {}
                for oid, m in fm.items():
                    if oid > 1 or m is None:
                        continue
                    if m.shape[:2] != frame.shape[:2]:
                        m = cv2.resize(m.astype(np.uint8), (frame.shape[1], frame.shape[0]),
                                       interpolation=cv2.INTER_NEAREST).astype(bool)
                    rm[oid] = m
                    kp = _pose_in_mask(frame, m, pose)
                    if kp is not None:
                        poses[si, oid] = kp
                        cols[oid].append(pf._torso_bgr(frame, kp))
                if 0 in rm and 1 in rm:
                    smaller = min(int(rm[0].sum()), int(rm[1].sum()))
                    if smaller > 50:
                        overlap[si] = float(np.logical_and(rm[0], rm[1]).sum()) / smaller

            if not _tracks_collapsed(poses):
                break
            print("[sam2] tracks collapsed onto one person; retrying from another seed",
                  flush=True)
            poses = None
        if poses is None:
            raise RuntimeError("every SAM2 seeding attempt tracked one person twice")

        return RawTracks(poses=poses, cols=cols, fps=fps, stride=stride,
                         tag=TAG_SAM2, frame_size=size, mask_overlap=overlap)
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)