"""
DirectRS Classification
=======================

This example demonstrates DirectRS post-processing on a pre-trained XGBoost
classifier. DirectRS works on the raw score (logit) space, providing exact
additive decomposition of the decision function while improving or maintaining
classification accuracy.

We use the TaiwanCredit dataset with a depth-2 XGBoost base model.
"""

# %%
# Setup
# -----
# Import libraries and suppress warnings.

import warnings
warnings.filterwarnings("ignore")

import numpy as np
from modeva import DataSet, TestSuite
from modeva.models import MoXGBClassifier

# %%
# Load Dataset
# ------------
# Load the TaiwanCredit dataset and create a random train/test split.

ds = DataSet()
ds.load(name="TaiwanCredit")
ds.set_random_split()

# %%
# Train Base Model
# ----------------
# Train an XGBoost classifier with depth 2.

model = MoXGBClassifier(
    name="XGB-cls-depth2",
    n_estimators=200, max_depth=2, learning_rate=0.1,
    random_state=42, verbosity=0
)
model.fit(ds.train_x, ds.train_y.ravel())

ts = TestSuite(ds, model)
ts.diagnose_accuracy_table().table

# %%
# Fit DirectRS
# ------------
# Post-process the trained XGBoost classifier with DirectRS.

from modeva.models import MoDirectRSClassifier

drs = MoDirectRSClassifier(
    base_model=model, ridge_alpha=100.0, n_passes=1
)
drs.fit(ds.train_x, ds.train_y.ravel(), verbose=True)

# %%
# Accuracy Comparison
# -------------------
# Compare AUC and accuracy between the base XGBoost model and DirectRS.

from sklearn.metrics import roc_auc_score, accuracy_score

y_test = ds.test_y.ravel()
base_proba = model.predict_proba(ds.test_x)[:, 1]
drs_proba = drs.predict_proba(ds.test_x)[:, 1]
base_acc = accuracy_score(y_test, model.predict(ds.test_x))
drs_acc = accuracy_score(y_test, drs.predict(ds.test_x))

print(f"{'Metric':<10s} {'Base XGB':>10s} {'DirectRS':>10s}")
print("-" * 32)
print(f"{'AUC':<10s} {roc_auc_score(y_test, base_proba):>10.4f} {roc_auc_score(y_test, drs_proba):>10.4f}")
print(f"{'Accuracy':<10s} {base_acc:>10.4f} {drs_acc:>10.4f}")

# %%
# S' Stretch Analysis
# -------------------
# Analyze the global stretch matrix S' extracted from tree geometry.

result = drs.get_global_stretch_analysis(ds.feature_names)

# %%
# Eigenvalue spectrum of S'.
result.plot("eigenvalue_spectrum")

# %%
# Feature activity scores.
result.plot("feature_activity")

# %%
# Local Explanation
# -----------------
# For classification, the decomposition operates on raw scores (logits).
# We verify using ``drs._core.predict`` which returns raw scores.

result = drs.explain_local(ds.test_x, feature_names=ds.feature_names)
local = result.value

raw_pred = drs._core.predict(ds.test_x)
recon = local['intercept'] + local['contributions'].sum(axis=1)
max_err = np.max(np.abs(raw_pred - recon))

print(f"Max |raw_score - (intercept + sum contributions)|: {max_err:.2e}")
print(f"Decomposition exact to machine precision: {max_err < 1e-10}")

# %%
# Local explanation waterfall plot.
result.plot()

# %%
# Global Feature Importance
# -------------------------
# Compute global feature importance using the default ``slope`` mode.

result = drs.importance_global(feature_names=ds.feature_names)
result.plot()

# %%
# Main/Interaction Decomposition
# ------------------------------
# Decompose model variance into main effects and interactions.

result = drs.importance_main_interaction(ds.test_x, feature_names=ds.feature_names)
mi = result.value

print(f"Orthogonalized variance split:")
print(f"  eta2_main = {mi['eta2_main']:.4f}  ({mi['eta2_main']*100:.1f}%)")
print(f"  eta2_int  = {mi['eta2_int']:.4f}  ({mi['eta2_int']*100:.1f}%)")
print(f"  rho(g, r) = {mi['rho']:.4f}")

# %%
# Main vs interaction importance bar chart.
result.plot()

# %%
# Geometric Interaction Traces
# ----------------------------
# Trace feature interactions through the adjacency matrix A.

result = drs.geometric_interaction_traces(
    feature_names=ds.feature_names, K=4, gamma=0.5
)
traces = result.value

print("Interaction spectrum:")
for k in range(len(traces['T'])):
    print(f"  k={k+1}: T_k = {traces['T'][k]:.6f},  E_k = {traces['E'][k]:.6f}")

# %%
# Adjacency matrix heatmap.
result.plot("adjacency")

# %%
# Interaction energy spectrum.
result.plot("spectrum")

# %%
# Top Feature Interactions
# ------------------------
# Show the strongest off-diagonal entries in the stretch matrix.

result = drs.get_off_diagonal_analysis(ds.feature_names, top_k=15)
result.plot()

# %%
# FANOVA Comparison
# -----------------
# Compare DirectRS feature importance with FANOVA decomposition.

results = ts.interpret_fi()
results.plot()

# %%
results = ts.interpret_ei()
results.plot()

# %%
# Side-by-side comparison of importance scores.

result_drs = drs.importance_global(feature_names=ds.feature_names, mode="slope")
result_fanova = ts.interpret_fi()

drs_imp = result_drs.value['importance']
fanova_table = result_fanova.table
fanova_imp = dict(zip(fanova_table["Name"], fanova_table["Score"]))

print(f"{'Feature':<16s} {'DirectRS':>10s} {'FANOVA':>10s}")
print("-" * 38)
for i, feat in enumerate(ds.feature_names):
    print(f"{feat:<16s} {drs_imp[i]:>10.4f} {fanova_imp.get(feat, 0.0):>10.4f}")