from matplotlib import pyplot as plt
from urllib.request import urlopen
import pandas

import datetime
import pandas as pd
import io
import numpy as np
import matplotlib.ticker as ticker

URL = "https://api.coronavirus.data.gov.uk/v2/data?areaType=overview&metric=newAntibodyTestsByPublishDate&format=csv"
data = urlopen(URL).read()
df = pd.read_csv(io.BytesIO(data))
df = df.assign(date=pd.to_datetime(df["date"]))
def add_moving_average(df):
    df["moving_average"] = moving_average(df["newAntibodyTestsByPublishDate"])

def moving_average(data):
    window_size = 7
    return np.concatenate([[float('nan')] * 3, np.convolve(data, np.ones(window_size)/window_size, mode='valid'), [float('nan')] * 3])

def plot_cases(df, start, end):
    df_2021_h1 = df[df.date >= start][df.date < end].sort_values("date")
    add_moving_average(df_2021_h1)
    plt.bar(df_2021_h1["date"], df_2021_h1["newAntibodyTestsByPublishDate"], color="green", label="Daily")
    plt.plot(df_2021_h1["date"], df_2021_h1["moving_average"], color="red", label="7-day roling average")
    plt.gca().yaxis.set_minor_locator(ticker.MultipleLocator(1000))
    plt.ylabel("Daily cases")
    plt.xlim((list(df_2021_h1["date"])[0], list(df_2021_h1["date"])[-1]))

    plt.grid(axis="y", which="both")
    plt.title(f"COVID cases from {start.date().isoformat()} until {end.date().isoformat()} (remade)")
    plt.xticks(rotation=90)
    plt.savefig(f"Chart of UK Positive COVID-19 Antibody Tests from {start.date().isoformat()} to {end.date().isoformat()} (remade).svg")
    
plot_cases(df, datetime.datetime(2021, 1, 1), datetime.datetime(2021, 7, 1))
plt.clf()
plot_cases(df, datetime.datetime(2021, 7, 1), datetime.datetime(2022, 1, 1))
plt.clf()
plot_cases(df, datetime.datetime(2022, 1, 1), datetime.datetime(2022, 7, 1))
plt.clf()
plot_cases(df, datetime.datetime(2022, 7, 1), datetime.datetime(2023, 1, 1))