import gradio as gr
import plotly.graph_objects as go
import json
import requests
import os
from PIL import Image
import hopsworks
import joblib
import pandas as pd
import numpy as np

API_KEY = os.environ['API_KEY_TOMTOM']
print("API KEY", API_KEY)

# Log into hopsworks
project = hopsworks.login()
fs = project.get_feature_store()
mr = project.get_model_registry()

model = mr.get_model("sthlm_incidents_model", version=4)
model_dir = model.download()
model = joblib.load(model_dir + "/sthlm_model.pkl")
print("Model downloaded")

api_params_incidents = {
    'base_url': 'api.tomtom.com',
    'API_KEY': API_KEY,
    'min_lon': 18.00,
    'max_lon': 18.16,
    'min_lat': 59.25,
    'max_lat': 59.40,
    'version_number': 5,
    'category_filter': '0%2C1%2C2%2C3%2C4%2C5%2C6%2C7%2C8%2C9%2C10%2C11%2C14',
    'time_validity_filter': 'present',
    'language': 'en-GB',
    'fields': '%7Bincidents%7Btype%2Cgeometry%7Bcoordinates%7D%2Cproperties%7Bid%2CmagnitudeOfDelay%2Cevents%7Bdescription%2Ccode%2CiconCategory%7D%2CstartTime%2CendTime%7D%7D%7D'
}

def get_incident_details(params):
    url = f"https://{params['base_url']}/traffic/services/{params['version_number']}/incidentDetails?bbox={params['min_lon']}%2C{params['min_lat']}%2C{params['max_lon']}%2C{params['max_lat']}&fields={params['fields']}&language={params['language']}&categoryFilter={params['category_filter']}&timeValidityFilter={params['time_validity_filter']}&key={params['API_KEY']}"
    return json.loads(requests.get(url).text)

# Function that handles an incident
def handle_incident(incident):
    hour_value = incident['properties']['startTime'][11:13]

    # Create a dictionary with the incident details
    dict_incident = {
        'id': incident['properties']['id'],
        'magnitudeOfDelay': incident['properties']['magnitudeOfDelay'],
        'startTime': incident['properties']['startTime'],
        'endTime': incident['properties']['endTime'],
        'type': incident['type'],
        'code': incident['properties']['events'][0]['code'],
        'iconCategory': incident['properties']['events'][0]['iconCategory'],
        'description': incident['properties']['events'][0]['description'],
        'month': incident['properties']['startTime'][5:7],
        'hour': hour_value,
    }

    # Get the coordinates of the incident, if statement is for if the incident only has one coordinate
    if any(isinstance(j, list) for j in incident['geometry']['coordinates']):
        dict_incident['longitude'] = incident['geometry']['coordinates'][0][0]
        dict_incident['latitude'] = incident['geometry']['coordinates'][0][1]
    else:
        dict_incident['longitude'] = incident['geometry']['coordinates'][0]
        dict_incident['latitude'] = incident['geometry']['coordinates'][1]

    return dict_incident

def filter_map():
    latitude_list = []
    longitude_list = []
    index_list = []
    # Get the data from the json
    incidents = get_incident_details(api_params_incidents)['incidents']

    for (_, i) in enumerate(incidents):
        # Add the coordinates to the list
        if any(isinstance(j, list) for j in i['geometry']['coordinates']):
            latitude_list.append(i['geometry']['coordinates'][0][1])
            longitude_list.append(i['geometry']['coordinates'][0][0])
        else:
            latitude_list.append(i['geometry']['coordinates'][1])
            longitude_list.append(i['geometry']['coordinates'][0])

        if type(i['properties']['events']) == list:
            description = i['properties']['events'][0]['description']
        else:
            description = i['properties']['events']['description']

        starting_date = i['properties']['startTime'].split('T')[0]
        starting_time = i['properties']['startTime'].split('T')[1].split('+')[0]
        starting_string = f"started at {starting_time} on {starting_date}"

        # Check if it has an end time
        if i['properties']['endTime'] is not None:
            ending_date = i['properties']['endTime'].split('T')[0]
            ending_time = i['properties']['endTime'].split('T')[1].split('+')[0]
            ending_string = f"will end at {ending_time} on {ending_date}<br><br>"
        else:
            # Use the model
            row = handle_incident(i)

            # Create a df from the row
            df_row = pd.DataFrame(row, index=[row['id']])
            df_row = df_row.drop(["startTime", "endTime", "type", "id", "description"], axis=1)

            # change the order to code	hour	iconCategory	latitude	longitude	magnitudeOfDelay	month	duration
            df_row = df_row[['code', 'hour', 'iconCategory', 'latitude', 'longitude', 'magnitudeOfDelay', 'month']]

            # make the features lower case
            df_row.columns = df_row.columns.str.lower()
            df_row.columns = df_row.columns.str.replace(' ', '_')

            # Get the prediction
            prediction = model.predict(df_row)
            prediction = prediction[0]
            print(prediction)

            # add the prediction to the starting time
            ending_time = pd.to_datetime(starting_time) + pd.Timedelta(seconds=int(min(prediction, 100000)))
            ending_date = str(ending_time).split(' ')[0]
            ending_time = str(ending_time).split(' ')[1].split('+')[0]
            ending_string = f"<b>Predicted ending time</b>: {ending_time} on {ending_date}<br><br>"
        index_list.append((starting_string, ending_string, description))

    fig = go.Figure(go.Scattermapbox(
            customdata=index_list,
            lat=latitude_list,
            lon=longitude_list,
            mode='markers',
            marker=go.scattermapbox.Marker(
                size=6
            ),
            hoverinfo="text",
            hovertemplate="""
                <b>Incident in Stockholm started at %{customdata[0]}<br>
                %{customdata[1]}
                <b>Additional Information</b>: %{customdata[2]}<br>
            """
        ))

    fig.update_layout(
        mapbox_style="open-street-map",
        hoverlabel=dict(
            bgcolor="rgba(255, 255, 255, 0)",
            font_size=16,
            font_family="Arial"
        ),
        hovermode='closest',
        mapbox=dict(
            bearing=0,
            center=go.layout.mapbox.Center(
                lat=59.32,
                lon=18.08
            ),
            pitch=0,
            zoom=9
        ),
    )
    return fig

with gr.Blocks() as demo:
    with gr.Column():
        btn = gr.Button(value="Update Filter")
        map = gr.Plot()
    demo.load(filter_map, outputs=map)
    btn.click(filter_map, outputs=map)

demo.launch()