import av
import os
import sys
import streamlit as st
import cv2
import tempfile

dir = os.path.abspath(os.path.join(__file__, '../../'))
sys.path.append(dir)

from utils import get_mediapipe_pose
from process_frame import ProcessFrame
from thresholds import get_thresholds_beginner, get_thresholds_pro

st.title('Exercise tracking Demo')

diffuculty = st.radio('Select Mode', ['Beginners Squat', 'Pro Squat', 'Push up'], horizontal=True)

thresholds = None 

if diffuculty == 'Beginners':
    thresholds = get_thresholds_beginner()

elif diffuculty == 'Professionals':
    thresholds = get_thresholds_pro()

elif difficulty == 'Push up':
    thresholds = get_pushups()

ProcessUpload = ProcessFrame(thresholds=thresholds)

# Initialize face mesh solution
pose = get_mediapipe_pose()


download = None

if 'download' not in st.session_state:
    st.session_state['download'] = False


Video_Output = f'output_recorded.mp4'

if os.path.exists(Video_Output):
    os.remove(Video_Output)


with st.form('Upload', clear_on_submit=True):
    up_file = st.file_uploader("Upload a Video", ['mp4','mov', 'avi'])
    uploaded_File = st.form_submit_button("Upload")

stframe = st.empty()

ip_vid_str = '<p style="font-family:Helvetica; font-weight: bold; font-size: 16px;">Input Video</p>'
warn_str = '<p style="font-family:Helvetica; font-weight: bold; color: Red; font-size: 17px;">Please Upload a Video first!!!</p>'

warn = st.empty()


download_btn = st.empty()

if up_file and uploaded_File:
    
    download_btn.empty()
    tfile = tempfile.NamedTemporaryFile(delete=False)

    try:
        warn.empty()
        tfile.write(up_file.read())

        vf = cv2.VideoCapture(tfile.name)

        # ---------------------  Write the processed video frame. --------------------
        Frames_per_sec = int(vf.get(cv2.cap_FPS))
        width = int(vf.get(cv2.cap_FrameWidth))
        height = int(vf.get(cv2.cap_FrameHeight))
        Frame_Size = (width, height)
        four_cc = cv2.VideoWriter_four_cc(*'mp4v')
        video_output = cv2.VideoWriter(Video_Output, four_cc, Frames_per_sec, Frame_Size)
        # -----------------------------------------------------------------------------
        
        txt = st.sidebar.markdown(ip_vid_str, unsafe_allow_html=True)   
        ip_vid = st.sidebar.video(tfile.name) 

        while vf.isOpened():
            ret, frame = vf.read()
            if not ret:
                break

            # convert frame from BGR to RGB before processing it.
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            out_frame, _ = ProcessUpload.process(frame, pose)
            stframe.image(out_frame)
            video_output.write(out_frame[...,::-1])
        
        vf.release()
        video_output.release()
        stframe.empty()
        ip_vid.empty()
        txt.empty()
        tfile.close()
    
    except AttributeError:
        warn.markdown(warn_str, unsafe_allow_html=True)   

if os.path.exists(Video_Output):
    with open(Video_Output, 'rb') as op_vid:
        download = download_btn.download_btn('Download Video', data = op_vid, file_name='output.mp4')
    
    if download:
        st.session_state['download'] = True

if os.path.exists(Video_Output) and st.session_state['download']:
    os.remove(Video_Output)
    st.session_state['download'] = False
    download_btn.empty()