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

# Load data
df = pd.read_csv("../data/united_kingdom_covid.csv")
df = df[df["Year"] >= 0].copy()  # filter out invalid days
start_date = pd.to_datetime('2020-01-01')
df['Date'] = start_date + pd.to_timedelta(df['Year'], unit='D')

# Fill missing values with zeros
df['Cases'] = df['Daily new confirmed cases due to COVID-19 (rolling 7-day average, right-aligned)'].fillna(0)
df['Deaths'] = df['Daily new confirmed deaths due to COVID-19 (rolling 7-day average, right-aligned)'].fillna(0)

# Ensure output directory exists
os.makedirs('plots', exist_ok=True)

# Set up figure and axes
fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlim(df['Date'].min(), df['Date'].max())
ax.set_ylim(0, max(df['Cases'].max(), df['Deaths'].max()) * 1.1)

line1, = ax.plot([], [], label='Cases', color='blue', lw=2)
line2, = ax.plot([], [], label='Deaths', color='red', lw=2)

ax.set_title('7-Day Rolling Average of COVID-19 Cases and Deaths')
ax.set_xlabel('Date')
ax.set_ylabel('Count')
ax.legend()

# Format x-axis
ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2))
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
plt.xticks(rotation=45)

# Animation function
def animate(i):
    line1.set_data(df['Date'][:i], df['Cases'][:i])
    line2.set_data(df['Date'][:i], df['Deaths'][:i])
    return line1, line2

# Create and save animation
ani = FuncAnimation(fig, animate, frames=len(df), interval=30, blit=True)
ani.save("../plots/rolling_avg_animation1.gif", writer='pillow')
plt.close()
