"""A Hugging Face Space that returns PaGeR's *raw metric depth*.

The official demo (`prs-eth/PaGeR`) renders colour-mapped previews and GLB point
clouds, and never exposes the underlying array. Placement needs metres, so this
Space does the minimum: one forward pass, one `.npy` of float32 metric depth at
panoramic resolution, plus a preview and a few statistics for sanity.

It also deliberately skips the point-cloud path, which is most of the original
demo's GPU time. On a free account ZeroGPU allows 5 minutes of GPU per day, so
the difference decides whether a run of four panoramas fits in the budget.

The call sequence below mirrors the official Space's `_run_inference`, because
the model card's abbreviated snippet is wrong in ways that fail loudly (numpy
where a tensor is wanted) and, worse, in one way that would not have:

  * the backbone wants input centred to **[-1, 1]**, not [0, 1];
  * `Pager` stores whatever `device` it is given and later reads
    `self.device.type`, so it needs a `torch.device`, not the string "cuda";
  * the model must be cast to `pager.weight_dtype`, not merely moved;
  * `prepare_depth_for_logging` *calls* its cmap argument — it takes a colormap
    object — and needs `log_scale=pred["scale"]` to produce metric output;
  * the metric scale comes from one of two heads, indoor or outdoor, selected
    per scene via `skip_heads`.

DEPLOYING
    1. Create a Space, SDK "Gradio", hardware "ZeroGPU". A free account in good
       standing may host two.
    2. Upload this file, `requirements.txt` and `README.md`.
    3. `python scripts/push_pager_space.py --name pager-raw --yes` does all of
       that except the hardware switch, which needs the *Manage Spaces* scope.

Weights are CC BY-NC 4.0, inherited from the Depth Anything 3 backbone:
academic and non-commercial use only.
"""

from __future__ import annotations

import os
import subprocess
import sys
import tempfile
import traceback

import matplotlib
import numpy as np
import spaces
import torch
import gradio as gr
from PIL import Image

# PaGeR imports live under a top-level `src` package, so the repo root has to be
# importable. requirements.txt installs it; the clone is a fallback because that
# layout does not always survive a wheel build.
try:
    from src.pager import Pager
except ModuleNotFoundError:                                   # pragma: no cover
    REPO = os.path.join(os.path.dirname(__file__), "PaGeR")
    if not os.path.isdir(REPO):
        subprocess.run(
            ["git", "clone", "--depth", "1", "https://github.com/prs-eth/PaGeR", REPO],
            check=True,
        )
    sys.path.insert(0, REPO)
    from src.pager import Pager

from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf

from src.utils.geometry_utils import erp_to_cubemap
from src.utils.utils import prepare_depth_for_logging

CHECKPOINT = os.environ.get("PAGER_CHECKPOINT", "prs-eth/PaGeR-metric-depth")
MAX_WIDTH = 2048          # our panoramas are 2K native; PaGeR accepts up to 3K
SCENE_TO_SCALE_HEAD = {"Indoor": "scale_indoor", "Outdoor": "scale_outdoor"}
CMAP = matplotlib.colormaps["Spectral"]

# A torch.device, not the string "cuda". ZeroGPU also wants the model placed at
# module import rather than lazily inside the GPU function.
device = torch.device("cuda")
_cfg = OmegaConf.load(hf_hub_download(repo_id=CHECKPOINT, filename="config.yaml"))
_face = int(_cfg.face_size)
_fov = float(getattr(_cfg, "cube_fov", 90.0))

# Only the unified checkpoint ships both scale heads. With a single-head
# checkpoint `skip_heads` has nothing to choose between, so it is omitted rather
# than passed empty — and the scene selector is disabled to say so honestly.
_modalities = set(getattr(_cfg, "modalities", []) or [])
HAS_BOTH_HEADS = {"scale_indoor", "scale_outdoor"} <= _modalities

pager = Pager(CHECKPOINT, cfg=_cfg, device=device)
pager.get_intrinsics_extrinsics(image_size=_face, fov=_fov)
pager.model.to(device, dtype=getattr(pager, "weight_dtype", torch.float32))
pager.model.eval()

print(f"[pager-raw] checkpoint={CHECKPOINT} face={_face} fov={_fov} "
      f"modalities={sorted(_modalities)} dual_heads={HAS_BOTH_HEADS}")


