# ref: https://github.com/ShunyuYao/DFA-NeRF
from numpy.core.numeric import require
from numpy.lib.function_base import quantile
import torch
import numpy as np
from facemodel import Face_3DMM
from data_loader import load_dir
from util import *
import os
import sys
import cv2
import imageio
import argparse

dir_path = os.path.dirname(os.path.realpath(__file__))


def set_requires_grad(tensor_list):
    for tensor in tensor_list:
        tensor.requires_grad = True


parser = argparse.ArgumentParser()
parser.add_argument(
    "--path", type=str, default="obama/ori_imgs", help="idname of target person")
parser.add_argument('--img_h', type=int, default=512, help='image height')
parser.add_argument('--img_w', type=int, default=512, help='image width')
parser.add_argument('--frame_num', type=int,
                    default=11000, help='image number')
args = parser.parse_args()
start_id = 0
end_id = args.frame_num

lms = load_dir(args.path, start_id, end_id)
num_frames = lms.shape[0]
h, w = args.img_h, args.img_w
cxy = torch.tensor((w/2.0, h/2.0), dtype=torch.float).cuda()
id_dim, exp_dim, tex_dim, point_num = 100, 79, 100, 34650
model_3dmm = Face_3DMM(os.path.join(dir_path, '3DMM'),
                       id_dim, exp_dim, tex_dim, point_num)
lands_info = np.loadtxt(os.path.join(
    dir_path, '3DMM', 'lands_info.txt'), dtype=np.int32)
lands_info = torch.as_tensor(lands_info).cuda()
# mesh = openmesh.read_trimesh(os.path.join(dir_path, '3DMM', 'template.obj'))
focal = 1150

id_para = lms.new_zeros((1, id_dim), requires_grad=True)
exp_para = lms.new_zeros((num_frames, exp_dim), requires_grad=True)
tex_para = lms.new_zeros((1, tex_dim), requires_grad=True)
euler_angle = lms.new_zeros((num_frames, 3), requires_grad=True)
trans = lms.new_zeros((num_frames, 3), requires_grad=True)
light_para = lms.new_zeros((num_frames, 27), requires_grad=True)
trans.data[:, 2] -= 600
focal_length = lms.new_zeros(1, requires_grad=True)
focal_length.data += focal

set_requires_grad([id_para, exp_para, tex_para,
                   euler_angle, trans, light_para])

sel_ids = np.arange(0, num_frames, 10)
sel_num = sel_ids.shape[0]
arg_focal = 0.0
arg_landis = 1e5
for focal in range(500, 1500, 50):
    id_para = lms.new_zeros((1, id_dim), requires_grad=True)
    exp_para = lms.new_zeros((sel_num, exp_dim), requires_grad=True)
    euler_angle = lms.new_zeros((sel_num, 3), requires_grad=True)
    trans = lms.new_zeros((sel_num, 3), requires_grad=True)
    trans.data[:, 2] -= 600
    focal_length = lms.new_zeros(1, requires_grad=False)
    focal_length.data += focal
    set_requires_grad([id_para, exp_para, euler_angle, trans])

    optimizer_id = torch.optim.Adam([id_para], lr=.3)
    optimizer_exp = torch.optim.Adam([exp_para], lr=.3)
    optimizer_frame = torch.optim.Adam(
        [euler_angle, trans], lr=.3)
    iter_num = 2000

    for iter in range(iter_num):
        id_para_batch = id_para.expand(sel_num, -1)
        geometry = model_3dmm.forward_geo_sub(
            id_para_batch, exp_para, lands_info[-51:].long())
        proj_geo = forward_transform(
            geometry, euler_angle, trans, focal_length, cxy)
        loss_lan = cal_lan_loss(
            proj_geo[:, :, :2], lms[sel_ids, -51:, :].detach())
        loss_regid = torch.mean(id_para*id_para)*8
        loss_regexp = torch.mean(exp_para*exp_para)*0.5
        loss = loss_lan + loss_regid + loss_regexp
        optimizer_id.zero_grad()
        optimizer_exp.zero_grad()
        optimizer_frame.zero_grad()
        loss.backward()
        if iter > 1000:
            optimizer_id.step()
            optimizer_exp.step()
        optimizer_frame.step()
    print(focal, loss_lan.item(), torch.mean(trans[:, 2]).item())
    if loss_lan.item() < arg_landis:
        arg_landis = loss_lan.item()
        arg_focal = focal

sel_ids = np.arange(0, num_frames)
sel_num = sel_ids.shape[0]
id_para = lms.new_zeros((1, id_dim), requires_grad=True)
exp_para = lms.new_zeros((sel_num, exp_dim), requires_grad=True)
euler_angle = lms.new_zeros((sel_num, 3), requires_grad=True)
trans = lms.new_zeros((sel_num, 3), requires_grad=True)
trans.data[:, 2] -= 600
focal_length = lms.new_zeros(1, requires_grad=False)
focal_length.data += arg_focal
set_requires_grad([id_para, exp_para, euler_angle, trans])

optimizer_id = torch.optim.Adam([id_para], lr=.3)
optimizer_exp = torch.optim.Adam([exp_para], lr=.3)
optimizer_frame = torch.optim.Adam(
    [euler_angle, trans], lr=.3)
iter_num = 2000

for iter in range(iter_num):
    id_para_batch = id_para.expand(sel_num, -1)
    geometry = model_3dmm.forward_geo_sub(
        id_para_batch, exp_para, lands_info[-51:].long())
    proj_geo = forward_transform(
        geometry, euler_angle, trans, focal_length, cxy)
    loss_lan = cal_lan_loss(
        proj_geo[:, :, :2], lms[sel_ids, -51:, :].detach())
    loss_regid = torch.mean(id_para*id_para)*8
    loss_regexp = torch.mean(exp_para*exp_para)*0.5
    loss = loss_lan + loss_regid + loss_regexp
    optimizer_id.zero_grad()
    optimizer_exp.zero_grad()
    optimizer_frame.zero_grad()
    loss.backward()
    if iter > 1000:
        optimizer_id.step()
        optimizer_exp.step()
    optimizer_frame.step()
print(arg_focal, loss_lan.item(), torch.mean(trans[:, 2]).item())


torch.save({'id': id_para.detach().cpu(), 'exp': exp_para.detach().cpu(),
            'euler': euler_angle.detach().cpu(), 'trans': trans.detach().cpu(),
            'focal': focal_length.detach().cpu()}, os.path.join(os.path.dirname(args.path), 'track_params.pt'))
print('face tracking params saved')