from __future__ import annotations

import argparse
from pathlib import Path
from typing import Optional

import numpy as np

from infer.io import load_array, load_frames_from_dir, save_image, save_sequence_grid
from infer.predictor import Predictor


def _parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Black-box SavedModel inference: output predicted frames.")
    ap.add_argument("--model_dir", type=str, default="savedmodel", help="Path to SavedModel directory")
    ap.add_argument("--frames_dir", type=str, default=None, help="Directory containing input frames (images)")
    ap.add_argument("--array", type=str, default=None, help="Path to .npy/.npz containing frames")
    ap.add_argument(
        "--pad_last_frame",
        type=str,
        default="none",
        choices=["none", "zero", "one", "repeat"],
        help="If model expects 4 frames but you provide 3, pad the last frame with: zero/one/repeat",
    )
    ap.add_argument("--out_dir", type=str, default="outputs", help="Output directory")
    ap.add_argument("--save_sequence_grid", action="store_true", help="Save a grid of the predicted sequence")
    ap.add_argument("--grid_cols", type=int, default=8, help="Columns for sequence grid")
    return ap.parse_args()


def main() -> None:
    args = _parse_args()
    if (args.frames_dir is None) == (args.array is None):
        raise SystemExit("Provide exactly one of --frames_dir or --array")

    if args.frames_dir is not None:
        frames = load_frames_from_dir(args.frames_dir)  # [T,H,W,C] 0..255
    else:
        frames = load_array(args.array)

    pred = Predictor(args.model_dir)
    seq = pred.predict_sequence(frames, pad_last_frame=args.pad_last_frame)  # [B,T,H,W,C]
    if seq.ndim == 5:
        seq0 = seq[0]
        last = seq[0, -1]
    elif seq.ndim == 4:
        seq0 = seq
        last = seq[-1]
    else:
        raise RuntimeError(f"Unexpected prediction shape: {seq.shape}")

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    save_image(out_dir / "pred_last.png", last)

    if args.save_sequence_grid:
        save_sequence_grid(out_dir / "pred_sequence_grid.png", seq0, cols=args.grid_cols)

    print(f"Wrote: {out_dir / 'pred_last.png'}")
    if args.save_sequence_grid:
        print(f"Wrote: {out_dir / 'pred_sequence_grid.png'}")


if __name__ == "__main__":
    main()