import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import os

def load_data(path):
    return pd.read_csv(path)

def animate_cases(df, country_name):
    # Filter for the specified country
    df = df[df['Country'] == country_name].copy()

    # Create a date index
    df['Date'] = pd.to_datetime('2020-01-01') + pd.to_timedelta(df.index, unit='D')

    # Set up the figure
    fig, ax = plt.subplots(figsize=(14, 6))

    # Initialize lines
    line1, = ax.plot([], [], label='Daily New Cases', color='red')
    line2, = ax.plot([], [], label='7-Day Rolling Average', color='blue')

    # Axis limits
    ax.set_xlim(df['Date'].min(), df['Date'].max())
    ymax = df['Daily new confirmed cases of COVID-19'].max() * 1.1
    ax.set_ylim(0, ymax)

    # Titles and grid
    ax.set_title(f'Daily COVID-19 Cases – {country_name}')
    ax.set_xlabel('Date')
    ax.set_ylabel('Number of Cases')
    ax.grid(True)
    ax.legend()

    # Update function for animation
    def update(frame):
        line1.set_data(df['Date'][:frame], df['Daily new confirmed cases of COVID-19'][:frame])
        line2.set_data(df['Date'][:frame], df["Daily new confirmed cases due to COVID-19 (rolling 7-day average, right-aligned)"][:frame])
        return line1, line2

    # Create animation
    os.makedirs('plots', exist_ok=True)
    gif_path = '../plots/uk_daily_cases.gif'
    anim = FuncAnimation(fig, update, frames=len(df), interval=30, blit=True)
    anim.save(gif_path, writer='pillow', fps=25)
    plt.close()

    # Generate HTML file
    html_path = '../plots/animated_plot.html'
    with open(html_path, 'w', encoding='utf-8') as f:
        f.write(f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>COVID-19 Cases Animation – {country_name}</title>
</head>
<body>
    <h1>COVID-19 Cases Animation – {country_name}</h1>
    <p>Wykres przedstawia dzienne przypadki COVID-19 oraz 7-dniową średnią kroczącą.</p>
    <img src="uk_daily_cases.gif" alt="Animated COVID-19 Cases Chart">
</body>
</html>""")
    print(f"✔ GIF saved as: {gif_path}")
    print(f"✔ HTML generated: {html_path}")

if __name__ == '__main__':
    uk_df = load_data('../data/united_kingdom_covid.csv')
    animate_cases(uk_df, 'United Kingdom')
