"""Spatial regression in Python (Part 1 of 2): OLS diagnostics, spatial lag and spatial error models with spreg,
on median house prices in West Yorkshire's 301 MSOAs.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/spatial-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. See wy_data.py for the downloads.
"""
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 LogNorm, TwoSlopeNorm
from matplotlib.ticker import FixedLocator, FuncFormatter, NullLocator
import libpysal
from libpysal.weights import Queen, KNN
import esda
import spreg

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)
SEED = 42
PERMS = 9999
rng = np.random.default_rng(SEED)
np.random.seed(SEED)   # esda.Moran has no seed argument


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


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

# ---------------------------------------------------------------- Step 1: data
section("data")
gdf = load_west_yorkshire()
print(f"{len(gdf)} MSOAs with a median price and census data: " +
      ", ".join(f"{k} {v}" for k, v in gdf.lad.value_counts().items()))
X_NAMES = ["level4", "detached", "no_car", "ln_density"]
LABELS = {"level4": "Level 4+ qualifications (%)", "detached": "Detached houses (%)",
          "no_car": "Households with no car (%)", "ln_density": "ln(population density)"}
print(gdf[["price", "level4", "detached", "no_car", "density"]].describe().loc[["mean", "min", "50%", "max"]]
      .round(1).to_string())
lo, hi = gdf.loc[gdf.price.idxmin()], gdf.loc[gdf.price.idxmax()]
print(f"cheapest: {lo['name']} ({lo.lad}) GBP {lo.price:,.0f} | dearest: {hi['name']} ({hi.lad}) GBP {hi.price:,.0f}")
print("correlations with ln(price): " +
      ", ".join(f"{c} {gdf.ln_price.corr(gdf[c]):+.2f}" for c in X_NAMES))

lads = gdf.dissolve("lad").reset_index()


def outline(ax, labels=False):
    lads.boundary.plot(ax=ax, color="#222222", linewidth=0.6)
    if labels:
        for _, r in lads.iterrows():
            p = r.geometry.representative_point()
            ax.annotate(r.lad, (p.x, p.y), ha="center", fontsize=7.5, color="#111111",
                        bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.75))
    ax.set_axis_off()


fig, ax = plt.subplots(figsize=(7.2, 5.4))
gdf.plot("price", ax=ax, cmap="viridis", norm=LogNorm(gdf.price.min(), gdf.price.max()), edgecolor="white",
         linewidth=0.15, legend=True,
         legend_kwds={"label": "Median price paid, year to March 2026 (GBP, log scale)", "shrink": 0.75})
cb = ax.get_figure().axes[-1]
cb.yaxis.set_major_locator(FixedLocator([80e3, 100e3, 150e3, 200e3, 300e3, 400e3, 600e3]))
cb.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"£{v / 1000:,.0f}k"))
cb.yaxis.set_minor_locator(NullLocator())
outline(ax, labels=True)
ax.set_title("Median house price by MSOA, West Yorkshire (301 MSOAs)")
fig.savefig(OUT / "01-median-price-msoa.png"); plt.close(fig)

# ---------------------------------------------------------------- Step 2: weights and OLS
section("spatial weights")
w_queen = Queen.from_dataframe(gdf, use_index=False)
print(f"queen contiguity: mean {w_queen.mean_neighbors:.2f} neighbours (min {w_queen.min_neighbors}, "
      f"max {w_queen.max_neighbors}), islands {len(w_queen.islands)}")
w_queen.transform = "r"

y = gdf[["ln_price"]].to_numpy()
X = gdf[X_NAMES].to_numpy()

section("OLS with spatial diagnostics (queen weights)")
ols = spreg.OLS(y, X, w=w_queen, spat_diag=True, moran=True,
                name_y="ln_price", name_x=X_NAMES, name_w="queen", name_ds="West Yorkshire MSOAs")
print(ols.summary)
vif = spreg.vif(ols)
print("VIF: " + ", ".join(f"{n} {v[0]:.2f}" for n, v in zip(X_NAMES, vif[1:])))

