from shared.utils.phase_progress import control_video_encoding
import hashlib
import os
import shutil
import time
import zipfile
from contextlib import nullcontext
from pathlib import Path

import cv2
import imageio
import numpy as np
import torch
import torch.nn.functional as F
from einops import rearrange
from PIL import Image

from .camera import get_plucker_embedding

VISTA4D_REUSE_GENERATED_MAP_IN_MEMORY = True
VISTA4D_DA3_PROCESS_RES = 896
VISTA4D_DA3_CHUNK_SIZE = -1
VISTA4D_DA3_CHUNK_OVERLAP = 8
_VISTA4D_MAP_CACHE = {}
_VISTA4D_MAP_CACHE_MAX_ITEMS = 2


def _center_slice(length, frame_num):
    if length < frame_num:
        raise ValueError(f"Vista4D input needs at least {frame_num} frames, got {length}.")
    start = (length - frame_num) // 2
    return slice(start, start + frame_num)


def _tensor_to_video_np(frames):
    video = frames.detach().cpu().float().permute(1, 2, 3, 0).numpy()
    return ((video + 1.0) * 127.5).clip(0, 255).astype(np.uint8)


def _video_np_to_tensor(video, device, dtype):
    video = torch.from_numpy(video).permute(3, 0, 1, 2).float().div_(127.5).sub_(1.0)
    return video.to(device=device, dtype=dtype)


def _crop_resize_video(video, height, width, resample=Image.Resampling.LANCZOS):
    if video.shape[1:3] == (height, width):
        return video.copy()
    frames = []
    for frame in video:
        image = Image.fromarray(frame)
        in_w, in_h = image.size
        if in_h / in_w > height / width:
            crop_h = int(in_w * height / width)
            top = (in_h - crop_h) // 2
            image = image.crop((0, top, in_w, top + crop_h))
        else:
            crop_w = int(in_h * width / height)
            left = (in_w - crop_w) // 2
            image = image.crop((left, 0, left + crop_w, in_h))
        frames.append(np.asarray(image.resize((width, height), resample)))
    return np.stack(frames, axis=0)


def _resize_intrinsics(intrinsics, height, width, height_input, width_input):
    intrinsics = intrinsics.copy()
    if height_input / width_input > height / width:
        crop_h = int(width_input * height / width)
        intrinsics[..., 3] -= (height_input - crop_h) / 2
        height_input = crop_h
    else:
        crop_w = int(height_input * width / height)
        intrinsics[..., 2] -= (width_input - crop_w) / 2
        width_input = crop_w
    intrinsics[..., 0] *= width / width_input
    intrinsics[..., 1] *= height / height_input
    intrinsics[..., 2] *= width / width_input
    intrinsics[..., 3] *= height / height_input
    return intrinsics


def _infer_intrinsics_size(intrinsics, fallback_height, fallback_width):
    cx = float(np.nanmedian(intrinsics[..., 2]))
    cy = float(np.nanmedian(intrinsics[..., 3]))
    if np.isfinite(cx) and np.isfinite(cy) and cx > 0 and cy > 0:
        return max(1, int(round(cy * 2))), max(1, int(round(cx * 2)))
    return fallback_height, fallback_width


def _read_video(path):
    cap = cv2.VideoCapture(os.fspath(path))
    if not cap.isOpened():
        raise ValueError(f"Could not open Vista4D video: {path}")
    frames = []
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    cap.release()
    if not frames:
        raise ValueError(f"Vista4D video has no frames: {path}")
    return np.stack(frames, axis=0)


def _save_video(path, video, fps):
    os.makedirs(os.path.dirname(os.fspath(path)), exist_ok=True)
    writer = imageio.get_writer(os.fspath(path), fps=fps, quality=9, macro_block_size=1)
    try:
        for frame in video:
            writer.append_data(frame)
    finally:
        writer.close()


def _load_masks(folder, frame_num, height, width, default):
    folder = Path(folder)
    if not folder.is_dir():
        fill = np.ones if default else np.zeros
        return fill((frame_num, height, width), dtype=np.bool_)
    files = sorted(folder.glob("*.png"))
    if len(files) == 0:
        raise ValueError(f"Vista4D mask folder is empty: {folder}")
    masks = []
    for file in files:
        frame = cv2.imread(os.fspath(file), cv2.IMREAD_GRAYSCALE)
        if frame is None:
            raise ValueError(f"Could not load Vista4D mask: {file}")
        masks.append(frame > 127)
    masks = np.stack(masks, axis=0)
    masks = masks[_center_slice(masks.shape[0], frame_num)]
    if masks.shape[1:3] != (height, width):
        masks = F.interpolate(torch.from_numpy(masks[:, None].astype(np.float32)), size=(height, width), mode="nearest")[:, 0].numpy() > 0.5
    return masks


