import onnxruntime as rt
import numpy as np
import json
import torch
import cv2
import os
from torch.utils.data.dataset import Dataset
import random
import math
import argparse

# Constants and paths defining model, image, and dataset specifics
MODEL_DIR = './movenet_int8.onnx'  # Path to the MoveNet model
IMG_SIZE = 192  # Image size used for processing
FEATURE_MAP_SIZE = 48  # Feature map size used in the model
CENTER_WEIGHT_ORIGIN_PATH = './center_weight_origin.npy'  # Path to center weight origin file
DATASET_PATH = 'your_dataset_path'  # Base path for the dataset
EVAL_LABLE_PATH = os.path.join(DATASET_PATH, "val2017.json")  # Path to validation labels JSON file
EVAL_IMG_PATH = os.path.join(DATASET_PATH, 'imgs')  # Path to validation images


def getDist(pre, labels):
    """
    Calculate the Euclidean distance between predicted and labeled keypoints.

    Args:
        pre: Predicted keypoints [batchsize, 14]
        labels: Labeled keypoints [batchsize, 14]

    Returns:
        dist: Distance between keypoints [batchsize, 7]
    """
    pre = pre.reshape([-1, 17, 2])
    labels = labels.reshape([-1, 17, 2])
    res = np.power(pre[:,:,0]-labels[:,:,0],2)+np.power(pre[:,:,1]-labels[:,:,1],2)
    return res


def getAccRight(dist, th = 5/IMG_SIZE):
    """
    Compute accuracy for each keypoint based on a threshold.

    Args:
        dist: Distance between keypoints [batchsize, 7]
        th: Threshold for accuracy computation

    Returns:
        res: Accuracy per keypoint [7,] representing the count of correct predictions
    """
    res = np.zeros(dist.shape[1], dtype=np.int64)
    for i in range(dist.shape[1]):
            res[i] = sum(dist[:,i]<th)
    return res

def myAcc(output, target):
    '''
    Compute accuracy across keypoints.

    Args:
        output: Predicted keypoints
        target: Labeled keypoints

    Returns:
        cate_acc: Categorical accuracy [7,] representing the count of correct predictions per keypoint
    '''
    # [h, ls, rs, lb, rb, lr, rr]
    # output[:,6:10] = output[:,6:10]+output[:,2:6]
    # output[:,10:14] = output[:,10:14]+output[:,6:10]
    # Calculate distance between predicted and labeled keypoints
    dist = getDist(output, target)
    # Calculate accuracy for each keypoint
    cate_acc = getAccRight(dist)
    return cate_acc

# Predefined numpy arrays and weights for calculations
_range_weight_x = np.array([[x for x in range(FEATURE_MAP_SIZE)] for _ in range(FEATURE_MAP_SIZE)])
_range_weight_y = _range_weight_x.T
_center_weight = np.load(CENTER_WEIGHT_ORIGIN_PATH).reshape(FEATURE_MAP_SIZE,FEATURE_MAP_SIZE)

def maxPoint(heatmap, center=True):
    """
    Find the coordinates of maximum values in a heatmap.

    Args:
        heatmap: Input heatmap data
        center: Flag to indicate whether to consider center-weighted points

    Returns:
        x, y: Coordinates of maximum values in the heatmap
    """
    if len(heatmap.shape) == 3:
        batch_size,h,w = heatmap.shape
        c = 1
    elif len(heatmap.shape) == 4:
        # n,c,h,w
        batch_size,c,h,w = heatmap.shape
    if center:
        heatmap = heatmap*_center_weight
    heatmap = heatmap.reshape((batch_size,c, -1)) #64,c, cfg['feature_map_size']xcfg['feature_map_size']
    max_id = np.argmax(heatmap,2)#64,c, 1
    y = max_id//w
    x = max_id%w
    # bv
    return x,y

