!pip install -U pip

!pip install git+https://github.com/Techtonique/genbooster.git

!pip install mlsauce

import multiprocessing as mp
import resource
import time
import warnings

import lightgbm as lgb
import matplotlib.pyplot as plt
import mlsauce as ms
import numpy as np
import pandas as pd
import xgboost as xgb
from sklearn.base import clone
from sklearn.datasets import (load_breast_cancer, load_diabetes, load_iris,
                              load_wine, make_friedman1)
from sklearn.linear_model import Ridge
from sklearn.metrics import balanced_accuracy_score, mean_absolute_error, mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.utils import all_estimators
from tqdm.auto import tqdm

from genbooster import BoosterClassifier, BoosterRegressor

warnings.filterwarnings("ignore")

# Shared, untuned settings (genbooster defaults, mirrored in mlsauce)
SEED = 42
N_ESTIMATORS = 100
LEARNING_RATE = 0.1
N_HIDDEN_FEATURES = 5
TOLERANCE = 1e-4
TIMEOUT_S = 120     # per run, wall clock; slower runs are killed and recorded as failures
MEM_BUDGET_GB = 3   # per run; guards against base learners that store large fitted state
TEST_SIZE = 0.2

EXCLUDED = [
    "MultiOutputRegressor", "MultiOutputClassifier", "StackingRegressor", "StackingClassifier",
    "VotingRegressor", "VotingClassifier", "TransformedTargetRegressor", "RegressorChain",
    "GradientBoostingRegressor", "HistGradientBoostingRegressor", "RandomForestRegressor",
    "ExtraTreesRegressor", "MLPRegressor",
]

REG_DATASETS = {"diabetes": load_diabetes(return_X_y=True)}
CLF_DATASETS = {
    "iris": load_iris(return_X_y=True),
    "wine": load_wine(return_X_y=True),
    "breast_cancer": load_breast_cancer(return_X_y=True),
}

COMMON = dict(n_estimators=N_ESTIMATORS, learning_rate=LEARNING_RATE,
              n_hidden_features=N_HIDDEN_FEATURES, direct_link=True,
              dropout=0.0, tolerance=TOLERANCE)
GB_ONLY = dict(weights_distribution="uniform", random_state=SEED)
MS_ONLY = dict(weights_distr="uniform", activation="relu", row_sample=1.0, col_sample=1.0,
               seed=SEED, verbose=0)

BOOSTERS = {
    "regression": {
        "genbooster": lambda base: BoosterRegressor(base_estimator=base, **COMMON, **GB_ONLY),
        "mlsauce": lambda base: ms.GenericBoostingRegressor(base_model=base, **COMMON, **MS_ONLY),
    },
    "classification": {
        "genbooster": lambda base: BoosterClassifier(base_estimator=base, **COMMON, **GB_ONLY),
        "mlsauce": lambda base: ms.GenericBoostingClassifier(base_model=base, **COMMON, **MS_ONLY),
    },
}

# Untuned references: library defaults, only seed and logging set
REFERENCES = {
    "regression": {
        "XGBoost": lambda: xgb.XGBRegressor(random_state=SEED, verbosity=0),
        "LightGBM": lambda: lgb.LGBMRegressor(random_state=SEED, verbose=-1),
    },
    "classification": {
        "XGBoost": lambda: xgb.XGBClassifier(random_state=SEED, verbosity=0),
        "LightGBM": lambda: lgb.LGBMClassifier(random_state=SEED, verbose=-1),
    },
}


def seeded(base):
    """Same base-learner seed for both boosters (genbooster would otherwise set it itself)."""
    base = clone(base)
    if "random_state" in base.get_params():
        base.set_params(random_state=SEED)
    return base


def scores(task, y_te, pred):
    if task == "regression":
        if not np.isfinite(pred).all():
            raise ValueError("non-finite predictions")
        return {"rmse": mean_squared_error(y_te, pred) ** 0.5,
                "mae": mean_absolute_error(y_te, pred)}
    return {"balanced_accuracy": balanced_accuracy_score(y_te, pred)}


def evaluate(task, model, X_tr, X_te, y_tr, y_te):
    """Fit + predict; return the task's metrics and elapsed seconds."""
    t0 = time.perf_counter()
    pred = np.asarray(model.fit(X_tr, y_tr).predict(X_te)).ravel()
    elapsed = time.perf_counter() - t0
    if pred.shape != y_te.shape:
        raise ValueError(f"predictions have shape {pred.shape}, expected {y_te.shape}")
    return {**scores(task, y_te, pred), "time_s": elapsed}


def _vm_data_bytes():
    with open("/proc/self/status") as f:
        return next(int(l.split()[1]) * 1024 for l in f if l.startswith("VmData"))


def _die_with_parent():
    """Linux: kill this child if the notebook kernel dies or is interrupted, so orphans
    don't keep using the CPU and skew later timings."""
    try:
        import ctypes
        ctypes.CDLL("libc.so.6").prctl(1, 9)  # PR_SET_PDEATHSIG, SIGKILL
    except Exception:
        pass


