import os
import tqdm
import random
import numpy as np
import pandas as pd
import plotly.graph_objects as go

from sklearn.cluster import KMeans, MiniBatchKMeans


def generate_clusters(folder, n_cluster):
    frames = os.listdir(os.path.join(folder, 'pose_shape'))
    frames = [frame.split('.')[0] for frame in frames]
    # random.shuffle(frames)
    # frames = frames[:5000]

    all_poses = np.zeros((len(frames), 69), np.float32)

    i = 0
    for frame in tqdm.tqdm(frames):
        # id, frame_nr = frame.split('_')
        # frame_nr = frame_nr.split('.')[0]

        pose_shape = np.load(os.path.join(folder, 'pose_shape', frame + '.npy'))
        pose = pose_shape[3:72]

        all_poses[i, :] = pose

        i+= 1


    # kmeans = KMeans(n_clusters=20, random_state=0).fit(all_poses)
    kmeans = MiniBatchKMeans(n_clusters=n_cluster, random_state=0, verbose=1).fit(all_poses)

    data = {'frame': frames, 'label': kmeans.labels_}

    df = pd.DataFrame(data)
    df.to_csv('custom_data/' + folder.split('/')[-1] + '_pose_cluster_' + str(n_cluster) + '.csv', index=False)


def sample_from_cluster(file, n_sample, output_file):
    df = pd.read_csv(file)
    samples = df.groupby('label', group_keys=False).apply(lambda x: x.sample(n_sample))

    sampled_frames = np.asarray(samples['frame'].array)

    with open(os.path.join(folder, output_file), 'w') as f:
        for i, frame in enumerate(sampled_frames):
            f.write(frame + '\n')


if __name__ == '__main__':
    folder = '/home/workspace/data/dataset/test_ready'

    is_skirt = False
    F = np.load("custom_utils/smpl_faces" + ("_skirt" if is_skirt == 1 else "") + ".npy")

    # generate_clusters(folder, 20)
    sample_from_cluster('custom_data/test_ready_pose_cluster_20.csv', 10, '/home/workspace/data/dataset/test_ready/frames_pose_cluster_20_sample_10.txt')


    # samples = df.groupby('label', group_keys=False).apply(lambda x: x.sample())

    # sampled_frames = np.asarray(samples['frame'].array)
    # sampled_labels = np.asarray(samples['label'].array)

    # for i in range(sampled_frames.shape[0]):
    #     body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', sampled_frames[i] + '.npz')); body_no_orient = body_no_orient['mesh_posed_body_no_orient']

    #     x = body_no_orient
    #     fig = go.Figure(data=[
    #             go.Mesh3d(
    #                 # 8 vertices of a cube
    #                 x=x[:,0],
    #                 y=x[:,1],
    #                 z=x[:,2],
    #                 # i, j and k give the vertices of triangles
    #                 i = F[:,0],
    #                 j = F[:,1],
    #                 k = F[:,2],
    #                 text='cluster: ' + str(sampled_labels[i])
    #             ),
    #         ])

    #     fig.update_layout(scene_aspectmode='manual',
    #                     scene_aspectratio=dict(x=1, y=1, z=1))
    #     fig.update_layout(
    #         scene = dict(
    #             xaxis = dict(nticks=4, range=[-1.5,1.5],),
    #             yaxis = dict(nticks=4, range=[-1.5,1.5],),
    #             zaxis = dict(nticks=4, range=[-1.5,1.5],),)
    #             )

    #     fig.show()

    #     print()


    print()