def movenetDecode(data, kps_mask=None,mode='output', num_joints = 17, 
                img_size=192, hm_th=0.1):

    '''
    Decode MoveNet output data to predicted keypoints.

    Args:
        data: MoveNet output data
        kps_mask: Keypoints mask
        mode: Mode of decoding ('output' or 'label')
        num_joints: Number of joints/keypoints
        img_size: Image size
        hm_th: Threshold for heatmap processing

    Returns:
        res: Decoded keypoints
    '''
    
    ##data [64, 7, 48, 48] [64, 1, 48, 48] [64, 14, 48, 48] [64, 14, 48, 48]
    #kps_mask [n, 7]
    if mode == 'output':
        batch_size = data[0].shape[0]
        heatmaps = data[0]
        heatmaps[heatmaps < hm_th] = 0
        centers = data[1]
        regs = data[2]
        offsets = data[3]
        cx,cy = maxPoint(centers)
        dim0 = np.arange(batch_size,dtype=np.int32).reshape(batch_size,1)
        dim1 = np.zeros((batch_size,1),dtype=np.int32)
        res = []
        for n in range(num_joints):
            reg_x_origin = (regs[dim0,dim1+n*2,cy,cx]+0.5).astype(np.int32)
            reg_y_origin = (regs[dim0,dim1+n*2+1,cy,cx]+0.5).astype(np.int32)
            reg_x = reg_x_origin+cx
            reg_y = reg_y_origin+cy
            ### for post process
            reg_x = np.reshape(reg_x, (reg_x.shape[0],1,1))
            reg_y = np.reshape(reg_y, (reg_y.shape[0],1,1))
            reg_x = reg_x.repeat(FEATURE_MAP_SIZE,1).repeat(FEATURE_MAP_SIZE,2)
            reg_y = reg_y.repeat(FEATURE_MAP_SIZE,1).repeat(FEATURE_MAP_SIZE,2)
            range_weight_x = np.reshape(_range_weight_x,(1,FEATURE_MAP_SIZE,FEATURE_MAP_SIZE)).repeat(reg_x.shape[0],0)
            range_weight_y = np.reshape(_range_weight_y,(1,FEATURE_MAP_SIZE,FEATURE_MAP_SIZE)).repeat(reg_x.shape[0],0)
            tmp_reg_x = (range_weight_x-reg_x)**2
            tmp_reg_y = (range_weight_y-reg_y)**2
            tmp_reg = (tmp_reg_x+tmp_reg_y)**0.5+1.8#origin 1.8
            tmp_reg = heatmaps[:,n,...]/tmp_reg
            tmp_reg = tmp_reg[:,np.newaxis,:,:]
            reg_x,reg_y = maxPoint(tmp_reg, center=False)
            reg_x[reg_x>47] = 47
            reg_x[reg_x<0] = 0
            reg_y[reg_y>47] = 47
            reg_y[reg_y<0] = 0
            score = heatmaps[dim0,dim1+n,reg_y,reg_x]
            offset_x = offsets[dim0,dim1+n*2,reg_y,reg_x]#*img_size//4
            offset_y = offsets[dim0,dim1+n*2+1,reg_y,reg_x]#*img_size//4
            res_x = (reg_x+offset_x)/(img_size//4)
            res_y = (reg_y+offset_y)/(img_size//4)
            res_x[score<hm_th] = -1
            res_y[score<hm_th] = -1
            res.extend([res_x, res_y])
        res = np.concatenate(res,axis=1) #bs*14
    elif mode == 'label':
        kps_mask = kps_mask.detach().cpu().numpy()
        data = data.detach().cpu().numpy()
        batch_size = data.shape[0]
        heatmaps = data[:,:17,:,:]
        centers = data[:,17:18,:,:]
        regs = data[:,18:52,:,:]
        offsets = data[:,52:,:,:]
        cx,cy = maxPoint(centers)
        dim0 = np.arange(batch_size,dtype=np.int32).reshape(batch_size,1)
        dim1 = np.zeros((batch_size,1),dtype=np.int32)
        res = []
        for n in range(num_joints):
            reg_x_origin = (regs[dim0,dim1+n*2,cy,cx]+0.5).astype(np.int32)
            reg_y_origin = (regs[dim0,dim1+n*2+1,cy,cx]+0.5).astype(np.int32)
            reg_x = reg_x_origin+cx
            reg_y = reg_y_origin+cy
            reg_x[reg_x>47] = 47
            reg_x[reg_x<0] = 0
            reg_y[reg_y>47] = 47
            reg_y[reg_y<0] = 0
            offset_x = offsets[dim0,dim1+n*2,reg_y,reg_x]#*img_size//4
            offset_y = offsets[dim0,dim1+n*2+1,reg_y,reg_x]#*img_size//4
            res_x = (reg_x+offset_x)/(img_size//4)
            res_y = (reg_y+offset_y)/(img_size//4)
            res_x[kps_mask[:,n]==0] = -1
            res_y[kps_mask[:,n]==0] = -1
            res.extend([res_x, res_y])
        res = np.concatenate(res,axis=1) #bs*14
    return res

def label2heatmap(keypoints, other_keypoints, img_size):
    '''
    Convert labeled keypoints to heatmaps for keypoints.

    Args:
        keypoints: Target person's keypoints
        other_keypoints: Other people's keypoints
        img_size: Size of the image

    Returns:
        heatmaps: Heatmaps for keypoints
        sigma: Value used for heatmap generation
    '''
    #keypoints: target person
    #other_keypoints: other people's keypoints need to be add to the heatmap
    heatmaps = []
    keypoints_range = np.reshape(keypoints,(-1,3))
    keypoints_range = keypoints_range[keypoints_range[:,2]>0]
    min_x = np.min(keypoints_range[:,0])
    min_y = np.min(keypoints_range[:,1])
    max_x = np.max(keypoints_range[:,0])
    max_y = np.max(keypoints_range[:,1])
    area = (max_y-min_y)*(max_x-min_x)
    sigma = 3
    if area < 0.16:
        sigma = 3
    elif area < 0.3:
        sigma = 5
    else:
        sigma = 7
    for i in range(0,len(keypoints),3):
        if keypoints[i+2]==0:
            heatmaps.append(np.zeros((img_size//4, img_size//4)))
            continue
        x = int(keypoints[i]*img_size//4) 
        y = int(keypoints[i+1]*img_size//4)
        if x==img_size//4:x=(img_size//4-1)
        if y==img_size//4:y=(img_size//4-1)
        if x>img_size//4 or x<0:x=-1
        if y>img_size//4 or y<0:y=-1
        heatmap = generate_heatmap(x, y, other_keypoints[i//3], (img_size//4, img_size//4),sigma)
        heatmaps.append(heatmap)
    heatmaps = np.array(heatmaps, dtype=np.float32)    
    return heatmaps,sigma

def generate_heatmap(x, y, other_keypoints, size, sigma):
    '''
    Generate a heatmap for a specific keypoint.

    Args:
        x, y: Absolute position of the keypoint
        other_keypoints: Position of other keypoints
        size: Size of the heatmap
        sigma: Value used for heatmap generation

    Returns:
        heatmap: Generated heatmap for the keypoint
    '''
    #x,y  abs postion
    #other_keypoints   positive position
    sigma+=6
    heatmap = np.zeros(size)
    if x<0 or y<0 or x>=size[0] or y>=size[1]:
        return heatmap
    tops = [[x,y]]
    if len(other_keypoints)>0:
        #add other people's keypoints
        for i in range(len(other_keypoints)):
            x = int(other_keypoints[i][0]*size[0])
            y = int(other_keypoints[i][1]*size[1])
            if x==size[0]:x=(size[0]-1)
            if y==size[1]:y=(size[1]-1)
            if x>size[0] or x<0 or  y>size[1] or y<0: continue
            tops.append([x,y])
    for top in tops:
        #heatmap[top[1]][top[0]] = 1
        x,y = top
        x0 = max(0,x-sigma//2)
        x1 = min(size[0],x+sigma//2)
        y0 = max(0,y-sigma//2)
        y1 = min(size[1],y+sigma//2)
        for map_y in range(y0, y1):
            for map_x in range(x0, x1):
                d2 = ((map_x  - x) ** 2 + (map_y  - y) ** 2)**0.5
                if d2<=sigma//2:
                    heatmap[map_y, map_x] += math.exp(-d2/(sigma//2)*3)
                if heatmap[map_y, map_x] > 1:
                    heatmap[map_y, map_x] = 1
    # heatmap[heatmap<0.1] = 0
    return heatmap

def label2center(cx, cy, other_centers, img_size, sigma):
    '''
    Convert labeled keypoints to a center heatmap.

    Args:
        cx, cy: Center coordinates
        other_centers: Other people's centers
        img_size: Size of the image
        sigma: Value used for heatmap generation

    Returns:
        heatmaps: Heatmap representing the center
    '''
    heatmaps = []
    heatmap = generate_heatmap(cx, cy, other_centers, (img_size//4, img_size//4),sigma+2)
    heatmaps.append(heatmap)
    heatmaps = np.array(heatmaps, dtype=np.float32)
    return heatmaps

def label2reg(keypoints, cx, cy, img_size):
    '''
    Convert labeled keypoints to regression maps.

    Args:
        keypoints: Labeled keypoints
        cx, cy: Center coordinates
        img_size: Size of the image

    Returns:
        heatmaps: Regression maps for keypoints
    '''

    heatmaps = np.zeros((len(keypoints)//3*2, img_size//4, img_size//4), dtype=np.float32)
    for i in range(len(keypoints)//3):
        if keypoints[i*3+2]==0:
            continue
        x = keypoints[i*3]*img_size//4
        y = keypoints[i*3+1]*img_size//4
        if x==img_size//4:x=(img_size//4-1)
        if y==img_size//4:y=(img_size//4-1)
        if x>img_size//4 or x<0 or y>img_size//4 or y<0:
            continue
        reg_x = x-cx
        reg_y = y-cy
        for j in range(cy-2,cy+3):
            if j<0 or j>img_size//4-1:
                continue
            for k in range(cx-2,cx+3): 
                if k<0 or k>img_size//4-1:
                    continue
                if cx<img_size//4/2-1:
                    heatmaps[i*2][j][k] = reg_x-(cx-k)#/(img_size//4)
                else:
                    heatmaps[i*2][j][k] = reg_x+(cx-k)#/(img_size//4)
                if cy<img_size//4/2-1:
                    heatmaps[i*2+1][j][k] = reg_y-(cy-j)#/(img_size//4)
                else:
                    heatmaps[i*2+1][j][k] = reg_y+(cy-j)
    return heatmaps

def label2offset(keypoints, cx, cy, regs, img_size):
    '''
    Convert labeled keypoints to offset maps.

    Args:
        keypoints: Labeled keypoints
        cx, cy: Center coordinates
        regs: Regression maps
        img_size: Size of the image

    Returns:
        heatmaps: Offset maps for keypoints
    '''
    heatmaps = np.zeros((len(keypoints)//3*2, img_size//4, img_size//4), dtype=np.float32)
    for i in range(len(keypoints)//3):
        if keypoints[i*3+2]==0:
            continue
        large_x = int(keypoints[i*3]*img_size)
        large_y = int(keypoints[i*3+1]*img_size)
        small_x = int(regs[i*2,cy,cx]+cx)
        small_y = int(regs[i*2+1,cy,cx]+cy)
        offset_x = large_x/4-small_x
        offset_y = large_y/4-small_y
        if small_x==img_size//4:small_x=(img_size//4-1)
        if small_y==img_size//4:small_y=(img_size//4-1)
        if small_x>img_size//4 or small_x<0 or small_y>img_size//4 or small_y<0:
            continue
        heatmaps[i*2][small_y][small_x] = offset_x#/(img_size//4)
        heatmaps[i*2+1][small_y][small_x] = offset_y#/(img_size//4)
    return heatmaps

class TensorDataset(Dataset):
    '''
    Custom Dataset class for handling data loading and preprocessing
    '''

    def __init__(self, data_labels, img_dir, img_size, data_aug=None):
        self.data_labels = data_labels
        self.img_dir = img_dir
        self.data_aug = data_aug
        self.img_size = img_size
        self.interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, 
                                cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]


    def __getitem__(self, index):
        item = self.data_labels[index]
        """
        item = {
                     "img_name":save_name,
                     "keypoints":save_keypoints,
                     "center":save_center,
                     "other_centers":other_centers,
                     "other_keypoints":other_keypoints,
                    }
        """
        # [name,h,w,keypoints...]
        img_path = os.path.join(self.img_dir, item["img_name"])
        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (self.img_size, self.img_size),
                        interpolation=random.choice(self.interp_methods))
        #### Data Augmentation
        if self.data_aug is not None:
            img, item = self.data_aug(img, item)
        img = img.astype(np.float32)
        img = np.transpose(img,axes=[2,0,1])
        keypoints = item["keypoints"]
        center = item['center']
        other_centers = item["other_centers"]
        other_keypoints = item["other_keypoints"]
        kps_mask = np.ones(len(keypoints)//3)
        for i in range(len(keypoints)//3):
            if keypoints[i*3+2]==0:
                kps_mask[i] = 0
        heatmaps,sigma = label2heatmap(keypoints, other_keypoints, self.img_size) #(17, 48, 48)
        cx = min(max(0,int(center[0]*self.img_size//4)),self.img_size//4-1)
        cy = min(max(0,int(center[1]*self.img_size//4)),self.img_size//4-1)
        centers = label2center(cx, cy, other_centers, self.img_size, sigma) #(1, 48, 48)
        regs = label2reg(keypoints, cx, cy, self.img_size) #(14, 48, 48)
        offsets = label2offset(keypoints, cx, cy, regs, self.img_size)#(14, 48, 48)
        labels = np.concatenate([heatmaps,centers,regs,offsets],axis=0)
        img = img / 127.5 - 1.0
        return img, labels, kps_mask, img_path

    def __len__(self):
        return len(self.data_labels)

# Function to get data loader based on mode (e.g., evaluation)
def getDataLoader(mode, input_data):
    '''
    Function to get data loader based on mode (e.g., evaluation).

    Args:
        mode: Mode of data loader (e.g., 'eval')
        input_data: Input data

    Returns:
        data_loader: DataLoader for specified mode
    '''

    if mode=="eval":
        val_loader = torch.utils.data.DataLoader(
                                        TensorDataset(input_data[0],
                                            EVAL_IMG_PATH,
                                            IMG_SIZE,
                                        ),
                                        batch_size=1, 
                                        shuffle=False, 
                                        num_workers=0, 
                                        pin_memory=False)
        return val_loader

# Class for managing data and obtaining evaluation data loader
class Data():
    '''
    Class for managing data and obtaining evaluation data loader.
    '''
    def __init__(self):
        pass

    def getEvalDataloader(self):
        with open(EVAL_LABLE_PATH, 'r') as f:
            data_label_list = json.loads(f.readlines()[0])
        print("[INFO] Total images: ", len(data_label_list))
        input_data = [data_label_list]
        data_loader = getDataLoader("eval", 
                                        input_data)
        return data_loader

# Configs for onnx inference session
def make_parser():
    '''
    Create parser for MoveNet ONNX runtime inference.

    Returns:
        parser: Argument parser for MoveNet inference
    '''
    parser = argparse.ArgumentParser("movenet onnxruntime inference")
    parser.add_argument(
        "--ipu",
        action="store_true",
        help="Use IPU for inference.",
    )
    parser.add_argument(
        "--provider_config",
        type=str,
        default="vaip_config.json",
        help="Path of the config file for seting provider_options.",
    )
    return parser.parse_args()

if __name__ == '__main__':

    args = make_parser()
    
    if args.ipu:
        providers = ["VitisAIExecutionProvider"]
        provider_options = [{"config_file": args.provider_config}]
    else:
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        provider_options = None
    # Get evaluation data loader using the Data class
    data = Data()
    data_loader = data.getEvalDataloader()
    # Load MoveNet model using ONNX runtime
    model = rt.InferenceSession(MODEL_DIR, providers=providers, provider_options=provider_options)
    
    correct = 0
    total = 0
    # Loop through the data loader for evaluation
    for batch_idx, (imgs, labels, kps_mask, img_names) in enumerate(data_loader):
        
        if batch_idx%100 == 0:
            print('Finish ',batch_idx)
    
        imgs = imgs.detach().cpu().numpy()
        imgs = imgs.transpose((0,2,3,1))
        output = model.run(['1548_transpose','1607_transpose','1665_transpose','1723_transpose'],{'blob.1':imgs})
        output[0] = output[0].transpose((0,3,1,2))
        output[1] = output[1].transpose((0,3,1,2))
        output[2] = output[2].transpose((0,3,1,2))
        output[3] = output[3].transpose((0,3,1,2))
        pre = movenetDecode(output, kps_mask,mode='output',img_size=IMG_SIZE)
        gt = movenetDecode(labels, kps_mask,mode='label',img_size=IMG_SIZE)
        
        #n
        acc = myAcc(pre, gt)
        
        correct += sum(acc)
        total += len(acc)
    # Compute and print accuracy based on evaluated data
    acc = correct/total
    print('[Info] acc: {:.3f}% \n'.format(100. * acc))