"""Incremental multiclass scikit-learn ensembles in Pyodide."""
import inspect
import json
import warnings
import numpy as np
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier, ExtraTreesClassifier
from sklearn.linear_model import LogisticRegression, SGDClassifier, Perceptron, RidgeClassifier
from sklearn.svm import SVC, LinearSVC
from sklearn.naive_bayes import GaussianNB
from sklearn.neighbors import KNeighborsClassifier
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis, QuadraticDiscriminantAnalysis
from sklearn.neural_network import MLPClassifier

CLASSIFIERS = {cls.__name__: cls for cls in (
    DecisionTreeClassifier, LogisticRegression, SGDClassifier, Perceptron,
    RidgeClassifier, SVC, LinearSVC, GaussianNB, KNeighborsClassifier,
    LinearDiscriminantAnalysis, QuadraticDiscriminantAnalysis, MLPClassifier,
)}
FORESTS = {"random_forest": RandomForestClassifier, "extra_trees": ExtraTreesClassifier}
GRID_SIZE = 160
axis = np.linspace(0, 1, GRID_SIZE)
gx, gy = np.meshgrid(axis, axis)
GRID = np.column_stack((gx.ravel(), gy.ravel()))
history = []
configuration = None


def catalog():
    return [{"name": name, "weighted": "sample_weight" in inspect.signature(cls.fit).parameters}
            for name, cls in CLASSIFIERS.items()]


def train(request):
    global configuration, history
    seed = request.get("seed", 42)
    if type(seed) is not int or not 0 <= seed <= 2**32 - 1:
        raise ValueError("Random seed must be an integer from 0 to 4294967295.")
    params = request.get("params", {})
    if not isinstance(params, dict):
        raise ValueError('Arguments must be a JSON object, for example {"max_depth": 2}.')
    method = request["method"]
    forest = method in FORESTS
    boosting = method == "boosting"
    if method not in (*FORESTS, "boosting", "bagging"):
        raise ValueError("Unknown ensemble method.")
    if forest and any(k in params for k in ("n_estimators", "warm_start")):
        raise ValueError("Use Add/Remove to set the number of trees; omit n_estimators and warm_start.")
    cls = FORESTS[method] if forest else CLASSIFIERS[request["classifier"]]
    estimator = cls(**params)
    if boosting and "sample_weight" not in inspect.signature(cls.fit).parameters:
        raise ValueError(f"{cls.__name__} does not support sample_weight. Select bagging or another classifier.")
    X = np.array([[p["x"], p["y"]] for p in request["points"]], dtype=float)
    y = np.array([p["label"] for p in request["points"]])
    classes = np.unique(y)
    if len(y) > 1500 or len(classes) < 2 or not set(classes).issubset({-1, 1, 2}):
        raise ValueError("Draw at least two classes, with at most 1,500 points.")
    key = json.dumps({**{k: request[k] for k in ("points", "classifier", "params", "method")}, "seed": seed}, sort_keys=True)
    if key != configuration:
        history, configuration = [], key
    count = request["count"]
    if not 0 <= count <= len(history) or count >= 100:
        raise ValueError("Ensemble state changed. Reset the dataset and try again.")
    history = history[:count]
    # Each model has a repeatable stream, including after Remove then Add.
    rng = np.random.default_rng(np.random.SeedSequence([seed, count]))
    if not forest and "random_state" in estimator.get_params() and params.get("random_state") is None:
        estimator.set_params(random_state=int(rng.integers(2**31 - 1)))
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        if forest:
            options = dict(params)
            if options.get("random_state") is None:
                options["random_state"] = seed
            estimator = cls(n_estimators=count + 1, warm_start=True, **options)
            estimator.estimators_ = [m["estimator"] for m in history]
            estimator.fit(X, y)
            learner = estimator.estimators_[-1]
            training_probabilities = learner.predict_proba(X)
            grid_probabilities = learner.predict_proba(GRID)
            predictions = classes[training_probabilities.argmax(axis=1)]
            grid = classes[grid_probabilities.argmax(axis=1)]
            alpha, perfect = 1.0, False
        else:
            # SAMME: only mistakes increase a sample's log weight.
            logs = sum((m["alpha"] * (m["predictions"] != y) for m in history), np.zeros(len(y)))
            weights = np.exp(logs - logs.max())
            weights /= weights.sum()
            if boosting:
                estimator.fit(X, y, sample_weight=weights)
            else:
                indices = rng.integers(len(y), size=len(y))
                estimator.fit(X[indices], y[indices])
            learner = estimator
            predictions = learner.predict(X)
            error = float(weights @ (predictions != y))
            if boosting and error >= 1 - 1 / len(classes) - 1e-12:
                raise ValueError("This learner cannot beat chance with the current weights. Try different classifier arguments or another classifier.")
            clipped = max(error, 1e-12)
            alpha = float(np.log((1 - clipped) / clipped) + np.log(len(classes) - 1)) if boosting else 1.0
            perfect = boosting and error < 1e-12
            grid = learner.predict(GRID)
    result = {"classes": classes.astype(int).tolist(), "forest": forest, "alpha": alpha,
              "perfect": perfect, "predictions": predictions.astype(int).tolist(),
              "tree": {"size": GRID_SIZE, "grid": grid.astype(int).tolist()},
              "warning": " ".join(dict.fromkeys(str(w.message) for w in caught))}
    if forest:
        result["trainingProbabilities"] = training_probabilities.tolist()
        result["tree"]["probabilities"] = grid_probabilities.tolist()
    # Commit only after all prediction and serialization work succeeds.
    json.dumps(result, allow_nan=False)
    history.append({"estimator": learner, "alpha": alpha, "predictions": predictions})
    return result


def handle_request(payload):
    try:
        return json.dumps({"model": train(json.loads(payload))}, allow_nan=False)
    except Exception as exc:
        return json.dumps({"error": str(exc)})