import os
import cv2
import torch
import shutil
from datetime import datetime
import glob
import gc
import gradio as gr
import numpy as np
import open3d as o3d
import concerto
from scipy.spatial.transform import Rotation as R
import trimesh
import time
from typing import List, Tuple
from pathlib import Path
from einops import rearrange
from tqdm import tqdm
import camtools as ct
from PIL import Image
from torchvision import transforms as TF
try:
    import flash_attn
except ImportError:
    flash_attn = None

from visual_util import predictions_to_glb
from vggt.models.vggt import VGGT
from vggt.utils.load_fn import load_and_preprocess_images
from vggt.utils.pose_enc import pose_encoding_to_extri_intri
from vggt.utils.geometry import unproject_depth_map_to_point_map

device = "cuda" if torch.cuda.is_available() else "cpu"

def run_model(target_dir, model) -> dict:
    """
    Run the VGGT model on images in the 'target_dir/images' folder and return predictions.
    """
    print(f"Processing images from {target_dir}")

    # if not torch.cuda.is_available():
    #     raise ValueError("CUDA is not available. Check your environment.")

    # Move model to device
    model = model.to(device)
    model.eval()

    # Load and preprocess images
    image_names = glob.glob(os.path.join(target_dir, "images", "*"))
    image_names = sorted(image_names)
    print(f"Found {len(image_names)} images")
    if len(image_names) == 0:
        raise ValueError("No images found. Check your upload.")

    images = load_and_preprocess_images(image_names).to(device)
    print(f"Preprocessed images shape: {images.shape}")

    # Run inference
    print("Running inference...")
    with torch.no_grad():
        if device == "cuda":
            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                predictions = model(images)
        else:
            predictions = model(images)

    # Convert pose encoding to extrinsic and intrinsic matrices
    print("Converting pose encoding to extrinsic and intrinsic matrices...")
    extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:])
    predictions["extrinsic"] = extrinsic
    predictions["intrinsic"] = intrinsic

    # Convert tensors to numpy
    for key in predictions.keys():
        if isinstance(predictions[key], torch.Tensor):
            predictions[key] = predictions[key].cpu().numpy().squeeze(0)  # remove batch dimension

    # Generate world points from depth map
    print("Computing world points from depth map...")
    depth_map = predictions["depth"]  # (S, H, W, 1)
    world_points = unproject_depth_map_to_point_map(depth_map, predictions["extrinsic"], predictions["intrinsic"])
    predictions["world_points_from_depth"] = world_points

    # Clean up
    torch.cuda.empty_cache()
    return predictions

def handle_uploads(input_file,input_video,conf_thres,frame_slider,prediction_mode,if_TSDF):
    """
    Create a new 'target_dir' + 'images' subfolder, and place user-uploaded
    images or extracted frames from video into it. Return (target_dir, image_paths).
    """
    start_time = time.time()
    gc.collect()
    torch.cuda.empty_cache()

    # Create a unique folder name
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    target_dir = f"demo_output/inputs_{timestamp}"
    target_dir_images = os.path.join(target_dir, "images")
    target_dir_pcds = os.path.join(target_dir, "pcds")

    # Clean up if somehow that folder already exists
    if os.path.exists(target_dir):
        shutil.rmtree(target_dir)
    os.makedirs(target_dir)
    os.makedirs(target_dir_images)
    os.makedirs(target_dir_pcds)
    # --- Handle video ---
    if input_video is not None:
        print("processing video")
        if isinstance(input_video, dict) and "name" in input_video:
            video_path = input_video["name"]
        else:
            video_path = input_video

        vs = cv2.VideoCapture(video_path)
        fps = vs.get(cv2.CAP_PROP_FPS)
        frame_interval = int(fps * frame_slider)  # 1 frame/sec

        count = 0
        video_frame_num = 0
        image_paths = []
        while True:
            gotit, frame = vs.read()
            if not gotit:
                break
            count += 1
            if count % frame_interval == 0:
                image_path = os.path.join(target_dir_images, f"{video_frame_num:06}.png")
                cv2.imwrite(image_path, frame)
                image_paths.append(image_path)
                video_frame_num += 1
        # Sort final images for gallery
        image_paths = sorted(image_paths)
        original_points, original_colors, original_normals = parse_frames(target_dir,conf_thres,prediction_mode,if_TSDF)
    if input_file is not None:
        print("processing ply")
        pcd = o3d.io.read_point_cloud(input_file.name) 
        pcd.estimate_normals()
        original_points = np.asarray(pcd.points)
        original_colors = np.asarray(pcd.colors)
        original_normals = np.asarray(pcd.normals)
        image_paths = None
    scene_3d = trimesh.Scene()
    point_cloud_data = trimesh.PointCloud(vertices=original_points, colors=original_colors, vertex_normals=original_normals)
    scene_3d.add_geometry(point_cloud_data)
    original_temp = os.path.join(target_dir_pcds,"original.glb")
    scene_3d.export(file_obj=original_temp)
    np.save(os.path.join(target_dir_pcds, f"points.npy"), original_points)
    np.save(os.path.join(target_dir_pcds, f"colors.npy"), original_colors)
    np.save(os.path.join(target_dir_pcds, f"normals.npy"), original_normals)
    end_time = time.time()
    print(f"Files copied to {target_dir}; took {end_time - start_time:.3f} seconds")
    return target_dir, image_paths,original_temp, end_time - start_time

