# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
GLB 3D export utilities for GCT predictions.
"""

import os
import copy
from typing import Optional, Tuple

import numpy as np
import cv2
import matplotlib
from scipy.spatial.transform import Rotation

from lingbot_map.vis.sky_segmentation import (
    _SKYSEG_INPUT_SIZE,
    _SKYSEG_SOFT_THRESHOLD,
    _mask_to_float,
    _mask_to_uint8,
    _result_map_to_non_sky_conf,
)

try:
    import trimesh
except ImportError:
    trimesh = None
    print("trimesh not found. GLB export will not work.")


def predictions_to_glb(
    predictions: dict,
    conf_thres: float = 50.0,
    filter_by_frames: str = "all",
    mask_black_bg: bool = False,
    mask_white_bg: bool = False,
    show_cam: bool = True,
    mask_sky: bool = False,
    target_dir: Optional[str] = None,
    prediction_mode: str = "Predicted Pointmap",
) -> "trimesh.Scene":
    """
    Converts GCT predictions to a 3D scene represented as a GLB file.

    Args:
        predictions: Dictionary containing model predictions with keys:
            - world_points: 3D point coordinates (S, H, W, 3)
            - world_points_conf: Confidence scores (S, H, W)
            - images: Input images (S, H, W, 3) or (S, 3, H, W)
            - extrinsic: Camera extrinsic matrices (S, 3, 4)
        conf_thres: Percentage of low-confidence points to filter out
        filter_by_frames: Frame filter specification ("all" or frame index)
        mask_black_bg: Mask out black background pixels
        mask_white_bg: Mask out white background pixels
        show_cam: Include camera visualization
        mask_sky: Apply sky segmentation mask
        target_dir: Output directory for intermediate files
        prediction_mode: "Predicted Pointmap" or "Predicted Depthmap"

    Returns:
        trimesh.Scene: Processed 3D scene containing point cloud and cameras

    Raises:
        ValueError: If input predictions structure is invalid
        ImportError: If trimesh is not available
    """
    if trimesh is None:
        raise ImportError("trimesh is required for GLB export. Install with: pip install trimesh")

    if not isinstance(predictions, dict):
        raise ValueError("predictions must be a dictionary")

    if conf_thres is None:
        conf_thres = 10.0

    print("Building GLB scene")

    # Parse frame filter
    selected_frame_idx = None
    if filter_by_frames != "all" and filter_by_frames != "All":
        try:
            selected_frame_idx = int(filter_by_frames.split(":")[0])
        except (ValueError, IndexError):
            pass

    # Select prediction source
    if "Pointmap" in prediction_mode:
        print("Using Pointmap Branch")
        if "world_points" in predictions:
            pred_world_points = predictions["world_points"]
            pred_world_points_conf = predictions.get(
                "world_points_conf", np.ones_like(pred_world_points[..., 0])
            )
        else:
            print("Warning: world_points not found, falling back to depth-based points")
            pred_world_points = predictions["world_points_from_depth"]
            pred_world_points_conf = predictions.get(
                "depth_conf", np.ones_like(pred_world_points[..., 0])
            )
    else:
        print("Using Depthmap and Camera Branch")
        pred_world_points = predictions["world_points_from_depth"]
        pred_world_points_conf = predictions.get(
            "depth_conf", np.ones_like(pred_world_points[..., 0])
        )

    images = predictions["images"]
    camera_matrices = predictions["extrinsic"]

    # Apply sky segmentation if enabled
    if mask_sky and target_dir is not None:
        pred_world_points_conf = _apply_sky_mask(
            pred_world_points_conf, target_dir, images
        )

    # Apply frame filter
    if selected_frame_idx is not None:
        pred_world_points = pred_world_points[selected_frame_idx][None]
        pred_world_points_conf = pred_world_points_conf[selected_frame_idx][None]
        images = images[selected_frame_idx][None]
        camera_matrices = camera_matrices[selected_frame_idx][None]

    # Prepare vertices and colors
    vertices_3d = pred_world_points.reshape(-1, 3)

    # Handle different image formats
    if images.ndim == 4 and images.shape[1] == 3:  # NCHW format
        colors_rgb = np.transpose(images, (0, 2, 3, 1))
    else:
        colors_rgb = images
    colors_rgb = (colors_rgb.reshape(-1, 3) * 255).astype(np.uint8)

    # Apply confidence filtering
    conf = pred_world_points_conf.reshape(-1)
    conf_threshold = np.percentile(conf, conf_thres) if conf_thres > 0 else 0.0
    conf_mask = (conf >= conf_threshold) & (conf > 1e-5)

    # Apply background masking
    if mask_black_bg:
        black_bg_mask = colors_rgb.sum(axis=1) >= 16
        conf_mask = conf_mask & black_bg_mask

    if mask_white_bg:
        white_bg_mask = ~(
            (colors_rgb[:, 0] > 240) &
            (colors_rgb[:, 1] > 240) &
            (colors_rgb[:, 2] > 240)
        )
        conf_mask = conf_mask & white_bg_mask

    vertices_3d = vertices_3d[conf_mask]
    colors_rgb = colors_rgb[conf_mask]

    # Handle empty point cloud
    if vertices_3d is None or np.asarray(vertices_3d).size == 0:
        vertices_3d = np.array([[1, 0, 0]])
        colors_rgb = np.array([[255, 255, 255]])
        scene_scale = 1
    else:
        lower_percentile = np.percentile(vertices_3d, 5, axis=0)
        upper_percentile = np.percentile(vertices_3d, 95, axis=0)
        scene_scale = np.linalg.norm(upper_percentile - lower_percentile)

    colormap = matplotlib.colormaps.get_cmap("gist_rainbow")

    # Build scene
    scene_3d = trimesh.Scene()
    point_cloud_data = trimesh.PointCloud(vertices=vertices_3d, colors=colors_rgb)
    scene_3d.add_geometry(point_cloud_data)

    # Prepare camera matrices
    num_cameras = len(camera_matrices)
    extrinsics_matrices = np.zeros((num_cameras, 4, 4))
    extrinsics_matrices[:, :3, :4] = camera_matrices
    extrinsics_matrices[:, 3, 3] = 1

    # Add cameras
    if show_cam:
        for i in range(num_cameras):
            world_to_camera = extrinsics_matrices[i]
            camera_to_world = np.linalg.inv(world_to_camera)
            rgba_color = colormap(i / num_cameras)
            current_color = tuple(int(255 * x) for x in rgba_color[:3])
            integrate_camera_into_scene(scene_3d, camera_to_world, current_color, scene_scale)

    # Align scene
    scene_3d = apply_scene_alignment(scene_3d, extrinsics_matrices)

    print("GLB Scene built")
    return scene_3d


def _apply_sky_mask(
    conf: np.ndarray,
    target_dir: str,
    images: np.ndarray
) -> np.ndarray:
    """Apply sky segmentation mask to confidence scores."""
    try:
        import onnxruntime
    except ImportError:
        print("Warning: onnxruntime not available, skipping sky masking")
        return conf

    target_dir_images = os.path.join(target_dir, "images")
    if not os.path.exists(target_dir_images):
        print(f"Warning: Images directory not found at {target_dir_images}")
        return conf

    image_list = sorted(os.listdir(target_dir_images))
    S, H, W = conf.shape if hasattr(conf, "shape") else (len(images), images.shape[1], images.shape[2])

    skyseg_model_path = "skyseg.onnx"
    if not os.path.exists(skyseg_model_path):
        print("Downloading skyseg.onnx...")
        download_file_from_url(
            "https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx",
            skyseg_model_path
        )

    skyseg_session = onnxruntime.InferenceSession(skyseg_model_path)
    sky_mask_list = []

    for i, image_name in enumerate(image_list[:S]):
        image_filepath = os.path.join(target_dir_images, image_name)
        mask_filepath = os.path.join(target_dir, "sky_masks", image_name)

        if os.path.exists(mask_filepath):
            sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE)
        else:
            sky_mask = segment_sky(image_filepath, skyseg_session, mask_filepath)

        if sky_mask.shape[0] != H or sky_mask.shape[1] != W:
            sky_mask = cv2.resize(sky_mask, (W, H), interpolation=cv2.INTER_LINEAR)

        sky_mask_list.append(_mask_to_float(sky_mask))

    sky_mask_array = np.array(sky_mask_list)
    sky_mask_binary = (sky_mask_array > _SKYSEG_SOFT_THRESHOLD).astype(np.float32)
    return conf * sky_mask_binary


def integrate_camera_into_scene(
    scene: "trimesh.Scene",
    transform: np.ndarray,
    face_colors: Tuple[int, int, int],
    scene_scale: float,
    frustum_thickness: float = 1.0,
):
    """
    Integrates a camera mesh into the 3D scene.

    Args:
        scene: The 3D scene to add the camera model
        transform: Transformation matrix for camera positioning
        face_colors: RGB color tuple for the camera
        scene_scale: Scale of the scene
        frustum_thickness: Multiplier for frustum edge thickness (>1 = thicker)
    """
    cam_width = scene_scale * 0.05
    cam_height = scene_scale * 0.1

    rot_45_degree = np.eye(4)
    rot_45_degree[:3, :3] = Rotation.from_euler("z", 45, degrees=True).as_matrix()
    rot_45_degree[2, 3] = -cam_height

    opengl_transform = get_opengl_conversion_matrix()
    complete_transform = transform @ opengl_transform @ rot_45_degree
    camera_cone_shape = trimesh.creation.cone(cam_width, cam_height, sections=4)

    # Build thicker frustum by stacking rotated copies
    slight_rotation = np.eye(4)
    slight_rotation[:3, :3] = Rotation.from_euler("z", 2, degrees=True).as_matrix()

    shell_scales = [1.0, 0.95]
    shell_transforms = [np.eye(4), slight_rotation]
    # Add extra shells for thickness
    if frustum_thickness > 1.0:
        n_extra = max(1, int(frustum_thickness - 1))
        for k in range(1, n_extra + 1):
            # Progressively rotated and scaled copies
            angle = 2.0 + k * 2.0
            scale = 1.0 + k * 0.02
            rot = np.eye(4)
            rot[:3, :3] = Rotation.from_euler("z", angle, degrees=True).as_matrix()
            shell_scales.append(scale)
            shell_transforms.append(rot)
            rot_neg = np.eye(4)
            rot_neg[:3, :3] = Rotation.from_euler("z", -angle, degrees=True).as_matrix()
            shell_scales.append(scale)
            shell_transforms.append(rot_neg)

    vertices_parts = []
    for s, t_mat in zip(shell_scales, shell_transforms):
        vertices_parts.append(
            transform_points(t_mat, s * camera_cone_shape.vertices)
        )
    vertices_combined = np.concatenate(vertices_parts)
    vertices_transformed = transform_points(complete_transform, vertices_combined)

    mesh_faces = compute_camera_faces_multi(camera_cone_shape, len(shell_scales))
    camera_mesh = trimesh.Trimesh(vertices=vertices_transformed, faces=mesh_faces)
    camera_mesh.visual.face_colors[:, :3] = face_colors
    scene.add_geometry(camera_mesh)


def apply_scene_alignment(
    scene_3d: "trimesh.Scene",
    extrinsics_matrices: np.ndarray
) -> "trimesh.Scene":
    """
    Aligns the 3D scene based on the extrinsics of the first camera.

    Args:
        scene_3d: The 3D scene to be aligned
        extrinsics_matrices: Camera extrinsic matrices

    Returns:
        Aligned 3D scene
    """
    opengl_conversion_matrix = get_opengl_conversion_matrix()

    align_rotation = np.eye(4)
    align_rotation[:3, :3] = Rotation.from_euler("y", 180, degrees=True).as_matrix()

    initial_transformation = (
        np.linalg.inv(extrinsics_matrices[0]) @ opengl_conversion_matrix @ align_rotation
    )
    scene_3d.apply_transform(initial_transformation)
    return scene_3d


def get_opengl_conversion_matrix() -> np.ndarray:
    """Returns the OpenGL conversion matrix (flips Y and Z axes)."""
    matrix = np.identity(4)
    matrix[1, 1] = -1
    matrix[2, 2] = -1
    return matrix


def transform_points(
    transformation: np.ndarray,
    points: np.ndarray,
    dim: Optional[int] = None
) -> np.ndarray:
    """
    Applies a 4x4 transformation to a set of points.

    Args:
        transformation: Transformation matrix
        points: Points to be transformed
        dim: Dimension for reshaping the result

    Returns:
        Transformed points
    """
    points = np.asarray(points)
    initial_shape = points.shape[:-1]
    dim = dim or points.shape[-1]

    transformation = transformation.swapaxes(-1, -2)
    points = points @ transformation[..., :-1, :] + transformation[..., -1:, :]

    return points[..., :dim].reshape(*initial_shape, dim)


def compute_camera_faces(cone_shape: "trimesh.Trimesh") -> np.ndarray:
    """Computes the faces for the camera mesh."""
    faces_list = []
    num_vertices_cone = len(cone_shape.vertices)

    for face in cone_shape.faces:
        if 0 in face:
            continue
        v1, v2, v3 = face
        v1_offset, v2_offset, v3_offset = face + num_vertices_cone
        v1_offset_2, v2_offset_2, v3_offset_2 = face + 2 * num_vertices_cone

        faces_list.extend([
            (v1, v2, v2_offset),
            (v1, v1_offset, v3),
            (v3_offset, v2, v3),
            (v1, v2, v2_offset_2),
            (v1, v1_offset_2, v3),
            (v3_offset_2, v2, v3),
        ])

    faces_list += [(v3, v2, v1) for v1, v2, v3 in faces_list]
    return np.array(faces_list)


def compute_camera_faces_multi(cone_shape: "trimesh.Trimesh", num_shells: int) -> np.ndarray:
    """Computes faces for a camera mesh with multiple shells (for thicker frustums).

    Connects each consecutive pair of vertex shells to form the frustum edges.
    """
    faces_list = []
    nv = len(cone_shape.vertices)

    for s in range(num_shells - 1):
        off_a = s * nv
        off_b = (s + 1) * nv
        for face in cone_shape.faces:
            if 0 in face:
                continue
            v1, v2, v3 = face
            faces_list.extend([
                (v1 + off_a, v2 + off_a, v2 + off_b),
                (v1 + off_a, v1 + off_b, v3 + off_a),
                (v3 + off_b, v2 + off_a, v3 + off_a),
            ])

    faces_list += [(v3, v2, v1) for v1, v2, v3 in faces_list]
    return np.array(faces_list)


def segment_sky(
    image_path: str,
    onnx_session,
    mask_filename: str
) -> np.ndarray:
    """
    Segments sky from an image using an ONNX model.

    Args:
        image_path: Path to input image
        onnx_session: ONNX runtime session with loaded model
        mask_filename: Path to save the output mask

    Returns:
        Continuous non-sky confidence map in [0, 1]
    """
    image = cv2.imread(image_path)
    result_map = run_skyseg(onnx_session, _SKYSEG_INPUT_SIZE, image)
    result_map_original = cv2.resize(
        result_map, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_LINEAR
    )
    output_mask = _result_map_to_non_sky_conf(result_map_original)

    os.makedirs(os.path.dirname(mask_filename), exist_ok=True)
    cv2.imwrite(mask_filename, _mask_to_uint8(output_mask))
    return output_mask


def run_skyseg(
    onnx_session,
    input_size: Tuple[int, int],
    image: np.ndarray
) -> np.ndarray:
    """
    Runs sky segmentation inference using ONNX model.

    Args:
        onnx_session: ONNX runtime session
        input_size: Target size for model input (width, height)
        image: Input image in BGR format

    Returns:
        Segmentation mask
    """
    temp_image = copy.deepcopy(image)
    resize_image = cv2.resize(temp_image, dsize=(input_size[0], input_size[1]))
    x = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB)
    x = np.array(x, dtype=np.float32)
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    x = (x / 255 - mean) / std
    x = x.transpose(2, 0, 1)
    x = x.reshape(-1, 3, input_size[0], input_size[1]).astype("float32")

    input_name = onnx_session.get_inputs()[0].name
    output_name = onnx_session.get_outputs()[0].name
    onnx_result = onnx_session.run([output_name], {input_name: x})

    onnx_result = np.array(onnx_result).squeeze()
    min_value = np.min(onnx_result)
    max_value = np.max(onnx_result)
    onnx_result = (onnx_result - min_value) / (max_value - min_value)
    onnx_result *= 255
    return onnx_result.astype("uint8")


def download_file_from_url(url: str, filename: str):
    """Downloads a file from a URL, handling redirects."""
    import requests

    try:
        response = requests.get(url, allow_redirects=False)
        response.raise_for_status()

        if response.status_code == 302:
            redirect_url = response.headers["Location"]
            response = requests.get(redirect_url, stream=True)
            response.raise_for_status()
        else:
            print(f"Unexpected status code: {response.status_code}")
            return

        with open(filename, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        print(f"Downloaded {filename} successfully.")

    except requests.exceptions.RequestException as e:
        print(f"Error downloading file: {e}")