####################################################################
#
# 1. extract raw data
# root : S1 / S5 / S6 / S7 / S8 / S9 / S11
# 2. extract video frames
#   python GEN_Human3.6M.py create ./cache/extract_Human3.6M.sh
# 3. generate the list
#   python GEN_Human3.6M.py genlist
# 4. delete all information generated by this script
#   python GEN_Human3.6M.py clean  ./cache/delete_Human3.6M.sh
#
########################################################################
# https://github.com/microsoft/multiview-human-pose-estimation-pytorch #
########################################################################
import os, sys, time, torch, copy
from os import path as osp
import xml.etree.ElementTree as ET
import numpy as np, math
from pathlib import Path
import cdflib


class H36M_Metadata:
  def __init__(self, metadata_file):
    self.subjects = []
    self.sequence_mappings = {}
    self.action_names = {}
    self.camera_ids = []

    tree = ET.parse(metadata_file)
    root = tree.getroot()

    self.resolution = []
    for i, tr in enumerate(root.find('mapping')):
      if i == 0:
        _, _, *self.subjects = [td.text for td in tr]
        self.sequence_mappings = {subject: {} for subject in self.subjects}
      elif i < 33:
        action_id, subaction_id, *prefixes = [td.text for td in tr]
        for subject, prefix in zip(self.subjects, prefixes):
          self.sequence_mappings[subject][(action_id, subaction_id)] = prefix
      elif i == 34:
        for j, x in enumerate(tr):
          if j <= 1: continue
          self.resolution.append( x.text[1:-1].split(';') )
      else:
        print('seems useless [{:}] : {:}'.format(i, tr.text))

    for i, elem in enumerate(root.find('actionnames')):
      action_id = str(i + 1)
      self.action_names[action_id] = elem.text

    self.camera_ids = [elem.text for elem in root.find('dbcameras/index2id')]
    print ('There are {:} camers : {:}'.format(len(self.camera_ids), self.camera_ids))
    print ('There are {:} resolutions'.format(len(self.resolution)))
    self.camera_projectionParams = []
    for elem in root.find('dbcameras/projectionParams'):
      xparams = np.array([float(x) for x in elem.text[1:-1].split(' ')])
      self.camera_projectionParams.append( xparams )
    self.w0 = [float(x) for x in root.find('w0').text[1:-1].split(' ')]

  def get_base_filename(self, subject, action, subaction, camera):
    return '{}.{}'.format(self.sequence_mappings[subject][(action, subaction)], camera)

  # subject in [1,11], camera in [1,4]
  def loadCamera(self, subject, camera):
    start = 6*((camera-1)*11 + (subject-1))
    A = self.w0[start:start+6]
    B = self.w0[265+(camera-1)*9-1:264+camera*9]
    Params = np.array(A+B)
    R = self.rotationMatrix(Params[0], Params[1], Params[2], 'xyz')
    T = Params[3:6]
    f = Params[6:8]
    c = Params[8:10]
    k = Params[10:13]
    p = Params[13:15]
    resolution = self.resolution[subject-1][camera-1]
    resolution = [int(x) for x in resolution.split(' ')]
    return {'subject': subject,
            'camera' : camera,
            'camera_name': self.camera_ids[camera-1],
            'K': k, 'R': R, 'T': T,
            'f': f, 'c': c, 'p': p}
  
  def rotationMatrix(self, xangle, yangle, zangle, order):
    # Here we assume we rotate z, then x then y.
    c1 = math.cos(xangle) # The x angle
    c2 = math.cos(yangle) # The y angle
    c3 = math.cos(zangle) # the z angle
    s1 = math.sin(xangle)
    s2 = math.sin(yangle)
    s3 = math.sin(zangle)
    if order == 'zxy':
      rotMat = [[c2*c3-s1*s2*s3, c2*s3+s1*s2*c3, -s2*c1],
                [-c1*s3        , c1*c3, s1],
                [s2*c3+c2*s1*s3, s2*s3-c2*s1*c3, c2*c1]]
    else:
      rotMat = np.eye(3)
      for o in order:
        if o == 'x':
          temp = [[ 1, 0,0],
                  [ 0,c1,s1],
                  [ 0,-s1,c1]]
        elif o == 'y':
          temp = [[c2,0,-s2],
                  [0 ,1,0],
                  [s2,0,c2]]
        elif o == 'z':
          temp = [[c3 ,s3, 0],
                  [-s3,c3, 0],
                  [0  , 0, 1]]
        else: raise ValueError('unknow order : {:}'.format(order))
        rotMat = np.array(temp) * rotMat
    return rotMat


