import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from shapely.geometry import Point

boundaries = pd.read_csv('data/Local_Authority_Districts_December_2021_UK_BGC_2022_-6374591305423020363.csv')
geometry = [Point(xy) for xy in zip(boundaries['LONG'], boundaries['LAT'])]
gdf_boundaries = gpd.GeoDataFrame(boundaries, geometry=geometry, crs='EPSG:4326')
gdf_boundaries = gdf_boundaries.rename(columns={'LAD21CD': 'AreaCode'})

covid = pd.read_csv('data/england-cases.csv')
covid['Date'] = pd.to_datetime(covid['Date'])

dates = sorted(covid['Date'].unique())

fig, ax = plt.subplots(figsize=(12, 10))
plt.axis('off')

def update(day):
    ax.clear()
    plt.axis('off')
    current_date = dates[day]
    covid_today = covid[covid['Date'] == current_date]

    merged = gdf_boundaries.merge(covid_today, on='AreaCode', how='left').fillna(0)

    merged.plot(ax=ax, color='lightgray', markersize=20, alpha=0.5)

    sizes = merged['DailyCases'] * 15
    sizes[sizes < 10] = 10 

    merged.plot(ax=ax, color='red', markersize=sizes, alpha=0.7)

    ax.set_title(f'Daily COVID-19 Cases by Region\nDate: {current_date.strftime("%Y-%m-%d")}', fontsize=18)
    ax.set_xlim(-6, 2)
    ax.set_ylim(49.5, 59)

    for size, label, y in zip([10, 150, 300], ['0–5 cases', '10 cases', '20 cases'], [50, 51, 52]):
        ax.scatter(-5.5, y, s=size, color='red', alpha=0.7)
        ax.text(-5.3, y, label, verticalalignment='center', fontsize=12)

    ax.text(-5.7, 53, 'Legend: Bubble size = Daily Cases', fontsize=14, weight='bold')

ani = FuncAnimation(fig, update, frames=len(dates), repeat=False)
ani.save('plots/covid_cases_bubble_animation.mp4', writer='ffmpeg', fps=5)
plt.close()
