"""
Weakness Region Analysis (Classification)
==========================================

This example demonstrates AMIF (Adversarial Mutual Information Forest) weakness
region diagnostics for classification models. In addition to the regression
diagnostics, classification includes confusion matrix comparisons between the
full test set and weak regions.

We use the TaiwanCredit dataset with an XGBoost classifier.
"""

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

import warnings
warnings.filterwarnings("ignore")

from modeva import DataSet, TestSuite
from modeva.models import MoXGBClassifier

# %%
# Load Dataset
# ------------
# Load the TaiwanCredit dataset and set up for classification.

ds = DataSet()
ds.load("TaiwanCredit")
ds.set_target(feature="FlagDefault")
ds.set_task_type("Classification")
ds.set_random_split(test_ratio=0.2)

print(f"Train: {ds.train_x.shape}, Test: {ds.test_x.shape}")

# %%
# Train Model
# -----------
# Train an XGBoost classifier.

model = MoXGBClassifier(
    name="XGBoost", n_estimators=200, max_depth=4, learning_rate=0.1
)
model.fit(ds.train_x, ds.train_y.ravel())

# %%
# Run Weakness Region Diagnostics
# --------------------------------
# Use ``diagnose_weakness_region`` with the CARF geometry method and ACC metric.
# The ``min_count`` parameter sets the minimum number of samples required
# per region bin to compute a valid metric.

ts = TestSuite(ds, model)

result = ts.diagnose_weakness_region(
    geometry_method="carf",
    metric="ACC",
    bins=8,
    weak_fraction=0.2,
    top_n_features=10,
    min_count=15,
)

# %%
# Result Table
# ------------
# The result table summarizes performance across all 2D grid regions.

result.table

# %%
# Print the weak region summary statistics.

v = result.value
print(f"Weak test samples: {v['n_weak_samples']} / {v['n_total_samples']} "
      f"({100*v['n_weak_samples']/v['n_total_samples']:.1f}%)")
print(f"Metric: {v['metric']}, Cutoff: {v['cutoff']:.4f}")

# %%
# Region Performance Heatmaps
# ----------------------------
# Heatmaps show model performance across the 2D geometry-MI grid.

result.plot(name="region_performance_train", figsize=(6.5, 5))

# %%
# Test performance heatmap.
result.plot(name="region_performance_test", figsize=(6.5, 5))

# %%
# Sample count per region.
result.plot(name="region_sample_count", figsize=(6.5, 5))

# %%
# Feature Rankings
# ----------------
# JS divergence ranking shows which features differ most between weak and
# non-weak regions.

result.plot(name="js_divergence_ranking", figsize=(6.5, 4))

# %%
# MI importance ranking.
result.plot(name="mi_importance_ranking", figsize=(6.5, 4))

# %%
# Score Distributions
# -------------------
# Distribution of geometry and MI scores across the dataset.

result.plot(name="geometry_score_distribution", figsize=(6.5, 4))

# %%
# MI score distribution.
result.plot(name="mi_score_distribution", figsize=(6.5, 4))

# %%
# Confusion Matrices
# ------------------
# Compare confusion matrices between the full test set and the weak region.
# This reveals whether certain classes are disproportionately affected.

result.plot(name=("confusion_matrix", "all_test"), figsize=(5, 4))

# %%
# Confusion matrix for weak region samples only.
result.plot(name=("confusion_matrix", "weak_test"), figsize=(5, 4))

# %%
# Feature Distributions (Weak vs. Rest)
# --------------------------------------
# Compare feature distributions between weak and non-weak samples
# for the top 3 most divergent features.

feat_figs = [
    f for f in result.get_figure_names()
    if isinstance(f, tuple) and f[0] == "feature_distribution"
]
for fig_name in feat_figs[:3]:
    result.plot(name=fig_name, figsize=(6.5, 4))