import pickle
from tqdm import tqdm
from pathlib import Path
import os
import os.path as osp
import argparse
import logging

'''
    gernerate the 17 Number of Pose Points Format from 18 Number of Pose Points
    OUMVLP 17
               # keypoints = {
            #     0: "nose",
            #     1: "left_eye",
            #     2: "right_eye",
            #     3: "left_ear",
            #     4: "right_ear",
            #     5: "left_shoulder",
            #     6: "right_shoulder",
            #     7: "left_elbow",
            #     8: "right_elbow",
            #     9: "left_wrist",
            #     10: "right_wrist",
            #     11: "left_hip",
            #     12: "right_hip",
            #     13: "left_knee",
            #     14: "right_knee",
            #     15: "left_ankle",
            #     16: "right_ankle"
            # }
    OUMVLP 18
    mask=[0,15,14,17,16,5,2,6,3,7,4,11,8,12,9,13,10]
           # keypoints = {
            #     0: "nose",  
            #     1: "neck",
            #     2: "Rshoulder",
            #     3: "Relbow",
            #     4: "Rwrist",
            #     5: "Lshoudler",
            #     6: "Lelbow",
            #     7: "Lwrist",
            #     8: "Rhip",
            #     9: "Rknee",
            #     10: "Rankle",
            #     11: "Lhip",
            #     12: "Lknee",
            #     13: "Lankle",
            #     14: "Reye",
            #     15: "Leye",
            #     16: "Rear",
            #     17: "Lear"
            # }
'''

def ToOUMVLP17(input_path: Path, output_path: Path):
    mask=[0,15,14,17,16,5,2,6,3,7,4,11,8,12,9,13,10]
    TOTAL_SUBJECTS = 10307
    progress = tqdm(total=TOTAL_SUBJECTS)

    for subject in input_path.iterdir():
        output_subject = subject.name
        for seq in subject.iterdir():
            output_seq = seq.name
            for view in seq.iterdir():
                src = os.path.join(view, f"{view.name}.pkl")
                dst = os.path.join(output_path, output_subject, output_seq, view.name)
                os.makedirs(dst, exist_ok=True)
                with open(src,'rb') as f:
                    srcdata = pickle.load(f)
                    #[T,18,3]
                data = srcdata[...,mask,:].copy()
                # #[T,17,3]
                pkl_path = os.path.join(dst,f'{view.name}.pkl')
                pickle.dump(data,open(pkl_path,'wb')) 
        progress.update(1)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='OpenGait dataset pretreatment module.')
    parser.add_argument('-i', '--input_path', default='', type=str, help='Root path of raw dataset.')
    parser.add_argument('-o', '--output_path', default='', type=str, help='Output path of pickled dataset.')
    parser.add_argument('-l', '--log_to_file', default='./pretreatment.log', type=str, help='Log file path. Default: ./pretreatment.log')
    args = parser.parse_args()
    logging.info('Begin')
    ToOUMVLP17(input_path=Path(args.input_path), output_path=Path(args.output_path))
    logging.info('Done')