#!/usr/bin/env python3
"""Create a small ScanNet motion-blur benchmark from decoded ScanNet frames.

The script creates three synchronized views of the same short sequence:

1. sharp RGB-D sequence for GauS-SLAM
2. blurred RGB-D sequence for GauS-SLAM
3. VD-Diff style test folders: test/blur/<scene>, test/sharp/<scene>

Blur is synthesized by averaging a temporal window of neighboring RGB frames.
Depth, pose, and intrinsics are copied unchanged so the SLAM drop measures the
effect of color blur in a controlled prototype.
"""

from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import numpy as np
from PIL import Image


def list_images(path: Path) -> list[Path]:
    return sorted(path.glob("*.jpg")) + sorted(path.glob("*.png"))


def load_rgb(path: Path) -> np.ndarray:
    return np.asarray(Image.open(path).convert("RGB"), dtype=np.float32)


def save_rgb(array: np.ndarray, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(np.clip(array, 0, 255).astype(np.uint8)).save(path, quality=95)


def copy_tree(src: Path, dst: Path) -> None:
    if dst.exists():
        shutil.rmtree(dst)
    shutil.copytree(src, dst)


def copy_file(src: Path, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src, dst)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--scene", default="scene0004_00")
    parser.add_argument("--src-root", type=Path, default=Path("/home/szha0669/storage/blur_slam_exp/data/scannet"))
    parser.add_argument("--out-root", type=Path, default=Path("/home/szha0669/storage/blur_slam_exp/data/scannet_blur_proto"))
    parser.add_argument("--start", type=int, default=0)
    parser.add_argument("--end", type=int, default=120)
    parser.add_argument("--stride", type=int, default=2)
    parser.add_argument("--blur-radius", type=int, default=3)
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()

    if args.blur_radius < 1:
        raise ValueError("--blur-radius must be >= 1")

    src_scene = args.src_root / args.scene
    color_paths = list_images(src_scene / "color")
    if not color_paths:
        raise FileNotFoundError(f"No color frames found in {src_scene / 'color'}")

    selected = list(range(args.start, min(args.end, len(color_paths)), args.stride))
    if len(selected) < 2 * args.blur_radius + 1:
        raise ValueError("Selected sequence is too short for the requested blur radius")

    sharp_scene = args.out_root / "gaus_sharp" / args.scene
    blur_scene = args.out_root / "gaus_blur" / args.scene
    vddiff_blur = args.out_root / "vddiff" / "test" / "blur" / args.scene
    vddiff_sharp = args.out_root / "vddiff" / "test" / "sharp" / args.scene

    if args.overwrite:
        for path in [sharp_scene, blur_scene, vddiff_blur, vddiff_sharp]:
            if path.exists():
                shutil.rmtree(path)

    for base in [sharp_scene, blur_scene]:
        (base / "color").mkdir(parents=True, exist_ok=True)
        (base / "depth").mkdir(parents=True, exist_ok=True)
        (base / "pose").mkdir(parents=True, exist_ok=True)
        copy_tree(src_scene / "intrinsic", base / "intrinsic")

    vddiff_blur.mkdir(parents=True, exist_ok=True)
    vddiff_sharp.mkdir(parents=True, exist_ok=True)

    src_depth_paths = sorted((src_scene / "depth").glob("*.png"))
    src_pose_paths = sorted((src_scene / "pose").glob("*.txt"))
    index_map = {}

    for out_idx, src_idx in enumerate(selected):
        out_name = f"{out_idx:06d}"
        src_color = color_paths[src_idx]
        sharp = load_rgb(src_color)

        lo = max(0, src_idx - args.blur_radius)
        hi = min(len(color_paths), src_idx + args.blur_radius + 1)
        stack = np.stack([load_rgb(color_paths[i]) for i in range(lo, hi)], axis=0)
        blurred = stack.mean(axis=0)

        save_rgb(sharp, sharp_scene / "color" / f"{out_name}.jpg")
        save_rgb(blurred, blur_scene / "color" / f"{out_name}.jpg")
        save_rgb(sharp, vddiff_sharp / f"{out_name}.png")
        save_rgb(blurred, vddiff_blur / f"{out_name}.png")

        copy_file(src_depth_paths[src_idx], sharp_scene / "depth" / f"{out_name}.png")
        copy_file(src_depth_paths[src_idx], blur_scene / "depth" / f"{out_name}.png")
        copy_file(src_pose_paths[src_idx], sharp_scene / "pose" / f"{out_name}.txt")
        copy_file(src_pose_paths[src_idx], blur_scene / "pose" / f"{out_name}.txt")
        index_map[out_name] = int(src_idx)

    manifest = {
        "scene": args.scene,
        "source": str(src_scene),
        "frames": len(selected),
        "start": args.start,
        "end": args.end,
        "stride": args.stride,
        "blur_radius": args.blur_radius,
        "sharp_gaus": str(sharp_scene),
        "blur_gaus": str(blur_scene),
        "vddiff_blur": str(vddiff_blur),
        "vddiff_sharp": str(vddiff_sharp),
        "index_map": index_map,
    }
    args.out_root.mkdir(parents=True, exist_ok=True)
    (args.out_root / f"{args.scene}_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()