"""Geographically weighted regression in Python (Part 2 of 2): GWR and MGWR with mgwr on median house prices
in West Yorkshire's 301 MSOAs, compared with the global OLS, spatial lag and spatial error models of Part 1.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/geographically-weighted-regression-python/
Data: ONS / HM Land Registry median price paid by MSOA, Census 2021 (Nomis) and ONS Open Geography boundaries,
all under the Open Government Licence v3.0. wy_data.py (shared with Part 1) downloads and joins them.
"""
import contextlib
import io
import json
import sys
import time
import warnings
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")   # render to files only
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
from matplotlib.patches import Patch
from libpysal.weights import Queen
import esda
import spreg
from mgwr.gwr import GWR, MGWR
from mgwr.sel_bw import Sel_BW

from wy_data import load_west_yorkshire

warnings.filterwarnings("ignore")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9, "axes.titlesize": 10})
OUT = Path("output"); OUT.mkdir(exist_ok=True)
PERMS = 9999
np.random.seed(42)   # esda.Moran has no seed argument
GREY = "#cfcfcf"


def section(title):
    print(f"\n=== {title} ===")


section("versions")
print(f"python {sys.version.split()[0]} | mgwr {version('mgwr')} | spreg {version('spreg')} | "
      f"libpysal {version('libpysal')} | esda {version('esda')} | geopandas {version('geopandas')} | "
      f"numpy {version('numpy')}")

# ---------------------------------------------------------------- Step 1: data, standardised
section("data")
gdf = load_west_yorkshire()
X_NAMES = ["level4", "detached", "no_car", "ln_density"]
NAMES = ["Intercept"] + X_NAMES
LABELS = {"Intercept": "Intercept (local price level)", "level4": "Level 4+ qualifications",
          "detached": "Detached houses", "no_car": "Households with no car", "ln_density": "ln(population density)"}
coords = np.column_stack([gdf.centroid.x, gdf.centroid.y])
y_raw = gdf[["ln_price"]].to_numpy()
X_raw = gdf[X_NAMES].to_numpy()
y = (y_raw - y_raw.mean()) / y_raw.std()
X = (X_raw - X_raw.mean(axis=0)) / X_raw.std(axis=0)
n = len(gdf)
print(f"{n} MSOAs; y = standardised ln(median price); X = standardised {', '.join(X_NAMES)}")
print(f"SD of ln(price) = {y_raw.std():.3f}; SDs of X: " +
      ", ".join(f"{k} {s:.2f}" for k, s in zip(X_NAMES, X_raw.std(axis=0))))

w = Queen.from_dataframe(gdf, use_index=False)
w.transform = "r"
lads = gdf.dissolve("lad").reset_index()


def outline(ax):
    lads.boundary.plot(ax=ax, color="#222222", linewidth=0.5)
    ax.set_axis_off()


def aicc(logl, p):
    """AICc with p parameters including the error variance; for OLS this equals mgwr's GWR AICc formula."""
    return -2 * logl + 2 * n * p / (n - p - 1)


def resid_moran(e):
    m = esda.Moran(np.asarray(e).flatten(), w, permutations=PERMS)
    return m.I, m.p_sim


# ---------------------------------------------------------------- Step 2: the global models from Part 1
section("global models on the standardised data (as in Part 1)")
ols = spreg.OLS(y, X, name_x=X_NAMES)
with contextlib.redirect_stdout(io.StringIO()):   # spreg prints the class name of each ML fit
    lag = spreg.ML_Lag(y, X, w=w, name_x=X_NAMES)
    err = spreg.ML_Error(y, X, w=w, name_x=X_NAMES)
glob = {"OLS": (ols.betas[:, 0], ols.r2, aicc(ols.logll, 6), resid_moran(ols.u)),
        "Spatial lag (ML)": (lag.betas[:5, 0], lag.pr2, aicc(lag.logll, 7), resid_moran(lag.u)),
        "Spatial error (ML)": (err.betas[:5, 0], err.pr2, aicc(err.logll, 7), resid_moran(err.e_filtered))}
for k, (b, r2, a, (mi, p)) in glob.items():
    print(f"{k:<19} R2 {r2:.3f}  AICc {a:7.1f}  residual Moran's I {mi:+.3f} (p {p:.4f})  betas " +
          " ".join(f"{v:+.3f}" for v in b))