section("Moran's I of the OLS residuals")
e_ols = ols.u.flatten()
mi_ols = esda.Moran(e_ols, w_queen, permutations=PERMS)
print(f"I = {mi_ols.I:.3f}  E[I] = {mi_ols.EI:.4f}  z(permutation) = {mi_ols.z_sim:.1f}  pseudo p = {mi_ols.p_sim:.4f}")
print(f"largest of {PERMS} permuted I values: {mi_ols.sim.max():.3f}")
gdf["e_ols"] = e_ols
worst = gdf.reindex(gdf.e_ols.abs().sort_values(ascending=False).index)[:4]
for _, r in worst.iterrows():
    print(f"  {r['name']} ({r.lad}): residual {r.e_ols:+.3f} = price {100 * (np.exp(r.e_ols) - 1):+.0f}% vs the OLS fit")

lim = np.abs(e_ols).max()
fig, axes = plt.subplots(1, 2, figsize=(11, 4.6), gridspec_kw={"width_ratios": [1.25, 1]})
gdf.plot("e_ols", ax=axes[0], cmap="RdBu_r", norm=TwoSlopeNorm(0, -lim, lim), edgecolor="white", linewidth=0.15,
         legend=True, legend_kwds={"label": "OLS residual (log points; + = dearer than predicted)", "shrink": 0.8})
outline(axes[0], labels=True)
axes[0].set_title(f"OLS residuals (Moran's I = {mi_ols.I:.3f})")
z = (e_ols - e_ols.mean()) / e_ols.std()
lag_z = libpysal.weights.lag_spatial(w_queen, z)
ax = axes[1]
ax.scatter(z, lag_z, s=10, color="#0072B2", alpha=0.7, edgecolor="none")
b = np.polyfit(z, lag_z, 1)
xx = np.linspace(z.min(), z.max(), 50)
ax.plot(xx, b[0] * xx + b[1], color="#D55E00", lw=1.6, label=f"slope = Moran's I = {mi_ols.I:.3f}")
ax.axhline(0, color="#888888", lw=0.6); ax.axvline(0, color="#888888", lw=0.6)
ax.set_xlabel("OLS residual (z-score)"); ax.set_ylabel("Average residual of neighbours (spatial lag, z-score)")
ax.set_title(f"Moran scatterplot of the OLS residuals (slope = I = {mi_ols.I:.3f})")
for t, (px, py) in {"HH": (0.93, 0.95), "LL": (0.07, 0.04), "LH": (0.07, 0.95), "HL": (0.93, 0.04)}.items():
    ax.text(px, py, t, transform=ax.transAxes, ha="center", va="center", color="#555555", fontsize=9)
fig.tight_layout(); fig.savefig(OUT / "02-ols-residuals-moran.png"); plt.close(fig)


# ---------------------------------------------------------------- Step 3: spatial lag and spatial error models
def lag_impacts(betas, rho, eig):
    """Direct, indirect and total impacts of a spatial lag model (LeSage and Pace 2009).
    Direct = beta * mean(diag((I - rho W)^-1)), computed from the eigenvalues of W; total = beta / (1 - rho)."""
    m_dir = np.mean(1.0 / (1.0 - rho * eig)).real
    direct = betas * m_dir
    total = betas / (1.0 - rho)
    return direct, total - direct, total


def lag_impact_ci(model, eig, draws=2000):
    """Point estimates plus 95% intervals by simulating (beta, rho) from their asymptotic normal distribution."""
    params = model.betas.flatten()   # constant, x..., rho
    k = len(X_NAMES)
    point = lag_impacts(params[1:1 + k], params[-1], eig)
    sims = rng.multivariate_normal(params, model.vm, size=draws)
    sims = sims[np.abs(sims[:, -1]) < 1]
    s = np.array([lag_impacts(d[1:1 + k], d[-1], eig) for d in sims])   # draws x 3 x k
    lo_, hi_ = np.percentile(s, 2.5, axis=0), np.percentile(s, 97.5, axis=0)
    return point, lo_, hi_


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


section("spatial lag model (ML)")
t0 = time.perf_counter()
lag = spreg.ML_Lag(y, X, w=w_queen, name_y="ln_price", name_x=X_NAMES, name_w="queen", spat_impacts="full")
print(f"fitted in {time.perf_counter() - t0:.2f} s")
print(lag.summary)

