import os
import sys

import math
import shutil
import numpy as np
import scipy.io as sio

from tqdm import tqdm

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)


    pass



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)