# Copyright (c) OpenMMLab. All rights reserved.
import logging
import time

import cv2
from argparse import ArgumentParser

from mmengine.logging import print_log
from mmpose.apis import inference_topdown, init_model
from mmpose.registry import VISUALIZERS
from mmpose.structures import merge_data_samples
import numpy as np

class Position:
    def __init__(self, person_kpts):
        if len(person_kpts) == 8:
            self.left_elbow_angle = person_kpts[0]
            self.right_elbow_angle = person_kpts[1]
            self.left_shoulder_angle = person_kpts[2]
            self.right_shoulder_angle = person_kpts[3]
            self.left_leg_angle = person_kpts[4]
            self.right_leg_angle = person_kpts[5]
            self.left_hip_angle = person_kpts[6]
            self.right_hip_angle = person_kpts[7]
        else:
            self.left_hip = person_kpts[11]
            self.left_knee = person_kpts[13]
            self.left_ankle = person_kpts[15]

            self.right_hip = person_kpts[12]
            self.right_knee = person_kpts[14]
            self.right_ankle = person_kpts[16]

            self.left_shoulder = person_kpts[5]
            self.left_elbow = person_kpts[7]
            self.left_wrist = person_kpts[9]

            self.right_shoulder = person_kpts[6]
            self.right_elbow = person_kpts[8]
            self.right_wrist = person_kpts[10]

            self.left_elbow_angle = angle_between(self.left_shoulder - self.left_elbow, self.left_wrist - self.left_elbow)
            self.right_elbow_angle = angle_between(self.right_shoulder - self.right_elbow, self.right_wrist - self.right_elbow)

            self.left_shoulder_angle = angle_between(self.right_shoulder - self.left_shoulder, self.left_elbow - self.left_shoulder)
            self.right_shoulder_angle = angle_between(self.left_shoulder - self.right_shoulder, self.right_elbow - self.right_shoulder)

            self.left_leg_angle = angle_between(self.left_hip - self.left_knee, self.left_ankle - self.left_knee)
            self.right_leg_angle = angle_between(self.right_hip - self.right_knee, self.right_ankle - self.right_knee)

            self.left_hip_angle = angle_between(self.left_shoulder - self.left_hip, self.left_knee - self.left_hip)
            self.right_hip_angle = angle_between(self.right_shoulder - self.right_hip, self.right_knee - self.right_hip)

    def to_array(self):
        return [
            self.left_elbow_angle,
            self.right_elbow_angle,
            self.left_shoulder_angle,
            self.right_shoulder_angle,
            self.left_leg_angle,
            self.right_leg_angle,
            self.left_hip_angle,
            self.right_hip_angle
        ]

    def distance_to(self, position):
        error = 0
        success = True

        for i, x in enumerate(self.to_array()):
            if 3 < i < 6:
                continue

            y = position.to_array()[i]

            if abs(y) > 165:
                x = abs(x)
                y = abs(y)

            dist = abs(y - x)

            if dist > 20:
                success = False
                # print(f"{i} nie jest ok: moje: {x}, cel: {y}")
                error += abs(y - x)
            else:
                pass
                # print(f"{i} jest ok: moje: {x}, cel: {position.to_array()[i]}")

        return (success, error)

def parse_args():
    parser = ArgumentParser()
    # parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--draw-heatmap',
        action='store_true',
        help='Visualize the predicted heatmap')
    parser.add_argument(
        '--show-kpt-idx',
        action='store_true',
        default=False,
        help='Whether to show the index of keypoints')
    parser.add_argument(
        '--skeleton-style',
        default='mmpose',
        type=str,
        choices=['mmpose', 'openpose'],
        help='Skeleton style selection')
    parser.add_argument(
        '--kpt-thr',
        type=float,
        default=0.3,
        help='Visualizing keypoint thresholds')
    parser.add_argument(
        '--radius',
        type=int,
        default=3,
        help='Keypoint radius for visualization')
    parser.add_argument(
        '--thickness',
        type=int,
        default=1,
        help='Link thickness for visualization')
    parser.add_argument(
        '--alpha', type=float, default=0.8, help='The transparency of bboxes')
    parser.add_argument(
        '--camera-id', type=int, default=0, help='Camera device ID')
    args = parser.parse_args()
    return args


def angle_between(v1, v2):
    """
    Liczy kąt między wektorami v1 i v2 w stopniach.
    Znak kąta zależy od kierunku (przeciwnie do ruchu wskazówek zegara jest dodatni).
    """
    # kąt w radianach
    angle = np.arccos(np.clip(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)), -1.0, 1.0))

    # obliczamy znak kąta w 2D: + jeśli v2 jest "po lewej" od v1
    sign = np.sign(v1[0] * v2[1] - v1[1] * v2[0])

    return np.degrees(angle) * sign

def is_visible(kpt, threshold=0.3):
    # kpt = [x, y, score]
    return kpt[2] > threshold