eig_q = np.linalg.eigvals(w_queen.full()[0])
(dir_, ind_, tot_), lo_i, hi_i = lag_impact_ci(lag, eig_q)
section("spatial lag impacts with 95% simulation intervals (2,000 draws)")
print(f"rho = {lag.rho:.3f}; spatial multiplier 1/(1-rho) = {1 / (1 - lag.rho):.2f}")
for j, n in enumerate(X_NAMES):
    print(f"{n:<11} direct {dir_[j]:+.4f} [{lo_i[0, j]:+.4f}, {hi_i[0, j]:+.4f}]  "
          f"indirect {ind_[j]:+.4f} [{lo_i[1, j]:+.4f}, {hi_i[1, j]:+.4f}]  "
          f"total {tot_[j]:+.4f} [{lo_i[2, j]:+.4f}, {hi_i[2, j]:+.4f}]")
j4 = X_NAMES.index("level4")
print(f"level4: +10 percentage points everywhere -> total {100 * (np.exp(10 * tot_[j4]) - 1):+.1f}% on prices "
      f"(direct {100 * (np.exp(10 * dir_[j4]) - 1):+.1f}%); OLS said {100 * (np.exp(10 * ols.betas[1, 0]) - 1):+.1f}%")

section("spatial error model (ML)")
err = spreg.ML_Error(y, X, w=w_queen, name_y="ln_price", name_x=X_NAMES, name_w="queen")
print(err.summary)

section("spatial Durbin model (ML lag with WX), for comparison")
sdm = spreg.ML_Lag(y, X, w=w_queen, slx_lags=1, name_y="ln_price", name_x=X_NAMES, name_w="queen",
                   spat_impacts="full")
print(sdm.summary)

section("robustness: GM estimators that allow heteroskedasticity")
gm_err = spreg.GM_Error_Het(y, X, w=w_queen, name_y="ln_price", name_x=X_NAMES)
gm_lag = spreg.GM_Lag(y, X, w=w_queen, robust="white", name_y="ln_price", name_x=X_NAMES)
for n, j in zip(X_NAMES, range(1, 5)):
    print(f"{n:<11} GM error (het) {gm_err.betas[j, 0]:+.4f} (z {gm_err.z_stat[j][0]:+.1f})   "
          f"GM lag (White) {gm_lag.betas[j, 0]:+.4f} (z {gm_lag.z_stat[j][0]:+.1f})")
print(f"GM error lambda {gm_err.betas[-1, 0]:.3f} | GM lag rho {np.ravel(gm_lag.rho)[0]:.3f}")

# ---------------------------------------------------------------- Step 4: compare the models
section("model comparison (queen weights)")
mi_lag = resid_moran(lag.u, w_queen)
mi_err = resid_moran(err.e_filtered, w_queen)
mi_sdm = resid_moran(sdm.u, w_queen)
rows = [("OLS", ols.logll, ols.aic, ols.schwarz, len(X_NAMES) + 1, mi_ols.I, mi_ols.p_sim),
        ("Spatial lag (ML)", lag.logll, lag.aic, lag.schwarz, len(X_NAMES) + 2, *mi_lag),
        ("Spatial error (ML)", err.logll, err.aic, err.schwarz, len(X_NAMES) + 2, *mi_err),
        ("Spatial Durbin (ML)", sdm.logll, sdm.aic, sdm.schwarz, 2 * len(X_NAMES) + 2, *mi_sdm)]
comp = pd.DataFrame(rows, columns=["model", "logL", "AIC", "SC", "params", "resid_I", "resid_p"])
print(comp.round(3).to_string(index=False))
from scipy import stats
for name, ll, df in [("lag vs OLS", lag.logll, 1), ("error vs OLS", err.logll, 1)]:
    lr = 2 * (ll - ols.logll)
    print(f"LR {name}: {lr:.1f} (df {df}, p = {stats.chi2.sf(lr, df):.1e})")
lr = 2 * (sdm.logll - lag.logll)
print(f"LR Durbin vs lag: {lr:.1f} (df 4, p = {stats.chi2.sf(lr, 4):.1e})")
lr = 2 * (sdm.logll - err.logll)
print(f"LR Durbin vs error (common factor): {lr:.1f} (df 4, p = {stats.chi2.sf(lr, 4):.1e})")

# coefficient comparison: % change in price for a one-standard-deviation increase in each variable
sd = gdf[X_NAMES].std().to_numpy()


def pct(b):
    return 100 * (np.exp(b * sd) - 1)