# ---------------------------------------------------------------- Step 3: GWR bandwidth by AICc
section("GWR bandwidth selection (adaptive bisquare, AICc)")
t0 = time.perf_counter()
gwr_selector = Sel_BW(coords, y, X)   # defaults: adaptive (nearest neighbours), bisquare kernel, golden section
gwr_bw = gwr_selector.search(criterion="AICc")
t_gwr_bw = time.perf_counter() - t0
print(f"optimal bandwidth: {gwr_bw:.0f} nearest neighbours ({100 * gwr_bw / n:.0f}% of the MSOAs); "
      f"search took {t_gwr_bw:.1f} s")

grid = np.arange(30, 301, 6)
curve = np.array([GWR(coords, y, X, int(b)).fit().aicc for b in grid])
print("AICc profile: " + ", ".join(f"{b}: {a:.1f}" for b, a in zip(grid[::5], curve[::5])))

t0 = time.perf_counter()
gwr = GWR(coords, y, X, gwr_bw).fit()
print(f"GWR fit in {time.perf_counter() - t0:.2f} s")
gwr.summary()

# ---------------------------------------------------------------- Step 4: local t-values and multiple testing
section("GWR local significance with the corrected critical t-value")
# no arguments = the corrected 5% level; passing alpha=0.05 would give the UNcorrected critical value
crit = gwr.critical_tval()
naive_crit = gwr.critical_tval(alpha=0.05)
print(f"effective number of parameters ENP = {gwr.ENP:.1f}; adjusted alpha = 0.05 * k / ENP = {gwr.adj_alpha[1]:.5f}; "
      f"critical |t| = {crit:.2f} (vs {naive_crit:.2f} uncorrected)")
t_filt = gwr.filter_tvals()   # t-values with the non-significant ones set to 0
for j, nm in enumerate(NAMES):
    b = gwr.params[:, j]
    sig = t_filt[:, j] != 0
    naive = np.abs(gwr.tvalues[:, j]) > naive_crit
    print(f"{nm:<11} beta min {b.min():+.3f} median {np.median(b):+.3f} max {b.max():+.3f} | significant: "
          f"{sig.sum():3d} ({(b[sig] > 0).sum()} +, {(b[sig] < 0).sum()} -) | uncorrected: {naive.sum():3d}")

# ---------------------------------------------------------------- Step 5: MGWR
section("MGWR bandwidth selection (back-fitting, AICc)")
t0 = time.perf_counter()
mgwr_selector = Sel_BW(coords, y, X, multi=True)
mgwr_bws = mgwr_selector.search(criterion="AICc")
t_mgwr_bw = time.perf_counter() - t0
mgwr = MGWR(coords, y, X, mgwr_selector).fit()
t_mgwr = time.perf_counter() - t0
print(f"bandwidth search {t_mgwr_bw:.1f} s; search + fit {t_mgwr:.1f} s (GWR search took {t_gwr_bw:.1f} s)")
mgwr.summary()
with contextlib.redirect_stdout(io.StringIO()):
    bw_ci = mgwr.get_bws_intervals(mgwr_selector)
crit_j = mgwr.critical_tval()   # one corrected critical value per covariate
mt_filt = mgwr.filter_tvals()
for j, nm in enumerate(NAMES):
    b = mgwr.params[:, j]
    sig = mt_filt[:, j] != 0
    print(f"{nm:<11} bandwidth {mgwr_bws[j]:3.0f} (95% interval {bw_ci[j][0]:.0f}-{bw_ci[j][1]:.0f}) ENP_j {mgwr.ENP_j[j]:5.1f} "
          f"critical |t| {crit_j[j]:.2f} | beta min {b.min():+.3f} median {np.median(b):+.3f} max {b.max():+.3f} | "
          f"significant {sig.sum():3d} ({(b[sig] > 0).sum()} +, {(b[sig] < 0).sum()} -)")

# ---------------------------------------------------------------- Step 6: local R2 and collinearity
section("local R2 and local collinearity")
gdf["gwr_r2"] = gwr.localR2.flatten()
print(f"GWR local R2: min {gdf.gwr_r2.min():.2f}, median {gdf.gwr_r2.median():.2f}, max {gdf.gwr_r2.max():.2f}; "
      f"global GWR R2 {gwr.R2:.3f}, MGWR R2 {mgwr.R2:.3f}")
