"""Core LingBot-Map reconstruction pipeline (device-agnostic building blocks).

Shared by ``app.py`` (Gradio / ZeroGPU front-end) and ``smoke_test.py`` (local CLI).
Nothing here is ZeroGPU-specific — the GPU attach / bf16 dance lives in ``app.py`` —
so this module can be imported and exercised on a plain CPU box for verification.

Attribution
-----------
Built on LingBot-Map by the Robbyant Team (Apache-2.0):
  https://github.com/Robbyant/lingbot-map
``postprocess``, ``prepare_for_visualization`` and the video-frame extraction logic
are adapted (lightly) from the upstream ``demo.py``. The model itself builds on
VGGT and DINOv2. See the Space README for full credits.
"""

from __future__ import annotations

import os
import tempfile
from typing import List, Optional, Sequence, Tuple

import numpy as np
import torch

from lingbot_map.models.gct_stream import GCTStream
from lingbot_map.utils.pose_enc import pose_encoding_to_extri_intri
from lingbot_map.utils.geometry import unproject_depth_map_to_point_map
# (w2c kept for the GLB path → closed_form_inverse_se3_general intentionally NOT imported)
# NB: preprocess_paths() reimplements upstream load_and_preprocess_images WITHOUT torchvision
# (the ZeroGPU base image ships torch but not torchvision), so nothing here imports torchvision.

# NB: import the exporter from the *submodule*, not ``lingbot_map.vis``. The package
# __init__ imports the viser-based PointCloudViewer at load time, which would force a
# viser install (and a crash without it). glb_export only needs numpy/cv2/matplotlib/
# scipy/trimesh.
from lingbot_map.vis.glb_export import predictions_to_glb


# =============================================================================
# Model construction
# =============================================================================

def build_model(
    ckpt_path: Optional[str] = None,
    device: str = "cpu",
    camera_num_iterations: int = 4,
    use_sdpa: bool = True,
) -> GCTStream:
    """Build GCTStream with the same config as upstream ``demo.load_model`` and
    (optionally) load a checkpoint.

    ``use_sdpa=True`` selects the built-in scaled-dot-product-attention fallback so
    no FlashInfer source build is needed — the right choice for ZeroGPU.
    """
    model = GCTStream(
        img_size=518,
        patch_size=14,
        enable_3d_rope=True,
        max_frame_num=1024,
        kv_cache_sliding_window=64,
        kv_cache_scale_frames=8,
        kv_cache_cross_frame_special=True,
        kv_cache_include_scale_frames=True,
        use_sdpa=use_sdpa,
        camera_num_iterations=camera_num_iterations,
    )
    if ckpt_path:
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
        state_dict = ckpt.get("model", ckpt)
        missing, unexpected = model.load_state_dict(state_dict, strict=False)
        if missing:
            print(f"[build_model] missing keys: {len(missing)}")
        if unexpected:
            print(f"[build_model] unexpected keys: {len(unexpected)}")
    return model.to(device).eval()


# =============================================================================
# Input gathering / preprocessing
# =============================================================================

_IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".webp", ".JPG", ".JPEG", ".PNG")


def list_image_paths(files: Optional[Sequence]) -> List[str]:
    """Normalise a Gradio file upload (paths or file-like objs) to a sorted list of
    image paths. Sorted by basename so frames named 000.jpg, 001.jpg, ... keep order.
    """
    if not files:
        return []
    paths = []
    for f in files:
        p = f if isinstance(f, str) else getattr(f, "name", None)
        if p and p.lower().endswith(tuple(e.lower() for e in _IMAGE_EXTS)):
            paths.append(p)
    return sorted(paths, key=lambda p: os.path.basename(p))


