import os
# --- FIX: Set environment variable to prevent Protobuf/TensorFlow crash ---
os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
# -----------------------------------------------------------------------

import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from tensorflow.keras.utils import to_categorical
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# --- 1. Data Loading and Setup ---
DATA_PATH = os.path.join('ML_Data')
ACTIONS = np.array(sorted(os.listdir(DATA_PATH)))
SEQUENCE_LENGTH = 20
label_map = {label:num for num, label in enumerate(ACTIONS)}
NUM_CLASSES = len(ACTIONS)

sequences, labels = [], []
for action in ACTIONS:
    for sequence in range(15):
        window = []
        for frame_num in range(SEQUENCE_LENGTH):
            try:
                # Load the 63 features (21 landmarks * 3 coords)
                res = np.load(os.path.join(DATA_PATH, action, str(sequence), f"{frame_num}.npy"))
                window.append(res)
            except:
                continue
        if len(window) == SEQUENCE_LENGTH:
            sequences.append(window)
            labels.append(label_map[action])

X = np.array(sequences)
y = to_categorical(labels).astype(int)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.05)

# --- 2. Define the Neural Network Model (LSTM) ---
model = Sequential()
model.add(LSTM(64, return_sequences=True, activation='relu', input_shape=(SEQUENCE_LENGTH, 63)))
model.add(LSTM(128, return_sequences=True, activation='relu'))
model.add(LSTM(64, return_sequences=False, activation='relu'))
model.add(Dense(64, activation='relu'))
model.add(Dense(32, activation='relu'))
model.add(Dense(NUM_CLASSES, activation='softmax'))

# --- 3. Compile, Train, and Save ---
model.compile(optimizer='Adam', loss='categorical_crossentropy', metrics=['categorical_accuracy'])
print(f"Starting Training for {NUM_CLASSES} signs...")

model.fit(X_train, y_train, epochs=100, callbacks=[])

res = model.predict(X_test)
y_pred = np.argmax(res, axis=1)
y_true = np.argmax(y_test, axis=1)
print(f"\nModel Accuracy on Test Data: {accuracy_score(y_true, y_pred)*100:.2f}%")
model.save('/tmp/sign_language_model.h5')
print("Model trained and saved as sign_language_model.h5")