def _save_masks(folder, masks):
    folder = Path(folder)
    folder.mkdir(parents=True, exist_ok=True)
    for idx, mask in enumerate(masks):
        cv2.imwrite(os.fspath(folder / f"{idx:05d}.png"), mask.astype(np.uint8) * 255)


def _pixel_grid(height, width):
    ys, xs = np.meshgrid(np.arange(height, dtype=np.float32), np.arange(width, dtype=np.float32), indexing="ij")
    return xs.reshape(-1), ys.reshape(-1)


def _frame_points(video, depths, sky_mask, dynamic_mask, cam_c2w, intrinsics, frame_idx, xs, ys):
    depth = depths[frame_idx].reshape(-1).astype(np.float32)
    valid = np.isfinite(depth) & (depth > 1e-5)
    if sky_mask is not None:
        valid &= ~sky_mask[frame_idx].reshape(-1).astype(np.bool_)
    if not np.any(valid):
        return np.empty((0, 3), dtype=np.float32), np.empty((0, 3), dtype=np.uint8), np.empty((0,), dtype=np.bool_)
    fx, fy, cx, cy = intrinsics[frame_idx]
    z = depth[valid]
    cam_points = np.stack(((xs[valid] - cx) * z / fx, (ys[valid] - cy) * z / fy, z), axis=1)
    world_points = cam_points @ cam_c2w[frame_idx, :3, :3].T + cam_c2w[frame_idx, :3, 3]
    colors = video[frame_idx].reshape(-1, 3)[valid]
    point_dynamic = dynamic_mask[frame_idx].reshape(-1)[valid] if dynamic_mask is not None else np.zeros(colors.shape[0], dtype=np.bool_)
    return world_points.astype(np.float32), colors, point_dynamic.astype(np.bool_)


def _project_points(points, colors, cam_w2c, intrinsics, height, width):
    if points.shape[0] == 0:
        return None
    cam_points = points @ cam_w2c[:3, :3].T + cam_w2c[:3, 3]
    z = cam_points[:, 2]
    valid = z > 1e-5
    if not np.any(valid):
        return None
    fx, fy, cx, cy = intrinsics
    valid_idx = np.nonzero(valid)[0]
    x = np.rint(cam_points[valid, 0] * fx / z[valid] + cx).astype(np.int32)
    y = np.rint(cam_points[valid, 1] * fy / z[valid] + cy).astype(np.int32)
    inside = (x >= 0) & (x < width) & (y >= 0) & (y < height)
    if not np.any(inside):
        return None
    return y[inside] * width + x[inside], z[valid][inside], colors[valid][inside], valid_idx[inside]


def _render_point_chunks(point_chunks, cam_w2c, intrinsics, height, width):
    empty_frame = np.zeros((height, width, 3), dtype=np.uint8)
    empty_mask = np.zeros((height, width), dtype=np.bool_)
    zbuf = np.full(height * width, np.inf, dtype=np.float32)
    projected = []
    for points, colors, is_dynamic in point_chunks:
        item = _project_points(points, colors, cam_w2c, intrinsics, height, width)
        if item is None:
            continue
        pix, z, colors, _ = item
        np.minimum.at(zbuf, pix, z)
        projected.append((pix, z, colors, is_dynamic))
    if len(projected) == 0:
        return empty_frame, empty_mask, empty_mask.copy()
    output = np.zeros((height * width, 3), dtype=np.uint8)
    dynamic = np.zeros(height * width, dtype=np.bool_)
    for pix, z, colors, is_dynamic in projected:
        keep = z <= zbuf[pix] + 1e-4
        output[pix[keep]] = colors[keep]
        if is_dynamic:
            dynamic[pix[keep]] = True
    mask = np.isfinite(zbuf).reshape(height, width)
    return output.reshape(height, width, 3), mask, dynamic.reshape(height, width) & mask


def _to_torch_point_cache(point_cache, device):
    return [(torch.from_numpy(points).to(device=device, dtype=torch.float32), torch.from_numpy(colors).to(device=device)) for points, colors in point_cache]


