import os
import numpy as np
import plotly.express as ex
from simpleicp import PointCloud, SimpleICP
from scipy.spatial.distance import cdist

folder = '/home/workspace/data/dataset/tailornet'

def get_data_split():
    data = np.load(os.path.join(folder, 'split_static_pose_shape.npz'))
    return data['train'], data['test']


def get_garment_info(garment_class):
    data = np.load(os.path.join(folder, 'garment_class_info.pkl'), allow_pickle=True)
    return data[garment_class]


def get_smpl():
    data = np.load(os.path.join('TailorNet_dataset/smpl/smpl_female.npz'), allow_pickle=True)
    data_high_res = np.load(os.path.join('TailorNet_dataset/smpl/smpl_hres_female.npz'), allow_pickle=True)
    supersmpl = np.load(os.path.join('custom_utils/super_model_f.pkl'), allow_pickle=True)

    return data_high_res['v_template'], supersmpl['v_template']

    print()


def register_smpl(verts_1, verts_2):

    # Create point cloud objects
    pc_fix = PointCloud(verts_1, columns=["x", "y", "z"])
    pc_mov = PointCloud(verts_2, columns=["x", "y", "z"])

    # Create simpleICP object, add point clouds, and run algorithm!
    icp = SimpleICP()
    icp.add_point_clouds(pc_fix, pc_mov)
    H, X_mov_transformed, rigid_body_transformation_params = icp.run(correspondences=10000, min_change=0.01, neighbors=25)

    dist = cdist(verts_1, X_mov_transformed)
    matching_dist = np.min(dist, axis=1)
    matching = np.argmin(dist, axis=1)
    

    print()


if __name__ == '__main__':
    train_split, test_split = get_data_split()
    garment_info = get_garment_info('old-t-shirt')
    smpl = get_smpl()
    tailornet_smpl, super_smpl = get_smpl()
    # registration = register_smpl(super_smpl, tailornet_smpl)

    print()