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_path):
    image = cv2.imread(image_path)
    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]
        predicted_class = model.classes_[np.argmax(probabilities)]
        
        return predicted_class, probabilities, landmarks
    
    return None, 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")

# Define available alphabets
all_alphabets = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
excluded_alphabets = 'DMNPTUVXZ'
available_alphabets = ''.join(set(all_alphabets) - set(excluded_alphabets))

# Multi-select for alphabets
selected_alphabets = st.multiselect("Select alphabets to recognize:", list(available_alphabets))

if selected_alphabets:
    num_selected = len(selected_alphabets)
    cols = min(3, num_selected)  # Maximum 3 columns
    rows = (num_selected + cols - 1) // cols
    
    for i in range(0, num_selected, cols):
        row_alphabets = selected_alphabets[i:i+cols]
        cols_in_row = st.columns(len(row_alphabets))
        
        for j, alphabet in enumerate(row_alphabets):
            with cols_in_row[j]:
                image_path = f'asl test set/{alphabet}.jpg'
                if os.path.exists(image_path):
                    predicted_class, probabilities, landmarks = process_and_predict(image_path)
                    if predicted_class is not None:
                        st.write(f"Alphabet: {alphabet}")
                        st.write(f"Predicted: {predicted_class}")
                        st.write(f"Confidence: {probabilities[model.classes_.tolist().index(predicted_class)]:.2f}")
                        fig = plot_hand_landmarks(landmarks, f"Hand Landmarks for {alphabet}")
                        st.pyplot(fig)
                    else:
                        st.write(f"No hand detected for {alphabet}")
                else:
                    st.write(f"Image not found for {alphabet}")

else:
    st.write("Please select at least one alphabet to recognize.")