"""Model zoo for the ZDPShift Space: six backbones, zero-shot and corrected.

Weights resolve from ZDP_WEIGHTS_ROOT if set (local development), otherwise
from the shijianjian/ZDPShift model repo via snapshot_download. Image models
are fine-tuned on translated SceneFlow; video models on translated SceneFlow +
Dynamic Replica, matching the data their released baselines were trained on.

The three image repos and BiDAStereo's vendored RAFT all ship a top-level
`core` package; each family's modules and sys.path entries are swapped in and
out when the active family changes, so all six coexist in one process.
"""
import os
import shutil
import sys

import numpy as np
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
CODE = os.path.join(HERE, "zdp")
sys.path.insert(0, CODE)
sys.path.insert(0, os.path.join(CODE, "third_party"))


_REQUIRED = [f"{m}_{v}.pth" for m in ("raft", "igev", "fs", "dynamic", "bida", "sav")
             for v in ("zs", "ours")] + ["fs_cfg.yaml", "video_depth_anything_vits.pth"]


def _weights_root() -> str:
    """env var > ./weights next to the app > hub download — validated."""
    candidates = []
    env = os.environ.get("ZDP_WEIGHTS_ROOT")
    if env:
        candidates.append(("ZDP_WEIGHTS_ROOT", env))
    local = os.path.join(HERE, "weights")
    if os.path.isdir(local):
        candidates.append(("./weights", local))
    for label, root in candidates:
        missing = [f for f in _REQUIRED if not os.path.exists(os.path.join(root, f))]
        if not missing:
            return root
        raise RuntimeError(f"weights root {label}={root} is missing: {missing}")
    from huggingface_hub import snapshot_download
    root = snapshot_download("shijianjian/ZDPShift", allow_patterns=_REQUIRED)
    missing = [f for f in _REQUIRED if not os.path.exists(os.path.join(root, f))]
    if missing:
        raise RuntimeError(
            f"hub repo shijianjian/ZDPShift is missing {missing} — "
            "upload the demo weight set (see hf_weights_stage/upload.sh) or set "
            "ZDP_WEIGHTS_ROOT / place a weights/ folder next to app.py")
    return root


ROOT = _weights_root()
W = {  # (zero-shot, corrected)
    "raft": (f"{ROOT}/raft_zs.pth", f"{ROOT}/raft_ours.pth"),
    "igev": (f"{ROOT}/igev_zs.pth", f"{ROOT}/igev_ours.pth"),
    "fs": (f"{ROOT}/fs_zs.pth", f"{ROOT}/fs_ours.pth"),
    "dynamic": (f"{ROOT}/dynamic_zs.pth", f"{ROOT}/dynamic_ours.pth"),
    "bida": (f"{ROOT}/bida_zs.pth", f"{ROOT}/bida_ours.pth"),
    "sav": (f"{ROOT}/sav_zs.pth", f"{ROOT}/sav_ours.pth"),
}
FS_CFG = f"{ROOT}/fs_cfg.yaml"
NAMES = {"raft": "RAFT-Stereo", "igev": "IGEV-Stereo", "fs": "FoundationStereo",
         "dynamic": "DynamicStereo", "bida": "BiDAStereo", "sav": "StereoAnyVideo"}
IMAGE, VIDEO = ("raft", "igev", "fs"), ("dynamic", "bida", "sav")

# StereoAnyVideo constructs its VDA prior from a repo-relative path; the actual
# values are then overwritten by the checkpoint (the prior is frozen), but the
# file must exist where the constructor looks for it.
_VDA_DST = os.path.join(CODE, "third_party/StereoAnyVideo/models",
                        "Video-Depth-Anything/checkpoints/video_depth_anything_vits.pth")


def _ensure_vda():
    if not os.path.exists(_VDA_DST):
        os.makedirs(os.path.dirname(_VDA_DST), exist_ok=True)
        shutil.copyfile(f"{ROOT}/video_depth_anything_vits.pth", _VDA_DST)


# ---- image-family namespace juggling ------------------------------------
_family_modules: dict[str, dict] = {}
_active_family: str | None = None
_CLASH = ("core", "models", "raft_stereo", "igev_stereo", "utils")
_FAMILY_DIRS = {
    "raft": [os.path.join(CODE, "third_party/RAFT-Stereo"),
             os.path.join(CODE, "third_party/RAFT-Stereo/core")],
    "igev": [os.path.join(CODE, "third_party/IGEV/IGEV-Stereo"),
             os.path.join(CODE, "third_party/IGEV/IGEV-Stereo/core")],
    "fs": [os.path.join(CODE, "third_party/FoundationStereo")],
    "bida": [os.path.join(CODE, "third_party/bidastereo/third_party/RAFT"),
             os.path.join(CODE, "third_party/bidastereo/third_party/RAFT/core")],
}


def _swap_family(fam: str):
    global _active_family
    if fam == _active_family:
        return
    if _active_family is not None:
        stash = {}
        for name in list(sys.modules):
            if name.split(".")[0] in _CLASH:
                stash[name] = sys.modules.pop(name)
        _family_modules[_active_family] = stash
    for name, mod in _family_modules.pop(fam, {}).items():
        sys.modules[name] = mod
    others = {d for f, dirs in _FAMILY_DIRS.items() if f != fam for d in dirs}
    sys.path[:] = [p for p in sys.path if p not in others]
    for d in reversed(_FAMILY_DIRS[fam]):
        if d in sys.path:
            sys.path.remove(d)
        sys.path.insert(0, d)
    _active_family = fam


