import cv2
import torch

import data.MPII.ref as ds
import myutils.img
from myutils.group import HeatmapParser
from myutils.posture import *

# 歪头的斜率阈值
CROOKED_HEAD_THRE=8
# 斜肩的斜率阈值
OBLIQUE_SHOULDER_THRE=2

# 人体骨骼各坐标点的下标对应图

class PersonPosture():
    def __init__(self, pred,orig_img):
        self.pred = pred[0]['keypoints']
        self.orin_img = orig_img
    def analysis(self, visual_angle=True):
        # 肩斜率
        Shoulder_slope = segment_slope(self.pred[12], self.pred[13])
        head_slope = segment_slope(self.pred[8], self.pred[9])
        # 左腿的大腿根和膝盖斜率
        left_thigh_knee = segment_slope(self.pred[1], self.pred[2])
        # 左腿脚跟和膝盖斜率
        left_calf_knee = segment_slope(self.pred[1], self.pred[0])
        # 左腿大腿根和脚跟的斜率
        left_lag_slope = segment_slope(self.pred[2], self.pred[0])
        # 右腿的大腿根和膝盖斜率
        right_lag_slope = segment_slope(self.pred[3], self.pred[5])
        # 右腿脚跟和膝盖斜率
        right_thigh_knee = segment_slope(self.pred[3], self.pred[4])
        # 右腿大腿根和脚跟的斜率
        right_calf_knee = segment_slope(self.pred[4], self.pred[5])


        # 正面照
        if visual_angle:
            # 两个膝盖的间距
            knee_distance = absolute_distance(self.pred[4],self.pred[1])

            # 两个脚的间距
            feet_distance = absolute_distance(self.pred[0],self.pred[5])

            # 检查斜肩
            if Shoulder_slope!=None and abs(Shoulder_slope) > OBLIQUE_SHOULDER_THRE:
                cv2.line(self.orin_img,tuple(map(int,self.pred[12][:2])),tuple(map(int,self.pred[13][:2])),(0,0,255),thickness=3)
                print("斜肩")
            # 检查O型腿
            flag = 0
            if feet_distance != None and knee_distance != None:
                if feet_distance / knee_distance < 0.8:
                    cv2.line(self.orin_img, tuple(map(int,self.pred[2][:2])), tuple(map(int,self.pred[1][:2])), (0, 0, 255), thickness=3)
                    cv2.line(self.orin_img, tuple(map(int,self.pred[1][:2])), tuple(map(int,self.pred[0][:2])), (0, 0, 255), thickness=3)
                    cv2.line(self.orin_img, tuple(map(int,self.pred[3][:2])), tuple(map(int,self.pred[4][:2])), (0, 0, 255), thickness=3)
                    cv2.line(self.orin_img, tuple(map(int,self.pred[4][:2])), tuple(map(int,self.pred[5][:2])), (0, 0, 255), thickness=3)
                    print("膝盖间距较大，O形腿")
                    flag = 1
            # 检查 X型腿
            # 如果两个膝盖和两个脚四个点都检测到了
            if feet_distance!=None and knee_distance!=None:
                # 如果脚之间的距离比膝盖的距离大于1.3,(假设用户是自然笔直站立，而不是双脚叉开站立)，则判定为X型腿
                if feet_distance/knee_distance>1.3:
                    cv2.line(self.orin_img, tuple(map(int,self.pred[2][:2])), tuple(map(int,self.pred[1][:2])), (0, 0, 255), thickness=3)
                    cv2.line(self.orin_img, tuple(map(int,self.pred[1][:2])), tuple(map(int,self.pred[0][:2])), (0, 0, 255), thickness=3)
                    cv2.line(self.orin_img, tuple(map(int,self.pred[3][:2])), tuple(map(int,self.pred[4][:2])), (0, 0, 255), thickness=3)
                    cv2.line(self.orin_img, tuple(map(int,self.pred[4][:2])), tuple(map(int,self.pred[5][:2])), (0, 0, 255), thickness=3)
                    print("双脚间距较大,X形腿")
                    flag = 1
            if flag == 0:
                print("腿部没有检测到O型腿或X型腿，正常 ")
            # 检查歪头
            if head_slope!=None and abs(head_slope)<CROOKED_HEAD_THRE:
                cv2.line(self.orin_img, self.pred[8][:2], self.pred[9][:2], (0, 0, 255), thickness=3)
                print("歪头")
            else:
                print("头部正常，无明显歪头现象")
        # 侧面照
        else:
            # 检查头部前倾
            if head_slope != np.Inf and  0 < head_slope < CROOKED_HEAD_THRE:
                cv2.line(self.orin_img, tuple(map(int,self.pred[8][:2])), tuple(map(int,self.pred[9][:2])), (0, 0, 255), thickness=3)
                print("头部前倾")
            else:
                print("头部正常，无明显倾斜")
            # 检查膝盖超伸
            if ((left_lag_slope!=np.Inf or left_calf_knee!=np.Inf or left_thigh_knee!=np.Inf)
                or (right_lag_slope!=np.Inf or right_calf_knee!=np.Inf or right_thigh_knee!=np.Inf)):

                cv2.line(self.orin_img, tuple(map(int,self.pred[2][:2])), tuple(map(int,self.pred[1][:2])), (0, 0, 255), thickness=3)
                cv2.line(self.orin_img, tuple(map(int,self.pred[1][:2])), tuple(map(int,self.pred[0][:2])), (0, 0, 255), thickness=3)
                cv2.line(self.orin_img, tuple(map(int,self.pred[3][:2])), tuple(map(int,self.pred[4][:2])), (0, 0, 255), thickness=3)
                cv2.line(self.orin_img, tuple(map(int,self.pred[4][:2])), tuple(map(int,self.pred[5][:2])), (0, 0, 255), thickness=3)
                print("膝盖超伸")
            else:
                print("膝盖正常，无明显超伸")
        return self.orin_img