def extract_video_frames(
    video_path: str,
    fps: int = 10,
    max_frames: int = 60,
    out_dir: Optional[str] = None,
) -> Tuple[List[str], str]:
    """Sample up to ``max_frames`` frames spread EVENLY across the WHOLE video.

    This is the key to reconstructing a full walkthrough (e.g. an apartment tour through
    several rooms): we pick frame indices uniformly over the entire duration instead of
    grabbing the first N from the start (which would only ever show the opening room).
    ``fps`` bounds the density for short clips; for long clips ``max_frames`` is the
    binding cap and the chosen frames still span end-to-end. Memory-safe (reads
    frame-by-frame, O(1) memory). Frames are written to a fresh temp dir.
    """
    import cv2

    if out_dir is None:
        out_dir = tempfile.mkdtemp(prefix="lingbot_frames_")
    cap = cv2.VideoCapture(video_path)
    src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    if total <= 0:  # CAP_PROP_FRAME_COUNT is unreliable for some codecs → count by grabbing
        total = 0
        while cap.grab():
            total += 1
        cap.release()
        cap = cv2.VideoCapture(video_path)
    total = max(total, 1)

    # how many frames the target fps yields over the whole clip; cap by max_frames + availability
    n_at_fps = max(1, int(round(total * float(fps) / max(1.0, src_fps))))
    n = max(2, min(int(max_frames), n_at_fps, total))
    want = {int(round(x)) for x in np.linspace(0, total - 1, n)}

    saved, idx = [], 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if idx in want:
            path = os.path.join(out_dir, f"{len(saved):06d}.jpg")
            cv2.imwrite(path, frame)
            saved.append(path)
        idx += 1
    cap.release()
    return saved, out_dir


def preprocess_paths(
    paths: Sequence[str],
    image_size: int = 518,
    patch_size: int = 14,
) -> torch.Tensor:
    """Preprocess images to the model's canonical input — a **torchvision-free**, faithful
    reimplementation of upstream ``load_and_preprocess_images`` (crop mode).

    The ZeroGPU base image ships torch but NOT torchvision, and upstream's loader imports
    ``torchvision.transforms`` only for ``ToTensor``. We replicate the exact logic (EXIF
    transpose, alpha-on-white, width=518, height→multiple of patch_size, center-crop height,
    pad mismatched shapes with white) and swap ToTensor for a numpy→torch equivalent.

    Returns a tensor of shape [S, 3, H, W] in range [0, 1].
    """
    from PIL import Image, ImageOps

    def to_tensor_rgb(pil_img):
        # torchvision.transforms.ToTensor() equivalent for an 8-bit RGB image:
        # HWC uint8 [0,255] -> CHW float32 [0,1]
        arr = np.array(pil_img, dtype=np.uint8)                   # [H, W, 3] (copy → writable)
        return torch.from_numpy(arr).permute(2, 0, 1).contiguous().float().div_(255.0)

    target = image_size
    imgs = []
    for p in paths:
        img = Image.open(p)
        img = ImageOps.exif_transpose(img)                        # honor phone/camera EXIF
        if img.mode == "RGBA":
            bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
            img = Image.alpha_composite(bg, img)
        img = img.convert("RGB")
        w, h = img.size
        new_w = target
        new_h = round(h * (new_w / w) / patch_size) * patch_size
        img = img.resize((new_w, new_h), Image.Resampling.BICUBIC)
        t = to_tensor_rgb(img)                                    # [3, new_h, new_w]
        if new_h > target:                                        # center-crop height
            sy = (new_h - target) // 2
            t = t[:, sy:sy + target, :]
        imgs.append(t)

    # If frames differ in shape (mixed aspect ratios), pad to the max with white (1.0).
    shapes = {(t.shape[1], t.shape[2]) for t in imgs}
    if len(shapes) > 1:
        max_h = max(s[0] for s in shapes)
        max_w = max(s[1] for s in shapes)
        out = []
        for t in imgs:
            ph, pw = max_h - t.shape[1], max_w - t.shape[2]
            if ph > 0 or pw > 0:
                pt, pb = ph // 2, ph - ph // 2
                pl, pr = pw // 2, pw - pw // 2
                t = torch.nn.functional.pad(t, (pl, pr, pt, pb), mode="constant", value=1.0)
            out.append(t)
        imgs = out

    return torch.stack(imgs)


# =============================================================================
# Inference + post-processing
# =============================================================================

