import pandas as pd
import folium
from folium import IFrame
import plotly.graph_objects as go
from APIs.geolocation import get_geolocation, get_risques
from APIs.meteo import get_info_meteo
import os
from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding
from llama_index.core import Settings
from llama_index.core.query_engine import RetrieverQueryEngine

load_dotenv()

Settings.llm = MistralAI(
    max_tokens=2048,
    api_key=os.environ.get('MISTRAL_API_KEY')
)

Settings.embed_model = MistralAIEmbedding(
    model_name="mistral-embed",
    api_key=os.environ.get('MISTRAL_API_KEY')
)

documents = SimpleDirectoryReader("documents").load_data()
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine(similarity_top_k=15)


def on_init(lat, lon, address):
    map_html, fig = show_map(lat, lon, address)
    df = get_info_meteo(lat, lon)
    avg_temp = df.loc[:, 'daily_temperature_2m_max_C'].mean()
    avg_humidity = df.loc[:, 'avg_hourly_soil_moisture_0_to_1cm_m³'].mean()
    # avg_snowfall = df.loc[:, 'daily_precipitation_sum_mm'].mean()
    avg_cloud_cover = df.loc[:, 'med_cloud_cover_%'].mean()
    avg_precipitations = df.loc[:, 'daily_rain_sum_mm'].mean()

    return map_html, fig, avg_temp, avg_humidity, avg_cloud_cover, avg_precipitations


def print_conversation(chat_history):
    for question, response in chat_history:
        print('Question: ', end='\n')
        print(question, end='\n')
        print('Response: ', end='\n')
        print(response, end='\n')


def respond(elevage, temperature, humidite, chutes_de_neige, couverture_nuageuse, precipitations, message, chat_history):
    PROMPT = f"""
    Tu es un chatbot qui répond en français, qui doit aider à déterminer les risques parasitaires d'élevage en fonction des conditions météorologiques.
    Voici les conditions: 
    - élevage: {elevage}
    - température moyenne : {temperature} °C
    - humidité moyenne : {humidite} m^3

    Fais moi un tableau avec chaque parasite et/ou tiques susceptible de toucher mon élevage dans ces conditions météo, avec les conditions de prolifération, la vitesse de propagation, les troubles causés par ces parasites
    Essaie d'écrire des données numériques précises dans le tableau
    Si un parasite et/ou tique ne rentre pas dans les conditions, ne le met pas dans le tableau.
    Donne-moi ensuite les solutions de prévention et de traitement pour chacun d'eux indépendamment du tableau.

    Réponds également au message suivant:
    {message}
    
    Tu dois impérativement répondre en français
    """

    reponse = query_engine.query(PROMPT)
    chat_history.append((message, str(reponse)))

    print_conversation(chat_history)
    return '', chat_history


def get_and_plot_meteo(lat, lon):
    df = get_info_meteo(lat, lon)
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df['date'], y=df['daily_temperature_2m_max_C'],
                  name="Température °C", mode="lines+markers"))
    fig.add_trace(go.Scatter(x=df['date'], y=df['avg_hourly_relative_humidity_2m_%'],
                  name="Humidité moyenne en %", mode="lines+markers", yaxis="y2"))
    fig.update_layout(
        title="Température maximale et humidité moyenne journalière",
        xaxis_title="Date",
        yaxis_title="Température (°C)",
        yaxis2=dict(
            title="Humidity (%)",
            overlaying="y",
            side="right"
        )
    )
    return fig


def show_map(lat, lon, address):
    if address:
        lat_tmp, lon_tmp, code_insee = get_geolocation(address, None, None)
        risques = get_risques(code_insee=code_insee)
        if lat_tmp or lon_tmp:
            lat, lon = lat_tmp, lon_tmp
        else:
            return "Adress not found. Please enter a valid address", ""
    if lat and lon:
        lat_tmp, lon_tmp, code_insee = get_geolocation(None, lat, lon)
        risques = get_risques(code_insee=code_insee)
        # Create a map centered at the input coordinates
        location_map = folium.Map(location=[lat, lon], zoom_start=14)
        folium.Marker([lat, lon], popup=risques).add_to(location_map)
        map_html = location_map._repr_html_()
        fig = get_and_plot_meteo(lat, lon)
    return map_html, fig