import os
import numpy as np
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.callbacks import ModelCheckpoint

# Directory where posture data is stored
data_dir = "posture_data"
# Define the posture classes
classes = ["Standing", "WallSit", "Pushup", "Squat", "BicepCurl"]
# Map each class to an integer label
label_map = {label: i for i, label in enumerate(classes)}

# Initialize lists to store the data and labels
X = []
y = []

# Load and preprocess the data
for label in classes:
    folder_path = os.path.join(data_dir, label)
    for filename in os.listdir(folder_path):
        if filename.endswith('.npy'):
            filepath = os.path.join(folder_path, filename)
            # Load the landmarks and flatten to (99,)
            data = np.load(filepath).flatten()
            X.append(data)
            y.append(label_map[label])

# Convert lists to numpy arrays
X = np.array(X)  # Shape: (num_samples, 99)
y = to_categorical(y, num_classes=len(classes))  # One-hot encode the labels

print("Data loaded and processed:")
print("X shape:", X.shape)
print("y shape:", y.shape)

### Step 2: Define the Model
model = Sequential([
    Dense(256, activation='relu', input_shape=(99,)),
    Dropout(0.5),
    Dense(128, activation='relu'),
    Dropout(0.5),
    Dense(64, activation='relu'),
    Dropout(0.5),
    Dense(len(classes), activation='softmax')  # Output layer for classification
])

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# Display model summary
model.summary()

### Step 3: Set up Checkpoints
checkpoint_path = "best_pose_classification_model.keras"
model_checkpoint = ModelCheckpoint(checkpoint_path, save_best_only=True, monitor='val_loss', mode='min')

### Step 4: Train the Model
# Split data into training and validation sets (e.g., 80% train, 20% validate)
from sklearn.model_selection import train_test_split
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

history = model.fit(
    X_train, y_train,
    validation_data=(X_val, y_val),
    epochs=50,  # Adjust the number of epochs as needed
    batch_size=32,
    callbacks=[model_checkpoint]
)

print("Model training complete. Best model saved to:", checkpoint_path)