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

def animate_deaths(df, country_name):
    df = df[df['Country'] == country_name].copy()
    df['Date'] = pd.to_datetime('2020-01-01') + pd.to_timedelta(df.index, unit='D')

    # Prepare the plot
    fig, ax = plt.subplots(figsize=(14, 6))

    line1, = ax.plot([], [], label='Daily Deaths', color='red')
    line2, = ax.plot([], [], label='7-Day Average', color='black')

    ax.set_xlim(df['Date'].min(), df['Date'].max())
    ymax = df['Daily new confirmed deaths due to COVID-19'].max() * 1.1
    ax.set_ylim(0, ymax)

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

    def update(frame):
        line1.set_data(df['Date'][:frame], df['Daily new confirmed deaths due to COVID-19'][:frame])
        line2.set_data(df['Date'][:frame], df['Daily new confirmed deaths due to COVID-19 (rolling 7-day average, right-aligned)'][:frame])
        return line1, line2

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

    # Generate HTML with the gif
    html_path = '../plots/animated_deaths_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 Deaths Animation – {country_name}</title>
</head>
<body>
    <h1>COVID-19 Daily Deaths – {country_name}</h1>
    <p>Animacja pokazuje dzienne zgony z powodu COVID-19 oraz 7-dniową średnią kroczącą.</p>
    <img src="uk_daily_deaths.gif" alt="Animated COVID-19 Deaths Chart">
</body>
</html>""")
    print(f"✔ GIF saved as: {gif_path}")
    print(f"✔ HTML generated: {html_path}")

if __name__ == '__main__':
    uk_df = pd.read_csv('../data/united_kingdom_covid.csv')
    animate_deaths(uk_df, 'United Kingdom')
