#!/usr/bin/env python3
"""
movienet_prep.py — extract probe crops from MovieNet-PS for actors in our gallery.

Usage:
  python scripts/movienet_prep.py \
    --movienet <movienet_root> \
    --gallery  gallery.h5 \
    --output   eval/ \
    [--split   Train_app10] \
    [--margin  0.2] \
    [--max-per-actor 50]

MovieNet-PS format (annotation.zip + Image.zip):
  annotation/test/train_test/Train_app<N>.mat  — N annotations per actor
    Train[i] = [imdb_id (nm...), count, [[img_path, bbox[x,y,w,h], label], ...]]
  Image/<movie_tt_id>/shot_XXXX_img_Y.jpg      — source frames

Output:
  eval/probe/         face crops (jpg)
  eval/gt.json        [{crop, imdb_id, actor_name, source_frame}]
"""

import argparse
import json
import sys
import zipfile
from io import BytesIO
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from sae_gallery import load_gallery_hdf5  # noqa: E402

try:
    import cv2
    import numpy as np
    import scipy.io as sio
except ImportError as e:
    print(f"[error] missing dependency: {e}", file=sys.stderr)
    print("Install: pip install opencv-python scipy numpy", file=sys.stderr)
    sys.exit(1)


# ── MovieNet-PS loader ────────────────────────────────────────────────────────

def load_movienet_annotations(movienet_root: Path, split: str) -> list[dict]:
    """
    Parse a MovieNet-PS Train_app<N>.mat split.
    Returns flat list of {"imdb_id", "img_path", "bbox": [x,y,w,h]}.
    img_path is relative to Image/ inside Image.zip, e.g. tt0047396/shot_0004_img_1.jpg
    """
    mat_path = movienet_root / "annotation" / "test" / "train_test" / f"{split}.mat"
    if not mat_path.exists():
        # try extracting from annotation.zip
        zip_path = movienet_root / "annotation.zip"
        if not zip_path.exists():
            raise FileNotFoundError(f"annotation.zip not found in {movienet_root}")
        inner = f"annotation/test/train_test/{split}.mat"
        print(f"[prep] extracting {inner} from annotation.zip…", file=sys.stderr)
        with zipfile.ZipFile(zip_path) as z:
            z.extract(inner, movienet_root)
        mat_path = movienet_root / inner

    data = sio.loadmat(str(mat_path))["Train"]
    annotations = []
    for row in data:
        imdb_id  = str(row[0].flat[0])        # e.g. "nm0000023"
        entries  = row[2]                       # array of [path, bbox, label]
        for entry in entries:
            img_path = str(entry[0].flat[0])   # e.g. "tt0032138/shot_0003_img_1.jpg"
            bbox     = [float(v) for v in entry[1].flat]  # [x, y, w, h]
            annotations.append({"imdb_id": imdb_id, "img_path": img_path, "bbox": bbox})
    return annotations


# ── Gallery loader ────────────────────────────────────────────────────────────

def load_gallery_ids(gallery_path: str) -> dict[str, str]:
    """Return {imdb_id: actor_name} for all actors in the gallery."""
    data = load_gallery_hdf5(Path(gallery_path))
    return {a["imdb_id"]: a["name"] for a in data["actors"]}


# ── Crop + save ───────────────────────────────────────────────────────────────

