import os

import geopandas as gpd
import pandas as pd
import numpy as np
import shapely
import matplotlib.pyplot as plt
import geopy
import urllib
from shapely.geometry import shape
import contextily as ctx
from tqdm.auto import tqdm

with_plotly = True # Preferably to True, but in case plotly cannot be installed, this allows the notebook to work
if with_plotly:
    import plotly.express as px
    import plotly.graph_objects as go

crs =     'epsg:3857'
osm_crs=  'epsg:4326'

def download_if_nexist(url, filename):
    """
    If the (local) file <filename> does not exists, download it from <url> 

    Parameters
    ----------
    url: str
       url to fetch
    filename: str
       local file to save

    Returns
    -------

    None
    """
    if not os.path.isfile(filename):
        with urllib.request.urlopen(url) as response:
            with open(filename, "wb") as f:
                f.write(response.read())

def get_osm_polyg(name):
    """
    Request Nominatim (OpenStreetMap) for "name", and returns the polygon (or point) given in the (first) result
    
    Parameters
    ----------
    name: str
        
    """
    url = f"http://nominatim.openstreetmap.org/search.php?q={name}&format=json&limit=1&polygon_geojson=1"
    with urllib.request.urlopen(url) as response:
        json_res = json.loads(response.read())
        if len(json_res)==0:
            return None
        return shape(json_res[0]["geojson"])

zipcodes_normal = zipcodes[~zipcodes.is_special].drop("is_special", axis=1)

division_counts = pd.DataFrame(columns = ["#zipcodes", "#localities", "desc", "query"],
             data=[
 ["=1", "=1",   "Only one locality within the commune",         "(zipcode==1) & (locality == 1)"],
 ["=1", ">1",   "Zipcode share by all the localities",          "(zipcode==1) & (locality>1)"],
 [">1", "=zip", "A distinct zipcode for each locality",         "(zipcode >1) & (zipcode==locality)"],
 [">1", ">zip", "Some localities share zipcodes",               "(zipcode >1) & (zipcode<locality)"],
# [">1", "<zip", "(more zip than localities, does not happen)",  "(zipcode >1) & (zipcode>locality)"] 
])

zipcode_boundaries_filename = "data/zipcode_boundaries_shapefile_3812.zip"

# https://www.geo.be/catalog/details/9738c7c0-5255-11ea-8895-34e12d0f0423?l=fr
download_if_nexist("https://bgu.bpost.be/assets/9738c7c0-5255-11ea-8895-34e12d0f0423_x-shapefile_3812.zip",
                  zipcode_boundaries_filename)

# https://statbel.fgov.be/fr/open-data/secteurs-statistiques-2020

download_if_nexist("https://statbel.fgov.be/sites/default/files/files/opendata/Statistische%20sectoren/sh_statbel_statistical_sectors_31370_20200101.shp.zip",
                   "data/stat_sectors_2020.zip")
statistical_sectors = gpd.read_file("zip://data/stat_sectors_2020.zip/sh_statbel_statistical_sectors_20200101.shp")
statistical_sectors["CNIS5_2020"] = statistical_sectors["CNIS5_2020"].astype(str)

# Group (with "dissolve") sectors per NIS code
nis_boundaries = statistical_sectors[["CNIS5_2020", "T_MUN_FR", "T_MUN_NL", "geometry"]].dissolve(by="CNIS5_2020").reset_index()
nis_boundaries = nis_boundaries.rename({"CNIS5_2020": "niscode"}, axis=1)
nis_boundaries = nis_boundaries.to_crs(crs)

def plot_mismatch_zip(sjoin_mismatches, i):
    print(f"Mismatch record {i}:")
    display(sjoin_mismatches.iloc[[i]])

    zip1 = sjoin_mismatches.iloc[i].zipcode_left
    zip2 = sjoin_mismatches.iloc[i].zipcode_right

    print("From BPost:")
    display( zipcodes[ zipcodes.zipcode.isin([zip1, zip2])])

    try: 
        print("Know by OSM at the given point:")
        pt = zipcodes_centers.to_crs(osm_crs).loc[sjoin_mismatches.iloc[i].index_right].geometry
        geolocator = geopy.geocoders.Nominatim(user_agent="smalsresearch")
        print(geolocator.reverse(f"{pt.y}, {pt.x}").address)

        print("Know by OSM for the given locality:")
        loc = zipcodes_centers.loc[sjoin_mismatches.iloc[i].index_right].locality
        print(geolocator.geocode(f"{loc}, Belgium").address)
    except Exception as e:
        print("Cannot get OSM data:", e)

    query = f"zipcode in ('{zip1}', '{zip2}')"


    df_bnd = zipcodes_boundaries.assign(label=zipcodes_boundaries.zipcode).query(query).to_crs(osm_crs)
    df_pnt = zipcodes_centers.assign(label=zipcodes_centers.zipcode+" "+zipcodes_centers.locality).query(query).to_crs(osm_crs)

    df_pnt["mismatch"] = False
    df_pnt.loc[sjoin_mismatches.iloc[i].index_right, "mismatch"] = True

    if with_plotly:
        fig = px.choropleth_mapbox(df_bnd, 
                                   geojson=df_bnd.geometry, 
                                   locations=df_bnd.index, 
                                   hover_name="label",
                                   color=df_bnd['label'],
                                   center={"lat": df_bnd.bounds.mean()[["miny", "maxy"]].mean(), 
                                           "lon": df_bnd.bounds.mean()[["minx", "maxx"]].mean()},
                                   mapbox_style="open-street-map",
                                   #zoom=10,
                                   opacity=0.5,

                           )

        for i, bloc in df_pnt.groupby("mismatch"): 
            fig.add_trace(go.Scattermapbox(lat=bloc.geometry.y,
                                           lon=bloc.geometry.x,
                                           mode='markers+text', # "markers+text",
                                           text=bloc.label, 
                                           name="Zipcode centers" + (" (mismatch)" if i else ""),
                                           textposition="bottom center"
                                          )
                                    )
        fig.show()
    else: 
        df=pd.concat([df_bnd, df_pnt]).to_crs(crs)
        df["label"] = df["label"] + np.where(df.mismatch==True, "*", "")
        ax = df.plot("label", figsize=(10,15), legend=True)
        try:
            ctx.add_basemap(ax)
        except Exception as  e:
            print("Cannot add basemap: ", e)
        plt.show()
        display(df[["zipcode", "geometry","label"]])