"""Depth-model bake-off — runs ON a RunPod GPU pod (see launch_depth_bakeoff.py).

Candidates (licenses verified 2026-07-28, all product-viable):
  - Depth Anything V2 Small (Apache-2.0, relative)     via transformers
  - MoGe-2 ViT-S / ViT-B (+normals) (MIT, metric)      via moge
  - Metric3D v2 ViT-S / ViT-L (BSD-2, metric)          via torch.hub

For each of 8 living-room rig views (floor-rich / seam / wall-starved / ceiling):
predict depth, robust-fit to the SfM sparse depths (scale+shift on disparity for
relative models, scale-only for metric ones), report median absolute relative
residual + colorized depth maps. Results zip -> Gitea; DONE marker signals driver.
"""
import io
import json
import os
import struct
import time
import traceback
import urllib.request
import zipfile
from pathlib import Path

import numpy as np

WORK = Path("/tmp/bake")
RIG_URL = os.environ["RIG_URL"]
SFM_URL = os.environ["SFM_URL"]
OUT_BASE = os.environ["OUT_BASE"]  # gitea generic package base for results
GITEA_TOKEN = os.environ["GITEA_TOKEN"]
RINGS = ["pano_camera00", "pano_camera05", "pano_camera06", "pano_camera10"]
FOCAL = 720.0


def fetch(url, dest):
    dest.parent.mkdir(parents=True, exist_ok=True)
    urllib.request.urlretrieve(url, dest)


def gitea_put(rel, data):
    req = urllib.request.Request(f"{OUT_BASE}/{rel}", data=data, method="PUT")
    req.add_header("Authorization", f"token {GITEA_TOKEN}")
    req.add_header("Content-Type", "application/octet-stream")
    urllib.request.urlopen(req, timeout=300)


def read_images_bin(path):
    imgs = {}
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        for _ in range(n):
            struct.unpack("<I", f.read(4))
            q = struct.unpack("<4d", f.read(32))
            t = struct.unpack("<3d", f.read(24))
            struct.unpack("<I", f.read(4))
            name = b""
            while True:
                c = f.read(1)
                if c == b"\x00":
                    break
                name += c
            npts = struct.unpack("<Q", f.read(8))[0]
            data = np.frombuffer(f.read(24 * npts), dtype=np.float64).reshape(-1, 3)
            imgs[name.decode()] = (q, t, data[:, :2].copy(),
                                   np.frombuffer(data[:, 2].tobytes(), dtype=np.int64))
    return imgs


def read_points3D_bin(path):
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        ids = np.empty(n, np.int64)
        xyz = np.empty((n, 3))
        for i in range(n):
            ids[i] = struct.unpack("<Q", f.read(8))[0]
            xyz[i] = struct.unpack("<3d", f.read(24))
            f.read(3 + 8)
            tl = struct.unpack("<Q", f.read(8))[0]
            f.read(8 * tl)
    return dict(zip(ids.tolist(), xyz))


