#!/usr/bin/env python
# coding: utf-8

# # Electrify_Clusters

# ### All the necessary Python imports

# In[2]:


get_ipython().run_line_magic('matplotlib', 'inline')
from pathlib import Path

from electrificationplanner import clustering, electrify

import sys
sys.path.insert(0, '../minigrid-optimiser')
from mgo import mgo


# ### If clusters not already created, create it now

# In[ ]:


folder_input = Path('/home/chris/Documents/GIS')
ghs_in = folder_input / 'GHS-POP/GHS_POP_GPW42015_GLOBE_R2015A_54009_250_v1_0.tif'
clip_boundary = folder_input / 'gadm_uganda.gpkg'
clip_boundary_layer = 'gadm36_UGA_0'
grid_in = folder_input / 'uganda_grid.gpkg'
clusters_file = folder_input / 'clusters.gpkg'

print('Clipping raster...', end='', flush=True)
clipped, affine, crs = clustering.clip_raster(ghs_in, clip_boundary, clip_boundary_layer)
print('\t\tDoneCreating clusters...', end='', flush=True)
clusters = clustering.create_clusters(clipped, affine, crs)
print('\t\tDoneFiltering and merging...', end='', flush=True)
clusters = clustering.filter_merge_clusters(clusters)
print('\tDone\nGetting population...', end='', flush=True)
clusters = clustering.cluster_pops(clusters, ghs_in)
print('\t\tDoneGetting grid dists...', end='', flush=True)
clusters = clustering.cluster_grid_distance(clusters, grid_in, clipped[0].shape, affine)
print(f'\t\tDoneSaving to {str(clusters_out)}...', end='', flush=True)
clustering.save_clusters(clusters, clusters_file)
print('\t\tDone')


# ### Enter all input data here

# In[3]:


folder_input = Path('/home/chris/Documents/GIS')

clusters_file = folder_input / 'clusters.gpkg' # must be polygons with attributes pop_sum, area_m2, grid_dist
clusters_out = folder_input / 'clusters_out1.gpkg'
network_out = folder_input / 'network_out1.gpkg'

grid_dist_connected = 1000  # clusters within this distance of grid are considered connected

minimum_pop = 500 # exclude any population below this

# off-grid costs
demand_per_person_kwh_month = 6 # 6kWh/month = MTF Tier 2
demand_per_person_kw_peak = demand_per_person_kwh_month / (4*30)  # 130 4hours/day*30days/month based on MTF numbers, should use a real demand curve
mg_gen_cost_per_kw = 4000
mg_cost_per_m2 = 2

# grid costs
cost_wire_per_m = 50
grid_cost_per_m2 = 2


# ### Read in the clusters file, convert to desired CRS (ostensibly better for distances) and convert to points, filter on population along the way

# In[8]:


clusters = electrify.load_clusters(clusters_file, grid_dist_connected=grid_dist_connected,
                                   minimum_pop=minimum_pop)


# ### We then take all the clusters and calculate the optimum network that connects them all together. The ML model returns T_x and T_y containing the start and end points of each new arc created

# In[9]:


network, nodes = electrify.create_network(clusters)


# ### Then we're ready to calculate the optimum grid extension.
# This is done by expanding out from each already connected node, finding the optimum connection of nearby nodes. This is then compared to the off-grid cost and if better, these nodes are marked as connected. Then the loop continues until no new connections are found.

# In[12]:


network, nodes = electrify.run_model(network,
                           nodes,
                           demand_per_person_kw_peak=demand_per_person_kw_peak,
                           mg_gen_cost_per_kw=mg_gen_cost_per_kw,
                           mg_cost_per_m2=mg_cost_per_m2,
                           cost_wire_per_m=cost_wire_per_m,
                           grid_cost_per_m2=grid_cost_per_m2)


# ### And then do a join to get the results back into a polygon shapefile

# In[14]:


network_gdf, clusters_joined = electrify.spatialise(network, nodes, clusters)


# In[ ]:


clusters_joined.to_file(str(clusters_out), driver='GPKG')
network_gdf.to_file(str(network_out), driver='GPKG')


# ### And display some summary results

# In[16]:


new_conns = clusters_joined.loc[clusters_joined['conn_end'] == 1].loc[clusters_joined['conn_start'] == 0]
og = clusters_joined.loc[clusters_joined['conn_end'] == 0]
arcs = network_gdf.loc[network_gdf['existing'] == 0].loc[network_gdf['enabled'] == 1]
cost = og['og_cost'].sum() + cost_wire_per_m * arcs['length'].sum() + grid_cost_per_m2 * new_conns['area_m2'].sum()

total_modelled_pop = clusters['pop_sum'].sum()
urban_elec_rate = 0.6
currently_electrified = clusters.loc[clusters['conn_start'] == 1, 'pop_sum'].sum() * urban_elec_rate
new_conn_pop = clusters_joined.loc[clusters_joined['conn_start'] == 0].loc[clusters_joined['conn_end'] == 1, 'pop_sum'].sum()
off_grid_pop = clusters_joined.loc[clusters_joined['conn_start'] == 0].loc[clusters_joined['conn_end'] == 0, 'pop_sum'].sum()

print(f'{len(new_conns)} connected')
print(f'{len(og)} off-grid')
print(f'{len(arcs)} lines')
print()
print(f'Cost ${cost:,.0f}')
print()
print(f'Modelled pop: {total_modelled_pop:,.0f}')
print(f'Currently electrified: {currently_electrified:,.0f}')
print(f'New connections: {new_conn_pop:,.0f}')
print(f'Off-grid connections {off_grid_pop:,.0f}')