@spaces.GPU(duration=90)
def predict(image_path: str, scene: str = "Indoor"):
    """Equirectangular RGB in, raw metric depth out."""
    if image_path is None:
        raise gr.Error("Upload an equirectangular panorama first.")

    img = Image.open(image_path).convert("RGB")
    if img.width > MAX_WIDTH:
        img = img.resize((MAX_WIDTH, MAX_WIDTH // 2), Image.LANCZOS)
    if img.width != 2 * img.height:
        raise gr.Error(
            f"Expected a 2:1 equirectangular panorama, got {img.width}x{img.height}."
        )

    rgb = np.asarray(img, dtype=np.uint8)
    height, width = rgb.shape[:2]

    try:
        erp = torch.from_numpy(rgb).permute(2, 0, 1).to(torch.float32) / 255.0
        erp = erp * 2.0 - 1.0                    # centred input for the backbone
        cubemap = erp_to_cubemap(erp, face_w=_face, fov=_fov).unsqueeze(0).to(device)

        kwargs = {}
        if HAS_BOTH_HEADS:
            active = SCENE_TO_SCALE_HEAD[scene]
            kwargs["skip_heads"] = {h for h in SCENE_TO_SCALE_HEAD.values()
                                    if h != active}

        with torch.inference_mode():
            pred = pager(cubemap, dtype=torch.float16, **kwargs)

        depth, _viz = prepare_depth_for_logging(
            pager,
            pred["depth"][0],
            pred["sky"][0] if "sky" in pred else None,
            (height, width),
            CMAP,
            log_scale=pred.get("scale", None),
        )
    except Exception:
        # Surfaced to the caller rather than buried in the Space log, so a
        # remote client can diagnose without log access.
        raise gr.Error("PaGeR inference failed:\n" + traceback.format_exc()[-1500:])

    depth = np.squeeze(np.asarray(depth, dtype=np.float32))
    if depth.shape != (height, width):
        raise gr.Error(f"depth is {depth.shape}, expected {(height, width)}")

    path = os.path.join(tempfile.mkdtemp(), "pager_depth.npy")
    np.save(path, depth)

    finite = depth[np.isfinite(depth) & (depth > 0)]
    head = SCENE_TO_SCALE_HEAD[scene] if HAS_BOTH_HEADS else "single-head checkpoint"
    stats = (
        f"shape {depth.shape[0]}x{depth.shape[1]}  ·  metres  ·  {head}\n"
        f"min {finite.min():.3f}   p05 {np.percentile(finite, 5):.3f}   "
        f"median {np.median(finite):.3f}   "
        f"p95 {np.percentile(finite, 95):.3f}   max {finite.max():.3f}\n"
        f"non-finite or non-positive: {depth.size - finite.size} px"
    )

    lo, hi = np.percentile(finite, [2, 98]) if finite.size else (0.0, 1.0)
    norm = np.clip((depth - lo) / max(hi - lo, 1e-6), 0, 1)
    preview = (CMAP(1.0 - norm)[..., :3] * 255).astype(np.uint8)
    return path, preview, stats


with gr.Blocks(title="PaGeR raw metric depth") as demo:
    gr.Markdown(
        "## PaGeR — raw metric depth\n"
        "Returns the float32 metric depth array as `.npy`, not a colour map. "
        "Upload a 2:1 equirectangular panorama.\n\n"
        f"Model: `{CHECKPOINT}` (CC BY-NC 4.0, non-commercial)."
    )
    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="filepath", label="Equirectangular panorama")
            mode = gr.Radio(
                list(SCENE_TO_SCALE_HEAD), value="Indoor", label="Metric scale head",
                interactive=HAS_BOTH_HEADS,
                info=("Indoor and outdoor scales come from different heads."
                      if HAS_BOTH_HEADS else
                      "This checkpoint ships a single scale head; ignored."),
            )
            run = gr.Button("Estimate depth", variant="primary")
        with gr.Column():
            out_file = gr.File(label="Metric depth (.npy, float32, metres)")
            out_prev = gr.Image(label="Preview")
            out_txt = gr.Textbox(label="Statistics", lines=3)

    # A stable api_name keeps scripts/fetch_pager_depth.py working when the UI
    # is rearranged.
    run.click(predict, inputs=[inp, mode], outputs=[out_file, out_prev, out_txt],
              api_name="predict")

demo.queue(max_size=8).launch(show_error=True)