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

# Load data
df = pd.read_csv("../data/england_vaccinations.csv")
df["date"] = pd.to_datetime(df["date"])

# Create figure
fig, ax = plt.subplots(figsize=(12, 6))

# Initialize empty lines
line1, = ax.plot([], [], label="Total vaccinations", color="blue")
line2, = ax.plot([], [], label="At least 1 dose", color="green")
line3, = ax.plot([], [], label="Fully vaccinated", color="orange")
line4, = ax.plot([], [], label="Boosters", color="purple")

# Configure plot
ax.set_xlim(df["date"].min(), df["date"].max())
ax.set_ylim(0, df[["total_vaccinations", "people_vaccinated", "people_fully_vaccinated", "total_boosters"]].max().max() * 1.05)
ax.set_title("COVID-19 Vaccination Progress in England")
ax.set_xlabel("Date")
ax.set_ylabel("Number of people")
ax.grid(True)
ax.legend()

# Update function
def update(frame):
    line1.set_data(df["date"][:frame], df["total_vaccinations"][:frame])
    line2.set_data(df["date"][:frame], df["people_vaccinated"][:frame])
    line3.set_data(df["date"][:frame], df["people_fully_vaccinated"][:frame])
    line4.set_data(df["date"][:frame], df["total_boosters"][:frame])
    return line1, line2, line3, line4

# Create animation
anim = FuncAnimation(fig, update, frames=len(df), interval=50, blit=True)

# Save to HTML
html_path = "../plots/vaccination_progress1.html"
os.makedirs("../plots", exist_ok=True)
with open(html_path, "w") as f:
    f.write(anim.to_html5_video())

plt.close()
print(f"Saved HTML animation to: {html_path}")