def _child(queue, task, model_fn, data):
    _die_with_parent()
    try:
        # Cap new allocations so one runaway base learner raises MemoryError instead of
        # getting the whole kernel OOM-killed
        limit = _vm_data_bytes() + int(MEM_BUDGET_GB * 2**30)
        resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
        queue.put(evaluate(task, model_fn(), *data))
    except MemoryError:
        queue.put({"error": f"memory >{MEM_BUDGET_GB} GB"})
    except Exception as e:
        queue.put({"error": f"{type(e).__name__}: {e}"[:80]})


# Each run happens in a forked child: a hard timeout, a memory cap, and memory returned
# after every run. Without fork (Windows, macOS default) runs happen in-process, unguarded.
FORK = "fork" in mp.get_all_start_methods() and hasattr(resource, "RLIMIT_DATA")


def safe_evaluate(row, model_fn, *data):
    if not FORK:
        try:
            row.update(evaluate(row["task"], model_fn(), *data))
        except Exception as e:
            row["error"] = f"{type(e).__name__}: {e}"[:80]
        return row
    ctx = mp.get_context("fork")
    queue = ctx.Queue()
    p = ctx.Process(target=_child, args=(queue, row["task"], model_fn, data))
    p.start()
    try:
        row.update(queue.get(timeout=TIMEOUT_S))
    except Exception:  # queue.Empty: still running, or the child died
        row["error"] = f"timeout >{TIMEOUT_S}s" if p.is_alive() else f"crashed (exit {p.exitcode})"
    p.kill()
    p.join()
    return row


def split(task, X, y):
    if task == "classification":
        # Both packages' predict() returns class indices, so encode labels as 0..k-1
        y = LabelEncoder().fit_transform(y)
        return train_test_split(X, y, test_size=TEST_SIZE, random_state=SEED, stratify=y)
    return train_test_split(X, y, test_size=TEST_SIZE, random_state=SEED)

base_learners = [(n, c) for n, c in all_estimators(type_filter="regressor") if n not in EXCLUDED]

# Untimed warm-up so one-off import costs don't land on the first timed model
Xw, yw = make_friedman1(n_samples=60, random_state=0)
for task in BOOSTERS:
    yt = yw if task == "regression" else (yw > np.median(yw)).astype(int)
    for make in BOOSTERS[task].values():
        make(Ridge()).fit(Xw, yt).predict(Xw)

rows = []
for task, datasets in [("regression", REG_DATASETS), ("classification", CLF_DATASETS)]:
    for dname, (X, y) in datasets.items():
        data = split(task, X, y)
        for ref_name, ref_fn in REFERENCES[task].items():
            rows.append(safe_evaluate(dict(task=task, dataset=dname, base="-", model=ref_name),
                                      ref_fn, *data))
        for name, RegressorClass in tqdm(base_learners, desc=f"{task}: {dname}"):
            try:
                base = RegressorClass()
            except Exception as e:  # needs required constructor args
                rows.append(dict(task=task, dataset=dname, base=name, model="-",
                                 error=f"init: {e}"[:80]))
                continue
            for bname, make in BOOSTERS[task].items():
                rows.append(safe_evaluate(dict(task=task, dataset=dname, base=name, model=bname),
                                          lambda: make(seeded(base)), *data))

results = pd.DataFrame(rows)
results.to_csv("benchmark_vs_mlsauce.csv", index=False)
ok_mask = results["time_s"].notna()
print(f"{ok_mask.sum()} successful runs, {(~ok_mask).sum()} failed")

ok = results[results["time_s"].notna()]
METRICS = {"regression": ["rmse", "mae"], "classification": ["balanced_accuracy"]}
refs = {task: ok[(ok["task"] == task) & ok["model"].isin(REFERENCES[task])]
                .set_index(["dataset", "model"])[METRICS[task] + ["time_s"]]
        for task in METRICS}
pd.concat(refs, names=["task"]).round(3)

PRIMARY = {"regression": ("rmse", False), "classification": ("balanced_accuracy", True)}
wide = {task: (ok[(ok["task"] == task) & ok["model"].isin(BOOSTERS[task])]
               .pivot_table(index=["dataset", "base"], columns="model",
                            values=METRICS[task] + ["time_s"])
               .dropna())
        for task in METRICS}

summary = []
for task, datasets in [("regression", REG_DATASETS), ("classification", CLF_DATASETS)]:
    metric, higher_better = PRIMARY[task]
    for dname in datasets:
        t = wide[task].loc[dname]
        g, m = t[metric]["genbooster"], t[metric]["mlsauce"]
        tie = np.isclose(g, m, rtol=1e-3, atol=1e-9)
        g_better = (g > m) if higher_better else (g < m)
        r = refs[task]
        summary.append({
            "task": task, "dataset": dname, "metric": metric,
            "base learners": len(t),
            "genbooster better": int((g_better & ~tie).sum()),
            "mlsauce better": int((~g_better & ~tie).sum()),
            "ties": int(tie.sum()),
            "best genbooster": f"{(g.max() if higher_better else g.min()):.3f} "
                               f"({g.idxmax() if higher_better else g.idxmin()})",
            "best mlsauce": f"{(m.max() if higher_better else m.min()):.3f} "
                            f"({m.idxmax() if higher_better else m.idxmin()})",
            "XGBoost": round(r.loc[(dname, "XGBoost"), metric], 3),
            "LightGBM": round(r.loc[(dname, "LightGBM"), metric], 3),
            "total time gb / ms (s)": f"{t['time_s']['genbooster'].sum():.0f} / "
                                      f"{t['time_s']['mlsauce'].sum():.0f}",
        })
