"""
Benchmark: insurance-glm-tools (FactorClusterer / R2VF)
=========================================================

Scenario: A UK motor insurer has 30 postcode districts. The true risk structure
has 7 territory bands. The actuary needs to group these 30 districts into
rating territories.

We compare:
  Baseline: Manual quintile banding — sort districts by fitted GLM relativity,
            cut into 5 equal-exposure groups. Standard actuarial practice.

  Library:  R2VF via FactorClusterer (BIC-selected fused lasso). Merges adjacent
            districts (ordered by fitted relativity) where the difference shrinks
            to zero. Produces data-driven band boundaries.

IMPORTANT: R2VF operates on *adjacent* levels — it requires the factor to be
ordered. For postcode districts this means pre-ordering by fitted GLM
relativity (ascending) before calling FactorClusterer.

Metrics:
  - Test Poisson deviance
  - Band recovery accuracy (fraction of districts correctly assigned to true band)
  - Train BIC
  - Number of bands produced

Note: uses scipy from the insurance-glm-tools dependency tree. No sklearn import
(avoids numpy/scipy version conflicts in Databricks serverless).

Seed: 42, n=20,000 policies, 30 districts.
"""

import time
import numpy as np
import statsmodels.api as sm
import pandas as pd

from insurance_glm_tools.cluster import FactorClusterer

print("=" * 60)
print("Benchmark: insurance-glm-tools (R2VF factor clustering)")
print("=" * 60)

# ---------------------------------------------------------------------------
# Data generating process
# ---------------------------------------------------------------------------
rng = np.random.default_rng(42)
n = 20_000

