# train.py
import pandas as pd, numpy as np, pickle
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.multioutput import MultiOutputRegressor
from sklearn.metrics import mean_absolute_error

CSV_PATH = "student_allinone_300_padded.csv"  # change path if needed
FEATURES = ["Attendance","StudyHours","ParentalSupport","SleepHours",
            "ReadingHours","BehaviorScore","PretestScore",
            "HomeworkCompletion","Participation"]
TARGETS = ["AssignmentAvg","TestScore"]

def main():
    df = pd.read_csv(CSV_PATH).copy()

    # === Highly recommended: make targets depend on inputs (if your CSV targets were random) ===
    rng = np.random.default_rng(42)
    if ("AssignmentAvg" in df.columns) and ("TestScore" in df.columns):
        # Always recompute to ensure consistency
        df["AssignmentAvg"] = (
            df["PretestScore"] * 0.5
            + df["StudyHours"] * 3
            + df["HomeworkCompletion"] * 0.20
            + df["Participation"] * 2
            + rng.integers(-5, 6, size=len(df))
        ).clip(0, 100).round(2)

        df["TestScore"] = (
            df["PretestScore"] * 0.6
            + df["Attendance"] * 0.20
            + df["ParentalSupport"] * 3
            + df["SleepHours"] * 2
            + df["ReadingHours"] * 2
            + df["BehaviorScore"] * 2
            + rng.integers(-5, 6, size=len(df))
        ).clip(0, 100).round(2)

    X = df[FEATURES]
    y = df[TARGETS]

    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, random_state=42)
    model = MultiOutputRegressor(RandomForestRegressor(n_estimators=200, random_state=42)).fit(Xtr, ytr)
    mae = mean_absolute_error(yte, model.predict(Xte))
    print("MAE:", round(mae, 3))

    # Save feature bounds so app can clip
    feature_mins = X.min().to_dict()
    feature_maxs = X.max().to_dict()

    with open("student_model.pkl", "wb") as f:
        pickle.dump({
            "model": model,
            "features": FEATURES,
            "targets": TARGETS,
            "feature_mins": feature_mins,
            "feature_maxs": feature_maxs
        }, f)

    print("Saved student_model.pkl with bounds.")

if __name__ == "__main__":
    main()