"""Leave-one-out comparison of salary models on the shipped dataset.

Run offline with ``python baselines.py``. It writes ``data/baselines.json``,
which the app shows in its method notes only when the dataset fingerprint in
the file matches the loaded data. The comparison motivates two design choices:
fitting on market-contract players only, and using TabICLv2 rather than a
classical regressor.
"""

from __future__ import annotations

import json
import os
import time
from typing import Any, Callable

import numpy as np
from scipy import stats
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LogisticRegression, RidgeCV
from sklearn.metrics import roc_auc_score
from sklearn.neighbors import KNeighborsRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

import analysis as A

BASELINES_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data", "baselines.json")


def _loo_regression(X: np.ndarray, y_log: np.ndarray, predict_one: Callable) -> dict[str, float]:
    n = len(y_log)
    preds = np.zeros(n)
    for i in range(n):
        train = np.ones(n, dtype=bool)
        train[i] = False
        preds[i] = predict_one(X[train], y_log[train], X[[i]])
    ss_res = float(((y_log - preds) ** 2).sum())
    ss_tot = float(((y_log - y_log.mean()) ** 2).sum())
    rho, _ = stats.spearmanr(y_log, preds)
    return {
        "r2_log_salary": round(1 - ss_res / ss_tot, 3),
        "spearman_rho": round(float(rho), 3),
        "mean_abs_error_log": round(float(np.abs(y_log - preds).mean()), 3),
    }


def _sklearn_predictor(make_model: Callable) -> Callable:
    def predict_one(X_train, y_train, X_test):
        model = make_model()
        model.fit(X_train, y_train)
        return float(model.predict(X_test)[0])
    return predict_one


def _tabicl_predictor(X_train, y_train, X_test) -> float:
    quantiles = A.fit_predict_log_salary(X_train, y_train, X_test)[0]
    return float(quantiles[A.QUANTILE_GRID.index(0.50)])


REGRESSORS: dict[str, Callable] = {
    "Ridge regression (standardized inputs, alpha chosen by inner cross-validation)":
        _sklearn_predictor(lambda: make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-2, 3, 12)))),
    "k-nearest neighbours (k = 10, standardized inputs)":
        _sklearn_predictor(lambda: make_pipeline(StandardScaler(), KNeighborsRegressor(n_neighbors=10))),
    "Random forest (200 trees)":
        _sklearn_predictor(lambda: RandomForestRegressor(n_estimators=200, random_state=0)),
    "TabICLv2 (median of the predicted distribution)": _tabicl_predictor,
}


def _loo_classification(X: np.ndarray, label: np.ndarray, predict_proba_one: Callable) -> dict[str, float]:
    n = len(label)
    probs = np.zeros(n)
    for i in range(n):
        train = np.ones(n, dtype=bool)
        train[i] = False
        probs[i] = predict_proba_one(X[train], label[train], X[[i]])
    return {
        "auc": round(float(roc_auc_score(label, probs)), 3),
        "accuracy": round(float(np.mean((probs >= 0.5) == label)), 3),
    }


def _logistic_proba(X_train, y_train, X_test) -> float:
    model = make_pipeline(StandardScaler(), LogisticRegression(C=1.0, max_iter=1000))
    model.fit(X_train, y_train)
    return float(model.predict_proba(X_test)[0, 1])


def _tabicl_proba(X_train, y_train, X_test) -> float:
    from tabicl import TabICLClassifier

    model = TabICLClassifier(n_estimators=A._model_n_estimators(), device=os.environ.get("TABICL_DEVICE", "cpu"),
                             random_state=42, verbose=False)
    model.fit(X_train, y_train)
    return float(model.predict_proba(X_test)[0, 1])


def run(df) -> dict[str, Any]:
    t0 = time.time()
    X_all = df[A.STAT_COLUMNS].to_numpy(dtype=float)
    y_all = np.log(df[A.SALARY_COLUMN].to_numpy(dtype=float))
    market = A.model_training_mask(df)
    subsets = {"all players": np.ones(len(df), dtype=bool), "market contracts": market}
    regression = []
    for subset_name, mask in subsets.items():
        for method, predict_one in REGRESSORS.items():
            metrics = _loo_regression(X_all[mask], y_all[mask], predict_one)
            regression.append({"method": method, "subset": subset_name, "n": int(mask.sum()), **metrics})
            print(subset_name, "|", method, metrics, flush=True)
    label = market.astype(int)
    classification = []
    for method, proba_one in {
        "Logistic regression (standardized inputs)": _logistic_proba,
        "TabICLv2 classifier": _tabicl_proba,
    }.items():
        metrics = _loo_classification(X_all, label, proba_one)
        classification.append({"method": method, **metrics})
        print("classifier |", method, metrics, flush=True)
    return {
        "fingerprint": A.dataset_fingerprint(df),
        "n_estimators": A._model_n_estimators(),
        "market_contract_threshold": A.MARKET_CONTRACT_THRESHOLD,
        "regression": regression,
        "classification": classification,
        "majority_class_accuracy": round(float(max(label.mean(), 1 - label.mean())), 3),
        "seconds": round(time.time() - t0),
    }


def load_cached_baselines(df, path: str = BASELINES_PATH) -> dict[str, Any] | None:
    if not os.path.exists(path):
        return None
    try:
        with open(path, "r", encoding="utf-8") as fh:
            cached = json.load(fh)
    except (OSError, ValueError):
        return None
    if cached.get("fingerprint") != A.dataset_fingerprint(df):
        return None
    return cached


if __name__ == "__main__":
    data = A.load_players()
    result = run(data)
    os.makedirs(os.path.dirname(BASELINES_PATH), exist_ok=True)
    with open(BASELINES_PATH, "w", encoding="utf-8") as fh:
        json.dump(result, fh, indent=2)
    print(json.dumps(result, indent=2))