def process(root, metadata, subjects, info):
  commands, total = [], 0
  for camera_ID, camera in enumerate(metadata.camera_ids):
    base_filename = metadata.get_base_filename(info['subject'], info['action'], info['subaction'], camera)
    CDF_path = root / info['subject'] / 'D2_Positions' / '{:}.cdf'.format(base_filename)
    CDF_path = CDF_path.resolve()
    cdf_file = cdflib.CDF( str(CDF_path) )
    Poses_2d = cdf_file.varget("Pose")
    Poses_2d = Poses_2d.reshape(Poses_2d.shape[1], 32, 2)
    video_path = root / info['subject'] / 'Videos' / '{:}.mp4'.format(base_filename)
    frame_xdir = root / info['subject'] / 'Frames' / '{:}'.format(base_filename)
    video_path, frame_xdir = video_path.resolve(), frame_xdir.resolve()
    frame_xdir.mkdir(parents=True, exist_ok=True)
    #command = "ffmpeg -nostats -loglevel 0 -i {:} -qscale:v 3 {:}/image-%06d.png".format(video_path, frame_xdir)
    command = "ffmpeg -i {:} -q:v 1 -f image2 -start_number 0 {:}/image-%06d.png".format(str(video_path).replace(' ', '\ '), str(frame_xdir).replace(' ', '\ '))
    camera_info = metadata.loadCamera(subjects[info['subject']], camera_ID+1)
    info['pose-{:}'.format(camera)] = Poses_2d
    commands.append( command )
    total += Poses_2d.shape[0]
  return commands, total


def create_frame_script(metadata, root, script):
  print ('load the raw data from {:}'.format(root))
  infos, subject_set, filtered = [], set(), 0
  subjects = {'S{:}'.format(i) : i for i in [1,5,6,7,8,9,11]}
  for sbjK in subjects.keys():
    data = metadata.sequence_mappings[sbjK]
    for action, subaction in data.keys():
      if int(action) <= 1: continue
      if sbjK == 'S11' and action == '2' and subaction == '2': continue
      # check whether this video exist
      does_exist = True
      for camera in metadata.camera_ids:
        base_filename = metadata.get_base_filename(sbjK, action, subaction, camera)
        video_path = root / sbjK / 'Videos' / '{:}.mp4'.format(base_filename)
        if not video_path.exists(): does_exist = False
      if not does_exist:
        filtered += 1
        continue
      infos.append( {'subject'  : sbjK,
                     'action'   : action,
                     'subaction': subaction,
                     'name'     : data[(action, subaction)]}
                  )
      subject_set.add( sbjK )
  print ('There are {:} subject-sub-actions from {:} subjects ({:} are filtered).'.format(len(infos), len(subject_set), filtered))
  subject_set, total = sorted(list(subject_set)), 0
  if script is None: save_file = None
  else             : save_file = open(script, 'w')
  # information
  for xinfo in infos:
    commands, num_frames = process(root, metadata, subjects, xinfo)
    total += num_frames
    for command in commands:
      if save_file: save_file.write('{:}\n'.format(command))
  print ('We save the script into {:} ({:} frames in total).'.format(script, total))
  save_meta_path = root / 'meta-info.pth'
  torch.save(infos, save_meta_path)
  print ('save the meta information into {:}'.format(save_meta_path))


def delete_all(root, script):
  assert root.exists(), 'root : {:} not exist'.format(root)
  meta_path = root / 'meta-info.pth'
  if meta_path.exists(): cmds = ['rm {:}'.format(meta_path.resolve())]
  else                 : cmds = []
  subjects = {'S{:}'.format(i) : i for i in [1,5,6,7,8,9,11]}
  for subject in subjects:
    xpath = root / subject / 'Frames'
    if not xpath.exists(): continue
    for x in xpath.glob('*'):
      if x.is_dir():
        cmds.append('rm -rf {:}'.format( str(x.resolve()).replace(' ', '\ ') ))
  if script is None: save_file = None
  else             : save_file = open(script, 'w')
  for cmd in cmds:
    if save_file: save_file.write('{:}\n'.format(cmd))
 

