import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
import joblib

# Load your dataset
df = pd.read_csv("mask_dataset.csv")

# Encode features
face_shape_encoder = LabelEncoder()
skin_tone_encoder = LabelEncoder()
style_encoder = LabelEncoder()

X = pd.DataFrame({
    'face_shape': face_shape_encoder.fit_transform(df['face_shape']),
    'skin_tone': skin_tone_encoder.fit_transform(df['skin_tone'])
})

y = style_encoder.fit_transform(df['recommended_style'])

# Train the model
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X, y)

# Save the model and encoders
joblib.dump(model, "mask_model.pkl")
joblib.dump(face_shape_encoder, "face_shape_encoder.pkl")
joblib.dump(skin_tone_encoder, "skin_tone_encoder.pkl")
joblib.dump(style_encoder, "style_encoder.pkl")

print("✅ Training complete. Model and encoders saved.")