import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from evo.tools.file_interface import read_kitti_poses_file
from pathlib import Path
from visualization.visualizer import visualize_simulation


def load_trajectory_data(traj_file, char_file, num_cams=30):
    trajectory = read_kitti_poses_file(traj_file)
    matrix_trajectory = torch.from_numpy(np.array(trajectory.poses_se3)).to(torch.float32)
    
    raw_trans = torch.clone(matrix_trajectory[:, :3, 3])
    raw_rot = matrix_trajectory[:, :3, :3]
    rot6d = raw_rot[:, :, :2].permute(0, 2, 1).reshape(-1, 6)
    
    trajectory_feature = torch.hstack([rot6d, raw_trans]).permute(1, 0)
    
    padded_trajectory_feature = F.pad(
        trajectory_feature, 
        (0, num_cams - trajectory_feature.shape[1])
    )
    
    padding_mask = torch.ones((num_cams))
    padding_mask[trajectory_feature.shape[1]:] = 0
    
    char_feature = torch.from_numpy(np.load(char_file)).to(torch.float32)
    padding_size = num_cams - char_feature.shape[0]
    padded_char_feature = F.pad(char_feature, (0, 0, 0, padding_size)).permute(1, 0)
    
    return {
        "traj_filename": Path(traj_file).name,
        "char_filename": Path(char_file).name,
        "traj_feat": padded_trajectory_feature,
        "char_feat": padded_char_feature,
        "padding_mask": padding_mask,
        "raw_matrix_trajectory": matrix_trajectory
    }

def create_trajectory_interface():
    """Create Gradio interface for loading trajectory data"""
    
    def process_files(traj_file, char_file):
        try:
            result = load_trajectory_data(traj_file.name, char_file.name)
            
            # Convert tensors to numpy for display
            info = {
                "Trajectory filename": result["traj_filename"],
                "Character filename": result["char_filename"],
                "Trajectory shape": result["traj_feat"].shape,
                "Character shape": result["char_feat"].shape,
                "Valid frames": int(result["padding_mask"].sum().item())
            }
            
            return str(info)
            
        except Exception as e:
            return f"Error processing files: {str(e)}"

    interface = gr.Interface(
        fn=process_files,
        inputs=[
            gr.File(label="Trajectory File (.txt)"),
            gr.File(label="Character File (.npy)")
        ],
        outputs=gr.Textbox(label="Results"),
        title="Trajectory Data Loader",
        description="Upload trajectory (.txt) and character (.npy) files to load and process them."
    )
    
    return interface

if __name__ == "__main__":
    interface = create_trajectory_interface()
    interface.launch()