def human36M_generate(metadata, root, save_dir):
  save_dir = Path(str(save_dir))
  save_dir.mkdir(parents=True, exist_ok=True)
  meta_file_path = root / 'meta-info.pth'
  assert meta_file_path.exists(), 'did not find the meta file : {:}'.format(meta_file_path)
  print('start processing Human 3.6M and will save it into {:}'.format(save_dir))
  meta_infos = torch.load( meta_file_path )
  subjects = {'S{:}'.format(i) : i for i in [1,5,6,7,8,9,11]}

  def create_cameras(meta, sub):
    cameras = {}
    for i, xID in enumerate(metadata.camera_ids):
      camera = {'INDEX': i+1, 'ID': xID, 'SUB': sub}
      info = meta.loadCamera(sub, i+1)
      camera['K'] = np.matrix(info['K'])
      camera['R'] = np.matrix(info['R'])
      camera['t'] = np.array (info['T']).reshape((3,1))
      cameras[xID] = camera
    return cameras

  def get_frame_path(frame, offsets):
    timestamp, ext = frame.name.split('-')[1].split('.')
    timestamp, returnx = int(timestamp), []
    for offset in offsets:
      xname = 'image-{:06d}.{:}'.format(timestamp+offset, ext)
      cpath = frame.parent / xname
      if cpath.exists(): returnx.append( str(cpath) )
      else             : returnx.append( None )
    return returnx

  Datas, all_cameras = [], {}
  for idx, xinfo in enumerate(meta_infos):
    subject, action, subaction = xinfo['subject'], xinfo['action'], xinfo['subaction']
    if subject not in all_cameras:
      all_cameras[subject] = create_cameras(metadata, subjects[subject])
    for camera_key in metadata.camera_ids:
      cameras_index = subject
      # all_cameras[cameras_index][camera_key]
      base_filename = metadata.get_base_filename(subject, action, subaction, camera_key)
      frame_path = root / subject / 'Frames' / '{:}'.format(base_filename)
      assert frame_path.exists(), 'can not find : {:}'.format(frame_path)
      frames = sorted(list(frame_path.glob('*.png')))
      for iF, frame in enumerate(frames):
        try:
          points2D = xinfo['pose-{:}'.format(camera_key)][iF].T
          has_pose = np.ones((1,32))
          points2D = np.concatenate((points2D, has_pose), axis=0)
        except:
          print ('This frame failed : {:}'.format(frame))
          continue
        Past, Now, Next = get_frame_path(frame, [-1,0,1])
        box = (float(points2D[0].min()), float(points2D[1].min()), float(points2D[0].max()), float(points2D[1].max()))
        data = {'cameras_index' : cameras_index,
                'camera_key'    : camera_key,
                'camera_name'   : camera_key,
                'previous_frame': Past,
                'current_frame' : Now,
                'next_frame'    : Next,
                'points'        : points2D,
                'points-pose'   : points2D,
                'box-default'   : box,
                'normalizeL-default': None}
        Datas.append( data )
    print('Process {:03d}/{:03d} videos : {:} in total'.format(idx, len(meta_infos), len(Datas)))  

  def clear_data(xxdatas, is_pts=False, is_vid=False):
    new = []
    for x in xxdatas:
      x = copy.deepcopy(x)
      if is_pts: x['points'] = None
      if is_vid: x['previous_frame'] = None
      if is_vid: x['next_frame'] = None
      new.append( x )
    return new
  # -------- save --------
  base_save_name = root / 'human'
  torch.save({'name': 'Human3.6M-Pose', 'datas': Datas, 'all_cameras': all_cameras}, str(base_save_name) + '.pth')
  torch.save({'name': 'Human3.6M-Pose', 'datas': clear_data(Datas, True, False), 'all_cameras': all_cameras}, str(base_save_name) + '-nopts.pth')
  torch.save({'name': 'Human3.6M-Pose', 'datas': clear_data(Datas, False, True), 'all_cameras': all_cameras}, str(base_save_name) + '-novid.pth')
  torch.save({'name': 'Human3.6M-Pose', 'datas': clear_data(Datas, True, True ), 'all_cameras': all_cameras}, str(base_save_name) + '-nopts-novid.pth')
  
  print ('save into {:}'.format(base_save_name))


if __name__ == '__main__':
  HOME_STR = 'DOME_HOME'
  if HOME_STR not in os.environ: HOME_STR = 'HOME'
  assert HOME_STR in os.environ, 'Doest not find the HOME dir : {}'.format(HOME_STR)
  this_dir = osp.dirname(os.path.abspath(__file__))
  SAVE_DIR = Path(this_dir) / 'lists' / 'Human3.6M'
  print ('This dir : {:}, HOME : [{:}] : {:}'.format(this_dir, HOME_STR, os.environ[HOME_STR]))
  Human3_6M = Path(osp.join(os.environ[HOME_STR], 'datasets', 'landmark-datasets', 'Human-3.6M'))
  assert Human3_6M.exists(), 'Can not find Human 3.6M directory : {:}'.format(Human3_6M)

  metadata = H36M_Metadata('cache/Human3.6M-metadata.xml')
  if sys.argv[1] == 'create':
    assert len(sys.argv) == 3
    script_path = sys.argv[2]
    create_frame_script(metadata, Human3_6M, script_path)
  elif sys.argv[1] == 'genlist':
    human36M_generate(metadata, Human3_6M, SAVE_DIR)
  elif sys.argv[1] == 'clean':
    assert len(sys.argv) == 3
    script_path = sys.argv[2]
    delete_all(Human3_6M, script_path)
  else:
    raise ValueError('Invalid argv : {:}'.format( sys.argv ))
