#!/usr/bin/env python3
# coding=utf-8

import _init_paths
import os 
import os.path as osp
import cv2
import numpy as np
import torch
import argparse 
import pickle as pkl 
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
plt.switch_backend('agg')
plt.ioff()

from tqdm import tqdm
from torch.utils.data import DataLoader
from lib.models.model import get_pose_model, get_discriminator
from lib.core.config import  config, update_config, update_dir

pairs = [(0, 1), (1, 2), (2, 13), (3, 13), (3, 4), (4, 5), (6, 7), (7, 8), (8, 12), (9, 10),(9, 12), (10, 11),(12, 14), (12, 15), (13, 14), (15, 16)]
pairs_left = [(3, 13), (3, 4), (4, 5), (9, 10), (9, 12), (10, 11)]
pairs_right = [(0, 1), (1, 2), (2, 13), (6, 7), (7, 8), (8, 12)]

colors = {
    'pink': np.array([197, 27, 125]),  # L lower leg
    'light_pink': np.array([233, 163, 201]),  # L upper leg
    'light_green': np.array([161, 215, 106]),  # L lower arm
    'green': np.array([77, 146, 33]),  # L upper arm
    'red': np.array([215, 48, 39]),  # head
    'light_red': np.array([252, 146, 114]),  # head
    'light_orange': np.array([252, 141, 89]),  # chest
    'purple': np.array([118, 42, 131]),  # R lower leg
    'light_purple': np.array([175, 141, 195]),  # R upper
    'light_blue': np.array([145, 191, 219]),  # R lower arm
    'blue': np.array([69, 117, 180]),  # R upper arm
    'gray': np.array([130, 130, 130]),  #
    'white': np.array([255, 255, 255]),  #
}

parser = argparse.ArgumentParser()
parser.add_argument('--indir', default='./demo_input')
parser.add_argument('--outdir', default='./demo_output')
parser.add_argument('--cfg', default='../cfg/h36m_gt/tmc_klbone.yaml')
parser.add_argument('--pretrain', default='../models/tmc_klbone.pth.tar')
# parser.add_argument('--cfg', default='../cfg/best_samenorm.yaml')
args = parser.parse_args()

indir = args.indir 
outdir = args.outdir
os.makedirs(outdir, exist_ok=True)

update_config(args.cfg)

pose_model = get_pose_model(config)
print(pose_model)
assert osp.exists(args.pretrain), "Can not find pretrained model at {}".format(args.pretrain)
state_dict = torch.load(args.pretrain)['pose_model_state_dict']
pose_model.load_state_dict(state_dict)
pose_model.eval()
torch.set_grad_enabled(False)

img_files = [osp.join(indir, _file) for _file in os.listdir(indir) if _file.endswith(".jpg") and osp.exists(osp.join(indir, _file.replace(".jpg", ".pkl")))]
kps_files = [_file.replace(".jpg", ".pkl") for _file in img_files]

all_joints_2d = []
all_joints_3d_pre = []

for kps_file in kps_files: 
    with open(kps_file, "rb") as f: 
        data = pkl.load(f)
    joints_2d = torch.from_numpy(data['joints_2d'][np.newaxis]).float()
    original_joints_2d = torch.from_numpy(data['original_joints_2d'][np.newaxis]).float()
    joints_3d, _, _ = pose_model(joints_2d, is_train=False)
    all_joints_2d.append(original_joints_2d.numpy())
    all_joints_3d_pre.append(joints_3d.numpy())

all_joints_2d = np.concatenate(all_joints_2d, axis=0)
all_joints_3d_pre = np.concatenate(all_joints_3d_pre, axis=0)
# print(all_joints_3d_pre.shape)


for idx, joints_3d_pre in tqdm(enumerate(all_joints_3d_pre), total=len(all_joints_3d_pre)):
    joints_2d = all_joints_2d[idx]
    joints_3d_pre = joints_3d_pre - joints_3d_pre[13:14]
    image_path = img_files[idx]
    save_path = osp.join(outdir, osp.basename(image_path).replace(".jpg", "_out.jpg"))
    image = cv2.imread(image_path)

    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection='3d')
    ax.view_init(elev=10., azim=45.)

    for pair in pairs: 
        i, j = pair
        if pair in pairs_left: 
            color = "blue"
            cv_color = (255, 0, 0)
        elif pair in pairs_right: 
            color = "green"
            cv_color = (0, 255, 0)
        else: 
            color = "darkorange"
            cv_color = (89, 141, 252)
        x1, y1 = joints_2d[i].astype(np.int)
        x2, y2 = joints_2d[j].astype(np.int)
        
        cv2.line(image, (x1, y1), (x2, y2), cv_color, 2)
        x1, y1, z1 = joints_3d_pre[i]
        x2, y2, z2 = joints_3d_pre[j]
        ax.plot([z1, z2], [x1, x2], [-y1, -y2], c=color, linewidth=3)

    image = image[::-1, :, ::-1].copy().astype(np.float) / 255.
    r = 0.95
    xroot = yroot = zroot = 0.
    # radius = max(4, (np.mean(image.shape[:2]) * 0.01).astype(int))
    radius = 0.75
    xx = np.linspace(-r * radius + xroot, r * radius + xroot, image.shape[1])
    yy = np.linspace(-r * radius + yroot, r * radius + yroot, image.shape[0])
    xx, yy = np.meshgrid(xx, yy)
    zz = np.ones_like(xx) * (-3.2* radius + zroot)
    ax.plot_surface(zz, xx, yy, rstride=1, cstride=1, facecolors=image, shade=False)
    ax.set_xlabel('Z', fontsize=13)
    ax.set_ylabel("X", fontsize=13)
    ax.set_zlabel("Y", fontsize=13)
    ax.set_ylim3d([-radius+xroot, radius+xroot])
    ax.set_zlim3d([-radius+yroot, radius+yroot])
    ax.set_xlim3d([-2.5 * radius+zroot, radius+zroot])
    ax.get_xaxis().set_ticklabels([])
    ax.get_yaxis().set_ticklabels([])
    ax.set_zticklabels([])

    white = (1.0, 1.0, 1.0, 0.0)
    ax.w_xaxis.set_pane_color(white)
    ax.w_yaxis.set_pane_color(white)

    ax.w_xaxis.line.set_color(white)
    ax.w_yaxis.line.set_color(white)
    ax.w_zaxis.line.set_color(white)
        
    plt.savefig(save_path)
    plt.close()
