import spaces  # MUST be the very first import for ZeroGPU to work

import io
import numpy as np
import cv2
from PIL import Image, UnidentifiedImageError
import torch
import gradio as gr
from fastapi import File, UploadFile, HTTPException
from fastapi.responses import Response
from transformers import DepthProImageProcessor, DepthProForDepthEstimation


# =====================================================================
# 1. Load Apple Depth Pro at module level.
#    ZeroGPU safely emulates CUDA during startup so .to("cuda") works here.
# =====================================================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype  = torch.float16 if torch.cuda.is_available() else torch.float32

print("Loading Apple Depth Pro components...")
image_processor = DepthProImageProcessor.from_pretrained("apple/DepthPro-hf")
model = DepthProForDepthEstimation.from_pretrained(
    "apple/DepthPro-hf",
    torch_dtype=dtype,
).to(device)
model.eval()
print(f"Model successfully loaded on target device: {device}")


# =====================================================================
# 2. Shared helpers
# =====================================================================
def _run_depth_inference(image: Image.Image):
    """Returns post-processed depth output list. Caller handles indexing."""
    inputs = image_processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    return image_processor.post_process_depth_estimation(
        outputs,
        target_sizes=[(image.height, image.width)],
    )


def _write_binary_ply(x, y, z, rgb_ints, path: str) -> str:
    """Pack arrays into a binary little-endian PLY and write to path."""
    num_points = len(z)
    vertex_dtype = np.dtype([
        ('x', '<f4'), ('y', '<f4'), ('z', '<f4'),
        ('red', 'u1'), ('green', 'u1'), ('blue', 'u1'),
    ])
    vertices = np.empty(num_points, dtype=vertex_dtype)
    vertices['x']     = x.astype(np.float32)
    vertices['y']     = (-y).astype(np.float32)   # invert Y → right-side up
    vertices['z']     = z.astype(np.float32)
    vertices['red']   = rgb_ints[:, 0]
    vertices['green'] = rgb_ints[:, 1]
    vertices['blue']  = rgb_ints[:, 2]

    header = (
        "ply\n"
        "format binary_little_endian 1.0\n"
        f"element vertex {num_points}\n"
        "property float x\n"
        "property float y\n"
        "property float z\n"
        "property uchar red\n"
        "property uchar green\n"
        "property uchar blue\n"
        "end_header\n"
    )
    with open(path, "wb") as f:
        f.write(header.encode("ascii"))
        f.write(vertices.tobytes())
    return path


def _frame_to_pointcloud(image: Image.Image, depth: np.ndarray, max_pts: int = 50_000):
    """Project a single frame into (x, y, z, rgb) arrays in camera space."""
    h, w = depth.shape
    color_np = np.array(image, dtype=np.float32) / 255.0
    focal = float(max(w, h))
    cx, cy = w / 2.0, h / 2.0

    u_g, v_g = np.meshgrid(np.arange(w, dtype=np.float32),
                             np.arange(h, dtype=np.float32))
    z = depth.flatten()
    x = ((u_g.flatten() - cx) * z) / focal
    y = ((v_g.flatten() - cy) * z) / focal

    valid = (z > 0.1) & (z < 30.0)
    x, y, z = x[valid], y[valid], z[valid]
    rgb = (color_np.reshape(-1, 3)[valid] * 255).astype(np.uint8)

    if len(z) > max_pts:
        idx = np.random.choice(len(z), max_pts, replace=False)
        x, y, z, rgb = x[idx], y[idx], z[idx], rgb[idx]

    return x, y, z, rgb


# =====================================================================
# 3a. Single-image function (original feature)
# =====================================================================
@spaces.GPU(duration=120)
def create_3d_point_cloud(input_pil_image):
    """Single photo → coloured 3D point cloud PLY."""
    if input_pil_image is None:
        return None

    image = input_pil_image.convert("RGB")
    if max(image.width, image.height) > 512:
        image.thumbnail((512, 512), Image.LANCZOS)

    post = _run_depth_inference(image)
    depth = post[0]["predicted_depth"].detach().cpu().numpy().astype(np.float32)

    x, y, z, rgb = _frame_to_pointcloud(image, depth, max_pts=200_000)
    return _write_binary_ply(x, y, z, rgb, "/tmp/output_scene.ply")


