# Copyright (c) 2009-2019 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause License.

from hoomd import _hoomd
#from hoomd.jit import _jit
import hoomd
import tempfile
import shutil
import subprocess
import os

import numpy as np
from hoomd.swapmc import _swapmc as _plugin_patch
#from hoomd.jit import patch

#This python file faciltates the creation of relevant patch energies
#we'll do shifted lennard jones and polydisperse's system for now
#let's make different classes for different pair potentials. Of course, we'll let the user defined path energy to be available as well
class lj(object):
    R''' Define the patch energy of a lennard jones particle.

    Args:
        scaledr_cut (float): the scale of cutoff radius relative to particle center to center, beyond which all pair interactions are assumed 0.
        eps (float): uniform energetics
        llvm_ir_fname (str): File name of the llvm IR file to load.
        clang_exec (str): The Clang executable to use

    Patch energies define energetic interactions between pairs of shapes in :py:mod:`hpmc <hoomd.hpmc>` integrators.
    Shapes within a cutoff distance of *r_cut* are potentially interacting and the energy of interaction is a function
    the type and orientation of the particles and the vector pointing from the *i* particle to the *j* particle center.

    The :py:class:`user` patch energy takes C++ code, JIT compiles it at run time and executes the code natively
    in the MC loop at with full performance. It enables researchers to quickly and easily implement custom energetic
    interactions without the need to modify and recompile HOOMD.

    .. rubric:: C++ code

    Supply C++ code to the *code* argument and :py:class:`user` will compile the code and call it to evaluate
    patch energies. Compilation assumes that a recent ``clang`` installation is on your PATH. This is convenient
    when the energy evaluation is simple or needs to be modified in python. More complex code (i.e. code that
    requires auxiliary functions or initialization of static data arrays) should be compiled outside of HOOMD
    and provided via the *llvm_ir_file* input (see below).

    The text provided in *code* is the body of a function with the following signature:

    .. code::

        float eval(const vec3<float>& r_ij,
                   unsigned int type_i,
                   const quat<float>& q_i,
                   float d_i,
                   float charge_i,
                   unsigned int type_j,
                   const quat<float>& q_j,
                   float d_j,
                   float charge_j)

    * ``vec3`` and ``quat`` are defined in HOOMDMath.h.
    * *r_ij* is a vector pointing from the center of particle *i* to the center of particle *j*.
    * *type_i* is the integer type of particle *i*
    * *q_i* is the quaternion orientation of particle *i*
    * *d_i* is the diameter of particle *i*
    * *charge_i* is the charge of particle *i*
    * *type_j* is the integer type of particle *j*
    * *q_j* is the quaternion orientation of particle *j*
    * *d_j* is the diameter of particle *j*
    * *charge_j* is the charge of particle *j*
    * Your code *must* return a value.
    * When \|r_ij\| is greater than *r_cut*, the energy *must* be 0. This *r_cut* is applied between
      the centers of the two particles: compute it accordingly based on the maximum range of the anisotropic
      interaction that you implement.

    Example:

    .. code-block:: python

        square_well = """float rsq = dot(r_ij, r_ij);
                            if (rsq < 1.21f)
                                return -1.0f;
                            else
                                return 0.0f;
                      """
        patch = hoomd.jit.patch.user(mc=mc, r_cut=1.1, code=square_well)

    .. rubric:: LLVM IR code

    You can compile outside of HOOMD and provide a direct link
    to the LLVM IR file in *llvm_ir_file*. A compatible file contains an extern "C" eval function with this signature:

    .. code::

        float eval(const vec3<float>& r_ij,
                   unsigned int type_i,
                   const quat<float>& q_i,
                   float d_i,
                   float charge_i,
                   unsigned int type_j,
                   const quat<float>& q_j,
                   float d_j,
                   float charge_j)

    ``vec3`` and ``quat`` are defined in HOOMDMath.h.

    Compile the file with clang: ``clang -O3 --std=c++11 -DHOOMD_LLVMJIT_BUILD -I /path/to/hoomd/include -S -emit-llvm code.cc`` to produce
    the LLVM IR in ``code.ll``.

    .. versionadded:: 2.3
    '''
    def __init__(self, mc, kT, scaledr_cut, eps, mode='shifted', llvm_ir_file=None, clang_exec=None):
        hoomd.util.print_status_line();
        #there is no need to write your own code. By definition, we're
        if (mode == 'truncated'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return 4.0f*eps*r6inv*(r6inv-1.0f);
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        elif (mode == 'shifted'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;

                               float rsqinv_cut = sigmasq /(rcut*rcut);
                               float r6inv_cut = rsqinv_cut*rsqinv_cut*rsqinv_cut;
                               return 4.0f*eps*(r6inv*(r6inv-1.0f)-r6inv_cut*(r6inv_cut-1.0f));
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        elif (mode == 'force-shift'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;

                               float rsqinv_cut = sigmasq /(rcut*rcut);
                               float r6inv_cut = rsqinv_cut*rsqinv_cut*rsqinv_cut;
                               return 4.0f*eps*( r6inv*(r6inv-1.0f)-r6inv_cut*(r6inv_cut-1.0f) +6*(sqrtf(rsq)/rcut-1.0f)*r6inv_cut*(2*r6inv_cut-1.0f) );
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        # check if initialization has occurred
        if hoomd.context.exec_conf is None:
            raise RuntimeError('Error creating Lennard-Jones patch energy, call context.initialize() first');

        # raise an error if this run is on the GPU
        if hoomd.context.exec_conf.isCUDAEnabled():
            hoomd.context.msg.error("Patch energies are not supported on the GPU\n");
            raise RuntimeError("Error initializing patch energy");

        # Find a clang executable if none is provided
        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang'

        if code is not None and llvm_ir_file is None:
            llvm_ir = self.compile_user(code, clang)
        else:
            # IR is a text file
            with open(llvm_ir_file,'r') as f:
                llvm_ir = f.read()

        self.compute_name = "patch"
        self.cpp_evaluator = _plugin_patch.PatchEnergyJITCustom(hoomd.context.exec_conf, llvm_ir, scaledr_cut);
        mc.set_PatchEnergyEvaluator(self);

        self.mc = mc
        self.enabled = True
        self.log = False

    def compile_user(self, code, clang_exec, fn=None):
        R'''Helper function to compile the provided code into an executable

        Args:
            code (str): C++ code to compile
            clang_exec (str): The Clang executable to use
            fn (str): If provided, the code will be written to a file.

        .. versionadded:: 2.3
        '''
        cpp_function = """
#include "hoomd/HOOMDMath.h"
#include "hoomd/VectorMath.h"

extern "C"
{
float eval(const vec3<float>& r_ij,
    unsigned int type_i,
    const quat<float>& q_i,
    float d_i,
    float charge_i,
    unsigned int type_j,
    const quat<float>& q_j,
    float d_j,
    float charge_j)
    {
"""
        cpp_function += code
        cpp_function += """
    }
}
"""

        include_path = os.path.dirname(hoomd.__file__) + '/include';
        include_path_source = hoomd._hoomd.__hoomd_source_dir__;

        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang';

        if fn is not None:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o',fn,'-']
        else:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o','-','-']
        p = subprocess.Popen(cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)

        # pass C++ function to stdin
        output = p.communicate(cpp_function.encode('utf-8'))
        llvm_ir = output[0].decode()

        f = open("llvm_ir_file",'w')
        f.write(llvm_ir)
        f.close()

        if p.returncode != 0:
            hoomd.context.msg.error("Error compiling provided code\n");
            hoomd.context.msg.error("Command "+' '.join(cmd)+"\n");
            hoomd.context.msg.error(output[1].decode()+"\n");
            raise RuntimeError("Error initializing patch energy");

        return llvm_ir

    R''' Disable the patch energy and optionally enable it only for logging

    Args:
        log (bool): If true, only use patch energy as a log quantity

    '''
    def disable(self,log=None):
        hoomd.util.print_status_line();

        if log:
            # enable only for logging purposes
            self.mc.cpp_integrator.disablePatchEnergyLogOnly(log)
            self.log = True
        else:
            # disable completely
            self.mc.cpp_integrator.setPatchEnergy(None);
            self.log = False

        self.enabled = False

    R''' (Re-)Enable the patch energy

    '''
    def enable(self):
        hoomd.util.print_status_line()
        self.mc.cpp_integrator.setPatchEnergy(self.cpp_evaluator);

class softrepulsive(object):
    R''' Define the patch energy of a soft-repulsive particle.
    '''
    def __init__(self, mc, kT, scaledr_cut, eps, mode='shifted', llvm_ir_file=None, clang_exec=None):
        hoomd.util.print_status_line();
        #there is no need to write your own code. By definition, we're
        if (mode == 'truncated'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return 4.0f*eps*r6inv*r6inv;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        elif (mode == 'shifted'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;

                               float rsqinv_cut = sigmasq /(rcut*rcut);
                               float r6inv_cut = rsqinv_cut*rsqinv_cut*rsqinv_cut;
                               return 4.0f*eps*(r6inv*r6inv-r6inv_cut*r6inv_cut);
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        elif (mode == 'force-shift'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;

                               float rsqinv_cut = sigmasq /(rcut*rcut);
                               float r6inv_cut = rsqinv_cut*rsqinv_cut*rsqinv_cut;
                               return 4.0f*eps*( r6inv*(r6inv)-r6inv_cut*(r6inv_cut) +6*(sqrtf(rsq)/rcut-1.0f)*r6inv_cut*(2*r6inv_cut) );
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        # check if initialization has occurred
        if hoomd.context.exec_conf is None:
            raise RuntimeError('Error creating Lennard-Jones patch energy, call context.initialize() first');

        # raise an error if this run is on the GPU
        if hoomd.context.exec_conf.isCUDAEnabled():
            hoomd.context.msg.error("Patch energies are not supported on the GPU\n");
            raise RuntimeError("Error initializing patch energy");

        # Find a clang executable if none is provided
        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang'

        if code is not None and llvm_ir_file is None:
            llvm_ir = self.compile_user(code, clang)
        else:
            # IR is a text file
            with open(llvm_ir_file,'r') as f:
                llvm_ir = f.read()

        self.compute_name = "patch"
        self.cpp_evaluator = _plugin_patch.PatchEnergyJITCustom(hoomd.context.exec_conf, llvm_ir, scaledr_cut);
        mc.set_PatchEnergyEvaluator(self);

        self.mc = mc
        self.enabled = True
        self.log = False

    def compile_user(self, code, clang_exec, fn=None):
        R'''Helper function to compile the provided code into an executable

        Args:
            code (str): C++ code to compile
            clang_exec (str): The Clang executable to use
            fn (str): If provided, the code will be written to a file.

        .. versionadded:: 2.3
        '''
        cpp_function = """
#include "hoomd/HOOMDMath.h"
#include "hoomd/VectorMath.h"

extern "C"
{
float eval(const vec3<float>& r_ij,
    unsigned int type_i,
    const quat<float>& q_i,
    float d_i,
    float charge_i,
    unsigned int type_j,
    const quat<float>& q_j,
    float d_j,
    float charge_j)
    {
"""
        cpp_function += code
        cpp_function += """
    }
}
"""

        include_path = os.path.dirname(hoomd.__file__) + '/include';
        include_path_source = hoomd._hoomd.__hoomd_source_dir__;

        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang';

        if fn is not None:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o',fn,'-']
        else:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o','-','-']
        p = subprocess.Popen(cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)

        # pass C++ function to stdin
        output = p.communicate(cpp_function.encode('utf-8'))
        llvm_ir = output[0].decode()

        f = open("llvm_ir_file",'w')
        f.write(llvm_ir)
        f.close()

        if p.returncode != 0:
            hoomd.context.msg.error("Error compiling provided code\n");
            hoomd.context.msg.error("Command "+' '.join(cmd)+"\n");
            hoomd.context.msg.error(output[1].decode()+"\n");
            raise RuntimeError("Error initializing patch energy");

        return llvm_ir

    R''' Disable the patch energy and optionally enable it only for logging

    Args:
        log (bool): If true, only use patch energy as a log quantity

    '''
    def disable(self,log=None):
        hoomd.util.print_status_line();

        if log:
            # enable only for logging purposes
            self.mc.cpp_integrator.disablePatchEnergyLogOnly(log)
            self.log = True
        else:
            # disable completely
            self.mc.cpp_integrator.setPatchEnergy(None);
            self.log = False

        self.enabled = False

    R''' (Re-)Enable the patch energy

    '''
    def enable(self):
        hoomd.util.print_status_line()
        self.mc.cpp_integrator.setPatchEnergy(self.cpp_evaluator);

class polydisperse(object):
    R''' Define the patch energy of a polydisperse soft-repulsive
    '''
    def __init__(self, mc, kT, scaledr_cut, v0, eps, model, llvm_ir_file=None, clang_exec=None,kappa=3.0):
        hoomd.util.print_status_line();

        if (model == "polydisperse12"):
            a = scaledr_cut
            A = np.array([  [1,a**2,a**4],
                            [0,2*a,4*a**3],
                            [0,2,12*a**2]
                            ])
            b = np.array([-v0/a**12,12*v0/a**13,-156*v0/a**14])
            c = np.linalg.solve(A,b)
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return v0*r6inv*r6inv+c0+c1*_rsq+c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c[0]/kT,c[1]/kT,c[2]/kT);
        elif (model == "lennardjones"):
            a = scaledr_cut
            A = np.array([  [1,a**2,a**4],
                            [0,2*a,4*a**3],
                            [0,2,12*a**2]
                            ])
            b = np.array([-v0*(1/a**12-1/a**6),v0*(12*v0/a**13-6/a**7),-v0*(156/a**14-42/a**8)])
            c = np.linalg.solve(A,b)
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return v0*(r6inv*r6inv-r6inv)+c0+c1*_rsq+c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c[0]/kT,c[1]/kT,c[2]/kT);
        elif (model == "polydisperse18"):
            a = scaledr_cut
            A = np.array([  [1,a**2,a**4],
                            [0,2*a,4*a**3],
                            [0,2,12*a**2]
                            ])
            b = np.array([-v0*(1/a**18),(18*v0/a**19),-v0*(342/a**20)])
            c = np.linalg.solve(A,b)
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return v0*(r6inv*r6inv*r6inv)+c0+c1*_rsq+c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c[0]/kT,c[1]/kT,c[2]/kT);
        elif (model == "polydisperse10"):
            a = scaledr_cut
            c0 =  -(56.0)*v0/(a**10);
            c1 =  (140.0)*v0/(a**12);
            c2 =  -(120.0)*v0/(a**14);
            c3 =  (35.0)*v0/(a**16);
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float c3   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r10inv = rsqinv*rsqinv*rsqinv*rsqinv*rsqinv;
                               return v0*r10inv+c0+c1*_rsq+c2*_rsq*_rsq+c3*_rsq*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c0/kT,c1/kT,c2/kT,c3/kT);
        elif (model == "polydisperse106"):
            a = scaledr_cut
            c0 = (-21 + 10*a**4)*v0/a**10
            c1 =  -((5*(-7 + 3*a**4)*v0)/a**12)
            c2 = (3*(-5 + 2*a**4)*v0)/a**14
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r10inv = rsqinv*rsqinv*rsqinv*rsqinv*rsqinv;
                               return v0*(r10inv-rsqinv*rsqinv*rsqinv)+c0+c1*_rsq + c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c0/kT,c1/kT,c2/kT);
        elif (model == "polydisperseyukawa"):
            a = scaledr_cut
            c0 = -((np.exp(-kappa*a)*(15 + 7*kappa*a + kappa**2*a**2)*v0)/(8*a))
            c1 = (np.exp(-kappa*a)*(5 + 5*kappa*a + kappa**2*a**2)*v0)/(4*a**3)
            c2 = -((np.exp(-kappa*a)*(3 + 3*kappa*a + kappa**2*a**2)*v0)/(8*a**5))
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float kappa = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float rinv = sqrtf(rsqinv);
                               float _rsq = rsq / sigmasq;
                               float _r = sqrtf(_rsq);
                               float r10inv = rsqinv*rsqinv*rsqinv*rsqinv*rsqinv;
                               return v0*exp(-_r*kappa)*rinv+c0+c1*_rsq + c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,kappa,c0/kT,c1/kT,c2/kT);
        else:
            raise RuntimeError('Error creating Polydisperse patch energy. Not one of the available models. Perhaps theres a typo?');

        # check if initialization has occurred
        if hoomd.context.exec_conf is None:
            raise RuntimeError('Error creating Lennard-Jones patch energy, call context.initialize() first');

        # raise an error if this run is on the GPU
        if hoomd.context.exec_conf.isCUDAEnabled():
            hoomd.context.msg.error("Patch energies are not supported on the GPU\n");
            raise RuntimeError("Error initializing patch energy");

        # Find a clang executable if none is provided
        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang'

        if code is not None and llvm_ir_file is None:
            llvm_ir = self.compile_user(code, clang)
        else:
            # IR is a text file
            with open(llvm_ir_file,'r') as f:
                llvm_ir = f.read()

        self.compute_name = "patch"
        self.cpp_evaluator = _plugin_patch.PatchEnergyJITCustom(hoomd.context.exec_conf, llvm_ir, scaledr_cut);
        mc.set_PatchEnergyEvaluator(self);

        self.mc = mc
        self.enabled = True
        self.log = False

    def compile_user(self, code, clang_exec, fn=None):
        R'''Helper function to compile the provided code into an executable

        Args:
            code (str): C++ code to compile
            clang_exec (str): The Clang executable to use
            fn (str): If provided, the code will be written to a file.

        .. versionadded:: 2.3
        '''
        cpp_function = """
#include "hoomd/HOOMDMath.h"
#include "hoomd/VectorMath.h"

extern "C"
{
float eval(const vec3<float>& r_ij,
    unsigned int type_i,
    const quat<float>& q_i,
    float d_i,
    float charge_i,
    unsigned int type_j,
    const quat<float>& q_j,
    float d_j,
    float charge_j)
    {
"""
        cpp_function += code
        cpp_function += """
    }
}
"""

        include_path = os.path.dirname(hoomd.__file__) + '/include';
        include_path_source = hoomd._hoomd.__hoomd_source_dir__;

        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang';

        if fn is not None:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o',fn,'-']
        else:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o','-','-']
        p = subprocess.Popen(cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)

        # pass C++ function to stdin
        output = p.communicate(cpp_function.encode('utf-8'))
        llvm_ir = output[0].decode()

        f = open("llvm_ir_file",'w')
        f.write(llvm_ir)
        f.close()

        if p.returncode != 0:
            hoomd.context.msg.error("Error compiling provided code\n");
            hoomd.context.msg.error("Command "+' '.join(cmd)+"\n");
            hoomd.context.msg.error(output[1].decode()+"\n");
            raise RuntimeError("Error initializing patch energy");

        return llvm_ir

    R''' Disable the patch energy and optionally enable it only for logging

    Args:
        log (bool): If true, only use patch energy as a log quantity

    '''
    def disable(self,log=None):
        hoomd.util.print_status_line();

        if log:
            # enable only for logging purposes
            self.mc.cpp_integrator.disablePatchEnergyLogOnly(log)
            self.log = True
        else:
            # disable completely
            self.mc.cpp_integrator.setPatchEnergy(None);
            self.log = False

        self.enabled = False

    R''' (Re-)Enable the patch energy

    '''
    def enable(self):
        hoomd.util.print_status_line()
        self.mc.cpp_integrator.setPatchEnergy(self.cpp_evaluator);


