# -*- coding: utf-8 -*-
#
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# Using this computer program means that you agree to the terms
# in the LICENSE file included with this software distribution.
# Any use not explicitly granted by the LICENSE is prohibited.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# For comments or questions, please email us at deca@tue.mpg.de
# For commercial licensing contact, please contact ps-license@tuebingen.mpg.de

import pickle
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import jittor.nn as nn
import jittor as jt
from .lbs import lbs, batch_rodrigues, vertices2landmarks, rot_mat_to_euler
from jittor import einops

def to_tensor(array, dtype=jt.float32):
    if 'jt.array' not in str(type(array)):
        return jt.array(array, dtype=dtype)


def to_np(array, dtype=np.float32):
    if 'scipy.sparse' in str(type(array)):
        array = array.todense()
    return np.array(array, dtype=dtype)


class Struct(object):
    def __init__(self, **kwargs):
        for key, val in kwargs.items():
            setattr(self, key, val)


__dir__ = Path(__file__).parent.absolute()
FLAMEConfig = SimpleNamespace(
    flame_model_path=str(__dir__ / 'data/FLAME2020/generic_model.pkl'),
    n_shape=100,
    n_exp=50,
    n_tex=50,
    tex_type='BFM',
    tex_path=str(__dir__ / 'data/FLAME2020/FLAME_albedo_from_BFM.npz'),
    flame_lmk_embedding_path=str(__dir__ / 'data/landmark_embedding.npy')
)


