"""Add classical models to the saved Jev Choice/Noul experiment without API calls.

Run: python combine_model_comparison.py
All models use the same seed-18 test split. The 16-example variants fit
the exact rows supplied to Jev; full variants fit the remaining train split.
"""

import json
from pathlib import Path

import numpy as np
from sklearn.ensemble import ExtraTreesClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

from benchmark import datasets
from choice_threshold_benchmark import EXAMPLES_PER_CLASS, SEED, metrics


def model_factories():
    return {
        "gradient_boosting": lambda: GradientBoostingClassifier(random_state=SEED),
        "extra_trees": lambda: ExtraTreesClassifier(
            n_estimators=200, random_state=SEED, n_jobs=-1
        ),
        "logistic_regression": lambda: make_pipeline(
            StandardScaler(), LogisticRegression(max_iter=5000, random_state=SEED)
        ),
    }


def main():
    saved = json.loads(Path("choice_threshold_results.json").read_text())
    combined = {}
    for name, X, y, _names, _positive, _negative in datasets():
        X_dev, X_test, y_dev, y_test = train_test_split(
            X, y, test_size=0.25, stratify=y, random_state=SEED
        )
        X_train, X_val, y_train, y_val = train_test_split(
            X_dev, y_dev, test_size=0.25, stratify=y_dev, random_state=SEED
        )
        rng = np.random.default_rng(SEED)
        chosen = np.concatenate([
            rng.choice(np.flatnonzero(y_train == label), EXAMPLES_PER_CLASS, replace=False)
            for label in (0, 1)
        ])
        previous = saved[name]
        assert len(X_val) == previous["n_val"]
        assert len(X_test) == previous["n_test"]
        assert y_test.tolist() == [r["true_label"] for r in previous["test_rows"]]

        results = dict(previous["test_metrics"])
        for model_name, factory in model_factories().items():
            for variant, subset in (("16_examples", chosen), ("full_train", None)):
                model = factory()
                model.fit(
                    X_train if subset is None else X_train[subset],
                    y_train if subset is None else y_train[subset],
                )
                probability = model.predict_proba(X_test)[:, 1]
                results[f"{model_name}_{variant}"] = metrics(y_test, probability, 0.5)

        combined[name] = {
            "positive": previous["positive"], "negative": previous["negative"],
            "n_train": len(X_train), "n_val": len(X_val), "n_test": len(X_test),
            "jev_selected_cutoff": previous["cutoffs"]["jev_selected"],
            "results": results,
        }
    Path("all_model_results.json").write_text(json.dumps(combined, indent=2) + "\n")
    for name, item in combined.items():
        print(f"\n{name} (positive: {item['positive']}; test n={item['n_test']})")
        for model, result in item["results"].items():
            print(f"  {model:<38} acc={result['accuracy']:.4f} "
                  f"auc={result['auc']:.4f} sensitivity={result['sensitivity']:.4f} "
                  f"specificity={result['specificity']:.4f} "
                  f"matrix={result['confusion_matrix']}")


if __name__ == "__main__":
    main()