import os
import sys

import math
import shutil
import warnings
import tables as tb
import numpy as np
import scipy.io as sio

from tqdm import tqdm

from tables import NaturalNameWarning
warnings.filterwarnings('ignore', category=NaturalNameWarning)

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import readPC2, loadInfo
from data.utils import Frame




def clean_dataset(folder):
    ids = os.listdir(folder)

    for id in tqdm(ids):

        if not os.path.isdir(os.path.join(folder, id)):
            continue

        files = os.listdir(os.path.join(folder, id))

        for file in files:
            if file.startswith('glmarks') or file.find('nricp') != -1:
                os.remove(os.path.join(folder, id, file))

 
def generate_frames_file(folder, file_name):
    ids = os.listdir(folder)
    frames = []

    for id in tqdm(ids):

        if not os.path.isdir(os.path.join(folder, id)):
            continue

        garments = [f.replace('.obj','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('.obj')]
        info = loadInfo(os.path.join(folder, id, 'info.mat'))

        for garment in garments:
            for i in range(info['poses'].shape[1]):
                frames.append(Frame(id, i, garment))

    with open(os.path.join(folder, file_name), 'w') as f:
        for i, frame in enumerate(frames):
            f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n')


def split_data(folder, ratio):
    with open(os.path.join(folder, 'frames_all.txt'), "r") as f:
        frames = f.readlines()

    ids = sorted(list(set([frame.split()[0] for frame in frames])))
    train_last_id = round(len(ids) * ratio)

    train_frames = []
    valid_frames = []

    for frame in tqdm(frames):
        id = frame.split()[0]
        if id in ids[:train_last_id]:
            train_frames.append(frame)
        else:
            valid_frames.append(frame)

    with open(os.path.join(folder, f'frames_train_{ratio}.txt'), 'w') as f:
        f.writelines(train_frames)

    with open(os.path.join(folder, f'frames_valid_{round(1-ratio,1)}.txt'), 'w') as f:
        f.writelines(valid_frames)


def genereated_registered_and_mask(folder, is_skirt = False):
    ids = os.listdir(folder)

    for id in tqdm(ids):
        garment_dir = os.path.join(folder, id)
        if not os.path.isdir(garment_dir):
            continue
        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')]
        
        assert len(garments) == 1

        if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')):
            continue
        
        garment = garments[0]

        verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
        map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1))
        verts = verts[:, map, :] # ---------------> This is the garment after registration
        raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1))
        mask = raw_mask > 0
        # boundaries = raw_mask == 2

        writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True)
        np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask)


def create_extras_file(folder):
    # create a new pytables file
    root = '/'.join(folder.split('/')[:-1])
    type = folder.split('/')[-1]
    h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='w', title='Extras')

    frames = []
    with open(os.path.join(folder, 'frames_all.txt'), 'r') as f:
        lines = f.readlines()
        lines = [line.strip() for line in lines]
    frames = lines

    for frame in tqdm(frames[:3000]):
        # add a new empty group to the root group for each frame
        id, frame_nr, garment = frame.split()
        if not id in h5file.root:
            h5file.create_group('/', id)
        if not frame_nr in h5file.root[id]:
            h5file.create_group(f'/{id}', frame_nr)
        h5file.create_group(f'/{id}/{frame_nr}', f'{garment}')


    h5file.close()



if __name__ == '__main__':

    folder = '/home/workspace/data/dataset/cloth3d/train'

    # clean_dataset(folder)
    # generate_frames_file(folder, 'frames_all.txt')
    # split_data(folder, 0.9)
    create_extras_file(folder)