ols_ci = (ols.betas[1:, 0] - 1.96 * ols.std_err[1:], ols.betas[1:, 0] + 1.96 * ols.std_err[1:])
err_se = np.sqrt(np.diag(err.vm))[1:1 + len(X_NAMES)]
err_ci = (err.betas[1:1 + len(X_NAMES), 0] - 1.96 * err_se, err.betas[1:1 + len(X_NAMES), 0] + 1.96 * err_se)
series = [("OLS", ols.betas[1:, 0], ols_ci, "#0072B2"),
          ("Spatial error", err.betas[1:1 + len(X_NAMES), 0], err_ci, "#E69F00"),
          ("Spatial lag: direct", dir_, (lo_i[0], hi_i[0]), "#009E73"),
          ("Spatial lag: total", tot_, (lo_i[2], hi_i[2]), "#D55E00")]
section("% change in median price for a +1 SD change in each variable")
print("SD: " + ", ".join(f"{n} {s:.2f}" for n, s in zip(X_NAMES, sd)))
for name, b, _, _ in series:
    print(f"{name:<20} " + "  ".join(f"{n} {v:+.1f}%" for n, v in zip(X_NAMES, pct(b))))

fig, ax = plt.subplots(figsize=(8, 4.4))
ypos = np.arange(len(X_NAMES))
for i, (name, b, (l, h), col) in enumerate(series):
    off = (i - 1.5) * 0.18
    ax.errorbar(pct(b), ypos + off, xerr=[pct(b) - pct(l), pct(h) - pct(b)], fmt="o", color=col, ms=5,
                capsize=2.5, lw=1.4, label=name, markeredgecolor="#333333", markeredgewidth=0.4)
ax.axvline(0, color="#666666", lw=0.8)
ax.set_yticks(ypos, [f"{LABELS[n]}\n(1 SD = {s:.1f})" for n, s in zip(X_NAMES, sd)])
ax.invert_yaxis()
ax.set_xlabel("Change in median house price for a 1 SD increase (%), with 95% interval")
ax.set_title("Coefficients by model: OLS, spatial error, and spatial lag impacts")
ax.legend(loc="lower right", fontsize=8, frameon=True)
ax.grid(axis="x", color="#dddddd", lw=0.6)
fig.tight_layout(); fig.savefig(OUT / "03-coefficient-comparison.png"); plt.close(fig)

# residual maps for the three models
gdf["e_lag"] = lag.u.flatten()
gdf["e_err"] = np.asarray(err.e_filtered).flatten()
lim = max(np.abs(gdf[c]).max() for c in ["e_ols", "e_lag", "e_err"])
fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.4))
for ax, col, title, (mi, p) in zip(axes, ["e_ols", "e_lag", "e_err"],
                                   ["OLS", "Spatial lag", "Spatial error (filtered)"],
                                   [(mi_ols.I, mi_ols.p_sim), mi_lag, mi_err]):
    gdf.plot(col, ax=ax, cmap="RdBu_r", norm=TwoSlopeNorm(0, -lim, lim), edgecolor="white", linewidth=0.15)
    outline(ax)
    ax.set_title(f"{title}\nresidual Moran's I = {mi:.3f} (p = {p:.4f})", fontsize=9)
sm = plt.cm.ScalarMappable(cmap="RdBu_r", norm=TwoSlopeNorm(0, -lim, lim))
fig.colorbar(sm, ax=axes, shrink=0.7, label="Residual (log points)")
fig.savefig(OUT / "04-residual-maps-by-model.png"); plt.close(fig)

# direct and indirect impacts
fig, ax = plt.subplots(figsize=(7.5, 3.8))
d10, i10 = pct(dir_), pct(tot_) - pct(dir_)
ax.barh(ypos, d10, color="#0072B2", label="Direct (own MSOA, incl. feedback)")
ax.barh(ypos, i10, left=d10, color="#E69F00", label="Indirect (spillover via neighbours)")
ax.errorbar(pct(tot_), ypos, xerr=[pct(tot_) - pct(lo_i[2]), pct(hi_i[2]) - pct(tot_)], fmt="none",
            ecolor="#222222", capsize=3, lw=1.1, label="95% interval of the total")
