#!/usr/bin/env python
"""
HSMR smoke test on Apple Silicon (MPS) — Path B engine verification.

Answers the Lane B smoke-test questions on HSMR's bundled demo clips (no John footage,
no SKEL license needed beyond the model files John downloaded):
  1. Does the HSMR forward run on MPS?  (device + CPU-fallback)
  2. Timing (target < 2x realtime).
  3. Pose sanity (shape, NaNs) + an overlay render to eyeball quality.

Detection uses the torchvision person detector (lib.modeling.pipelines.vitdet, the
detectron2-free drop-in). The one remaining gate is SKEL: build_inference_pipeline
instantiates the SKEL body model, so this completes only once John's SKEL files are in place.

Run (once SKEL files are in data_inputs/body_models/skel/):
  cd ~/Documents/MOCAPGOD/.engine/HSMR && source .venv/bin/activate
  PYTHONPATH=. python smoke_test.py -i data_inputs/demo/example_videos/gymnasts.mp4
  PYTHONPATH=. python smoke_test.py -i data_inputs/demo/example_imgs   # image folder
"""
import os, time, argparse
os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')  # some ops lack MPS kernels

import _headless  # noqa: F401  fake pyrender/OpenGL for headless macOS — MUST precede lib.* imports

# One import line: every helper is an attribute of hsmr_demo (via its star-imports).
from lib.kits.hsmr_demo import (
    load_inputs, imgs_det2patches, build_inference_pipeline, prepare_mesh,
    visualize_full_img, IMG_MEAN_255, IMG_STD_255, asb, assemble_dict,
    get_logger, save_video, save_img, np, torch, Path, DEFAULT_HSMR_ROOT,
    build_detector,
)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('-i', '--input_path', required=True)
    ap.add_argument('-o', '--output_path', default='data_outputs/smoke')
    ap.add_argument('-d', '--device', default='mps')
    ap.add_argument('-m', '--model_root', default=None)
    ap.add_argument('--rec_bs', type=int, default=64)
    ap.add_argument('--ignore_skel', action='store_true')
    args = ap.parse_args()

    log = get_logger(brief=True)
    model_root = Path(args.model_root) if args.model_root else DEFAULT_HSMR_ROOT
    out = Path(args.output_path); out.mkdir(parents=True, exist_ok=True)

    # Inputs (reuse the demo loader; it only reads .input_path / .input_type).
    fake = argparse.Namespace(input_path=args.input_path, input_type='auto')
    raw_imgs, meta = load_inputs(fake)
    log.info(f'device={args.device}  type={meta["type"]}  frames={len(raw_imgs)}')

    # Detection -> patches (real torchvision person detector, on CPU).
    dets, ratios = build_detector(device=args.device)(raw_imgs)
    patches, det_meta = imgs_det2patches(raw_imgs, dets, ratios, 5)
    log.info(f'patches={len(patches)}')

    # Build pipeline — instantiates SKEL, so this is where a missing SKEL file fails.
    pipe = build_inference_pipeline(model_root=model_root, device=args.device)

    # Recovery loop (mirrors run_demo), timed.
    t0 = time.time()
    pd_params, pd_cam_t = [], []
    for bw in asb(total=len(patches), bs_scope=args.rec_bs, enable_tqdm=True):
        p = patches[bw.sid:bw.eid]
        pn = ((p - IMG_MEAN_255) / IMG_STD_255).transpose(0, 3, 1, 2)
        with torch.no_grad():
            o = pipe(pn)
        pd_params.append({k: v.detach().cpu().clone() for k, v in o['pd_params'].items()})
        pd_cam_t.append(o['pd_cam_t'].detach().cpu().clone())
    dt = time.time() - t0
    pd_params = assemble_dict(pd_params, expand_dim=False)
    pd_cam_t = torch.cat(pd_cam_t, dim=0)

    poses = pd_params['poses']
    fps = len(patches) / dt if dt else 0
    log.info(f'RECOVERY ok: {len(patches)} patches in {dt:.2f}s ({fps:.2f} fps)  '
             f'poses={tuple(poses.shape)}  nan={bool(torch.isnan(poses).any())}')

    # Dump raw params regardless (this is the Path-B payload: SKEL rotations).
    np.savez(out / 'params.npz', poses=poses.numpy(),
             betas=pd_params['betas'].numpy(), cam_t=pd_cam_t.numpy())
    log.info(f'params -> {out/"params.npz"}')

    # Mesh geometry for visual QC in Blender (pyrender/EGL is dead headless on macOS, and
    # Blender is the fleet renderer + Lane C's tool). prepare_mesh uses skel_model, no GL.
    try:
        m_skin, m_skel = prepare_mesh(pipe, pd_params)
        np.savez(out / 'mesh_qc.npz',
                 skin_v=m_skin['v'].detach().cpu().numpy(), skin_f=np.asarray(m_skin['f']),
                 skel_v=m_skel['v'].detach().cpu().numpy(), skel_f=np.asarray(m_skel['f']))
        log.info(f'mesh geometry -> {out/"mesh_qc.npz"} (import in Blender to eyeball)')
    except Exception as e:
        log.warning(f'prepare_mesh failed (params still saved): {e!r}')

    log.info('smoke test complete.')


if __name__ == '__main__':
    main()