# Copyright (c) 2020 Graphcore Ltd. All rights reserved.
import argparse
import numpy as np
from functools import partial
import tensorflow.compat.v1 as tf
from tensorflow.python import ipu
from ipu_sparse_ops import sparse, layers

tf.disable_eager_execution()
tf.disable_v2_behavior()

def get_program_arguments():
    parser = argparse.ArgumentParser(description='Sparse matmul tool.')
    parser.add_argument("--batch-size", type=int,
                        help="Number of vectors in a mini-batch")
    parser.add_argument("--input-size", type=int,
                        help="Input size (rows of weight matrix)")
    parser.add_argument("--output-size", type=int,
                        help="Output size (cols of weight matrix)")
    parser.add_argument("--pattern", type=str,
                        help="Choose how the sparsity pattern is generated.",
                        choices=['random', 'fixed'])
    parser.add_argument("--density", type=float,
                        help="Non-zero density (only used if --pattern is random).")
    parser.set_defaults(
        batch_size=2, input_size=16, output_size=8, pattern='fixed', density=0.1)
    return parser.parse_args()

def make_triplets_test_inputs(args, data_type):
    input_size = args.input_size
    output_size = args.output_size
    batch_size = args.batch_size
    num_groups = 1
    topk_ratio = 0.5
    if args.pattern == 'fixed':
        rhs_values = np.random.rand(input_size, output_size)
        sparse_mask = np.identity(input_size)
        sparse_mask[1, 3] = 1
        sparse_mask[0, 7] = 1
        masked_rhs = np.multiply(sparse_mask[:, 0:output_size], rhs_values)
        triplets = sparse.triplets_from_dense(masked_rhs)
        fc = layers.SparseFcLayer.from_triplets(
            args.output_size, [args.batch_size, args.input_size], topk_ratio, *triplets)
    else:
        random_gen = np.random.default_rng(seed=random_seed)
        fc = layers.SparseFcLayer.from_random_generator(
            args.output_size, [args.batch_size, args.input_size], args.density, topk_ratio,
            random_gen.standard_normal, random_seed)
        masked_rhs = sparse.dense_from_triplets(fc.spec, *fc.triplets)
    return fc, masked_rhs