def _project_points_torch(points, colors, cam_w2c, intrinsics, height, width):
    if points.shape[0] == 0:
        return None
    cam_points = points @ cam_w2c[:3, :3].T + cam_w2c[:3, 3]
    z_all = cam_points[:, 2]
    valid = z_all > 1e-5
    if not bool(valid.any()):
        return None
    fx, fy, cx, cy = intrinsics.unbind(0)
    valid_idx = torch.nonzero(valid, as_tuple=False).squeeze(1)
    z = z_all[valid_idx]
    x = torch.round(cam_points[valid_idx, 0] * fx / z + cx).to(torch.int64)
    y = torch.round(cam_points[valid_idx, 1] * fy / z + cy).to(torch.int64)
    inside = (x >= 0) & (x < width) & (y >= 0) & (y < height)
    if not bool(inside.any()):
        return None
    return y[inside] * width + x[inside], z[inside], colors[valid_idx[inside]]


def _render_point_chunks_torch(point_chunks, cam_w2c, intrinsics, height, width):
    device = cam_w2c.device
    zbuf = torch.full((height * width,), float("inf"), dtype=torch.float32, device=device)
    projected = []
    order_start = 0
    for points, colors, is_dynamic in point_chunks:
        item = _project_points_torch(points, colors, cam_w2c, intrinsics, height, width)
        if item is None:
            order_start += points.shape[0]
            continue
        pix, z, colors = item
        order = torch.arange(order_start, order_start + pix.shape[0], dtype=torch.int32, device=device)
        zbuf.scatter_reduce_(0, pix, z, reduce="amin", include_self=True)
        projected.append((pix, z, colors, order, is_dynamic))
        order_start += points.shape[0]

    frame = torch.zeros((height * width, 3), dtype=torch.uint8, device=device)
    dynamic = torch.zeros((height * width,), dtype=torch.bool, device=device)
    if len(projected) > 0:
        winner_order = torch.full((height * width,), -1, dtype=torch.int32, device=device)
        for pix, z, _, order, is_dynamic in projected:
            keep = z <= zbuf[pix] + 1e-4
            if keep.any():
                winner_order.scatter_reduce_(0, pix[keep], order[keep], reduce="amax", include_self=True)
                if is_dynamic:
                    dynamic[pix[keep]] = True
        for pix, _, colors, order, _ in projected:
            selected = order == winner_order[pix]
            if selected.any():
                frame[pix[selected]] = colors[selected]

    mask = torch.isfinite(zbuf)
    return frame.reshape(height, width, 3).cpu().numpy(), mask.reshape(height, width).cpu().numpy(), (dynamic & mask).reshape(height, width).cpu().numpy()


def _render_point_cloud_video_torch(static_cache, dynamic_cache, target_cam_w2c, target_intrinsics, height, width, window_radius):
    device = torch.device("cuda")
    static_cache = _to_torch_point_cache(static_cache, device)
    dynamic_cache = _to_torch_point_cache(dynamic_cache, device)
    target_cam_w2c = torch.from_numpy(target_cam_w2c).to(device=device, dtype=torch.float32)
    target_intrinsics = torch.from_numpy(target_intrinsics).to(device=device, dtype=torch.float32)
    rendered = []
    alpha_masks = []
    dynamic_masks = []
    context = torch.autocast(device_type="cuda", enabled=False) if device.type == "cuda" else nullcontext()
    with torch.no_grad(), context:
        for target_idx in range(len(dynamic_cache)):
            if window_radius < 0:
                point_chunks = [(points, colors, False) for points, colors in static_cache]
            else:
                start = max(0, target_idx - window_radius)
                end = min(len(dynamic_cache), target_idx + window_radius + 1)
                point_chunks = [(static_cache[idx][0], static_cache[idx][1], False) for idx in range(start, end) if static_cache[idx][0].shape[0] > 0]
            dynamic_points, dynamic_colors = dynamic_cache[target_idx]
            if dynamic_points.shape[0] > 0:
                point_chunks.append((dynamic_points, dynamic_colors, True))
            frame, mask, rendered_dynamic = _render_point_chunks_torch(point_chunks, target_cam_w2c[target_idx], target_intrinsics[target_idx], height, width)
            rendered.append(frame)
            alpha_masks.append(mask)
            dynamic_masks.append(rendered_dynamic)
    return np.stack(rendered, axis=0), np.stack(alpha_masks, axis=0), np.stack(dynamic_masks, axis=0)


