ielts_model/src/model_comparison.py

"""Compare tabular baselines for next-week IELTS band prediction."""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from typing import Any

import joblib
import numpy as np
import pandas as pd
from lightgbm import LGBMRegressor
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import GroupShuffleSplit
from sklearn.multioutput import MultiOutputRegressor
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from xgboost import XGBRegressor


SKILLS = ["listening", "reading", "writing", "speaking"]
TARGET_COLUMNS = ["target_next_overall_band", *[f"target_next_{skill}_band" for skill in SKILLS]]
CURRENT_TARGET_COLUMNS = ["overall_band", *[f"{skill}_band" for skill in SKILLS]]
TARGET_NAMES = ["overall_band", *[f"{skill}_band" for skill in SKILLS]]

CATEGORICAL_FEATURES = ["learner_archetype"]
NUMERIC_FEATURES = [
    "week",
    "days_elapsed",
    "weekly_study_hours",
    "accumulated_study_hours",
    "weeks_elapsed",
    "avg_weekly_hours",
    "target_band",
    "ceiling_band",
    "start_overall_band",
    "current_overall_band",
    "overall_band",
    "study_consistency",
    "learning_rate",
    "motivation",
    "noise_level",
    "distance_to_ceiling",
    "distance_to_target",
    "is_high_band",
    "start_listening_band",
    "start_reading_band",
    "start_writing_band",
    "start_speaking_band",
    "listening_band",
    "reading_band",
    "writing_band",
    "speaking_band",
    "mock_overall",
    "mock_listening",
    "mock_reading",
    "mock_writing",
    "mock_speaking",
    "mock_overall_rolling3_mean",
    "mock_overall_rolling3_std",
    "overall_band_trend4",
    "mock_listening_rolling3_mean",
    "mock_listening_rolling3_std",
    "listening_band_trend4",
    "mock_reading_rolling3_mean",
    "mock_reading_rolling3_std",
    "reading_band_trend4",
    "mock_writing_rolling3_mean",
    "mock_writing_rolling3_std",
    "writing_band_trend4",
    "mock_speaking_rolling3_mean",
    "mock_speaking_rolling3_std",
    "speaking_band_trend4",
]


def split_by_student(data: pd.DataFrame, test_size: float, random_state: int) -> tuple[pd.DataFrame, pd.DataFrame]:
    splitter = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=random_state)
    train_idx, test_idx = next(splitter.split(data, groups=data["student_id"]))
    return data.iloc[train_idx].copy(), data.iloc[test_idx].copy()


def make_preprocessor(scale_numeric: bool = False) -> ColumnTransformer:
    numeric_step: str | StandardScaler = StandardScaler() if scale_numeric else "passthrough"
    return ColumnTransformer(
        transformers=[
            ("categorical", OneHotEncoder(handle_unknown="ignore"), CATEGORICAL_FEATURES),
            ("numeric", numeric_step, NUMERIC_FEATURES),
        ],
        remainder="drop",
    )


def build_models(seed: int, n_jobs: int) -> dict[str, Pipeline]:
    return {
        "ridge": Pipeline(
            steps=[
                ("preprocess", make_preprocessor(scale_numeric=True)),
                ("model", Ridge(alpha=1.0)),
            ]
        ),
        "random_forest": Pipeline(
            steps=[
                ("preprocess", make_preprocessor()),
                (
                    "model",
                    RandomForestRegressor(
                        n_estimators=300,
                        min_samples_leaf=3,
                        random_state=seed,
                        n_jobs=n_jobs,
                    ),
                ),
            ]
        ),
        "xgboost": Pipeline(
            steps=[
                ("preprocess", make_preprocessor()),
                (
                    "model",
                    MultiOutputRegressor(
                        XGBRegressor(
                            n_estimators=350,
                            max_depth=4,
                            learning_rate=0.045,
                            subsample=0.9,
                            colsample_bytree=0.9,
                            objective="reg:squarederror",
                            tree_method="hist",
                            random_state=seed,
                            n_jobs=n_jobs,
                        ),
                        n_jobs=1,
                    ),
                ),
            ]
        ),
        "lightgbm": Pipeline(
            steps=[
                ("preprocess", make_preprocessor()),
                (
                    "model",
                    MultiOutputRegressor(
                        LGBMRegressor(
                            n_estimators=450,
                            learning_rate=0.035,
                            num_leaves=31,
                            min_child_samples=30,
                            subsample=0.9,
                            colsample_bytree=0.9,
                            random_state=seed,
                            n_jobs=n_jobs,
                            verbose=-1,
                        ),
                        n_jobs=1,
                    ),
                ),
            ]
        ),
    }