low = gdf.nsmallest(3, "gwr_r2")
print("lowest local R2: " + "; ".join(f"{r['name']} ({r.lad}) {r.gwr_r2:.2f}" for _, r in low.iterrows()))
corr, vifs, cn, vdp = gwr.local_collinearity()
gdf["gwr_cn"] = cn.flatten()
print(f"GWR local condition number: min {cn.min():.1f}, median {np.median(cn):.1f}, max {cn.max():.1f}; "
      f"> 30 in {(cn > 30).sum()} MSOAs")
print("GWR local VIF max per variable: " + ", ".join(f"{k} {v:.1f}" for k, v in zip(X_NAMES, vifs.max(axis=0))) +
      f"; MSOAs with any VIF > 10: {(vifs > 10).any(axis=1).sum()}")
pairs = [(a, b) for i, a in enumerate(X_NAMES) for b in X_NAMES[i + 1:]]
jmax = np.argmax(np.abs(corr).max(axis=0))
print(f"strongest local correlation: {pairs[jmax][0]} vs {pairs[jmax][1]}, from {corr[:, jmax].min():+.2f} to "
      f"{corr[:, jmax].max():+.2f} (global {np.corrcoef(X[:, X_NAMES.index(pairs[jmax][0])], X[:, X_NAMES.index(pairs[jmax][1])])[0, 1]:+.2f})")
mcn, _ = mgwr.local_collinearity()
print(f"MGWR local condition number: median {np.median(mcn):.1f}, max {mcn.max():.1f}; > 30 in {(mcn > 30).sum()} MSOAs")

# ---------------------------------------------------------------- Step 7: comparison
section("comparison with the global models")
mi_gwr = resid_moran(gwr.resid_response)
mi_mgwr = resid_moran(mgwr.resid_response)
rows = [(k, v[1], v[2], v[3][0], v[3][1]) for k, v in glob.items()]
rows += [(f"GWR (bw {gwr_bw:.0f})", gwr.R2, gwr.aicc, *mi_gwr), ("MGWR", mgwr.R2, mgwr.aicc, *mi_mgwr)]
comp = pd.DataFrame(rows, columns=["model", "R2", "AICc", "resid_I", "resid_p"])
comp["ENP"] = [5, 6, 6, gwr.ENP, mgwr.ENP]   # coefficients (+ rho or lambda); GWR/MGWR: trace of the hat matrix
print("(R2 for the lag model is spreg's pseudo R2; for the error model it is the R2 of the non-spatial part)")
print(comp.round(3).to_string(index=False))

# ---------------------------------------------------------------- figures
# 1. bandwidth selection
fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), gridspec_kw={"width_ratios": [1.2, 1]})
ax = axes[0]
ax.plot(grid, curve, "-", color="#0072B2", lw=1.6)
ax.scatter([gwr_bw], [gwr.aicc], color="#D55E00", zorder=3, s=40,
           label=f"golden-section optimum: {gwr_bw:.0f} neighbours (AICc {gwr.aicc:.1f})")
ax.axhline(glob["OLS"][2], color="#777777", ls="--", lw=1, label=f"OLS (AICc {glob['OLS'][2]:.1f})")
ax.set_xlabel("Adaptive bandwidth (nearest neighbours)"); ax.set_ylabel("AICc (lower is better)")
ax.set_title("GWR: AICc against bandwidth"); ax.legend(frameon=False, fontsize=8); ax.grid(color="#e5e5e5", lw=0.6)
ax = axes[1]
yy = np.arange(len(NAMES))
lo_ = [mgwr_bws[j] - bw_ci[j][0] for j in range(len(NAMES))]
hi_ = [bw_ci[j][1] - mgwr_bws[j] for j in range(len(NAMES))]
ax.errorbar(mgwr_bws, yy, xerr=[lo_, hi_], fmt="o", color="#009E73", capsize=3, ms=6, label="MGWR bandwidth (95% interval)")
ax.axvline(gwr_bw, color="#D55E00", ls=":", lw=1.4, label=f"GWR single bandwidth ({gwr_bw:.0f})")
ax.axvline(n - 1, color="#777777", ls="--", lw=1, label=f"all {n - 1} neighbours (global)")
ax.set_yticks(yy, [LABELS[k] for k in NAMES]); ax.invert_yaxis()
ax.set_xlim(0, 320); ax.set_xlabel("Bandwidth (nearest neighbours)")
ax.set_title("MGWR: one bandwidth per variable"); ax.legend(frameon=False, fontsize=7.5, loc="lower right")
ax.grid(axis="x", color="#e5e5e5", lw=0.6)
fig.tight_layout(); fig.savefig(OUT / "01-bandwidth-aicc.png"); plt.close(fig)