# =====================================================================
# 3b. Video → environment map (new feature)
# =====================================================================
@spaces.GPU(duration=300)
def create_map_from_video(video_path, frame_stride: int = 10, max_frames: int = 20):
    """
    Walk-around video → merged 3D environment map.

    Pipeline
    --------
    1. Extract keyframes (every `frame_stride` video frames, up to `max_frames`).
    2. Downscale each frame to 512 px and run DepthPro → metric depth.
    3. SIFT feature detection + FLANN matching between consecutive frames.
    4. Essential Matrix → R, t̂  (unit translation from recoverPose).
    5. Use metric depth at matched inliers to recover the true translation scale.
    6. Accumulate 4×4 camera poses in a world coordinate frame.
    7. Project each frame's pixels into world space and merge point clouds.
    8. Subsample to ≤ 500k points and write binary PLY.
    """
    if video_path is None:
        return None, "No video provided."

    # ------------------------------------------------------------------
    # Step 1: Extract keyframes
    # ------------------------------------------------------------------
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, "Could not open video file."

    frames_pil = []
    total_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_indices = range(0, total_video_frames, int(frame_stride))

    for fi in frame_indices:
        if len(frames_pil) >= int(max_frames):
            break
        cap.set(cv2.CAP_PROP_POS_FRAMES, fi)
        ret, bgr = cap.read()
        if not ret:
            break
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        pil = Image.fromarray(rgb)
        if max(pil.width, pil.height) > 512:
            pil.thumbnail((512, 512), Image.LANCZOS)
        frames_pil.append(pil)
    cap.release()

    n = len(frames_pil)
    if n < 2:
        return None, f"Need at least 2 frames; got {n}."

    W, H = frames_pil[0].width, frames_pil[0].height
    focal = float(max(W, H))
    cx, cy = W / 2.0, H / 2.0
    K = np.array([[focal, 0, cx], [0, focal, cy], [0, 0, 1]], dtype=np.float64)

    # ------------------------------------------------------------------
    # Step 2: DepthPro inference on every keyframe
    # ------------------------------------------------------------------
    depth_maps = []
    for pil in frames_pil:
        post = _run_depth_inference(pil)
        d = post[0]["predicted_depth"].detach().cpu().numpy().astype(np.float32)
        depth_maps.append(d)

    # ------------------------------------------------------------------
    # Step 3: SIFT keypoints for every frame
    # ------------------------------------------------------------------
    sift = cv2.SIFT_create(nfeatures=2000)
    kps_list, desc_list = [], []
    for pil in frames_pil:
        gray = cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2GRAY)
        kps, desc = sift.detectAndCompute(gray, None)
        kps_list.append(kps)
        desc_list.append(desc)

    # ------------------------------------------------------------------
    # Step 4-5: Relative pose estimation with metric-depth scale recovery
    # ------------------------------------------------------------------
    flann = cv2.FlannBasedMatcher({"algorithm": 1, "trees": 5}, {"checks": 50})

    # poses[i] = 4×4 world-to-camera transform for frame i
    poses = [np.eye(4, dtype=np.float64)]

    for i in range(n - 1):
        d_i   = desc_list[i]
        d_ip1 = desc_list[i + 1]

        # Fall back to identity step if SIFT failed on either frame
        if d_i is None or d_ip1 is None or len(d_i) < 10 or len(d_ip1) < 10:
            poses.append(poses[-1].copy())
            continue

        try:
            matches = flann.knnMatch(d_i, d_ip1, k=2)
        except cv2.error:
            poses.append(poses[-1].copy())
            continue

        # Lowe's ratio test
        good = [m for m, nn in matches if m.distance < 0.7 * nn.distance]
        if len(good) < 10:
            poses.append(poses[-1].copy())
            continue

        pts1 = np.float32([kps_list[i][m.queryIdx].pt for m in good])
        pts2 = np.float32([kps_list[i + 1][m.trainIdx].pt for m in good])

        E, mask_E = cv2.findEssentialMat(
            pts1, pts2, K, method=cv2.RANSAC, prob=0.999, threshold=1.0
        )
        if E is None:
            poses.append(poses[-1].copy())
            continue

        _, R, t_unit, mask_pose = cv2.recoverPose(E, pts1, pts2, K, mask=mask_E)

        # --- Scale recovery via metric depth ---
        # recoverPose gives ||t|| = 1. We use DepthPro depth at inlier
        # correspondences to estimate the true translation magnitude.
        inliers = mask_pose.ravel() == 255
        pts1_in = pts1[inliers]
        pts2_in = pts2[inliers]

        depth_i   = depth_maps[i]
        depth_ip1 = depth_maps[i + 1]
        h_d, w_d  = depth_i.shape

        scales = []
        for p1, p2 in zip(pts1_in, pts2_in):
            u1, v1 = int(np.clip(p1[0], 0, w_d - 1)), int(np.clip(p1[1], 0, h_d - 1))
            u2, v2 = int(np.clip(p2[0], 0, w_d - 1)), int(np.clip(p2[1], 0, h_d - 1))

            d1 = float(depth_i[v1, u1])
            d2 = float(depth_ip1[v2, u2])
            if not (0.1 < d1 < 30.0 and 0.1 < d2 < 30.0):
                continue

            # 3-D positions in camera-i and camera-(i+1) coords
            P1 = np.array([(p1[0] - cx) * d1 / focal,
                            (p1[1] - cy) * d1 / focal,
                            d1])
            P2 = np.array([(p2[0] - cx) * d2 / focal,
                            (p2[1] - cy) * d2 / focal,
                            d2])
            # Actual translation ≈ P2 - R @ P1
            t_actual = P2 - R @ P1
            scales.append(np.linalg.norm(t_actual))

        scale = float(np.median(scales)) if scales else 1.0
        t_scaled = t_unit * scale

        # Build incremental 4×4 transform
        T_rel = np.eye(4, dtype=np.float64)
        T_rel[:3, :3] = R
        T_rel[:3, 3]  = t_scaled.ravel()

        poses.append(T_rel @ poses[-1])

    # ------------------------------------------------------------------
    # Step 6-7: Project each frame → world space and accumulate
    # ------------------------------------------------------------------
    all_pts, all_rgb = [], []

    for i, (pil, depth, pose) in enumerate(zip(frames_pil, depth_maps, poses)):
        x_c, y_c, z_c, rgb = _frame_to_pointcloud(pil, depth, max_pts=50_000)

        # Camera → world: pts_world = inv(pose) @ pts_cam_homogeneous
        pts_h = np.vstack([x_c, y_c, z_c, np.ones_like(z_c)])   # 4 × N
        pts_w = (np.linalg.inv(pose) @ pts_h)[:3]                # 3 × N

        all_pts.append(pts_w.T)   # N × 3
        all_rgb.append(rgb)

    XYZ = np.concatenate(all_pts, axis=0)
    RGB = np.concatenate(all_rgb, axis=0)

    # ------------------------------------------------------------------
    # Step 8: Global subsample → binary PLY
    # ------------------------------------------------------------------
    MAX_TOTAL = 500_000
    if len(XYZ) > MAX_TOTAL:
        idx = np.random.choice(len(XYZ), MAX_TOTAL, replace=False)
        XYZ = XYZ[idx]
        RGB = RGB[idx]

    ply_path = _write_binary_ply(XYZ[:, 0], XYZ[:, 1], XYZ[:, 2],
                                  RGB, "/tmp/environment_map.ply")
    msg = (
        f"✅ Processed {n} frames · "
        f"{len(XYZ):,} points in world map"
    )
    return ply_path, msg


