import cv2
import mediapipe as mp
import numpy as np
import torch
import torch.nn as nn
import gradio as gr

# ----------------------------
# Labels
# ----------------------------
GESTURE_LABELS = {
    0: "A", 1: "B", 2: "L", 3: "U", 4: "V", 5: "W",
    6: "Z", 7: "F", 8: "five", 9: "one", 10: "three",
    11: "two", 12: "six", 13: "seven", 14: "eight",
    15: "nine", 16: "ten", 17: "E", 18: "four",
    19: "i", 20: "k", 21: "r", 22: "zero",
    23: "m", 24: "s"
}
CONF_THRESHOLD = 0.6
BLOCK_SIZE = 10   # frames per block

# ----------------------------
# Model
# ----------------------------
class GestureNet(nn.Module):
    def __init__(self, input_size=126, num_classes=len(GESTURE_LABELS)):
        super().__init__()
        self.fc1 = nn.Linear(input_size, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, num_classes)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.dropout(x)
        return self.fc3(x)

model = GestureNet()
model.load_state_dict(torch.load("gesture_model1.pth", map_location="cpu"))
model.eval()

# ----------------------------
# MediaPipe
# ----------------------------
mp_hands = mp.solutions.hands

# ----------------------------
# Landmark extraction + prediction
# ----------------------------
def extract_coords(results):
    coords = []
    if results.multi_hand_landmarks:
        for hand_landmarks in results.multi_hand_landmarks:
            hand_coords = np.array([[lm.x, lm.y, lm.z] for lm in hand_landmarks.landmark])
            hand_coords -= hand_coords[0]
            max_val = np.max(np.linalg.norm(hand_coords, axis=1))
            if max_val > 0:
                hand_coords /= max_val
            coords.extend(hand_coords.flatten())
    return coords

def run_model(coords):
    """Return label string (either gesture or 'Unknown')"""
    if len(coords) < 126:
        coords = coords + [0.0] * (126 - len(coords))
    elif len(coords) > 126:
        coords = coords[:126]

    input_tensor = torch.tensor(coords, dtype=torch.float32).unsqueeze(0)
    with torch.no_grad():
        outputs = model(input_tensor)
        probs = torch.softmax(outputs, dim=1)
        pred_class = torch.argmax(probs, dim=1).item()
        confidence = probs[0][pred_class].item()

    if confidence >= CONF_THRESHOLD:
        return GESTURE_LABELS[pred_class]
    return "Unknown"

# ----------------------------
# Image prediction (unchanged format)
# ----------------------------
def predict_image_from_rgb(image_rgb):
    with mp_hands.Hands(static_image_mode=True, max_num_hands=2) as hands:
        results = hands.process(image_rgb)
    coords = extract_coords(results)
    if not coords:
        return "No hand detected"
    label = run_model(coords)
    # run_model returns label or "Unknown" – we can keep the simple label
    return label

# ----------------------------
# 🆕 Video prediction → build word from 10-frame blocks
# ----------------------------
def predict_video_word(video_path):
    """
    Divides video into non-overlapping blocks of BLOCK_SIZE frames.
    If every frame in a block yields the *same* valid gesture,
    that gesture is added to the output word.
    Returns the concatenated word, or 'Unknown' if no block is valid.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return "Cannot open video file"

    frame_labels = []   # list of strings (gesture, "Unknown", or "No hand detected")

    with mp_hands.Hands(
        static_image_mode=False,
        max_num_hands=2,
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5,
    ) as hands:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = hands.process(frame_rgb)
            coords = extract_coords(results)
            if coords:
                label = run_model(coords)  # gesture name or "Unknown"
                frame_labels.append(label)
            else:
                frame_labels.append("No hand detected")

    cap.release()

    # Build word from non-overlapping blocks
    word = []
    for i in range(0, len(frame_labels) - BLOCK_SIZE + 1, BLOCK_SIZE):
        block = frame_labels[i : i + BLOCK_SIZE]
        # Check: all labels identical AND the label is not invalid
        if (len(set(block)) == 1 and
            block[0] not in ("Unknown", "No hand detected")):
            word.append(block[0])

    if word:
        return "".join(word)
    return "Unknown"

# ----------------------------
# Universal predictor
# ----------------------------
def predict_file(file_path, is_video=False):
    if file_path is None:
        return "No file uploaded"

    # Direct numpy array (image)
    if isinstance(file_path, np.ndarray):
        if len(file_path.shape) == 3 and file_path.shape[2] == 3:
            image_rgb = cv2.cvtColor(file_path, cv2.COLOR_BGR2RGB)
            return predict_image_from_rgb(image_rgb)
        else:
            return "Invalid image format"

    file_ext = str(file_path).lower()
    is_video_file = is_video or file_ext.endswith(('.mp4', '.avi', '.mov', '.mkv', '.webm'))

    if is_video_file:
        return predict_video_word(file_path)   # returns the word
    else:
        image = cv2.imread(file_path)
        if image is None:
            return "Could not read image file"
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return predict_image_from_rgb(image_rgb)

# ----------------------------
# Gradio UI
# ----------------------------
with gr.Blocks(title="Hand Gesture Recognition") as app:
    gr.Markdown("# ✋ Hand Gesture Recognition (Word‑Builder)")
    gr.Markdown(
        "Upload an image or video. For videos, non‑overlapping blocks of **10 frames** "
        "are checked. If all 10 frames in a block show the same gesture, that letter is "
        "added to the output word. Blocks with mixed or invalid gestures are ignored."
    )
    
    with gr.Row():
        input_file = gr.File(
            label="Upload Image or Video",
            file_types=["image", "video"],
            type="filepath"
        )
        output_text = gr.Textbox(label="Prediction", lines=3)
    
    predict_btn = gr.Button("Predict", variant="primary")
    predict_btn.click(
        predict_file,
        inputs=input_file,
        outputs=output_text
    )

app.launch()