summary = pd.DataFrame(summary).set_index(["task", "dataset"]).T
summary

INK, INK_2, MUTED = "#0b0b0b", "#52514e", "#898781"
GRID, SURFACE, AXIS = "#e1e0d9", "#fcfcfb", "#c3c2b7"
STYLE = {  # color + marker shape: identity never rests on color alone
    # genbooster: larger ring drawn first, so a tie shows as a square inside a ring
    "genbooster": dict(color="#2a78d6", marker="o", s=90, zorder=3),
    "mlsauce": dict(color="#eb6834", marker="s", s=34, zorder=4),
}
REF_STYLE = {"XGBoost": "--", "LightGBM": ":"}


def dumbbell(ax, t, metric, higher_better, ref_x, title, top_n):
    r = t[metric]
    best = r.max(axis=1) if higher_better else r.min(axis=1)
    r = r.assign(best=best).sort_values("best", ascending=not higher_better).head(top_n)[::-1]
    ypos = np.arange(len(r))
    ax.set_facecolor(SURFACE)
    ax.hlines(ypos, r["genbooster"], r["mlsauce"], color=GRID, lw=2, zorder=1)
    for bname, st in STYLE.items():
        ax.scatter(r[bname], ypos, label=bname, edgecolor=SURFACE, linewidth=1.2, **st)
    for ref, ls in REF_STYLE.items():
        ax.axvline(ref_x[ref], color=INK_2, ls=ls, lw=1.2, zorder=2)
    # Reference labels sit above the plot area, stacked so they never hit the data
    for k, (ref, x) in enumerate(sorted(ref_x.items(), key=lambda kv: kv[1])):
        ax.annotate(f"{ref} {x:.3g}", (x, 1.0), xycoords=("data", "axes fraction"),
                    xytext=(0, 4 + 12 * k), textcoords="offset points",
                    fontsize=8.5, color=INK_2, ha="center", va="bottom")
    ax.margins(y=0.02)
    ax.set_yticks(ypos, r.index, fontsize=9, color=INK)
    ax.set_title(title, loc="left", fontsize=11.5, color=INK, pad=34)
    label = metric.replace("_", " ")
    ax.set_xlabel(f"test {label} ({'higher' if higher_better else 'lower'} is better)",
                  fontsize=9.5, color=INK_2)
    ax.grid(axis="x", color=GRID, lw=0.6)
    ax.set_axisbelow(True)
    ax.tick_params(axis="x", colors=MUTED, labelsize=8.5)
    ax.tick_params(axis="y", length=0)
    for side in ("top", "right", "left"):
        ax.spines[side].set_visible(False)
    ax.spines["bottom"].set_color(AXIS)


def plot_task(task, datasets, ncols, top_n, suptitle):
    metric, higher_better = PRIMARY[task]
    nrows = int(np.ceil(len(datasets) / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=(6.5 * ncols, nrows * (0.3 * top_n + 1.6) + 0.6),
                             facecolor=SURFACE, layout="constrained", squeeze=False)
    for ax, dname in zip(axes.flat, datasets):
        t = wide[task].loc[dname]
        ref_x = {ref: refs[task].loc[(dname, ref), metric] for ref in REF_STYLE}
        dumbbell(ax, t, metric, higher_better, ref_x,
                 f"{dname} — top {min(top_n, len(t))} base learners", top_n)
    for ax in axes.flat[len(datasets):]:
        ax.set_visible(False)
    handles, labels = axes.flat[0].get_legend_handles_labels()
    ref_handles = [plt.Line2D([], [], color=INK_2, ls=ls, lw=1.2) for ls in REF_STYLE.values()]
    fig.legend(handles + ref_handles, labels + [f"{r} (untuned)" for r in REF_STYLE],
               loc="outside lower center", ncol=4, frameon=False, fontsize=9.5, labelcolor=INK)
    fig.suptitle(suptitle, x=0.01, ha="left", fontsize=13.5, color=INK)
    return fig


fig = plot_task("regression", REG_DATASETS, ncols=1, top_n=25,
                suptitle="Regression (load_diabetes): untuned boosting of scikit-learn regressors")
fig.savefig("benchmark_regression.png", dpi=150, bbox_inches="tight", facecolor=SURFACE)
plt.show()

fig = plot_task("classification", CLF_DATASETS, ncols=3, top_n=15,
                suptitle="Classification: untuned boosting of scikit-learn regressors "
                         "(balanced accuracy)")
fig.savefig("benchmark_classification.png", dpi=150, bbox_inches="tight", facecolor=SURFACE)
plt.show()