import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import torch.nn.functional as F
import gradio as gr
import numpy as np
import cv2
import yaml
import sys
import tempfile
from pathlib import Path
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf

# Add the app root to sys.path so core.* and Utils are importable
CODE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, CODE_DIR)

from core.utils.utils import InputPadder
from Utils import vis_disparity, AMP_DTYPE

MODEL_REPO = "nvidia/c-fast-foundationstereo"
WEIGHTS_FILENAME = "model_best_bp2_serialize.pth"
CFG_FILENAME = "cfg.yaml"


def _load_model():
    """Download the checkpoint and cfg from the HF Hub, then torch.load the pickled model."""
    weights_path = hf_hub_download(MODEL_REPO, WEIGHTS_FILENAME, repo_type="model")
    cfg_path = hf_hub_download(MODEL_REPO, CFG_FILENAME, repo_type="model")
    with open(cfg_path, "r") as f:
        cfg: dict = yaml.safe_load(f)
    cfg.setdefault("valid_iters", 8)
    cfg.setdefault("max_disp", 192)
    cfg.setdefault("normalize", True)
    cfg.setdefault("cv_group", 8)
    cfg.setdefault("volume_dim", 28)
    args = OmegaConf.create(cfg)
    # Monkey-patch torch.load to allow weights_only=False for this pickled model
    _orig_load = torch.load
    torch.load = lambda *a, **k: _orig_load(*a, **{**k, "weights_only": k.get("weights_only", False)})
    model = _orig_load(weights_path, map_location="cpu", weights_only=False)
    torch.load = _orig_load
    model.args = args
    model.cuda().eval()
    return model, args


model, model_args = _load_model()
model = model.to("cuda")


@spaces.GPU(duration=60)
def predict_disparity(
    left_image: np.ndarray,
    right_image: np.ndarray,
    valid_iters: int = 8,
    max_disp: int = 192,
) -> tuple:
    """Estimate disparity from a rectified binocular stereo image pair.

    Given a rectified left and right stereo image (RGB), produces a color-mapped
    disparity visualization alongside the raw 16-bit disparity map.

    Args:
        left_image: Rectified left stereo image (RGB), as a numpy array (H, W, 3).
        right_image: Rectified right stereo image (RGB), as a numpy array (H, W, 3).
        valid_iters: Number of GRU refinement iterations (more = better quality, less = faster).
        max_disp: Maximum disparity for volume encoding. 192 is sufficient for most scenes.

    Returns:
        A tuple of (color-mapped disparity visualization, 16-bit disparity map).
    """
    if left_image is None or right_image is None:
        raise ValueError("Both left and right images are required.")

    # Ensure RGB
    img0 = left_image[..., :3]
    img1 = right_image[..., :3]
    if img0.dtype != np.uint8:
        img0 = (img0 * 255).astype(np.uint8) if img0.max() <= 1.0 else img0.astype(np.uint8)
        img1 = (img1 * 255).astype(np.uint8) if img1.max() <= 1.0 else img1.astype(np.uint8)

    H, W = img0.shape[:2]
    img0_ori = img0.copy()
    img1_ori = img1.copy()

    # Resize right to match left if needed
    if img1.shape[:2] != (H, W):
        img1 = cv2.resize(img1, (W, H))

    img0_t = torch.as_tensor(img0).cuda().float()[None].permute(0, 3, 1, 2)
    img1_t = torch.as_tensor(img1).cuda().float()[None].permute(0, 3, 1, 2)

    padder = InputPadder(img0_t.shape, divis_by=32, force_square=False)
    img0_t, img1_t = padder.pad(img0_t, img1_t)

    model.args.valid_iters = valid_iters
    model.args.max_disp = max_disp

    with torch.amp.autocast("cuda", enabled=True, dtype=AMP_DTYPE):
        disp = model.forward(img0_t, img1_t, iters=valid_iters, test_mode=True, optimize_build_volume="pytorch1")
    disp = padder.unpad(disp.float())
    disp_np = disp.data.cpu().numpy().reshape(H, W).clip(0, None)

    # Color-mapped visualization
    vis = vis_disparity(disp_np, color_map=cv2.COLORMAP_TURBO)

    # Side-by-side: left | right | disparity
    combined = np.concatenate([img0_ori, img1_ori, vis], axis=1)

    # Save 16-bit disparity to a temp file
    disp_16 = (disp_np * 256).clip(0, 65535).astype(np.uint16)
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    cv2.imwrite(tmp.name, disp_16)

    return combined, tmp.name


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# Fast-FoundationStereo: Real-Time Zero-Shot Stereo Disparity")
        gr.Markdown(
            "Upload a **rectified** binocular stereo image pair to estimate a disparity map. "
            "This is [nvidia/c-fast-foundationstereo](https://huggingface.co/nvidia/c-fast-foundationstereo), "
            "a 14.6M-parameter real-time stereo foundation model from NVIDIA (CVPR 2026)."
        )

        with gr.Row():
            with gr.Column():
                left_input = gr.Image(label="Left Image (rectified RGB)", type="numpy")
                right_input = gr.Image(label="Right Image (rectified RGB)", type="numpy")
            with gr.Column():
                disparity_vis = gr.Image(label="Disparity Visualization (left | right | color-mapped disparity)")
                disparity_raw = gr.Image(label="16-bit Disparity Map", type="filepath")

        with gr.Accordion("Advanced Settings", open=False):
            valid_iters_slider = gr.Slider(
                minimum=2, maximum=16, value=8, step=1,
                label="Valid Iterations (GRU refinement steps)",
                info="More iterations = better quality but slower. 8 is default, 4 is faster."
            )
            max_disp_slider = gr.Slider(
                minimum=64, maximum=416, value=192, step=32,
                label="Max Disparity",
                info="Maximum disparity for volume encoding. 192 is enough for most scenes."
            )

        run_btn = gr.Button("Estimate Disparity", variant="primary")

        gr.Examples(
            examples=[
                ["examples/left.png", "examples/right.png", 8, 192],
                ["examples/left.png", "examples/right.png", 4, 192],
            ],
            inputs=[left_input, right_input, valid_iters_slider, max_disp_slider],
            outputs=[disparity_vis, disparity_raw],
            fn=predict_disparity,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=predict_disparity,
        inputs=[left_input, right_input, valid_iters_slider, max_disp_slider],
        outputs=[disparity_vis, disparity_raw],
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)