# -*- 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 pixie@tue.mpg.de
# For commercial licensing contact, please contact ps-license@tuebingen.mpg.de

import pickle

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


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.0
            texture_basis = tex_space[pc_key].reshape(-1, n_pc) / 255.0
        else:
            print("texture type ", config.tex_type, "not exist!")
            raise NotImplementedError

        n_tex = config.n_tex
        num_components = texture_basis.shape[1]
        texture_mean = torch.from_numpy(texture_mean).float()[None, ...]
        texture_basis = torch.from_numpy(texture_basis[:, :n_tex]).float()[None, ...]
        self.register_buffer("texture_mean", texture_mean)
        self.register_buffer("texture_basis", texture_basis)

    def forward(self, texcode=None):
        """
        texcode: [batchsize, n_tex]
        texture: [bz, 3, 256, 256], range: 0-1
        """
        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 = F.interpolate(texture, [256, 256])
        texture = texture[:, [2, 1, 0], :, :]
        return texture


def texture_flame2smplx(cached_data, flame_texture, smplx_texture):
    """Convert flame texture map (face-only) into smplx texture map (includes body texture)
    TODO: pytorch version ==> grid sample
    """
    if smplx_texture.shape[0] != smplx_texture.shape[1]:
        print("SMPL-X texture not squared (%d != %d)" % (smplx_texture[0], smplx_texture[1]))
        return
    if smplx_texture.shape[0] != cached_data["target_resolution"]:
        print(
            "SMPL-X texture size does not match cached image resolution (%d != %d)" %
            (smplx_texture.shape[0], cached_data["target_resolution"])
        )
        return
    x_coords = cached_data["x_coords"]
    y_coords = cached_data["y_coords"]
    target_pixel_ids = cached_data["target_pixel_ids"]
    source_uv_points = cached_data["source_uv_points"]

    source_tex_coords = np.zeros_like((source_uv_points)).astype(int)
    source_tex_coords[:, 0] = np.clip(
        flame_texture.shape[0] * (1.0 - source_uv_points[:, 1]),
        0.0,
        flame_texture.shape[0],
    ).astype(int)
    source_tex_coords[:, 1] = np.clip(
        flame_texture.shape[1] * (source_uv_points[:, 0]), 0.0, flame_texture.shape[1]
    ).astype(int)

    smplx_texture[y_coords[target_pixel_ids].astype(int),
                  x_coords[target_pixel_ids].astype(int), :, ] = flame_texture[source_tex_coords[:,
                                                                                                 0],
                                                                               source_tex_coords[:,
                                                                                                 1]]

    return smplx_texture