EasyMocap/easymocap/multistage/fitting.py

'''
  @ Date: 2022-03-22 16:11:44
  @ Author: Qing Shuai
  @ Mail: s_q@zju.edu.cn
  @ LastEditors: Qing Shuai
  @ LastEditTime: 2022-07-25 11:51:50
  @ FilePath: /EasyMocapPublic/easymocap/multistage/fitting.py
'''
# This function provides a realtime fitting interface
from collections import namedtuple
from time import time, sleep
import numpy as np
import cv2
import torch
import copy

from ..config.baseconfig import load_object_from_cmd
from ..mytools.debug_utils import log, mywarn
from ..mytools import Timer
from ..config import Config
from ..mytools.triangulator import iterative_triangulate
from ..bodymodel.base import Params
from .torchgeometry import axis_angle_to_euler, euler_to_axis_angle

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

from scipy.spatial.transform import Rotation
def aa2euler(aa):
    aa = np.array(aa)
    R = cv2.Rodrigues(aa)[0]
    # res = Rotation.from_dcm(R).as_euler('XYZ', degrees=True)
    res = Rotation.from_matrix(R).as_euler('XYZ', degrees=False)
    return np.round(res, 2).tolist()

def rotmat2euler(rot):
    res = Rotation.from_matrix(rot).as_euler('XYZ', degrees=True)
    return res

def euler2rotmat(euler):
    res = Rotation.from_euler('XYZ', euler, degrees=True)
    return res.as_matrix()

def batch_rodrigues_jacobi(rvec):
    shape = rvec.shape
    rvec = rvec.view(-1, 3)
    device = rvec.device
    dSkew = torch.zeros(3, 9, device=device)
    dSkew[0, 5] = -1
    dSkew[1, 6] = -1
    dSkew[2, 1] = -1
    dSkew[0, 7] =  1
    dSkew[1, 2] =  1
    dSkew[2, 3] =  1
    dSkew = dSkew[None]
    theta = torch.norm(rvec, dim=-1, keepdim=True) + 1e-5
    c = torch.cos(theta)
    s = torch.sin(theta)
    c1 = 1 - c
    itheta = 1 / theta
    r = rvec / theta
    zeros = torch.zeros_like(r[:, :1])
    rx, ry, rz = torch.split(r, 1, dim=1)
    rrt = torch.matmul(r[:, :, None], r[:, None, :])
    skew = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((r.shape[0], 3, 3))
    I = torch.eye(3, device=rvec.device, dtype=rvec.dtype)[None]
    rot_mat = I + s[:, None] * skew + c1[:, None] * torch.bmm(skew, skew)
    
    drrt = torch.stack([
        rx + rx, ry, rz, ry, zeros, zeros, rz, zeros, zeros,
        zeros, rx, zeros, rx, ry + ry, rz, zeros, rz, zeros,
        zeros, zeros, rx, zeros, zeros, ry, rx, ry, rz + rz
    ], dim=-1).view((r.shape[0], 3, 9))
    jacobi = torch.zeros((r.shape[0], 3, 9), device=rvec.device, dtype=rvec.dtype)
    for i in range(3):
        ri = r[:, i:i+1]
        a0 = -s * ri
        a1 = (s - 2*c1*itheta)*ri
        a2 = c1 * itheta
        a3 = (c-s*itheta)*ri
        a4 = s * itheta
        jaco = a0[:, None] * I + a1[:, None] * rrt + a2[:, None] * drrt[:, i].view(-1, 3, 3) + a3[:, None] * skew + a4[:, None] * dSkew[:, i].view(-1, 3, 3)
        jacobi[:, i] = jaco.view(-1, 9)
    rot_mat = rot_mat.view(*shape[:-1], 3, 3)
    jacobi = jacobi.view(*shape[:-1], 3, 9)
    return rot_mat, jacobi

def getJacobianOfRT(rvec, tvec, joints):
    # joints: (bn, nJ, 3)
    dtype, device = rvec.dtype, rvec.device
    bn, nJoints = joints.shape[:2]
    # jacobiToRvec: (bn, 3, 9) // tested by OpenCV and PyTorch 
    Rot, jacobiToRvec = batch_rodrigues_jacobi(rvec)
    I3 = torch.eye(3, dtype=dtype, device=device)[None]
    # jacobiJ_R: (bn, nJ, 3, 3+3+3) => (bn, nJ, 3, 9)
    # // flat by column:
    # // x, 0, 0 | y, 0, 0 | z, 0, 0
    # // 0, x, 0 | 0, y, 0 | 0, z, 0
    # // 0, 0, x | 0, 0, y | 0, 0, z
    jacobi_J_R = torch.zeros((bn, nJoints, 3, 9), dtype=dtype, device=device)
    jacobi_J_R[:, :, 0, :3] = joints
    jacobi_J_R[:, :, 1, 3:6] = joints
    jacobi_J_R[:, :, 2, 6:9] = joints
    # jacobi_J_rvec: (bn, nJ, 3, 3)
    jacobi_J_rvec = torch.matmul(jacobi_J_R, jacobiToRvec[:, None].transpose(-1, -2))
    # if True: # 测试自动梯度
    #     def test_func(rvec):
    #         Rot = batch_rodrigues(rvec[None])[0]
    #         joints_new = joints[0] @ Rot.t()
    #         return joints_new
    #     jac_J_rvec = torch.autograd.functional.jacobian(test_func, rvec[0])
    #     my_j = jacobi_joints_RT[0, ..., :3]
    # jacobi_J_tvec: (bn, nJx3, 3)
    jacobi_J_tvec = I3[None].expand(bn, nJoints, -1, -1)
    jacobi_J_rt = torch.cat([jacobi_J_rvec, jacobi_J_tvec], dim=-1)
    return Rot, jacobiToRvec, jacobi_J_rt

class Model:
    rootIdx = 0
    parents = []


INDEX_HALF = [0,1,2,3,4,5,6,7,15,16,17,18]

class LowPassFilter:
    def __init__(self):
        self.prev_raw_value = None
        self.prev_filtered_value = None
    
    def process(self, value, alpha):
        if self.prev_raw_value is None:
            s = value
        else:
            s = alpha * value + (1.0 - alpha) * self.prev_filtered_value
        self.prev_raw_value = value
        self.prev_filtered_value = s
        return s

class OneEuroFilter:
    def __init__(self, mincutoff=1.0, beta=0.0, dcutoff=1.0, freq=30):
        self.freq = freq
        self.mincutoff = mincutoff
        self.beta = beta
        self.dcutoff = dcutoff
        self.x_filter = LowPassFilter()
        self.dx_filter = LowPassFilter()
    
    def compute_alpha(self, cutoff):
        te = 1.0 / self.freq
        tau = 1.0 / (2 * np.pi * cutoff)
        return 1.0 / (1.0 + tau / te)

    def process(self, x):
        prev_x = self.x_filter.prev_raw_value
        dx = 0.0 if prev_x is None else (x - prev_x) * self.freq
        edx = self.dx_filter.process(dx, self.compute_alpha(self.dcutoff))
        cutoff = self.mincutoff + self.beta * np.abs(edx)
        return self.x_filter.process(x, self.compute_alpha(cutoff))

class BaseBody:
    def __init__(self, cfg_triangulator, cfg_model, cfg) -> None:
        self.triangulator = load_object_from_cmd(cfg_triangulator, [])
        self.body_model = load_object_from_cmd(cfg_model, ['args.use_pose_blending', False, 'args.device', 'cpu'])
        self.cfg = cfg
        self.register_from_lbs(self.body_model)

    def register_from_lbs(self, body_model):
        kintree_shape = np.array(self.cfg.shape.kintree)
        self.nJoints = body_model.J_regressor.shape[0]
        self.k_shapeBlend = body_model.j_shapedirs[self.nJoints:]
        self.j_shapeBlend = body_model.j_shapedirs[:self.nJoints]
        self.jacobian_limb_shapes = self.k_shapeBlend[kintree_shape[:, 1]] - self.k_shapeBlend[kintree_shape[:, 0]]
        self.k_template = body_model.j_v_template[self.nJoints:]
        self.j_template = body_model.j_v_template[:self.nJoints]
        self.k_weights = body_model.j_weights[self.nJoints:]
        self.j_weights = body_model.j_weights[:self.nJoints]
        parents = body_model.parents[1:].cpu().numpy()
        child = np.arange(1, parents.shape[0]+1, dtype=np.int64)
        self.kintree = np.stack([parents, child], axis=1)
        self.parents = np.zeros(parents.shape[0]+1, dtype=int) - 1
        self.parents[self.kintree[:, 1]] = self.kintree[:, 0]
        self.rootIdx = 0
        self.time = time()

def rotation_matrix_from_3x3(A):
    U, s, Vt = np.linalg.svd(A, full_matrices=False)
    V = Vt.T
    T = np.dot(V, U.T)
    # does the current solution use a reflection?
    have_reflection = np.linalg.det(T) < 0

    # if that's not what was specified, force another reflection
    if have_reflection:
        V[:,-1] *= -1
        s[-1] *= -1
        T = np.dot(V, U.T)
    return T

def svd_rot(src, tgt, reflection=False, debug=True):
    # optimum rotation matrix of Y
    A = np.dot(src.T, tgt)
    U, s, Vt = np.linalg.svd(A, full_matrices=False)
    V = Vt.T
    T = np.dot(V, U.T)
    # does the current solution use a reflection?
    have_reflection = np.linalg.det(T) < 0

    # if that's not what was specified, force another reflection
    if reflection != have_reflection:
        V[:,-1] *= -1
        s[-1] *= -1
        T = np.dot(V, U.T)
    if debug:
        err = np.linalg.norm(tgt - src @ T.T, axis=1)
        print('[svd] ', err)
    return T

def normalize(vector):
    return vector/np.linalg.norm(vector)

def rad_from_2vec(vec1, vec2):
    return np.arccos((normalize(vec1)*normalize(vec2)).sum())

def smoothing_factor(t_e, cutoff):
    r = 2 * 3.14 * cutoff * t_e
    return r / (r + 1)

def exponential_smoothing(a, x, x_prev):
    return a * x + (1 - a) * x_prev

FilterResult = namedtuple('FilterResult', ['x', 'dx', 'v', 't'])

