import os
# Must be set BEFORE importing cv2 to enable the EXR codec.
os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1")

import tempfile
import spaces
import gradio as gr
import numpy as np
import torch
import cv2
import tifffile
import matplotlib
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForDepthEstimation

MODEL_ID = "LiheYoung/depth-anything-large-hf"

processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForDepthEstimation.from_pretrained(MODEL_ID)
model.eval()

cmap = matplotlib.colormaps.get_cmap("Spectral_r")


def _write_exr(path, depth_f32):
    """Write single-value float32 depth to EXR (RGB channels all = depth).
    Try OpenCV first, then the OpenEXR package. Return path or None."""
    try:
        ok = cv2.imwrite(path, depth_f32)  # cv2 writes float32 -> EXR
        if ok and os.path.exists(path) and os.path.getsize(path) > 0:
            return path
    except Exception as e:
        print("cv2 EXR write failed:", e)
    try:
        import OpenEXR, Imath
        h, w = depth_f32.shape
        header = OpenEXR.Header(w, h)
        ftype = Imath.Channel(Imath.PixelType(Imath.PixelType.FLOAT))
        header["channels"] = {c: ftype for c in ("R", "G", "B")}
        out = OpenEXR.OutputFile(path, header)
        buf = np.ascontiguousarray(depth_f32).tobytes()
        out.writePixels({"R": buf, "G": buf, "B": buf})
        out.close()
        return path
    except Exception as e:
        print("OpenEXR write failed:", e)
    return None


@spaces.GPU
@torch.no_grad()
def predict_depth(image: Image.Image):
    if image is None:
        return None, []
    image = image.convert("RGB")
    model.to("cuda")
    inputs = processor(images=image, return_tensors="pt").to("cuda")
    outputs = model(**inputs)
    predicted_depth = outputs.predicted_depth

    prediction = torch.nn.functional.interpolate(
        predicted_depth.unsqueeze(1),
        size=image.size[::-1],  # (H, W)
        mode="bicubic",
        align_corners=False,
    ).squeeze()

    # Raw float32 depth (full precision, NOT normalized) for TIFF/EXR.
    depth = prediction.detach().cpu().numpy().astype(np.float32)

    dmin, dmax = float(depth.min()), float(depth.max())
    norm = (depth - dmin) / (dmax - dmin + 1e-8)  # 0..1 for preview + 16-bit PNG

    # 8-bit colored preview (display only — precision doesn't matter here).
    colored = Image.fromarray((cmap(norm)[:, :, :3] * 255).astype(np.uint8))

    tmp = tempfile.mkdtemp(prefix="depth_")
    png16_path = os.path.join(tmp, "depth_16bit.png")
    tiff_path = os.path.join(tmp, "depth_float32.tiff")
    exr_path = os.path.join(tmp, "depth_float32.exr")

    # 16-bit PNG: normalized depth mapped to the full 0..65535 range.
    cv2.imwrite(png16_path, (norm * 65535.0).astype(np.uint16))

    # 32-bit float TIFF: raw depth, fully lossless.
    tifffile.imwrite(tiff_path, depth)

    # 32-bit float EXR: raw depth (best-effort).
    exr_ok = _write_exr(exr_path, depth)

    files = [png16_path, tiff_path]
    if exr_ok:
        files.append(exr_ok)
    return colored, files


title = "# Depth Anything (ZeroGPU) — high-bit-depth export"
description = (
    "Monocular depth with Depth-Anything-Large on ZeroGPU.\n\n"
    "**Downloads preserve full precision** — the colored image is an 8-bit preview only:\n"
    "- `depth_16bit.png` — 16-bit grayscale, depth normalized to 0–65535\n"
    "- `depth_float32.tiff` — 32-bit float, raw depth values (lossless)\n"
    "- `depth_float32.exr` — 32-bit float OpenEXR, raw depth values"
)

with gr.Blocks(title="Depth Anything") as demo:
    gr.Markdown(title)
    gr.Markdown(description)
    with gr.Row():
        input_image = gr.Image(label="Input", type="pil")
        depth_color = gr.Image(label="Depth (colored preview, 8-bit)", type="pil")
    downloads = gr.Files(label="High-bit-depth downloads (PNG16 / TIFF32 / EXR32)")
    run = gr.Button("Compute depth", variant="primary")
    run.click(fn=predict_depth, inputs=input_image, outputs=[depth_color, downloads])

if __name__ == "__main__":
    demo.queue().launch()