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=1024,
    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):
    print('titi')
    map_html, fig = show_map(lat, lon, address)
    df = get_info_meteo(lat, lon)
    avg_temp = df.loc[:, 'avg_hourly_soil_temperature_0cm_C'].mean()
    avg_humidity = df.loc[:, 'avg_hourly_soil_moisture_0_to_1cm_m³'].mean()

    print('toto')
    return map_html, fig, avg_temp, avg_humidity


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, meteo, 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}
    - humidité moyenne : {humidite}
    - météo : {meteo}

    {message}

    Donne-moi ensuite les solutions de prévention et de traitement pour chacun d'eux indépendamment du tableau.
    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