def infer(
    model: GCTStream,
    images: torch.Tensor,
    num_scale_frames: int = 8,
    keyframe_interval: int = 1,
    output_device: Optional[torch.device] = None,
) -> dict:
    """Run streaming inference under ``no_grad`` (+ autocast on CUDA).

    The model is expected to already be on its target device. dtype is chosen from
    that device: bf16 on capability>=8 GPUs (all ZeroGPU hardware), fp16 on older
    GPUs, fp32 on CPU.
    """
    device = next(model.parameters()).device
    images = images.to(device)

    if device.type == "cuda":
        dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
        with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
            return model.inference_streaming(
                images,
                num_scale_frames=num_scale_frames,
                keyframe_interval=keyframe_interval,
                output_device=output_device,
            )
    with torch.no_grad():
        return model.inference_streaming(
            images,
            num_scale_frames=num_scale_frames,
            keyframe_interval=keyframe_interval,
            output_device=output_device,
        )


# ---- copied verbatim from upstream demo.py (Apache-2.0) so the Space is
# ---- self-contained and does not depend on demo.py being importable. ----

_BATCHED_NDIMS = {
    "pose_enc": 3,
    "depth": 5,
    "depth_conf": 4,
    "world_points": 5,
    "world_points_conf": 4,
    "extrinsic": 4,
    "intrinsic": 4,
    "chunk_scales": 2,
    "chunk_transforms": 4,
    "images": 5,
}


def _squeeze_single_batch(key, value):
    """Drop the leading batch dimension for single-sequence demo outputs."""
    batched_ndim = _BATCHED_NDIMS.get(key)
    if batched_ndim is None or not hasattr(value, "ndim"):
        return value
    if value.ndim == batched_ndim and value.shape[0] == 1:
        return value[0]
    return value


def postprocess(predictions, images):
    """Decode pose encoding → (w2c) extrinsics + intrinsics, move everything to CPU.

    COORDINATE CONVENTION (important): ``pose_encoding_to_extri_intri`` returns
    *world-to-camera* (w2c) extrinsics (OpenCV "camera-from-world"). We KEEP w2c —
    unlike upstream demo.py, which inverts to c2w for its **viser** viewer.
    ``predictions_to_glb`` expects **w2c** (it inverts internally to place camera
    frustums and to align the scene), so handing it c2w would misplace every camera
    and flip the scene's orientation. Keeping w2c is the correct input for the GLB path.
    """
    extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:])

    predictions["extrinsic"] = extrinsic      # w2c, [B, S, 3, 4]
    predictions["intrinsic"] = intrinsic
    predictions.pop("pose_enc_list", None)
    predictions.pop("images", None)

    for k in list(predictions.keys()):
        if isinstance(predictions[k], torch.Tensor):
            predictions[k] = _squeeze_single_batch(
                k, predictions[k].to("cpu", non_blocking=True)
            )
    images_cpu = images.to("cpu", non_blocking=True)
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    return predictions, images_cpu


def prepare_for_visualization(predictions, images=None):
    """Convert predictions to the unbatched NumPy format used by vis code."""
    vis_predictions = {}
    for k, v in predictions.items():
        if isinstance(v, torch.Tensor):
            v = _squeeze_single_batch(k, v.detach().cpu())
            vis_predictions[k] = v.numpy()
        elif isinstance(v, np.ndarray):
            vis_predictions[k] = _squeeze_single_batch(k, v)
        else:
            vis_predictions[k] = v

    if images is None:
        images = predictions.get("images")

    if isinstance(images, torch.Tensor):
        images = images.detach().cpu()
    if isinstance(images, np.ndarray):
        images = _squeeze_single_batch("images", images)
    elif isinstance(images, torch.Tensor):
        images = _squeeze_single_batch("images", images).numpy()

    if isinstance(images, torch.Tensor):
        images = images.numpy()

    if images is not None:
        vis_predictions["images"] = images

    return vis_predictions


def reconstruct_predictions(
    model: GCTStream,
    images: torch.Tensor,
    num_scale_frames: int = 8,
    keyframe_interval: int = 1,
    offload: bool = True,
) -> dict:
    """End-to-end GPU path: inference → postprocess → unbatched NumPy vis dict.

    Returns a dict with (at least) world_points [S,H,W,3], world_points_conf [S,H,W],
    extrinsic [S,3,4], images [S,3,H,W] — everything ``predictions_to_glb`` needs.
    Designed to be the body of the ``@spaces.GPU`` call; its NumPy output is cheap to
    cache so the GLB can be re-exported (CPU-only) when the user tweaks conf/show_cam.
    """
    output_device = torch.device("cpu") if offload else None
    preds = infer(
        model, images,
        num_scale_frames=num_scale_frames,
        keyframe_interval=keyframe_interval,
        output_device=output_device,
    )
    images_for_post = preds["images"] if offload else images
    preds, images_cpu = postprocess(preds, images_for_post)
    return prepare_for_visualization(preds, images_cpu)


