import numpy as np
from pytorch3d.structures import Meshes
import os

from models.VAEs.animate_any_mesh_testnms import VariationalAutoEncoder
from models.Diffusions.transformer_modified import BsplineVarianceDiT
from omegaconf import OmegaConf
import torch
from b_spline import make_clamped_uniform_knots, bspline_basis_matrix_cox, B_spline_compute
from utils.rf_diffusion import rf_sample
from utils.render import clear_scene, import_glb, get_all_vertices, get_all_faces,drive_mesh_with_trajs_frames
from utils.mesh_utils import merge_vertices_with_indices

# Config paths
vae_config_path = "configs/vae_animate.yml"
diffusion_config_path = "configs/diffusion.yml"


def normalize_vertices_faces(
    vertices: torch.Tensor,
    faces: torch.Tensor,
    scale: float = 1.8,
    eps: float = 1e-12,
):
    """
    Normalize a single mesh (torch tensors) to fit roughly inside a box of size `scale`.

    Args:
        vertices: torch.Tensor (N,3) or (N,3) float32/64
        faces: torch.LongTensor (F,3)
        scale: target extent after normalization (same meaning as your original scale=1.8)
        eps: small value to avoid division by zero

    Returns:
        norm_vertices: torch.Tensor (N,3) normalized vertices
        faces: torch.LongTensor (F,3) (returned as-is, dtype long)
        vertex_normals: torch.Tensor (N,3) computed on the normalized mesh
        transform: dict with keys:
            - center (torch.Tensor (3,))
            - extent (float)
            - scale (float)
            - scale_factor (float)
    """
    if vertices.dim() != 2 or vertices.shape[1] != 3:
        raise ValueError("vertices must be shape (N,3)")
    if faces.dim() != 2 or faces.shape[1] != 3:
        raise ValueError("faces must be shape (F,3)")

    device = vertices.device
    dtype = vertices.dtype

    # bbox
    b_min = torch.min(vertices, dim=0).values
    b_max = torch.max(vertices, dim=0).values
    extent_vec = b_max - b_min
    extent = float(torch.max(extent_vec).item())

    # avoid division by zero
    if extent <= eps:
        scale_factor = 1.0
    else:
        scale_factor = float(scale) / extent

    center = (b_max + b_min) / 2.0

    # normalize
    norm_vertices = (vertices - center.to(dtype=dtype, device=device)) * float(scale_factor)

    # build PyTorch3D mesh (single mesh in a batch)
    mesh = Meshes(verts=[norm_vertices], faces=[faces.long().to(device)])

    # compute vertex normals using pytorch3d ops (returns packed normals)
    verts_packed = mesh.verts_packed()      # (N,3)
    faces_packed = mesh.faces_packed()      # (F,3)

    # get packed normals (shape (N,3) for the packed verts of the single mesh)
    v_normals = mesh.verts_normals_packed()   # <- use this (or .verts_normals_list()[0])

    transform = {
        "center": center.detach().cpu(),         # keep on cpu to be easy to serialize
        "extent": extent,
        "scale": float(scale),
        "scale_factor": float(scale_factor)
    }

    return norm_vertices, faces.long(), v_normals, transform


def denormalize_vertices(norm_vertices: torch.Tensor, transform) -> torch.Tensor:
    """
    Denormalize normalized vertices back to original coordinates.

    Args:
        norm_vertices: torch.Tensor, can be shape
            - (N,3) single frame
            - (T,N,3) sequence of frames
            - (B,N,3) batched meshes
        transform: dict returned by normalize_vertices_faces (must contain 'scale_factor' and 'center')

    Returns:
        orig_vertices: torch.Tensor with same leading dims as norm_vertices, restored to original scale & position
    """
    sf = float(transform["scale_factor"])
    center = torch.as_tensor(transform["center"], dtype=norm_vertices.dtype, device=norm_vertices.device)

    # broadcasting: handle (N,3), (T,N,3), (B,N,3) etc.
    # put center into shape (1,...,1,3) that broadcasts to norm_vertices
    if norm_vertices.dim() == 2:         # (N,3)
        return norm_vertices / sf + center
    else:
        # create shape of ones for broadcasting (1, ..., 1, 3)
        expand_shape = [1] * (norm_vertices.dim() - 1) + [3]
        center_view = center.view(*expand_shape)
        return norm_vertices / sf + center_view


# Load checkpoint with possible 'module.' prefix handling
def load_compatible_checkpoint(model, ckpt_path, device):
    if not os.path.exists(ckpt_path):
        raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}")
    print(f"Loading checkpoint from: {ckpt_path}")
    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
        model_weights = checkpoint['model_state_dict']
    else:
        model_weights = checkpoint
    if list(model_weights.keys())[0].startswith('module.'):
        model_weights = {k.replace('module.', ''): v for k, v in model_weights.items()}
    model.load_state_dict(model_weights)
    print("Model weights loaded successfully.")
    return model

# Set random seed for reproducibility
def set_seed(seed):
    import random
    import torch
    import numpy as np
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