parser = HeatmapParser()
def post_process(det, mat_, trainval, c=None, s=None, resolution=None):
    mat = np.linalg.pinv(np.array(mat_).tolist() + [[0, 0, 1]])[:2]
    res = det.shape[1:3]
    cropped_preds = parser.parse(np.float32([det]))[0]

    if len(cropped_preds) > 0:
        cropped_preds[:, :, :2] = myutils.img.kpt_affine(cropped_preds[:, :, :2] * 4, mat)  # size 1x16x3

    preds = np.copy(cropped_preds)
    ##for inverting predictions from input res on cropped to original image
    if trainval != 'cropped':
        for j in range(preds.shape[1]):
            preds[0, j, :2] = myutils.img.transform(preds[0, j, :2], c, s, resolution, invert=1)
    return preds


def inference(img, func, config, c, s):
    """
    forward pass at test time
    calls post_process to post process results
    """
    # 256,256
    # 128.0,128.0
    # 1.28
    # (256,256)
    height, width = img.shape[0:2]
    center = (width / 2, height / 2)
    scale = max(height, width) / 200
    res = (config['train']['input_res'], config['train']['input_res'])

    # [[1,0,0],[0,1,0]]
    mat_ = myutils.img.get_transform(center, scale, res)[:2]
    inp = img / 255

    def array2dict(tmp):
        return {
            'det': tmp[0][:, :, :16],
        }

    tmp1 = array2dict(func([inp]))
    tmp2 = array2dict(func([inp[:, ::-1]]))

    tmp = {}
    for ii in tmp1:
        tmp[ii] = np.concatenate((tmp1[ii], tmp2[ii]), axis=0)

    det = tmp['det'][0, -1] + tmp['det'][1, -1, :, :, ::-1][ds.flipped_parts['mpii']]
    if det is None:
        return [], []
    det = det / 2

    det = np.minimum(det, 1)

    return post_process(det, mat_, 'valid', c, s, res)
def main():
    from train import init
    func, config = init()

    def runner(imgs):
        return func(0, config, 'inference', imgs=torch.Tensor(np.float32(imgs)))['preds']

    def do(img, c, s):
        ans = inference(img, runner, config, c, s)
        if len(ans) > 0:
            ans = ans[:,:,:3]

        ## ans has shape N,16,3 (num preds, joints, x/y/visible)
        pred = []
        for i in range(ans.shape[0]):
            pred.append({'keypoints': ans[i,:,:]})
        return pred
if __name__ == '__main__':
    image_path = "data/custom/O型腿3.jpeg"
    from train import init
    func, config = init()

    def runner(imgs):
        return func(0, config, 'inference', imgs=torch.Tensor(np.float32(imgs)))['preds']
    def do(img, c, s):
        ans = inference(img, runner, config, c, s)
        if len(ans) > 0:
            ans = ans[:,:,:3]
        ## ans has shape N,16,3 (num preds, joints, x/y/visible)
        pred = []
        for i in range(ans.shape[0]):
            pred.append({'keypoints': ans[i,:,:]})
        return pred

    input_res = 256
    orig_img = cv2.imread(image_path)
    orig_img_reverse = cv2.imread(image_path)[:,:,::-1]

    shape = orig_img_reverse.shape[0:2]

    # 这里需要使用yolo检测图片中的人体矩形区域，返回矩形的中心和矩形较长边的边长/200
    # 这里要起一个yolo5的服务。因为yolo5的模型的相对路径已经定死了，无法直接引入到项目中。除非通过命令行的方式

    # import requests
    # DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s"
    # image_data = open(image_path, "rb").read()
    # response = requests.post(DETECTION_URL, files={"image": image_data}).json()
    # for box in response:
    #     if box['class']!=0:
    #         continue
    #     else:
    #         c = [(box['xmax']+box['xmin'])/2,(box['ymax']+box['ymin'])/2]
    #         s = max((box['xmax']-box['xmin'])/200,(box['ymax']-box['ymin'])/200)+1

    # Model
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
    response = model(image_path).pandas().xyxy[0].to_dict(orient='records')
    for box in response:
        if box['class']!=0:
            continue
        else:
            c = [(box['xmax']+box['xmin'])/2,(box['ymax']+box['ymin'])/2]
            s = max((box['xmax']-box['xmin'])/200,(box['ymax']-box['ymin'])/200)+1

    # # 矩形中心坐标点
    # # 缩放比例  实验结果：720*1280的图片,人物占据大部分的情况下，7比较好。
    # s = 7
    im = myutils.img.crop(orig_img_reverse,c,s,(input_res,input_res))


    pred = do(im,c,s)
    points = pred[0]['keypoints']
    for index,point in enumerate(points):
        print(point[2])
        if point[2]>0.3:
            cv2.circle(orig_img, (int(point[0]), int(point[1])), 5, color=(0, 0, 255), thickness=-1)
            cv2.putText(orig_img,str(index),(int(point[0]), int(point[1])),cv2.FONT_ITALIC, 1, (0,255,0), 1)


    person = PersonPosture(pred,orig_img)
    orig_img = person.analysis(True)
    cv2.imshow("test", orig_img)
    cv2.waitKey()

