#!/usr/bin/env python3
"""Minimal example: onboard karting video -> speed + lap times.

    python example.py my_onboard_clip.mp4
"""
import sys
import numpy as np
import torch
from modeling_kartnet import KartNet, extract_features, predict

video = sys.argv[1]
dev = 'cuda' if torch.cuda.is_available() else 'cpu'
model = KartNet()
model.load_state_dict(torch.load('kartnet_v3.pt', map_location=dev))

print('extracting features ...')
feats = extract_features(video)
out = predict(model, feats, device=dev)

print(f"frames: {len(out['t'])} ({out['t'][-1]:.0f}s)")
print(f"speed:  median {np.median(out['speed_mps'])*3.6:.1f} km/h, "
      f"p95 {np.percentile(out['speed_mps'],95)*3.6:.1f} km/h")
if len(out['lap_times_s']):
    print('lap times:', ' '.join(f"{x:.2f}" for x in out['lap_times_s']))
np.savez(video.rsplit('.', 1)[0] + '_kartnet.npz', **out)
print('saved', video.rsplit('.', 1)[0] + '_kartnet.npz')