tracking.py

"""Tracking-data import / export.

File format (CSV, one row per entity per time sample; this is what a real tracking feed is converted to):

    play_id, t, entity, side, pos, num, x, y, z
    1, 0.0, QB,   O, QB, 12, 40.0, 26.7,
    1, 0.0, CB1,  D, CB, 24, 41.1, 8.2,
    1, 0.0, ball, ,  ,   ,   40.0, 26.7, 0.0

  t      seconds; the first sample of a play is the snap (pass snap_t to shift it)
  entity any unique id per player on that play ("ball" for the ball)
  side   O (offense) or D (defense); blank for the ball
  x      yards, 0-120 end zone to end zone, the offense attacking toward LARGER x
  y      yards, 0-53.33, measured from the offense's LEFT sideline (top of the screen)
  z      ball height in yards (optional)

If the source data has the offense moving toward x=0, mirror it first (x -> 120 - x, y -> 53.33 - y).
Samples may be at any rate; the loader resamples to the viewer's 10 frames per second.
"""
from __future__ import annotations

import csv
import io
import math

from rules import FIELD_WIDTH_YD as W

FPS = 10
POS_VMAX = {"QB": 8.4, "RB": 9.8, "FB": 8.6, "WR": 10.0, "TE": 9.0, "OL": 7.5, "LT": 7.5, "LG": 7.5, "C": 7.5, "RG": 7.5,
            "RT": 7.5, "DE": 8.8, "DT": 8.0, "NT": 7.8, "OLB": 9.3, "ILB": 9.0, "MLB": 9.0, "LB": 9.1, "CB": 10.3,
            "S": 10.0, "FS": 10.0, "SS": 9.8, "NB": 10.2}
COLS = ["play_id", "t", "entity", "side", "pos", "num", "x", "y", "z"]


def to_csv(rep, play_id=1) -> str:
    """Write a replay in the tracking format (used to test the round trip and to hand clients a sample)."""
    out = io.StringIO()
    w = csv.writer(out)
    w.writerow(COLS)
    ids, meta = rep["ids"], {m["id"]: m for m in rep["meta"]}
    seen_t = set()
    for f in rep["frames"]:
        key = round(f["t"], 3)
        if key in seen_t:
            continue
        seen_t.add(key)
        for k, pid in enumerate(ids):
            m = meta[pid]
            w.writerow([play_id, f["t"], pid, m["side"], m["pos"], m["num"], f["p"][2 * k], f["p"][2 * k + 1], ""])
        w.writerow([play_id, f["t"], "ball", "", "", "", f["b"][0], f["b"][1], f["b"][2]])
    return out.getvalue()


def plays_in(text) -> list:
    rows = list(csv.DictReader(io.StringIO(text)))
    if not rows or "play_id" not in rows[0]:
        raise ValueError("Not a tracking file (no play_id column). See tracking.py for the format.")
    return sorted({r["play_id"] for r in rows})


def _interp(ts, vs, t):
    if t <= ts[0]:
        return vs[0]
    if t >= ts[-1]:
        return vs[-1]
    lo, hi = 0, len(ts) - 1
    while hi - lo > 1:
        mid = (lo + hi) // 2
        if ts[mid] <= t:
            lo = mid
        else:
            hi = mid
    f = (t - ts[lo]) / max(1e-9, ts[hi] - ts[lo])
    return vs[lo] + (vs[hi] - vs[lo]) * f