tf.logging.set_verbosity(tf.logging.ERROR)
np.set_printoptions(linewidth=200)
random_seed = 1
np.random.seed(random_seed)

# Make input data and dense and sparse weights:
args = get_program_arguments()
data_type = tf.float32
lhs_values = np.random.rand(args.batch_size, args.input_size)
fc, masked_rhs = make_triplets_test_inputs(args, data_type)

# Configure the IPU:
cfg = ipu.utils.create_ipu_config()
cfg = ipu.utils.auto_select_ipus(cfg, 1)
ipu.utils.configure_ipu_system(cfg)

with tf.device("cpu"):
    # Placeholders for the dense matmul operands:
    lhs = tf.placeholder(data_type, lhs_values.shape)
    rhs = tf.placeholder(data_type, masked_rhs.shape)
    # Placeholders for the sparse representation:
    fc.create_placeholders(data_type)
    compute_dense_grad_w = tf.placeholder(tf.bool, shape=[])

def matmul_with_grad(input, weights):
    z = tf.matmul(input, weights)
    s = tf.reduce_sum(z)
    input_grad = tf.gradients(s, input)
    weights_grad = tf.gradients(s, weights)
    return z, input_grad, weights_grad

def sparse_matmul_with_grad(fc, do_topk, input):
    with tf.variable_scope("fc", reuse=tf.AUTO_REUSE, use_resource=True):
        z = fc(input, do_topk)
        s = tf.reduce_sum(z)
        input_grad = tf.gradients(s, input)
        weights_grad = tf.gradients(s, fc.get_values_var())
        return z, input_grad, weights_grad, fc.get_dense_grad_w(s)

# Build the IPU graphs:
sparse_matmul = partial(sparse_matmul_with_grad, fc)
with ipu.scopes.ipu_scope("/device:IPU:0"):
    reference_fetches = ipu.ipu_compiler.compile(matmul_with_grad, [lhs, rhs])
    sparse_fetches = ipu.ipu_compiler.compile(sparse_matmul, [compute_dense_grad_w, lhs])
    with tf.variable_scope("fc", reuse=tf.AUTO_REUSE, use_resource=True):
        sparse_data_update_op = fc.update_sparsity_op()

ipu.utils.move_variable_initialization_to_cpu()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    reference_result, reference_input_grad, reference_weight_grad = sess.run(
        reference_fetches,
        feed_dict={lhs: lhs_values, rhs: masked_rhs})

    sess.run(sparse_data_update_op, feed_dict=fc.feed_dict())
    sparse_result, sparse_input_grad, sparse_weight_grad, dense_grad_w = sess.run(
        sparse_fetches, feed_dict={lhs: lhs_values, compute_dense_grad_w: True})

# Check all the results:

# Convert the sparse gradient metainfo back to triplets and then use those row and col indices
# to index the dense reference weight gradient:
sparse_data = sparse.SparseRepresentation(fc.data.metainfo_state, sparse_weight_grad[0])
triplets = sparse.triplets_from_representation(fc.spec, sparse_data)
reference_grad_nzvalues = sparse.values_at_indices(triplets[0], triplets[1], reference_weight_grad[0])

# Convert the dense reference weight gradient to a sparse one using the same mask
# that we used for the weights so we can compare the nzvalues against the sparse grad:
_, _, values = sparse.triplets_from_dense(reference_weight_grad[0])
sparse_data = sparse.representation_from_triplets(fc.spec, *triplets)
reference_grad_nzvalues = sparse_data.nz_values

# Need to set tolerances for fp32 as numpy is set for doubles by default:
rtol = 1e-05
atol = 1e-06

if not np.allclose(reference_result, sparse_result, rtol=rtol, atol=atol, equal_nan=True):
    print(f"Reference result:\n{reference_result}")
    print(f"Sparse result:\n{sparse_result}")
    diff = reference_result-sparse_result
    print(f"Difference:\n{diff}")
    diff_triplet = sparse.triplets_from_dense(diff)
    print(f"Difference triplets:\nrows: {diff_triplet[0]}\ncols: {diff_triplet[1]}\nvalues: {diff_triplet[2]}")
    raise RuntimeError("Sparse and reference results do not match.")

if not np.allclose(reference_input_grad, sparse_input_grad, rtol=rtol, atol=atol, equal_nan=True):
    raise RuntimeError("Sparse and reference input gradients do not match.")

if not np.allclose(sparse_data.nz_values, sparse_weight_grad, rtol=rtol, atol=atol, equal_nan=True):
    print(f"Reference weight grad (dense):\n{reference_weight_grad}")
    print(f"Reference weight grad (sparse):\n{sparse_data.nz_values}")
    print(f"Sparse weight grad:\n{sparse_weight_grad[0]}")
    raise RuntimeError("Sparse and reference weight gradients do not match.")

if not np.array_equal(dense_grad_w, reference_weight_grad):
    print(f"Reference weight grad (dense):\n{reference_weight_grad}")
    print(f"Dense grad W from custom op:\n{dense_grad_w}")
    raise ValueError("Custom Op's Dense Weight Grad does not match reference.")

print("Results match.")


# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# 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.
#
# Contact: ps-license@tuebingen.mpg.de

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import numpy as np
import cv2
import time
# from numba import jit
# import numba

import torch
import torch.nn.functional as F

from .utils import rot_mat_to_euler

def find_dynamic_lmk_idx_and_bcoords_projective(vertices, neck_joint_loc, pose, dynamic_lmk_faces_idx,
                                     dynamic_lmk_b_coords,
                                     neck_kin_chain, dtype=torch.float32):
    batch_size = vertices.shape[0]
    # print('pose shape len')
    # print(len(pose.shape))
    if len(pose.shape) == 2:
        aa_pose = torch.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 = torch.index_select(pose.view(batch_size, -1, 3, 3), 1,
                                     neck_kin_chain)

    b = rot_mats.shape[0]

    rel_rot_mat = torch.eye(3, device=vertices.device,
                            dtype=dtype).reshape(1, 3, 3).repeat(b, 1, 1)
    for idx in range(len(neck_kin_chain)):
        rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)
    neck_cam_loc = -torch.bmm(torch.transpose(rel_rot_mat, 1, 2), neck_joint_loc.reshape(-1, 3, 1)).reshape(-1, 3)
    neck_cam_loc_len = torch.norm(neck_cam_loc, dim=1).reshape(-1, 1).repeat(1, 3)
    neck_cam_dir = neck_cam_loc / neck_cam_loc_len
    y_ang = -torch.atan2(neck_cam_dir[:, 0], neck_cam_dir[:, 2])

    # print('neck cam dir')
    # print(neck_cam_dir)

    y_rot_angle = torch.round(
        torch.clamp(y_ang * 180.0 / np.pi,
                    max=39)).to(dtype=torch.long)

    # print('face cam ang')
    # print(y_rot_angle)

    neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
    mask = y_rot_angle.lt(-39).to(dtype=torch.long)
    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 = torch.index_select(dynamic_lmk_faces_idx,
                                           0, y_rot_angle)
    dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords,
                                          0, y_rot_angle)

    return dyn_lmk_faces_idx, dyn_lmk_b_coords

