import gradio as gr
from folium import Map
import numpy as np
from ast import literal_eval
import pandas as pd

from gradio_folium import Folium
import folium
from huggingface_hub import InferenceClient
from geopy.geocoders import Nominatim

from examples import (
    description_sf,
    output_example_sf,
    description_loire,
    output_example_loire,
    df_examples
)

geolocator = Nominatim(user_agent="HF-trip-planner")

def get_coordinates(address):
    location = geolocator.geocode(address)
    if location:
        return (location.latitude, location.longitude)
    else:
        return None

repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
llm_client = InferenceClient(model=repo_id, timeout=180)


def generate_key_points(text):
    prompt = f"""             
    Please generate a set of key geographical points for the following description: {text}, as a json list of less than 10 dictionnaries with the following keys: 'name', 'description'. Precise the full location in the 'name' if there is a possible ambiguity.
    Generally try to minimze the distance between locations. Always think of the transportation means that you want to use, and the timing: morning, afternoon, where to sleep.
    Only generate a 'Thought:' and a 'Key points:' sections, nothing else.

    For instance:
    Description: {description_sf}
    Thought: {output_example_sf}

    Description: {description_loire}
    Thought: {output_example_loire}

    Now begin. You can make the descriptions a bit more verbose than in the examples.

    Description: {text}
    Thought: 
    """
    return llm_client.text_generation(prompt, max_new_tokens=2000)


def parse_llm_output(output):
    rationale = "Thought: " + output.split("Key points:")[0]
    key_points = output.split("Key points:")[1]
    output = key_points.replace("    ", "")
    parsed_output = literal_eval(output)
    dataframe = pd.DataFrame.from_dict(parsed_output)
    return dataframe, rationale


def get_coordinates_row(row):
    coords = get_coordinates(row["name"])
    if coords is not None:
        row["lat"], row["lon"] = coords
    return row


def create_map_from_markers(dataframe):
    dataframe = dataframe.apply(get_coordinates_row, axis=1)
    f_map = Map(
        location=[dataframe["lat"].mean(), dataframe["lon"].mean()],
        zoom_start=5,
        tiles="CartoDB Voyager",
    )
    for _, row in dataframe.iterrows():
        if np.isnan(row["lat"]) or np.isnan(row["lon"]):
            continue
        marker = folium.CircleMarker(
            location=[row["lat"], row["lon"]],
            radius=10,
            popup=folium.Popup(
                f"<h4>{row['name']}</h4><p>{row['description']}</p>", max_width=450
            ),
            fill=True,
            fill_color="blue",
            fill_opacity=0.6,
            color="blue",
            weight=1,
        )
        marker.add_to(f_map),

    bounds = [[row["lat"], row["lon"]] for _, row in dataframe.iterrows()]
    f_map.fit_bounds(bounds, padding=(100, 100))
    return f_map


def run_display(text):
    output = generate_key_points(text)
    dataframe, rationale = parse_llm_output(output)
    map = create_map_from_markers(dataframe)
    return map, rationale


df_examples = pd.DataFrame.from_dict(
    [
        {"description": description_loire, "output": output_example_loire},
        {"description": description_aligned, "output": output_example_aligned},
        {"description": description_chinatown, "output": output_example_chinatown},
        {"description": description_taiwan, "output": output_example_taiwan},
    ]
)


def select_example(df, data: gr.SelectData):
    row = df.iloc[data.index[0], :]
    dataframe, rationale = parse_llm_output(row["output"])
    return row["description"], create_map_from_markers(dataframe), rationale


with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue=gr.themes.colors.yellow,
        secondary_hue=gr.themes.colors.blue,
    )
) as demo:
    gr.Markdown("# 🗺️ LLM trip planner (based on Mixtral)")
    text = gr.Textbox(
        label="Describe your trip here:",
        value=description_sf,
    )
    button = gr.Button()
    gr.Markdown("### LLM Output 👇\n_Click the map to see information about the places._")

    # Get initial map and rationale
    example_dataframe, example_rationale = parse_llm_output(output_example_sf)
    display_rationale = gr.Markdown(example_rationale)
    starting_map = create_map_from_markers(example_dataframe)
    map = Folium(value=starting_map, height=700, label="Chosen locations")
    button.click(run_display, inputs=[text], outputs=[map, display_rationale])

    gr.Markdown("### Other examples")
    clickable_examples = gr.DataFrame(value=df_examples, height=200)
    clickable_examples.select(
        select_example, clickable_examples, outputs=[text, map, display_rationale]
    )

if __name__ == "__main__":
    demo.launch()