ax.axvline(0, color="#666666", lw=0.8)
ax.set_yticks(ypos, [LABELS[n] for n in X_NAMES]); ax.invert_yaxis()
ax.set_xlabel("Change in median house price for a 1 SD increase in every MSOA (%)")
ax.set_title(f"Spatial lag model impacts (rho = {lag.rho:.2f}, queen weights)")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.17), ncol=3, fontsize=8, frameon=False)
ax.grid(axis="x", color="#dddddd", lw=0.6)
fig.tight_layout(); fig.savefig(OUT / "05-lag-impacts.png"); plt.close(fig)

# ---------------------------------------------------------------- Step 5: sensitivity to the weights
section("sensitivity to the spatial weights")
cents = gdf.copy(); cents["geometry"] = gdf.representative_point()
weights = {"queen": w_queen}
weights["rook"] = libpysal.weights.Rook.from_dataframe(gdf, use_index=False)
for k in (4, 6, 8, 12):
    weights[f"knn{k}"] = KNN.from_dataframe(cents, k=k, use_index=False)
sens = []
for name, w in weights.items():
    w.transform = "r"
    o = spreg.OLS(y, X, w=w, spat_diag=True, moran=True)
    with contextlib.redirect_stdout(io.StringIO()):   # spreg prints the class name of each ML fit
        lg = spreg.ML_Lag(y, X, w=w)
        er = spreg.ML_Error(y, X, w=w)
    eig = np.linalg.eigvals(w.full()[0])
    d, ind, tot = lag_impacts(lg.betas[1:1 + len(X_NAMES), 0], lg.rho, eig)
    sens.append({"weights": name, "mean_nbrs": w.mean_neighbors, "resid_I": o.moran_res[0],
                 "rLM_lag": o.rlm_lag[0], "rLM_err": o.rlm_error[0], "rho": lg.rho, "lambda": er.lam,
                 "AIC_lag": lg.aic, "AIC_err": er.aic, "lag_total_level4": tot[0], "lag_direct_level4": d[0],
                 "err_level4": er.betas[1, 0], "err_no_car": er.betas[3, 0], "lag_total_no_car": tot[2]})
sens = pd.DataFrame(sens)
print(sens.round(4).to_string(index=False))

fig, axes = plt.subplots(1, 3, figsize=(12, 3.6))
xs = np.arange(len(sens))
axes[0].plot(xs, sens.rho, "o-", color="#D55E00", label="rho (lag)")
axes[0].plot(xs, sens["lambda"], "s-", color="#0072B2", label="lambda (error)")
axes[0].set_ylabel("Spatial parameter"); axes[0].set_title("Spatial parameters"); axes[0].legend(frameon=False)
axes[1].plot(xs, sens.AIC_lag, "o-", color="#D55E00", label="spatial lag")
axes[1].plot(xs, sens.AIC_err, "s-", color="#0072B2", label="spatial error")
axes[1].axhline(ols.aic, color="#777777", ls="--", lw=1, label="OLS")
axes[1].set_ylabel("AIC (lower is better)"); axes[1].set_title("Model fit"); axes[1].legend(frameon=False)
for col, lab, mk, c in [("lag_total_level4", "lag: total impact", "o", "#D55E00"),
                        ("lag_direct_level4", "lag: direct impact", "^", "#E69F00"),
                        ("err_level4", "error: coefficient", "s", "#0072B2")]:
    axes[2].plot(xs, 100 * (np.exp(10 * sens[col]) - 1), mk + "-", color=c, label=lab)
axes[2].set_ylabel("% price change per +10 pp Level 4+"); axes[2].set_title("Level 4+ qualifications effect")
axes[2].legend(frameon=False, fontsize=8)
for ax in axes:
    ax.set_xticks(xs, sens.weights); ax.grid(color="#e5e5e5", lw=0.6)
fig.tight_layout(); fig.savefig(OUT / "06-weights-sensitivity.png"); plt.close(fig)

metrics = {"n": len(gdf), "ols_r2": ols.r2, "ols_aic": ols.aic, "lag_aic": lag.aic, "err_aic": err.aic,
           "sdm_aic": sdm.aic, "rho": lag.rho, "lambda": err.lam, "resid_moran_ols": mi_ols.I}
(OUT / "metrics.json").write_text(json.dumps({k: round(float(v), 4) for k, v in metrics.items()}, indent=2))
print("\nwrote figures and output/metrics.json")