def _render_point_cloud_video(video, depths, sky_mask, dynamic_mask, cam_c2w, intrinsics, target_cam_c2w=None, target_intrinsics=None, window_radius=-1):
    frame_num, height, width = video.shape[:3]
    if target_cam_c2w is None:
        target_cam_c2w = cam_c2w
    if target_intrinsics is None:
        target_intrinsics = intrinsics
    xs, ys = _pixel_grid(height, width)
    static_cache = []
    dynamic_cache = []
    for idx in range(frame_num):
        points, colors, point_dynamic = _frame_points(video, depths, sky_mask, dynamic_mask, cam_c2w, intrinsics, idx, xs, ys)
        static = ~point_dynamic
        static_cache.append((points[static], colors[static]))
        dynamic_cache.append((points[point_dynamic], colors[point_dynamic]))
    target_cam_w2c = np.linalg.inv(target_cam_c2w).astype(np.float32)

    rendered = []
    alpha_masks = []
    dynamic_masks = []
    window_radius = int(window_radius)
    if window_radius < 0:
        static_points = [points for points, _ in static_cache if points.shape[0] > 0]
        static_colors = [colors for points, colors in static_cache if points.shape[0] > 0]
        static_cache = [(np.concatenate(static_points, axis=0), np.concatenate(static_colors, axis=0))] if static_points else []
    if torch.cuda.is_available():
        return _render_point_cloud_video_torch(static_cache, dynamic_cache, target_cam_w2c, target_intrinsics, height, width, window_radius)
    for target_idx in range(frame_num):
        if window_radius < 0:
            point_chunks = [(points, colors, False) for points, colors in static_cache]
        else:
            start = max(0, target_idx - window_radius)
            end = min(frame_num, target_idx + window_radius + 1)
            point_chunks = [(static_cache[idx][0], static_cache[idx][1], False) for idx in range(start, end) if static_cache[idx][0].shape[0] > 0]
        dynamic_points, dynamic_colors = dynamic_cache[target_idx]
        if dynamic_points.shape[0] > 0:
            point_chunks.append((dynamic_points, dynamic_colors, True))
        frame, mask, rendered_dynamic = _render_point_chunks(point_chunks, target_cam_w2c[target_idx], target_intrinsics[target_idx], height, width)
        rendered.append(frame)
        alpha_masks.append(mask)
        dynamic_masks.append(rendered_dynamic)
    return np.stack(rendered, axis=0), np.stack(alpha_masks, axis=0), np.stack(dynamic_masks, axis=0)


def _load_cameras(path, frame_num):
    data = np.load(path)
    cam_c2w = data["cam_c2w"][_center_slice(data["cam_c2w"].shape[0], frame_num)]
    intrinsics = data["intrinsics"][_center_slice(data["intrinsics"].shape[0], frame_num)]
    return cam_c2w.astype(np.float32), intrinsics.astype(np.float32)


def _save_cameras(path, cam_c2w, intrinsics):
    os.makedirs(os.path.dirname(os.fspath(path)), exist_ok=True)
    np.savez(path, cam_c2w=cam_c2w.astype(np.float32), intrinsics=intrinsics.astype(np.float32))


def _is_camera_npz(path):
    if Path(path).suffix.lower() != ".npz":
        return False
    try:
        with np.load(path) as data:
            return "cam_c2w" in data.files and "intrinsics" in data.files
    except Exception:
        return False


def _extract_zip(path):
    digest = hashlib.sha1(os.path.abspath(path).encode("utf-8")).hexdigest()[:12]
    target = Path("ckpts") / "temp" / "vista4d_maps" / digest
    if target.is_dir():
        return target
    target.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(path) as archive:
        archive.extractall(target)
    roots = [item for item in target.iterdir() if item.is_dir()]
    if len(roots) == 1 and (roots[0] / "video_src.mp4").is_file():
        return roots[0]
    return target


def _resolve_custom_input(input_custom):
    if input_custom is None or str(input_custom).strip() == "":
        return None, None
    path = Path(input_custom)
    if path.is_file() and _is_camera_npz(path):
        return None, path
    if path.is_file() and path.suffix.lower() == ".zip":
        return _extract_zip(path), None
    if path.is_file():
        path = path.parent
    if not path.is_dir():
        raise ValueError(f"Vista4D custom_guide must be a preprocessed map folder/zip or target camera .npz, got: {input_custom}")
    return path, None


def _default_output_dir():
    from wgp import save_path
    return save_path


def _hash_array(hasher, array):
    array = np.ascontiguousarray(array)
    hasher.update(str(array.shape).encode("utf-8"))
    hasher.update(str(array.dtype).encode("utf-8"))
    hasher.update(array.view(np.uint8))


def _hash_file(path):
    hasher = hashlib.sha1()
    with open(path, "rb") as reader:
        for chunk in iter(lambda: reader.read(1024 * 1024), b""):
            hasher.update(chunk)
    return hasher.hexdigest()