def find_dynamic_lmk_idx_and_bcoords(vertices, pose, dynamic_lmk_faces_idx,
                                     dynamic_lmk_b_coords,
                                     neck_kin_chain, dtype=torch.float32):
    ''' Compute the faces, barycentric coordinates for the dynamic landmarks

        To do so, we first compute the rotation of the neck around the y-axis
        and then use a pre-computed look-up table to find the faces and the
        barycentric coordinates that will be used.

        Special thanks to Soubhik Sanyal (soubhik.sanyal@tuebingen.mpg.de)
        for providing the original TensorFlow implementation and for the LUT.

        Parameters
        ----------
        vertices: torch.tensor BxVx3, dtype = torch.float32
            The tensor of input vertices
        pose: torch.tensor Bx(Jx3), dtype = torch.float32
            The current pose of the body model
        dynamic_lmk_faces_idx: torch.tensor L, dtype = torch.long
            The look-up table from neck rotation to faces
        dynamic_lmk_b_coords: torch.tensor Lx3, dtype = torch.float32
            The look-up table from neck rotation to barycentric coordinates
        neck_kin_chain: list
            A python list that contains the indices of the joints that form the
            kinematic chain of the neck.
        dtype: torch.dtype, optional

        Returns
        -------
        dyn_lmk_faces_idx: torch.tensor, dtype = torch.long
            A tensor of size BxL that contains the indices of the faces that
            will be used to compute the current dynamic landmarks.
        dyn_lmk_b_coords: torch.tensor, dtype = torch.float32
            A tensor of size BxL that contains the indices of the faces that
            will be used to compute the current dynamic landmarks.
    '''

    batch_size = vertices.shape[0]

    if len(pose.shape) == 2:
        aa_pose = torch.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 = torch.index_select(pose.view(batch_size, -1, 3, 3), 1,
                                      neck_kin_chain)

    b = rot_mats.shape[0]

    rel_rot_mat = torch.eye(3, device=vertices.device,
                            dtype=dtype).reshape(1, 3, 3).repeat(b, 1, 1)
    for idx in range(len(neck_kin_chain)):
        rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)

    y_rot_angle = torch.round(
        torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
                    max=39)).to(dtype=torch.long)
    neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
    mask = y_rot_angle.lt(-39).to(dtype=torch.long)
    neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
    y_rot_angle = (neg_mask * neg_vals +
                   (1 - neg_mask) * y_rot_angle)
#     print('std yrot ang')
#     print(y_rot_angle)

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

    return dyn_lmk_faces_idx, dyn_lmk_b_coords

def vertices2landmarks(vertices, faces, lmk_faces_idx, lmk_bary_coords):
    ''' Calculates landmarks by barycentric interpolation

        Parameters
        ----------
        vertices: torch.tensor BxVx3, dtype = torch.float32
            The tensor of input vertices
        faces: torch.tensor Fx3, dtype = torch.long
            The faces of the mesh
        lmk_faces_idx: torch.tensor L, dtype = torch.long
            The tensor with the indices of the faces used to calculate the
            landmarks.
        lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32
            The tensor of barycentric coordinates that are used to interpolate
            the landmarks

        Returns
        -------
        landmarks: torch.tensor BxLx3, dtype = torch.float32
            The coordinates of the landmarks for each mesh in the batch
    '''
    # Extract the indices of the vertices for each face
    # BxLx3
    batch_size, num_verts = vertices.shape[:2]
    device = vertices.device

    lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view(
        batch_size, -1, 3)

    lmk_faces += torch.arange(
        batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts

    lmk_vertices = vertices.view(-1, 3)[lmk_faces].view(
        batch_size, -1, 3, 3)

    landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords])
    return landmarks

def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents,
        lbs_weights, pose2rot=True, dtype=torch.float32):
    ''' Performs Linear Blend Skinning with the given shape and pose parameters

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        pose : torch.tensor Bx(J + 1) * 3
            The pose parameters in axis-angle format
        v_template torch.tensor BxVx3
            The template mesh that will be deformed
        shapedirs : torch.tensor 1xNB
            The tensor of PCA shape displacements
        posedirs : torch.tensor Px(V * 3)
            The pose PCA coefficients
        J_regressor : torch.tensor JxV
            The regressor array that is used to calculate the joints from
            the position of the vertices
        parents: torch.tensor J
            The array that describes the kinematic tree for the model
        lbs_weights: torch.tensor N x V x (J + 1)
            The linear blend skinning weights that represent how much the
            rotation matrix of each part affects each vertex
        pose2rot: bool, optional
            Flag on whether to convert the input pose tensor to rotation
            matrices. The default value is True. If False, then the pose tensor
            should already contain rotation matrices and have a size of
            Bx(J + 1)x9
        dtype: torch.dtype, optional

        Returns
        -------
        verts: torch.tensor BxVx3
            The vertices of the mesh after applying the shape and pose
            displacements.
        joints: torch.tensor BxJx3
            The joints of the model
    '''

    batch_size = max(betas.shape[0], pose.shape[0])
    device = betas.device

    # Add shape contribution
    v_shaped = v_template + blend_shapes(betas, shapedirs)

    # Get the joints
    # NxJx3 array
    J = vertices2joints(J_regressor, v_shaped)

    # print('j1t')
    # print(J[0][1])

    # 3. Add pose blend shapes
    # N x J x 3 x 3
    ident = torch.eye(3, dtype=dtype, device=device)
    if pose2rot:
        rot_mats = batch_rodrigues(
            pose.view(-1, 3), dtype=dtype).view([batch_size, -1, 3, 3])

        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
        # (N x P) x (P, V * 3) -> N x V x 3
        pose_offsets = torch.matmul(pose_feature, posedirs) \
            .view(batch_size, -1, 3)
    else:
        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
        rot_mats = pose.view(batch_size, -1, 3, 3)

        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1),
                                    posedirs).view(batch_size, -1, 3)

    v_posed = pose_offsets + v_shaped
    # 4. Get the global joint location

    # print(rot_mats[0][1])

    J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)

    # 5. Do skinning:
    # W is N x V x (J + 1)
    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
    # (N x V x (J + 1)) x (N x (J + 1) x 16)
    num_joints = J_regressor.shape[0]
    T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \
        .view(batch_size, -1, 4, 4)

    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1],
                               dtype=dtype, device=device)
    v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
    v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))

    verts = v_homo[:, :, :3, 0]

    return verts, J_transformed, A, J

def vertices2joints(J_regressor, vertices):
    ''' Calculates the 3D joint locations from the vertices

    Parameters
    ----------
    J_regressor : torch.tensor JxV
        The regressor array that is used to calculate the joints from the
        position of the vertices
    vertices : torch.tensor BxVx3
        The tensor of mesh vertices

    Returns
    -------
    torch.tensor BxJx3
        The location of the joints
    '''

    return torch.einsum('bik,ji->bjk', [vertices, J_regressor])

def vertices2joints_np(J_regressor_np, vertices_np):
    ''' Calculates the 3D joint locations from the vertices

    Parameters
    ----------
    J_regressor : torch.tensor JxV
        The regressor array that is used to calculate the joints from the
        position of the vertices
    vertices : torch.tensor BxVx3
        The tensor of mesh vertices

    Returns
    -------
    torch.tensor BxJx3
        The location of the joints
    '''
    n_j, n_v = J_regressor_np.shape
    b = vertices_np.shape[0]
    return np.matmul(np.tile(J_regressor_np.reshape(1, n_j, n_v), (b, 1, 1)), vertices_np)

def blend_shapes(betas, shape_disps):
    ''' Calculates the per vertex displacement due to the blend shapes

    Parameters
    ----------
    betas : torch.tensor Bx(num_betas)
        Blend shape coefficients
    shape_disps: torch.tensor Vx3x(num_betas)
        Blend shapes

    Returns
    -------
    torch.tensor BxVx3
        The per-vertex displacement due to shape deformation
    '''

    # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
    # i.e. Multiply each shape displacement by its corresponding beta and
    # then sum them.
    blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps])
    return blend_shape

def blend_shapes_np(betas, shape_disps):
    ''' Calculates the per vertex displacement due to the blend shapes

    Parameters
    ----------
    betas : torch.tensor Bx(num_betas)
        Blend shape coefficients
    shape_disps: torch.tensor Vx3x(num_betas)
        Blend shapes

    Returns
    -------
    torch.tensor BxVx3
        The per-vertex displacement due to shape deformation
    '''

    # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
    # i.e. Multiply each shape displacement by its corresponding beta and
    # then sum them.
    # blend_shape = betas, shape_disps
    b = betas.shape[0]
    m, k, l = shape_disps.shape
    shape_disps_r = np.tile(shape_disps.reshape(1, m, k, l), (b, 1, 1, 1))
    betas_r = np.tile(betas.reshape(b, 1, -1, 1), (1, m, 1, 1))
    blend_shape = np.matmul(shape_disps_r, betas_r)

    return blend_shape.reshape((b, m, 3))

def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    ''' Calculates the rotation matrices for a batch of rotation vectors
        Parameters
        ----------
        rot_vecs: torch.tensor Nx3
            array of N axis-angle vectors
        Returns
        -------
        R: torch.tensor Nx3x3
            The rotation matrices for the given axis-angle parameters
    '''

    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)

    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))

    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
    return rot_mat

def rodnumba(p):
    rot, jac = cv2.Rodrigues(p)
    return rot, jac

