"""
==================================================
Weak-Cluster Detection and Repair (FuseKernel + MoE)
==================================================

Detect the weak regions of a model with FuseKernel, then repair them with a
residual-driven Mixture of Experts. FuseKernel clusters samples in its learned
kernel geometry and breaks performance down per cluster (automatic weakness
detection); ``MoMoERegressor`` with ``cluster_method="ltc"`` then fits a
specialised expert per learning-trajectory cluster, improving accuracy most where
the base model was weakest.

For fast detection the FuseKernel uses only the XGBoost tree co-membership kernel
(``use_rbf=False``, ``use_spectral=False``) at the base model's depth, decoded with
the Nystrom solver -- no dense RBF or spectral fit is needed just to locate the
weak clusters. A subsample of CaliforniaHousing keeps the example quick.
"""

# %%
# Installation

# To install the required package, use the following command:
# !pip install modeva

# %%
# Authentication

# To get authentication, use the following command: (To get full access please replace the token to your own token)
# from modeva.utils.authenticate import authenticate
# authenticate(auth_code='eaaa4301-b140-484c-8e93-f9f633c8bacb')

# %%
# Import required modules
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
from sklearn.metrics import r2_score
from modeva import DataSet
from modeva.models import MoXGBRegressor, MoFuseKernelRegressor, MoMoERegressor

# %%
# Load and subsample the dataset
# ----------------------------------------------------------
ds = DataSet()
ds.load(name="CaliforniaHousing")
ds.set_active_samples()
sub = ds.subsample_random(dataset="main", sample_size=3000)
ds.set_active_samples(dataset="main", sample_idx=sub.value["sample_idx"])
ds.set_random_split()

# %%
# Base model
# ----------------------------------------------------------
# A depth-3 XGBoost whose weak regions we will detect and repair.
base = MoXGBRegressor(name="base", max_depth=3, n_estimators=100,
                      learning_rate=0.1, random_state=0)
base.fit(ds.train_x, ds.train_y.ravel())

# %%
# Detect weak clusters with FuseKernel
# ----------------------------------------------------------
# FuseKernel clusters samples in its learned kernel geometry; ``diagnose_weak_clusters``
# breaks train/test performance down per cluster so the weak regions surface
# automatically. For fast detection we use only the tree co-membership kernel at the
# base model's depth (``use_rbf=False``, ``use_spectral=False``), decoded with Nystrom.
fk = MoFuseKernelRegressor(name="FuseKernel", use_xgb=True, use_rbf=False,
                           use_spectral=False, fit_method="grid", solver="nystrom",
                           gbdt_params={"n_estimators": 100, "max_depth": 3})
fk.fit(ds.train_x, ds.train_y.ravel())

diag = fk.diagnose_weak_clusters(ds, n_clusters=6)
diag.table

# %%
# The weakest cluster and its test-row mask.
worst = diag.value["worst_clusters"]
worst_id = int(worst.iloc[0]["cluster"])
weak_mask = np.asarray(diag.value["labels_test"]) == worst_id
print(f"weak cluster = {worst_id}, test samples in it = {int(weak_mask.sum())}")
worst

# %%
# Repair with a residual-driven Mixture of Experts
# ----------------------------------------------------------
# ``MoMoERegressor`` with ``cluster_method="ltc"`` fits a baseline, clusters samples
# by their learning trajectories (where the base model struggles), and fits a
# specialised expert per cluster with a gate routing between them.
moe = MoMoERegressor(name="MoE-repair", n_clusters=5, cluster_method="ltc",
                     expert="xgboost", max_depth=3, n_estimators=100, random_state=0)
moe.fit(ds.train_x, ds.train_y.ravel())

# %%
# Before vs after
# ----------------------------------------------------------
# The repair helps overall, and most on the FuseKernel-detected weak cluster.
Xte, yte = np.asarray(ds.test_x), np.asarray(ds.test_y).ravel()

def r2(model, mask=None):
    p = model.predict(Xte)
    if mask is None:
        mask = np.ones(len(yte), dtype=bool)
    return round(r2_score(yte[mask], p[mask]), 4)

pd.DataFrame({
    ("overall", "base"):                          {"R2": r2(base)},
    ("overall", "MoE repair"):                     {"R2": r2(moe)},
    (f"weak cluster {worst_id}", "base"):          {"R2": r2(base, weak_mask)},
    (f"weak cluster {worst_id}", "MoE repair"):    {"R2": r2(moe, weak_mask)},
}).T