import streamlit as st
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
import torch
import utils


@st.cache_resource()
def load_model():
    print('Loading model...')
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model_pose = YOLO('yolov8l-pose.pt')
    model_pose.to(device)
    return model_pose


def draw_output(image_pil: Image.Image, keypoints: dict):    
    draw = ImageDraw.Draw(image_pil)
    line_width = 10
    font = ImageFont.truetype("DejaVuSerif-Bold.ttf", 70)
    
    ear, eye = None, None
    if keypoints["left_ear"] and keypoints["left_eye"]:
        ear = keypoints["left_ear"]
        eye = keypoints["left_eye"]
    elif keypoints["right_ear"] and keypoints["right_eye"]:
        ear = keypoints["right_ear"]
        eye = keypoints["right_eye"]
    
    # draw extended left and right eye lines
    if ear and eye:
        left_new_point = utils.extend_line(ear, eye, 3)
        l1 = [ear, left_new_point]
        draw.line(l1, fill='red', width=line_width)
        # draw a horizontal line from ear forwards
        ear = np.array(ear)
        l1 = np.array(l1)
        l1_vector = l1[1] - l1[0]
        x_s = np.sign(l1_vector)[0]
        length_l1 = np.linalg.norm(l1_vector)
        p2 = ear + np.array([length_l1*x_s, 0])
        ear = tuple(ear.tolist())
        l = [ear, tuple(p2.tolist())]
        draw.line(l, fill='gray', width=line_width//2)
        # draw angle
        angle = utils.calculate_angle_to_horizontal(l1_vector)
        draw.text(ear, f'{angle:.2f}', fill='red', font=font)

    # draw elbow angles
    left_elbow_angle, right_elbow_angle = utils.get_elbow_angles(keypoints)
    if left_elbow_angle:
        draw.text(keypoints['left_elbow'], f'{left_elbow_angle:.2f}', fill='red', font=font)
        # draw polyline for left arm
        draw.line([keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist']], fill='blue', width=line_width)
    if right_elbow_angle:
        draw.text(keypoints['right_elbow'], f'{right_elbow_angle:.2f}', fill='red', font=font)
        # draw polyline for right arm
        draw.line([keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist']], fill='blue', width=line_width)

    return image_pil


st.title('Pose Estimation App')

device = 'cuda' if torch.cuda.is_available() else 'cpu'
st.caption(f'Using device: {device}')

mode = st.radio('Select mode:', ['Upload an Image', 'Webcam Capture'])

if mode == 'Upload an Image':
    img_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
elif mode == 'Webcam Capture':
    img_file = st.camera_input("Take a picture")

img = None
if img_file is not None:
    img = Image.open(img_file)

st.divider()
if img is not None:
    # predict
    with st.spinner('Predicting...'):
        model = load_model()
        pred = model(img)[0]
    st.markdown('**Results:**')
    keypoints = utils.get_keypoints(pred)
    if keypoints is not None:
        img = draw_output(img, keypoints)
        st.image(img, caption='Predicted image', use_column_width=True)
        lea, rea = utils.get_eye_angles(keypoints)
        lba, rba = utils.get_elbow_angles(keypoints)
        st.write('Angles:')
        st.json({'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba})
        st.write('Raw keypoints:')
        st.json(keypoints)
    else:
        st.error('No keypoints detected!')
        st.image(img, caption='Original image', use_column_width=True)