# one colour scale for the intercept and one shared by all slopes, in both figures, so GWR and MGWR are comparable
LIM_INT = max(np.abs(gwr.params[:, 0]).max(), np.abs(mgwr.params[:, 0]).max())
LIM_SLOPE = max(np.abs(gwr.params[:, 1:]).max(), np.abs(mgwr.params[:, 1:]).max())


def coef_maps(params, tfilt, title, fname, bws=None, crits=None):
    fig, axes = plt.subplots(2, 3, figsize=(12.5, 7.6))
    for j, (nm, ax) in enumerate(zip(NAMES, axes.flat)):
        b = params[:, j]
        lim = LIM_INT if j == 0 else LIM_SLOPE
        norm = TwoSlopeNorm(0, -lim, lim)
        sig = tfilt[:, j] != 0
        gdf.assign(b=b).plot("b", ax=ax, cmap="RdBu_r", norm=norm, edgecolor="white", linewidth=0.1)
        if (~sig).any():
            gdf[~sig].plot(ax=ax, color=GREY, edgecolor="white", linewidth=0.1)
        outline(ax)
        sm = plt.cm.ScalarMappable(cmap="RdBu_r", norm=norm)
        fig.colorbar(sm, ax=ax, shrink=0.7, label="Local coefficient (SD of ln price per SD)")
        extra = f"bandwidth {bws[j]:.0f}, |t| > {crits[j]:.2f}" if bws is not None else f"|t| > {crits:.2f}"
        ax.set_title(f"{LABELS[nm]}\n{extra}: {sig.sum()} of {n} significant", fontsize=9)
    ax = axes.flat[-1]; ax.set_axis_off()
    ax.legend(handles=[Patch(facecolor=GREY, edgecolor="#999999", label="Not significant after correction")],
              loc="center", frameon=False)
    ax.text(0.5, 0.25, "Red: positive local coefficient\nBlue: negative local coefficient\n"
            "Coefficients are standardised; the four slopes\nshare one colour scale in the GWR and MGWR maps", ha="center", va="center", transform=ax.transAxes, fontsize=8.5)
    fig.suptitle(title, fontsize=11)
    fig.tight_layout(); fig.savefig(OUT / fname); plt.close(fig)


coef_maps(gwr.params, t_filt, f"GWR local coefficients (adaptive bisquare, {gwr_bw:.0f} neighbours)",
          "02-gwr-coefficients.png", crits=crit)
coef_maps(mgwr.params, mt_filt, "MGWR local coefficients (one bandwidth per variable)", "03-mgwr-coefficients.png",
          bws=mgwr_bws, crits=crit_j)

# 4. local R2 and local condition number
fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.3))
gdf.plot("gwr_r2", ax=axes[0], cmap="viridis", edgecolor="white", linewidth=0.1, legend=True,
         legend_kwds={"label": "GWR local R²", "shrink": 0.75})
outline(axes[0]); axes[0].set_title(f"GWR local R² (median {gdf.gwr_r2.median():.2f})")
gdf.plot("gwr_cn", ax=axes[1], cmap="cividis", edgecolor="white", linewidth=0.1, legend=True,
         legend_kwds={"label": "Local condition number", "shrink": 0.75})
if (gdf.gwr_cn > 30).any():
    gdf[gdf.gwr_cn > 30].boundary.plot(ax=axes[1], color="#D55E00", linewidth=1.2)
