import sys
import os
from torchvision import transforms
import torch
import cv2
import numpy as np
import requests

yolov7_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../vendor/yolov7')
sys.path.append(yolov7_path)

from models.experimental import attempt_load
from utils.datasets import letterbox
from utils.general import non_max_suppression_kpt, strip_optimizer, xyxy2xywh
from utils.plots import colors, output_to_keypoint, plot_one_box_kpt, plot_skeleton_kpts
from utils.torch_utils import select_device




HOST = 'localhost'
PORT = '8420'
def predict_http_request(payload):
    """
    Args:
        - payload: {'array': (1, 10, 50) shape (10 frames)}
    Returns:
        - score: Value between 0 and 1
        - status: Good or bad posture (depending on threshold:0.7)
    """
    response = requests.post(f"http://{HOST}:{PORT}/predict", json=payload)
    if response.status_code == 200:
        return response.json()
    else:
        print("Error:", response.status_code)
        print(response.text)


model = None
device = None
def model_initialization(device_ref, w):
    global model, device
    device = select_device(device_ref)
    model = attempt_load(w, map_location=device)
    return model


separation =0
length = 10
frame_count = 0
# logic for multiple persons
people = {}
next_object_id = 0
# logic for single persons
current_sequence = []
current_score = 0
current_status = 'good'
previous_status = "None"
longevity = 0 # frames spent in the current status
separation = 1
multiple = False
frame_count = 0
empty = False


# for audio playing
iterations_in_bad_posture = 0 
max_iterations_in_bad_posture = 5
should_alert = False


@torch.no_grad()
def on_update(frame, recently_alerted, threshold = 0.7):
    global separation, frame_count, current_sequence, empty, current_score, current_status, should_alert,\
            iterations_in_bad_posture, max_iterations_in_bad_posture
    if recently_alerted:
        should_alert = False

    orig_image = frame 
    image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB) 
    image = letterbox(image, (frame.shape[1]), stride=64, auto=True)[0]
    image = transforms.ToTensor()(image)
    image = torch.tensor(np.array([image.numpy()]))
    image = image.to(device)
    image = image.float()

    with torch.no_grad():  #get predictions
        output_data, _ = model(image)

    output_data = non_max_suppression_kpt(output_data,   #Apply non max suppression
                                0.25,   # Conf. Threshold.
                                0.65, # IoU Threshold.
                                nc=model.yaml['nc'], # Number of classes.
                                nkpt=model.yaml['nkpt'], # Number of keypoints.
                                kpt_label=True)

    output = output_to_keypoint(output_data)
    if multiple:
        if len(output) == 0:
            if not empty:
                print("Wiping data, waiting for objects to appear in frame")
            people = {}
            next_object_id = 0
            empty = True
        else:
            empty = False
    else:
        if output.shape[0] > 0:
            if frame_count % separation == 0:
                landmarks = output[0, 7:].T
                current_sequence += [landmarks[:-1]]

            if len(current_sequence) == 10:
                current_sequence = np.array([current_sequence])
                payload = {'array': current_sequence.tolist() }
                response = predict_http_request(payload)

                current_score = response['score']
                response['status'] = 'good' if current_score > threshold else 'bad'
                if response['status'] == 'bad':
                    iterations_in_bad_posture += 1
                current_status = response['status']
                current_sequence = []

    im0 = image[0].permute(1, 2, 0) * 255 # Change format [b, c, h, w] to [h, w, c] for displaying the image.
    im0 = im0.cpu().numpy().astype(np.uint8)
    
    im0 = cv2.cvtColor(im0, cv2.COLOR_RGB2BGR) #reshape image format to (BGR)
    gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh

    for i, pose in enumerate(output_data):  # detections per image
        if empty: break
    
        if len(output_data) == 0:
            continue
        for det_index, (*xyxy, conf, cls) in enumerate(reversed(pose[:,:6])): #loop over poses for drawing on frame
            c = int(cls)  # integer class
            kpts = pose[det_index, 6:]


            if multiple:
                # get the centroid (cx, cy) for the current rectangle
                rect = [tensor.cpu().numpy() for tensor in xyxy]
                cx, cy = (rect[0] + rect[2]) / 2, (rect[1] + rect[3]) / 2
                matched_object_id = None

                # iterating through known people
                for object_id, data in people.items():
                    distance = np.sqrt((cx - data['centroid'][0]) ** 2 + (cy - data['centroid'][1]) ** 2)
                    print(distance)
                    if distance < 300:  # Adjust the threshold as needed
                        matched_object_id = object_id
                        break

                if matched_object_id is None:
                    matched_object_id = next_object_id
                    next_object_id += 1

                if matched_object_id not in people:
                    people[matched_object_id] = {'centroid': (cx, cy), 'yoloid': det_index, 'status': 'good', 'score': 0, 'sequence' : []}
                else:
                    people[matched_object_id]['centroid'] = (cx, cy)
                    people[matched_object_id]['yoloid'] = det_index

                obj = people[matched_object_id]
                label = f"ID #{obj['yoloid']} Score: {obj['score']:.2f}"
                plot_one_box_kpt(xyxy, im0, label=label, color=colors(c, True), 
                            line_thickness=3, kpt_label=True, kpts=kpts, steps=3, 
                            cmap=people[matched_object_id]['status'])
            else:
                label = f"ID #{0} Score: {current_score:.2f}"
                plot_one_box_kpt(xyxy, im0, label=label, color=colors(c, True), 
                            line_thickness=3,kpt_label=True, kpts=kpts, steps=3, 
                            cmap=current_status)


    if frame_count % separation == 0 and multiple:
        for _, data in people.items():
            if data['yoloid'] < output.shape[0]:
                yoloid = data['yoloid']
                landmarks = output[yoloid, 7:].T
                data['sequence'] += [landmarks[:-1]]
            
                if len(data['sequence']) == length:
                    payload = {'array': np.array([data['sequence']]).tolist()}
                    response = predict_http_request(payload)

                    data['score'] = response['score']
                    data['status'] = 'good' if score > THRESHOLD else 'bad'
                    data['sequence'] = []

                # print(f"{data['yoloid']} -> {data['status']}", end=' ')
            else:
                data['sequence'] = []

        
        statuses = [(people[p]['yoloid'], people[p]['status']) for p in people]
        # for id, status in statuses:
        #     print(f'{id}: {status}', end='\t')
        # print()
    if iterations_in_bad_posture >= max_iterations_in_bad_posture:
        iterations_in_bad_posture = 0
        should_alert = True

    frame_count += 1
    return im0, current_status, current_score, should_alert