import pandas as pd
import matplotlib.pyplot as plt
import os
import numpy as np
from tkinter import Tk, Label, Listbox, Button, MULTIPLE, Frame, StringVar, messagebox

def load_data_and_plot():
    # Get the current directory where the script is located
    current_dir = os.path.dirname(os.path.abspath(__file__))

    # Load the CO2 emissions data
    emissions_file = os.path.join(current_dir, 'annual-co2-emissions-per-country.csv')
    emissions_data = pd.read_csv(emissions_file)

    # Load the population data
    population_file = os.path.join(current_dir, 'population.csv')
    population_data = pd.read_csv(population_file)

    # Create GUI for entity selection
    root = Tk()
    root.title("CO2 Emissions Per Capita Plotter")
    root.geometry("500x400")

    # Get all unique entities (countries)
    all_entities = sorted(list(set(emissions_data['Entity']) & set(population_data['Entity'])))

    # Frame for entity selection
    selection_frame = Frame(root)
    selection_frame.pack(pady=10, fill="both", expand=True)

    Label(selection_frame, text="Select countries to plot:", font=("Arial", 12)).pack(pady=5)

    # Create listbox for entity selection
    entity_listbox = Listbox(selection_frame, selectmode=MULTIPLE, height=15, width=40)
    entity_listbox.pack(fill="both", expand=True, padx=20, pady=5)

    # Add entities to listbox
    for entity in all_entities:
        entity_listbox.insert('end', entity)

    # Button frame
    button_frame = Frame(root)
    button_frame.pack(pady=10)

    # Function to plot selected entities
    def plot_selected_entities():
        # Get selected indices
        selected_indices = entity_listbox.curselection()

        if not selected_indices:
            messagebox.showwarning("Warning", "No countries selected. Please select at least one country.")
            return

        # Get selected entity names
        selected_entities = [entity_listbox.get(i) for i in selected_indices]

        # Colors for different entities
        colors = ['#004488', '#BB5566', '#DDAA33', '#009988', '#EE3377', '#44BB99', '#AA3377', '#882255']

        # Create plot
        plt.figure(figsize=(14, 7))

        # For legend
        plot_handles = []
        plot_labels = []

        # Min and max years for consistent x-axis
        min_year = 2023
        max_year = 1950

        # For y-axis scaling
        max_value = 0

        # Plot each selected entity with a different color
        for i, entity in enumerate(selected_entities):
            color_idx = i % len(colors)  # Cycle through colors if more entities than colors

            # Filter data for the entity
            entity_emissions = emissions_data[emissions_data['Entity'] == entity]
            entity_population = population_data[population_data['Entity'] == entity]

            # Filter for years from 1950 onwards
            entity_emissions_filtered = entity_emissions[entity_emissions['Year'] >= 1950]
            entity_population_filtered = entity_population[entity_population['Year'] >= 1950]

            # Merge the datasets on 'Year'
            merged_data = pd.merge(entity_emissions_filtered, entity_population_filtered, on='Year')

            if merged_data.empty:
                messagebox.showinfo("Info", f"No data found for {entity} after 1950")
                continue

            # Update min/max year for x-axis limits
            if merged_data['Year'].min() < min_year:
                min_year = merged_data['Year'].min()
            if merged_data['Year'].max() > max_year:
                max_year = merged_data['Year'].max()

            # Get column names
            emissions_col = [col for col in merged_data.columns if 'CO2' in col or 'emissions' in col.lower()][0]
            population_col = 'all years'  # Column name for population

            # Calculate per capita emissions (tons of CO2 per person)
            merged_data['CO2_per_capita'] = merged_data[emissions_col] / merged_data[population_col]

            # Update max value for y-axis scaling
            if merged_data['CO2_per_capita'].max() > max_value:
                max_value = merged_data['CO2_per_capita'].max()

            # Plot the data
            line, = plt.plot(merged_data['Year'], merged_data['CO2_per_capita'],
                          marker='o', linestyle='-', color=colors[color_idx], linewidth=2,
                          label=entity)

            plot_handles.append(line)
            plot_labels.append(entity)

            # Add annotation for peak value
            max_year_entity = merged_data.loc[merged_data['CO2_per_capita'].idxmax(), 'Year']
            max_value_entity = merged_data['CO2_per_capita'].max()

            plt.annotate(f'{entity} peak: {max_value_entity:.2f}',
                         xy=(max_year_entity, max_value_entity),
                         xytext=(max_year_entity+2, max_value_entity+0.3),
                         arrowprops=dict(arrowstyle='->', color=colors[color_idx]))

        # Add title and labels
        plt.title('CO2 Emissions Per Capita (1950-Present)', fontsize=16)
        plt.xlabel('Year', fontsize=12)
        plt.ylabel('CO2 Emissions (tons per person)', fontsize=12)
        plt.grid(True, alpha=0.3)

        # Set axis limits
        plt.xlim(min_year-1, max_year+1)
        plt.ylim(0, max_value * 1.15)  # Add 15% padding at the top

        # Add legend
        plt.legend(handles=plot_handles, labels=plot_labels, loc='best')

        # Improve formatting
        plt.tight_layout()

        # Save the figure
        save_path = os.path.join(current_dir, 'co2_per_capita_comparison.png')
        plt.savefig(save_path, dpi=300)

        # Show the plot
        plt.show()

        print(f"Plot created and saved as 'co2_per_capita_comparison.png'")

        # Close the GUI
        root.destroy()

    # Plot button
    Button(button_frame, text="Plot Selected Countries", command=plot_selected_entities,
           font=("Arial", 12), bg="#4CAF50", fg="white", padx=10).pack(pady=5)

    # Run the GUI
    root.mainloop()

if __name__ == "__main__":
    load_data_and_plot()