class FLAME(nn.Module):
    """
    borrowed from https://github.com/soubhiksanyal/FLAME_PyTorch/blob/master/FLAME.py
    Given flame parameters this class generates a differentiable FLAME function
    which outputs the a mesh and 2D/3D facial landmarks
    """

    def __init__(self, config):
        super(FLAME, self).__init__()
        # print('creating the FLAME Decoder')
        with open(config.flame_model_path, 'rb') as f:
            ss = pickle.load(f, encoding='latin1')
            flame_model = Struct(**ss)

        self.dtype = jt.float32
        self.faces_tensor = jt.array(to_np(flame_model.f, dtype=np.int64)).int64()
        self.v_template = jt.array(to_np(flame_model.v_template), dtype=self.dtype)
        # The shape components and expression
        shapedirs = jt.array(to_np(flame_model.shapedirs), dtype=self.dtype)
        shapedirs = jt.concat([shapedirs[:, :, :config.n_shape], shapedirs[:, :, 300:300 + config.n_exp]], 2)
        self.shapedirs = shapedirs

        # The pose components
        num_pose_basis = flame_model.posedirs.shape[-1]
        posedirs = np.reshape(flame_model.posedirs, [-1, num_pose_basis]).T
        self.posedirs = jt.array(to_np(posedirs), dtype=self.dtype)

        # J_regressor
        self.J_regressor = jt.array(to_np(flame_model.J_regressor), dtype=self.dtype)
        parents = jt.array(to_np(flame_model.kintree_table[0])).int64()
        parents[0] = -1
        self.parents = parents
        self.lbs_weights = jt.array(to_np(flame_model.weights), dtype=self.dtype)

        # Fixing Eyeball and neck rotation
        default_eyeball_pose = jt.zeros([1, 6], dtype=self.dtype)
        self.eye_pose = jt.array(default_eyeball_pose)
        default_eyeball_pose_mat = jt.init.eye(3, dtype=self.dtype).view(1, 9).repeat(1, 2)
        self.eye_pose_mat = jt.array(default_eyeball_pose_mat)
        default_neck_pose = jt.zeros([1, 3], dtype=self.dtype)
        self.neck_pose = jt.array(default_neck_pose)
        default_neck_pose_mat = jt.init.eye(3, dtype=self.dtype).view(1, 9)
        self.neck_pose_mat = jt.array(default_neck_pose_mat)

        # Static and Dynamic Landmark embeddings for FLAME
        lmk_embeddings = np.load(config.flame_lmk_embedding_path, allow_pickle=True, encoding='latin1')
        lmk_embeddings = lmk_embeddings[()]
        self.lmk_faces_idx = jt.array(lmk_embeddings['static_lmk_faces_idx']).int64()
        self.lmk_bary_coords = jt.array(lmk_embeddings['static_lmk_bary_coords'], dtype=self.dtype)
        self.dynamic_lmk_faces_idx = jt.array(lmk_embeddings['dynamic_lmk_faces_idx']).int64()
        self.dynamic_lmk_bary_coords = jt.array(lmk_embeddings['dynamic_lmk_bary_coords'], dtype=self.dtype)
        self.full_lmk_faces_idx = jt.array(lmk_embeddings['full_lmk_faces_idx']).int64()
        self.full_lmk_bary_coords = jt.array(lmk_embeddings['full_lmk_bary_coords'], dtype=self.dtype)

        # Neck kinematic chain
        neck_kin_chain = []
        NECK_IDX = 1
        curr_idx = jt.array([NECK_IDX], dtype=jt.int64)
        while curr_idx[0] != -1:
            neck_kin_chain.append(curr_idx)
            curr_idx = self.parents[curr_idx[0]]
        self.neck_kin_chain = jt.stack(neck_kin_chain)

       
        
    def _find_dynamic_lmk_idx_and_bcoords(self, pose, dynamic_lmk_faces_idx,
                                          dynamic_lmk_b_coords,
                                          neck_kin_chain, dtype=jt.float32, pose2rot=True):
        """
            Selects the face contour depending on the reletive position of the head
            Input:
                vertices: N X num_of_vertices X 3
                pose: N X full pose
                dynamic_lmk_faces_idx: The list of contour face indexes
                dynamic_lmk_b_coords: The list of contour barycentric weights
                neck_kin_chain: The tree to consider for the relative rotation
                dtype: Data type
            return:
                The contour face indexes and the corresponding barycentric weights
        """

        batch_size = pose.shape[0]

        if pose2rot:
            aa_pose = jt.index_select(pose.view(batch_size, -1, 3), 1,
                                         neck_kin_chain)
            rot_mats = batch_rodrigues(
                aa_pose.view(-1, 3), dtype=dtype).view(batch_size, -1, 3, 3)
        else:
            rot_mats = jt.index_select(pose.view(batch_size, -1, 9), 1,
                                          neck_kin_chain).view(batch_size, -1, 3, 3)
        rel_rot_mat = jt.init.eye(3, dtype=dtype).unsqueeze(0).expand(batch_size, -1, -1)
        for idx in range(len(neck_kin_chain)):
            rel_rot_mat = jt.bmm(rot_mats[:, idx], rel_rot_mat)

        y_rot_angle = jt.round(
            jt.clamp(rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
                        max=39)).to(dtype=jt.int64)

        neg_mask = y_rot_angle.less(0).int64()
        mask = y_rot_angle.less(-39).int64()
        
        neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
        y_rot_angle = (neg_mask * neg_vals +
                       (1 - neg_mask) * y_rot_angle)

        dyn_lmk_faces_idx = jt.index_select(dynamic_lmk_faces_idx,
                                               0, y_rot_angle)
        dyn_lmk_b_coords = jt.index_select(dynamic_lmk_b_coords,
                                              0, y_rot_angle)
        return dyn_lmk_faces_idx, dyn_lmk_b_coords
  

    def seletec_3d68(self, vertices):
        repeated_faces_idx = einops.repeat(self.full_lmk_faces_idx, 'f -> b f', b=vertices.shape[0])
        repeated_bary_coords = einops.repeat(self.full_lmk_bary_coords, 'f c -> b f c', b=vertices.shape[0])

        landmarks3d = vertices2landmarks(vertices, self.faces_tensor, repeated_faces_idx, repeated_bary_coords)

        return landmarks3d



    def execute(self, shape_params=None, expression_params=None, pose_params=None, eye_pose_params=None, pose2rot=True,
                ignore_global_rot=False, return_lm2d=True, return_lm3d=True):
        """
            Input:
                shape_params: N X number of shape parameters
                expression_params: N X number of expression parameters
                pose_params: N X number of pose parameters (6)
            return:d
                vertices: N X V X 3
                landmarks: N X number of landmarks X 3
        """
        batch_size = shape_params.shape[0]
        betas = jt.concat([shape_params, expression_params], dim=1)

        if pose2rot:
            # The pose parameters:
            #   global rotation (3), head rotation around the neck (3), jaw rotation (3), each of the eyeballs (3+3)
            if pose_params is None:
                pose_params = self.eye_pose.expand(batch_size, -1)
            if eye_pose_params is None:
                eye_pose_params = self.eye_pose.expand(batch_size, -1)
            head_pose = pose_params[:, :3] if not ignore_global_rot else jt.zeros_like(pose_params[:, :3])
            full_pose = jt.concat(
                [head_pose, self.neck_pose.expand(batch_size, -1), pose_params[:, 3:], eye_pose_params], dim=1)
        else:
            if pose_params is None:
                pose_params = self.eye_pose_mat.expand(batch_size, -1)
            if eye_pose_params is None:
                eye_pose_params = self.eye_pose_mat.expand(batch_size, -1)
            head_pose = pose_params[:, :9] if not ignore_global_rot else self.eye_pose_mat.expand(batch_size, -1)[:, :9]
            full_pose = jt.concat(
                [head_pose, self.neck_pose_mat.expand(batch_size, -1), pose_params[:, 9:], eye_pose_params], dim=1)
        template_vertices = self.v_template.unsqueeze(0).expand(batch_size, -1, -1)

        vertices, _ = lbs(betas, full_pose, template_vertices,
                          self.shapedirs, self.posedirs,
                          self.J_regressor, self.parents,
                          self.lbs_weights, pose2rot, self.dtype)

        if return_lm2d:
            lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(batch_size, -1)
            lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).expand(batch_size, -1, -1)

            dyn_lmk_faces_idx, dyn_lmk_bary_coords = self._find_dynamic_lmk_idx_and_bcoords(
                full_pose, self.dynamic_lmk_faces_idx,
                self.dynamic_lmk_bary_coords,
                self.neck_kin_chain, dtype=self.dtype, pose2rot=pose2rot)
            lmk_faces_idx = jt.concat([dyn_lmk_faces_idx, lmk_faces_idx], 1)
            lmk_bary_coords = jt.concat([dyn_lmk_bary_coords, lmk_bary_coords], 1)

            landmarks2d = vertices2landmarks(vertices, self.faces_tensor,
                                             lmk_faces_idx,
                                             lmk_bary_coords)
        else:
            landmarks2d = None

        if return_lm3d:
            bz = vertices.shape[0]
            landmarks3d = vertices2landmarks(vertices, self.faces_tensor,
                                             self.full_lmk_faces_idx.repeat(bz, 1),
                                             self.full_lmk_bary_coords.repeat(bz, 1, 1))
        else:
            landmarks3d = None

        return vertices, landmarks2d, landmarks3d


