# Script for optimizing the ratio parameters

import torch
import torch.nn as nn
import os
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import pickle
import numpy as np
import time

POSE_MODEL = 'OpenPose'  # ViTPose_large, ViTPose_base, OpenPose

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


class Planar_Euclidean_Loss(nn.Module):
    '''
    Euclidean distance of x,y coordinates between 3D points, ignore z coordinate
    '''

    def __init__(self):
        super(Planar_Euclidean_Loss, self).__init__()

    def forward(self, pred, target):
        loss = torch.nn.functional.mse_loss(pred[0:2,:], target[0:2,:])  # only consider xy-axis
        return loss


class Target4_Model(nn.Module):
    '''
    target computation model
    '''

    def __init__(self, r1, r2):
        super(Target4_Model, self).__init__()
        self.r1 = nn.Parameter(torch.tensor(r1), requires_grad=True)
        self.r2 = nn.Parameter(torch.tensor(r2), requires_grad=True)

    def forward(self, X1, X2, t2):
        X3 = X1 + self.r1 * (X2 - X1)
        pred = X3 + self.r2 * torch.norm(X3 - X1) * t2
        return pred


class PositionDataset(Dataset):
    def __init__(self, X1, X2, t2, target):
        self.X1 = X1
        self.X2 = X2
        self.t2 = t2
        self.target = target

    def __len__(self):
        return len(self.X1)

    def __getitem__(self, i):
        return self.X1[i], self.X2[i], self.t2[i], self.target[i]


def optimize_side(data, target, params=[0.35, 0.1], epoch=1000, lr=0.01, use_gpu=False):
    if use_gpu and torch.cuda.is_available():
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')

    X1, X2, t2 = data
    dataset = PositionDataset(X1, X2, t2, target)
    data_loader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=0)

    model = Target4_Model(*params)
    model.to(device)
    print('model.state_dict():', model.state_dict())
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    loss_fn = Planar_Euclidean_Loss()

    for ep in range(epoch):
        total_loss = 0
        for i, data in enumerate(data_loader):
            X1, X2, t2, target = data

            X1 = X1.to(device)
            X2 = X2.to(device)
            t2 = t2.to(device)
            target = target.to(device)

            opt.zero_grad()
            pred = model.forward(X1, X2, t2)
            loss = loss_fn(pred, target)
            total_loss += loss.item()
            loss.backward()
            opt.step()

        if ep % 100 == 0:
            print('epoch: {}, loss: {}, r1: {}, r2: {}'.format(ep, total_loss, model.state_dict()['r1'], model.state_dict()['r2']))


def optimize_front_linear(data, target):
    '''
    Optimize the two ratios using least square.
    Only use x&y values.
     :X1, X2 --3D coordinates of the right shoulder and the right hip
     :t2 -- the direction vector of the line connecting X3 and target
    '''
    X1, X2, t2 = data[0], data[1], data[2]

    # Onlu use x&y values
    X1_xy = np.hstack(X1)[0:2,:].T.reshape(-1, 1)   # (2nx1)
    X2_xy = np.hstack(X2)[0:2,:].T.reshape(-1, 1)

    r1_coeff = X2_xy - X1_xy    # 2nx1
    r2_coeff = (np.hstack(t2)[0:2,:].T * np.sqrt(((np.hstack(X1)-np.hstack(X2))**2).sum(axis=0)).reshape(-1,1)).reshape(-1,1)  # 2nx1
    A = np.hstack([r1_coeff, r2_coeff])
    b = np.hstack(target)[0:2,:].T.reshape(-1,1) - X1_xy   # reshape (nx2) to (2nx1)
    r = np.linalg.pinv(A) @ b

    return r


if __name__ == '__main__':

    print("HPE model: ", POSE_MODEL)
    # collect data
    target12_data = [[], [], []]   # list of list of np.array: [X1_list, X2_list, t2_list]
    target1_GT = []               # list of np.array
    target2_GT = []

    target4_data = [[], [], []]
    target4_GT = []

    for SUBJECT_NAME in os.listdir('data'):
        subject_folder_path = os.path.join('data', SUBJECT_NAME)
        if os.path.isfile(subject_folder_path):
            continue

        scan_pose = 'front'
        with open(subject_folder_path + '/' + scan_pose + '/' + POSE_MODEL + '/position_data.pickle', 'rb') as f:
            position_data = pickle.load(f)

        target12_data[0].append(position_data[scan_pose][0])  # X1
        target12_data[1].append(position_data[scan_pose][1])  # X2
        target12_data[2].append(position_data[scan_pose][2])  # t2

        with open(subject_folder_path + '/' + scan_pose + '/two_cam_gt.pickle', 'rb') as f:
            ground_truth = pickle.load(f)
        target1_GT.append(ground_truth['target_1'])
        target2_GT.append(ground_truth['target_2'])

        # skip outlier for openpose target 4:
        if POSE_MODEL == 'OpenPose' and (SUBJECT_NAME == 'charles_xu' or SUBJECT_NAME == 'jingyu_wu'):
            continue

        scan_pose = 'side'
        with open(subject_folder_path + '/' + scan_pose + '/' + POSE_MODEL + '/position_data.pickle', 'rb') as f:
            position_data = pickle.load(f)

        target4_data[0].append(position_data[scan_pose][0])  # X1
        target4_data[1].append(position_data[scan_pose][1])  # X2
        target4_data[2].append(position_data[scan_pose][2])  # t2

        with open(subject_folder_path + '/' + scan_pose + '/two_cam_gt.pickle', 'rb') as f:
            ground_truth = pickle.load(f)
        target4_GT.append(ground_truth['target_4'])

    start_time = time.time()
    optimize_side(target4_data, target4_GT)
    print('training used {:.3f} s'.format(time.time() - start_time))

    target1_ratio = optimize_front_linear(target12_data, target1_GT)
    target2_ratio = optimize_front_linear(target12_data, target2_GT)
    print("target1_ratio: \n", target1_ratio)
    print("target2_ratio: \n", target2_ratio)