def rounded_band_predictions(model: Pipeline, x_test: pd.DataFrame) -> np.ndarray:
    predictions = model.predict(x_test)
    return np.clip(np.round(predictions * 2) / 2, 0.0, 9.0)


def evaluate_predictions(y_true: pd.DataFrame, predictions: np.ndarray) -> dict[str, float]:
    per_target = {}
    for idx, target_name in enumerate(TARGET_NAMES):
        per_target[f"mae_{target_name}"] = round(
            float(mean_absolute_error(y_true[TARGET_COLUMNS[idx]], predictions[:, idx])),
            4,
        )
    per_target["macro_mae"] = round(float(np.mean(list(per_target.values()))), 4)
    return per_target


def evaluate_naive(test_data: pd.DataFrame) -> dict[str, float]:
    per_target = {}
    for idx, target_name in enumerate(TARGET_NAMES):
        per_target[f"mae_{target_name}"] = round(
            float(mean_absolute_error(test_data[TARGET_COLUMNS[idx]], test_data[CURRENT_TARGET_COLUMNS[idx]])),
            4,
        )
    per_target["macro_mae"] = round(float(np.mean(list(per_target.values()))), 4)
    return per_target


def stability_score(metrics: dict[str, float]) -> float:
    target_maes = [metrics[f"mae_{target_name}"] for target_name in TARGET_NAMES]
    return round(float(np.std(target_maes)), 4)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Compare IELTS prediction models.")
    parser.add_argument("--training-data", type=Path, default=Path("data/synthetic_training_rows.csv"))
    parser.add_argument("--reports-dir", type=Path, default=Path("reports"))
    parser.add_argument("--model-dir", type=Path, default=Path("models"))
    parser.add_argument("--test-size", type=float, default=0.2)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--n-jobs", type=int, default=-1)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    args.reports_dir.mkdir(parents=True, exist_ok=True)
    args.model_dir.mkdir(parents=True, exist_ok=True)

    data = pd.read_csv(args.training_data)
    train_data, test_data = split_by_student(data, test_size=args.test_size, random_state=args.seed)

    x_train = train_data[CATEGORICAL_FEATURES + NUMERIC_FEATURES]
    y_train = train_data[TARGET_COLUMNS]
    x_test = test_data[CATEGORICAL_FEATURES + NUMERIC_FEATURES]
    y_test = test_data[TARGET_COLUMNS]

    rows: list[dict[str, Any]] = []
    trained_models: dict[str, Pipeline] = {}

    naive_metrics = evaluate_naive(test_data)
    rows.append(
        {
            "model": "naive_current_band",
            "fit_seconds": 0.0,
            "stability_mae_std": stability_score(naive_metrics),
            **naive_metrics,
        }
    )

    for model_name, model in build_models(seed=args.seed, n_jobs=args.n_jobs).items():
        start = time.perf_counter()
        model.fit(x_train, y_train)
        fit_seconds = time.perf_counter() - start
        predictions = rounded_band_predictions(model, x_test)
        metrics = evaluate_predictions(y_test, predictions)
        rows.append(
            {
                "model": model_name,
                "fit_seconds": round(fit_seconds, 2),
                "stability_mae_std": stability_score(metrics),
                **metrics,
            }
        )
        trained_models[model_name] = model

    results = pd.DataFrame(rows).sort_values(["macro_mae", "stability_mae_std", "fit_seconds"])
    best_model_name = str(results.iloc[0]["model"])

    csv_path = args.reports_dir / "model_comparison.csv"
    json_path = args.reports_dir / "model_comparison.json"
    results.to_csv(csv_path, index=False)
    json_path.write_text(
        json.dumps(
            {
                "best_model": best_model_name,
                "selection_rule": "lowest macro_mae, then lowest target MAE std, then fastest fit time",
                "train_rows": int(len(train_data)),
                "test_rows": int(len(test_data)),
                "train_students": int(train_data["student_id"].nunique()),
                "test_students": int(test_data["student_id"].nunique()),
                "results": results.to_dict(orient="records"),
            },
            indent=2,
        ),
        encoding="utf-8",
    )

    if best_model_name in trained_models:
        best_model_path = args.model_dir / "ielts_best_model.joblib"
        joblib.dump(trained_models[best_model_name], best_model_path)
        print(f"Saved best model: {best_model_path}")
    else:
        print("Best result is the naive baseline; no trained model was saved as best.")

    print(f"Saved comparison CSV: {csv_path}")
    print(f"Saved comparison JSON: {json_path}")
    print(results.to_string(index=False))


if __name__ == "__main__":
    main()