def load_csv(text, play_id=None, snap_t=None):
    """Parse one play from tracking CSV text into the replay structure the film viewer reads."""
    rows = list(csv.DictReader(io.StringIO(text)))
    missing = [c for c in ("t", "entity", "x", "y") if rows and c not in rows[0]]
    if not rows or missing:
        raise ValueError(f"Not a tracking file (missing columns: {missing or 'all rows'}). See tracking.py for the format.")
    if play_id is None:
        play_id = rows[0].get("play_id")
    rows = [r for r in rows if r.get("play_id", play_id) == play_id]
    tracks, info = {}, {}
    for r in rows:
        e = r["entity"].strip()
        tracks.setdefault(e, []).append((float(r["t"]), float(r["x"]), float(r["y"]), float(r["z"]) if r.get("z") not in (None, "") else 0.0))
        if e != "ball":
            info.setdefault(e, dict(side=(r.get("side") or "O").strip().upper()[:1] or "O", pos=(r.get("pos") or e).strip(),
                                    num=int(float(r["num"])) if r.get("num") not in (None, "") else 0))
    if not info:
        raise ValueError("No players found for that play.")
    for e in tracks:
        tracks[e].sort()
    t0 = snap_t if snap_t is not None else min(v[0][0] for v in tracks.values())
    t1 = max(v[-1][0] for v in tracks.values())
    n = int(round((t1 - t0) * FPS)) + 1
    ids = list(info)
    if sum(1 for e in ids if info[e]["side"] == "O") < 5 or sum(1 for e in ids if info[e]["side"] == "D") < 5:
        raise ValueError("Expected offense and defense players (side column O / D).")

    def series(e):
        tr = tracks[e]
        ts = [v[0] - t0 for v in tr]
        return ts, [v[1] for v in tr], [v[2] for v in tr], [v[3] for v in tr]

    S = {e: series(e) for e in ids}
    ball = series("ball") if "ball" in tracks else None
    frames = []
    for i in range(n):
        t = i / FPS
        p = []
        for e in ids:
            ts, xs, ys, _ = S[e]
            p += [round(_interp(ts, xs, t), 2), round(_interp(ts, ys, t), 2)]
        if ball:
            ts, xs, ys, zs = ball
            b = [round(_interp(ts, xs, t), 2), round(_interp(ts, ys, t), 2), round(max(0.0, _interp(ts, zs, t)), 2)]
        else:
            b = [p[0], p[1], 0.0]
        frames.append(dict(t=round(t, 2), p=p, b=b))
    # ball spot, and what the tracks say about speed and the pass
    bx, by = frames[0]["b"][0], frames[0]["b"][1]
    if not ball:
        ox = [frames[0]["p"][2 * k] for k, e in enumerate(ids) if info[e]["side"] == "O"]
        bx, by = sorted(ox)[len(ox) // 2], W / 2
    meta = []
    for k, e in enumerate(ids):
        sp = [math.hypot(frames[j + 3]["p"][2 * k] - frames[j]["p"][2 * k],
                         frames[j + 3]["p"][2 * k + 1] - frames[j]["p"][2 * k + 1]) / 0.3 for j in range(max(0, n - 3))] or [8.0]
        default = POS_VMAX.get(info[e]["pos"].upper(), 9.0)          # a short play rarely shows a player's true top speed
        vmax = round(min(11.5, max(default, max(sp) * 1.04)), 2)
        meta.append(dict(id=e, side=info[e]["side"], pos=info[e]["pos"], num=info[e]["num"], vmax=vmax,
                         mph=round(vmax * 2.045, 1), ht=None, wt=None, forty=None, react=None))
    marks = dict(throw=None, land=None, end=round(frames[-1]["t"], 2), target=None)
    events = [dict(t=0.0, text="Snap", kind="play")]
    if ball and n > 6:
        bs = [math.hypot(frames[j + 1]["b"][0] - frames[j]["b"][0], frames[j + 1]["b"][1] - frames[j]["b"][1]) * FPS for j in range(n - 1)]
        j = 0
        while j < len(bs):
            if bs[j] > 11:
                k = j
                while k < len(bs) and bs[k] > 8:
                    k += 1
                if (k - j) / FPS >= 0.4:
                    marks["throw"], marks["land"] = round(j / FPS, 2), round(k / FPS, 2)
                    fl = frames[min(k, n - 1)]
                    cand = [(math.hypot(fl["p"][2 * q] - fl["b"][0], fl["p"][2 * q + 1] - fl["b"][1]), ids[q])
                            for q in range(len(ids)) if info[ids[q]]["side"] == "O"]
                    marks["target"] = min(cand)[1]
                    events += [dict(t=marks["throw"], text="Ball thrown", kind="play"), dict(t=marks["land"], text="Ball arrives", kind="play")]
                    break
                j = k
            j += 1
    events.append(dict(t=marks["end"], text="Whistle", kind="result"))
    yards = round(frames[-1]["b"][0] - bx, 1)
    return dict(ids=ids, meta=meta, frames=frames, events=events, marks=marks, bx=round(bx, 2), by=round(by, 2),
                result=dict(outcome="recorded", yards=yards, td=False, turnover=False, penalty=None),
                summary=f"Recorded play {play_id}: ball moved {yards:+.0f} yds", source="recorded")