import gradio as gr
import pandas as pd
from datetime import datetime
import matplotlib.pyplot as plt
from pywaffle import Waffle
import math

# Load the life expectancy data from the World Bank
# This is a direct link to the CSV data for life expectancy at birth, total (years)
file = "life_expectancy_2023_world_back_data.csv"

# Read the CSV data, skipping the first 4 rows of metadata
try:
    df = pd.read_csv(file)

    life_expectancy_data = df[['Country','Life Expectancy at Birth']].copy()
    # life_expectancy_data.columns = ['Country', 'Life Expectancy']
    life_expectancy_data.dropna(inplace=True)
    countries = sorted(life_expectancy_data['Country'].unique())
except Exception as e:
    print(f"Error loading data: {e}")
    countries = ["Error loading country data"]
    life_expectancy_data = pd.DataFrame(columns=['Country', 'Life Expectancy'])


def create_life_calendar(name, birth_year, country):
    """
    Creates a graphical calendar of life, with each year represented
    as a row of 52 dots (weeks).
    """
    # --- Input Validation and Data Fetching ---
    if not all([name, birth_year, country]) or country.startswith("Error"):
        return None, "Please provide your name, birth year, and select a country."

    if life_expectancy_data.empty:
         return None, "Life expectancy data is unavailable. Cannot generate calendar."

    try:
        birth_year = int(birth_year)
        country_data = life_expectancy_data[life_expectancy_data['Country'] == country]
        if country_data.empty:
            return None, f"Sorry, life expectancy data for {country} is not available."
        life_expectancy = int(country_data['Life Expectancy at Birth'].iloc[0])

    except (ValueError, TypeError):
        return None, "Please enter a valid birth year."

    # --- Calculations ---
    now = datetime.now()
    current_year = now.year
    current_week = now.isocalendar()[1]
    age = current_year - birth_year

    # --- Plotting Setup ---
    fig, ax = plt.subplots(figsize=(10, life_expectancy / 10), dpi=120)
    
    # Define colors
    past_color = '#d9534f'  # Red
    future_color = '#5cb85c' # Green
    
    dot_size = 10
    weeks_in_year = 52

    # --- Main Loop to Plot Dots ---
    for year in range(1, life_expectancy + 1):
        for week in range(1, weeks_in_year + 1):
            # Determine color
            if year < age:
                color = past_color
            elif year == age:
                color = past_color if week < current_week else future_color
            else:
                color = future_color
            
            # Plot a single dot (week)
            # We plot from top to bottom, so we reverse the y-coordinate
            ax.scatter(week, life_expectancy - year, c=color, s=dot_size)

    # --- Formatting and Labels ---
    # Add year/age labels to the y-axis
    ax.set_yticks(range(0, life_expectancy, 5))
    ax.set_yticklabels([str(life_expectancy - i) for i in range(0, life_expectancy, 5)])
    ax.set_ylabel("Age", fontsize=12)
    
    # Configure x-axis for weeks
    ax.set_xticks([1, 13, 26, 39, 52])
    ax.set_xticklabels(["Week 1", "Week 13", "Week 26", "Week 39", "Week 52"])
    ax.set_xlabel("Week of the Year", fontsize=12)
    ax.xaxis.tick_top()
    ax.xaxis.set_label_position('top')

    # Remove the plot frame/spines for a cleaner look
    for spine in ['left', 'right', 'bottom']:
        ax.spines[spine].set_visible(False)

    # Set plot limits
    ax.set_xlim(0, weeks_in_year + 1)
    ax.set_ylim(-1, life_expectancy)

    # Create a custom legend
    legend_elements = [
        plt.Line2D([0], [0], marker='o', color='w', label='Past Week', markerfacecolor=past_color, markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='Future Week', markerfacecolor=future_color, markersize=10)
    ]
    ax.legend(handles=legend_elements, loc='center left', bbox_to_anchor=(1.02, 0.5), fontsize=12)
    
    fig.tight_layout()
    
    weeks_left = (life_expectancy - age) * 52 - current_week
    message = f"Hello {name}, based on a life expectancy of {life_expectancy} in {country}, you have approximately {weeks_left:,} weeks remaining."

    return fig, message


# Create the Gradio interface
with gr.Blocks(css=".center-text {text-align: center;}",theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Last Sunday
        Get reminded of how many Sundays remain :)

        This app shows you a visualization of how many Sunday are remaining in your life. It's created to remind oneself that time in life is limited and is not to be wasted.

        It's very easy to use. You simply enter your name and date of birth. Taking life expectancy as 80 years, it tells you how many weeks are left until you die.

        Inspired by the PARAS CHOPRA [Last Sunday](https://chromewebstore.google.com/detail/the-last-sunday-reminder/aiojhapcgfgmiacbbjfgedhlcchmpelh?hl=en) app.
        """
        ,elem_classes="center-text"
    )
    with gr.Row():
        name_input = gr.Textbox(label="Your Name")
        birth_year_input = gr.Number(label="Your Birth Year", minimum=1900, maximum=datetime.now().year)
        country_input = gr.Dropdown(choices=countries, label="Your Country")

    # output_text = gr.Textbox(label="Your Remaining Sundays")

    # with gr.Row():
    #     expired_input = gr.Number(label="Expired Sundays (Red Dots)", value=1560)
    #     remaining_input = gr.Number(label="Remaining Sundays (Green Dots)", value=2600)

    generate_button = gr.Button("Generate My Life Calendar", variant="primary")

    with gr.Column():
        output_plot = gr.Plot()
        output_text = gr.Label()

    generate_button.click(
        fn=create_life_calendar,
        inputs=[name_input, birth_year_input, country_input],
        outputs=[output_plot, output_text]
    )

demo.launch()