def batch_rodrigues_np(pose_body):
    ''' Calculates the rotation matrices for a batch of rotation vectors
        Parameters
        ----------
        rot_vecs: ndarray Nx3
            array of N axis-angle vectors
        Returns
        -------
        R: ndarray Nx3x3
            The rotation matrices for the given axis-angle parameters
        R_jac: ndarray Nx3x3x3
            Jacobians of the rotation matrices
    '''
    batch_size = pose_body.shape[0]
    n_j = int(pose_body.shape[1]/3)
    dt = pose_body.dtype
    rot_mats = np.zeros((batch_size, n_j, 3, 3), dtype=dt)
    rot_mats_jac = np.zeros((batch_size, n_j, 3, 3, 3), dtype=dt)
    for b in range(0, pose_body.shape[0]):
        for i in range(0, n_j):
            # rot, jac = cv2.Rodrigues(pose_body[b][3 * (i): 3 * (i + 1)].reshape(-1))
            rot, jac = rodnumba(pose_body[b][3 * (i): 3 * (i + 1)].reshape(-1))
            # print(numba.typeof(rot))
            # print(numba.typeof(jac))
            rot_mats[0, i] = rot
            rot_mats_jac[0, i] = jac.reshape(3, 3, 3)

    return rot_mats, rot_mats_jac

def transform_mat(R, t):
    ''' Creates a batch of transformation matrices
        Args:
            - R: Bx3x3 array of a batch of rotation matrices
            - t: Bx3x1 array of a batch of translation vectors
        Returns:
            - T: Bx4x4 Transformation matrix
    '''
    # No padding left or right, only add an extra row
    return torch.cat([F.pad(R, [0, 0, 0, 1]),
                      F.pad(t, [0, 0, 0, 1], value=1)], dim=2)

def transform_mat_np(R, t):
    ''' Creates a batch of transformation matrices
        Args:
            - R: Bx3x3 array of a batch of rotation matrices
            - t: Bx3x1 array of a batch of translation vectors
        Returns:
            - T: Bx4x4 Transformation matrix
    '''
    # No padding left or right, only add an extra row
    b = R.shape[0]
    T = np.zeros((b, 4, 4), dtype=R.dtype)
    T[:, :3, :3] = R
    T[:, :3, 3] = t.reshape(b, 3)
    T[:, 3, 3] = 1
    return T

def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : torch.tensor BxNx3x3
        Tensor of rotation matrices
    joints : torch.tensor BxNx3
        Locations of joints
    parents : torch.tensor BxN
        The kinematic tree of each object
    dtype : torch.dtype, optional:
        The data type of the created tensors, the default is torch.float32

    Returns
    -------
    posed_joints : torch.tensor BxNx3
        The locations of the joints after applying the pose rotations
    rel_transforms : torch.tensor BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    """

    joints = torch.unsqueeze(joints, dim=-1)

    rel_joints = joints.clone()
    rel_joints[:, 1:] -= joints[:, parents[1:]]

    transforms_mat = transform_mat(
        rot_mats.reshape(-1, 3, 3),
        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)

    # print(transforms_mat[0][0])
    # print(transforms_mat[0][1])

    transform_chain = [transforms_mat[:, 0]]
    for i in range(1, parents.shape[0]):
        # Subtract the joint location at the rest pose
        # No need for rotation, since it's identity when at rest
        curr_res = torch.matmul(transform_chain[parents[i]],
                                transforms_mat[:, i])
        transform_chain.append(curr_res)

    transforms = torch.stack(transform_chain, dim=1)

    # print(transforms[0][1])

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]

    joints_homogen = F.pad(joints, [0, 0, 0, 1])

    rel_transforms = transforms - F.pad(
        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0])

    return posed_joints, rel_transforms

# @jit
def batch_rigid_transform_diff(rot_mats, rot_mats_jac, transform_jac_chain, joints, parents, is_jac=True):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : ndarray BxNx3x3
        Tensor of rotation matrices
    rot_mats_jac : ndarray BxNx3x3x3
        Tensor of rotation matrix Jacobians
    joints : ndarray BxNx3
        Locations of joints
    parents : ndarray BxN
        The kinematic tree of each object

    Returns
    -------
    posed_joints : ndarray BxNx3
        The locations of the joints after applying the pose rotations
    pose_joints_jac : ndarray BxNxNx3x3
        Jacobian of pose_joints w.r.t. pose
    rel_transforms : ndarray BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    rel_transforms_jac : jacobian w.r.t. pose
    """
    # joints = joints.reshape()
    # joints = torch.unsqueeze(joints, dim=-1)

    t_0 = time.time()

    rel_joints = joints.copy()
    rel_joints[:, 1:] -= joints[:, parents[1:]]

    transforms_mat = transform_mat_np(
        rot_mats.reshape(-1, 3, 3),
        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)

    b = joints.shape[0]
    n_j = joints.shape[1]

    t_1 = time.time()

    if is_jac:
        ainds = np.arange(0, parents.shape[0])
        transform_jac_chain.fill(0)
        transform_jac_chain[:, ainds, ainds, :, 0:3, 0:3] = rot_mats_jac[:, ainds]

    transform_chain = np.zeros((b, parents.shape[0], 4, 4), dtype=rot_mats.dtype)
    transform_chain[:, 0] = transforms_mat[:, 0]
    for i in range(1, parents.shape[0]):
        transform_chain[:, i] = np.matmul(transform_chain[:, parents[i]],
                                transforms_mat[:, i])

    if is_jac:
        inds = np.arange(1, len(parents))
        trans_parent = transform_chain[:, parents[inds]].reshape(b, len(inds), 1, 1, 4, 4)
        trans_jac = transform_jac_chain[:, inds, inds, :].reshape(1, len(inds), 1, 3, 4, 4)
        transform_jac_chain[:, inds, inds, :] = np.matmul(trans_parent, trans_jac).reshape(len(inds), 3, 4, 4)

        m = np.eye(parents.shape[0], dtype=np.bool)
        b = transforms_mat.shape[0]
        for i in range(1, parents.shape[0]):
            m[i] = m[i] | m[parents[i]]
            # transform_jac_chain[:, i, :, :] += np.matmul(transform_jac_chain[:, parents[i], :, :], transforms_mat[:, i])
            tr_jac_ch_sel = transform_jac_chain[:, parents[i], m[i], :]
            transform_jac_chain[:, i, m[i], :] += np.matmul(tr_jac_ch_sel, transforms_mat[:, i]).reshape(b, -1, 3, 4, 4)

    t_2 = time.time()

    transforms = transform_chain
    posed_joints = transforms[:, :, :3, 3]
    joints_rot = np.matmul(transforms[:, :, 0:3, 0:3], joints.reshape(b, n_j, 3, 1)).reshape((b, n_j, 3))
    rel_transforms = transforms.copy()
    rel_transforms[:, :, 0:3, 3] = rel_transforms[:, :, 0:3, 3] - joints_rot
    if is_jac:
        transforms_jac = np.transpose(transform_jac_chain, (0, 2, 3, 1, 4, 5))
        posed_joints_jac = transforms_jac[:, :, :, :, :3, 3]
        tjhj = np.matmul(transforms_jac[:, :, :, :, 0:3, 0:3], joints.reshape((b, 1, 1, n_j, 3, 1))).reshape(
            (b, n_j, 3, n_j, 3))
        rel_transforms_jac = transforms_jac.copy()
        rel_transforms_jac[:, :, :, :, 0:3, 3] = rel_transforms_jac[:, :, :, :, 0:3, 3] - tjhj
    else:
        posed_joints_jac = None
        rel_transforms_jac = None
    t_3 = time.time()

    # print('brgd breakdown {} {} {} '.format(t_1-t_0, t_2-t_1, t_3-t_2))

    return posed_joints, posed_joints_jac, rel_transforms, rel_transforms_jac