def update_gallery_on_upload(input_file,input_video,conf_thres,frame_slider,prediction_mode,TSDF_mode):
    """
    Whenever user uploads or changes files, immediately handle them
    and show in the gallery. Return (target_dir, image_paths).
    If nothing is uploaded, returns "None" and empty list.
    """
    if not input_video and not input_file:
        return None, None, None, None
    if_TSDF = True if TSDF_mode=="Yes" else False
    target_dir, image_paths,original_view, reconstruction_time = handle_uploads(input_file,input_video,conf_thres,frame_slider,prediction_mode,if_TSDF)
    if input_file is not None:
        return original_view, target_dir, [], f"Upload and preprocess complete with {reconstruction_time:.3f} sec. Click \"PCA Generate\" to begin PCA processing."
    if input_video is not None:
        return original_view, target_dir, image_paths, f"Upload and preprocess complete with {reconstruction_time:.3f} sec. Click \"PCA Generate\" to begin PCA processing."

def clear_fields():
    """
    Clears the 3D viewer, the stored target_dir, and empties the gallery.
    """
    return None

def PCAing_log(is_example, log_output):
    """
    Display a quick log message while waiting.
    """
    if is_example:
        return log_output
    return "Loading for Doing PCA..."

def reset_log():
    """
    Reset a quick log message.
    """
    return "A new point cloud file or video is uploading and preprocessing..."

def parse_frames(
    target_dir,
    conf_thres=3.0,
    prediction_mode="Pointmap Regression",
    if_TSDF=True,
):
    """
    Perform reconstruction using the already-created target_dir/images.
    """
    if not os.path.isdir(target_dir) or target_dir == "None":
        return None, "No valid target directory found. Please upload first.", None, None

    start_time = time.time()
    gc.collect()
    torch.cuda.empty_cache()

    # Prepare frame_filter dropdown
    target_dir_images = os.path.join(target_dir, "images")
    target_dir_pcds = os.path.join(target_dir, "pcds")
    all_files = sorted(os.listdir(target_dir_images)) if os.path.isdir(target_dir_images) else []
    all_files = [f"{i}: {filename}" for i, filename in enumerate(all_files)]
    frame_filter_choices = ["All"] + all_files

    print("Running run_model...")
    with torch.no_grad():
        predictions = run_model(target_dir, VGGT_model)

    # Save predictions
    prediction_save_path = os.path.join(target_dir, "predictions.npz")
    np.savez(prediction_save_path, **predictions)

    # Convert pose encoding to extrinsic and intrinsic matrices
    images = predictions["images"]
    Ts, Ks = predictions["extrinsic"],predictions["intrinsic"]
    Ts = ct.convert.pad_0001(Ts)
    Ts_inv = np.linalg.inv(Ts)
    Cs = np.array([ct.convert.T_to_C(T) for T in Ts])  # (n, 3)

    # [1, 8, 294, 518, 3]
    world_points = predictions["world_points"]

    # Compute view direction for each pixel
    # (b n h w c) - (n, 3)
    view_dirs = world_points - rearrange(Cs, "n c -> n 1 1 c")
    view_dirs = rearrange(view_dirs, "n h w c -> (n h w) c")
    view_dirs = view_dirs / np.linalg.norm(view_dirs, axis=-1, keepdims=True)

    # Extract points and colors
    # [1, 8, 3, 294, 518]
    img_num = world_points.shape[1]
    images = predictions["images"]
    points = rearrange(world_points, "n h w c -> (n h w) c")
    colors = rearrange(images, "n c h w -> (n h w) c")

    if prediction_mode=="Pointmap Branch":
        world_points_conf = predictions["world_points_conf"]
        conf = world_points_conf.reshape(-1)
        points,Ts_inv,_ = Coord2zup(points, Ts_inv)
        scale = 3 / (points[:, 2].max() - points[:, 2].min())
        points *= scale
        Ts_inv[:, :3, 3] *= scale

        # Create a point cloud
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        pcd.colors = o3d.utility.Vector3dVector(colors)
        pcd.estimate_normals()
        # o3d.io.write_point_cloud("pcd.ply", pcd)
        try:
            pcd, inliers, rotation_matrix, offset = extract_and_align_ground_plane(pcd)
        except Exception as e:
            print(f"cannot find ground, err:{e}")
        # Filp normals such that normals always point to camera
        # Compute the dot product between the normal and the view direction
        # If the dot product is less than 0, flip the normal
        normals = np.asarray(pcd.normals)
        view_dirs = np.asarray(view_dirs)
        dot_product = np.sum(normals * view_dirs, axis=-1)
        flip_mask = dot_product > 0
        normals[flip_mask] = -normals[flip_mask]

        # Normalize normals a nd m
        normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True)
        pcd.normals = o3d.utility.Vector3dVector(normals)
        if conf_thres == 0.0:
            conf_threshold = 0.0
        else:
            conf_threshold = np.percentile(conf, conf_thres)
        conf_mask = (conf >= conf_threshold) & (conf > 1e-5)
        points = points[conf_mask]
        colors = colors[conf_mask]
        normals = normals[conf_mask]
    elif prediction_mode=="Depthmap Branch":
        # Integrate RGBD images into a TSDF volume and extract a mesh
        # (n, h, w, 3)
        im_colors = rearrange(images, "n c h w -> (n) h w c")
        # (b, n, h, w, 3)
        im_dists = world_points - rearrange(Cs, "n c -> n 1 1 c")
        im_dists = np.linalg.norm(im_dists, axis=-1, keepdims=False)

        # Convert distance to depth
        im_depths = []  # (n, h, w, c)
        for im_dist, K in zip(im_dists, Ks):
            im_depth = ct.convert.im_distance_to_im_depth(im_dist, K)
            im_depths.append(im_depth)
        im_depths = np.stack(im_depths, axis=0)
        if if_TSDF:
            mesh = integrate_rgbd_to_mesh(
                Ks=Ks,
                Ts=Ts,
                im_depths=im_depths,
                im_colors=im_colors,
                voxel_size=1 / 512,
            )
            rotation_angle = -np.pi / 2
            rotation_axis = np.array([1, 0, 0])  # X 轴
            mesh.rotate(
                o3d.geometry.get_rotation_matrix_from_axis_angle(rotation_axis * rotation_angle),
                center=(0,0,0)
            )
            vertices = np.asarray(mesh.vertices)
            scale_factor = 3./(np.max(vertices[:,2])-np.min(vertices[:,2]))
            mesh.scale(scale_factor, center=(0,0,0))
            points = np.asarray(mesh.vertices)
            colors = np.asarray(mesh.vertex_colors) if mesh.has_vertex_colors() else np.zeros_like(vertices)
            if not mesh.has_vertex_normals():
                mesh.compute_vertex_normals()
            normals = np.asarray(mesh.vertex_normals)
            Ts_inv = rotx(Ts_inv, theta=-90)
            Ts_inv[:, :3, 3] *= scale_factor
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(points)
            pcd.colors = o3d.utility.Vector3dVector(colors)
            pcd.normals = o3d.utility.Vector3dVector(normals)
        else:
            points=[]
            for K, T, im_depth in zip(Ks, Ts, im_depths):
                point = ct.project.im_depth_to_point_cloud(
                    im_depth=im_depth,
                    K=K,
                    T=T,
                    to_image=False,
                    ignore_invalid=False,
                )
                points.append(point)
            points = np.vstack(points)
            colors = im_colors.reshape(-1,3)
            world_points_conf = predictions["depth_conf"]
            conf = world_points_conf.reshape(-1)
            if conf_thres == 0.0:
                conf_threshold = 0.0
            else:
                conf_threshold = np.percentile(conf, conf_thres)
            conf_mask = (conf >= conf_threshold) & (conf > 1e-5)
            points = points[conf_mask]
            colors = colors[conf_mask]
            points,Ts_inv,_ = Coord2zup(points, Ts_inv)
            scale_factor = 3./(np.max(points[:,2])-np.min(points[:,2]))
            points *= scale_factor
            Ts_inv[:, :3, 3] *= scale_factor
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(points)
            pcd.colors = o3d.utility.Vector3dVector(colors)
            pcd.estimate_normals()
        try:
            pcd, inliers, rotation_matrix, offset = extract_and_align_ground_plane(pcd)
        except Exception as e:
            print(f"cannot find ground, err:{e}")
    original_points = np.asarray(pcd.points)
    original_colors = np.asarray(pcd.colors)
    original_normals = np.asarray(pcd.normals)
    # Cleanup
    del predictions
    gc.collect()
    torch.cuda.empty_cache()
    end_time = time.time()
    print(f"Total time: {end_time - start_time:.2f} seconds")
    return original_points, original_colors, original_normals