def sparse_obs(imgs, points, name):
    """(pixel xy, depth) of triangulated points observed in this view."""
    q, t, xys, p3d = imgs[name]
    w, x, y, z = q
    R = np.array([
        [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
        [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
        [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)]])
    valid = p3d >= 0
    pw = np.array([points[p] for p in p3d[valid]])
    pc = pw @ R.T + np.array(t)
    ok = pc[:, 2] > 0.05
    return xys[valid][ok], pc[:, 2][ok]


def robust_fit_relative(pred_disp, xy, depth_gt):
    """fit 1/z = a*disp + b (relative models); return per-point z residuals."""
    d = pred_disp[xy[:, 1].astype(int), xy[:, 0].astype(int)].astype(np.float64)
    g = 1.0 / depth_gt
    A = np.stack([d, np.ones_like(d)], 1)
    sel = np.ones(len(d), bool)
    for _ in range(3):  # IRLS-ish trim
        coef, *_ = np.linalg.lstsq(A[sel], g[sel], rcond=None)
        r = np.abs(A @ coef - g)
        sel = r < 3 * np.median(r[sel]) + 1e-12
    inv = A @ coef
    z = np.where(inv > 1e-6, 1.0 / inv, np.inf)
    return np.abs(z - depth_gt) / depth_gt, float(coef[0])


def robust_fit_metric(pred_depth, xy, depth_gt):
    """scale-only fit for metric models; also report raw (unfitted) residual."""
    z = pred_depth[xy[:, 1].astype(int), xy[:, 0].astype(int)].astype(np.float64)
    ok = z > 1e-6
    s = np.median(depth_gt[ok] / z[ok])
    fitted = np.abs(z * s - depth_gt) / depth_gt
    raw = np.abs(z - depth_gt) / depth_gt
    return fitted, raw, float(s)


def colorize(depth, path):
    import cv2
    d = depth.copy()
    lo, hi = np.percentile(d[np.isfinite(d)], [2, 98])
    d = np.clip((d - lo) / max(hi - lo, 1e-6), 0, 1)
    cv2.imwrite(str(path), cv2.applyColorMap((d * 255).astype(np.uint8),
                                             cv2.COLORMAP_TURBO))


def main():
    import cv2
    import torch

    WORK.mkdir(parents=True, exist_ok=True)
    out = WORK / "out"
    out.mkdir(exist_ok=True)
    print("[bake] downloading inputs", flush=True)
    fetch(RIG_URL, WORK / "rig.zip")
    fetch(SFM_URL, WORK / "sfm.zip")
    with zipfile.ZipFile(WORK / "rig.zip") as zf:
        # pick 2 frames that actually exist, evenly spaced across the capture
        in_ring0 = sorted(n.split("/")[-1] for n in zf.namelist()
                          if n.startswith(f"images/{RINGS[0]}/") and n.endswith(".jpg"))
        frames = [in_ring0[len(in_ring0) // 4], in_ring0[3 * len(in_ring0) // 4]]
        print(f"[bake] frames: {frames}", flush=True)
        names = [f"{r}/{f}" for f in frames for r in RINGS]
        for n in names:
            zf.extract(f"images/{n}", WORK)
    with zipfile.ZipFile(WORK / "sfm.zip") as zf:
        zf.extractall(WORK / "sfm")
    imgs_bin = read_images_bin(WORK / "sfm/sparse/0/images.bin")
    points = read_points3D_bin(WORK / "sfm/sparse/0/points3D.bin")

    images = {n: cv2.cvtColor(cv2.imread(str(WORK / "images" / n)), cv2.COLOR_BGR2RGB)
              for n in names}
    obs = {}
    for n in names:
        if n in imgs_bin:
            obs[n] = sparse_obs(imgs_bin, points, n)
            print(f"[gt] {n}: {len(obs[n][1])} sparse obs", flush=True)

    device = "cuda"
    results = {}

    def record(model, name, fitted, extra):
        results.setdefault(model, {})[name] = {
            "median_absrel": float(np.median(fitted)),
            "p90_absrel": float(np.percentile(fitted, 90)),
            "n_obs": int(len(fitted)), **extra}

    # --- Depth Anything V2 Small (relative disparity) ---
    try:
        t0 = time.time()
        from transformers import pipeline
        pipe = pipeline("depth-estimation",
                        model="depth-anything/Depth-Anything-V2-Small-hf", device=0)
        load_s = time.time() - t0
        for n in names:
            if n not in obs:
                continue
            t1 = time.time()
            pred = pipe(__import__("PIL.Image", fromlist=["fromarray"]).fromarray(images[n]))
            disp = np.array(pred["predicted_depth"])
            disp = cv2.resize(disp, images[n].shape[:2][::-1])
            fitted, slope = robust_fit_relative(disp, *obs[n])
            record("dav2_small", n, fitted,
                   {"fit": "scale+shift(disp)", "infer_s": round(time.time() - t1, 2)})
            colorize(disp, out / f"dav2_small__{n.replace('/', '_')}.png")
        results["dav2_small"]["_load_s"] = round(load_s, 1)
        del pipe
        torch.cuda.empty_cache()
    except Exception:
        results["dav2_small"] = {"_error": traceback.format_exc()[-1500:]}

    # --- MoGe-2 (metric point maps) ---
    for tag, hf in [("moge2_vits", "Ruicheng/moge-2-vits-normal"),
                    ("moge2_vitb", "Ruicheng/moge-2-vitb-normal")]:
        try:
            t0 = time.time()
            from moge.model.v2 import MoGeModel
            model = MoGeModel.from_pretrained(hf).to(device).eval()
            load_s = time.time() - t0
            for n in names:
                if n not in obs:
                    continue
                t1 = time.time()
                ten = torch.tensor(images[n] / 255.0, dtype=torch.float32,
                                   device=device).permute(2, 0, 1)
                o = model.infer(ten)
                depth = o["depth"].cpu().numpy().astype(np.float64)
                depth = cv2.resize(depth, images[n].shape[:2][::-1])
                fitted, raw, scale = robust_fit_metric(depth, *obs[n])
                record(tag, n, fitted,
                       {"fit": "scale-only", "raw_median_absrel": float(np.median(raw)),
                        "metric_scale_fit": scale, "infer_s": round(time.time() - t1, 2)})
                colorize(depth, out / f"{tag}__{n.replace('/', '_')}.png")
            results[tag]["_load_s"] = round(load_s, 1)
            del model
            torch.cuda.empty_cache()
        except Exception:
            results[tag] = {"_error": traceback.format_exc()[-1500:]}

    # --- Metric3D v2 (metric, wants intrinsics) ---
    # its hub code imports mmcv 1.x utils; shim them from mmengine (no 1.x wheels
    # exist for this torch, and a source build blows the time box)
    try:
        import sys
        import types
        import mmengine
        from mmengine.utils.dl_utils import collect_env as _collect_env
        fake_mmcv = types.ModuleType("mmcv")
        fake_utils = types.ModuleType("mmcv.utils")
        fake_utils.collect_env = _collect_env
        fake_utils.Config = mmengine.Config
        fake_utils.DictAction = mmengine.DictAction
        fake_mmcv.utils = fake_utils
        sys.modules.setdefault("mmcv", fake_mmcv)
        sys.modules.setdefault("mmcv.utils", fake_utils)
    except Exception:
        print("[bake] mmcv shim failed:\n" + traceback.format_exc()[-800:], flush=True)

    for tag, hub in [("metric3d_vits", "metric3d_vit_small"),
                     ("metric3d_vitl", "metric3d_vit_large")]:
        try:
            t0 = time.time()
            model = torch.hub.load("yvanyin/metric3d", hub, pretrain=True)
            model = model.to(device).eval()
            load_s = time.time() - t0
            for n in names:
                if n not in obs:
                    continue
                t1 = time.time()
                rgb = images[n]
                # their vit recipe: resize to fit 616x1064 canonical, scale intrinsics
                h, w = rgb.shape[:2]
                scale_in = min(616 / h, 1064 / w)
                nh, nw = int(h * scale_in), int(w * scale_in)
                r = cv2.resize(rgb, (nw, nh))
                pad_h, pad_w = 616 - nh, 1064 - nw
                r = cv2.copyMakeBorder(r, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT,
                                       value=[123.675, 116.28, 103.53])
                mean = torch.tensor([123.675, 116.28, 103.53]).float()[:, None, None]
                std = torch.tensor([58.395, 57.12, 57.375]).float()[:, None, None]
                ten = torch.from_numpy(r.transpose(2, 0, 1)).float()
                ten = torch.div(ten - mean, std)[None].to(device)
                with torch.no_grad():
                    pred, conf, od = model.inference({"input": ten})
                pred = pred.squeeze().cpu().numpy()[:nh, :nw]
                pred = pred * (FOCAL * scale_in) / 1000.0  # canonical -> real
                depth = cv2.resize(pred, (w, h)).astype(np.float64)
                fitted, raw, scale = robust_fit_metric(depth, *obs[n])
                record(tag, n, fitted,
                       {"fit": "scale-only", "raw_median_absrel": float(np.median(raw)),
                        "metric_scale_fit": scale, "infer_s": round(time.time() - t1, 2)})
                colorize(depth, out / f"{tag}__{n.replace('/', '_')}.png")
            results[tag]["_load_s"] = round(load_s, 1)
            del model
            torch.cuda.empty_cache()
        except Exception:
            results[tag] = {"_error": traceback.format_exc()[-1500:]}

    (out / "results.json").write_text(json.dumps(results, indent=2))
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
        for p in sorted(out.iterdir()):
            z.write(p, p.name)
    gitea_put("bakeoff.zip", buf.getvalue())
    gitea_put("DONE", b"done")
    print("[bake] uploaded, done", flush=True)


if __name__ == "__main__":
    main()