def _make_generated_map_cache_key(video, frame_num, height, width, fps, process_res, chunk_size, chunk_overlap, custom_settings, target_camera_path, model_mode):
    hasher = hashlib.sha1()
    hasher.update(b"vista4d-memory-map-v2")
    hasher.update(f"{frame_num}|{height}|{width}|{fps}|{process_res}|{chunk_size}|{chunk_overlap}".encode("utf-8"))
    if target_camera_path is None:
        hasher.update(str(model_mode or "").encode("utf-8"))
    if isinstance(custom_settings, dict):
        hasher.update(str(custom_settings.get("vista4d_seg_keywords", "_all_")).encode("utf-8"))
        hasher.update(str(custom_settings.get("vista4d_scene_scale", 1.0)).encode("utf-8"))
        if target_camera_path is None:
            hasher.update(str(custom_settings.get("vista4d_camera_strength", 100.0)).encode("utf-8"))
    if target_camera_path is not None:
        hasher.update(os.path.abspath(os.fspath(target_camera_path)).encode("utf-8"))
        hasher.update(_hash_file(target_camera_path).encode("utf-8"))
    _hash_array(hasher, video)
    return hasher.hexdigest()


def _cache_generated_map(cache_key, data):
    _VISTA4D_MAP_CACHE[cache_key] = data
    while len(_VISTA4D_MAP_CACHE) > _VISTA4D_MAP_CACHE_MAX_ITEMS:
        _VISTA4D_MAP_CACHE.pop(next(iter(_VISTA4D_MAP_CACHE)))


def _parse_keywords(custom_settings):
    raw = "_all_"
    if isinstance(custom_settings, dict):
        raw = str(custom_settings.get("vista4d_seg_keywords", raw))
    return [item.strip() for item in raw.replace(";", ",").split(",") if item.strip()]


def _camera_strength(custom_settings):
    if not isinstance(custom_settings, dict):
        return 1.0
    return float(custom_settings.get("vista4d_camera_strength", 100.0) or 0.0) / 100.0


def _get_da3_settings():
    return VISTA4D_DA3_PROCESS_RES, VISTA4D_DA3_CHUNK_SIZE, VISTA4D_DA3_CHUNK_OVERLAP


def _save_generated_map_to_disk(data, source_cam_c2w, source_intrinsics, depths, sky_mask, fps):
    map_dir = Path(_default_output_dir()) / f"vista4d_map_{time.strftime('%Y%m%d_%H%M%S')}"
    map_dir.mkdir(parents=True, exist_ok=True)
    _save_video(map_dir / "video_src.mp4", data["source_video"], fps=fps)
    _save_video(map_dir / "video_pc.mp4", data["point_cloud_video"], fps=fps)
    _save_masks(map_dir / "alpha_mask_src", data["source_alpha_mask"])
    _save_masks(map_dir / "dynamic_mask_src", data["source_motion_mask"])
    _save_masks(map_dir / "static_mask_src", data["source_alpha_mask"] & ~data["source_motion_mask"])
    _save_masks(map_dir / "alpha_mask_pc", data["point_cloud_alpha_mask"])
    _save_masks(map_dir / "dynamic_mask_pc", data["point_cloud_motion_mask"])
    _save_masks(map_dir / "static_mask_pc", data["point_cloud_alpha_mask"] & ~data["point_cloud_motion_mask"])
    _save_masks(map_dir / "sky_mask_src", sky_mask.astype(np.bool_))
    _save_cameras(map_dir / "cameras_src.npz", source_cam_c2w, source_intrinsics)
    _save_cameras(map_dir / "cameras_tgt.npz", data["cam_c2w"], data["intrinsics"])
    np.save(map_dir / "depths_da3.npy", depths.astype(np.float16))
    shutil.make_archive(os.fspath(map_dir), "zip", root_dir=map_dir)
    print(f"Vista4D map saved to: {map_dir}")
    return map_dir


def _rotation_x(angle):
    c, s = np.cos(angle), np.sin(angle)
    return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float32)


def _rotation_y(angle):
    c, s = np.cos(angle), np.sin(angle)
    return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float32)


def _estimate_focus_depth(depths, sky_mask):
    valid = np.isfinite(depths) & (depths > 1e-5)
    if sky_mask is not None:
        valid &= ~sky_mask.astype(np.bool_)
    if not np.any(valid):
        return 1.0
    samples = depths[valid]
    if samples.size > 200000:
        samples = samples[np.linspace(0, samples.size - 1, 200000, dtype=np.int64)]
    return max(float(np.median(samples)), 1e-3)