def extract_and_align_ground_plane(pcd, 
                                   height_percentile=20, 
                                   ransac_distance_threshold=0.01, 
                                   ransac_n=3, 
                                   ransac_iterations=1000,
                                   max_angle_degree=40,
                                   max_trials=6):
    points = np.asarray(pcd.points)
    z_vals = points[:, 2]
    z_thresh = np.percentile(z_vals, height_percentile)
    low_indices = np.where(z_vals <= z_thresh)[0]

    remaining_indices = low_indices.copy()

    for trial in range(max_trials):
        if len(remaining_indices) < ransac_n:
            raise ValueError("Not enough points left to fit a plane.")

        low_pcd = pcd.select_by_index(remaining_indices)

        plane_model, inliers = low_pcd.segment_plane(
            distance_threshold=ransac_distance_threshold,
            ransac_n=ransac_n,
            num_iterations=ransac_iterations)
        a, b, c, d = plane_model
        normal = np.array([a, b, c])
        normal /= np.linalg.norm(normal)
        
        # current_plane_pcd = pcd.select_by_index(remaining_indices[inliers])
        # o3d.io.write_point_cloud("plane.ply",current_plane_pcd)
        # exit()

        angle = np.arccos(np.clip(np.dot(normal, [0, 0, 1]), -1.0, 1.0)) * 180 / np.pi
        if angle <= max_angle_degree:
            inliers_global = remaining_indices[inliers]

            target = np.array([0, 0, 1])
            axis = np.cross(normal, target)
            axis_norm = np.linalg.norm(axis)

            if axis_norm < 1e-6:
                rotation_matrix = np.eye(3)
            else:
                axis /= axis_norm
                rot_angle = np.arccos(np.clip(np.dot(normal, target), -1.0, 1.0))
                rotation = R.from_rotvec(axis * rot_angle)
                rotation_matrix = rotation.as_matrix()

            rotated_points = points @ rotation_matrix.T
            ground_points_z = rotated_points[inliers_global, 2]
            offset = np.mean(ground_points_z)
            rotated_points[:, 2] -= offset

            aligned_pcd = o3d.geometry.PointCloud()
            aligned_pcd.points = o3d.utility.Vector3dVector(rotated_points)
            if pcd.has_colors():
                aligned_pcd.colors = pcd.colors
            if pcd.has_normals():
                rotated_normals = np.asarray(pcd.normals) @ rotation_matrix.T
                aligned_pcd.normals = o3d.utility.Vector3dVector(rotated_normals)

            return aligned_pcd, inliers_global, rotation_matrix, offset

        else:
            rejected_indices = remaining_indices[inliers]
            remaining_indices = np.setdiff1d(remaining_indices, rejected_indices)

    raise ValueError("Failed to find a valid ground plane within max trials.")

