import streamlit as st
import cv2
import numpy as np
import mediapipe as mp
import pickle
import os

# Define the models directory
models_dir = 'models'

# Load models
def load_model(file_path):
    try:
        with open(file_path, 'rb') as file:
            return pickle.load(file)
    except FileNotFoundError:
        st.error(f"Error: Model file not found at {file_path}")
    except Exception as e:
        st.error(f"Error loading model from {file_path}: {e}")
    return None

models = {
    'Random Forest': load_model(os.path.join(models_dir, 'random_forest_model.pkl')),
    'SVM': load_model(os.path.join(models_dir, 'svm_model.pkl')),
    'Hard Voting': load_model(os.path.join(models_dir, 'hard_voting_classifier.pkl')),
    'Soft Voting': load_model(os.path.join(models_dir, 'soft_voting_classifier.pkl'))
}

# Load label encoder
label_encoder = load_model(os.path.join(models_dir, 'label_encoder.pkl'))
if label_encoder is None:
    st.error("Failed to load label encoder. Exiting.")
    st.stop()

# Initialize MediaPipe Hands
mp_hands = mp.solutions.hands
hands = mp_hands.Hands(static_image_mode=True, max_num_hands=1, min_detection_confidence=0.5)

def process_landmarks(hand_landmarks):
    x_ = [landmark.x for landmark in hand_landmarks.landmark]
    y_ = [landmark.y for landmark in hand_landmarks.landmark]
    min_x, min_y = min(x_), min(y_)
    return [coord - min_val for landmark in hand_landmarks.landmark for coord, min_val in zip((landmark.x, landmark.y), (min_x, min_y))]

def get_model_predictions(data_aux, model):
    if hasattr(model, 'predict_proba'):
        probabilities = model.predict_proba([data_aux])[0]
    elif hasattr(model, 'decision_function'):
        scores = model.decision_function([data_aux])[0]
        probabilities = (scores - scores.min()) / (scores.max() - scores.min())
    else:
        class_pred = model.predict([data_aux])[0]
        probabilities = np.zeros(len(label_encoder.classes_))
        probabilities[class_pred] = 1

    top_index = np.argmax(probabilities)
    return label_encoder.inverse_transform([top_index])[0], probabilities[top_index]

def predict_alphabet(image):
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    results = hands.process(image_rgb)

    if not results.multi_hand_landmarks:
        return None, None

    data_aux = process_landmarks(results.multi_hand_landmarks[0])
    predictions = {}

    for model_name, model in models.items():
        if model is None:
            continue
        prediction, probability = get_model_predictions(data_aux, model)
        predictions[model_name] = (prediction, probability)

    return predictions, results.multi_hand_landmarks[0]

def draw_hand_landmarks(image, hand_landmarks):
    for landmark in hand_landmarks.landmark:
        h, w, _ = image.shape
        cx, cy = int(landmark.x * w), int(landmark.y * h)
        cv2.circle(image, (cx, cy), 5, (0, 255, 0), -1)
    return image

# Streamlit App
st.title("ASL Hand Shape Recognition App")

uploaded_file = st.file_uploader("Choose a hand image...", type=["jpg", "jpeg", "png"])

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

    predictions, hand_landmarks = predict_alphabet(image)

    if predictions and hand_landmarks:
        st.subheader("Predictions:")
        for model_name, (prediction, probability) in predictions.items():
            st.write(f"{model_name}: {prediction} (Probability: {probability:.2f})")

        st.subheader("Hand Landmarks:")
        landmark_image = draw_hand_landmarks(image.copy(), hand_landmarks)
        st.image(landmark_image, channels="BGR", caption="Hand Landmarks")
    else:
        st.write("No hand detected in the image.")

# User input to draw hand sign (placeholder)
user_input = st.text_input("Enter alphabet to draw hand sign:")
if user_input:
    if len(user_input) == 1 and user_input.isalpha():
        st.write(f"Drawing hand sign for alphabet: {user_input.upper()}")
        # Placeholder for drawing hand sign
        st.write("(Hand sign drawing functionality not implemented)")
    else:
        st.error("Please enter a single alphabet.")

# Release MediaPipe resources
hands.close()