import os
from cerebras.cloud.sdk import Cerebras
import gradio as gr
import folium
from geopy.geocoders import Nominatim
from IPython.display import HTML
import plotly.express as px

# Initialize Cerebras client
client = Cerebras(api_key="csk-vy9xw6mmk38ytjwyn34cty6e99f5j3c3m28xdphrv6nxj363")

def generate_country_info(country_name):
    messages = [
        {
            "role": "system",
            "content": "You have wide and specialised knowledge about each country of the world. Write comprehensive general information about any country entered in the prompt. like details of independence if any, population, most popular cities, tropical information, cultural information, GDP, the capital and numbers of states or provinces, and tourist attractions"
        },
        {
            "role": "user",
            "content": country_name
        }
    ]
    stream = client.chat.completions.create(
        messages=messages,
        model="llama-3.3-70b",
        stream=True,
        max_completion_tokens=4096,
        temperature=0.54,
        top_p=1
    )
    response = ""
    for chunk in stream:
        response += chunk.choices[0].delta.content or ""
    return response

def generate_choropleth_map(country_name):
    data = {
        "Country": [country_name],
        "Values": [100]
    }
    fig = px.choropleth(
        data,
        locations="Country",
        locationmode="country names",
        color="Values",
        color_continuous_scale="Inferno",
        title=f"Country Map Highlighting {country_name}"
    )
    return fig

def generate_country_map(country_name):
    geolocator = Nominatim(user_agent="geoapi")
    location = geolocator.geocode(country_name)
    if location:
        latitude, longitude = location.latitude, location.longitude
        country_map = folium.Map(location=[latitude, longitude], zoom_start=6)
        marker = folium.Marker([latitude, longitude], popup=country_name)
        marker.add_to(country_map)
        return country_map._repr_html_()
    else:
        return "Location not found. Please try again."

def explore_country(country_name):
    country_info = generate_country_info(country_name)
    choropleth_map = generate_choropleth_map(country_name)
    country_map = generate_country_map(country_name)
    return choropleth_map, country_map, country_info

with gr.Blocks() as app:
    gr.Markdown("# Countries Xplorer\nExplore detailed information about any country of the world.\nEnter the name of a country to view its location on maps and learn more about it.")

    with gr.Row():
        country_input = gr.Textbox(label="Enter the Country You Want to Explore", placeholder="E.g., Nigeria")

    with gr.Row():
        with gr.Column():
            choropleth_plot = gr.Plot(label="Choropleth Map")
            country_map_html = gr.HTML(label="Country Map")
        with gr.Column():
            country_info_box = gr.Markdown(label="Country Information")

    def app_interface(country_name):
        choropleth_map, country_map, country_info = explore_country(country_name)
        return (choropleth_map, country_map, country_info)

    submit_button = gr.Button("Explore")
    submit_button.click(
        app_interface,
        inputs=country_input,
        outputs=[choropleth_plot, country_map_html, country_info_box]
    )

app.launch()