def rotx(x, theta=90):
    """
    Rotate x by theta degrees around the x-axis
    """
    theta = np.deg2rad(theta)
    rot_matrix = np.array(
        [
            [1, 0, 0, 0],
            [0, np.cos(theta), -np.sin(theta), 0],
            [0, np.sin(theta), np.cos(theta), 0],
            [0, 0, 0, 1],
        ]
    )
    return rot_matrix@ x


def Coord2zup(points, extrinsics, normals = None):
    """
    Convert the dust3r coordinate system to the z-up coordinate system
    """
    points = np.concatenate([points, np.ones([points.shape[0], 1])], axis=1).T
    points = rotx(points, -90)[:3].T
    if normals is not None:
        normals = np.concatenate([normals, np.ones([normals.shape[0], 1])], axis=1).T
        normals = rotx(normals, -90)[:3].T
        normals = normals / np.linalg.norm(normals, axis=1, keepdims=True)
    t = np.min(points,axis=0)
    points -= t
    extrinsics = rotx(extrinsics, -90)
    extrinsics[:, :3, 3] -= t.T
    return points, extrinsics, normals
def integrate_rgbd_to_mesh(
    Ks,
    Ts,
    im_depths,
    im_colors,
    voxel_size,
    bbox=None,
):
    """
    Integrate RGBD images into a TSDF volume and extract a mesh.

    Args:
        Ks: (N, 3, 3) camera intrinsics.
        Ts: (N, 4, 4) camera extrinsics.
        im_depths: (N, H, W) depth images, already in world scale.
        im_colors: (N, H, W, 3) color images, float range in [0, 1].
        voxel_size: TSDF voxel size, in meters, e.g. 3 / 512.
        bbox: Open3D axis-aligned bounding box, for cropping.

    Per Open3D convention, invalid depth values shall be set to 0.
    """
    num_images = len(Ks)
    if (
        len(Ts) != num_images
        or len(im_depths) != num_images
        or len(im_colors) != num_images
    ):
        raise ValueError("Ks, Ts, im_depths, im_colors must have the same length.")

    # Constants.
    trunc_voxel_multiplier = 8.0
    sdf_trunc = trunc_voxel_multiplier * voxel_size

    volume = o3d.pipelines.integration.ScalableTSDFVolume(
        voxel_length=voxel_size,
        sdf_trunc=sdf_trunc,
        color_type=o3d.pipelines.integration.TSDFVolumeColorType.RGB8,
    )

    for K, T, im_depth, im_color in tqdm(
        zip(Ks, Ts, im_depths, im_colors),
        total=len(Ks),
        desc="Integrating RGBD frames",
    ):
        # Set invalid depth values to 0, based on bounding box.
        if bbox is not None:
            points = ct.project.im_depth_to_point_cloud(
                im_depth=im_depth,
                K=K,
                T=T,
                to_image=False,
                ignore_invalid=False,
            )
            assert len(points) == im_depth.shape[0] * im_depth.shape[1]
            point_indices_inside_bbox = bbox.get_point_indices_within_bounding_box(
                o3d.utility.Vector3dVector(points)
            )
            point_indices_outside_bbox = np.setdiff1d(
                np.arange(len(points)), point_indices_inside_bbox
            )
            im_depth.ravel()[point_indices_outside_bbox] = 0

        im_color_uint8 = np.ascontiguousarray((im_color * 255).astype(np.uint8))
        im_depth_uint16 = np.ascontiguousarray((im_depth * 1000).astype(np.uint16))
        im_color_o3d = o3d.geometry.Image(im_color_uint8)
        im_depth_o3d = o3d.geometry.Image(im_depth_uint16)
        im_rgbd_o3d = o3d.geometry.RGBDImage.create_from_color_and_depth(
            im_color_o3d,
            im_depth_o3d,
            depth_scale=1000.0,
            depth_trunc=10.0,
            convert_rgb_to_intensity=False,
        )
        o3d_intrinsic = o3d.camera.PinholeCameraIntrinsic(
            width=im_depth.shape[1],
            height=im_depth.shape[0],
            fx=K[0, 0],
            fy=K[1, 1],
            cx=K[0, 2],
            cy=K[1, 2],
        )
        o3d_extrinsic = T
        volume.integrate(
            im_rgbd_o3d,
            o3d_intrinsic,
            o3d_extrinsic,
        )

    mesh = volume.extract_triangle_mesh()
    return mesh