# =====================================================================
# 4. Gradio UI — two tabs
# =====================================================================
with gr.Blocks(title="Software-ToF 3D Cloud Visualizer") as demo:
    gr.Markdown("# 📷 Software-ToF Camera Simulation")

    with gr.Tab("📷 Single Image → 3D Snapshot"):
        gr.Markdown(
            "Upload any photo to get an interactive 3D point cloud "
            "with absolute metric depths from Apple DepthPro."
        )
        with gr.Row():
            with gr.Column():
                input_img  = gr.Image(type="pil", label="Input Scene Photo")
                img_btn    = gr.Button("Generate ToF 3D Structure", variant="primary")
            with gr.Column():
                output_3d  = gr.Model3D(
                    label="Interactive 3D Point Cloud",
                    display_mode="point_cloud",
                )
        img_btn.click(fn=create_3d_point_cloud,
                      inputs=[input_img],
                      outputs=[output_3d])

    with gr.Tab("🎥 Video → Environment Map"):
        gr.Markdown(
            "**Walk around a room while recording, then upload the video.**  \n"
            "Frames are matched with SIFT feature correspondences. "
            "DepthPro provides metric depth per frame; the scale of each "
            "relative camera translation is recovered from matched 3-D points, "
            "then all frames are fused into a single world-coordinate map."
        )
        with gr.Row():
            with gr.Column():
                video_input   = gr.Video(label="Walk-around Video")
                frame_stride  = gr.Slider(5, 30, value=10, step=5,
                                           label="Frame Stride (sample every N frames)")
                max_frames    = gr.Slider(5, 30, value=20, step=5,
                                           label="Max Frames to Process")
                video_btn     = gr.Button("Generate Environment Map", variant="primary")
            with gr.Column():
                map_output    = gr.Model3D(
                    label="3D Environment Map",
                    display_mode="point_cloud",
                )
                status_box    = gr.Textbox(label="Status", interactive=False)
        video_btn.click(
            fn=create_map_from_video,
            inputs=[video_input, frame_stride, max_frames],
            outputs=[map_output, status_box],
        )


# =====================================================================
# 5. REST endpoints via Gradio's internal FastAPI app
# =====================================================================
@demo.app.get("/status")
def status():
    return {"status": "running", "device": str(device)}


@demo.app.post("/get_metric_depth")
async def get_metric_depth(file: UploadFile = File(...)):
    if not (file.content_type or "").startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")
    try:
        image = Image.open(io.BytesIO(await file.read())).convert("RGB")
    except (UnidentifiedImageError, OSError):
        raise HTTPException(status_code=400, detail="Could not decode image.")
    try:
        post = _run_depth_inference(image)
        depth = post[0]["predicted_depth"].detach().cpu().numpy().astype(np.float32)
        buf = io.BytesIO()
        np.save(buf, depth)
        return Response(content=buf.getvalue(), media_type="application/octet-stream")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")


# =====================================================================
# 6. Launch
# =====================================================================
if __name__ == "__main__":
    demo.launch()