def _apply_local_camera_delta(cam_c2w, translations, rotations):
    target = cam_c2w.copy()
    for idx in range(cam_c2w.shape[0]):
        delta = np.eye(4, dtype=np.float32)
        delta[:3, :3] = rotations[idx]
        delta[:3, 3] = translations[idx]
        target[idx] = cam_c2w[idx] @ delta
    return target


def _generate_target_camera_trajectory(cam_c2w, intrinsics, depths, sky_mask, model_mode, strength):
    mode = str(model_mode or "dolly_zoom")
    if mode not in {
        "dolly_zoom", "left_front_zoom", "right_front_zoom", "close_crane_above", "close_crane_below", "arc_right_45", "arc_left_45",
        "push_in", "pull_back", "truck_right", "truck_left", "pedestal_up", "pedestal_down",
        "pan_right_45", "pan_left_45", "tilt_up_45", "tilt_down_45", "zoom_in", "zoom_out", "bird_view",
        "crane_above_right", "crane_above_left", "crane_below_right", "crane_below_left",
    }:
        mode = "dolly_zoom"
    focus = _estimate_focus_depth(depths, sky_mask)
    frame_num = cam_c2w.shape[0]
    progress = np.linspace(0.0, 1.0, frame_num, dtype=np.float32) * np.float32(strength)
    translations = np.zeros((frame_num, 3), dtype=np.float32)
    rotations = np.repeat(np.eye(3, dtype=np.float32)[None], frame_num, axis=0)
    target_intrinsics = intrinsics.copy()
    focal_scale = np.ones(frame_num, dtype=np.float32)

    if mode == "dolly_zoom":
        translations[:, 2] = 0.35 * focus * progress
        focal_scale = np.clip((focus - translations[:, 2]) / focus, 0.25, 4.0).astype(np.float32)
    elif mode in ("left_front_zoom", "right_front_zoom"):
        translations[:, 0] = (-0.22 if mode == "left_front_zoom" else 0.22) * focus * progress
        translations[:, 2] = 0.20 * focus * progress
        focal_scale = np.clip(1.0 + 0.25 * progress, 0.25, 4.0).astype(np.float32)
    elif mode in ("close_crane_above", "close_crane_below"):
        translations[:, 1] = (-0.22 if mode == "close_crane_above" else 0.22) * focus * progress
        translations[:, 2] = 0.18 * focus * progress
        for idx, amount in enumerate(progress):
            rotations[idx] = _rotation_x(np.deg2rad(-15.0 if mode == "close_crane_above" else 15.0) * amount)
    elif mode in ("arc_right_45", "arc_left_45"):
        sign = -1.0 if mode == "arc_right_45" else 1.0
        angles = np.deg2rad(45.0) * progress * sign
        translations[:, 0] = -focus * np.sin(angles)
        translations[:, 2] = focus * (1.0 - np.cos(angles))
        for idx, angle in enumerate(angles):
            rotations[idx] = _rotation_y(angle)
    elif mode == "push_in":
        translations[:, 2] = 0.25 * focus * progress
    elif mode == "pull_back":
        translations[:, 2] = -0.25 * focus * progress
    elif mode == "truck_right":
        translations[:, 0] = 0.25 * focus * progress
    elif mode == "truck_left":
        translations[:, 0] = -0.25 * focus * progress
    elif mode == "pedestal_up":
        translations[:, 1] = -0.25 * focus * progress
    elif mode == "pedestal_down":
        translations[:, 1] = 0.25 * focus * progress
    elif mode == "pan_right_45":
        for idx, amount in enumerate(progress):
            rotations[idx] = _rotation_y(np.deg2rad(45.0) * amount)
    elif mode == "pan_left_45":
        for idx, amount in enumerate(progress):
            rotations[idx] = _rotation_y(np.deg2rad(-45.0) * amount)
    elif mode == "tilt_up_45":
        for idx, amount in enumerate(progress):
            rotations[idx] = _rotation_x(np.deg2rad(45.0) * amount)
    elif mode == "tilt_down_45":
        for idx, amount in enumerate(progress):
            rotations[idx] = _rotation_x(np.deg2rad(-45.0) * amount)
    elif mode == "zoom_in":
        focal_scale = np.clip(1.0 + 0.60 * progress, 0.25, 4.0).astype(np.float32)
    elif mode == "zoom_out":
        focal_scale = np.clip(1.0 - 0.35 * progress, 0.25, 4.0).astype(np.float32)
    elif mode == "bird_view":
        translations[:, 1] = -0.70 * focus * progress
        translations[:, 2] = -0.12 * focus * progress
        for idx, amount in enumerate(progress):
            rotations[idx] = _rotation_x(np.deg2rad(-60.0) * amount)
    elif mode in ("crane_above_right", "crane_above_left", "crane_below_right", "crane_below_left"):
        above = "above" in mode
        right = "right" in mode
        translations[:, 0] = (0.18 if right else -0.18) * focus * progress
        translations[:, 1] = (-0.24 if above else 0.24) * focus * progress
        translations[:, 2] = 0.16 * focus * progress
        yaw = np.deg2rad(-12.0 if right else 12.0)
        pitch = np.deg2rad(-18.0 if above else 18.0)
        for idx, amount in enumerate(progress):
            rotations[idx] = _rotation_y(yaw * amount) @ _rotation_x(pitch * amount)

    target_intrinsics[:, 0:2] *= focal_scale[:, None]
    return _apply_local_camera_delta(cam_c2w, translations, rotations), target_intrinsics