def get_pca_color(feat, start = 0, brightness=1.25, center=True):
    u, s, v = torch.pca_lowrank(feat, center=center, q=3*(start+1), niter=5)
    projection = feat @ v
    projection = projection[:, 3*start:3*(start+1)] * 0.6 + projection[:, 3*start:3*(start+1)] * 0.4
    min_val = projection.min(dim=-2, keepdim=True)[0]
    max_val = projection.max(dim=-2, keepdim=True)[0]
    div = torch.clamp(max_val - min_val, min=1e-6)
    color = (projection - min_val) / div * brightness
    color = color.clamp(0.0, 1.0)
    return color

def Concerto_process(target_dir, original_points, original_colors, original_normals, slider_value, bright_value, model_type):
    gc.collect()
    torch.cuda.empty_cache()
    target_dir_pcds = os.path.join(target_dir, "pcds")

    point = {"coord": original_points, "color": original_colors, "normal":original_normals}
    original_coord = point["coord"].copy()
    original_color = point["color"].copy()
    point = transform(point)

    with torch.inference_mode():
        for key in point.keys():
            if isinstance(point[key], torch.Tensor) and device=="cuda":
                point[key] = point[key].cuda(non_blocking=True)
        # model forward:
        concerto_start_time = time.time()
        if model_type =="Concerto":
            point = concerto_model(point)
        elif model_type == "Sonata":
            point = sonata_model(point)
        concerto_end_time = time.time()
        # upcast point feature
        # Point is a structure contains all the information during forward
        for _ in range(2):
            assert "pooling_parent" in point.keys()
            assert "pooling_inverse" in point.keys()
            parent = point.pop("pooling_parent")
            inverse = point.pop("pooling_inverse")
            parent.feat = torch.cat([parent.feat, point.feat[inverse]], dim=-1)
            point = parent
        while "pooling_parent" in point.keys():
            assert "pooling_inverse" in point.keys()
            parent = point.pop("pooling_parent")
            inverse = point.pop("pooling_inverse")
            parent.feat = point.feat[inverse]
            point = parent

        # here point is down-sampled by GridSampling in default transform pipeline
        # feature of point cloud in original scale can be acquired by:
        _ = point.feat[point.inverse]

        # PCA
        point_feat = point.feat.cpu().detach().numpy()
        np.save(os.path.join(target_dir_pcds,"feat.npy"),point_feat)
        pca_start_time = time.time()
        pca_color = get_pca_color(point.feat,start = slider_value, brightness=bright_value, center=True)
        pca_end_time = time.time()

    # inverse back to original scale before grid sampling
    # point.inverse is acquired from the GirdSampling transform
    point_inverse = point.inverse.cpu().detach().numpy()
    np.save(os.path.join(target_dir_pcds,"inverse.npy"),point_inverse)
    original_pca_color = pca_color[point.inverse]
    points = original_coord
    colors = original_pca_color.cpu().detach().numpy()

    end_time = time.time()
    return points, colors, concerto_end_time - concerto_start_time, pca_end_time - pca_start_time

def gradio_demo(target_dir,pca_slider,bright_slider, model_type, if_color=True, if_normal=True):
    target_dir_pcds = os.path.join(target_dir, "pcds")
    if not os.path.isfile(os.path.join(target_dir_pcds,"points.npy")):
        return None, "No point cloud available. Please upload data first."
    original_points = np.load(os.path.join(target_dir_pcds,"points.npy"))
    if if_color:
        original_colors = np.load(os.path.join(target_dir_pcds,"colors.npy"))
    else:
        original_colors = np.zeros_like(original_points)
    if if_normal:
        original_normals = np.load(os.path.join(target_dir_pcds,"normals.npy"))
    else:
        original_normals = np.zeros_like(original_points)
    processed_temp = (os.path.join(target_dir_pcds,"processed.glb"))
    processed_points, processed_colors, concerto_time, pca_time = Concerto_process(target_dir,original_points, original_colors,original_normals, pca_slider, bright_slider, model_type)
    feat_3d = trimesh.Scene()
    feat_data = trimesh.PointCloud(vertices=processed_points, colors=processed_colors, vertex_normals=original_normals)
    feat_3d.add_geometry(feat_data)
    feat_3d.export(processed_temp)

    return processed_temp, f"Feature visualization process finished with {concerto_time:.3f} seconds using Concerto inference and {pca_time:.3f} seconds using PCA. Updating visualization."

def concerto_slider_update(target_dir,pca_slider,bright_slider,is_example,log_output):
    if is_example == "True":
        return None, log_output
    else:
        target_dir_pcds = os.path.join(target_dir, "pcds")
        if os.path.isfile(os.path.join(target_dir_pcds,"feat.npy")):
            feat = np.load(os.path.join(target_dir_pcds,"feat.npy"))
            inverse = np.load(os.path.join(target_dir_pcds,"inverse.npy"))
            feat = torch.tensor(feat, device = device)
            inverse = torch.tensor(inverse, device = device)
            pca_start_time = time.time()
            pca_colors = get_pca_color(feat,start = pca_slider, brightness=bright_slider, center=True)
            processed_colors = pca_colors[inverse].cpu().detach().numpy()
            pca_end_time = time.time()
            pca_time = pca_end_time - pca_start_time
            processed_points = np.load(os.path.join(target_dir_pcds,"points.npy"))
            processed_normals = np.load(os.path.join(target_dir_pcds,"normals.npy"))
            processed_temp = (os.path.join(target_dir_pcds,"processed.glb"))
            feat_3d = trimesh.Scene()
            feat_data = trimesh.PointCloud(vertices=processed_points, colors=processed_colors, vertex_normals=processed_normals)
            feat_3d.add_geometry(feat_data)
            feat_3d.export(processed_temp)
            log_output = f"Feature visualization process finished with{pca_time:.3f} seconds using PCA. Updating visualization."
        else:
            processed_temp = None
            log_output = "No representations saved, please click PCA generate first."
            # processed_temp, log_output = gradio_demo(target_dir,pca_slider,bright_slider)
    return processed_temp, log_output

