import numpy as np
import geopandas as gp
import multiprocessing as mp
import libpysal as ps
import sys
sys.path.append('/Users/Ziqi/Desktop/mgwr/')
from mgwr.gwr import GWR,MGWR
from mgwr.sel_bw import Sel_BW

#Load Berlin example
prenz = gp.read_file(ps.examples.get_path('prenzlauer.zip'))

b_y = np.log(prenz['price'].values.reshape((-1, 1)))
b_X = prenz[['review_sco','accommodat','bathrooms']].values 
b_X = (b_X - b_X.mean(axis=0)) / b_X.std(axis=0)
b_y = (b_y - b_y.mean(axis=0)) / b_y.std(axis=0)
u = prenz['X']
v = prenz['Y']
b_coords = list(zip(u, v))

#This might be needed to turn off the OpenMP multi-threading
%env OMP_NUM_THREADS = 1

%%time
gwr_selector = Sel_BW(b_coords, b_y, b_X)
gwr_bw = gwr_selector.search()
print(gwr_bw)
gwr_results = GWR(b_coords, b_y, b_X, gwr_bw).fit()

%%time
mgwr_selector = Sel_BW(b_coords, b_y, b_X, multi=True)
mgwr_bw = mgwr_selector.search()
print(mgwr_bw)
mgwr_results = MGWR(b_coords, b_y, b_X, selector=mgwr_selector).fit()

#Parrallelization is more favored when you your data are large and/or your machine have many many cores.
#mgwr has soft dependency of numba, please install numba if you need better performance (pip install numba).

n_proc = 2 #two processors
pool = mp.Pool(n_proc) 

%%time
gwr_selector = Sel_BW(b_coords, b_y, b_X)
gwr_bw = gwr_selector.search(pool=pool) #add pool to Sel_BW.search
print(gwr_bw)
gwr_results = GWR(b_coords, b_y, b_X, gwr_bw).fit(pool=pool) #add pool to GWR.fit

%%time
mgwr_selector = Sel_BW(b_coords, b_y, b_X, multi=True)
mgwr_bw = mgwr_selector.search(pool=pool) #add pool to Sel_BW.search
print(mgwr_bw)
mgwr_results = MGWR(b_coords, b_y, b_X, selector=mgwr_selector).fit(pool=pool) #add pool to MGWR.fit

pool.close() # Close the pool when you finish
pool.join()