import torch
import torch.nn as nn
import numpy as np
import pickle
from scipy.io import loadmat
from pytorch3d.transforms import so3_exponential_map


class FVMModule(nn.Module):
    def __init__(self, batch_size):
        super(FVMModule, self).__init__()

        self.id_dims = 150
        self.exp_dims = 52

        # a = loadmat('assets/BFM/BFM09_model_info.mat')
        model_dict = np.load('assets/FVM/faceverse_simple_v2.npy', allow_pickle=True).item()
        self.register_buffer('skinmask', torch.tensor(model_dict['skinmask']))
        kp_inds = torch.tensor(model_dict['keypoints']).squeeze().long()
        #kp_inds = torch.cat([kp_inds[0:48], kp_inds[49:54], kp_inds[55:68]])
        self.register_buffer('kp_inds', kp_inds)

        meanshape = torch.tensor(model_dict['meanshape'])
        meanshape[:, 1:] = -meanshape[:, 1:]
        self.register_buffer('meanshape', meanshape.view(1, -1).float())

        idBase = torch.tensor(model_dict['idBase']).view(-1, 3, self.id_dims).float()
        idBase[:, 1:, :] = -idBase[:, 1:, :]
        self.register_buffer('idBase', idBase.view(-1, self.id_dims))

        exBase = torch.tensor(model_dict['exBase']).view(-1, 3, self.exp_dims).float()
        exBase[:, 1:, :] = -exBase[:, 1:, :]
        self.register_buffer('exBase', exBase.view(-1, self.exp_dims))

        self.register_buffer('faces', torch.tensor(model_dict['tri']).long())

        self.batch_size = batch_size
        self.id_coeff = nn.Parameter(torch.zeros(1, self.id_dims).float())
        self.exp_coeff = nn.Parameter(torch.zeros(self.batch_size, self.exp_dims).float())
        self.scale = nn.Parameter(torch.ones(1).float() * 0.3)
        self.pose = nn.Parameter(torch.zeros(self.batch_size, 6).float())

    def set_id_param(self, id_coeff, scale):
        self.id_coeff.data = id_coeff
        self.scale.data = scale
        self.id_coeff.requires_grad = False
        self.scale.requires_grad = False

    def get_lms(self, vs):
        lms = vs[:, self.kp_inds, :]
        return lms

    def get_vs(self, id_coeff, exp_coeff):
        n_b = id_coeff.size(0)

        face_shape = torch.einsum('ij,aj->ai', self.idBase, id_coeff) + \
            torch.einsum('ij,aj->ai', self.exBase, exp_coeff) + self.meanshape

        face_shape = face_shape.view(n_b, -1, 3)
        face_shape = face_shape - \
            self.meanshape.view(1, -1, 3).mean(dim=1, keepdim=True)

        return face_shape

    def forward(self):
        id_coeff = self.id_coeff.repeat(self.batch_size, 1)
        vertices = self.get_vs(id_coeff, self.exp_coeff)
        R = so3_exponential_map(self.pose[:, :3])
        T = self.pose[:, 3:]
        vertices = torch.bmm(vertices * self.scale, R.permute(0,2,1)) + T[:, None, :]
        landmarks = self.get_lms(vertices)
        return vertices, landmarks

    def reg_loss(self, id_weight, exp_weight):
        id_reg_loss = (self.id_coeff ** 2).sum()
        exp_reg_loss = (self.exp_coeff ** 2).sum(-1).mean()
        return id_reg_loss * id_weight + exp_reg_loss * exp_weight

    def temporal_smooth_loss(self, smo_weight):
        return ((self.exp_coeff[1:] - self.exp_coeff[:-1]) ** 2).sum(-1).mean() * smo_weight

    def save(self, path, batch_id=-1):
        if batch_id < 0:
            id_coeff = self.id_coeff.detach().cpu().numpy()
            exp_coeff = self.exp_coeff.detach().cpu().numpy()
            scale = self.scale.detach().cpu().numpy()
            pose = self.pose.detach().cpu().numpy()
            np.savez(path, id_coeff=id_coeff, exp_coeff=exp_coeff, scale=scale, pose=pose)
        else:
            id_coeff = self.id_coeff.detach().cpu().numpy()
            exp_coeff = self.exp_coeff[batch_id:batch_id+1].detach().cpu().numpy()
            scale = self.scale.detach().cpu().numpy()
            pose = self.pose[batch_id:batch_id+1].detach().cpu().numpy()
            np.savez(path, id_coeff=id_coeff, exp_coeff=exp_coeff, scale=scale, pose=pose)