# set random seed
# (random seed affect pca color, yet change random seed need manual adjustment kmeans)
# (the pca prevent in paper is with another version of cuda and pytorch environment)
concerto.utils.set_seed(53124)
# Load model
if device == 'cuda' and flash_attn is not None:
    print("Loading model with Flash Attention on GPU.")
    concerto_model = concerto.load("concerto_large", repo_id="Pointcept/Concerto").to(device)
    sonata_model = concerto.model.load("sonata", repo_id="facebook/sonata").to(device)
else:
    print("Loading model on CPU or without Flash Attention.")
    custom_config = dict(
        # enc_patch_size=[1024 for _ in range(5)],  # reduce patch size if necessary
        enable_flash=False,
    )
    concerto_model = concerto.load(
        "concerto_large", repo_id="Pointcept/Concerto", custom_config=custom_config
    ).to(device)
    sonata_model = concerto.load("sonata", repo_id="facebook/sonata", custom_config=custom_config).to(device)

transform = concerto.transform.default()

VGGT_model = VGGT().to(device)
_URL = "https://huggingface.co/facebook/VGGT-1B/resolve/main/model.pt"
VGGT_model.load_state_dict(torch.hub.load_state_dict_from_url(_URL))
# VGGT_model.load_state_dict(torch.load("vggt/ckpt/model.pt",weights_only=True))

examples_video = [
    # ["example/video/conference_room.mp4", 0.0, 1, "Depthmap Branch", "Yes",0,1.2, "True"],
    # ["example/video/office.mp4", 0.0, 1, "Pointmap Branch", "Yes",2,1.1, "True"],
    ["example/video/re10k_1.mp4", 10.0, 1, "Depthmap Branch", "No",2,1.2, "True"],
    ["example/video/re10k_2.mp4", 30.0, 1, "Depthmap Branch", "Yes",1,1.2, "True"],
    ["example/video/re10k_3.mp4", 10.0, 1, "Depthmap Branch", "Yes",1,1.2, "True"],
    ["example/video/re10k_4.mp4", 10.0, 1, "Depthmap Branch", "Yes",1,1., "True"],
]

examples_pcd = [
    ["example/pcd/scannet_0024.png","example/pcd/scannet_0024.ply",2,1.2, "True"],
    ["example/pcd/scannet_0603.png","example/pcd/scannet_0603.ply",0,1.2, "True"],
    # ["example/pcd/hm3d_00012_kDgLKdMd5X8_2.png","example/pcd/hm3d_00012_kDgLKdMd5X8_2.ply",0,1.2, "True"],
    ["example/pcd/hm3d_00113_3goH1WRaCYC.png","example/pcd/hm3d_00113_3goH1WRaCYC.ply",0,1.2, "True"],
    ["example/pcd/s3dis_Area2_auditorium1.png","example/pcd/s3dis_Area2_auditorium1.ply",0,1.2, "True"],
    # ["example/pcd/s3dis_Area4_lobby1.png","example/pcd/s3dis_Area4_lobby1.ply",1,1., "True"],
]

# ["example/pcd/scannetpp_2a1b555966.png","example/pcd/scannetpp_2a1b555966.ply",1,1.1, "True"],
    # ["example/pcd/hm3d_00012_kDgLKdMd5X8_1.png","example/pcd/hm3d_00012_kDgLKdMd5X8_1.ply",0,1.0, "True"],
    # ["example/pcd/s3dis_Area2_conferenceRoom1.png","example/pcd/s3dis_Area2_conferenceRoom1.ply",0,1.2, "True"],
    # ["example/pcd/s3dis_Area4_hallway3.png","example/pcd/s3dis_Area4_hallway3.ply",0,1.2, "True"],

