# train_model.py
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
import joblib
import os

# 1. Create Sample Dataset (REPLACE WITH YOUR ACTUAL DATA)
data = {
    'face_shape': ['Oval', 'Round', 'Square'] * 50,
    'skin_tone': ['Fair', 'Medium', 'Dark'] * 50,
    'face_size': ['Small', 'Medium', 'Large'] * 50,
    'mask_style': ['Glitter', 'Animal', 'Floral'] * 50
}
df = pd.DataFrame(data)

# 2. Initialize Label Encoders
encoders = {
    'face_shape': LabelEncoder().fit(df['face_shape'].unique()),
    'skin_tone': LabelEncoder().fit(df['skin_tone'].unique()),
    'face_size': LabelEncoder().fit(df['face_size'].unique()),
    'mask_style': LabelEncoder().fit(df['mask_style'].unique())
}

# 3. Encode Features
X = pd.DataFrame({
    'face_shape': encoders['face_shape'].transform(df['face_shape']),
    'skin_tone': encoders['skin_tone'].transform(df['skin_tone']),
    'face_size': encoders['face_size'].transform(df['face_size'])
})
y = encoders['mask_style'].transform(df['mask_style'])

# 4. Train/Test Split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 5. Train Model
model = RandomForestClassifier(
    n_estimators=100,
    max_depth=5,
    random_state=42
)
model.fit(X_train, y_train)

# 6. Evaluate
print(f"Training Accuracy: {model.score(X_train, y_train):.2f}")
print(f"Test Accuracy: {model.score(X_test, y_test):.2f}")

# 7. Save to model/ Directory
os.makedirs('model', exist_ok=True)
joblib.dump(model, 'model/random_forest.pkl')
joblib.dump(encoders, 'model/label_encoders.pkl')

print("\nModel and encoders saved to model/ directory!")
print("Face Shape Classes:", encoders['face_shape'].classes_)
print("Mask Style Classes:", encoders['mask_style'].classes_)