# Argument parser
def set_parse():
  import argparse
  parser = argparse.ArgumentParser()
  parser.add_argument("--seed", type=int, default=3407)
  parser.add_argument("--input_glb_path", type=str, default="") # the input glb file path
  parser.add_argument("--glb_name", type=str, default="dragon",help="the object name without .glb")
  parser.add_argument("--generated_seq_length", type=int, default=16)
  parser.add_argument("--vae_path", type=str, default="") # the trained VAE model path
  parser.add_argument("--dit_path", type=str, default="") # the trained diffusion model path
  parser.add_argument("--save_path", type=str, default="") # the output folder path
  parser.add_argument("--prompt", type=str, default="") # the text prompt for motion generation
  parser.add_argument("--guidance_scale", type=float, default=3.0)
  parser.add_argument("--azi", type=float, default=0.0) # the azimuth angle for rendering
  parser.add_argument("--ele", type=float, default=0.0) # the elevation angle for rendering
  parser.add_argument("--export_format", type=str, default="none", choices=["none", "abc", "fbx","glb"])
  opt = parser.parse_args()
  return opt
  
if __name__ == "__main__":


    opt = set_parse()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    set_seed(opt.seed)

    import time

    # Loading scaling files
    deformation_mean = torch.load("dataset/data_statistics/motion_mean.pt",weights_only=True).to(torch.float32).to(device)
    deformation_std  = torch.load("dataset/data_statistics/motion_std.pt",weights_only=True).to(torch.float32).to(device)
    static_mean = torch.load("dataset/data_statistics/static_mean.pt",weights_only=True).to(torch.float32).to(device)
    static_std = torch.load("dataset/data_statistics/static_std.pt",weights_only=True).to(torch.float32).to(device)

    # Loading VAE and Diffusion
    VAE_config = OmegaConf.load(vae_config_path)
    VAE_model = VariationalAutoEncoder(**VAE_config['model']).to(device)
    VAE_model = load_compatible_checkpoint(VAE_model, opt.vae_path, device)
    VAE_model.eval()
    print("DyMeshVAE loaded Successful!!!")

    DIT_config = OmegaConf.load(diffusion_config_path)
    diffusion_model = BsplineVarianceDiT(**DIT_config['model'],training=False).to(device)
    diffusion_model = load_compatible_checkpoint(diffusion_model, opt.dit_path, device)
    diffusion_model.eval()
    print("RF Model loaded Successful!!!")

  
    # Building the Bspline basis for trajectory prediction
    knots = make_clamped_uniform_knots(16, 3)
    denom_clamp=1e-6
    full_timesteps = torch.linspace(0.0, 1.0, steps = opt.generated_seq_length)           
    B_s = bspline_basis_matrix_cox(3, knots, full_timesteps).cuda()


    # creating save dir
    save_video_folder = os.path.join(opt.save_path,opt.glb_name)
    if not os.path.exists(save_video_folder):
        os.makedirs(save_video_folder)


    t1 = time.time()
    # Loading glb file and preprocess:
    glb_file_path = os.path.join(opt.input_glb_path,opt.glb_name+".glb")
    clear_scene()
    mesh_objects = import_glb(glb_file_path)
    all_vertices = get_all_vertices(mesh_objects)
    all_faces = get_all_faces(mesh_objects)

    merged_verts, merged_faces, all_indices = merge_vertices_with_indices(all_vertices, all_faces)


    print('merged_verts: ',merged_verts.shape, merged_faces.shape)

    norm_v, faces, norm_n, transform = normalize_vertices_faces(merged_verts, merged_faces)
    print('loading mesh success!!!')


    with torch.no_grad():
        vertices_tensor = torch.tensor(norm_v, dtype=torch.float32).unsqueeze(0).to(device)  # (1, N, 3)
        norm_tensor = torch.tensor(norm_n, dtype=torch.float32).unsqueeze(0).to(device)
        vertices_repeat = vertices_tensor.unsqueeze(0).repeat(1, 16, 1, 1) 
        # encode
        kl, x, idx, pc0_embed_ori,x0,posterior =  VAE_model.encode(vertices_tensor,norm_tensor,vertices_repeat)

        # diffusion sampling
        x0_latent, xt_latent = x.chunk(2, dim=-1)
        x0_latent = (x0_latent - static_mean) / static_std
        xt_latent = (xt_latent - deformation_mean) / deformation_std
        x_start = torch.cat([x0_latent, xt_latent], dim=-1)
        samples = rf_sample(model=diffusion_model, 
                            shape=x_start.shape, 
                            text_prmpt=[opt.prompt], 
                            guidance_scale=opt.guidance_scale, 
                            device=device, 
                            f0=x0_latent)

        x0_recon = samples[:, :, :x0_latent.shape[-1]]
        xt_recon = samples[:, :, x0_latent.shape[-1]:]
        
        # scaling back
        x0_recon = x0_recon * static_std + static_mean
        xt_recon = xt_recon * deformation_std + deformation_mean
        x = torch.cat([x0_recon, xt_recon], dim=-1)
        # decode
        outputs = VAE_model.decode(x,vertices_tensor,norm_tensor,pc0_embed_ori) # bspline control points
        
        # predicted bspline trajectories: T*N*3
        P_pred_bspline = B_spline_compute(B_s,outputs[0,:,:,:], torch.ones_like(outputs[0, :, :, 0]), denom_clamp)
        pred_deformated_pc_bspline = P_pred_bspline + vertices_tensor.repeat(opt.generated_seq_length,1,1)

        pred_deformated_pc_bspline = torch.cat([vertices_tensor,pred_deformated_pc_bspline],dim=0)

        trajs = [pred_deformated_pc_bspline[:, idx].detach().cpu() for idx in all_indices]

        # rendering the driven mesh with predicted trajectories
        drive_mesh_with_trajs_frames(mesh_objects, trajs, save_video_folder , azi=opt.azi, ele=opt.ele, export_format=opt.export_format)
    
    t2 = time.time()

    print('Total inference time:', t2 - t1)
    print(f"Current GPU memory usage: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
    print(f"Peak GPU memory usage: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")