with gr.Blocks(
    css="""
    .custom-log * {
        font-style: italic;
        font-size: 22px !important;
        background-image: linear-gradient(120deg, #0ea5e9 0%, #6ee7b7 60%, #34d399 100%);
        -webkit-background-clip: text;
        background-clip: text;
        font-weight: bold !important;
        color: transparent !important;
        text-align: center !important;
        width: 800px;  
        height: 100px; 
    }
    
    .example-log * {
        font-style: italic;
        font-size: 16px !important;
        background-image: linear-gradient(120deg, #0ea5e9 0%, #6ee7b7 60%, #34d399 100%);
        -webkit-background-clip: text;
        background-clip: text;
        color: transparent !important;
    }

    .common-markdown * {
        font-size: 22px !important;
        -webkit-background-clip: text;
        background-clip: text;
        font-weight: bold !important;
        color: #0ea5e9 !important;
        text-align: center !important;
    }

    #big-box {
        border: 3px solid #00bcd4; 
        padding: 20px;            
        background-color: transparent;  
        border-radius: 15px;   
    }
    
    #my_radio .wrap {
        display: flex;
        flex-wrap: nowrap;
        justify-content: center;
        align-items: center;
    }

    #my_radio .wrap label {
        display: flex;
        width: 50%;
        justify-content: center;
        align-items: center;
        margin: 0;
        padding: 10px 0;
        box-sizing: border-box;
    }
    """,
) as demo:
    gr.HTML(
    """
    <h1>Concerto: Joint 2D-3D Self-Supervised Learning for Emergent Spatial Representations</h1>
    <div style="font-size: 16px; line-height: 1.5;">
        <ol>
            <details style="display:inline;">
                <summary style="display:inline;"><h3>Getting Started:(<strong>Click to expand</strong>)</h3></summary>
            <li><strong>Before Start:</strong> We deploy the model on CPU, thus making the inference speed slow. Due to space limitations, you may encounter errors. If so, please deploy this demo locally.</li>
            <li><strong>Upload Your Data:</strong> Use the "Upload Video" or "Upload Point Cloud" blocks on the left to provide your input. If you upload a video, it will be automatically split into individual frames with the specified frame gap by VGGT.</li>
            <li>
                <strong>[Optional] Adjust Video-Lifted Point Cloud:</strong>
                Before reconstructing the video, you can fine-tune the VGGT lifting process using the options below
                <details style="display:inline;">
                    <summary style="display:inline;">(<strong>Click to expand</strong>)</summary>
                    <ul>
                        <li><em>Frame Gap / N Sec:</em> Adjust the frame interval.</li>
                        <li><em>Confidence Threshold:</em> Adjust the point filtering based on confidence levels.</li>
                        <li><em>Select Prediction Mode:</em> Choose between "Depthmap Branch" and "Pointmap Branch."</li>
                        <li><em>TSDF Integration (Depthmap Branch Mode):</em> Enable TSDF integration to reduce noise in the point cloud when using the "Depthmap Branch" mode. This procedure will cost a long time for refinement.</li>
                    </ul>
                </details>
            </li>
            <li><strong>PCA Generation:</strong> After reconstruction, click the "PCA Generate" button to start the representation extraction and PCA process.</li>
            <li><strong>Clear:</strong> Click the "Clear" button to reset all content in the blocks.</li>
            <li><strong>Point Cloud Preview:</strong> Your uploaded video or point cloud will be displayed in this block.</li>
            <li><strong>PCA Result:</strong> The PCA point cloud will appear here. You can rotate, drag, and zoom to explore the model, and download the GLB file.</li>
            <li>
                <strong>[Optional] Adjust the Point Cloud Input (pre-release feature of the next work): use the checkbox "Input with Point Cloud Color" and "Input with Point Cloud Normal".
            </li>
            <li>
                <strong>[Optional] Adjust PCA Visualization:</strong>
                Fine-tune the PCA visualization using the options below
                <details style="display:inline;">
                    <summary style="display:inline;">(<strong>Click to expand</strong>)</summary>
                    <ul>
                        <li><em>Model Type:</em> Choose the model from Concerto and Sonata.</li>
                        <li><em>PCA Start Dimension:</em> PCA reduces high-dimensional representations into 3D vectors. Adjust the PCA start dimension to change the range of the visualization. Increasing this value can help you see PCA visualization with less variance when the initial PCA dimension shows less diversity.</li>
                        <li><em>PCA Brightness:</em> Adjust the brightness of the PCA visualization results.</li>
                        <li><em>Notice:</em> As a linear dimension reduction method, PCA has its limitation. Sometimes, the visualization cannot fully exhibit the quality of representations.</li>
                    </ul>
                </details>
            </li>
            </details>
        </ol>
    </div>

    """
    )
    _ = gr.Textbox(label="_", visible=False, value="False")
    is_example = gr.Textbox(label="is_example", visible=False, value="False")
    target_dir = gr.Textbox(label="Target Dir", visible=False, value="None")
    preview_imgs = gr.Image(type="filepath",label="Preview Imgs", visible=False, value="None")
    with gr.Row():
        with gr.Column(scale=1,elem_id="big-box"):
            input_file = gr.File(label="Upload Point Cloud", file_types=[".ply"])
            input_video = gr.Video(label="Upload Video", interactive=True)
            image_gallery = gr.Gallery(
                label="Preview",
                columns=4,
                height="300px",
                show_download_button=True,
                object_fit="contain",
                preview=True,
            )

            frame_slider = gr.Slider(minimum=0.1, maximum=10, value=1, step=0.1, 
                             label="1 Frame/ N Sec", interactive=True)
            conf_thres = gr.Slider(minimum=0, maximum=100, value=10, step=0.1, 
                             label="Confidence", interactive=True)
            prediction_mode = gr.Radio(
                ["Depthmap Branch", "Pointmap Branch"],
                label="Select a Prediction Mode",
                value="Depthmap Branch",
                scale=1,
                elem_id="my_radio",
            )
            TSDF_mode = gr.Radio(
                ["Yes", "No"],
                label="TSDF integration under Depthmap Branch mode",
                value="Yes",
                scale=1,
                elem_id="my_radio",
            )
            reconstruction_btn = gr.Button("Video Reconstruct")
        with gr.Column(scale=2):
            log_output = gr.Markdown(
                "Please upload a video or point cloud ply file, then click \"PCA Generate\".", elem_classes=["custom-log"]
            )
            original_view = gr.Model3D(height=520, zoom_speed=0.5, pan_speed=0.5, label="Point Cloud Preview", camera_position = (90,None,None))
            processed_view = gr.Model3D(height=520, zoom_speed=0.5, pan_speed=0.5, label="PCA Result", camera_position = (90,None,None))
            with gr.Row():
                if_color = gr.Checkbox(label="Input with Point Cloud Color", value=True)
                if_normal = gr.Checkbox(label="Input with Point Cloud Normal", value=True)
            model_type = gr.Radio(
                ["Concerto", "Sonata"],
                label="Select a Model Type",
                value="Concerto",
                scale=1,
                elem_id="my_radio",
            )
            pca_slider = gr.Slider(minimum=0, maximum=5, value=0, step=1, 
                             label="PCA Start Dimension", interactive=True)
            bright_slider = gr.Slider(minimum=0.5, maximum=1.5, value=1.2, step=0.05, 
                             label="PCA Brightness", interactive=True)
            with gr.Row():
                submit_btn = gr.Button("PCA Generate")
                clear_btn = gr.ClearButton(
                        [input_video, input_file, original_view, processed_view, log_output, target_dir, image_gallery],
                        scale=1,
                        elem_id="my_clear",
                    )

    gr.Markdown("Click any row to load an example.", elem_classes=["example-log"])
    with gr.Row():
        def example_video_updated(
            inputs,
            conf_thres,
            frame_slider,
            prediction_mode,
            TSDF_mode,
            pca_slider,
            bright_slider,
            is_example,
        ):
            return inputs,conf_thres,frame_slider,prediction_mode,TSDF_mode,pca_slider,bright_slider,is_example
        gr.Examples(
            examples=examples_video,
            inputs=[
                input_video,
                conf_thres,
                frame_slider,
                prediction_mode,
                TSDF_mode,
                pca_slider,
                bright_slider,
                is_example,
            ],
            outputs=[
                input_video,
                conf_thres,
                frame_slider,
                prediction_mode,
                TSDF_mode,
                pca_slider,
                bright_slider,
                is_example,
            ],
            label = "Video Examples",
            fn=example_video_updated,
            cache_examples=False,
            examples_per_page=50,
            # examples_per_page=2
        )
    with gr.Row():
        def example_file_updated(
            preview_imgs,
            inputs,
            pca_slider,
            bright_slider,
            is_example,
        ):
            return inputs,pca_slider,bright_slider,is_example
        gr.Examples(
            examples=examples_pcd,
            inputs=[
                preview_imgs,
                input_file,
                pca_slider,
                bright_slider,
                is_example,
            ],
            outputs=[
                input_file,
                pca_slider,
                bright_slider,
                is_example,
            ],
            label = "Point Cloud Examples",
            fn=example_file_updated,
            cache_examples=False,
            examples_per_page=50,
            # examples_per_page=2
        )

    reconstruction_btn.click(
        fn = update_gallery_on_upload,
        inputs = [input_file,input_video,conf_thres,frame_slider,prediction_mode,TSDF_mode],
        outputs = [original_view, target_dir, image_gallery, log_output]
    )
    submit_btn.click(fn=clear_fields, inputs=[], outputs=[processed_view]).then(
        fn=PCAing_log, inputs=[is_example, log_output], outputs=[log_output]
    ).then(
        fn=gradio_demo,
        inputs=[target_dir,pca_slider,bright_slider, model_type, if_color, if_normal],
        outputs=[processed_view,log_output],
    ).then(
        fn=lambda: "False", inputs=[], outputs=[is_example]  # set is_example to "False"
    )

    pca_slider.change(fn=clear_fields, inputs=[], outputs=[processed_view]).then(
        fn=PCAing_log, inputs=[is_example, log_output], outputs=[log_output]
    ).then(
        fn=concerto_slider_update,
        inputs=[target_dir,pca_slider,bright_slider,is_example,log_output],
        outputs=[processed_view, log_output],
    ).then(
        fn=lambda: "False", inputs=[], outputs=[is_example]  # set is_example to "False"
    )
    bright_slider.change(fn=clear_fields, inputs=[], outputs=[processed_view]).then(
        fn=PCAing_log, inputs=[is_example, log_output], outputs=[log_output]
    ).then(
        fn=concerto_slider_update,
        inputs=[target_dir,pca_slider,bright_slider,is_example,log_output],
        outputs=[processed_view, log_output],
    ).then(
        fn=lambda: "False", inputs=[], outputs=[is_example]  # set is_example to "False"
    )
    model_type.change(fn=clear_fields, inputs=[], outputs=[processed_view]).then(
        fn=PCAing_log, inputs=[is_example, log_output], outputs=[log_output]
    ).then(
        fn=gradio_demo,
        inputs=[target_dir,pca_slider,bright_slider, model_type, if_color, if_normal],
        outputs=[processed_view,log_output],
    ).then(
        fn=lambda: "False", inputs=[], outputs=[is_example]  # set is_example to "False"
    )
    
    input_file.change(fn=reset_log, inputs=[], outputs=[log_output]).then(
        fn=update_gallery_on_upload,
        inputs=[input_file,input_video, conf_thres,frame_slider,prediction_mode,TSDF_mode],
        outputs=[original_view, target_dir, _, log_output],
    )

    input_video.change(fn=reset_log, inputs=[], outputs=[log_output]).then(
        fn=update_gallery_on_upload,
        inputs=[input_file,input_video, conf_thres,frame_slider,prediction_mode,TSDF_mode],
        outputs=[original_view, target_dir, image_gallery, log_output],
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_error=True, share=True)