#!/usr/bin/python

import requests
import csv
import datetime
import os
import matplotlib.pyplot as pp
import numpy as np


dataurl = "https://covid.ourworldindata.org/data/full_data.csv"
date = datetime.date.today()

datafile = f"{date}-full-data.csv"

if not os.path.isfile(datafile):
    r = requests.get(dataurl)
    with open(datafile, "wb") as f:
        f.write(r.content)
else:
    print(f"file found: {datafile}")

# processing
data = {}
with open(datafile, "r") as f:
    reader = csv.reader(f)
    for row in reader:
        date,location,new_cases,new_deaths,total_cases,total_deaths = row
        if location=="location":
            # table header
            continue
        if location not in data:
            data[location] = []
        year, month, day = date.split("-")
        def toint(a):
            try:
                return int(a)
            except:
                return np.nan
        data[location].append([datetime.date(int(year), int(month), int(day)), new_cases, new_deaths, total_cases, total_deaths])


data2 = {}

# reorganize data
for loc in data:
    time = []
    new_cases, new_deaths, total_cases, total_deaths = [], [], [], []
    for entry in data[loc]:
        t_, new_cases_, new_deaths_, total_cases_, total_deaths_ = entry
        time.append(t_)
        new_cases.append(toint(new_cases_))
        new_deaths.append(toint(new_deaths_))
        total_cases.append(toint(total_cases_))
        total_deaths.append(toint(total_deaths_))
    data2[loc] = [time, new_cases, new_deaths, total_cases, total_deaths]

# countries of interest
countries = ["Germany", "Italy", "Spain", "Denmark", "China", "Japan", "South Korea", "Iran", "Belgium", "World"]

# plotting
for loc in data2:
    if loc not in countries:
        continue
    time, new_cases, new_deaths, total_cases, total_deaths = data2[loc]
    
    # total cases
    pp.figure("total_cases")
    pp.plot(time, total_cases, label=f"{loc}")

    # death/case
    pp.figure("death_per_case")
    pp.plot(time, np.array(total_deaths)/np.array(total_cases), label=f"{loc}")
# figure postprocessing
for fig in ["total_cases", "death_per_case"]:
    pp.figure(fig)
    pp.yscale("log")
    pp.xticks(rotation=90)
    pp.legend(frameon=False)
    pp.tight_layout()

    pp.savefig(f"{fig}.png")