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

    anim = FuncAnimation(fig, update, frames=len(df), interval=30, blit=True)

    # Save as GIF
    os.makedirs('plots', exist_ok=True)
    anim.save('../plots/uk_daily_deaths1.gif', writer='pillow', fps=25)
    plt.close()

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

