"""Models and classification metrics.

Algorithms copied verbatim from ppi_benchmark_controls.py v1.0.0 (lines 665-716).
Do not edit the algorithms without re-running the golden tests."""
from __future__ import annotations

import warnings

import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score, average_precision_score, balanced_accuracy_score, brier_score_loss,
    confusion_matrix, f1_score, matthews_corrcoef, precision_score, recall_score, roc_auc_score,
)
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler

# Web-app addition: parallelism cap for shared servers. Does not change results
# (forests and XGBoost are seeded; n_jobs only affects scheduling).
N_JOBS = -1


def build_model(name: str, seed: int, control: bool = False):
    if name == "lr":
        return LogisticRegression(max_iter=2000, random_state=seed)
    if name == "rf":
        if control:
            return RandomForestClassifier(n_estimators=200, max_depth=6, random_state=seed)
        return RandomForestClassifier(n_estimators=300, n_jobs=N_JOBS, random_state=seed)
    if name == "mlp":
        return MLPClassifier(hidden_layer_sizes=(128, 64), max_iter=500, early_stopping=True, random_state=seed)
    if name == "xgb":
        from xgboost import XGBClassifier
        return XGBClassifier(n_estimators=300, eval_metric="logloss", random_state=seed, n_jobs=N_JOBS)
    raise ValueError(f"Unknown model {name!r}; choose from lr, rf, mlp, xgb")



def ece_equal_width(y: np.ndarray, p: np.ndarray, n_bins: int = 10) -> float:
    idx = np.clip(np.digitize(p, np.linspace(0, 1, n_bins + 1)[1:-1], right=True), 0, n_bins - 1)
    return float(sum((idx == b).mean() * abs(y[idx == b].mean() - p[idx == b].mean())
                     for b in range(n_bins) if (idx == b).any()))


def classification_metrics(y: np.ndarray, p: np.ndarray, threshold: float = 0.5) -> dict:
    """ROC-AUC, PR-AUC (average precision) and fixed-threshold metrics (never tuned on test)."""
    if len(np.unique(y)) < 2:
        return {"roc_auc": np.nan, "pr_auc": np.nan, "mcc": np.nan, "n_test": len(y)}
    pred = (p >= threshold).astype(int)
    tn, fp, fn, tp = confusion_matrix(y, pred, labels=[0, 1]).ravel()
    return {"roc_auc": roc_auc_score(y, p), "pr_auc": average_precision_score(y, p),
            "mcc": matthews_corrcoef(y, pred), "f1": f1_score(y, pred, zero_division=0),
            "balanced_accuracy": balanced_accuracy_score(y, pred), "accuracy": accuracy_score(y, pred),
            "precision": precision_score(y, pred, zero_division=0), "recall": recall_score(y, pred, zero_division=0),
            "specificity": tn / (tn + fp) if (tn + fp) else 0.0, "brier": brier_score_loss(y, p),
            "ece_10bin": ece_equal_width(y, p), "positive_prevalence": float(y.mean()), "n_test": int(len(y))}


def fit_predict(model_name, x_train, y_train, x_test, seed, control=False) -> np.ndarray:
    if len(np.unique(y_train)) < 2:
        raise ValueError("training data contain a single class")
    scaler = StandardScaler().fit(x_train)
    model = build_model(model_name, seed, control)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        model.fit(scaler.transform(x_train), y_train)
    return model.predict_proba(scaler.transform(x_test))[:, 1]


# =============================================================================
# 7. Graph-aware uncertainty
# =============================================================================