from __future__ import annotations

import os
import os.path as osp
from collections import defaultdict
import time
from mmpose.apis.inference import batch_inference_pose_model

import numpy as np
import torch
import torch.nn as nn
import scipy.signal as signal

from ultralytics import YOLO
from mmpose.apis import (
    init_pose_model,
    get_track_id,
    vis_pose_result,
)

ROOT_DIR = osp.abspath(f"{__file__}/../../")
VIT_DIR = osp.join(ROOT_DIR, "third-party/ViTPose")

VIS_THRESH = 0.5
BBOX_CONF = 0.5
TRACKING_THR = 0.1
MINIMUM_FRMAES = 15
MINIMUM_JOINTS = 6

class DetectionModel(object):
    def __init__(self, pose_model_ckpt, device, with_tracker=True):
        
        # ViTPose
        pose_model_cfg = osp.join(VIT_DIR, 'configs/wholebody/2d_kpt_sview_rgb_img/topdown_heatmap/coco-wholebody/ViTPose_huge_wholebody_256x192.py')
        #'vitpose-h-multi-coco.pth')
        self.pose_model = init_pose_model(pose_model_cfg, pose_model_ckpt, device=device)
        
        # YOLO
        bbox_model_ckpt = osp.join(ROOT_DIR, 'checkpoints', 'yolov8x.pt')
        if with_tracker:
            self.bbox_model = YOLO(bbox_model_ckpt)
        else:
            self.bbox_model = None
    
        self.device = device
        self.initialize_tracking()
        
    def initialize_tracking(self, ):
        self.next_id = 0
        self.frame_id = 0
        self.pose_results_last = []
        self.tracking_results = {
            'id': [],
            'frame_id': [],
            'bbox': [],
        }
        
    def xyxy_to_cxcys(self, bbox, s_factor=1.05):
        cx, cy = bbox[[0, 2]].mean(), bbox[[1, 3]].mean()
        scale = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 200 * s_factor
        return np.array([[cx, cy, scale]])
        
    def compute_bboxes_from_keypoints(self, s_factor=1.2):
        X = self.tracking_results['keypoints'].copy()
        mask = X[..., -1] > VIS_THRESH

        bbox = np.zeros((len(X), 3))
        for i, (kp, m) in enumerate(zip(X, mask)):
            bb = [kp[m, 0].min(), kp[m, 1].min(),
                  kp[m, 0].max(), kp[m, 1].max()]
            cx, cy = [(bb[2]+bb[0])/2, (bb[3]+bb[1])/2]
            bb_w = bb[2] - bb[0]
            bb_h = bb[3] - bb[1]
            s = np.stack((bb_w, bb_h)).max()
            bb = np.array((cx, cy, s))
            bbox[i] = bb
        
        bbox[:, 2] = bbox[:, 2] * s_factor / 200.0
        self.tracking_results['bbox'] = bbox
    
    def compute_bbox(self, img):
        bboxes = self.bbox_model.predict(
            img, device=self.device, classes=0, conf=BBOX_CONF, save=False, verbose=False
        )[0].boxes.xyxy.detach().cpu().numpy()

        bboxes = [{'bbox': bbox} for bbox in bboxes]
        imgs = [img for _ in range(len(bboxes))]
        return bboxes, imgs
    
    def batch_detection(self, bboxes, imgs, batch_size=32):
        all_poses = []
        all_bboxes = []
        for i in range(0, len(bboxes), batch_size):
            poses, bbox_xyxy = batch_inference_pose_model(
                self.pose_model,
                imgs[i:i+batch_size],
                bboxes[i:i+batch_size],
                return_heatmap=False)
            all_poses.append(poses)
            all_bboxes.append(bbox_xyxy)
        all_poses = np.concatenate(all_poses)
        all_bboxes = np.concatenate(all_bboxes)
        return all_poses, all_bboxes
        
    def track(self, img, fps, length):
        # bbox detection
        bboxes = self.bbox_model.predict(
            img, device=self.device, classes=0, conf=BBOX_CONF, save=False, verbose=False
        )[0].boxes.xyxy.detach().cpu().numpy()

        pose_results = [{'bbox': bbox} for bbox in bboxes]
        
       
        pose_results, self.next_id = get_track_id(
            pose_results,
            self.pose_results_last,
            self.next_id,
            use_oks=False,
            tracking_thr=TRACKING_THR,
            use_one_euro=True,
            fps=fps)
        
        for pose_result in pose_results:
            
            _id = pose_result['track_id']
            xyxy = pose_result['bbox']
            bbox = xyxy# self.xyxy_to_cxcys(xyxy)
            
            self.tracking_results['id'].append(_id)
            self.tracking_results['frame_id'].append(self.frame_id)
            self.tracking_results['bbox'].append(bbox)
        
        self.frame_id += 1
        self.pose_results_last = pose_results
    
    def process(self, fps):

        for key in ['id', 'frame_id', 'bbox']:
            self.tracking_results[key] = np.array(self.tracking_results[key])
        #self.compute_bboxes_from_keypoints()
            
        output = defaultdict(lambda: defaultdict(list))
        ids = np.unique(self.tracking_results['id'])

        for _id in ids:
            idxs = np.where(self.tracking_results['id'] == _id)[0]

            for key, val in self.tracking_results.items():
                if key == 'id': continue
                output[_id][key] = val[idxs]

        # Smooth bounding box detection
        ids = list(output.keys())
        for _id in ids:
            if len(output[_id]['bbox']) < MINIMUM_FRMAES:
                del output[_id]
                continue
            
            kernel = int(int(fps/2) / 2) * 2 + 1
            smoothed_bbox = np.array([signal.medfilt(param, kernel) for param in output[_id]['bbox'].T]).T
            output[_id]['bbox'] = smoothed_bbox
        
        return output
    
    def visualize(self, img, pose_results):
        vis_img = vis_pose_result(
            self.pose_model,
            img,
            pose_results,
            dataset=self.pose_model.cfg.data['test']['type'],
            dataset_info = None, #self.pose_model.cfg.data['test'].get('dataset_info', None),
            kpt_score_thr=0.3,
            radius=4,
            thickness=1,
            show=False
        )
        return vis_img