import os
import sys

import pickle
import scipy.io as sio
from math import *

import numpy as np
import tensorflow as tf

# Left-Right
eval_joints = {
        'All': [1, 2, 3, 4, 5, 7, 8, 9, 12, 15, 16, 17, 18, 19, 20, 21],
        'Head': [12, 15], # Neck, Head
        'Torso': [3, 6, 9], # Back joints
        'Shoulder': [13, 14, 16, 17], # Inner, Outer 
        'Elbow': [18, 19],
        'Wrist': [20, 21],
        'Hip': [1, 2],
        'Knee': [4, 5],
        'Foot': [7, 8, 10, 11]} # Ankle, Foot
        
class SMPL:

    def __init__(self, path):
        self.J_regressor = []
        self.v_template = []
        self.shapedirs = []
        self.posedirs = []
        self.weights = []
        # Load SMPL models
        # path = os.path.abspath(os.path.dirname(__file__))
        with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
            if sys.version_info[0]<3:
                female = pickle.load(f)
            else:
                female = pickle.load(f, encoding="latin1")
            self.J_regressor.append(tf.constant(np.array(female['J_regressor'].todense(), dtype=np.float32)))
            self.v_template.append(tf.constant(female['v_template'], dtype=np.float32))
            self.posedirs.append(tf.constant(female['posedirs'], dtype=np.float32))
            self.shapedirs.append(tf.constant(female['shapedirs'], dtype=np.float32))
            self.weights.append(tf.constant(female['weights'], dtype=np.float32))
        with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
            if sys.version_info[0]<3:
                male = pickle.load(f)
            else:
                male = pickle.load(f, encoding="latin1")
            self.J_regressor.append(tf.constant(np.array(male['J_regressor'].todense(), dtype=np.float32)))
            self.v_template.append(tf.constant(male['v_template'], dtype=np.float32))
            self.posedirs.append(tf.constant(male['posedirs'], dtype=np.float32))
            self.shapedirs.append(tf.constant(male['shapedirs'], dtype=np.float32))
            self.weights.append(tf.constant(male['weights'], dtype=np.float32))
            kintree_table = male['kintree_table']
            id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
            self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
        self.J_regressor = tf.stack(self.J_regressor, axis=0)
        self.v_template = tf.stack(self.v_template, axis=0)
        self.shapedirs = tf.stack(self.shapedirs, axis=0)
        self.posedirs = tf.stack(self.posedirs, axis=0)
        self.weights = tf.stack(self.weights, axis=0)
        self.weights_one_hot = tf.reduce_mean(self.weights, axis=0)
        self.weights_one_hot = tf.math.argmax(self.weights_one_hot, axis=-1)
        self.weights_one_hot = tf.one_hot(self.weights_one_hot, 24)

        # Topology no head/hands/feet
        F = sio.loadmat(path + '/ssmpl_faces.mat')
        self.f = np.array(F['faces_default'])
        self.f_skirt = np.array(F['faces_skirt'])
        self.nn = np.array(sio.loadmat(path + '/NN.mat'))
        self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)

        self.weights = tf.gather(self.weights, self.to_keep, axis=1)
        # self.shapedirs = tf.gather(self.shapedirs, self.to_keep, axis=1)
        # self.posedirs = tf.gather(self.posedirs, self.to_keep, axis=1)
        # self.J_regressor = tf.gather(self.J_regressor, self.to_keep, axis=2)
        # self.v_template = tf.gather(self.v_template, self.to_keep, axis=1)


    def Q2R(self, w, x, y, z):
        def diag(a, b):  # computes the diagonal entries,  1 - 2*a**2 - 2*b**2
            return 1 - 2 * tf.pow(a, 2) - 2 * tf.pow(b, 2)

        def tr_add(a, b, c, d):  # computes triangle entries with addition
            return 2 * a * b + 2 * c * d

        def tr_sub(a, b, c, d):  # computes triangle entries with subtraction
            return 2 * a * b - 2 * c * d

        m = [[diag(y, z), tr_sub(x, y, z, w), tr_add(x, z, y, w)],
                [tr_add(x, y, z, w), diag(x, z), tr_sub(y, z, x, w)],
                [tr_sub(x, z, y, w), tr_add(y, z, x, w), diag(x, y)]]
            
        R = tf.stack([tf.stack(m[i], axis=-1) for i in range(3)], axis=-2)
        return R
            
    def rodrigues(self, r):
        theta = tf.norm(r, axis=(1, 2), keepdims=True)
        theta = tf.maximum(theta, np.finfo(np.float32).eps)
        r_hat = r / theta
        z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
        m = tf.stack(
            (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
        m = tf.reshape(m, (-1, 3, 3))
        i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
            (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
        # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
        R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta))
        return R

    def with_zeros(self, x):
        return tf.concat((x, tf.ones((x.shape[0], 1, 4)) * tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)[tf.newaxis,:,:]), axis=1)

    def pack(self, x):
        return tf.concat((tf.zeros((x.shape[0], 24, 4, 3), dtype=tf.float32), x), axis=3)
        
    def with_zeros_old(self, x):
        return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0)

    def pack_old(self, x):
        return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
        

    def run_rodrigues(self, pose):
        batch_size = pose.shape[0]
        R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3)))
        R_out = tf.reshape(R_out, (batch_size, -1, 3, 3))
        #R_out = []
        #for i in range(batch_size):
        #    R_out.append(self.rodrigues(tf.reshape(pose[i,:], (-1, 1, 3))))
        #R_out = tf.stack(R_out, axis=0)
        return R_out

    def run(self, R, betas, remove=True):
        batch_size = betas.shape[0]

        #J_regressor = tf.gather(self.J_regressor, gender, axis=0)
        #shapedirs = tf.gather(self.shapedirs, gender, axis=0)
        #v_template = tf.gather(self.v_template, gender, axis=0)
        
        v_shaped = []
        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
        v_shaped = tf.stack(v_shaped, axis=-1)

        def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights, parent=self.parent):
            R = x[0]
            v_shaped = x[1]
            g = x[2]

            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
            R_cube_big = R
            R_cube = R_cube_big[1:]
            I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
            lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
            v_posed = v_shaped[:, :, g] + tf.tensordot(posedirs[g], lrotmin, axes=[[2], [0]])
            results = []
            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
            for j in range(1, 24):
                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
            stacked = tf.stack(results, axis=0)
            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
            
            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
            
            T = tf.tensordot(weights[g], results, axes=((1), (0)))
            rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
            v = tf.reshape(v, (-1, 4))[:, :3]

            return v, tf.squeeze(J), v_posed

        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
        T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32))
        
        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
        J_out = tf.subtract(J_out, J0)
        T_out = tf.subtract(T_out, J0)
        v_posed = tf.subtract(v_posed, J0)
        
        if remove:
            T_out = tf.gather(T_out, self.to_keep, axis=1)
            v_posed = tf.gather(v_posed, self.to_keep, axis=1)
        
        return T_out, J_out, v_posed

    def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):
        
        v_shaped = []
        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
        v_shaped = tf.stack(v_shaped, axis=-1)

        if weights is None:
            w = self.weights_one_hot[np.newaxis, :, :]
            if filter is not None:
                w = tf.gather(tf.gather(w, self.to_keep, axis=1), filter, axis=1)
            if v_cloth.shape[1] == 14475:
                w = tf.gather(w, self.to_keep, axis=1)
        else:
            w = weights

        def single_run(x, J_regressor=self.J_regressor, w=w, parent=self.parent):
            R = x[0]
            v_shaped = x[1]
            g = x[2]
            v_cloth = x[3]

            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
            R_cube_big = R
            results = []
            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
            for j in range(1, 24):
                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
            stacked = tf.stack(results, axis=0)
            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
            
            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
            
            T = tf.tensordot(w[0], results, axes=((1), (0)))
            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
            v = tf.reshape(v, (-1, 4))[:, :3]

            return v, tf.squeeze(J)

        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
        T_out, J_out = tf.map_fn(single_run, (R, v_shaped, gender, v_cloth), (tf.float32, tf.float32))
        
        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
        T_out = tf.subtract(T_out, J0)

        return T_out



    def get_joints_template(self, betas, gender, batch_size):
        v_shaped = tf.tensordot(betas[tf.newaxis,:], self.shapedirs[gender], axes=[[1], [2]]) + tf.reshape(self.v_template[gender], shape=(1,27578,3))
        J = tf.matmul(self.J_regressor[gender], v_shaped)
        J0 = tf.reshape(J[:,0,:], (-1, 1, 3))
        J_single = tf.subtract(J, J0)#[0]
        J = tf.ones((batch_size, 24, 3)) * J_single

        return J


    def run_garment(self, R, betas, gender, v_cloth, filter_idx=None):
        batch_size = v_cloth.shape[0]

        J = self.get_joints_template(betas, gender, batch_size)
        
        if filter_idx is not None:
            weights = tf.gather(self.weights, filter_idx, axis=1)

        R_cube_big = R
        results = []
        results.append(self.with_zeros(tf.concat((R_cube_big[:,0,:,:], J[:, 0, :, tf.newaxis]), axis=2)))
        for j in range(1, 24):
            results.append(tf.matmul(results[self.parent[j]], self.with_zeros(tf.concat((R_cube_big[:,j,:,:], (J[:, j, :] - J[:, self.parent[j], :])[:,:,tf.newaxis]), axis=2))))
        stacked = tf.stack(results, axis=1)
        results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((batch_size, 24, 1), dtype=tf.float32)), axis=2), (batch_size, 24, 4, 1))))
        
        J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([batch_size, 24, 1], tf.float32)], axis=2), axis=2), transpose_b=True)
        J = J[:, :, :3, :]
        
        T = tf.tensordot(weights[gender, :, :], results, axes=((1), (1)))
        T = tf.transpose(T, [1,0,2,3])
        
        rest_shape_h = tf.concat((v_cloth, tf.ones((batch_size, v_cloth.shape[1], 1), dtype=tf.float32)), axis=2)
        v = tf.matmul(T, tf.reshape(rest_shape_h, (batch_size, -1, 4, 1)))
        v = tf.reshape(v, (batch_size, -1, 4))[:, :, :3]

        J_out = tf.squeeze(J)
        if J.shape[0] == 1:
            J_out = J_out[tf.newaxis, :, :]
        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
        J_out = tf.subtract(J_out, J0)
        T_out = tf.subtract(v, J0)
        
        return T_out

    def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
        
        if J is None:
            J = self.get_joints_template(betas)
        
        def single_run(x, parent=self.parent):
            R = x[0]
            J = x[1]
            v_cloth = x[2]
            weights = x[3]

            R_cube_big = R
       
            results = []
            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
            for j in range(1, 24):
                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
            stacked = tf.stack(results, axis=0)
            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
            
            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
            
            T = tf.tensordot(weights, results, axes=((1), (0)))
            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
            v = tf.reshape(v, (-1, 4))[:, :3]

            return v, tf.squeeze(J)

        # gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
        T_out, J_out = tf.map_fn(single_run, (R, J, v_cloth, weights), (tf.float32, tf.float32))
        
        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
        J_out = tf.subtract(J_out, J0)
        T_out = tf.subtract(T_out, J0)
        
        return T_out

    def run_backward(self, R, J, v_posed, weights):
        batch_size = R.get_shape()[0]
        
        J_relative = J[:, 1:, :] - tf.gather(J, tf.constant(list(self.parent.values()), tf.int32), axis=1) # batch_size x 23 x 3

        J_out = []
        J_out.append(tf.expand_dims(J[:, 0, :], axis=-1)) # batch_size x 3 x 1
        results = []
        results.append(tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((batch_size, 3, 3), dtype=tf.float32))
        for j in range(1, 24):
            results.append(tf.matmul(R[:, self.parent[j], :, :], results[self.parent[j]]))
            J_out.append(tf.matmul(results[j], tf.expand_dims(J_relative[:, j-1, :], axis=-1)) + J_out[self.parent[j]]) # batch_size x 3 x 1
        results = tf.transpose(tf.stack(results, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 3)
        J_out = tf.transpose(tf.stack(J_out, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 1)

        results = tf.matmul(R, results)
        T = tf.transpose(tf.tensordot(results, weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 3)
        v_posed = tf.expand_dims(v_posed, axis=-1) # (batch_size x 6890 x 3 x 1)
        J = tf.expand_dims(J, axis=-1) # (batch_size x 24 x 3 x 1)
        v = tf.matmul(T, v_posed) + tf.transpose(tf.tensordot(J_out - tf.matmul(results, J), weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 1)
        v = tf.squeeze(v, -1) # (batch_size x 6890 x 3)
        J_out = tf.squeeze(J_out, -1) # (batch_size x 24 x 3)

        return v, J_out