# train_model.py
import numpy as np
import os
from glob import glob
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout, BatchNormalization
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
import joblib

DATA_DIR = "gesture_data"
SEQ_LEN = 30
MODEL_OUT = "gesture_lstm.h5"
LABELS_OUT = "labels.joblib"

def load_data(data_dir):
    X = []
    y = []
    for label in sorted(os.listdir(data_dir)):
        p = os.path.join(data_dir, label)
        if not os.path.isdir(p): continue
        files = glob(os.path.join(p, "*.npz"))
        for f in files:
            d = np.load(f)['data']
            if d.shape[0] != SEQ_LEN:
                continue
            X.append(d)  # shape (seq_len, features)
            y.append(label)
    X = np.array(X)
    y = np.array(y)
    return X, y

def build_model(input_shape, num_classes):
    model = Sequential([
        LSTM(128, return_sequences=True, input_shape=input_shape),
        Dropout(0.3),
        LSTM(64),
        BatchNormalization(),
        Dropout(0.3),
        Dense(64, activation="relu"),
        Dense(num_classes, activation="softmax")
    ])
    model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
    return model

def main():
    X, y = load_data(DATA_DIR)
    print("Loaded X:", X.shape, "y:", y.shape)
    le = LabelEncoder()
    y_enc = le.fit_transform(y)
    joblib.dump(le, LABELS_OUT)
    X_train, X_test, y_train, y_test = train_test_split(X, y_enc, test_size=0.2, random_state=42, stratify=y_enc)
    # normalize coords per sample (optional)
    # Flatten features along coordinates scale? We'll standardize by sample mean/std
    # Here we scale coords to zero mean unit std per sample for robustness
    def normalize_samples(arr):
        arr2 = arr.copy()
        for i in range(arr2.shape[0]):
            s = arr2[i]
            mean = s.mean(axis=0)
            std = s.std(axis=0) + 1e-8
            arr2[i] = (s - mean) / std
        return arr2
    X_train = normalize_samples(X_train)
    X_test = normalize_samples(X_test)
    input_shape = (X_train.shape[1], X_train.shape[2])  # (seq_len, features)
    model = build_model(input_shape, len(le.classes_))
    print(model.summary())
    ckpt = ModelCheckpoint(MODEL_OUT, save_best_only=True, monitor="val_loss", verbose=1)
    es = EarlyStopping(monitor="val_loss", patience=10, restore_best_weights=True)
    history = model.fit(X_train, y_train, validation_data=(X_test, y_test),
                        epochs=100, batch_size=16, callbacks=[ckpt, es])
    # final evaluation
    loss, acc = model.evaluate(X_test, y_test)
    print("Test loss/acc:", loss, acc)
    print("Saved model to", MODEL_OUT)
    print("Saved labels to", LABELS_OUT)

if __name__ == "__main__":
    main()