def batch_rigid_transform_fast_diff(rot_mats, rot_mats_jac, joints, parents):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : ndarray BxNx3x3
        Tensor of rotation matrices
    rot_mats_jac : ndarray BxNx3x3x3
        Tensor of rotation matrix Jacobians
    joints : ndarray BxNx3
        Locations of joints
    parents : ndarray BxN
        The kinematic tree of each object

    Returns
    -------
    posed_joints : ndarray BxNx3
        The locations of the joints after applying the pose rotations
    pose_joints_jac : ndarray BxNxNx3x3
        Jacobian of pose_joints w.r.t. pose
    rel_transforms : ndarray BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    rel_transforms_jac : jacobian w.r.t. pose
    """
    # joints = joints.reshape()
    # joints = torch.unsqueeze(joints, dim=-1)

    t_0 = time.time()

    rel_joints = joints.copy()
    rel_joints[:, 1:] -= joints[:, parents[1:]]

    transforms_mat = transform_mat_np(
        rot_mats.reshape(-1, 3, 3),
        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)

    # print(transforms_mat[0][0])
    # print(transforms_mat[0][1])

    b = joints.shape[0]
    n_j = joints.shape[1]

    t_1 = time.time()

    # transform_jac = np.zeros((b, n_j, 3, 4, 4))
    # transform_jac[:, 0, :, 0:3, 0:3] = rot_mats_jac[:, 0]
    transform_jac_chain = np.zeros((b, n_j, 3, parents.shape[0], 4, 4))
    for i in range(0, parents.shape[0]):
        transform_jac_chain[:, i, :, i, 0:3, 0:3] = rot_mats_jac[:, i]

    transform_chain = np.copy(transforms_mat)
    for i in range(1, parents.shape[0]):
        t_curr = np.matmul(transform_chain[:, parents[i], 0:3, 0:3],
                                transforms_mat[:, i, 0:3, 3].reshape(-1, 1, 3, 1)).reshape(-1, 1, 3)
        transform_chain[:, i, 0:3, 3] = t_curr + transform_chain[:, parents[i], 0:3, 3]
        t_i = np.tile(transforms_mat[:, i, 0:3, 3].reshape(-1, 1, 1, 3, 1), (1, n_j, 3, 1, 1))
        # print(transform_jac_chain[:, :, :, parents[i], 0:3, 0:3].shape)
        # print(t_i.shape)
        t_jac_curr = np.matmul(transform_jac_chain[:, :, :, parents[i], 0:3, 0:3], t_i)
        transform_jac_chain[:, :, :, i, 0:3, 3] = transform_jac_chain[:, :, :, parents[i], 0:3, 3] + \
                                                  t_jac_curr.reshape(-1, n_j, 3, 3)

    # transforms = np.stack(transform_chain, axis=1)
    transforms = transform_chain
    # transforms_jac = np.stack(transform_jac_chain, axis=3)
    transforms_jac = transform_jac_chain

    t_2 = time.time()
    # print(transforms[0][1])

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]
    posed_joints_jac = transforms_jac[:, :, :, :, :3, 3]

    # joints_homogen = np.zeros((b, n_j, 4, 1))
    # joints_homogen[:, :, 0:3, 0] = joints

    joints_rot = np.matmul(transforms[:,:,0:3,0:3], joints.reshape(b, n_j, 3, 1)).reshape((b, n_j, 3))
    # jht = np.pad(np.matmul(transforms, joints_homogen), [(0, 0), (0, 0), (0, 0), (3,0)])
    # rel_transforms = transforms - jht
    rel_transforms = transforms.copy()
    rel_transforms[:, :, 0:3, 3] = rel_transforms[:, :, 0:3, 3] - joints_rot

    # jhtj = np.matmul(transforms_jac, joints_homogen.reshape((b, 1, 1, n_j, 4, 1)))
    # jhtj = np.pad(jhtj, [(0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (3, 0)])

    tjhj = np.matmul(transforms_jac[:,:,:,:,0:3,0:3], joints.reshape((b, 1, 1, n_j, 3, 1))).reshape((b, n_j, 3, n_j, 3))

    rel_transforms_jac = transforms_jac.copy()
    rel_transforms_jac[:, :, :, :, 0:3, 3] = rel_transforms_jac[:, :, :, :, 0:3, 3] - tjhj
    # rel_transforms_jac = transforms_jac - jhtj

    t_3 = time.time()

    print('brgd breakdown {} {} {} '.format(t_1-t_0, t_2-t_1, t_3-t_2))

    return posed_joints, posed_joints_jac, rel_transforms, rel_transforms_jac

def prepare_J(betas, v_template, shapedirs, J_regressor, n_v):
    # Add shape contribution
    v_shaped = v_template + blend_shapes_np(betas, shapedirs)

    # Get the joints
    # NxJx3 array
    J = vertices2joints_np(J_regressor, v_shaped)

    n_j = J_regressor.shape[0]

    batch_size = betas.shape[0]

    homogen_coord = np.ones((batch_size, n_v, 1), dtype=v_template.dtype)

    transform_jac_chain = np.zeros((batch_size, n_j, n_j, 3, 4, 4), dtype=v_template.dtype)

    return J, v_shaped, homogen_coord, transform_jac_chain

def rel_to_direct(pose, parents):

    rot_mats, rot_mat_jacs = batch_rodrigues_np(pose)

    b = pose.shape[0]

    rot_chain = np.zeros((b, parents.shape[0], 3, 3))
    rot_chain[:, 0] = rot_mats[:, 0]
    for i in range(1, parents.shape[0]):
        rot_chain[:, i] = np.matmul(rot_chain[:, parents[i]], rot_mats[:, i])

    pose_dir = np.zeros_like(pose)
    n_j = int(pose.shape[1]/3)
    for b in range(0, pose.shape[0]):
        for i in range(0, n_j):
            rv, jac = cv2.Rodrigues(rot_chain[b, i])
            pose_dir[b, 3*i : 3*(i+1)] = rv.reshape(-1)

    return pose_dir

def lbs_diff_fast(pose, parents,
        J, v_shaped, W, W_j, homogen_coord, v_inds=None):
    t_0 = time.time()

    batch_size = pose.shape[0]

    # print('j1:')
    # print(J[0][1])

    # 3. Add pose blend shapes
    # N x J x 3 x 3

    t_1 = time.time()

    rot_mats, rot_mat_jacs = batch_rodrigues_np(pose)
    # rot_mats = rot_mats.reshape((batch_size, -1, 3, 3))
    n_j = rot_mats.shape[1]

    if v_inds is not None:
        v_shaped = v_shaped[:, v_inds, :]

    n_v = v_shaped.shape[1]

    v_posed = v_shaped

    t_2 = time.time()

    J_transformed, J_transformed_jac, A, A_jac = batch_rigid_transform_fast_diff(rot_mats, rot_mat_jacs, J, parents)

    t_3 = time.time()

    t_3 = time.time()

    # 5. Do skinning:
    # W is N x V x (J + 1)
    # W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    # (N x V x (J + 1)) x (N x (J + 1) x 16) = N x V x 16
    num_joints = n_j
    T = np.matmul(W, A.reshape(batch_size, num_joints, 16)) \
        .reshape((batch_size, -1, 4, 4))

    # W_j = np.tile(W.reshape((batch_size, 1, 1, n_v, n_j)), (1, n_j, 3, 1, 1))

    A_jact = A_jac  # .transpose(0, 2, 3, 1, 4, 5)
    T_jac = np.matmul(W_j, A_jact.reshape(batch_size, n_j, 3, n_j, -1))
    T_jac = T_jac.reshape((batch_size, n_j, 3, n_v, 4, 4))
    # N x n_j x 3 x V x 16

    v_posed_homo = np.concatenate([v_posed, homogen_coord], axis=2)

    v_homo = np.matmul(T, v_posed_homo.reshape((batch_size, n_v, 4, 1)))

    T_j = T.reshape((batch_size, 1, 1, n_v, 4, 4))

    v_posed_homo_j = v_posed_homo.reshape((batch_size, 1, 1, n_v, 4, 1))
    v_homo_jac2 = np.matmul(T_jac, v_posed_homo_j)

    verts = v_homo[:, :, :3, 0]

    v_homo_jac = v_homo_jac2[:, :, :, :, :3, :]

    verts_jac = v_homo_jac[:, :, :, :, :3, 0]

    t_4 = time.time()

    # print('breakdown b {} a {} rt {} f {}'.format(t_1-t_0, t_2-t_1, t_3-t_2, t_4-t_3))

    return verts, verts_jac, J_transformed, J_transformed_jac, A, A_jac, J  # , v_posed, v_posed_jac

# @jit
def lbs_diff(pose, posedirs, parents, J, v_shaped, lbs_weights, homogen_coord, transform_jac_chain, v_inds=None):
    ''' Performs Linear Blend Skinning with the given shape and pose parameters

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        pose : ndarray Bx(J + 1) * 3
            The pose parameters in axis-angle format
        v_template ndarray BxVx3
            The template mesh that will be deformed
        shapedirs : torch.tensor 1xNB
            The tensor of PCA shape displacements
        posedirs : ndarray Px(V * 3)
            The pose PCA coefficients
        J_regressor : ndarray JxV
            The regressor array that is used to calculate the joints from
            the position of the vertices
        parents: ndarray J
            The array that describes the kinematic tree for the model
        lbs_weights: ndarray N x V x (J + 1)
            The linear blend skinning weights that represent how much the
            rotation matrix of each part affects each vertex
        vinds: ndarray - list of required vertex indices (if None, all vertices will be processed)

        Returns
        -------
        verts: ndarray BxVx3
            The vertices of the mesh after applying the shape and pose
            displacements.
        verts_jac: ndarray BxVx(J+1)x3x3
        joints: ndarray BxJx3
            The joints of the model
        joints_jac: ndarray BxJxJx3x3
            Jacobian of joints' coordinates
    '''

    t_0 = time.time()

    batch_size = pose.shape[0]

    if v_inds is not None:
        lbs_weights = lbs_weights[v_inds]
        n_v = len(v_inds)
    else:
        n_v = 0

    n_j = J.shape[1]

    W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    # W_j = np.tile(W.reshape((batch_size, 1, 1, n_v, n_j)), (1, n_j, 3, 1, 1))
    W_j = W.reshape((batch_size, 1, 1, n_v, n_j))

    # print('j1:')
    # print(J[0][1])

    # 3. Add pose blend shapes
    # N x J x 3 x 3

    t_1 = time.time()

    rot_mats, rot_mat_jacs = batch_rodrigues_np(pose)
    # rot_mats = rot_mats.reshape((batch_size, -1, 3, 3))
    n_j = rot_mats.shape[1]
    pose_feature = (rot_mats[:, 1:, :, :] - np.tile(np.eye(3).reshape(1, 1, 3, 3), (batch_size, n_j-1, 1, 1)))\
        .reshape((batch_size, -1))

    if v_inds is not None:
        v_shaped = v_shaped[:, v_inds, :]
        inds = np.stack([3*v_inds, 3*v_inds+1, 3*v_inds+2], axis=1).reshape(-1)
        posedirs = posedirs[:, inds]
        # lbs_weights = lbs_weights[v_inds]

        # (N x P) x (P, V * 3) -> N x V x 3
    pose_offsets = np.matmul(pose_feature, posedirs) \
        .reshape((batch_size, -1, 3))

    n_v = v_shaped.shape[1]

    pose_offset_jacs = np.zeros((batch_size, n_j, 3, n_v, 3), dtype=pose.dtype)
    for i in range (1, n_j):
        pdi = np.matmul(rot_mat_jacs[:, i].reshape(-1, 9), posedirs[(i-1)*9 : i*9, :]).reshape((batch_size, 3, -1, 3))
        pose_offset_jacs[:, i] = pdi

    v_posed = pose_offsets + v_shaped

    #NxVxJ+1x3x3
    v_posed_jac = pose_offset_jacs

    # print(rot_mats[0][1])

    t_2 = time.time()

    J_transformed, J_transformed_jac, A, A_jac = batch_rigid_transform_diff(rot_mats, rot_mat_jacs, transform_jac_chain, J, parents)

    t_3 = time.time()

    # 5. Do skinning:
    # W is N x V x (J + 1)
    # W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    # (N x V x (J + 1)) x (N x (J + 1) x 16) = N x V x 16
    num_joints = n_j
    T = np.matmul(W, A.reshape(batch_size, num_joints, 16)) \
        .reshape((batch_size, -1, 4, 4))

    # W_j = np.tile(W.reshape((batch_size, 1, 1, n_v, n_j)), (1, n_j, 3, 1, 1))

    A_jact = A_jac #.transpose(0, 2, 3, 1, 4, 5)
    T_jac = np.matmul(W_j, A_jact.reshape(batch_size, n_j, 3, n_j, -1))
    T_jac = T_jac.reshape((batch_size, n_j, 3, n_v, 4, 4))
    #N x n_j x 3 x V x 16

    v_posed_homo = np.concatenate([v_posed, homogen_coord], axis=2)

    v_homo = np.matmul(T, v_posed_homo.reshape((batch_size, n_v, 4, 1)))

    # T_j = np.tile(T.reshape((batch_size, 1, 1, n_v, 4, 4)), (1, n_j, 3, 1, 1, 1))

    T_j = T.reshape((batch_size, 1, 1, n_v, 4, 4))

    # v_posed_jac_h = np.pad(v_posed_jac, ((0, 0), (0,0), (0,0), (0,0), (0,1))).reshape((batch_size, n_j, 3, n_v, 4, 1))
    # v_posed_jac_h = v_posed_jac_h.reshape((batch_size, n_j, 3, n_v, 4, 1))
    # v_homo_jac1 = np.matmul(T_j, v_posed_jac_h)

    v_homo_jac1 = np.matmul(T_j[:, :, :, :, 0:3, 0:3], v_posed_jac.reshape((batch_size, n_j, 3, n_v, 3, 1)))

    # v_posed_homo_j = np.tile(v_posed_homo.reshape((batch_size, 1, 1, n_v, 4, 1)), (1, n_j, 3, 1, 1, 1))
    v_posed_homo_j = v_posed_homo.reshape((batch_size, 1, 1, n_v, 4, 1))
    v_homo_jac2 = np.matmul(T_jac, v_posed_homo_j)
    # v_homo_jac2 = v_homo_jac2.transpose((0, 3, 1, 2, 4, 5))

    verts = v_homo[:, :, :3, 0]

    v_homo_jac = v_homo_jac1 + v_homo_jac2[:, :, :, :,  :3, :]

    verts_jac = v_homo_jac[:, :, :, :, :3, 0]

    t_4 = time.time()

    # print('breakdown b {} a {} rt {} f {}'.format(t_1-t_0, t_2-t_1, t_3-t_2, t_4-t_3))

    return verts, verts_jac, J_transformed, J_transformed_jac, A, A_jac, J #, v_posed, v_posed_jac

def lbs_diff_nopd(pose, posedirs_face, parents, J, v_shaped, lbs_weights, homogen_coord, transform_jac_chain, v_inds=None,
                  is_jac=True, bpm=None, face_expression=None, w_thr=0):
    ''' Performs Linear Blend Skinning with the given shape and pose parameters

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        pose : ndarray Bx(J + 1) * 3
            The pose parameters in axis-angle format
        v_template ndarray BxVx3
            The template mesh that will be deformed
        shapedirs : torch.tensor 1xNB
            The tensor of PCA shape displacements
        posedirs : ndarray Px(V * 3)
            The pose PCA coefficients
        J_regressor : ndarray JxV
            The regressor array that is used to calculate the joints from
            the position of the vertices
        parents: ndarray J
            The array that describes the kinematic tree for the model
        lbs_weights: ndarray N x V x (J + 1)
            The linear blend skinning weights that represent how much the
            rotation matrix of each part affects each vertex
        vinds: ndarray - list of required vertex indices (if None, all vertices will be processed)

        Returns
        -------
        verts: ndarray BxVx3
            The vertices of the mesh after applying the shape and pose
            displacements.
        verts_jac: ndarray BxVx(J+1)x3x3
        joints: ndarray BxJx3
            The joints of the model
        joints_jac: ndarray BxJxJx3x3
            Jacobian of joints' coordinates
    '''

    t_0 = time.time()

    batch_size = pose.shape[0]
    if v_inds is not None:

        if bpm is None:
            w_arr = np.sum(lbs_weights[v_inds] ** 2, axis=0)/len(v_inds)
            bpm = w_arr > w_thr

        lbs_weights = lbs_weights[v_inds]

    # 3. Add pose blend shapes
    # N x J x 3 x 3

    t_1 = time.time()

    rot_mats, rot_mat_jacs = batch_rodrigues_np(pose)

    n_j = rot_mats.shape[1]

    if v_inds is not None:
        v_shaped = v_shaped[:, v_inds, :]

        if face_expression is not None:
            posedirs_face = posedirs_face[v_inds]
            n_v = len(v_inds)
            dv = np.matmul(posedirs_face.reshape(1, n_v, 3, 1, -1), np.tile(face_expression.reshape(1, 1, 1, -1, 1),
                                                                            (1, n_v, 3, 1, 1)))
            v_shaped += dv.reshape(1, -1, 3)
            dv_jac = np.transpose(posedirs_face.reshape(n_v, 3, -1), (2, 0, 1)).reshape(-1, 3 * n_v)

    n_v = v_shaped.shape[1]

    v_posed = v_shaped

    t_2 = time.time()

    J_transformed, J_transformed_jac, A, A_jac = batch_rigid_transform_diff(rot_mats, rot_mat_jacs, transform_jac_chain, J, parents, is_jac)

    t_3 = time.time()

    # 5. Do skinning:
    # W is N x V x (J + 1)
    # W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    # (N x V x (J + 1)) x (N x (J + 1) x 16) = N x V x 16
    num_joints = n_j
    W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    T = np.matmul(W, A.reshape(batch_size, num_joints, 16)) \
        .reshape((batch_size, -1, 4, 4))
    v_posed_homo = np.concatenate([v_posed, homogen_coord], axis=2)
    v_homo = np.matmul(T, v_posed_homo.reshape((batch_size, n_v, 4, 1)))
    verts = v_homo[:, :, :3, 0]

    t_3_1 = time.time()

    if is_jac:

        Wred = W[:, :, bpm]
        n_bpm = Wred.shape[-1]
        W_j = Wred.reshape((batch_size, 1, 1, n_v, n_bpm))
        T_jac = np.matmul(W_j, A_jac.reshape(batch_size, n_j, 3, n_j, -1)[:, :, :, bpm])
        T_jac = T_jac.reshape((batch_size, n_j, 3, n_v, 4, 4))

        #N x n_j x 3 x V x 16

        t_3_2 = time.time()

        v_posed_homo_j = v_posed_homo.reshape((batch_size, 1, 1, n_v, 4, 1))
        v_homo_jac2 = np.matmul(T_jac, v_posed_homo_j)
        n_pars = pose.shape[1]
        verts_jac = v_homo_jac2[:, :, :, :,  :3, 0].reshape(n_pars, -1)
        if face_expression is not None:
            n_fp = 10
            dv_jac = np.matmul(T[:, :, 0:3, 0:3], dv_jac.reshape(n_fp, -1, 3, 1)).reshape(n_fp, -1)
            # dv_jac = np.zeros_like(dv_jac)
            verts_jac = np.concatenate([verts_jac, dv_jac], axis=0)
    else:
        verts_jac = None

    t_4 = time.time()

    # print('breakdown b {} a {} rt {} f1 {} f2 {} f3 {}'.format(t_1-t_0, t_2-t_1, t_3-t_2, t_3_1 - t_3, t_3_2 - t_3_1, t_4-t_3_2))

    return verts, verts_jac, J_transformed, J_transformed_jac, A, A_jac, J, bpm #, v_posed, v_posed_jac


# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This code is refer from:
https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/bbox/assigners/sim_ota_assigner.py
"""

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

