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

import asyncio
from gradio_folium import Folium
import folium
from huggingface_hub import InferenceClient
from geopy.geocoders import Nominatim
from collections import OrderedDict
from geopy.adapters import AioHTTPAdapter
import nest_asyncio
nest_asyncio.apply()

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

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


class AsyncLRUCache:
    def __init__(self, maxsize=100):
        self.cache = OrderedDict()
        self.maxsize = maxsize

    async def get(self, key):
        if key not in self.cache:
            return None
        self.cache.move_to_end(key)
        return self.cache[key]

    async def set(self, key, value):
        if key in self.cache:
            self.cache.move_to_end(key)
        self.cache[key] = value
        if len(self.cache) > self.maxsize:
            self.cache.popitem(last=False)


# Instantiate the cache
cache = AsyncLRUCache(maxsize=500)

async def geocode_address(address):
    # Check if the result is in cache
    cached_location = await cache.get(address)
    if cached_location:
        return cached_location

    # If not in cache, perform the geolocation request
    async with Nominatim(
        user_agent="HF-trip-planner",
        adapter_factory=AioHTTPAdapter,
    ) as geolocator:
        location = await geolocator.geocode(address, timeout=10)
        if location:
            # Save the result in cache for future use
            await cache.set(address, location)
        return location
    
async def ageocode_addresses(addresses):
    tasks = [geocode_address(address) for address in addresses]
    locations = await asyncio.gather(*tasks)
    return locations

def geocode_addresses(addresses):
    loop = asyncio.get_event_loop()
    result = loop.run_until_complete(ageocode_addresses(addresses))
    return result


def create_map_from_markers(dataframe):
    locations = geocode_addresses(dataframe["name"])
    dataframe["lat"] = [location.latitude if location else None for location in locations]
    dataframe["lon"] = [location.longitude if location else None for location in locations]

    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


def select_example(df, data: gr.SelectData):
    row = df.iloc[data.index[0], :]
    dataframe, rationale = parse_llm_output(row["output"])
    map = create_map_from_markers(dataframe)
    return row["description"], map, 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=600, 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()