from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np

from infer.io import load_array, load_frames_from_dir
from infer.predictor import Predictor
from infer.anomaly import compute_anomaly_from_4frames


def _parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Compute per-window anomaly scores (ahat_error and 1-SSIM vectors).")
    ap.add_argument("--model_dir", type=str, default="savedmodel")
    ap.add_argument("--frames_dir", type=str, default=None)
    ap.add_argument("--array", type=str, default=None)
    ap.add_argument("--out_json", type=str, default="anomaly.json")
    ap.add_argument("--out_csv", type=str, default="anomaly.csv")
    ap.add_argument("--start", type=int, default=0, help="Start frame index (inclusive) on the full sequence")
    ap.add_argument("--end", type=int, default=-1, help="End frame index (exclusive); -1 means to the end")
    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)

    if frames.ndim == 5:
        if frames.shape[0] != 1:
            raise SystemExit("Batch arrays are not supported in CLI; please provide a single sequence.")
        frames = frames[0]
    if frames.ndim != 4 or frames.shape[0] < 4:
        raise SystemExit(f"Expected frames with shape [T,H,W,C] and T>=4, got {frames.shape}")

    T = frames.shape[0]
    start = max(0, int(args.start))
    end = T if int(args.end) == -1 else min(T, int(args.end))
    if end - start < 4:
        raise SystemExit("Range too small; need at least 4 frames.")

    pred = Predictor(args.model_dir)
    if pred.sequence_len is not None and pred.sequence_len != 4:
        raise SystemExit(f"This tool assumes model sequence_len=4, got {pred.sequence_len}")

    rows: List[Dict[str, Any]] = []
    for t in range(start + 3, end):
        win = frames[t - 3 : t + 1]  # [4,H,W,C] where win[3] is GT frame4
        res = compute_anomaly_from_4frames(pred, win)
        rows.append(
            {
                "t": int(t),
                "ahat_error": res.ahat_error.tolist(),
                "dissim_1mssim": res.dissim_1mssim.tolist(),
            }
        )

    out_json = Path(args.out_json)
    out_json.parent.mkdir(parents=True, exist_ok=True)
    out_json.write_text(json.dumps({"items": rows}, ensure_ascii=False, indent=2), encoding="utf-8")

    out_csv = Path(args.out_csv)
    out_csv.parent.mkdir(parents=True, exist_ok=True)
    with out_csv.open("w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["t", "ahat_error_1", "ahat_error_2", "ahat_error_3", "ahat_error_4", "dissim_1", "dissim_2", "dissim_3", "dissim_4"])
        for r in rows:
            a = r["ahat_error"]
            d = r["dissim_1mssim"]
            w.writerow([r["t"], *a, *d])

    print(f"Wrote: {out_json}")
    print(f"Wrote: {out_csv}")


if __name__ == "__main__":
    main()