import streamlit as st
import cv2
import numpy as np
import mediapipe as mp
import joblib
import pandas as pd
from numpy.linalg import norm
import matplotlib.pyplot as plt
import os

# Function to load the Random Forest model
@st.cache_resource
def load_model():
    try:
        return joblib.load('best_random_forest_model.pkl')
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

# Load the model using the cached function
model = load_model()

# Ensure the model is loaded before proceeding
if model is None:
    st.stop()

# Initialize MediaPipe Hands
@st.cache_resource
def load_mediapipe_model():
    mp_hands = mp.solutions.hands
    return mp_hands.Hands(static_image_mode=True, max_num_hands=1, min_detection_confidence=0.5)

hands = load_mediapipe_model()
mp_drawing = mp.solutions.drawing_utils

# Function to normalize landmarks
def normalize_landmarks(landmarks):
    center = np.mean(landmarks, axis=0)
    landmarks_centered = landmarks - center
    std_dev = np.std(landmarks_centered, axis=0)
    landmarks_normalized = landmarks_centered / std_dev
    return np.nan_to_num(landmarks_normalized)

# Function to calculate angles between landmarks
def calculate_angles(landmarks):
    angles = []
    for i in range(20):
        for j in range(i + 1, 21):
            vector = landmarks[j] - landmarks[i]
            angle_x = np.arccos(np.clip(vector[0] / norm(vector), -1.0, 1.0))
            angle_y = np.arccos(np.clip(vector[1] / norm(vector), -1.0, 1.0))
            angles.extend([angle_x, angle_y])
    return angles

# Function to process image and predict alphabet
def process_and_predict(image):
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    results = hands.process(image_rgb)
    
    if results.multi_hand_landmarks:
        landmarks = np.array([[lm.x, lm.y] for lm in results.multi_hand_landmarks[0].landmark])
        landmarks_normalized = normalize_landmarks(landmarks)
        angles = calculate_angles(landmarks_normalized)
        
        angle_columns = [f'angle_{i}' for i in range(len(angles))]
        angles_df = pd.DataFrame([angles], columns=angle_columns)
        
        probabilities = model.predict_proba(angles_df)[0]
        return probabilities, landmarks
    
    return None, None

# Function to plot hand landmarks
def plot_hand_landmarks(landmarks, title):
    fig, ax = plt.subplots()
    ax.scatter(landmarks[:, 0], landmarks[:, 1], c='blue')
    for connection in mp.solutions.hands.HAND_CONNECTIONS:
        start_idx = connection[0]
        end_idx = connection[1]
        ax.plot([landmarks[start_idx, 0], landmarks[end_idx, 0]],
                [landmarks[start_idx, 1], landmarks[end_idx, 1]], 'r-')
    ax.invert_yaxis()
    ax.set_title(title)
    ax.axis('off')
    return fig

# Streamlit app
st.title("ASL Recognition App")

# Create two columns
col1, col2 = st.columns(2)

with col1:
    st.header("Predict ASL Sign")
    uploaded_file = st.file_uploader("Upload an image of an ASL sign", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
        st.image(image, caption="Uploaded Image", use_column_width=True)

        probabilities, landmarks = process_and_predict(image)
        
        if probabilities is not None:
            st.subheader("Top 5 Predictions:")
            top_indices = np.argsort(probabilities)[::-1][:5]
            for i in top_indices:
                st.write(f"{model.classes_[i]}: {probabilities[i]:.2f}")
        else:
            st.write("No hand detected in the image.")

with col2:
    st.header("Draw Hand Landmarks")
    all_alphabets = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
    excluded_alphabets = 'DMNPTUVXZ'
    available_alphabets = ''.join(set(all_alphabets) - set(excluded_alphabets))
    
    selected_alphabet = st.selectbox("Select an alphabet to draw landmarks:", list(available_alphabets))

    if selected_alphabet:
        image_path = f'asl test set/{selected_alphabet.lower()}.jpeg'
        if os.path.exists(image_path):
            image = cv2.imread(image_path)
            _, landmarks = process_and_predict(image)
            if landmarks is not None:
                fig = plot_hand_landmarks(landmarks, f"Hand Landmarks for {selected_alphabet}")
                st.pyplot(fig)
            else:
                st.write(f"No hand detected for {selected_alphabet}")
        else:
            st.write(f"Image not found for {selected_alphabet}")