n_hi = int((gdf.gwr_cn > 30).sum())
axes[1].set_title(f"GWR local condition number (max {gdf.gwr_cn.max():.1f}; "
                  + (f"{n_hi} MSOAs above 30 outlined)" if n_hi else "none above the warning level of 30)"))
outline(axes[1])
fig.tight_layout(); fig.savefig(OUT / "04-local-r2-condition-number.png"); plt.close(fig)

# 5. comparison with Part 1's global models
fig, axes = plt.subplots(1, 2, figsize=(12, 4.3), gridspec_kw={"width_ratios": [1, 1.35]})
ax = axes[0]
cols = ["#999999", "#E69F00", "#56B4E9", "#D55E00", "#009E73"]
ax.set_axisbelow(True); ax.grid(axis="x", color="#e5e5e5", lw=0.6)
ax.barh(comp.model, comp.AICc, color=cols)
for i, (a, mi) in enumerate(zip(comp.AICc, comp.resid_I)):
    ax.text(a + 8, i, f"AICc {a:.0f}\nresidual I {mi:+.2f}", va="center", ha="left", fontsize=8, color="#111111")
ax.set_xlim(0, comp.AICc.max() * 1.32)
ax.invert_yaxis(); ax.set_xlabel("AICc on standardised data (lower is better)")
ax.set_title("Fit and residual Moran's I by model")
ax = axes[1]
pos = np.arange(1, len(X_NAMES) + 1)
bp1 = ax.boxplot([gwr.params[:, j] for j in range(1, 5)], positions=pos - 0.18, widths=0.3, patch_artist=True,
                 showfliers=False)
bp2 = ax.boxplot([mgwr.params[:, j] for j in range(1, 5)], positions=pos + 0.18, widths=0.3, patch_artist=True,
                 showfliers=False)
for bp, c in [(bp1, "#D55E00"), (bp2, "#009E73")]:
    for patch in bp["boxes"]:
        patch.set_facecolor(c); patch.set_alpha(0.55)
    for med in bp["medians"]:
        med.set_color("#111111")
ax.scatter(pos, ols.betas[1:, 0], marker="D", color="#333333", zorder=4, s=28, label="OLS")
ax.scatter(pos, err.betas[1:5, 0], marker="s", color="#56B4E9", edgecolor="#333333", zorder=4, s=28,
           label="Spatial error")
ax.axhline(0, color="#777777", lw=0.8)
ax.set_xticks(pos, [LABELS[k].replace(" ", "\n", 1) for k in X_NAMES])
ax.set_ylabel("Coefficient (SD of ln price per SD)")
ax.legend(handles=[Patch(facecolor="#D55E00", alpha=0.55, label="GWR local (IQR, whiskers 1.5 IQR)"),
                   Patch(facecolor="#009E73", alpha=0.55, label="MGWR local"),
                   *ax.get_legend_handles_labels()[0]], fontsize=7.5, frameon=False, loc="lower left")
ax.set_title("Local coefficients against the global estimates"); ax.set_axisbelow(True); ax.grid(axis="y", color="#e5e5e5", lw=0.6)
fig.tight_layout(); fig.savefig(OUT / "05-model-comparison.png"); plt.close(fig)

gdf[["msoa", "name", "lad", "ln_price", "gwr_r2", "gwr_cn", "geometry"]].assign(
    **{f"gwr_{k}": gwr.params[:, j] for j, k in enumerate(NAMES)},
    **{f"mgwr_{k}": mgwr.params[:, j] for j, k in enumerate(NAMES)},
    **{f"mgwr_t_{k}": mt_filt[:, j] for j, k in enumerate(NAMES)}).to_file(OUT / "gwr_mgwr_results.gpkg",
                                                                            driver="GPKG")
metrics = {"gwr_bw": gwr_bw, "gwr_aicc": gwr.aicc, "gwr_r2": gwr.R2, "mgwr_aicc": mgwr.aicc, "mgwr_r2": mgwr.R2,
           "ols_aicc": glob["OLS"][2], "err_aicc": glob["Spatial error (ML)"][2], "mgwr_seconds": t_mgwr}
(OUT / "metrics.json").write_text(json.dumps({k: round(float(v), 4) for k, v in metrics.items()}, indent=2))
print("\nwrote figures, output/gwr_mgwr_results.gpkg and output/metrics.json")