import pandas as pd
import geopandas as gpd
import colorsys
import folium
from folium import plugins
from pyproj import CRS, Transformer
import logging

from folium.plugins import HeatMap
from matplotlib import pyplot as plt

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'
}

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(folium_map):

    # 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, folium_map)
    logger.info(f"Heat map time added to time map.")
    # interactive_map.save("test.html")

    add_signaled_speeds(folium_map)

    # Add bikes

    add_bike_heat_map_time(folium_map)

    # Pedestrian Part

    add_pedestrian_heat_map_time(folium_map)

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


def create_heat_map_toggle(folium_map):
    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, folium_map)

    add_bike_heat_toggle(folium_map)
    add_ped_heat_toggle(folium_map)
    # Add signald speeds data
    add_signaled_speeds(folium_map)

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


# Layer Adding Methods ================================================================================================
def add_bike_heat_map_time(folium_map):
    # Process heat map data
    bike_heat_view_data = get_view('bikeheat', 'latitude, longitude, year')
    bike_heat_df = gpd.GeoDataFrame(bike_heat_view_data, columns=['latitude', 'longitude', 'year'])

    assert not bike_heat_df.empty, f" Heat Dataframe is empty: {bike_heat_df.head(5)}"
    heat_data = [
        [[row['latitude'], row['longitude'], 0.1] for index, row in bike_heat_df[bike_heat_df['year'] == i].iterrows()]
        for
        i in range(2011, 2023)]
    logger.debug(f"First element of heat data: {heat_data[0]}")
    index = [2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022]
    AccidentType = "Bicycles: "
    index = [str(element) for element in index]
    index = [AccidentType + element for element in index]
    # plot heat map
    gradient = generate_hue_gradient(0.6, 5)
    hm = plugins.HeatMapWithTime(heat_data,
                                 auto_play=False,
                                 max_opacity=1,
                                 gradient=gradient,
                                 min_opacity=0.5,
                                 radius=9,
                                 use_local_extrema=False,
                                 blur=1,
                                 index=index,
                                 name="Accident Heatmap Bikes")
    hm.add_to(folium_map)


def add_pedestrian_heat_map_time(folium_map):
    # Process heat map data
    pedestrian_heat_view_data = get_view("pedestrianheat")
    heat_df = gpd.GeoDataFrame(pedestrian_heat_view_data, columns=['latitude', 'longitude', 'year'])

    assert not heat_df.empty, f" Heat Dataframe is empty: {heat_df.head(5)}"
    heat_data = [[[row['latitude'], row['longitude'], 0.5] for index, row in heat_df[heat_df['year'] == i].iterrows()]
                 for
                 i in range(2011, 2023)]
    logger.debug(f"First element of PED heat data: {heat_data[0]}")
    index = [2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022]
    AccidentType = "Pedestrians: "
    index = [str(element) for element in index]
    index = [AccidentType + element for element in index]
    # gradient =
    # plot heat map
    gradient = generate_hue_gradient(0.2, 5)
    hm = plugins.HeatMapWithTime(heat_data,
                                 auto_play=False,
                                 max_opacity=1,
                                 gradient=gradient,
                                 min_opacity=0.5,
                                 radius=9,
                                 use_local_extrema=False,
                                 blur=1,
                                 index=index,
                                 name="Accident Heatmap Pedestrian")
    hm.add_to(folium_map)


