import numpy as np
import libpysal as ps
from mgwr.gwr import GWR, MGWR
from mgwr.sel_bw import Sel_BW
import geopandas as gp
import matplotlib.pyplot as plt
import matplotlib as mpl

#Load Georgia dataset and generate plot of Georgia counties (figure 1)
georgia = gp.read_file(ps.examples.get_path('G_utm.shp'))
fig, ax = plt.subplots(figsize=(10,10))
georgia.plot(ax=ax, **{'edgecolor':'black', 'facecolor':'white'})
georgia.centroid.plot(ax=ax, c='black')

#Prepare Georgia dataset inputs
y = georgia['PctBach'].values.reshape((-1,1))
X = georgia[['PctFB', 'PctBlack', 'PctRural']].values
u = georgia['X']
v = georgia['Y']
coords = list(zip(u,v))

np.random.seed(908)
sample = np.random.choice(range(159), 10)
mask = np.ones_like(y,dtype=bool).flatten()
mask[sample] = False

cal_coords = coords[mask]
cal_y = y[mask]
cal_X = X[mask]

pred_coords = coords[~mask]
pred_y = y[~mask]
pred_X = X[~mask]

#Calibrate GWR model

gwr_selector = Sel_BW(cal_coords, cal_y, cal_X)
gwr_bw = gwr_selector.search(bw_min=2)
print(gwr_bw)
model = GWR(cal_coords, cal_y, cal_X, gwr_bw)
gwr_results = model.fit()

scale = gwr_results.scale
residuals = gwr_results.resid_response

pred_results = model.predict(pred_coords, pred_X, scale, residuals)

pred_results.predictions

pred_y

np.corrcoef(pred_results.predictions.flatten(), pred_y.flatten())[0][1]