def _create_map_from_control_video(input_frames, frame_num, height, width, fps, custom_settings, target_camera_path=None, model_mode=None):
    video = _tensor_to_video_np(input_frames)
    video = video[_center_slice(video.shape[0], frame_num)]
    height_input, width_input = video.shape[1:3]
    # video = _crop_resize_video(video, height, width)

    scene_scale = 1.0
    if isinstance(custom_settings, dict):
        scene_scale = float(custom_settings.get("vista4d_scene_scale", 1.0) or 1.0)
    process_res, chunk_size, chunk_overlap = _get_da3_settings()
    if process_res <= 0:
        process_res = width

    from preprocessing.depth_anything_v3.depth import resolve_da3_chunk_size, run_da3_reconstruction
    chunk_size = resolve_da3_chunk_size(chunk_size)
    cache_key = None
    if VISTA4D_REUSE_GENERATED_MAP_IN_MEMORY:
        cache_key = _make_generated_map_cache_key(video, frame_num, height, width, fps, process_res, chunk_size, chunk_overlap, custom_settings, target_camera_path, model_mode)
        cached = _VISTA4D_MAP_CACHE.get(cache_key)
        if cached is not None:
            print("Vista4D map reused from memory cache.")
            return cached

    depths, sky_mask, cam_c2w, intrinsics = run_da3_reconstruction(video, process_res=process_res, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    if scene_scale != 1.0:
        depths *= scene_scale
        cam_c2w[:, :3, 3] *= scene_scale
    source_cam_c2w = cam_c2w
    source_intrinsics = intrinsics
    target_cam_c2w = cam_c2w
    target_intrinsics = intrinsics
    if target_camera_path is not None:
        target_cam_c2w, target_intrinsics = _load_cameras(target_camera_path, frame_num)
        camera_height, camera_width = _infer_intrinsics_size(target_intrinsics, height_input, width_input)
        if (camera_height, camera_width) != (height, width):
            target_intrinsics = _resize_intrinsics(target_intrinsics, height, width, camera_height, camera_width)
        if scene_scale != 1.0:
            target_cam_c2w[:, :3, 3] *= scene_scale
    else:
        target_cam_c2w, target_intrinsics = _generate_target_camera_trajectory(cam_c2w, intrinsics, depths, sky_mask, model_mode, _camera_strength(custom_settings))

    keywords = _parse_keywords(custom_settings)
    if len(keywords) == 0:
        dynamic_mask = np.zeros((frame_num, height, width), dtype=np.bool_)
    elif "_all_" in keywords:
        dynamic_mask = np.ones((frame_num, height, width), dtype=np.bool_)
    else:
        from preprocessing.sam3.preprocessor import run_sam3_video
        dynamic_mask = run_sam3_video(video, keywords)

    source_alpha_mask = np.ones((frame_num, height, width), dtype=np.bool_)
    point_cloud_video, point_cloud_alpha_mask, point_cloud_dynamic_mask = _render_point_cloud_video(video, depths, sky_mask, dynamic_mask, cam_c2w, intrinsics, target_cam_c2w, target_intrinsics)
    data = {
        "source_video": video,
        "point_cloud_video": point_cloud_video,
        "source_alpha_mask": source_alpha_mask,
        "source_motion_mask": dynamic_mask,
        "point_cloud_alpha_mask": point_cloud_alpha_mask,
        "point_cloud_motion_mask": point_cloud_dynamic_mask,
        "cam_c2w": target_cam_c2w,
        "intrinsics": target_intrinsics,
    }
    if VISTA4D_REUSE_GENERATED_MAP_IN_MEMORY:
        _cache_generated_map(cache_key, data)
        print("Vista4D map kept in memory cache.")
        return data
    return _save_generated_map_to_disk(data, source_cam_c2w, source_intrinsics, depths, sky_mask, fps)


def _load_map(map_dir, frame_num, height, width):
    video_src = _read_video(map_dir / "video_src.mp4")
    video_pc = _read_video(map_dir / "video_pc.mp4")
    video_src = video_src[_center_slice(video_src.shape[0], frame_num)]
    video_pc = video_pc[_center_slice(video_pc.shape[0], frame_num)]
    src_h, src_w = video_src.shape[1:3]
    video_src = _crop_resize_video(video_src, height, width)
    video_pc = _crop_resize_video(video_pc, height, width)
    cam_c2w, intrinsics = _load_cameras(map_dir / "cameras_tgt.npz", frame_num)
    if (src_h, src_w) != (height, width):
        intrinsics = _resize_intrinsics(intrinsics, height, width, src_h, src_w)
    return {
        "source_video": video_src,
        "point_cloud_video": video_pc,
        "source_alpha_mask": _load_masks(map_dir / "alpha_mask_src", frame_num, height, width, True),
        "source_motion_mask": _load_masks(map_dir / "dynamic_mask_src", frame_num, height, width, False),
        "point_cloud_alpha_mask": _load_masks(map_dir / "alpha_mask_pc", frame_num, height, width, True),
        "point_cloud_motion_mask": _load_masks(map_dir / "dynamic_mask_pc", frame_num, height, width, False),
        "cam_c2w": cam_c2w,
        "intrinsics": intrinsics,
    }


def _pack_masks(alpha_mask, motion_mask, device, dtype):
    masks = np.stack((alpha_mask, motion_mask), axis=0)[None].astype(np.float32)
    masks = torch.from_numpy(masks).to(device=device, dtype=dtype)
    b, c, f, h, w = masks.shape
    masks = torch.cat((torch.repeat_interleave(masks[:, :, 0:1], repeats=4, dim=2), masks[:, :, 1:]), dim=2)
    masks = rearrange(masks, "b c (f sf) (h sh) (w sw) -> b (c sf sh sw) f h w", sf=4, sh=8, sw=8)
    return masks


def prepare_vista4d_condition(pipeline, input_frames, input_custom, frame_num, height, width, tile_size, fps=16, custom_settings=None, model_mode=None):
    map_dir, target_camera_path = _resolve_custom_input(input_custom)
    data = None
    if map_dir is None:
        if input_frames is None:
            raise ValueError("Vista4D needs a control video or a preprocessed map in custom_guide.")
        map_ref = _create_map_from_control_video(input_frames, frame_num, height, width, fps, custom_settings, target_camera_path, model_mode=model_mode)
        if isinstance(map_ref, dict):
            data = map_ref
        else:
            map_dir = map_ref

    if data is None:
        data = _load_map(map_dir, frame_num, height, width)
    device = pipeline.device
    dtype = pipeline.dtype
    vae_dtype = pipeline.VAE_dtype
    source_video = _video_np_to_tensor(data["source_video"], device, vae_dtype)
    point_video = _video_np_to_tensor(data["point_cloud_video"], device, vae_dtype)
    with control_video_encoding():
        source_latents = pipeline.vae.encode([source_video], tile_size=tile_size)[0].unsqueeze(0).to(device=device, dtype=dtype)
    with control_video_encoding():
        point_latents = pipeline.vae.encode([point_video], tile_size=tile_size)[0].unsqueeze(0).to(device=device, dtype=dtype)
    source_masks = _pack_masks(data["source_alpha_mask"], data["source_motion_mask"], device, dtype)
    point_masks = _pack_masks(data["point_cloud_alpha_mask"], data["point_cloud_motion_mask"], device, dtype)

    lat_h = height // pipeline.vae_stride[1]
    lat_w = width // pipeline.vae_stride[2]
    cam_c2w = torch.from_numpy(data["cam_c2w"][None]).to(device=device, dtype=dtype)
    intrinsics = torch.from_numpy(data["intrinsics"][None]).to(device=device, dtype=dtype)
    cam_emb = get_plucker_embedding(intrinsics, cam_c2w, height, width, height_dit=lat_h // 2, width_dit=lat_w // 2)
    cam_emb = cam_emb[:, ::pipeline.vae_stride[0]].to(dtype=dtype)

    return {
        "vista": {
            "source_latents": source_latents,
            "point_latents": point_latents,
            "source_mask_latents": source_masks,
            "point_mask_latents": point_masks,
            "cam_emb": cam_emb,
        }
    }