_cache: dict[tuple, object] = {}
_infer_fns: dict[str, object] = {}


def _dev() -> str:
    """Resolved per call: on ZeroGPU, CUDA exists only inside the GPU-decorated
    handler; on a CPU Space it never does; locally it always does."""
    return "cuda" if torch.cuda.is_available() else "cpu"


def _get_image(fam: str, variant: str):
    key = (fam, variant)
    if key not in _cache:
        _swap_family(fam)
        if fam == "raft":
            import eval_raft_stereo as ev
            net = ev.load_model(W[fam][variant == "ours"], symmetric=(variant == "ours"),
                                device=_dev())
        elif fam == "igev":
            import eval_igev_signed as ev
            net = (ev.load_model(W[fam][1], 64, 192, device=_dev()) if variant == "ours"
                   else ev.load_model(W[fam][0], 0, 192, device=_dev()))
        else:
            import eval_foundation_stereo as ev
            net = ev.load_model(W[fam][variant == "ours"], FS_CFG,
                                variant == "ours", 64, 192, device=_dev())
        _infer_fns[fam] = ev.infer
        _cache[key] = net
    else:
        _swap_family(fam)
    return _cache[key]


def _get_video(fam: str, variant: str):
    if fam == "bida":
        _swap_family("bida")
    if fam == "sav":
        _ensure_vda()
    key = (fam, variant)
    if key not in _cache:
        import vsm
        _cache[key] = vsm.build(fam, W[fam][variant == "ours"], device=_dev())
    return _cache[key]


# ---- public API ----------------------------------------------------------
def predict(fam: str, variant: str, left: np.ndarray, right: np.ndarray,
            width: int = 960, iters: int = 24) -> np.ndarray:
    """Signed disparity [H,W] at input resolution; positive in front of the
    screen plane, negative behind it."""
    import cv2
    if _dev() == "cpu":                      # keep CPU Spaces responsive
        width, iters = min(width, 640), min(iters, 16)
    H0, W0 = left.shape[:2]
    s = width / W0
    nw, nh = max(64, int(round(W0 * s)) // 32 * 32), max(64, int(round(H0 * s)) // 32 * 32)
    L = cv2.resize(left, (nw, nh), interpolation=cv2.INTER_AREA).astype(np.float32)
    R = cv2.resize(right, (nw, nh), interpolation=cv2.INTER_AREA).astype(np.float32)
    if fam in IMAGE:
        net = _get_image(fam, variant)
        d = _infer_fns[fam](net, L, R, iters, device=_dev())
    else:
        import vsm
        net = _get_video(fam, variant)
        T = 5
        video = np.stack([np.stack([L, R]).transpose(0, 3, 1, 2)] * T)
        sgn = {"dynamic": -1.0, "bida": -1.0, "sav": 1.0}[fam]
        d = sgn * vsm.infer_clip(fam, net, torch.from_numpy(video), iters=16,
                                 max_width=width, device=_dev())[T // 2]
        d = np.asarray(d)
    d = cv2.resize(d, (W0, H0), interpolation=cv2.INTER_LINEAR) / (nw / W0)
    return d


def colorize(d: np.ndarray, vmax: float | None = None) -> np.ndarray:
    """Signed disparity -> RdBu_r (red = in front, blue = behind screen)."""
    hexes = ["053061", "2166ac", "4393c3", "92c5de", "d1e5f0", "f7f7f7",
             "fddbc7", "f4a582", "d6604d", "b2182b", "67001f"]
    anchors = np.array([[int(h[i:i + 2], 16) for i in (0, 2, 4)] for h in hexes],
                       dtype=np.float32)
    if vmax is None:
        vmax = max(8.0, float(np.percentile(np.abs(d), 99)))
    x = (np.clip(d / vmax, -1, 1) * 0.5 + 0.5) * (len(hexes) - 1)
    i = np.clip(x.astype(int), 0, len(hexes) - 2)
    t = (x - i)[..., None]
    return (anchors[i] * (1 - t) + anchors[i + 1] * t).astype(np.uint8)


def predict_clip(fam: str, variant: str, lefts: list, rights: list,
                 width: int = 960, iters: int = 16) -> np.ndarray:
    """Signed disparity [T,H,W] for a temporal clip (video matchers only)."""
    import cv2
    assert fam in VIDEO, "predict_clip is for the video matchers"
    T = len(lefts)
    if _dev() == "cpu":                      # keep CPU Spaces responsive
        width, iters = min(width, 512), min(iters, 12)
        lefts, rights = lefts[:8], rights[:8]
        T = len(lefts)
    H0, W0 = lefts[0].shape[:2]
    s = width / W0
    nw, nh = max(64, int(round(W0 * s)) // 32 * 32), max(64, int(round(H0 * s)) // 32 * 32)
    frames = []
    for L, R in zip(lefts, rights):
        Ls = cv2.resize(L, (nw, nh), interpolation=cv2.INTER_AREA)
        Rs = cv2.resize(R, (nw, nh), interpolation=cv2.INTER_AREA)
        frames.append(np.stack([Ls, Rs]).transpose(0, 3, 1, 2))
    video = np.stack(frames).astype(np.float32)
    import vsm
    net = _get_video(fam, variant)
    sgn = {"dynamic": -1.0, "bida": -1.0, "sav": 1.0}[fam]
    d = sgn * np.asarray(vsm.infer_clip(fam, net, torch.from_numpy(video),
                                        iters=iters, max_width=width))
    out = np.stack([cv2.resize(f, (W0, H0), interpolation=cv2.INTER_LINEAR)
                    for f in d]) / (nw / W0)
    return out