'''
@Author: Jiapeng Zhou
@Email: jiapengzhou@link.cuhk.edu.cn
@Desc: script to process all data of RECMV. 
Before run this bash, what you only need to do is :
- run openpose to get keypoints.json file in the dir: (self.root, openpose)
- have a conda environment of videoavatar, named: "videoavatar"
Usage: python process_data_all.py --root  ~/DeepDxxx/xinyu_a/ --gid 0 --gender female
'''

import pdb, os, os.path as osp, sys, argparse
from typing import Any
import numpy as np, cv2

class RECMVDataProcessor():
    def __init__(self, args) -> None:
        self.args = args
        self.gender = args.gender

        self.root = args.root
        #out = osp.join(self.root, args.out)
        #os.makedirs(out, exist_ok=True)
        self.img_dir = osp.join(self.root, 'imgs')
        assert osp.exists(self.img_dir) and len(os.listdir(self.img_dir)) > 0, 'img file does not exist'
        self.gid = args.gid
    
    def __len__(self):
        return len(os.listdir(self.img_dir))
    
    def __call__(self):
        '''get normals, ATR-parsing, masks, camera, smpl_params, depth, featurelines'''

        # get ATR_parsing 
        self.get_ATR_parsing()
        
        # get masks
        self.get_masks()

        # get normal from pifuhd
        self.get_normals()

        '''cannot directly ca videoavatar
        # get reconstructed_poses.hdf5 from videoavatar
        self.get_poses()
        
        # get camera.npz and smpl_rec.npz
        self.get_camera_smpl_rec()

        '''
        
        # get depth
        #self.get_depth()
        
        # get featurelines
        #self.get_featurelines()

    def get_featurelines(self):
        fl_dir = osp.join(self.root, 'featurelines')
        if osp.exists(fl_dir) and len(os.listdir(fl_dir)) >= self.__len__():
            print('featurelines files exists')
        else:
            print('featurelines file does not exist')
    
    def get_depth(self):
        depth_dir = osp.join(self.root, 'depth')
        if osp.exists(depth_dir) and len(os.listdir(depth_dir)) >= self.__len__():
            print('depth files exists')
        else:
            print('depth file does not exist')
    
    def get_masks(self):
        mask_dir = osp.join(self.root, 'masks')
        if osp.exists(mask_dir) and len(os.listdir(mask_dir)) >= self.__len__():
            print('mask files exists')
            return
        os.makedirs(mask_dir, exist_ok=True)
        # mask是否需要按照单通道保存？
        print('\n-----------------Getting Masks-------------------------')
        os.system(f'cd ./RobustVideoMatting && python inference_itw_rotate.py \
                  --input_pth {self.img_dir} --output_pth {mask_dir} --device {self.gid}')

    
    def get_normals(self):
        normal_dir = osp.join(self.root, 'normals')
        if osp.exists(normal_dir) and len(os.listdir(normal_dir)) >= self.__len__():
            print('normal files exists')
            return

        print('\n-------------------Getting Normals----------------------')
        os.makedirs(normal_dir, exist_ok=True)
        os.system('cd ./pifuhd && python ./generate_normals.py --gid {} --imgpath {}'.format(self.gid, self.img_dir))

    def get_ATR_parsing(self):
        input_dir = osp.join(self.root, 'imgs')
        output_dir = osp.join(self.root, 'parsing_SCH_ATR')
        if osp.exists(output_dir) and len(os.listdir(output_dir)) >= self.__len__():
            print('ATR parsing files exists')
            return
        print('\n--------------------Getting ATR_Parsing------------------')
        os.makedirs(output_dir, exist_ok=True)

        os.system(f'cd ./Self-Correction-Human-Parsing && python simple_extractor.py \
                  --dataset atr --model-restore exp-schp-201908301523-atr.pth --gpu {self.gid} \
                  --input-dir {self.img_dir} --output-dir {output_dir} --logits')


    def get_poses(self):
        openpose_p = osp.join(self.root, 'openpose')
        assert osp.exists(openpose_p) and len(os.listdir(openpose_p)) >= self.__len__(), 'openpose keypoint folder doesnot exist or it is emtpy'

        output_dir = osp.join(self.root, 'videoavatars')
        if osp.exists(output_dir) and osp.exists(osp.join(output_dir, 'reconstructed_poses.hdf5')):
            print('reconstructed_poses.hdf5 files exists')
            return

        print('\n------------------Getting Poses---------------------------')
        os.makedirs(output_dir, exist_ok=True)
        # Note: Here needs to conda actiavte the environment of videoavatar
        os.system(f'cd ./videoavatars && bash get_reconstructed_poses.sh {self.root} {output_dir} {self.gender}')


    def get_camera_smpl_rec(self):
        if osp.exists(osp.join(self.root, 'camera.npz')) and osp.exists(osp.join(self.root, 'smpl_rec.npz')):
            print('camera.npz and smpl_rec.npz files exists')
            return

        print('\n---------------------Getting camera and smpl_rec--------------------')
        os.system(' python get_smpl_rec_camera.py --root {} --save_root {} --gender {}'.\
                      format(osp.join(self.root, 'videoavatars'), self.root, self.gender))


        

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--gid', type=int, default=0)
    parser.add_argument('--root', type=str, required=True)
    parser.add_argument('--gender', type=str, choices=['male', 'female'], required=True)
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    processor = RECMVDataProcessor(args)
    processor()
    
if __name__ == '__main__':
    main()