# 30 postcode districts, 7 true territory bands
n_districts = 30
n_true_bands = 7
# Districts 0-3 = band 0, 4-7 = band 1, ..., 28-29 = band 6
true_band = np.array([min(d // (n_districts // n_true_bands), n_true_bands - 1)
                       for d in range(n_districts)])
# True log relativities for each band
band_log_rel = np.array([-0.30, -0.15, 0.00, 0.15, 0.35, 0.55, 0.80])
district_log_rel = band_log_rel[true_band]

vehicle_group = rng.integers(1, 20, n).astype(float)
ncd_years     = rng.integers(0, 9, n).astype(float)
district      = rng.integers(0, n_districts, n)
exposure      = rng.uniform(0.3, 1.0, n)

log_rate = (
    -2.5
    + 0.015 * vehicle_group
    - 0.08 * ncd_years
    + district_log_rel[district]
)
y = rng.poisson(np.exp(log_rate) * exposure).astype(float)

df = pd.DataFrame({
    "vehicle_group": vehicle_group,
    "ncd_years":     ncd_years,
    "district":      district.astype(float),
    "exposure":      exposure,
})

# Train/test split
n_train = int(0.8 * n)
idx_all = np.arange(n)
rng.shuffle(idx_all)
train_idx, test_idx = idx_all[:n_train], idx_all[n_train:]

df_train = df.iloc[train_idx].reset_index(drop=True)
df_test  = df.iloc[test_idx].reset_index(drop=True)
y_train  = y[train_idx]
y_test   = y[test_idx]
exp_train = exposure[train_idx]
exp_test  = exposure[test_idx]

print(f"\nPolicies: {n} (train={n_train}, test={len(test_idx)})")
print(f"Districts: {n_districts}, True territory bands: {n_true_bands}")
print(f"Mean frequency: {(y / exposure).mean():.4f}")


def poisson_deviance_sum(y_true, y_pred):
    y_pred_c = np.maximum(y_pred, 1e-10)
    mask = y_true > 0
    d = np.zeros_like(y_true, dtype=float)
    d[mask] = y_true[mask] * np.log(y_true[mask] / y_pred_c[mask]) - (y_true[mask] - y_pred_c[mask])
    d[~mask] = y_pred_c[~mask]
    return 2.0 * d.sum()


def band_recovery_accuracy(pred_bands, true_bands):
    """
    Fraction of districts correctly grouped relative to true band structure.
    Two districts are 'correctly grouped' if they are in the same true band
    AND the same predicted band, or in different true bands AND different predicted bands.
    Uses the Hungarian-algorithm-free version: for each true band, assign it
    to the most common predicted band, then compute accuracy.
    """
    n = len(pred_bands)
    # Find best matching between predicted and true bands
    from collections import Counter
    true_unique = sorted(set(true_bands))
    pred_unique = sorted(set(pred_bands))

    # For each true band, find the most common predicted band
    best_match = {}
    for tb in true_unique:
        mask = np.array(true_bands) == tb
        if mask.sum() > 0:
            pred_in_true = np.array(pred_bands)[mask]
            best_match[tb] = Counter(pred_in_true).most_common(1)[0][0]

    # Accuracy: fraction where pred_band[i] == best_match[true_band[i]]
    correct = sum(1 for i in range(n) if pred_bands[i] == best_match.get(true_bands[i], -1))
    return correct / n


# ---------------------------------------------------------------------------
# Step 1: Initial GLM for district relativity ordering (both approaches need this)
# ---------------------------------------------------------------------------
print("\n--- Step 1: Initial GLM for district relativity ordering ---")

t0 = time.time()
district_dummies_train = pd.get_dummies(df_train["district"].astype(int), prefix="d", drop_first=True)
X_init_train = pd.concat([
    df_train[["vehicle_group", "ncd_years"]],
    district_dummies_train,
], axis=1).astype(float)

X_sm_init = sm.add_constant(X_init_train)
glm_init = sm.GLM(
    y_train, X_sm_init,
    family=sm.families.Poisson(),
    offset=np.log(np.maximum(exp_train, 1e-10)),
).fit(disp=False)
t_init = time.time() - t0
print(f"  Initial GLM fit time: {t_init:.2f}s")

# Extract district coefficients and ordering
d_cols = [c for c in X_init_train.columns if c.startswith("d_")]
d_coeffs = {int(c.split("_")[1]): float(glm_init.params.get(c, 0.0)) for c in d_cols}
d_coeffs[0] = 0.0  # reference district

district_order = sorted(range(n_districts), key=lambda d: d_coeffs.get(d, 0.0))
relativity_order_map = {d: i for i, d in enumerate(district_order)}
df_train["district_ordered"] = df_train["district"].astype(int).map(relativity_order_map)
df_test["district_ordered"]  = df_test["district"].astype(int).map(relativity_order_map)

# True band assignment in ordered space
true_band_ordered = np.array([true_band[d] for d in district_order])

# ---------------------------------------------------------------------------
# Baseline: Manual quintile banding
# ---------------------------------------------------------------------------
print("\n--- Baseline: Manual quintile banding (5 bands) ---")

n_manual_bands = 5
district_sorted = sorted(d_coeffs.items(), key=lambda x: x[1])
quintile_size = n_districts // n_manual_bands
quintile_map = {}
for i, (d, _) in enumerate(district_sorted):
    quintile_map[d] = min(i // quintile_size, n_manual_bands - 1)

df_train["district_band_manual"] = df_train["district"].astype(int).map(quintile_map).astype(float)
df_test["district_band_manual"]  = df_test["district"].astype(int).map(quintile_map).astype(float)

t0 = time.time()
band_dummies_train = pd.get_dummies(df_train["district_band_manual"].astype(int), prefix="band", drop_first=True)
X_manual_train = pd.concat([df_train[["vehicle_group", "ncd_years"]], band_dummies_train], axis=1).astype(float)
X_sm_manual_train = sm.add_constant(X_manual_train)
glm_manual = sm.GLM(
    y_train, X_sm_manual_train,
    family=sm.families.Poisson(),
    offset=np.log(np.maximum(exp_train, 1e-10)),
).fit(disp=False)

band_dummies_test = pd.get_dummies(df_test["district_band_manual"].astype(int), prefix="band", drop_first=True)
X_manual_test = pd.concat([df_test[["vehicle_group", "ncd_years"]], band_dummies_test], axis=1).astype(float)
X_manual_test = X_manual_test.reindex(columns=X_manual_train.columns, fill_value=0)
X_sm_manual_test = sm.add_constant(X_manual_test)
X_sm_manual_test = X_sm_manual_test.reindex(columns=X_sm_manual_train.columns, fill_value=0)

y_pred_manual = glm_manual.predict(X_sm_manual_test, offset=np.log(np.maximum(exp_test, 1e-10)))
t_manual = time.time() - t0

dev_manual = poisson_deviance_sum(y_test, y_pred_manual)
bic_manual = glm_manual.bic
n_manual_actual = int(df_train["district_band_manual"].nunique())

pred_bands_manual = np.array([quintile_map[d] for d in range(n_districts)])
bra_manual = band_recovery_accuracy(pred_bands_manual, true_band)

print(f"  Fit time: {t_manual:.2f}s")
print(f"  Bands produced: {n_manual_actual}")
print(f"  Test Poisson deviance: {dev_manual:.1f}")
print(f"  Train BIC: {bic_manual:.1f}")
print(f"  Band recovery accuracy: {bra_manual:.3f}")

# ---------------------------------------------------------------------------
# Library: R2VF via FactorClusterer (BIC-selected fused lasso)
# ---------------------------------------------------------------------------
print("\n--- Library: R2VF via FactorClusterer (BIC-selected) ---")

t0 = time.time()
fc = FactorClusterer(family="poisson", lambda_="bic", min_exposure=200)
fc.fit(
    df_train[["vehicle_group", "ncd_years", "district_ordered"]],
    y_train,
    exposure=exp_train,
    ordinal_factors=["district_ordered"],
)
t_r2vf = time.time() - t0

X_r2vf_train = fc.transform(df_train[["vehicle_group", "ncd_years", "district_ordered"]])
X_r2vf_test  = fc.transform(df_test[["vehicle_group", "ncd_years", "district_ordered"]])

result_r2vf = fc.refit_glm(X_r2vf_train, y_train, exposure=exp_train)

X_r2vf_train_sm = sm.add_constant(X_r2vf_train)
X_r2vf_test_sm  = sm.add_constant(X_r2vf_test)
X_r2vf_test_sm  = X_r2vf_test_sm.reindex(columns=X_r2vf_train_sm.columns, fill_value=0)

y_pred_r2vf = result_r2vf.predict(X_r2vf_test_sm, offset=np.log(np.maximum(exp_test, 1e-10)))

dev_r2vf = poisson_deviance_sum(y_test, y_pred_r2vf)
bic_r2vf = result_r2vf.bic

lm = fc.level_map("district_ordered")
lm_df = lm.to_df()
n_r2vf_bands = int(lm_df["merged_group"].nunique())

# Band recovery: map district (by relativity order) to merged_group
r2vf_assign = lm_df.set_index("original_level")["merged_group"].to_dict()
pred_bands_r2vf = np.array([r2vf_assign.get(i, 0) for i in range(n_districts)])
bra_r2vf = band_recovery_accuracy(pred_bands_r2vf, true_band_ordered)

print(f"  Fit time (IRLS + BIC grid): {t_r2vf:.1f}s")
print(f"  Bands produced: {n_r2vf_bands}")
print(f"  Test Poisson deviance: {dev_r2vf:.1f}")
print(f"  Train BIC: {bic_r2vf:.1f}")
print(f"  Band recovery accuracy: {bra_r2vf:.3f}")

# Show level map
print(f"\n  Level map (district_ordered -> merged band):")
for _, row in lm_df.iterrows():
    exp_val = row.get('group_exposure', row.get('exposure', 0.0))
    print(f"    ordered={int(row['original_level']):2d} -> band={int(row['merged_group'])}  "
          f"(true_band={true_band_ordered[int(row['original_level'])]})")

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)
print(f"{'Metric':<38} {'Manual quintiles':>16} {'R2VF (library)':>14}")
print("-" * 70)
print(f"{'Territory bands produced':<38} {n_manual_actual:>16} {n_r2vf_bands:>14}")
print(f"{'Test Poisson deviance':<38} {dev_manual:>16.1f} {dev_r2vf:>14.1f}")
print(f"{'Train BIC':<38} {bic_manual:>16.1f} {bic_r2vf:>14.1f}")
print(f"{'Band recovery accuracy (vs true DGP)':<38} {bra_manual:>16.3f} {bra_r2vf:>14.3f}")
print(f"{'Fit time':<38} {'<1s':>16} {t_r2vf:.0f}s")
print(f"\nTrue territory structure: {n_true_bands} bands")
print(f"Band recovery improvement (R2VF - manual): {bra_r2vf - bra_manual:+.3f}")
print(f"\nNote: R2VF requires pre-ordering districts by fitted GLM relativity.")
print("Without ordering, the fused lasso has no meaningful adjacency structure.")
