"""
Weakness Region Analysis (Regression)
======================================

This example demonstrates AMIF (Adversarial Mutual Information Forest) weakness
region diagnostics for regression models. AMIF identifies model weakness regions
by partitioning data into a 2D grid of geometry scores (data density via CARF)
and MI scores (predictive information), then finding regions where the model
performs worst.

We use the CaliforniaHousing dataset with an XGBoost regressor.
"""

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

import warnings
warnings.filterwarnings("ignore")

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

# %%
# Load Dataset
# ------------
# Load the CaliforniaHousing dataset and set up for regression.

ds = DataSet()
ds.load("CaliforniaHousing")
ds.set_target(feature="MedHouseVal")
ds.set_task_type("Regression")
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 regressor.

model = MoXGBRegressor(
    name="XGBoost", n_estimators=200, max_depth=5, 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 MSE metric.
# The ``weak_fraction`` parameter controls what fraction of the worst-performing
# regions are flagged as weak.

ts = TestSuite(ds, model)

result = ts.diagnose_weakness_region(
    geometry_method="carf",
    metric="MSE",
    bins=8,
    weak_fraction=0.2,
    top_n_features=8,
)

# %%
# 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.
# Comparing train vs test performance reveals overfitting patterns.

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. MI importance ranking shows which features carry the
# most predictive information.

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))

# %%
# 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))