# =============================================================================
# GLB export (CPU only — no GPU needed)
# =============================================================================

def _camera_centers_world(extrinsic):
    """(camera centers in world coords [S,3], w2c 4x4 [S,4,4]) from w2c [S,3,4]."""
    extr = np.asarray(extrinsic, dtype=np.float64)
    S = len(extr)
    w2c = np.repeat(np.eye(4)[None], S, axis=0)
    w2c[:, :3, :4] = extr
    centers = np.linalg.inv(w2c)[:, :3, 3]
    return centers, w2c


def _alignment_transform(w2c):
    """The same scene alignment predictions_to_glb applies (orient to the first camera +
    flip into the GLB/OpenGL convention). Applied to the WHOLE scene (points + trajectory).
    Verified to match predictions_to_glb's apply_scene_alignment."""
    from scipy.spatial.transform import Rotation
    opengl = np.eye(4); opengl[1, 1] = -1.0; opengl[2, 2] = -1.0
    align = np.eye(4); align[:3, :3] = Rotation.from_euler("y", 180, degrees=True).as_matrix()
    return np.linalg.inv(w2c[0]) @ opengl @ align


def _add_trajectory_tube(scene, centers_world):
    """Add a clean colored tube tracing the camera path (README "track" style) in WORLD
    coords; the caller applies the scene alignment afterwards. gist_rainbow start→end."""
    import trimesh
    import matplotlib
    S = len(centers_world)
    if S < 2:
        return
    ext = float(np.linalg.norm(centers_world.max(0) - centers_world.min(0))) or 1.0
    radius = max(ext * 0.004, 1e-5)
    cmap = matplotlib.colormaps.get_cmap("gist_rainbow")
    tubes = []
    for i in range(S - 1):
        if np.linalg.norm(centers_world[i + 1] - centers_world[i]) < 1e-9:
            continue
        seg = trimesh.creation.cylinder(radius=radius, segment=centers_world[i:i + 2], sections=8)
        seg.visual.vertex_colors = tuple(int(255 * x) for x in cmap(i / (S - 1))[:3]) + (255,)
        tubes.append(seg)
    if tubes:
        scene.add_geometry(trimesh.util.concatenate(tubes))


def _flatten_vis(vis_predictions: dict) -> dict:
    """Flatten a vis dict → a compact flat point set in WORLD coords (the format the GPU
    call returns and the GLB builder consumes). Unprojects depth at full resolution.
    Returns {points (N,3) f32, colors (N,3) u8, conf (N,) f32, extrinsic (S,3,4) f32}.
    """
    vis = vis_predictions
    if "world_points" in vis:
        world = np.asarray(vis["world_points"], np.float32)
        conf = np.asarray(vis.get("world_points_conf"), np.float32)
    else:
        world = np.asarray(
            unproject_depth_map_to_point_map(vis["depth"], vis["extrinsic"], vis["intrinsic"]),
            np.float32,
        )
        conf = np.asarray(vis["depth_conf"], np.float32)
    imgs = np.asarray(vis["images"])
    if imgs.ndim == 4 and imgs.shape[1] == 3:          # NCHW → NHWC
        imgs = np.transpose(imgs, (0, 2, 3, 1))
    return {
        "points": world.reshape(-1, 3),
        "colors": np.clip(imgs.reshape(-1, 3) * 255.0, 0, 255).astype(np.uint8),
        "conf": conf.reshape(-1),
        "extrinsic": np.asarray(vis["extrinsic"], np.float32),
    }


def _downsample_rec(rec: dict, max_points: int, seed: int = 0) -> dict:
    """Drop invalid points and randomly cap to max_points so the GPU→app payload stays
    bounded and independent of frame count."""
    pts, cols, conf = rec["points"], rec["colors"], rec["conf"]
    valid = conf > 1e-5
    pts, cols, conf = pts[valid], cols[valid], conf[valid]
    n_total = int(pts.shape[0])
    if n_total > max_points:
        sel = np.random.default_rng(seed).choice(n_total, max_points, replace=False)
        pts, cols, conf = pts[sel], cols[sel], conf[sel]
    return {"points": pts, "colors": cols, "conf": conf,
            "extrinsic": rec["extrinsic"], "n_total": n_total}


