import pandas as pd
import geopandas as gpd
import os
import folium
from folium import plugins
import logging

from folium.plugins import HeatMap

from db_connector import RemoteDB
import shapely
from shapely import wkb
import json

## MUST IMPORT otherwise contains the functions used in db interaction
from db_utils import *

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('map.py')
stream_handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
stream_handler.setFormatter(formatter)
logger.addHandler(stream_handler)


accidents_filepath = "../src/datasets/integrated/Accidents.geojson"
signaled_speeds_filepath = "../src/datasets/integrated/signaled_speeds.geojson.geojson"

# Map centered around zurich
zurich_coordinates = [47.368650,  	8.539183]
fixed_map_zurich_original_coords = folium.Map(
    location=zurich_coordinates,
    zoom_start=13,
    zoom_control=False,
    dragging=False,
    scrollWheelZoom=False,
    doubleClickZoom=False
)

gradient = {
    0.1: 'blue',
    0.3: 'cyan',
    0.5: 'lime',
    0.7: 'yellow',
    0.9: 'red'
}

interactive_map = folium.Map(
    location=zurich_coordinates,
    zoom_start=13,
    zoom_control=True,
    dragging=True,
    scrollWheelZoom=True,
    doubleClickZoom=False,
    tiles="cartodb positron"
)

interactive_map_toggle = folium.Map(
    location=zurich_coordinates,
    zoom_start=13,
    zoom_control=True,
    dragging=True,
    scrollWheelZoom=True,
    doubleClickZoom=False,
    tiles="cartodb positron"
)


speedLimits = ["T0","T20","T30","T50","T60","T80","T100"]
color_dict = {
    "T0": "red",
    "T20": "orange",
    "T30": "green",
    "T50": "yellow",
    "T60": "purple",
    "T80": "pink",
    "T100": "gray"
}


# Create Maps =========================================================================================================
def create_heat_map_with_time():


    # Process heat map data
    heat_view_data = get_view("heat")
    heat_df = gpd.GeoDataFrame(heat_view_data, columns=['latitude', 'longitude', 'year'])

    assert not heat_df.empty, f" Heat Dataframe is empty: {heat_df.head(5)}"
    add_heat_map_time(heat_df, interactive_map)
    #interactive_map.save("test.html")

    add_signaled_speeds(interactive_map)

    folium.LayerControl(collapsed=True).add_to(interactive_map)

def create_heat_map_toggle():

    heat_view_data = get_view("heat")
    heat_gdf = gpd.GeoDataFrame(heat_view_data, columns=['latitude', 'longitude', 'year'])

    assert not heat_gdf.empty, f" Heat Dataframe is empty: {heat_gdf.head(5)}"

    add_heat_year_toggle(heat_gdf, interactive_map_toggle)

    # Add signald speeds data
    add_signaled_speeds(interactive_map_toggle)

    folium.LayerControl(collapsed=True).add_to(interactive_map_toggle)


# Layer Adding Methods ================================================================================================
def add_heat_map_time(heat_df, map):
    heat_data = [[[row['latitude'], row['longitude'], 0.1] for index, row in heat_df[heat_df['year'] == i].iterrows()] for
                 i in range(2011, 2023)]
    index = [2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022]
    # create heat map
    logger.debug(f"First element of heat data: {heat_data[0]}")
    hm = plugins.HeatMapWithTime(heat_data,
                                 auto_play=False,
                                 max_opacity=0.8,
                                 gradient=gradient,
                                 min_opacity=0.3,
                                 radius=9,
                                 use_local_extrema=False,
                                 blur=1,
                                 index=index,
                                 name="Accident Heatmap")
    hm.add_to(map)


def add_signaled_speeds(map):
    # Add signald speeds data
    rows = """
            temporegime_technical as tempo, 
            wkb_geometry
    """
    sig_speeds_data = get_view("signaled_speeds", rows)
    sig_speed_df = pd.DataFrame(sig_speeds_data, columns=['tempo', 'wkb_geometry'])
    sig_speed_df['geometry'] = sig_speed_df['wkb_geometry'].apply(lambda x: wkb.loads(x, hex=True))
    logger.debug(f"{sig_speed_df.head()}")
    sig_speed_gdf = gpd.GeoDataFrame(sig_speed_df, geometry="geometry")
    for speedlimit in speedLimits:
        signal_speed = sig_speed_gdf[sig_speed_gdf["tempo"].str.contains(speedlimit, case=False)]
        geometries = json.loads(json.dumps(shapely.geometry.mapping(signal_speed['geometry'].unary_union)))

        folium.GeoJson(
            data=geometries,
            name=f'Signaled Speed {speedlimit}',
            color=color_dict[speedlimit],
            show=False,
            line_cap="butt",
        ).add_to(map)


def add_heat_year_toggle(heat_gdf, map):
    index = [2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022]
    # plot heat map
    for year in index:
        year_data = heat_gdf[heat_gdf['year'] == year]

        heatmap_layer = HeatMap(
            data=year_data[['latitude', 'longitude']],
            radius=8,
            gradient=gradient,
            min_opacity=0.5,
            max_opacity=0.8,
            blur=10,
            show=False,
            name=f'Accidents in {year}'
        )

        heatmap_layer.add_to(map)


def save_map_as_html(map, name):
    map.save(f"{name}.html")
    logger.info(f"Succesfully saved map {name}.")


if __name__ == "__main__":
    drop_view("heat")
    create_heat_view()
    create_heat_map_with_time()
    save_map_as_html(interactive_map, "heat_map_time")
    #create_heat_map_toggle()
    #save_map_as_html(interactive_map_toggle, "heat_map_toggle")