import pandas as pd
import plotly.graph_objects as go
import os

# --- 1. CARREGAR E PREPARAR OS DADOS ---

# Carregar o mesmo arquivo, pois ele contém múltiplos indicadores
df_raw = pd.read_csv('../data/API_FP.CPI.TOTL.ZG_DS2_en_csv_v2_2729.csv', skiprows=4)

# Transformar os dados do formato largo para o formato longo
years = [str(year) for year in range(2010, 2025)]
df_melt = df_raw.melt(
    id_vars=['Country Name', 'Country Code', 'Indicator Name', 'Indicator Code'],
    value_vars=years,
    var_name='Year',
    value_name='Value'
)

# Filtrar para os países de interesse
countries = ['Rwanda', 'South Africa', 'Poland']
df_filtered = df_melt[df_melt['Country Name'].isin(countries)]

# --- ALTERAÇÃO AQUI: Filtrar para o indicador de inflação ---
df_filtered = df_filtered[df_filtered['Indicator Name'].str.contains("Inflation")]

# Limpar os dados: remover valores nulos e converter o ano para inteiro
df_inflation = df_filtered.dropna(subset=['Value']).copy()
df_inflation['Year'] = df_inflation['Year'].astype(int)

print("Dados de inflação carregados e preparados com sucesso.")

# --- 2. CRIAR O GRÁFICO ESTÁTICO ---

# Cores para os países
color_map = {
    'Rwanda': 'blue',
    'Poland': 'red',
    'South Africa': 'green'
}

# Criar a figura
fig = go.Figure()

# Adicionar uma linha (trace) para cada país
for country in countries:
    country_data = df_inflation[df_inflation['Country Name'] == country]
    fig.add_trace(go.Scatter(
        x=country_data['Year'],
        y=country_data['Value'],
        mode='lines+markers',
        name=country,
        line=dict(color=color_map[country], width=2.5),
        marker=dict(size=5),
        hovertemplate=(
            "<b>" + country + "</b><br>" +
            "Year: %{x}<br>" +
            "Inflation: %{y:.2f}%" + # <-- Texto do hover alterado
            "<extra></extra>"
        )
    ))

# --- 3. CONFIGURAR O LAYOUT E ADICIONAR ANOTAÇÃO ---

# Aplicar o layout, ajustando os títulos e eixos para inflação
fig.update_layout(
    title=dict(
        text="<b>Inflation Rate Comparison Over Time</b>", # <-- Título alterado
        x=0.5,
        font=dict(size=26, family="Arial", color="black")
    ),
    xaxis=dict(
        title=dict(text="<b>Year</b>", font=dict(size=20, family="Arial", color="black")),
        tickfont=dict(size=14, family="Arial", color="black"),
        range=[2010, df_inflation['Year'].max() + 1],
        showgrid=True,
        gridcolor="lightgrey",
        showline=True,
        linecolor='black',
        linewidth=1
    ),
    yaxis=dict(
        title=dict(text="<b>Annual Inflation Rate (%)</b>", font=dict(size=20, family="Arial", color="black")), # <-- Título do eixo Y alterado
        tickfont=dict(size=14, family="Arial", color="black"),
        # Ajuste dinâmico do range do eixo Y, com uma pequena margem
        range=[df_inflation['Value'].min() - 2, 35], 
        showgrid=True,
        gridcolor="lightgrey",
        showline=True,
        linecolor='black',
        linewidth=1
    ),
    plot_bgcolor="white",
    margin=dict(l=120, r=50, t=90, b=80),
    legend=dict(
        font=dict(size=18, family="Arial", color="black"),
        yanchor="top",
        y=0.98,
        xanchor="right",
        x=0.98
    )
)

# Adicionar a linha vertical para a pandemia
fig.add_vline(
    x=2020, 
    line_width=2, 
    line_dash="dash", 
    line_color="black",
    annotation_text="Start of COVID-19",
    annotation_position="top left",
    annotation_font_size=12
)


# --- 4. EXIBIR E SALVAR O GRÁFICO ---

if __name__ == "__main__":
    output_dir = "../Plots/Economy"
    os.makedirs(output_dir, exist_ok=True) # Cria a pasta se não existir

    # --- Nomes dos arquivos de saída alterados ---
    output_path_html = os.path.join(output_dir, "inflation_comparison.html")
    output_path_png = os.path.join(output_dir, "inflation_comparison.png")

    fig.show()
    fig.write_html(output_path_html, include_plotlyjs="cdn")
    fig.write_image(output_path_png, width=1200, height=800)
    print(f"Gráfico estático de inflação salvo em: {output_path_html} e .png")