class FLAMETex(nn.Module):
    """
    FLAME texture:
    https://github.com/TimoBolkart/TF_FLAME/blob/ade0ab152300ec5f0e8555d6765411555c5ed43d/sample_texture.py#L64
    FLAME texture converted from BFM:
    https://github.com/TimoBolkart/BFM_to_FLAME
    """

    def __init__(self, config):
        super(FLAMETex, self).__init__()
        if config.tex_type == 'BFM':
            mu_key = 'MU'
            pc_key = 'PC'
            n_pc = 199
            tex_path = config.tex_path
            tex_space = np.load(tex_path)
            texture_mean = tex_space[mu_key].reshape(1, -1)
            texture_basis = tex_space[pc_key].reshape(-1, n_pc)

        elif config.tex_type == 'FLAME':
            mu_key = 'mean'
            pc_key = 'tex_dir'
            n_pc = 200
            tex_path = config.flame_tex_path
            tex_space = np.load(tex_path)
            texture_mean = tex_space[mu_key].reshape(1, -1) / 255.
            texture_basis = tex_space[pc_key].reshape(-1, n_pc) / 255.
        else:
            print('texture type ', config.tex_type, 'not exist!')
            raise NotImplementedError

        n_tex = config.n_tex
        num_components = texture_basis.shape[1]
        texture_mean = jt.array(texture_mean).float()[None, ...]
        texture_basis = jt.array(texture_basis[:, :n_tex]).float()[None, ...]
        self.register_buffer('texture_mean', texture_mean)
        self.register_buffer('texture_basis', texture_basis)

    def execute(self, texcode):
        '''
        texcode: [batchsize, n_tex]
        texture: [bz, 3, 256, 256], range: 0-1
        '''

        bs = texcode.shape[0]
        texcode = texcode[:1]

        # we use the same (first frame) texture for all frames

        texture = self.texture_mean + (self.texture_basis * texcode[:, None, :]).sum(-1)

        texture = texture.reshape(texcode.shape[0], 512, 512, 3).permute(0, 3, 1, 2)
        texture = jt.nn.interpolate(texture, [256, 256])
        texture = texture[:, [2, 1, 0], :, :].repeat(bs, 1, 1, 1)
        return texture
