import argparse
import glob
import json
import yaml
from pathlib import Path
import os
import re
import numpy as np

from scipy.spatial.transform import Rotation
from tqdm import tqdm
import torch
import trimesh

from primitive_anything.primitive_dataset import create_dataset
from primitive_anything.utils import torch_to, count_parameters
from primitive_anything.utils.logger import create_logger, print_log

CODE_SHAPE = {
    0: 'SM_GR_BS_CubeBevel_001.ply',
    1: 'SM_GR_BS_SphereSharp_001.ply',
    2: 'SM_GR_BS_CylinderSharp_001.ply',
}

shapename_map = {
    'SM_GR_BS_CubeBevel_001.ply': 1101002001034001,
    'SM_GR_BS_SphereSharp_001.ply': 1101002001034010,
    'SM_GR_BS_CylinderSharp_001.ply': 1101002001034002,
}

bs_dir = 'data/basic_shapes_norm'
mesh_bs = {}
for bs_path in glob.glob(os.path.join(bs_dir, '*.ply')):
    bs_name = os.path.basename(bs_path)
    bs = trimesh.load(bs_path)
    bs.visual.uv = np.clip(bs.visual.uv, 0, 1)
    bs.visual = bs.visual.to_color()
    mesh_bs[bs_name] = bs


def create_model(cfg_model):
    kwargs = cfg_model
    name = kwargs.pop('name')
    model = get_model(name)(**kwargs)
    print_log("Model '{}' init: nb_params={:,}, kwargs={}".format(name, count_parameters(model), kwargs))
    return model


from primitive_anything.primitive_transformer import PrimitiveTransformerDiscrete
def get_model(name):
    return {
        'discrete': PrimitiveTransformerDiscrete,
    }[name]


def euler_to_quat(euler):
    return Rotation.from_euler('XYZ', euler, degrees=True).as_quat()

def rotvec_to_quat(rotvec):
    return Rotation.from_rotvec(rotvec, degrees=True).as_quat()

def SRT_quat_to_matrix(scale, quat, translation):
    rotation_matrix = Rotation.from_quat(quat).as_matrix()
    transform_matrix = np.eye(4)
    transform_matrix[:3, :3] = rotation_matrix * scale
    transform_matrix[:3, 3] = translation
    return transform_matrix

def write_json(primitives, shapename_map, out_path):
    out_json = {}
    out_json['operation'] = 0
    out_json['type'] = 1
    out_json['scene_id'] = None

    new_group = []
    model_scene = trimesh.Scene()
    for scale, rotation, translation, type_code in zip(
        primitives['scale'].squeeze().cpu().numpy(),
        primitives['rotation'].squeeze().cpu().numpy(),
        primitives['translation'].squeeze().cpu().numpy(),
        primitives['type_code'].squeeze().cpu().numpy()
    ):
        if type_code == -1:
            break
        bs_name = CODE_SHAPE[type_code]
        new_block = {}
        new_block['type_id'] = shapename_map[bs_name]
        new_block['data'] = {}
        new_block['data']['location'] = translation.tolist()
        new_block['data']['rotation'] = euler_to_quat(rotation).tolist()
        new_block['data']['scale'] = scale.tolist()
        new_block['data']['color'] = ['808080']
        new_group.append(new_block)

        if new_block['type_id'] == 1101002001034001:
            cur_color = "#2FA9FF"
        elif new_block['type_id'] == 1101002001034002:
            cur_color = "#FFC203"
        elif new_block['type_id'] == 1101002001034010:
            cur_color = "#FF8A9C"

        def hex_to_rgb(hex_color):
            hex_color = hex_color.lstrip('#')
            return np.array([
                int(hex_color[0:2], 16),  # R
                int(hex_color[2:4], 16),  # G
                int(hex_color[4:6], 16),  # B
            ], dtype=np.uint8)[None]

        trans = SRT_quat_to_matrix(scale, euler_to_quat(rotation), translation)
        bs = mesh_bs[bs_name].copy().apply_transform(trans)
        new_vertex_colors = np.repeat(hex_to_rgb(cur_color), bs.visual.vertex_colors.shape[0], axis=0)
        bs.visual.vertex_colors[:, :3] = new_vertex_colors
        vertices = bs.vertices.copy()
        vertices[:, 1] = bs.vertices[:, 2]
        vertices[:, 2] = -bs.vertices[:, 1]
        bs.vertices = vertices
        model_scene.add_geometry(bs)

    out_json['group'] = new_group

    with open(out_path, 'w') as json_file:
        json.dump(out_json, json_file, indent=4)

    glb_path = out_path.replace('.json', '.glb')
    model_scene.export(glb_path)

    return glb_path, out_json


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, default='./configs/infer.yml', help='Config file path')
    parser.add_argument('-ck', '--AR_ckpt', type=str, default='./ckpt/mesh-transformer.ckpt.60.pt')
    parser.add_argument('-o', '--output', type=str, default='./results/infer')
    parser.add_argument('--bs_dir', type=str, default='data/basic_shapes_norm')
    parser.add_argument('--temperature', type=float, default=0.0)
    args = parser.parse_args()

    bs_names = []
    for bs_path in glob.glob(os.path.join(args.bs_dir, '*.ply')):
        bs_names.append(os.path.basename(bs_path))

    with open(args.config, mode='r') as fp:
        cfg = yaml.load(fp, Loader=yaml.FullLoader)

    AR_checkpoint = torch.load(args.AR_ckpt)
    
    os.makedirs(args.output, exist_ok=True)
    json_result_folder = os.path.join(args.output, 'JsonResults')
    os.makedirs(json_result_folder, exist_ok=True)

    create_logger(Path(args.output))

    dataset = create_dataset(cfg['dataset'])

    transformer = create_model(cfg['model'])
    transformer.load_state_dict(AR_checkpoint)

    for item_i, item in tqdm(enumerate(dataset)):
        pc = item.pop('pc')
    
        item_filename = dataset.data_filename[item_i]
        if torch.cuda.is_available():
            pc = pc.cuda()
            item = torch_to(item, torch.device('cuda'))
            transformer = transformer.cuda()

        recon_primitives, mask = transformer.generate(pc=pc.unsqueeze(0), temperature=args.temperature)

        out_path = os.path.join(json_result_folder, os.path.basename(item_filename).replace('.ply', '.json'))
        write_json(recon_primitives, shapename_map, out_path)