#!/usr/bin/env python3
"""Generate normalized depth assets for the cut-particle 2.5D study."""

from __future__ import annotations

import argparse
import dataclasses
import json
import time
from pathlib import Path

import cv2
import numpy as np
import torch
from PIL import Image


def normalize_depth(depth: np.ndarray) -> tuple[np.ndarray, dict[str, float]]:
    depth = np.asarray(depth, dtype=np.float32)
    finite = np.isfinite(depth) & (depth > 0)
    if not finite.any():
        raise RuntimeError("The model returned no finite positive depth values.")

    values = depth[finite]
    low = float(np.percentile(values, 1.0))
    high = float(np.percentile(values, 99.0))
    if high <= low:
        high = float(values.max())
        low = float(values.min())
    normalized = np.clip((depth - low) / max(high - low, 1e-8), 0, 1)
    normalized[~finite] = 1
    return normalized, {
        "raw_min": float(values.min()),
        "raw_max": float(values.max()),
        "percentile_1": low,
        "percentile_99": high,
    }


def save_outputs(
    depth: np.ndarray,
    input_path: Path,
    output_dir: Path,
    name: str,
    metadata: dict[str, object],
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    width, height = Image.open(input_path).size
    resized = cv2.resize(depth, (width, height), interpolation=cv2.INTER_CUBIC)
    normalized, stats = normalize_depth(resized)

    # Encoding: 0 = near, 65535 = far. The browser preview uses the 8-bit
    # companion; the 16-bit file remains the lossless integration source.
    depth_u16 = np.round(normalized * 65535).astype(np.uint16)
    depth_u8 = np.round(normalized * 255).astype(np.uint8)
    near_u8 = 255 - depth_u8

    cv2.imwrite(str(output_dir / f"depth-{name}-16.png"), depth_u16)
    cv2.imwrite(str(output_dir / f"depth-{name}.png"), depth_u8)
    cv2.imwrite(
        str(output_dir / f"depth-{name}-preview.png"),
        cv2.applyColorMap(near_u8, cv2.COLORMAP_TURBO),
    )
    np.save(output_dir / f"depth-{name}.npy", resized.astype(np.float32))

    payload = {
        **metadata,
        **stats,
        "source": str(input_path),
        "width": width,
        "height": height,
        "encoding": "0-near_1-far",
    }
    (output_dir / f"depth-{name}.json").write_text(
        json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
    )


def run_da3(args: argparse.Namespace) -> None:
    from depth_anything_3.api import DepthAnything3

    device = torch.device(args.device)
    started = time.perf_counter()
    model = DepthAnything3.from_pretrained(str(args.da3_checkpoint)).to(device)
    model.eval()
    prediction = model.inference(
        [str(args.input)],
        process_res=args.process_res,
        process_res_method="upper_bound_resize",
    )
    depth = prediction.depth[0]
    elapsed = time.perf_counter() - started
    save_outputs(
        depth,
        args.input,
        args.output_dir,
        "da3mono",
        {
            "model": "depth-anything/DA3MONO-LARGE",
            "device": str(device),
            "process_res": args.process_res,
            "elapsed_seconds": round(elapsed, 3),
        },
    )


def run_depth_pro(args: argparse.Namespace) -> None:
    import depth_pro
    from depth_pro.depth_pro import DEFAULT_MONODEPTH_CONFIG_DICT

    device = torch.device(args.device)
    precision = torch.float16 if device.type == "cuda" else torch.float32
    config = dataclasses.replace(
        DEFAULT_MONODEPTH_CONFIG_DICT,
        checkpoint_uri=str(args.depth_pro_checkpoint),
    )
    started = time.perf_counter()
    model, transform = depth_pro.create_model_and_transforms(
        config=config,
        device=device,
        precision=precision,
    )
    model.eval()
    image, _, focal_px = depth_pro.load_rgb(args.input)
    prediction = model.infer(transform(image), f_px=focal_px)
    depth = prediction["depth"].detach().float().cpu().numpy().squeeze()
    elapsed = time.perf_counter() - started
    estimated_focal = prediction.get("focallength_px")
    save_outputs(
        depth,
        args.input,
        args.output_dir,
        "depthpro",
        {
            "model": "apple/ml-depth-pro",
            "device": str(device),
            "precision": str(precision),
            "focal_px": (
                float(estimated_focal.detach().cpu().item())
                if estimated_focal is not None
                else focal_px
            ),
            "elapsed_seconds": round(elapsed, 3),
        },
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("model", choices=("da3", "depthpro"))
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--process-res", type=int, default=756)
    parser.add_argument(
        "--da3-checkpoint",
        type=Path,
        default=Path(".runtime/depth-eval/checkpoints/da3mono-large"),
    )
    parser.add_argument(
        "--depth-pro-checkpoint",
        type=Path,
        default=Path(".runtime/depth-eval/ml-depth-pro/checkpoints/depth_pro.pt"),
    )
    return parser.parse_args()


if __name__ == "__main__":
    arguments = parse_args()
    if arguments.model == "da3":
        run_da3(arguments)
    else:
        run_depth_pro(arguments)