class MyFilter:
    def __init__(self, key, filled, min_cutoff=1.0, d_cutoff=1.0,
        beta=0.1) -> None:
        self.key = key
        self.fill_result = filled
        self.min_cutoff = min_cutoff
        self.d_cutoff = d_cutoff
        self.beta = beta
        self.init = False
        self.records = []
        self.result = None
        self.counter = 0
        self.data = []
        self.conf = []
        self.smooth_record = []

    def fill(self, value, conf):
        filled = conf < 0.1
        if filled.sum() >= 1:
            value[filled] = self.fill_result[0][filled]
        if self.key == 'Rh':
            value = rotation_matrix_from_3x3(value.reshape(3, 3))
            value = cv2.Rodrigues(value)[0].reshape(3,)
            if (value < 0).all():
                value = -value
        return value[None]

    def __call__(self, value, conf):
        self.counter += 1
        x, v = value[0], conf[0]
        if self.key == 'Rh':
            x = cv2.Rodrigues(x)[0].reshape(-1)
            v = np.zeros((9,)) + v[0]
        t = np.zeros_like(x)
        t[v>0.1] = self.counter
        self.data.append(x)
        self.conf.append(v)
        if len(self.smooth_record) == 0:
            if self.key == 'Rh':
                start = x
            else:
                # start = self.fill_result[0]
                start = x
            smoothed = FilterResult(start, np.zeros_like(x), np.zeros_like(x), t)
            self.smooth_record.append(smoothed)
        if len(self.data) < 3:
            return self.fill(x, v)
        data = np.array(self.data)
        conf = np.array(self.conf)
        smoothed = self.smooth_record[-1]
        # 预计的速度
        dx_new = x - smoothed.x
        # 滤波器可见，当前可见
        flag_vis = (smoothed.v > 0.1) & (v > 0.1)
        # - 速度异常的，移除掉，认为当前帧不可见
        flag_outlier = (np.abs(smoothed.dx) > 0.05) & (np.abs(dx_new - smoothed.dx)/(1e-5 + smoothed.dx) > 2.)
        if self.key != 'Rh':
            v[flag_vis&flag_outlier] = 0.
        # 滤波器不可见，当前可见，速度打折，认为是新增的帧
        flag_new = (smoothed.v < 0.1)&(v>0.1)
        dx_new[flag_new] /= 3
        # 滤波器可见，当前不可见，速度使用滤波器的速度
        flag_unvis = (v<0.1) & (conf[-2] < 0.1)
        dx_new[flag_unvis] = smoothed.dx[flag_unvis]
        # 滤波器不可见，当前也不可见，速度清0
        dx_new[(v<0.1)&(smoothed.v<0.1)] = 0.
        # 实际估计出来的速度，这里要去掉不可见的地方
        # 混合的权重使用 0.7, 0.3默认全部使用新的一半
        weight_dx = np.zeros_like(dx_new) + 0.7
        dx_smoothed = smoothed.dx*(1-weight_dx) + dx_new*weight_dx
        smoothed_value = smoothed.x + dx_smoothed
        v_new = smoothed.v.copy()
        v_new = v_new * (1-weight_dx) + v*weight_dx
        t_new = smoothed.t.copy()
        t_new[v>0.1] = t[v>0.1]
        smooth_new = FilterResult(smoothed_value, dx_smoothed, v_new, t_new)
        self.smooth_record.append(smooth_new)
        if self.counter == 1000:
            if self.key == 'poses':
                import matplotlib.pyplot as plt
                xrange = np.arange(0, data.shape[0])
                smoothed = np.array([d.x for d in self.smooth_record])
                for nj in range(data.shape[1]):
                    valid = conf[:, nj] > 0.
                    plt.scatter(xrange[valid], data[valid, nj])
                    # yhat = savgol_filter(data[:, nj], data.shape[0], 3)
                    # plt.plot(yhat)
                plt.plot(smoothed)
                plt.show()
                import ipdb;ipdb.set_trace()
        # return self.fill(x, v)
        return self.fill(smooth_new.x, smooth_new.v)

    def __call__0(self, value, conf):
        self.counter += 1
        x, v = value[0], conf[0]
        if self.key == 'Rh':
            x = cv2.Rodrigues(x)[0].reshape(-1)
            v = np.zeros((9,)) + v[0]
        if self.result is None:
            result = FilterResult(x, np.zeros_like(x), v, (v>0)*self.counter)
            self.result = result
            self.records.append(result)
            return self.fill(result.x, result.v)
        # return self.fill(x, v)
        # 维护一个前一帧的，去除outlier
        prev = self.result
        t = prev.t.copy()
        t[v>0.] = self.counter
        dx = x - prev.x # 这里直接使用与之前的结果的差了，避免多帧不可见，然后速度过大
        MAX_DX = 1.
        WINDOW = 31
        not_valid = ((np.abs(dx) > MAX_DX) & (prev.v > 0.1))|\
            (t-prev.t > WINDOW)
        v[not_valid] = 0.
        x_all = np.stack([r.x for r in self.records[-WINDOW:]])
        v_all = np.stack([r.v for r in self.records[-WINDOW:]])
        dx_all = np.stack([r.dx for r in self.records[-WINDOW:]])
        v_sum = v_all.sum(axis=0)
        dx_mean = (dx_all * v_all).sum(axis=0)/(1e-5 + v_all.sum(axis=0))
        # if (x_all.shape[0] > 30) & (self.counter % 40 == 0):
        if True:
            x_mean = (x_all * v_all).sum(axis=0)/(1e-5 + v_all.sum(axis=0))
            x_pred = x_mean
            dx_pred = np.zeros_like(x_pred)
        elif x_all.shape[0] >= 5:
            # 进行smooth
            axt = np.zeros((2, x_all.shape[1]))
            xrange = np.arange(x_all.shape[0]).reshape(-1, 1)
            A0 = np.hstack([xrange, np.ones((x_all.shape[0], 1))])
            for nj in range(x_all.shape[1]):
                conf = v_all[:, nj:nj+1]
                if (conf>0.).sum() < 3:
                    continue
                A = conf * A0
                b = conf * (x_all[:, nj:nj+1])
                est = np.linalg.inv(A.T @ A) @ A.T @ b
                axt[:, nj] = est[:, 0]
            x_all_smoothed = xrange * axt[0:1] + axt[1:]
            x_pred = x_all_smoothed[x_all.shape[0]//2]
            dx_pred = axt[0]
        else:
            x_pred = x_all[x_all.shape[0]//2]
            dx_pred = dx_mean
        if x_all.shape[0] == 1:
            current = FilterResult(x, dx, v, t)
            self.records.append(current)
            self.result = current
        else: 
            # dx_hat = (dx * v + dx_mean * v_mean)/(v+v_mean+1e-5)
            # x_pred = x_mean + dx_hat
            # current = FilterResult(x_pred, dx_hat, v, t)
            current = FilterResult(x, dx, v, t)
            self.records.append(current)
            # 使用平均速度模型
            self.result = FilterResult(x_pred, dx_pred, v_sum, t)
        return self.fill(self.result.x, self.result.v)

    def __call__2(self, value, conf):
        self.counter += 1
        x, v = value[0], conf[0]
        if self.result is None:
            result = FilterResult(x, np.zeros_like(x), v, (v>0)*self.counter)
            self.result = result
            return self.fill(result.x, result.v)
        prev = self.result
        t = prev.t.copy()
        t[v>0.] = self.counter
        # update t
        # dx = (x - prev.x)/(np.maximum(t-prev.t, 1))
        dx = x - prev.x # 这里直接使用与之前的结果的差了，避免多帧不可见，然后速度过大
        dx_ = dx.copy()
        # 判断dx的大小
        large_dx = np.abs(dx) > 0.5
        if large_dx.sum() > 0:
            v[large_dx] = 0.
            t[large_dx] = prev.t[large_dx]
            dx[large_dx] = 0.
        missing_index = ((prev.v > 0.1) & (v < 0.1)) | (t - prev.t > 10)
        if missing_index.sum() > 0:
            print('missing', missing_index)
        new_index = (prev.v < 0.1) & (v > 0.1)
        if new_index.sum() > 0:
            print('new', new_index)
            dx[new_index] = 0.
        weight_dx = v/(1e-5+ 3*prev.v + 1*v)
        weight_x  = v/(1e-5+ 3*prev.v + 1*v)
        # 移除速度过大的点
        dx_hat = exponential_smoothing(weight_dx, dx, prev.dx)
        x_pred = prev.x + dx_hat
        x_hat = exponential_smoothing(weight_x, x, x_pred)
        dx_real = x_hat - prev.x
        # consider the unvisible v
        print_val = {
            't_pre': prev.t,
            'x_inp': x,
            'x_pre': prev.x,
            'x_new': x_hat,
            'dx_inp': dx_,
            'dx_pre': prev.dx,
            'dx_new': dx_hat,
            'dx_real': dx_real,
            'v': v,
            'v_pre': prev.v,
            'w_vel': weight_dx,
            'w_x': weight_x
        }
        for key in print_val.keys():
            print('{:7s}'.format(key), end='  ')
        print('')        
        for i in range(x.shape[0]):
            for key in print_val.keys():
                print('{:7.2f}'.format(print_val[key][i]), end='  ')
            print('')
        v[missing_index] = prev.v[missing_index] / 1.2 # 衰减系数
        result = FilterResult(x_hat, dx_hat, v, t)
        self.result = result
        return self.fill(result.x, result.v)
        if len(self.records) < 10:
            self.records.append([self.counter, value, conf])
            return self.fill(value[0], conf[0])
        if self.x is None:
            time = np.vstack([x[0] for x in self.records])
            value_pre = np.vstack([x[1] for x in self.records])
            conf_pre = np.vstack([x[2] for x in self.records])
            conf_sum = conf_pre.sum(axis=0)
            value_mean = (value_pre * conf_pre).sum(axis=0)/(conf_sum + 1e-5)
            self.x = value_mean
            self.x_conf = conf_sum
            t_prev = np.zeros_like(self.x, dtype=int) - 1
            t_prev[conf_sum>0] = self.counter
            self.t_prev = t_prev
            # 零速度初始化
            self.d_x = np.zeros_like(self.x)
            return self.fill(self.x, self.x_conf)
        # 假设每帧都传进来的吧
        return self.fill(self.x, self.x_conf)
        x_est, v_est, conf_est = self.x.copy(), self.d_x.copy(), self.x_conf.copy()
        value = value[0]
        conf = conf[0]
        d_x = value - self.x
        t_current = np.zeros_like(self.x, dtype=int) - 1
        t_current[conf>0.] = self.counter
        t_est = t_current - self.t_prev
        # 前一帧有观测，当前帧有观测，两帧之差在10帧以内。正常更新
        flag_vv = (t_current > 0) & (self.t_prev > 0) & \
            (t_current - self.t_prev < 10)
        # 前一帧无观测；当前帧有观测的；判断为新增的
        flag_iv = (self.t_prev < 0) & (t_current > 0)
        weight_vel = smoothing_factor(t_est, self.d_cutoff)
        # 将观测的速度权重置0
        weight_vel[flag_vv] = 0.
        vel_hat = exponential_smoothing(weight_vel, d_x, self.d_x)
        cutoff = self.min_cutoff + self.beta * np.abs(vel_hat)
        weight_value = smoothing_factor(t_est, cutoff)
        # 将观测的数值权重置0
        weight_value[flag_vv] = 0.
        weight_value[flag_iv] = 1. # 当前帧可见的，之前的帧不可见的，直接选择当前帧
        vel_hat[flag_iv] = 0.
        x_hat = exponential_smoothing(weight_value, value, self.x)
        flag_vi = (self.t_prev > 0) & (~flag_vv)
        flag_v = flag_vv | flag_vi | flag_iv
        # 前一帧有观测；当前帧无观测的；判断为丢失的
        x_est[flag_v] = x_hat[flag_v]
        v_est[flag_v] = vel_hat[flag_v]
        conf_est[flag_v] = (self.x_conf + conf)[flag_v]/2
        self.t_prev[flag_v] = self.counter
        self.x = x_est
        self.d_x = v_est
        self.x_conf = conf_est
        return self.fill(x_est, conf_est)

class IKBody(BaseBody):
    def __init__(self, key, cfg_triangulator, cfg_model, cfg, debug) -> None:
        super().__init__(cfg_triangulator, cfg_model, cfg)
        self.key = key
        self.frame_index = 0
        self.frame_latest = 0
        self.init = False
        self.records = []
        self.results = []
        self.blank_result = self.make_blank()
        self.fill_result = self.make_fill()
        self.results_newest = self.blank_result
        if True:
            self.lefthand = ManoFitterCPPCache('LEFT')
            self.righthand = ManoFitterCPPCache('RIGHT')
        self.up_vector = 'z'
        self.filter = {}
        for key in ['Rh', 'poses', 'handl', 'handr']:
            self.filter[key] = MyFilter(key, self.fill_result[key])

    def make_blank(self):
        raise NotImplementedError

    def make_fill(self):
        raise NotImplementedError

    def smooth_results(self, params=None):
        results = {'id': 0, 'type': 'smplh_half'}
        for key in ['Rh', 'poses', 'handl', 'handr']:
            value = self.filter[key](params[key], params[key+'_conf'])
            results[key] = value
        for key in ['shapes', 'Th']:
            results[key] = self.blank_result[key]
        return results

    def smooth_results_old(self, params=None):
        if params is not None:
            self.results.append(params)
            if len(self.results) < 10:
                return params
        else:
            if len(self.results) < 10:
                return self.fill_result
            else:
                params = self.fill_result
        results = {'id': 0}
        if False:
            for key in ['Rh', 'poses', 'handl', 'handr']:
                if not self.filter[key].init:
                    import ipdb;ipdb.set_trace()
                else:
                    value = self.filter[key](self.results[-1][key])
        if True:
            for key in ['Rh', 'poses', 'handl', 'handr']:
                find = False
                for WINDOW in [10, 20, 40]:
                    if WINDOW > len(self.results):
                        break
                    records = self.results[-WINDOW:]
                    value = np.vstack([r[key] for r in records])
                    conf = np.vstack([r[key+'_conf'] for r in records])
                    valid = conf[..., 0] > 0
                    if valid.sum() < WINDOW // 3:
                        import ipdb;ipdb.set_trace()
                    else:
                        value, conf = value[valid], conf[valid]
                        mean_value = value.mean(axis=0)
                        std_value = value.std(axis=0)
                        valid2 = (np.abs(value - mean_value) < WINDOW//3 * std_value).any(axis=-1)
                        if valid2.sum() < WINDOW // 4:
                            continue
                        find = True
                        value, conf = value[valid2], conf[valid2]
                        conf_sum = conf.sum(axis=0)
                        mean = (value*conf).sum(axis=0)/(conf_sum + 1e-5)
                        # 计算latest
                        break
                        if key in ['poses', 'handl', 'handr']:
                            conf_sum_p = conf.sum(axis=0)
                            mean_previous = (value*conf).sum(axis=0)/(conf_sum_p + 1e-5)
                            mean[conf_sum<0.01] = mean_previous[conf_sum<0.01]
                            conf_sum[conf_sum<0.01] = conf_sum_p[conf_sum<0.01]
                            # 使用fill的填值
                            mean[conf_sum<0.01] = self.fill_result[key][0][conf_sum<0.01]
                        break
                if find:
                    results[key] = mean[None]
                else:
                    results[key] = self.fill_result[key]
        if False: # 均值滤波
            for key in ['Rh', 'poses', 'handl', 'handr']:
                if key not in params.keys():
                    continue
                if key not in self.cfg.SMOOTH_SIZE.keys():
                    results[key] = params[key]
                records = self.results[-self.cfg.SMOOTH_SIZE[key]:]
                value = np.vstack([r[key] for r in records])
                conf = np.vstack([r[key+'_conf'] for r in records])
                conf_sum = conf.sum(axis=0)
                mean = (value*conf).sum(axis=0)/(conf_sum + 1e-5)
                # 计算latest
                if key in ['poses', 'handl', 'handr']:
                    records = self.results[-5*self.cfg.SMOOTH_SIZE[key]:]
                    value = np.vstack([r[key] for r in records])
                    conf = np.vstack([r[key+'_conf'] for r in records])
                    conf_sum_p = conf.sum(axis=0)
                    mean_previous = (value*conf).sum(axis=0)/(conf_sum_p + 1e-5)
                    mean[conf_sum<0.01] = mean_previous[conf_sum<0.01]
                    conf_sum[conf_sum<0.01] = conf_sum_p[conf_sum<0.01]
                    # 使用fill的填值
                    mean[conf_sum<0.01] = self.fill_result[key][0][conf_sum<0.01]
                results[key] = mean[None]
        results['Th'] = self.blank_result['Th']
        results['shapes'] = self.blank_result['shapes']
        return results

    def get_keypoints3d(self, records, key=None):
        if key is None:
            return np.stack([r[self.key] for r in records])
        else:
            return np.stack([r[key] for r in records])

    def check_keypoints(self, keypoints3d):
        flag = (keypoints3d[..., -1]>self.cfg.MIN_THRES).sum() > 5
        if len(self.records) > 1:
            pre = self.records[-1]
            k_pre = self.get_keypoints3d([pre])
            dist = np.linalg.norm(keypoints3d[..., :3] - k_pre[..., :3], axis=-1)
            conf = np.sqrt(keypoints3d[..., 3] * k_pre[..., 3])
            dist_mean = (dist * conf).sum()/conf.sum()
            flag = flag and dist_mean < 0.1
        return flag

    def __call__(self, data):
        self.frame_index += 1
        k3d = self.triangulator(data)[0]
        keypoints3d = self.get_keypoints3d([k3d])
        flag = self.check_keypoints(keypoints3d)
        if not flag:
            mywarn('Missing keypoints {} [{}->{}]'.format(keypoints3d[..., -1].sum(), self.frame_latest, self.frame_index))
            # 1. 初始化过了，但是超出帧数了，清零
            # 2. 没有初始化过，超出了，清零
            if (self.frame_index - self.frame_latest > 10 and self.init) or not self.init:
                mywarn('Missing keypoints, resetting...')
                self.init = False
                self.records = []
                self.results = []
                return [self.fill_result]
            else:
                return [self.smooth_results()]
        elif not self.init: # 暂时还没有初始化，先等待
            if len(self.records) < 10:
                self.records.append(k3d)
                return [self.fill_result]
        self.records.append(k3d)
        flag, params = self.fitting(keypoints3d, self.results_newest)
        if not flag:
            return [self.fill_result]
        self.frame_latest = self.frame_index
        # smooth results
        results = self.smooth_results(params)
        self.results_newest = results
        k3d['type'] = 'body25'
        return [results, k3d]

class HalfBodyIK(IKBody):
    def get_keypoints3d(self, records):
        THRES_WRIST = 0.2
        keypoints3d = super().get_keypoints3d(records)
        keypoints3d = keypoints3d[:, INDEX_HALF]
        handl = super().get_keypoints3d(records, key='handl3d')
        handr = super().get_keypoints3d(records, key='handr3d')
        dist_ll = np.linalg.norm(keypoints3d[:, 7, :3] - handl[:, 0, :3], axis=-1)
        dist_rr = np.linalg.norm(keypoints3d[:, 4, :3] - handr[:, 0, :3], axis=-1)
        log('Dist left = {}, right = {}'.format(dist_ll, dist_rr))
        handl[dist_ll>THRES_WRIST] = 0.
        handr[dist_rr>THRES_WRIST] = 0.
        keypoints3d = np.hstack([keypoints3d, handl, handr])
        conf = keypoints3d[..., 3:]
        keypoints3d = np.hstack([(keypoints3d[..., :3] * conf).sum(axis=0)/(1e-5 + conf.sum(axis=0)), conf.min(axis=0)])
        keypoints3d = keypoints3d[None]
        # if (keypoints3d.shape[0] == 10):
        return keypoints3d
    
    def _ik_shoulder(self, keypoints3d, params):
        SHOULDER_IDX = [2, 5]
        shoulder = keypoints3d[SHOULDER_IDX[1], :3] - keypoints3d[SHOULDER_IDX[0], :3]
        if self.up_vector == 'x':
            shoulder[..., 0] = 0.
            up_vector = np.array([1., 0., 0.], dtype=np.float32)
        elif self.up_vector == 'z':
            shoulder[..., 2] = 0.
            up_vector = np.array([0., 0., 1.], dtype=np.float32)        
        shoulder = shoulder/np.linalg.norm(shoulder, keepdims=True)
        # 限定一下角度范围
        theta = -np.rad2deg(np.arctan2(shoulder[1], shoulder[2]))
        if (theta < 30 or theta > 150) and False:
            return False, params
        front = np.cross(shoulder, up_vector)
        front = front/np.linalg.norm(front, keepdims=True)
        R = np.stack([shoulder, up_vector, front]).T
        Rh = cv2.Rodrigues(R)[0].reshape(1, 3)
        log('Shoulder:{}'.format(Rh))
        params['R'] = R
        params['Rh'] = Rh
        params['Rh_conf'] = np.zeros((1, 3)) + keypoints3d[SHOULDER_IDX, 3].min()
        return True, params

    def _ik_head(self, keypoints3d, params):
        HEAD_IDX = [0, 8, 9, 10, 11]
        HEAD_ROT_IDX = 0
        est_points = keypoints3d[HEAD_IDX, :3]
        valid = (keypoints3d[HEAD_IDX[0], 3] > self.cfg.MIN_THRES) and (keypoints3d[HEAD_IDX[1:], 3]>self.cfg.MIN_THRES).sum()>=2
        if not valid:
            params['poses_conf'][:, 3*HEAD_ROT_IDX:3*(HEAD_ROT_IDX+1)] = 0.
            return params
        params['poses_conf'][:, 3*HEAD_ROT_IDX:3*(HEAD_ROT_IDX+1)] = keypoints3d[HEAD_IDX, 3].sum()

        gt_points = self.k_template[HEAD_IDX].numpy()
        gt_points = gt_points - gt_points[:1]
        est_points = est_points - est_points[:1]
        # gt_points = gt_points / np.linalg.norm(gt_points, axis=-1, keepdims=True)
        # est_points = est_points / np.linalg.norm(est_points, axis=-1, keepdims=True)

        if True:
            R_global = svd_rot(gt_points, est_points)
            R_local = params['R'].T @ R_global
        elif False:
            est_points_inv = est_points @ params['R'].T.T
            R_local = svd_rot(gt_points, est_points_inv)
        else:
            gt_points = gt_points @ params['R'].T
            R_local = svd_rot(gt_points, est_points)
        euler = rotmat2euler(R_local)
        euler[0] = euler[0] - 25
        # log('euler before filter: {}'.format(euler))
        euler[0] = max(min(euler[0], 30), -30)
        euler[1] = max(min(euler[1], 60), -60)
        euler[2] = max(min(euler[2], 45), -45)
        # log('euler after filter: {}'.format(euler))
        R_local = euler2rotmat(euler)
        R_head = cv2.Rodrigues(R_local)[0].reshape(1, 3)
        params['poses'][:, 3*HEAD_ROT_IDX:3*(HEAD_ROT_IDX+1)] = R_head
        return params
    
    @staticmethod
    def _rad_from_twovec(keypoints3d, start, mid, end, MIN_THRES):
        start = keypoints3d[start]
        mid = keypoints3d[mid]
        end = keypoints3d[end]
        if isinstance(end, list):
            # dst is a list 
            if (end[:, 3] > MIN_THRES).sum() < 2:
                return 0, 0.
            end = np.sum(end * end[:, 3:], axis=0)/(end[:, 3:].sum())
            # use its mean to represent the points
        conf = [start[3], mid[3], end[3]]
        valid = (min(conf) > MIN_THRES).all()
        if not valid:
            return 0, 0.
        conf = sum(conf)
        dir_src = normalize(mid[:3] - start[:3])
        dir_dst = normalize(end[:3] - mid[:3])
        rad = rad_from_2vec(dir_src, dir_dst)
        return conf, rad

    def _ik_elbow(self, keypoints3d, params):
        for name, info in self.cfg.LEAF.items():
            conf, rad = self._rad_from_twovec(keypoints3d, info.start, info.mid, info.end, self.cfg.MIN_THRES)
            if conf <= 0.: continue
            params['poses_conf'][:, 3*info['index']:3*(info['index']+1)] = conf
            rad = np.clip(rad, *info['ranges'])
            rot = rad*np.array(info['axis']).reshape(1, 3)
            params['poses'][:, 3*info['index']:3*(info['index']+1)] = rot
        return params
    
    def _ik_arm(self, keypoints3d, params):
        # forward一遍获得关键点
        # 这里需要确保求解的关节点没有父节点了
        template = self.body_model.keypoints({'poses': params['poses'], 'shapes': params['shapes']}, return_tensor=False)[0]
        for name, info in self.cfg.NODE.items():
            idx = info['children']
            conf = keypoints3d[info['children'], 3]
            if not (conf >self.cfg.MIN_THRES).all():
                continue
            params['poses_conf'][:, 3*info['index']:3*(info['index']+1)] = conf.sum()
            est_points = keypoints3d[idx, :3]
            gt_points = template[idx]
            est_points = est_points - est_points[:1]
            gt_points = gt_points - gt_points[:1]
            R_children = svd_rot(gt_points, est_points)
            R_local = params['R'].T @ R_children
            euler = rotmat2euler(R_local)
            # log('euler {} before filter: {}'.format(name, euler))
            # euler[0] = max(min(euler[0], 90), -90)
            # euler[1] = max(min(euler[1], 90), -90)
            # euler[2] = max(min(euler[2], 90), -90)
            # log('euler {} after filter: {}'.format(name, euler))
            R_local = euler2rotmat(euler)
            params['poses'][:, 3*info['index']:3*(info['index']+1)] = cv2.Rodrigues(R_local)[0].reshape(1, 3)
        return params

    def _ik_palm(self, keypoints3d, params):
        # template = self.body_model.keypoints({'poses': params['poses'], 'shapes': params['shapes']}, return_tensor=False)[0]
        T_joints, _ = self.body_model.transform({'poses': params['poses'], 'shapes': params['shapes']}, return_vertices=False)
        T_joints = T_joints[0].cpu().numpy()
        for name, info in self.cfg.PALM.items():
            # 计算手掌的朝向
            est_points = keypoints3d[:, :3]
            est_conf = keypoints3d[:, 3]
            if est_conf[info.children].min() < self.cfg.MIN_THRES:
                continue
            # 计算朝向
            dir0 = normalize(est_points[info.children[1]] - est_points[info.children[0]])
            dir1 = normalize(est_points[info.children[-1]] - est_points[info.children[0]])
            normal = normalize(np.cross(dir0, dir1))
            dir_parent = normalize(est_points[info.parent[1]] - est_points[info.parent[0]])
            # 计算夹角
            rad = np.arccos((normal * dir_parent).sum()) - np.pi/2
            rad = np.clip(rad, *info['ranges'])
            rot = rad*np.array(info['axis']).reshape(1, 3)
            params['poses'][:, 3*info['index']:3*(info['index']+1)] = rot
            # 考虑手肘的朝向；这个时候还差一个绕手肘的朝向的方向的旋转；这个旋转是在手肘扭曲之前的
            # 先计算出这个朝向；再转化
            R_parents = params['R'] @ T_joints[info.index, :3, :3]
            normal_canonical = R_parents.T @ normal.reshape(3, 1)
            normal_canonical[0, 0] = 0
            normal_canonical = normalize(normal_canonical)
            # 在canonical下的投影
            # normal_T = np.array([0., -1., 0.])
            # trick: 旋转角度的正弦值等于在z轴上的坐标
            rad = np.arcsin(normal_canonical[2, 0])
            rot_x = np.array([-rad, 0., 0.])
            R_x = cv2.Rodrigues(rot_x)[0]
            R_elbow = cv2.Rodrigues(params['poses'][:, 3*info.index_elbow:3*info.index_elbow+3])[0]
            R_elbow = R_elbow @ R_x
            params['poses'][:, 3*info.index_elbow:3*info.index_elbow+3] = cv2.Rodrigues(R_elbow)[0].reshape(1, 3)
        return params

    def _ik_hand(self, template, keypoints3d, poses, conf, is_left):
        # 计算手的每一段的置信度
        poses_full = np.zeros((1, 45))
        conf_full = np.zeros((1, 45))
        y_axis = np.array([0., 1., 0.])
        log('_ik for left: {}'.format(is_left))
        for name, info in self.cfg.HAND.LEAF.items():
            conf, rad = self._rad_from_twovec(keypoints3d, *info.ranges, self.cfg.MIN_THRES)
            if conf <= 0.: 
                log('- skip: {}'.format(name))
                continue
            # trick: 手的朝向是反的
            rad = - rad
            if info.axis == 'auto':
                template_dir = template[info.ranges[2]] - template[info.ranges[1]]
                # y轴方向设成0
                template_dir[1] = 0.
                template_dir = normalize(template_dir)
                # 计算旋转轴，在与z轴的cross方向上
                rot_vec = normalize(np.cross(template_dir, y_axis)).reshape(1, 3)
            elif info.axis == 'root':
                template_dir0 = template[info.ranges[1]] - template[info.ranges[0]]
                template_dir1 = template[info.ranges[2]] - template[info.ranges[1]]
                template_dir0 = normalize(template_dir0)
                template_dir1 = normalize(template_dir1)
                costheta0 = (template_dir0 *template_dir1).sum()
                # 计算当前的夹角
                est_dir0 = keypoints3d[info.ranges[1], :3] - keypoints3d[info.ranges[0], :3]
                est_dir1 = keypoints3d[info.ranges[2], :3] - keypoints3d[info.ranges[1], :3]
                est_dir0 = normalize(est_dir0)
                est_dir1 = normalize(est_dir1)
                costheta1 = (est_dir0 * est_dir1).sum()
                # trick: 手的旋转角度都是相反的
                rad = - np.arccos(np.clip(costheta1/costheta0, 0., 1.))
                rot_vec = normalize(np.cross(template_dir1, y_axis)).reshape(1, 3)
            log('- get: {}: {:.1f}, {}'.format(name, np.rad2deg(rad), rot_vec))            
            poses_full[:, 3*info.index:3*info.index+3] = rot_vec * rad
            conf_full[:, 3*info.index:3*info.index+3] = conf
        # 求解
        usePCA = False
        if usePCA:
            ncomp = 24
            lamb = 0.05
            if is_left:
                A_full = self.body_model.components_full_l[:ncomp].T
                mean_full = self.body_model.mean_full_l
            else:
                A_full = self.body_model.components_full_r[:ncomp].T
                mean_full = self.body_model.mean_full_r
            valid = conf_full[0] > 0.
            A = A_full[valid, :]
            res = (poses_full[:, valid] - mean_full[:, valid]).T
            x = np.linalg.inv(A.T @ A + lamb * np.eye(ncomp)) @ A.T @ res
            poses_full = (A_full @ x).reshape(1, -1) + mean_full
            conf_full = np.zeros_like(poses_full) + valid.sum()
        return poses_full, conf_full

    def make_blank(self):
        params = self.body_model.init_params(1, ret_tensor=False)
        params['id'] = 0
        params['type'] = 'smplh_half'
        params['Th'][0, 0] = 1.
        params['Th'][0, 1] = -1
        params['Th'][0, 2] = 1.
        return params
    
    def make_fill(self):
        params = self.body_model.init_params(1, ret_tensor=False)
        params['id'] = 0
        params['type'] = 'smplh_half'
        params['Rh'][0, 2] = -np.pi/2
        params['handl'] = self.body_model.mean_full_l
        params['handr'] = self.body_model.mean_full_r
        params['poses'][0, 3*4+2] = -np.pi/4
        params['poses'][0, 3*5+2] =  np.pi/4
        return params    

    def fitting(self, keypoints3d, results_pre):
        # keypoints3d: (nFrames, nJoints, 4)
        # 根据肩膀计算身体朝向
        if len(keypoints3d.shape) == 3:
            keypoints3d = keypoints3d[0]
        params = self.body_model.init_params(1, ret_tensor=False)
        params['poses_conf'] = np.zeros_like(params['poses'])
        params['handl_conf'] = np.zeros_like(params['handl'])
        params['handr_conf'] = np.zeros_like(params['handr'])
        params['Rh_conf'] = 0.
        params['id'] = 0
        flag, params = self._ik_shoulder(keypoints3d, params)
        if (params['Rh_conf'] <= 0.01).all():
            return False, params
        params = self._ik_head(keypoints3d, params)
        params = self._ik_elbow(keypoints3d, params)
        params = self._ik_arm(keypoints3d, params)
        params = self._ik_palm(keypoints3d, params)
        if False:
            params['handl'], params['handl_conf'] = self._ik_hand(self.k_template[12:12+21].numpy(), keypoints3d[12:12+21], params['handl'], params['handl_conf'], is_left=True)
            params['handr'], params['handr_conf'] = self._ik_hand(self.k_template[12+21:12+21+21].numpy(), keypoints3d[12+21:12+21+21], params['handr'], params['handr_conf'], is_left=False)
        else:
            params_l = self.lefthand(keypoints3d[12:12+21])[0]
            params_r = self.righthand(keypoints3d[12+21:12+21+21])[0]
            # log('[{:06d}] {}'.format(self.frame_index, params_l['poses'][0]))
            # log('[{:06d}] {}'.format(self.frame_index, params_r['poses'][0]))
            ncomp = params_l['poses'].shape[1]
            A_full = self.body_model.components_full_l[:ncomp].T
            mean_full = self.body_model.mean_full_l
            poses_full = (A_full @ params_l['poses'].T).T + mean_full
            params['handl'] = poses_full            
            A_full = self.body_model.components_full_r[:ncomp].T
            mean_full = self.body_model.mean_full_r
            poses_full = (A_full @ params_r['poses'].T).T + mean_full
            params['handr'] = poses_full
            params['handl_conf'] = np.ones((1, 45))
            params['handr_conf'] = np.ones((1, 45))
        return True, params

class BaseFitter(BaseBody):
    def __init__(self, cfg_triangulator, cfg_model, 
        INIT_SIZE, WINDOW_SIZE, FITTING_SIZE, SMOOTH_SIZE,
        init_dict, fix_dict,
        cfg) -> None:
        super().__init__(cfg_triangulator, cfg_model, cfg)
        self.records = []
        self.results = []
        self.INIT_SIZE = INIT_SIZE
        self.WINDOW_SIZE = WINDOW_SIZE
        self.FITTING_SIZE = FITTING_SIZE
        self.SMOOTH_SIZE = SMOOTH_SIZE
        self.time = 0
        self.frame_latest = 0
        self.frame_index = 0
        self.init = False
        self.init_dict = init_dict
        self.fix_dict = fix_dict
        self.identity_cache = {}

    def get_keypoints3d(self, records):
        raise NotImplementedError
    
    def get_init_params(self, nFrames):
        params = self.body_model.init_params(nFrames, ret_tensor=True)
        for key, val in self.init_dict.items():
            if key == 'Rh':
                import cv2
                R = cv2.Rodrigues(params['Rh'][0].cpu().numpy())[0]
                for vec in self.init_dict['Rh']:
                    Rrel = cv2.Rodrigues(np.deg2rad(np.array(vec)))[0]
                    R = Rrel @ R
                Rh = cv2.Rodrigues(R)[0]
                params['Rh'] = torch.Tensor(Rh).reshape(-1, 3).repeat(nFrames, 1)
            else:
                params[key] = torch.Tensor(val).repeat(nFrames, 1)
        params['id'] = 0
        return params

    def add_any_reg(self, val, val0, JTJ, JTr, w):
        # loss: (val - val0)
        nVals = val.shape[0]
        if nVals not in self.identity_cache.keys():
            self.identity_cache[nVals] = torch.eye(nVals, device=val.device, dtype=val.dtype)
        identity = self.identity_cache[nVals]
        JTJ += w * identity
        JTr += -w*(val - val0).view(-1, 1)
        return 0
    
    def log(self, name, step, delta, res, keys_range=None):
        toc = (time() - self.time)*1000
        norm_d = torch.norm(delta).item()
        norm_f = torch.norm(res).item()
        text = '[{}:{:6.2f}]: step = {:3d}, ||delta|| = {:.4f}, ||res|| = {:.4f}'.format(name, toc, step, norm_d, norm_f)
        print(text)
        self.time = time()
        return norm_d, norm_f

    def fitShape(self, keypoints3d, params, weight, option):
        kintree = np.array(self.cfg.shape.kintree)
        nShapes = params['shapes'].shape[-1]
        # shapes: (1, 10)
        shapes = params['shapes'].T
        shapes0 = shapes.clone()
        device, dtype = shapes.device, shapes.dtype
        lengths3d_est = torch.norm(keypoints3d[:, kintree[:, 1], :3] - keypoints3d[:, kintree[:, 0], :3], dim=-1)
        conf = torch.sqrt(keypoints3d[:, kintree[:, 1], 3:] * keypoints3d[:, kintree[:, 0], 3:])
        conf = conf.repeat(1, 1, 3).reshape(-1, 1)
        nFrames = keypoints3d.shape[0]
        # jacobian: (nFrames, nLimbs, 3, nShapes)
        jacob_limb_shapes = self.jacobian_limb_shapes[None].repeat(nFrames, 1, 1, 1)
        jacob_limb_shapes = jacob_limb_shapes.reshape(-1, nShapes)
        # 注意：这里乘到雅克比的应该是 sqrt(conf)，这里把两个合并了
        JTJ_limb_shapes = jacob_limb_shapes.t() @ (jacob_limb_shapes * conf)
        lossnorm = 0
        self.time = time()
        for iter_ in range(option.max_iter):
            # perform shape blending
            shape_offset = self.k_shapeBlend @ shapes
            keyShaped = self.k_template + shape_offset[..., 0]
            JTJ = JTJ_limb_shapes
            JTr = torch.zeros((nShapes, 1), device=device, dtype=dtype)
            # 并行添加所有的骨架
            dir = keyShaped[kintree[:, 1]] - keyShaped[kintree[:, 0]]
            dir_normalized = dir / torch.norm(dir, dim=-1, keepdim=True)
            # res: (nFrames, nLimbs, 3)
            res = lengths3d_est[..., None] * dir_normalized[None] - dir[None]
            res = conf * res.reshape(-1, 1)
            JTr += jacob_limb_shapes.t() @ res
            self.add_any_reg(shapes, shapes0, JTJ, JTr, w=weight.init_shape)
            delta = torch.linalg.solve(JTJ, JTr)
            shapes += delta
            norm_d, norm_f = self.log('shape', iter_, delta, res)
            if torch.norm(delta) < option.gtol:
                break
            if iter_ > 0 and abs(norm_f - lossnorm)/norm_f < option.ftol:
                break
            lossnorm = norm_f
        shapes = shapes.t()
        params['shapes'] = shapes
        return params, keyShaped

    def fitRT(self, keypoints3d, params, weight, option, kpts_index=None):
        keys_optimized = ['Rh', 'Th']
        if kpts_index is not None:
            keypoints3d = keypoints3d[:, kpts_index]
        init_dict = {
            'Rh': params['Rh'].clone(),
            'Th': params['Th'].clone(),
        }
        init_dict['Rot'] = batch_rodrigues(init_dict['Rh'])
        params_dict = {
            'Rh': params['Rh'],
            'Th': params['Th'],
        }
        keys_range = {}
        for ikey, key in enumerate(keys_optimized):
            if ikey == 0:
                keys_range[key] = [0, init_dict[key].shape[-1]]
            else:
                start = keys_range[keys_optimized[ikey-1]][1]
                keys_range[key] = [start, start+init_dict[key].shape[-1]]
        NUM_FRAME = keys_range[keys_optimized[-1]][1]
        kpts = self.body_model.keypoints({'poses': params['poses'], 'shapes': params['shapes']})
        bn = keypoints3d.shape[0]
        conf = keypoints3d[..., -1:].repeat(1, 1, 3).reshape(bn, -1)
        dtype, device = kpts.dtype, kpts.device
        w_joints = 1./keypoints3d.shape[-2] * weight.joints
        self.time = time()
        for iter_ in range(option.max_iter):
            Rh, Th = params_dict['Rh'], params_dict['Th']
            rot, jacobi_R_rvec, jacobi_joints_RT = getJacobianOfRT(Rh, Th, kpts)
            kpts_rt = torch.matmul(kpts, rot.transpose(-1, -2)) + Th[:, None]
            # // loss: J_obs - (R @ jest + T) => -dR/drvec - dT/dtvec - Rx(djest/dtheta)
            jacobi_keypoints = -jacobi_joints_RT
            if kpts_index is not None:
                jacobi_keypoints = jacobi_keypoints[:, kpts_index]
                kpts_rt = kpts_rt[:, kpts_index]
            jacobi_keypoints_flat = jacobi_keypoints.view(bn, -1, jacobi_keypoints.shape[-1])
            JTJ_keypoints = jacobi_keypoints_flat.transpose(-1, -2) @ (jacobi_keypoints_flat * conf[..., None])
            res = conf[..., None] * ((keypoints3d[..., :3] - kpts_rt).view(bn, -1, 1))
            JTr_keypoints = jacobi_keypoints_flat.transpose(-1, -2) @ res
            # 
            JTJ = torch.eye(bn*NUM_FRAME, device=device, dtype=dtype) * 1e-4
            JTr = torch.zeros((bn*NUM_FRAME, 1), device=device, dtype=dtype)
            # accumulate loss
            for nf in range(bn):
                JTJ[nf*NUM_FRAME:(nf+1)*NUM_FRAME, nf*NUM_FRAME:(nf+1)*NUM_FRAME] += w_joints * JTJ_keypoints[nf]
            # add regularization for each parameter
            for nf in range(bn):
                for key in keys_optimized:
                    start, end = nf*NUM_FRAME + keys_range[key][0], nf*NUM_FRAME + keys_range[key][1]
                    if key == 'Rh':
                        # 增加初始化的loss
                        res_init = rot[nf] - init_dict['Rot'][nf]
                        JTJ[start:end, start:end] += weight['init_'+key] * jacobi_R_rvec[nf] @ jacobi_R_rvec[nf].T
                        JTr[start:end] += -weight.init_Rh * jacobi_R_rvec[nf] @ res_init.reshape(-1, 1)
                    else:
                        res_init = Th[nf] - init_dict['Th'][nf]
                        JTJ[start:end, start:end] += weight['init_'+key] * torch.eye(3)
                        JTr[start:end] += -weight['init_'+key] * res_init.reshape(-1, 1)
            JTr += - w_joints * JTr_keypoints.reshape(-1, 1)
            # solve
            delta = torch.linalg.solve(JTJ, JTr)
            norm_d, norm_f = self.log('pose', iter_, delta, res)
            if torch.norm(delta) < option.gtol:
                break
            if iter_ > 0 and abs(norm_f - lossnorm)/norm_f < option.ftol:
                break
            delta = delta.view(bn, NUM_FRAME)
            lossnorm = norm_f
            for key, _range in keys_range.items():
                if key not in params_dict.keys():continue
                params_dict[key] += delta[:, _range[0]:_range[1]]
                norm_key = torch.norm(delta[:, _range[0]:_range[1]])
        params.update(params_dict)
        return params

    @staticmethod    
    def localTransform(J_shaped, poses, rootIdx, kintree):
        bn = poses.shape[0]
        nThetas = poses.shape[1]//3
        localTrans = torch.eye(4, device=poses.device)[None, None].repeat(bn, nThetas, 1, 1)
        poses_flat = poses.reshape(-1, 3)
        rot_flat = batch_rodrigues(poses_flat)
        rot = rot_flat.view(bn, nThetas, 3, 3)
        localTrans[:, :, :3, :3] = rot
        # set the root
        localTrans[:, rootIdx, :3, 3] = J_shaped[rootIdx].view(1, 3)
        # relTrans: (nKintree, 3)
        relTrans = J_shaped[kintree[:, 1]] - J_shaped[kintree[:, 0]]
        localTrans[:, kintree[:, 1], :3, 3] = relTrans[None]
        return localTrans
    
    @staticmethod
    def globalTransform(localTrans, rootIdx, kintree):
        # localTrans: (bn, nJoints, 4, 4)
        globalTrans = localTrans.clone()
        # set the root
        for (parent, child) in kintree:
            globalTrans[:, child] = globalTrans[:, parent] @ localTrans[:, child]
        return globalTrans

    def jacobi_GlobalTrans_theta(self, poses, j_shaped, rootIdx, kintree,
        device, dtype):
        parents = self.parents
        start = time()
        tic = lambda x: print('-> [{:20s}]: {:.3f}ms'.format(x, 1000*(time()-start)))
        localTrans = self.localTransform(j_shaped, poses, rootIdx, kintree)
        # tic('local trans')
        globalTrans = self.globalTransform(localTrans, rootIdx, kintree)
        # tic('global trans')
        # 计算localTransformation
        poses_flat = poses.view(poses.shape[0], -1, 3)
        # jacobi_R_rvec: (bn, nJ, 3, 9)
        Rot, jacobi_R_rvec = batch_rodrigues_jacobi(poses_flat)
            
        bn, nJoints = localTrans.shape[:2]
        dGlobalTrans_template = torch.zeros((bn, nJoints, 4, 4), device=device, dtype=dtype)
        # compute global transformation
        # results: global transformation to each theta: (bn, nJ, 4, 4, nTheta)
        jacobi_GlobalTrans_theta = torch.zeros((bn, nJoints, 4, 4, nJoints*3), device=device, dtype=dtype)
        # tic('rodrigues')
        for djIdx in range(1, nJoints):
            if djIdx in self.cfg.IGNORE_JOINTS: continue
            # // 第djIdx个轴角的第dAxis个维度
            for dAxis in range(3):
                if dAxis in self.cfg.IGNORE_AXIS.get(str(djIdx), []): continue
                # if(model->frozenJoint[3*djIdx+dAxis])continue;
                # // 从上至下堆叠起来
                dGlobalTrans = dGlobalTrans_template.clone()
                # // 将local的映射过来
                dGlobalTrans[:, djIdx, :3, :3] = jacobi_R_rvec[:, djIdx, dAxis].view(bn, 3, 3)
                if djIdx != rootIdx:
                    # // d(R0 @ R1)/dt1 = R0 @ dR1/dt1, 这里的R0是上一级所有的累积，因此使用全局的
                    parent = parents[djIdx]
                    dGlobalTrans[:, djIdx] = globalTrans[:, parent] @ dGlobalTrans[:, djIdx]
                valid = np.zeros(nJoints, dtype=np.bool)
                valid[djIdx] = True
                # tic('current {}'.format(djIdx))
                # // 遍历骨架树: 将对当前theta的导数传递下去
                for (src, dst) in kintree:
                    # // 当前处理的关节为子节点的不用考虑
                    if dst == djIdx: continue
                    # if dst in self.cfg.IGNORE_JOINTS:continue
                    valid[dst] = valid[src]
                    if valid[src]:
                        # // 如果父节点是有效的: d(R0 @ R1)/dt0 = dR0/dt0 @ R1，这里的R1是当前的局部的，因此使用local的
                        dGlobalTrans[:, dst] = dGlobalTrans[:, src] @ localTrans[:, dst]
                # tic('forward {}'.format(djIdx))
                jacobi_GlobalTrans_theta[..., 3*djIdx+dAxis] = dGlobalTrans
        # tic('jacobia')
        return globalTrans, jacobi_GlobalTrans_theta

    def fitPose(self, keypoints3d, params, weight, option, kpts_index=None):
        # preprocess input data
        if kpts_index is not None:
            keypoints3d = keypoints3d[:, kpts_index]
        bn = keypoints3d.shape[0]
        conf = keypoints3d[..., -1:].repeat(1, 1, 3).reshape(bn, -1)
        if (conf > 0.3).sum() < 4: 
            print('skip')
            return params
        w_joints = 1./keypoints3d.shape[-2] * weight.joints
        # pre-calculate the shape
        Rh, Th, poses = params['Rh'], params['Th'], params['poses']
        init_dict = {
            'Rh': Rh.clone(),
            'Th': Th.clone(),
            'poses': poses.clone()
        }
        init_dict['Rot'] = batch_rodrigues(init_dict['Rh'])
        zero_dict = {key:torch.zeros_like(val) for key, val in init_dict.items()}
        keys_optimized = ['Rh', 'Th', 'poses']
        keys_range = {}
        for ikey, key in enumerate(keys_optimized):
            if ikey == 0:
                keys_range[key] = [0, init_dict[key].shape[-1]]
            else:
                start = keys_range[keys_optimized[ikey-1]][1]
                keys_range[key] = [start, start+init_dict[key].shape[-1]]
        NUM_FRAME = keys_range[keys_optimized[-1]][1]
        # calculate J_shaped
        shapes_t = params['shapes'].t()
        shape_offset = self.j_shapeBlend @ shapes_t
        # jshaped: (nJoints, 3)
        j_shaped = self.j_template + shape_offset[..., 0]
        shape_offset = self.k_shapeBlend @ shapes_t
        # kshaped: (nJoints, 3)
        k_shaped = self.k_template + shape_offset[..., 0]        
        # optimize parameters
        nJoints = j_shaped.shape[0]
        dtype, device = poses.dtype, poses.device
        lossnorm = 0
        self.time = time()
        for iter_ in range(option.max_iter):
            # forward the model
            # 0. poses => full poses
            def tic(name):
                print('[{:20}] {:.3f}ms'.format(name, 1000*(time()-self.time)))
            if 'handl' in params.keys():
                poses_full = self.body_model.extend_poses(poses, params['handl'], params['handr'])
                jacobi_posesful_poses = self.body_model.jacobian_posesfull_poses_
            else:
                poses_full = self.body_model.extend_poses(poses)
            jacobi_posesful_poses = self.body_model.jacobian_posesfull_poses(poses, poses_full)
            # tic('jacobian poses full')
            # 1. poses => local transformation  => global transformation(bn, nJ, 4, 4)
            globalTrans, jacobi_GlobalTrans_theta = self.jacobi_GlobalTrans_theta(poses_full, j_shaped, self.rootIdx, self.kintree, device, dtype)
            # tic('global transform')
            # 2. global transformation => relative transformation => final transformation
            relTrans = globalTrans.clone()
            relTrans[..., :3, 3:] -= torch.matmul(globalTrans[..., :3, :3], j_shaped[None, :, :, None])

            relTrans_weight = torch.einsum('kj,fjab->fkab', self.k_weights, relTrans)
            jacobi_relTrans_theta = jacobi_GlobalTrans_theta.clone()
            # // consider topRight: T - RT0: add -dRT0/dt
            # rot: (bn, nJoints, 3, 3, nThetas) @ (bn, nJoints, 1, 3, 1) => (bn, nJoints, 3, nThetas)
            rot = jacobi_GlobalTrans_theta[..., :3, :3, :]
            j0 = j_shaped.reshape(1, nJoints, 1, 3, 1).expand(bn, -1, -1, -1, -1)
            jacobi_relTrans_theta[..., :3, 3, :] -= torch.sum(rot*j0, dim=-2)
            jacobi_blendtrans_theta = torch.einsum('kj,fjabt->fkabt', self.k_weights, jacobi_relTrans_theta)
            kpts = torch.einsum('fkab,kb->fka', relTrans_weight[..., :3, :3], k_shaped) + relTrans_weight[..., :3, 3]
            # d(RJ0 + J1)/dtheta = d(R)/dtheta @ J0 + dJ1/dtheta
            rot = jacobi_blendtrans_theta[..., :3, :3, :]
            k0 = k_shaped.reshape(1, k_shaped.shape[0], 1, 3, 1).expand(bn, -1, -1, -1, -1)
            # jacobi_keypoints_theta: (bn, nKeypoints, 3, nThetas)
            jacobi_keypoints_theta = torch.sum(rot*k0, dim=-2) + jacobi_blendtrans_theta[..., :3, 3, :]
            # tic('keypoints')
            # // compute the jacobian of R T
            # // loss: J_obs - (R @ jest + T)
            rot, jacobi_R_rvec, jacobi_joints_RT = getJacobianOfRT(Rh, Th, kpts)
            kpts_rt = torch.matmul(kpts, rot.transpose(-1, -2)) + Th[:, None]
            rot_nk = rot[:, None].expand(-1, k_shaped.shape[0], -1, -1)
            jacobi_keypoints_theta = torch.matmul(rot_nk, jacobi_keypoints_theta)
            # compute jacobian of poses
            shape_jacobi = jacobi_keypoints_theta.shape[:-1]
            NUM_THETAS = jacobi_posesful_poses.shape[0]
            jacobi_keypoints_poses = (jacobi_keypoints_theta[..., :NUM_THETAS].view(-1, NUM_THETAS) @ jacobi_posesful_poses).reshape(*shape_jacobi, -1)
            # // loss: J_obs - (R @ jest + T) => -dR/drvec - dT/dtvec - Rx(djest/dtheta)
            jacobi_keypoints = torch.cat([-jacobi_joints_RT, -jacobi_keypoints_poses], dim=-1)
            if kpts_index is not None:
                jacobi_keypoints = jacobi_keypoints[:, kpts_index]
                kpts_rt = kpts_rt[:, kpts_index]
            jacobi_keypoints_flat = jacobi_keypoints.view(bn, -1, jacobi_keypoints.shape[-1])
            # tic('jacobian keypoints')
            JTJ_keypoints = jacobi_keypoints_flat.transpose(-1, -2) @ (jacobi_keypoints_flat * conf[..., None])
            res = conf[..., None] * ((keypoints3d[..., :3] - kpts_rt).view(bn, -1, 1))
            JTr_keypoints = jacobi_keypoints_flat.transpose(-1, -2) @ res
            cache_dict = {
                'Th': Th,
                'Rh': Rh,
                'poses': poses,
            }
            # 计算loss
            # JTJ = torch.zeros((bn*NUM_FRAME, bn*NUM_FRAME), device=device, dtype=dtype)
            JTJ = torch.eye(bn*NUM_FRAME, device=device, dtype=dtype) * 1e-4
            JTr = torch.zeros((bn*NUM_FRAME, 1), device=device, dtype=dtype)
            # add regularization for each parameter
            for nf in range(bn):
                for key in keys_optimized:
                    start, end = nf*NUM_FRAME + keys_range[key][0], nf*NUM_FRAME + keys_range[key][1]
                    JTJ[start:end, start:end] += weight['reg_{}'.format(key)] * torch.eye(end-start)
                    JTr[start:end] += -weight['reg_{}'.format(key)] * cache_dict[key][nf].view(-1, 1)
                    # add init for Rh
                    if key == 'Rh':
                        # 增加初始化的loss
                        res_init = rot[nf] - init_dict['Rot'][nf]
                        JTJ[start:end, start:end] += weight['init_'+key] * jacobi_R_rvec[nf] @ jacobi_R_rvec[nf].T
                        JTr[start:end] += -weight.init_Rh * jacobi_R_rvec[nf] @ res_init.reshape(-1, 1)
                    else:
                        res_init = cache_dict[key][nf] - init_dict[key][nf]
                        JTJ[start:end, start:end] += weight['init_'+key] * torch.eye(end-start)
                        JTr[start:end] += -weight['init_'+key] * res_init.reshape(-1, 1)
            # add keypoints loss
            for nf in range(bn):
                JTJ[nf*NUM_FRAME:(nf+1)*NUM_FRAME, nf*NUM_FRAME:(nf+1)*NUM_FRAME] += w_joints * JTJ_keypoints[nf]
            JTr += - w_joints * JTr_keypoints.reshape(-1, 1)
            # tic('add loss')
            delta = torch.linalg.solve(JTJ, JTr)
            # tic('solve')
            norm_d, norm_f = self.log('pose', iter_, delta, res, keys_range)
            if torch.norm(delta) < option.gtol:
                break
            if iter_ > 0 and abs(norm_f - lossnorm)/norm_f < option.ftol:
                break
            delta = delta.view(bn, NUM_FRAME)
            lossnorm = norm_f
            for key, _range in keys_range.items():
                if key not in cache_dict.keys():continue
                cache_dict[key] += delta[:, _range[0]:_range[1]]
        res = {
            'id': params['id'],
            'poses': poses,
            'shapes': params['shapes'],
            'Rh': Rh,
            'Th': Th,
        }
        for key, val in params.items():
            if key not in res.keys():
                res[key] = val
        return res

    def try_to_init(self, records):
        if self.init:
            return copy.deepcopy(self.params_newest)
        mywarn('>> Initialize')
        keypoints3d = self.get_keypoints3d(records)
        params, keypoints_template = self.fitShape(keypoints3d, self.get_init_params(keypoints3d.shape[0]), self.cfg.shape.weight, self.cfg.shape.option)
        Rot = batch_rodrigues(params['Rh'])
        keypoints_template = torch.matmul(keypoints_template, Rot[0].t())
        if self.cfg.initRT.mean_T:
            conf = keypoints3d[..., 3:]
            T = ((keypoints3d[..., :3] - keypoints_template[None])*conf).sum(dim=-2)/(conf.sum(dim=-2))
            params['Th'] = T
        params = self.fitRT(keypoints3d, params, self.cfg.initRT.weight, self.cfg.initRT.option,
            kpts_index=self.cfg.TORSO_INDEX)
        params = self.fitPose(keypoints3d, params, self.cfg.init_pose.weight, self.cfg.init_pose.option,
            kpts_index=self.cfg.TORSO_INDEX)
        params = self.fitPose(keypoints3d, params, self.cfg.init_pose.weight, self.cfg.init_pose.option,
            kpts_index=self.cfg.BODY_INDEX)
        mywarn('>> Initialize Rh = {}, Th = {}'.format(params['Rh'][0], params['Th'][0]))
        params = Params(params)[-1]
        self.init = True
        return params

    def fitting(self, params_init, records):
        keypoints3d = self.get_keypoints3d(records[-self.WINDOW_SIZE:])
        params = params_init
        params = self.fitRT(keypoints3d[-self.FITTING_SIZE:], params, self.cfg.RT.weight, self.cfg.RT.option)
        params = self.fitPose(keypoints3d[-self.FITTING_SIZE:], params, self.cfg.pose.weight, self.cfg.pose.option)
        return params

    def filter_result(self, result):
        poses = result['poses'].reshape(-1, 3)
        # TODO: 这里的xyz是scipy中的XYZ
        euler = axis_angle_to_euler(poses, order='xyz')
        # euler[euler>np.pi] = 0.
        poses = euler_to_axis_angle(euler, order='xyz')
        result['euler'] = euler
        return result

    def smooth_results(self):
        results_ = {}
        for key in self.results[0].keys():
            if key == 'id': continue
            if key not in self.SMOOTH_SIZE.keys():continue
            results_[key] = np.vstack([r[key] for r in self.results[-self.SMOOTH_SIZE[key]:]])
            results_[key] = np.mean(results_[key], axis=0, keepdims=True)
        results_['id'] = 0
        for key, val in self.fix_dict.items():
            results_[key] = np.array(val)
        # results_['Th'][:] = 0.
        return [results_]

    def check_keypoints(self, keypoints3d):
        flag = (keypoints3d[..., -1]>0.3).sum() > 5
        if len(self.records) > 1:
            pre = self.records[-1]
            k_pre = self.get_keypoints3d([pre])
            dist = torch.norm(keypoints3d[..., :3] - k_pre[..., :3], dim=-1)
            conf = torch.sqrt(keypoints3d[..., 3] * k_pre[..., 3])
            dist_mean = (dist * conf).sum()/conf.sum()
            flag = flag and dist_mean < 0.1
        return flag

    def __call__(self, data):
        self.frame_index += 1
        k3d = self.triangulator(data)[0]
        keypoints3d = self.get_keypoints3d([k3d])
        flag = self.check_keypoints(keypoints3d)
        if not flag:
            mywarn('Missing keypoints {} [{}->{}]'.format(keypoints3d[..., -1].sum(), self.frame_latest, self.frame_index))
            if self.frame_index - self.frame_latest > 10 and self.init:
                mywarn('Missing keypoints, resetting...')
                self.init = False
                self.records = []
                self.results = []
            return -1
        self.records.append(k3d)
        if len(self.records) < self.INIT_SIZE:
            return -1
        with Timer('fitting', True):
            params = self.try_to_init(self.records)
            params = self.fitting(params, self.records)
        params = self.filter_result(params)
        self.results.append(params)
        self.params_newest = params
        self.frame_latest = self.frame_index
        return self.smooth_results()

class FitterCPPCache:
    def __init__(self) -> None:
        self.init = False
        self.records = []
        self.frame_index = 0
        self.frame_latest = 0

    def try_to_init(self, keypoints3d):
        if self.init:
            return copy.deepcopy(self.params_newest)
        mywarn('>> Initialize')
        params = self.handmodel.init_params(1)
        if not (keypoints3d[..., -1] > 0.).all(): # 只有一半的点可见
            return params
        params = self.handmodel.fit3DShape(keypoints3d, params)
        params = self.handmodel.init3DRT(keypoints3d[-1:], params)
        params = self.handmodel.fit3DPose(keypoints3d[-1:], params)
        mywarn('>> Initialize Rh = {}, Th = {}'.format(params['Rh'][0], params['Th'][0]))
        self.init = True
        return params

    def check_keypoints(self, keypoints3d):
        flag = (keypoints3d[..., -1]>0.3).sum() > 5
        if len(self.records) > 1:
            k_pre = self.records[-1][None]
            dist = np.linalg.norm(keypoints3d[..., :3] - k_pre[..., :3], axis=-1)
            conf = np.sqrt(keypoints3d[..., 3] * k_pre[..., 3])
            dist_mean = (dist * conf).sum()/(1e-5+conf.sum())
            print(dist_mean)
            flag = flag and dist_mean < 0.1
        return flag

    def smooth_results(self, params=None):
        if params is None:
            params = self.params_newest.copy()
        params['poses'] = params['poses'][:, 3:]
        params['id'] = 0
        return [params]

    def __call__(self, keypoints3d):
        self.frame_index += 1
        flag = self.check_keypoints(keypoints3d)
        if not flag:
            mywarn('Missing keypoints {} [{}->{}]'.format(keypoints3d[..., -1].sum(), self.frame_latest, self.frame_index))
            if self.frame_index - self.frame_latest > 10:
                mywarn('Missing keypoints, resetting...')
                self.init = False
                self.records = []
            return self.smooth_results(self.handmodel.init_params(1))
        self.records.append(keypoints3d)
        with Timer('fitting'):
            params = self.try_to_init(keypoints3d[None])
            params = self.handmodel.fit3DPose(keypoints3d[None], params)
        self.params_newest = params
        self.frame_latest = self.frame_index
        return self.smooth_results()
class ManoFitterCPPCache(FitterCPPCache):
    def __init__(self, name='LEFT') -> None:
        super().__init__()
        self.handmodel = load_object_from_cmd('config/model/mano_full_cpp.yml', [
            'args.model_path', 'data/bodymodels/manov1.2/MANO_{}.pkl'.format(name),
            'args.regressor_path', 'data/smplx/J_regressor_mano_{}.txt'.format(name),
        ])

class SMPLFitterCPPCache(FitterCPPCache):
    def __init__(self, name='half') -> None:
        super().__init__()
        self.handmodel = load_object_from_cmd('config/model/mano_full_cpp.yml', [
            'args.model_path', 'data/bodymodels/manov1.2/MANO_{}.pkl'.format(name),
            'args.regressor_path', 'data/smplx/J_regressor_mano_{}.txt'.format(name),
        ])

class ManoFitterCPP:
    def __init__(self, cfg_triangulator, key) -> None:
        self.handmodel = load_object_from_cmd('config/model/mano_full_cpp.yml', [])
        self.triangulator = load_object_from_cmd(cfg_triangulator, [])
        self.time = 0
        self.frame_latest = 0
        self.frame_index = 0
        self.key = 'handl3d'
        self.init = False
        self.params_newest = None
        self.records, self.results = [], []
        self.INIT_SIZE = 10

    def get_keypoints3d(self, records, key=None):
        if key is None:
            return np.stack([r[self.key] for r in records])
        else:
            return np.stack([r[key] for r in records])
    
    def try_to_init(self, records):
        if self.init:
            return copy.deepcopy(self.params_newest)
        mywarn('>> Initialize')
        keypoints3d = self.get_keypoints3d(records)
        params = self.handmodel.init_params(1)
        params = self.handmodel.fit3DShape(keypoints3d, params)
        params = self.handmodel.init3DRT(keypoints3d[-1:], params)
        params = self.handmodel.fit3DPose(keypoints3d[-1:], params)
        mywarn('>> Initialize Rh = {}, Th = {}'.format(params['Rh'][0], params['Th'][0]))
        self.init = True
        return params

    def smooth_results(self):
        params = self.params_newest.copy()
        params['poses'] = params['poses'][:, 3:]
        params['id'] = 0
        return [params]

    def __call__(self, data):
        self.frame_index += 1
        k3d = self.triangulator(data)[0]
        keypoints3d = self.get_keypoints3d([k3d])
        # flag = self.check_keypoints(keypoints3d)
        flag = True
        if not flag:
            mywarn('Missing keypoints {} [{}->{}]'.format(keypoints3d[..., -1].sum(), self.frame_latest, self.frame_index))
            if self.frame_index - self.frame_latest > 10 and self.init:
                mywarn('Missing keypoints, resetting...')
                self.init = False
                self.records = []
                self.results = []
            return -1
        self.records.append(k3d)
        if len(self.records) < self.INIT_SIZE:
            return -1
        with Timer('fitting'):
            params = self.try_to_init(self.records)
            k3d = self.get_keypoints3d(self.records[-1:])
            params = self.handmodel.fit3DPose(k3d, params)
            # params['poses'] = torch.Tensor(params['poses'][:, 3:])
            # params['shapes'] = torch.Tensor(params['shapes'])
            # params['Rh'] = torch.Tensor(params['Rh'])
            # params['Th'] = torch.Tensor(params['Th'])
        # params = self.filter_result(params)
        self.results.append(params)
        self.params_newest = params
        self.frame_latest = self.frame_index
        return self.smooth_results()
        
class BodyFitter(BaseFitter):
    def __init__(self, key, **kwargs):
        super().__init__(**kwargs)
        self.key = key
    
    def get_keypoints3d(self, records, key=None):
        if key is None:
            return torch.Tensor(np.stack([r[self.key] for r in records]))
        else:
            return torch.Tensor(np.stack([r[key] for r in records]))

class ManoFitter(BodyFitter):
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

class HalfFitter(BodyFitter):
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.INDEX_HALF = [0,1,2,3,4,5,6,7,15,16,17,18]

    def get_keypoints3d(self, records):
        THRES_WRIST = 0.05
        keypoints3d = super().get_keypoints3d(records)
        keypoints3d = keypoints3d[:, self.INDEX_HALF]
        handl = super().get_keypoints3d(records, key='handl3d')
        handr = super().get_keypoints3d(records, key='handr3d')
        dist_ll = torch.norm(keypoints3d[:, 7, :3] - handl[:, 0, :3], dim=-1)
        dist_rr = torch.norm(keypoints3d[:, 4, :3] - handr[:, 0, :3], dim=-1)
        handl[dist_ll>THRES_WRIST] = 0.
        handr[dist_rr>THRES_WRIST] = 0.
        keypoints3d = np.hstack([keypoints3d, handl, handr])
        conf = keypoints3d[..., 3:]
        keypoints3d = np.hstack([(keypoints3d[..., :3] * conf).sum(axis=0)/(1e-5 + conf.sum(axis=0)), conf.min(axis=0)])
        keypoints3d = keypoints3d[None]
        # if (keypoints3d.shape[0] == 10):
        return torch.Tensor(keypoints3d)
    
    def filter_result(self, result):
        result = super().filter_result(result)
        # 限定一下关节旋转
        # 手肘
        # result['poses'][:, 5*3+1] = np.clip(result['poses'][:, 5*3+1], -2.5, 0.1)
        # result['poses'][:, 6*3+1] = np.clip(result['poses'][:, 5*3+1], -0.1, 2.5)
        # 手腕
        return result

class HalfHandFitter(HalfFitter):
    def __init__(self, cfg_handl, cfg_handr, **kwargs) -> None:
        super().__init__(**kwargs)
        self.handl = load_object_from_cmd(cfg_handl, [])
        self.handr = load_object_from_cmd(cfg_handr, [])

    def get_init_params(self, nFrames):
        params = super().get_init_params(nFrames)
        params_ = self.handl.get_init_params(nFrames)
        params['shapes_handl'] = params_['shapes']
        params['shapes_handr'] = params_['shapes'].clone()
        params['Rh_handl'] = torch.zeros((nFrames, 3))
        params['Rh_handr'] = torch.zeros((nFrames, 3))
        params['Th_handl'] = torch.zeros((nFrames, 3))
        params['Th_handr'] = torch.zeros((nFrames, 3))
        return params

    def fitPose(self, keypoints3d, params, weight, option):
        keypoints = {
            'handl': keypoints3d[:, -21-21:-21, :],
            'handr': keypoints3d[:, -21:, :]
        }
        for key in ['handl', 'handr']:
            kpts = keypoints[key]
            params_ = {
                'id': 0,
                'Rh': params['Rh_'+key],
                'Th': params['Th_'+key],
                'shapes': params['shapes_'+key],
                'poses': params[key],
            }
            if key == 'handl':
                params_ = self.handl.fitPose(kpts, params_, self.handl.cfg.pose.weight, self.handl.cfg.pose.option)
            else:
                params_ = self.handr.fitPose(kpts, params_, self.handr.cfg.pose.weight, self.handr.cfg.pose.option)
            params['Rh_'+key] = params_['Rh']
            params['Th_'+key] = params_['Th']
            params['shapes_'+key] = params_['shapes']
            params[key] = params_['poses']
        return super().fitPose(keypoints3d, params, weight, option, 
            kpts_index=[0,1,2,3,4,5,6,7,8,9,10,11,
                12, 17, 21, 25, 29,
                24, 29, 33, 37, 41])
    
    def try_to_init(self, records):
        if self.init:
            return self.params_newest
        params = super().try_to_init(records)
        self.handl.init = False
        self.handr.init = False
        key = 'handl'
        params_ = self.handl.try_to_init(records)
        params['handl'] = params_['poses']
        params['Rh_'+key] = params_['Rh']
        params['Th_'+key] = params_['Th']
        params['shapes_'+key] = params_['shapes']
        key = 'handr'
        params_ = self.handr.try_to_init(records)
        params[key] = params_['poses']
        params['Rh_'+key] = params_['Rh']
        params['Th_'+key] = params_['Th']
        params['shapes_'+key] = params_['shapes']
        return params

if __name__ == '__main__':
    from glob import glob
    from os.path import join
    from tqdm import tqdm
    from ..mytools.file_utils import read_json
    from ..config.baseconfig import load_object_from_cmd

    data= '/nas/datasets/EasyMocap/302'
    # data = '/home/qing/Dataset/handl'
    mode = 'half'
    # data = '/home/qing/DGPU/home/shuaiqing/zju-mocap-mp/female-jump'
    data = '/home/qing/Dataset/desktop/0402/test3'
    k3dnames = sorted(glob(join(data, 'output-keypoints3d', 'keypoints3d', '*.json')))
    if mode == 'handl':
        fitter = load_object_from_cmd('config/recon/fit_manol.yml', [])
    elif mode == 'half':
        fitter = load_object_from_cmd('config/recon/fit_half.yml', [])
    elif mode == 'smpl':
        fitter = load_object_from_cmd('config/recon/fit_smpl.yml', [])
    from easymocap.socket.base_client import BaseSocketClient
    client = BaseSocketClient('0.0.0.0', 9999)

    for k3dname in tqdm(k3dnames):
        k3ds = read_json(k3dname)
        if mode == 'handl':
            k3ds = np.array(k3ds[0]['handl3d'])
            data = {fitter.key3d: k3ds}
        elif mode == 'half':
            data = {
                'keypoints3d': np.array(k3ds[0]['keypoints3d']), 
                'handl3d': np.array(k3ds[0]['handl3d']), 
                'handr3d': np.array(k3ds[0]['handr3d'])
            }
        elif mode == 'smpl':
            k3ds = np.array(k3ds[0]['keypoints3d'])
            data = {fitter.key3d: k3ds}
        results = fitter(data)
        if results != -1:
            client.send_smpl(results)