from ppdet.modeling.losses.varifocal_loss import varifocal_loss
from ppdet.modeling.bbox_utils import batch_bbox_overlaps
from ppdet.core.workspace import register

@register
class SimOTAAssigner(object):
    """Computes matching between predictions and ground truth.

    Args:
        center_radius (int | float, optional): Ground truth center size
            to judge whether a prior is in center. Default 2.5.
        candidate_topk (int, optional): The candidate top-k which used to
            get top-k ious to calculate dynamic-k. Default 10.
        iou_weight (int | float, optional): The scale factor for regression
            iou cost. Default 3.0.
        cls_weight (int | float, optional): The scale factor for classification
            cost. Default 1.0.
        num_classes (int): The num_classes of dataset.
        use_vfl (int): Whether to use varifocal_loss when calculating the cost matrix.
    """
    __shared__ = ['num_classes']

    def __init__(self,
                 center_radius=2.5,
                 candidate_topk=10,
                 iou_weight=3.0,
                 cls_weight=1.0,
                 num_classes=80,
                 use_vfl=True):
        self.center_radius = center_radius
        self.candidate_topk = candidate_topk
        self.iou_weight = iou_weight
        self.cls_weight = cls_weight
        self.num_classes = num_classes
        self.use_vfl = use_vfl

    def get_in_gt_and_in_center_info(self, priors, gt_bboxes):
        num_gt = gt_bboxes.shape[0]

        repeated_x = priors[:, 0].unsqueeze(1).tile([1, num_gt])
        repeated_y = priors[:, 1].unsqueeze(1).tile([1, num_gt])
        repeated_stride_x = priors[:, 2].unsqueeze(1).tile([1, num_gt])
        repeated_stride_y = priors[:, 3].unsqueeze(1).tile([1, num_gt])

        # is prior centers in gt bboxes, shape: [n_prior, n_gt]
        l_ = repeated_x - gt_bboxes[:, 0]
        t_ = repeated_y - gt_bboxes[:, 1]
        r_ = gt_bboxes[:, 2] - repeated_x
        b_ = gt_bboxes[:, 3] - repeated_y

        deltas = paddle.stack([l_, t_, r_, b_], axis=1)
        is_in_gts = deltas.min(axis=1) > 0
        is_in_gts_all = is_in_gts.sum(axis=1) > 0

        # is prior centers in gt centers
        gt_cxs = (gt_bboxes[:, 0] + gt_bboxes[:, 2]) / 2.0
        gt_cys = (gt_bboxes[:, 1] + gt_bboxes[:, 3]) / 2.0
        ct_box_l = gt_cxs - self.center_radius * repeated_stride_x
        ct_box_t = gt_cys - self.center_radius * repeated_stride_y
        ct_box_r = gt_cxs + self.center_radius * repeated_stride_x
        ct_box_b = gt_cys + self.center_radius * repeated_stride_y

        cl_ = repeated_x - ct_box_l
        ct_ = repeated_y - ct_box_t
        cr_ = ct_box_r - repeated_x
        cb_ = ct_box_b - repeated_y

        ct_deltas = paddle.stack([cl_, ct_, cr_, cb_], axis=1)
        is_in_cts = ct_deltas.min(axis=1) > 0
        is_in_cts_all = is_in_cts.sum(axis=1) > 0

        # in boxes or in centers, shape: [num_priors]
        is_in_gts_or_centers = paddle.logical_or(is_in_gts_all, is_in_cts_all)

        is_in_gts_or_centers_inds = paddle.nonzero(
            is_in_gts_or_centers).squeeze(1)

        # both in boxes and centers, shape: [num_fg, num_gt]
        is_in_boxes_and_centers = paddle.logical_and(
            paddle.gather(
                is_in_gts.cast('int'), is_in_gts_or_centers_inds,
                axis=0).cast('bool'),
            paddle.gather(
                is_in_cts.cast('int'), is_in_gts_or_centers_inds,
                axis=0).cast('bool'))
        return is_in_gts_or_centers, is_in_boxes_and_centers

    def dynamic_k_matching(self, cost, pairwise_ious, num_gt, valid_mask):
        matching_matrix = np.zeros_like(cost.numpy())
        # select candidate topk ious for dynamic-k calculation
        topk_ious, _ = paddle.topk(pairwise_ious, self.candidate_topk, axis=0)
        # calculate dynamic k for each gt
        dynamic_ks = paddle.clip(topk_ious.sum(0).cast('int'), min=1)
        for gt_idx in range(num_gt):
            _, pos_idx = paddle.topk(
                cost[:, gt_idx], k=dynamic_ks[gt_idx], largest=False)
            matching_matrix[:, gt_idx][pos_idx.numpy()] = 1.0

        del topk_ious, dynamic_ks, pos_idx

        prior_match_gt_mask = matching_matrix.sum(1) > 1
        if prior_match_gt_mask.sum() > 0:
            cost = cost.numpy()
            cost_argmin = np.argmin(cost[prior_match_gt_mask, :], axis=1)
            matching_matrix[prior_match_gt_mask, :] *= 0.0
            matching_matrix[prior_match_gt_mask, cost_argmin] = 1.0
        # get foreground mask inside box and center prior
        fg_mask_inboxes = matching_matrix.sum(1) > 0.0
        valid_mask[valid_mask.copy()] = fg_mask_inboxes

        matched_gt_inds = matching_matrix[fg_mask_inboxes, :].argmax(1)

        matched_gt_inds = paddle.to_tensor(
            matched_gt_inds, place=pairwise_ious.place)

        return matched_gt_inds, valid_mask

    def get_sample(self, assign_gt_inds, gt_bboxes):
        pos_inds = np.unique(np.nonzero(assign_gt_inds > 0)[0])
        neg_inds = np.unique(np.nonzero(assign_gt_inds == 0)[0])
        pos_assigned_gt_inds = assign_gt_inds[pos_inds] - 1

        if gt_bboxes.size == 0:
            # hack for index error case
            assert pos_assigned_gt_inds.size == 0
            pos_gt_bboxes = np.empty_like(gt_bboxes).reshape(-1, 4)
        else:
            if len(gt_bboxes.shape) < 2:
                gt_bboxes = gt_bboxes.resize(-1, 4)
            pos_gt_bboxes = gt_bboxes[pos_assigned_gt_inds, :]
        return pos_inds, neg_inds, pos_gt_bboxes, pos_assigned_gt_inds

    def __call__(self,
                 pred_scores,
                 priors,
                 decoded_bboxes,
                 gt_bboxes,
                 gt_labels,
                 eps=1e-7):
        """Assign gt to priors using SimOTA.
        TODO: add comment.
        Returns:
            assign_result: The assigned result.
        """

        INF = 100000000
        num_gt = gt_bboxes.shape[0]
        num_bboxes = decoded_bboxes.shape[0]

        # assign 0 by default
        assigned_gt_inds = paddle.full(
            (num_bboxes, ), 0, dtype=paddle.int64).numpy()
        if num_gt == 0 or num_bboxes == 0:
            # No ground truth or boxes, return empty assignment
            priors = priors.numpy()
            labels = np.ones([num_bboxes], dtype=np.int64) * self.num_classes
            label_weights = np.ones([num_bboxes], dtype=np.float32)
            bbox_targets = np.zeros_like(priors)
            return priors, labels, label_weights, bbox_targets, 0

        valid_mask, is_in_boxes_and_center = self.get_in_gt_and_in_center_info(
            priors, gt_bboxes)

        valid_mask_inds = paddle.nonzero(valid_mask).squeeze(1)
        valid_decoded_bbox = decoded_bboxes[valid_mask_inds]
        valid_pred_scores = pred_scores[valid_mask_inds]
        num_valid = valid_decoded_bbox.shape[0]

        pairwise_ious = batch_bbox_overlaps(valid_decoded_bbox, gt_bboxes)
        if self.use_vfl:
            gt_vfl_labels = gt_labels.squeeze(-1).unsqueeze(0).tile(
                [num_valid, 1]).reshape([-1])
            valid_pred_scores = valid_pred_scores.unsqueeze(1).tile(
                [1, num_gt, 1]).reshape([-1, self.num_classes])
            vfl_score = np.zeros(valid_pred_scores.shape)
            vfl_score[np.arange(0, vfl_score.shape[0]), gt_vfl_labels.numpy(
            )] = pairwise_ious.reshape([-1])
            vfl_score = paddle.to_tensor(vfl_score)
            losses_vfl = varifocal_loss(
                valid_pred_scores, vfl_score,
                use_sigmoid=False).reshape([num_valid, num_gt])
            losses_giou = batch_bbox_overlaps(
                valid_decoded_bbox, gt_bboxes, mode='giou')
            cost_matrix = (
                losses_vfl * self.cls_weight + losses_giou * self.iou_weight +
                paddle.logical_not(is_in_boxes_and_center).cast('float32') * INF
            )
        else:
            iou_cost = -paddle.log(pairwise_ious + eps)
            gt_onehot_label = (F.one_hot(
                gt_labels.squeeze(-1).cast(paddle.int64),
                pred_scores.shape[-1]).cast('float32').unsqueeze(0).tile(
                    [num_valid, 1, 1]))

            valid_pred_scores = valid_pred_scores.unsqueeze(1).tile(
                [1, num_gt, 1])
            cls_cost = F.binary_cross_entropy(
                valid_pred_scores, gt_onehot_label, reduction='none').sum(-1)

            cost_matrix = (
                cls_cost * self.cls_weight + iou_cost * self.iou_weight +
                paddle.logical_not(is_in_boxes_and_center).cast('float32') * INF
            )

        matched_gt_inds, valid_mask = \
            self.dynamic_k_matching(
                cost_matrix, pairwise_ious, num_gt, valid_mask.numpy())

        # assign results
        gt_labels = gt_labels.numpy()
        priors = priors.numpy()
        matched_gt_inds = matched_gt_inds.numpy()
        gt_bboxes = gt_bboxes.numpy()

        assigned_gt_inds[valid_mask] = matched_gt_inds + 1

        pos_inds, neg_inds, pos_gt_bboxes, pos_assigned_gt_inds \
            = self.get_sample(assigned_gt_inds, gt_bboxes)

        num_cells = priors.shape[0]
        bbox_targets = np.zeros_like(priors)
        bbox_weights = np.zeros_like(priors)
        labels = np.ones([num_cells], dtype=np.int64) * self.num_classes
        label_weights = np.zeros([num_cells], dtype=np.float32)

        if len(pos_inds) > 0:
            pos_bbox_targets = pos_gt_bboxes
            bbox_targets[pos_inds, :] = pos_bbox_targets
            bbox_weights[pos_inds, :] = 1.0
            if not np.any(gt_labels):
                labels[pos_inds] = 0
            else:
                labels[pos_inds] = gt_labels.squeeze(-1)[pos_assigned_gt_inds]

            label_weights[pos_inds] = 1.0
        if len(neg_inds) > 0:
            label_weights[neg_inds] = 1.0

        pos_num = max(pos_inds.size, 1)

        return priors, labels, label_weights, bbox_targets, pos_num