from pathlib import Path
import argparse
import statistics
import time

import numpy as np

from highlighter.agent.capabilities.image_to_scalar import motion_measurer as mm


def _build_frames(height: int, width: int, count: int) -> list[np.ndarray]:
    rng = np.random.default_rng(0)
    base = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8)
    return [np.roll(base, shift=i + 1, axis=1) for i in range(count)]


def time_motion_measurer(
    tmpdir: Path | None = None,
    frames: int = 200,
    warmup: int = 10,
    height: int = 256,
    width: int = 256,
    use_gpu: bool = True,
    blur_kernel_size: int | None = 3,
) -> dict[str, float]:
    model_path = Path(tmpdir) / "neuflow_sintel.onnx" if tmpdir else mm.DEFAULT_MODEL_PATH

    measurer = mm.MotionMeasurer(
        score_type=["sum", "mean", "median"],
        window_size=2,
        model_path=model_path,
        blur_kernel_size=blur_kernel_size,
        use_gpu=use_gpu,
    )

    total_frames = warmup + frames + 1
    frame_inputs = _build_frames(height, width, total_frames)

    frame_idx = 0
    for i in range(warmup + 1):
        measurer.update(frame_inputs[i], frame_idx)
        frame_idx += 1

    timings = []
    for i in range(warmup + 1, total_frames):
        start = time.perf_counter()
        measurer.update(frame_inputs[i], frame_idx)
        timings.append(time.perf_counter() - start)
        frame_idx += 1

    total_s = sum(timings)
    return {
        "frames": float(frames),
        "total_s": total_s,
        "avg_ms": (total_s / frames) * 1000.0,
        "p50_ms": statistics.median(timings) * 1000.0,
        "p95_ms": float(np.percentile(timings, 95)) * 1000.0,
    }


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Time MotionMeasurer.update performance.")
    parser.add_argument("--frames", type=int, default=200)
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--height", type=int, default=256)
    parser.add_argument("--width", type=int, default=256)
    parser.add_argument("--cpu", action="store_true", help="Force CPU inference.")
    parser.add_argument(
        "--blur-kernel-size",
        type=int,
        default=3,
        help="Odd kernel size (3, 5, 7). Use 0 to disable blur.",
    )
    parser.add_argument("--tmpdir", type=Path, default=None)
    return parser.parse_args()


def main() -> None:
    args = _parse_args()
    blur_kernel_size = None if args.blur_kernel_size == 0 else args.blur_kernel_size
    stats = time_motion_measurer(
        tmpdir=args.tmpdir,
        frames=args.frames,
        warmup=args.warmup,
        height=args.height,
        width=args.width,
        use_gpu=not args.cpu,
        blur_kernel_size=blur_kernel_size,
    )
    print(
        "MotionMeasurer timing: "
        f"frames={int(stats['frames'])}, total={stats['total_s']:.3f}s, "
        f"avg={stats['avg_ms']:.2f}ms, p50={stats['p50_ms']:.2f}ms, "
        f"p95={stats['p95_ms']:.2f}ms"
    )


if __name__ == "__main__":
    main()