def export_glb(
    vis_predictions: dict,
    out_path: Optional[str] = None,
    conf_thres: float = 50.0,
    show_cam: bool = True,
    mask_sky: bool = False,
    target_dir: Optional[str] = None,
) -> str:
    """Convenience wrapper (smoke_test / local CPU runs): flatten a vis dict → build the GLB.
    The Space uses reconstruct_points() + build_glb_from_points() directly (compact payload)."""
    rec = _flatten_vis(vis_predictions)
    return build_glb_from_points(rec, out_path=out_path, conf_thres=conf_thres, show_cam=show_cam)


def build_glb_from_points(
    rec: dict,
    out_path: Optional[str] = None,
    conf_thres: float = 60.0,
    show_cam: bool = True,
    display_max: int = 2_000_000,
    seed: int = 0,
) -> str:
    """Build the .glb from a flat point set: percentile-filter by confidence, downsample for
    smooth in-browser rendering, add the camera-path line, and apply the scene alignment.

    ``conf_thres`` is a *percentile* — points below the conf_thres-th percentile of
    confidence are dropped. 0 keeps everything; higher = fewer points.
    """
    import trimesh
    if out_path is None:
        out_path = os.path.join(tempfile.mkdtemp(prefix="lingbot_glb_"), "scene.glb")

    pts = np.asarray(rec["points"], np.float64)
    cols = np.asarray(rec["colors"])
    conf = np.asarray(rec["conf"], np.float64)

    if conf.size:
        thr = np.percentile(conf, conf_thres) if conf_thres > 0 else 0.0
        mask = (conf >= thr) & (conf > 1e-5)
        pts, cols = pts[mask], cols[mask]
    if len(pts) > display_max:                          # browser viewer bogs down past a few M points
        sel = np.random.default_rng(seed).choice(len(pts), display_max, replace=False)
        pts, cols = pts[sel], cols[sel]
    if len(pts) == 0:                                   # degenerate guard
        pts, cols = np.zeros((1, 3)), np.array([[255, 255, 255]])

    scene = trimesh.Scene()
    scene.add_geometry(trimesh.PointCloud(vertices=pts, colors=cols))
    centers, w2c = _camera_centers_world(rec["extrinsic"])
    if show_cam:
        _add_trajectory_tube(scene, centers)
    scene.apply_transform(_alignment_transform(w2c))    # align points + trajectory together
    scene.export(out_path)
    return out_path


def count_points(rec: dict, conf_thres: float = 60.0) -> Tuple[int, int]:
    """(kept, total) after the same percentile filter build_glb_from_points uses."""
    conf = np.asarray(rec["conf"]).reshape(-1)
    total = conf.size
    if total == 0:
        return 0, 0
    thr = np.percentile(conf, conf_thres) if conf_thres > 0 else 0.0
    kept = int(np.count_nonzero((conf >= thr) & (conf > 1e-5)))
    return kept, total


def count_visible_points(vis_predictions: dict, conf_thres: float = 60.0) -> Tuple[int, int]:
    """Back-compat wrapper for the vis-dict form (smoke_test)."""
    return count_points(_flatten_vis(vis_predictions), conf_thres)


# =============================================================================
# GPU entry point that returns a COMPACT, frame-count-independent payload
# =============================================================================

def reconstruct_points(
    model: GCTStream,
    images: torch.Tensor,
    num_scale_frames: int = 8,
    keyframe_interval: int = 1,
    max_return_points: int = 6_000_000,
) -> dict:
    """End-to-end GPU path → a COMPACT flat point set whose size is bounded (independent of
    frame count). This is what lets the Space return many frames across the @spaces.GPU
    boundary without a huge per-pixel payload.

    Returns {points (M,3) f32, colors (M,3) u8, conf (M,) f32, extrinsic (S,3,4) f32, n_total}.
    """
    vis = reconstruct_predictions(
        model, images, num_scale_frames=num_scale_frames,
        keyframe_interval=keyframe_interval, offload=True,
    )
    return _downsample_rec(_flatten_vis(vis), max_return_points)