import math
import pandas as pd
import numpy as np
import folio
from utils import map_vals
from matplotlib import pyplot as plt

# NOTE
# This only needed to be ran once to generate the maps
# Maps are saved in the figures folder and loaded as html

service_data_pd = pd.read_csv("data/311-2016-2018.csv")
service_data_pd["Incident Zip"] = service_data_pd["Incident Zip"].astype("string")
service_data_pd["BBL"] = service_data_pd["BBL"].astype("string")
service_data_raw = pl.DataFrame(service_data_pd)
# service_data_raw = pl.read_csv("data/311-2016-2018.csv", null_values="", infer_schema_length=0)
# service_data_raw = service_data_raw.with_columns(
#     pl.col("Latitude").cast(pl.Float64),
#     pl.col("Longitude").cast(pl.Float64)
# )
# Clear some ram
del service_data_pd
gc.collect()

weather_data_raw = pd.read_csv("data/weather_NY_2010_2018Nov.csv")

def get_map_1():
    fig, weather_map = map_vals(
        weather_data_raw.loc[weather_data_raw["Year"] >= 2016],
        cols=["Latitude", "Longitude"],
        label_cols=["StationName"],
        sample_size=1000,
        color='red',
        radius=3,
        weight=4
    )
    fig, combined_map = map_vals(
        service_data_raw,
        cols=["Latitude", "Longitude"],
        color="blue", submap=weather_map,
        sample_size=1000,
        weight=2,
        radius=1
    )

    fig.save("figures/map1.html")

    return fig


def get_map_2():
    fig, service_map = map_vals(
        service_data_raw,
        cols=["Latitude", "Longitude"],
        color="blue",
        weight=2,
        radius=1,
        start_loc=[40.7128, -74.0060],
        sample_size=1000,
        zoom_start=10
    )
    fig, weather_map = map_vals(
        weather_data_raw.loc[weather_data_raw["Year"] >= 2016],
        cols=["Latitude", "Longitude"],
        submap=service_map,
        label_cols=["StationName"],
        color='red',
        radius=5,
        weight=2,
        sample_size=1000,
    )

    fig.save("figures/map2.html")

    return fig


def get_bounded_map():
    # Get prerecorded coords for the mins/max to maximize speed here
    # In notebook this is recorded via code
    lat_min = 40.49804421521046
    lat_max = 40.91294056699566
    long_min = -74.25521082506387
    long_max = -73.70038354802529

    fig = folium.Figure(height=500, width=750)
    service_bounds_map = folium.Map(
        location=[40.7128, -74.0060],
        zoom_start=10,
        tiles='cartodbpositron',
        zoom_control=False,
        scrollWheelZoom=False,
        dragging=False
    )

    kw = {
        "color": "#F1807E",
        "line_cap": "round",
        "fill": True,
        "fill_color": "blue",
        "weight": 3,
        "popup": "Service Data Coverage Zone",
    }

    folium.Rectangle(
        bounds=[[lat_min, long_min], [lat_max, long_max]],
        line_join="round",
        dash_array="5 5",
        **kw,
    ).add_to(service_bounds_map)

    fig.add_child(service_bounds_map)

    fig.save("figures/bounded_map.html")

    return fig

   
def get_final_map():
    lat_min = 40.49804421521046
    lat_max = 40.91294056699566
    long_min = -74.25521082506387
    long_max = -73.70038354802529

    mincon_lat = weather_data_raw["Latitude"] >= lat_min
    maxcon_lat = weather_data_raw["Latitude"] <= lat_max
    mincon_long = weather_data_raw["Longitude"] >= long_min
    maxcon_long = weather_data_raw["Longitude"] <= long_max

    service_bounds_map = folium.Map(
        location=[40.7128, -74.0060],
        zoom_start=10,
        tiles='cartodbpositron',
        zoom_control=False,
        scrollWheelZoom=False,
        dragging=False
    )

    kw = {
        "color": "#F1807E",
        "line_cap": "round",
        "fill": True,
        "fill_color": "blue",
        "weight": 3,
        "popup": "Service Data Coverage Zone",
    }

    folium.Rectangle(
        bounds=[[lat_min, long_min], [lat_max, long_max]],
        line_join="round",
        dash_array="5 5",
        **kw,
    ).add_to(service_bounds_map)

    wd_localized = weather_data_raw.loc[mincon_lat & maxcon_lat & mincon_long & maxcon_long]
    fig, wd_local_map = map_vals(
        wd_localized,
        submap=service_bounds_map,
        label_cols=["StationName"],
        color='red',
        radius=5, 
        weight=2,
        sample_size=1000,
    )

    fig.save("figures/final_map.html")
    
    return fig


def build_maps():
    get_map_1()
    get_map_2()
    get_bounded_map()
    get_final_map()

build_maps()