import streamlit as st
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import cv2
import numpy as np

import logging

@st.cache_resource
def get_model():
    repo_id = "danbiagini/hockey_breeds_v2"
    return hf_hub_download(repo_id=repo_id, filename="hockey_breeds-v2-101623.pt")


def run_inference(img, model, thresh=0.5):
    model = YOLO(model_f)
    st.session_state.results = model(img)
    return draw_hockey_boxes(img, st.session_state.results, thresh)


def draw_hockey_boxes(frame, results, thresh=0.5):
    colors = {0: (0, 255, 0), 1: (255, 0, 0), 2: (0, 0, 255), 3: (128, 0, 0), 4: (
        0, 128, 0), 5: (0, 0, 128), 6: (0, 64, 0), 7: (64, 0, 0), 8: (0, 0, 64)}
    font_scale = frame.shape[0] / 500
    objects = []

    for name in results:
        for box in name.boxes.data.tolist():
            x1, y1, x2, y2, score, class_id = box
            objects.append((name.names[int(class_id)], score))

            if score > thresh:
                cv2.rectangle(frame, (int(x1), int(y1)),
                              (int(x2), int(y2)), colors[(class_id % 9)], 3)
                cv2.putText(frame, f'{name.names[int(class_id)].upper()}: {score:.2f}', (int(x1), int(y1 - 10)),
                  cv2.FONT_HERSHEY_SIMPLEX, font_scale, colors[(class_id % 9)], 3, cv2.LINE_AA)
            else:
                print(
                    f'Found an object under confidence threshold {thresh} type: {name.names[class_id]}, score:{score}, x1, y2:{x1}, {y2}')
    return objects

def reset_image():
    st.session_state.img = None

def upload_img():
    if st.session_state.upload_img is not None:
        st.session_state.img = st.session_state.upload_img

def get_naked_image():
    if st.session_state.img is not None:
        img = st.session_state.img
        img.seek(0)
        return(cv2.imdecode(np.frombuffer(img.read(), np.uint8), 1))
    return None

def use_sample_image():
    st.session_state.img = open('src/images/samples/v2/net-chaos.jpg', 'rb')

# Init state
if 'results' not in st.session_state:
    st.session_state.results = []

if 'thresh' not in st.session_state:
    st.session_state.thresh = 0.5

if 'img' not in st.session_state:
    st.session_state.img = None


st.set_page_config(page_title='TopShelf POC', layout="wide",
                   page_icon="🥅")

st.title('Welcome To Top Shelf :goal_net:',
         help=':video_camera: + :ice_hockey_stick_and_puck: = :clipboard:')
st.subheader('Artificial Intelligence for Hockey Coaches and Players',
             help='Proof of concept application')

overview = '''**Top Shelf** helps coaches and players analyze their gameplay, providing helpful suggestions & recommendations on areas for improvement.  

We're starting with a focus on ice hockey, however this same technology could apply to other "invasion" games and sports, for example lacrosse, basketball, soccer, etc.

The special sauce behind **Top Shelf** is AI *Computer Vision*  technology that recognizes various hockey related objects in videos. 
The foundation of the technology is an AI model that can recognize players, nets, referees, pucks, and rink areas. 

**Top Shelf** uses this technology to analyze game play and provide insightful suggestions on areas for improvement. 
'''
st.markdown(overview)

st.subheader('Getting Started')
st.markdown('''We're currently in the training and testing phase of **Top Shelf** development.  This is a proof of concept application that friends of **Top Shelf** can use to help in development. 
To help us understand how our *Computer Vision* model is working you can upload hockey pictures and then the app will display what hockey objects were found. ''')

st.write("Upload an image file to try detecting hockey objects in your own hockey image, or use a sample image below.")


if st.session_state.img is None:
    st.file_uploader("Upload an image and Hockey Breeds v2 will find the hockey objects in the image",
            type=["jpg", "jpeg", "png"], key='upload_img', on_change=upload_img)

    with st.expander("Sample Images"):
        st.image('src/images/samples/v2/net-chaos.jpg')
        st.button("Use Sample", on_click=use_sample_image)

img = get_naked_image()
if img is not None:

    thresh = st.slider('Set the object confidence threshold', key='thresh',
                  min_value=0.0, max_value=1.0, value=0.5, step=0.05)

    with st.status("Detecting hockey objects..."):
        st.write("Loading model...")
        model_f = get_model()

        st.write("Running inference on image...")
        objects = run_inference(img, model_f, thresh)
        
        st.dataframe(objects, column_config={
                "0": "Object",
                "1": "Confidence"
        })

    # check if the results list is empty
    if len(st.session_state.results) == 0:
        st.write('**No hockey objects found in image!**')
        st.image(img, caption='Uploaded Image had no hockey objects')
    else:
        st.image(img, caption='Image with hockey object bounding boxes')

    st.button("Reset Image", on_click=reset_image)