def crop_face(img: "np.ndarray", bbox: list[float], margin: float) -> "np.ndarray | None":
    h, w = img.shape[:2]
    x, y, bw, bh = bbox
    # expand by margin
    pad_x = bw * margin
    pad_y = bh * margin
    x1 = max(0, int(x - pad_x))
    y1 = max(0, int(y - pad_y))
    x2 = min(w, int(x + bw + pad_x))
    y2 = min(h, int(y + bh + pad_y))
    crop = img[y1:y2, x1:x2]
    return crop if crop.size > 0 else None


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--movienet",       required=True,  help="MovieNet-PS root directory")
    p.add_argument("--gallery",        required=True,  help="gallery.h5 (for actor list)")
    p.add_argument("--output",         default="eval", help="output directory")
    p.add_argument("--split",          default="Train_app10",
                   help="annotation split to use (default: Train_app10)")
    p.add_argument("--margin",         type=float, default=0.2,
                   help="bbox expansion factor (default 0.2 = 20%%)")
    p.add_argument("--max-per-actor",  type=int,   default=50,
                   help="cap probe crops per actor (default 50)")
    args = p.parse_args()

    movienet_root = Path(args.movienet)
    out_dir = Path(args.output)
    probe_dir = out_dir / "probe"
    probe_dir.mkdir(parents=True, exist_ok=True)

    gallery_ids = load_gallery_ids(args.gallery)
    print(f"[prep] gallery actors: {len(gallery_ids)}", file=sys.stderr)

    annotations = load_movienet_annotations(movienet_root, args.split)
    print(f"[prep] total annotations in split: {len(annotations)}", file=sys.stderr)

    matched = [a for a in annotations if a["imdb_id"] in gallery_ids]
    print(f"[prep] annotations matching gallery: {len(matched)}", file=sys.stderr)

    if not matched:
        print("[error] no overlap between MovieNet and gallery — check IMDb ID format", file=sys.stderr)
        sys.exit(1)

    # Build set of image paths we actually need, then extract from Image.zip in one pass
    needed_paths = {a["img_path"] for a in matched}
    image_zip = movienet_root / "Image.zip"
    frame_cache: dict[str, np.ndarray] = {}

    n_missing_in_zip = 0
    print(f"[prep] extracting {len(needed_paths)} frames from Image.zip…", file=sys.stderr)
    with zipfile.ZipFile(image_zip) as zf:
        for img_path in needed_paths:
            zip_entry = f"Image/{img_path}"
            try:
                data = zf.read(zip_entry)
            except KeyError:
                n_missing_in_zip += 1  # frame referenced by an annotation but absent from Image.zip
                continue
            arr = np.frombuffer(data, dtype=np.uint8)
            img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
            if img is not None:
                frame_cache[img_path] = img

    print(f"[prep] frames loaded: {len(frame_cache)}/{len(needed_paths)}", file=sys.stderr)
    if n_missing_in_zip:
        print(f"[prep] frames absent from Image.zip: {n_missing_in_zip}", file=sys.stderr)

    per_actor_count: dict[str, int] = {}
    gt_entries = []
    n_failed = 0

    for ann in matched:
        imdb_id  = ann["imdb_id"]
        img_path = ann["img_path"]
        count    = per_actor_count.get(imdb_id, 0)
        if count >= args.max_per_actor:
            continue

        img = frame_cache.get(img_path)
        if img is None:
            n_failed += 1
            continue

        crop = crop_face(img, ann["bbox"], args.margin)
        if crop is None:
            n_failed += 1
            continue

        crop_name = f"{imdb_id}_{count:04d}.jpg"
        crop_path = probe_dir / crop_name
        cv2.imwrite(str(crop_path), crop)

        per_actor_count[imdb_id] = count + 1
        gt_entries.append({
            "crop":         str(crop_path),
            "imdb_id":      imdb_id,
            "actor_name":   gallery_ids[imdb_id],
            "source_frame": img_path,
        })

    gt_path = out_dir / "gt.json"
    with open(gt_path, "w") as f:
        json.dump(gt_entries, f, indent=2)

    print(f"[prep] crops saved:    {len(gt_entries)}", file=sys.stderr)
    print(f"[prep] crop failures:  {n_failed}", file=sys.stderr)
    print(f"[prep] actors covered: {len(per_actor_count)}/{len(gallery_ids)}", file=sys.stderr)
    for imdb_id, name in sorted(gallery_ids.items()):
        n = per_actor_count.get(imdb_id, 0)
        status = f"{n} crops" if n else "NO MATCH"
        print(f"  {name:30s} {status}", file=sys.stderr)
    print(f"[prep] gt.json → {gt_path}", file=sys.stderr)


if __name__ == "__main__":
    main()
