import cv2
import mediapipe as mp
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import LSTM
import streamlit as st

labels = np.array(['FALL', 'LYING', 'SIT', 'STAND', 'MOVE'])

n_time_steps = 25
mpPose = mp.solutions.pose
pose = mpPose.Pose()
mpDraw = mp.solutions.drawing_utils

def custom_lstm(*args, **kwargs):
    kwargs.pop('time_major', None)
    return LSTM(*args, **kwargs)

model = tf.keras.models.load_model('bro.h5', custom_objects={'LSTM': custom_lstm})

def make_landmark_timestep(results):
    c_lm = []
    for id, lm in enumerate(results.pose_landmarks.landmark):
        c_lm.append(lm.x)
        c_lm.append(lm.y)
        c_lm.append(lm.z)
        c_lm.append(lm.visibility)
    return c_lm

def draw_landmark_on_image(mpDraw, results, img, label):
    mpDraw.draw_landmarks(img, results.pose_landmarks, mpPose.POSE_CONNECTIONS)
    for id, lm in enumerate(results.pose_landmarks.landmark):
        h, w, c = img.shape
        cx, cy = int(lm.x * w), int(lm.y * h)
        if label != "FALL":
            cv2.circle(img, (cx, cy), 5, (0, 255, 0), cv2.FILLED)
        else:
            cv2.circle(img, (cx, cy), 5, (0, 0, 255), cv2.FILLED)
    return img

def draw_class_on_image(label, img):
    font = cv2.FONT_HERSHEY_SIMPLEX
    bottomLeftCornerOfText = (10, 30)
    fontScale = 1
    fontColor = (0, 255, 0)
    thickness = 2
    lineType = 2
    cv2.putText(img, label,
                bottomLeftCornerOfText,
                font,
                fontScale,
                fontColor,
                thickness,
                lineType)
    return img

def detect(model, lm_list):
    lm_list = np.array(lm_list)
    lm_list = np.expand_dims(lm_list, axis=0)
    results = model.predict(lm_list)
    if results[0][0] >= 0.5: 
        label = labels[0]
    elif results[0][1] >= 0.5:
        label = labels[1]  
    elif results[0][2] >= 0.5:
        label = labels[2]
    elif results[0][3] >= 0.5:
        label = labels[3]
    elif results[0][4] >= 0.5:
        label = labels[4]
    else:
        label = "NONE DETECTION"
    return label

def main():
    st.title("Pose Detection and Classification")
    
    run_type = st.sidebar.selectbox("Select input type", ("Camera", "Video File"))
    
    if run_type == "Camera":
        cap = cv2.VideoCapture(0)
    else:
        video_file = st.sidebar.file_uploader("Upload a video", type=["mp4", "mov", "avi"])
        if video_file is not None:
            # Temporarily save the uploaded video to disk to pass to cv2.VideoCapture
            with open("temp_video.mp4", "wb") as f:
                f.write(video_file.read())
            cap = cv2.VideoCapture("temp_video.mp4")
        else:
            st.write("Please upload a video file.")
            return
    
    stframe = st.empty()
    label = 'Starting...'
    lm_list = []

    while cap.isOpened():
        success, img = cap.read()
        if not success:
            break
        
        imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        results = pose.process(imgRGB)
        
        if results.pose_landmarks:
            c_lm = make_landmark_timestep(results)
            img = draw_landmark_on_image(mpDraw, results, img, label)
            img = draw_class_on_image(label, img)
            lm_list.append(c_lm)
            if len(lm_list) == n_time_steps:
                label = detect(model, lm_list)
                lm_list = []
        
        stframe.image(img, channels="BGR")
        
        if cv2.waitKey(1) == ord('q'):
            break
    
    cap.release()

if __name__ == '__main__':
    main()