def main():
    args = parse_args()

    # build the model from a config file and a checkpoint file
    if args.draw_heatmap:
        cfg_options = dict(model=dict(test_cfg=dict(output_heatmaps=True)))
    else:
        cfg_options = None

    model = init_model(
        "mmpose/configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w48_8xb32-210e_coco-256x192.py",
        "hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth",
        device=args.device,
        cfg_options=cfg_options)

    # init visualizer
    model.cfg.visualizer.radius = args.radius
    model.cfg.visualizer.alpha = args.alpha
    model.cfg.visualizer.line_width = args.thickness

    visualizer = VISUALIZERS.build(model.cfg.visualizer)
    visualizer.set_dataset_meta(
        model.dataset_meta, skeleton_style=args.skeleton_style)

    # start capturing video from camera
    cap = cv2.VideoCapture(args.camera_id)
    if not cap.isOpened():
        print("Error: Cannot open camera")
        return

    lastDegs = []
    lastMove = 0
    visible = 0
    maha = 0
    start = False

    oldpos = None
    waiting = 100

    positions = []
    numPositions = 0
    tPose = Position(
        [
            167.5568,
            -161.67027,
            -166.49443,
            168.22028,
            110.21745,
            166.41733,
            167.57822,
            -176.08066
        ]
    )

    krakowiak = Position(
        [np.float32(-98.226715), np.float32(106.39389), np.float32(-135.05656), np.float32(139.48904), np.float32(-149.9036), np.float32(8.216028), np.float32(174.70923), np.float32(-176.893)]
    )

    krakowiakRight = Position(
        [np.float32(-110.61421), np.float32(-174.2424), np.float32(-127.05916), np.float32(174.9463),
         np.float32(175.62007), np.float32(166.89127), np.float32(168.8219), np.float32(-178.02744)]
    )

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Error: Cannot read frame from camera")
            break

        # convert BGR to RGB
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # inference
        batch_results = inference_topdown(model, img)
        results = merge_data_samples(batch_results)

        person_kpts = results.pred_instances.keypoints[0]
        person_visible = results.pred_instances.keypoints_visible[0]

        left_hand_visible = person_visible[5] > 0.75 and person_visible[5] > 0.75 and person_visible[9] > 0.75
        right_hand_visible = person_visible[6] > 0.75 and person_visible[8] > 0.75 and person_visible[10] > 0.75

        left_leg_visible = person_visible[11] > 0.75 and person_visible[13] > 0.75 and person_visible[15] > 0.75
        right_leg_visible = person_visible[12] > 0.75 and person_visible[14] > 0.75 and person_visible[16] > 0.75

        position = Position(person_kpts)
        #
        if position.distance_to(tPose)[0]:
            print("\rT POSE!", end="")
        elif position.distance_to(krakowiak)[0]:
            print("\rKrakowiak", end="")
        elif position.distance_to(krakowiakRight)[0]:
            print("\rKrakowiak right", end="")
        else:
            print("\rNIC!", end="")

        # print(position.left_elbow_angle)
        # if oldpos is None:
        #     oldpos = position
        #
        # if oldpos.distance_to(position) < 7.5:
        #     print(f"\r{goalPosition.distance_to(position)}", end="")
        #
        # oldpos = position

        # if error > 200:

        # if start:
        #     if numPositions > 100:
        #         avgPosition = [0, 0, 0, 0, 0, 0, 0, 0]
        #
        #         for element in positions:
        #             for i, data in enumerate(element):
        #                 avgPosition[i] += data
        #
        #         for i, element in enumerate(avgPosition):
        #             avgPosition[i] = avgPosition[i] / numPositions
        #
        #         print(avgPosition)
        #
        #         break
        #     else:
        #         print(f"\r{numPositions}", end="")
        #         if oldpos is None:
        #             oldpos = position
        #
        #         if oldpos.distance_to(position)[0]:
        #             positions.insert(0, position.to_array())
        #             numPositions += 1
        #
        #         oldpos = position
        #
        # # else:
        # #     print(f"\rOK!", end="")
        #
        # if waiting != 0 and not start:
        #     print(f"\r{waiting}", end="")
        #     waiting -= 1
        # else:
        #     start = True

        # lastDegs.insert(0, (left_elbow_angle, right_elbow_angle))
        #
        # last = 0
        # lastCount = 0
        #
        # for element in lastDegs:
        #     last += element[1]
        #     lastCount += 1
        #
        # last = last / lastCount
        # dist = right_elbow_angle - last
        #
        # if not right_visible:
        #     visible = 0
        #     print("\rNie widać prawej ręki!!!!!!!!", end="")
        # else:
        #     if maha == 0:
        #         print("\rWidać rękę, nie maha!", end="")
        #     else:
        #         maha -= 1
        #
        #     visible += 1
        #
        # if 15 < abs(dist) < 60 and visible > 5:
        #     if lastMove != dist > 0:
        #         maha = 10
        #         print("\rmaha!", end="")
        #
        #     lastMove = dist > 0
        #
        # if len(lastDegs) > 5:
        #     lastDegs.pop()

        # visualize
        vis_img = visualizer.add_datasample(
            'result',
            img,
            data_sample=results,
            draw_gt=False,
            draw_bbox=True,
            kpt_thr=args.kpt_thr,
            draw_heatmap=args.draw_heatmap,
            show_kpt_idx=args.show_kpt_idx,
            skeleton_style=args.skeleton_style,
            show=False,
            out_file=None)

        # convert RGB back to BGR for OpenCV
        vis_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
        cv2.imshow('Live Pose Estimation', vis_img)

        # press 'q' to quit
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()