def add_heat_map_time(heat_df, folium_map):
    heat_data = [[[row['latitude'], row['longitude'], 0.5] 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 ALL")
    hm.add_to(folium_map)


def add_signaled_speeds(folium_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(folium_map)


def add_heat_year_toggle(heat_gdf, folium_map, name="All"):
    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'{name} in {year}'
        )

        heatmap_layer.add_to(folium_map)


def add_bike_heat_toggle(folium_map):
    bike_heat_view_data = get_view('bikeheat', 'latitude, longitude, year')
    heat_gdf = gpd.GeoDataFrame(bike_heat_view_data, columns=['latitude', 'longitude', 'year'])
    add_heat_year_toggle(heat_gdf, folium_map, name="motorcycles")


def add_ped_heat_toggle(folium_map):
    pedestrian_heat_view_data = get_view("pedestrianheat")
    heat_gdf = gpd.GeoDataFrame(pedestrian_heat_view_data, columns=['latitude', 'longitude', 'year'])
    add_heat_year_toggle(heat_gdf, folium_map, name="pedestrians")


# Utilities ===========================================================================================================

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


def setup_views():
    drop_view("heat")
    create_heat_view()
    drop_view("bikeheat")
    create_bike_heat_view()
    drop_view("pedestrianheat")
    create_pedestrian_heat_view()


def generate_hue_gradient(hue, num_colors):
    if num_colors < 2:
        num_colors = 2
    gradient = {}
    for i in range(num_colors):
        lightness = 0.1 + 0.8 * (i / (num_colors - 1))
        saturation = 0.1 + 0.8 * (i / (num_colors - 1))
        rgb = colorsys.hls_to_rgb(hue, lightness, saturation)
        gradient[i / (num_colors - 1)] = '#{:02x}{:02x}{:02x}'.format(int(rgb[0] * 255), int(rgb[1] * 255),
                                                                      int(rgb[2] * 255))
    return gradient


def generate_contrasting_gradient(num_colors):
    cmap = plt.get_cmap('viridis')  # viridis is a map with contrasting colors
    gradient = {}
    for i in range(num_colors):
        rgba = cmap(i / (num_colors - 1))
        gradient[i / (num_colors - 1)] = '#{:02x}{:02x}{:02x}'.format(int(rgba[0] * 255), int(rgba[1] * 255),
                                                                      int(rgba[2] * 255))
    return gradient


def add_miv_count_station_locations():
    get_data_mic_sql = """
    SELECT
        zsid, ekoord, nkoord,
        AVG(anzfahrzeuge) AS average_count
    FROM
        mivcount
    GROUP BY
        zsid, ekoord, nkoord
    """
    remote_db = RemoteDB()
    miv_result = remote_db.execute_query(get_data_mic_sql)
    miv_df = pd.DataFrame(miv_result)
    miv_df[['lon', 'lat']] = miv_df.apply(lambda row: convert_to_wgs84(row['ekoord'], row['nkoord']), axis=1)
    miv_df['average_count'] = miv_df['average_count'].apply(lambda x: round(float(x)))
    count_stations_layer = folium.FeatureGroup(name='Count-stations cars', show=False)
    for index, row in miv_df.iterrows():
        folium.Marker(location=[row['lat'], row['lon']], popup="avg. " + str(row['average_count']), show=False).add_to(count_stations_layer)
    count_stations_layer.add_to(toggle_map)
    remote_db.close()
def add_fb_count_station_locations():
    get_data_mic_sql = """
    SELECT DISTINCT  
        ost,
        nord,
        AVG(velo_total) as average_velo_count,
        AVG(fuss_total) as average_fuss_count
    FROM fbcount_copy
    GROUP BY ost,nord;
    """
    remote_db = RemoteDB()
    FB_result = remote_db.execute_query(get_data_mic_sql)
    FB_df = pd.DataFrame(FB_result)
    FB_df[['ost', 'nord']] = FB_df.apply(lambda row: convert_to_wgs84(row['ost'], row['nord']), axis=1)
    FB_df['average_velo_count'] = FB_df['average_velo_count'].apply(lambda x: round(float(x)))
    FB_df['average_velo_count'] = FB_df['average_velo_count'].astype(str)
    FB_df['average_fuss_count'] = FB_df['average_fuss_count'].apply(lambda x: round(float(x)))
    FB_df['average_fuss_count'] = FB_df['average_fuss_count'].astype(str)
    count_stations_layer = folium.FeatureGroup(name='Count-stations pedestrians and bicycles', show=False)
    for index, row in FB_df.iterrows():
        folium.Marker(location=[row['nord'], row['ost']], popup="Bicycle and pedestrian count station", show=False).add_to(count_stations_layer)
    count_stations_layer.add_to(toggle_map)
    remote_db.close()
def convert_to_wgs84(lon, lat):
    swiss_crs = CRS.from_epsg(2056)
    wgs84_crs = CRS.from_epsg(4326)

    transformer = Transformer.from_crs(swiss_crs, wgs84_crs, always_xy=True)
    lon, lat = transformer.transform(lon, lat)

    return pd.Series({'lon': lon, 'lat': lat})


if __name__ == "__main__":
    time_map = folium.Map(
        location=zurich_coordinates,
        zoom_start=13,
        zoom_control=True,
        dragging=True,
        scrollWheelZoom=True,
        doubleClickZoom=False,
        tiles="cartodb positron"
    )

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

    add_miv_count_station_locations()
    add_fb_count_station_locations()
    #setup_views()

    create_heat_map_with_time(time_map)
    create_heat_map_toggle(toggle_map)

    ## Save Maps ============================================================================================
    save_map_as_html(toggle_map, "html/heat_map_toggle")
    save_map_as_html(time_map, "html/heat_map_time")

    ## Create Maps with fixed center=============================================================================
    time_map_fix = folium.Map(
        location=zurich_coordinates,
        zoom_start=13,
        zoom_control=True,
        dragging=False,
        scrollWheelZoom=True,
        doubleClickZoom=False,
        tiles="cartodb positron"
    )

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

    #setup_views()

    create_heat_map_with_time(time_map_fix)
    create_heat_map_toggle(toggle_map_fix)

    ## Save Maps ============================================================================================
    save_map_as_html(toggle_map_fix, "html/heat_map_toggle_fix")
    save_map_as_html(time_map_fix, "html/heat_map_time_fix")