text: # Copyright (c) Hao Meng. All Rights Reserved.
  import time
  
  import numpy as np
  import torch
  from manopth.manolayer import ManoLayer
  
  from utils import bone
  
  
  class LM_Solver():
      def __init__(self, side = "right",num_Iter=500, th_beta=None, th_pose=None, lb_target=None,
                   weight=0.01):
          self.count = 0
          self.time_start = time.time()
          self.time_in_mano = 0
          self.minimal_loss = 9999
          self.best_beta = np.zeros([10, 1])
          self.num_Iter = num_Iter
  
          self.th_beta = th_beta
          self.th_pose = th_pose
  
          self.beta = th_beta.numpy()
          self.pose = th_pose.numpy()
  
          self.mano_layer = ManoLayer(side=side,
                                      mano_root='D:/code/manopth/mano/models', use_pca=False, flat_hand_mean=True)
  
          self.threshold_stop = 10 ** -13
          self.weight = weight
          self.residual_memory = []
  
          self.lb = np.zeros(21)
  
          _, self.joints = self.mano_layer(self.th_pose, self.th_beta)
          self.joints = self.joints.cpu().numpy().reshape(21, 3)
  
          self.lb_target = lb_target.reshape(15, 1)
          self.test_time = 0
  
      def update_target(self, target):
          self.lb_target = target.copy().reshape(15, 1)
  
      def update(self, beta_):
          beta = beta_.copy()
          self.count += 1
          now = time.time()
          my_th_beta = torch.from_numpy(beta).float().reshape(1, 10)
          _, joints = self.mano_layer(self.th_pose, my_th_beta)
          self.time_in_mano = time.time() - now
  
          useful_lb = bone.caculate_length(joints, label="useful")
          lb_ref = useful_lb[6]
          return useful_lb, lb_ref
  
      def new_cal_ref_bone(self, _shape):
          now = time.time()
          parent_index = [0,
                          0, 1, 2,
                          0, 4, 5,
                          0, 7, 8,
                          0, 10, 11,
                          0, 13, 14
                          ]
          index = [0,
                   1, 2, 3,  # index
                   4, 5, 6,  # middle
                   7, 8, 9,  # pinky
                   10, 11, 12,  # ring
                   13, 14, 15]  # thumb
          reoder_index = [
              13, 14, 15,
              1, 2, 3,
              4, 5, 6,
              10, 11, 12,
              7, 8, 9]
          shape = torch.Tensor(_shape.reshape((-1, 10)))
          th_v_shaped = torch.matmul(self.mano_layer.th_shapedirs,
                                     shape.transpose(1, 0)).permute(2, 0, 1) \
                        + self.mano_layer.th_v_template
          th_j = torch.matmul(self.mano_layer.th_J_regressor, th_v_shaped)
          temp1 = th_j.clone().detach()
          temp2 = th_j.clone().detach()[:, parent_index, :]
          result = temp1 - temp2
          result = torch.norm(result, dim=-1, keepdim=True)
          ref_len = result[:, [4]]
          result = result / ref_len
          self.time_in_mano = time.time() - now
          return torch.squeeze(result, dim=-1)[:, reoder_index].cpu().numpy()
  
      def get_residual(self, beta_):
          beta = beta_.copy()
          lb, lb_ref = self.update(beta)
          lb = lb.reshape(45, 1)
          return lb / lb_ref - self.lb_target
  
      def get_count(self):
          return self.count
  
      def get_bones(self, beta_):
          beta = beta_.copy()
          lb, _ = self.update(beta)
          lb = lb.reshape(15, 1)
          return lb
  
      def get_loss(self, beta_):
  
          beta = beta_.copy()
  
          lb, lb_ref = self.update(beta)
          lb = lb.reshape(15, 1)
  
          loss = np.linalg.norm(lb / lb_ref - self.lb_target) ** 2 + \
                 self.weight * np.linalg.norm(beta) ** 2
  
          return loss
  
      def new_get_loss(self, beta_):
          beta = beta_.copy()
          temp = self.new_cal_ref_bone(beta_)
          loss = temp.reshape((15, 1))
          loss = np.linalg.norm(loss - self.lb_target) ** 2 + \
                 self.weight * np.linalg.norm(beta_)
          return loss
  
      def get_derivative(self, beta_, n):
  
          beta = beta_.copy()
          params1 = np.array(beta)
          params2 = np.array(beta)
          step = 0.01
          params1[n] += step
          params2[n] -= step
  
          res1 = self.new_get_loss(params1)
          res2 = self.new_get_loss(params2)
  
          d = (res1 - res2) / (2 * step)
  
          return d.ravel()
  
      def batch_new_get_loss(self, beta_):
          weight = 1e-5
          beta = beta_.copy()
          temp = self.new_cal_ref_bone(beta)
          loss = np.transpose(temp)
          loss = np.linalg.norm(loss - self.lb_target, axis=0) ** 2 + \
                 weight * np.linalg.norm(beta, axis=-1)
          return loss
  
      def new_get_derivative(self, beta_):
          # params: beta_ 10*1
          # return: 1*10
          beta = beta_.copy().reshape((1, 10))
          temp_shape = np.zeros((20, beta.shape[1]))  # 20*10
          step = 0.01
          for t2 in range(10):  # 位置
              t3 = 10 + t2
              temp_shape[t2] = beta.copy()
              temp_shape[t3] = beta.copy()
              temp_shape[t2, t2] += step
              temp_shape[t3, t2] -= step
  
          res = self.batch_new_get_loss(temp_shape)
          d = res[0:10] - res[10:20]  # 10*1
          d = d.reshape((1, 10)) / (2 * step)
          return d
  
      # LM algorithm
      def LM(self):
          u = 1e-2
          v = 1.5
          beta = self.beta.reshape(10, 1)
  
          out_n = 1
          num_beta = np.shape(beta)[0]  # the number of beta
          # calculating the init Jocobian matrix
          Jacobian = np.zeros([out_n, beta.shape[0]])
  
          last_update = 0
          last_loss = 0
          self.test_time = 0
          for i in range(self.num_Iter):
              loss = self.new_get_loss(beta)
              if loss < self.minimal_loss:
                  self.minimal_loss = loss
                  self.best_beta = beta
  
              if abs(loss - last_loss) < self.threshold_stop:
                  self.time_total = time.time() - self.time_start
                  return beta
  
              # for k in range(num_beta):
              #     Jacobian[:, k] = self.get_derivative(beta, k)
              Jacobian = self.new_get_derivative(beta)
              jtj = np.matmul(Jacobian.T, Jacobian)
              jtj = jtj + u * np.eye(jtj.shape[0])
  
              update = last_loss - loss
              delta = (np.matmul(np.linalg.inv(jtj), Jacobian.T) * loss)
  
              beta -= delta
  
              if update > last_update and update > 0:
                  u /= v
              else:
                  u *= v
  
              last_update = update
              last_loss = loss
              self.residual_memory.append(loss)
  
          return beta
  
      def get_result(self):
          return self.residual_memory
  