#!/usr/bin/python

import requests
import csv
import datetime
import os
import matplotlib.pyplot as pp
import numpy as np
import sys
import importlib
import time
sys.path.append(".")

#### config
# countries of interest
countries = ["Germany", "Italy", "China", "Japan", "South Korea", "Iran", "United States", "World", "United Kingdom", "Sweden"]
# enabled plots
plots = ["basics", "death_per_case",
        #"normalized_to_first_death", "delay_from_china", "delay_from_usa", "normalized_to_ten_cases", "percent_increase",
        "doubling_time",
        "all_countries",
        ]
### manual data
# population: sourced ECDC data
from population_repository import pop

###
def get_data():
    """fetch data from remote, cache locally and reorganize internal data
    not beautiful (at all), but effective!!"""

    tries = 10
    delay = 10

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

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

    if not os.path.isfile(datafile):

        for n in tries:
            try:
                r = requests.get(dataurl)
            except:
                print(f"==> download failed, retrying after {delay}s up to another {tries-n} times…")
                time.sleep(delay)
                continue
            break

        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:
            if len(row) == 6:
                date,location,new_cases,new_deaths,total_cases,total_deaths = row
            elif len(row) == 10:
                date,location,new_cases,new_deaths,total_cases,total_deaths,weekly_cases,weekly_deaths,biweekly_cases,biweekly_deaths = row
            else:
                print(f"WARNING! Table format changed, new header:\n{row})")
                exit(1)

            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])
    # reorganize data
    data2 = {}
    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]
    return data2

data = get_data()

for plot in plots:
    i = importlib.import_module(plot)
    i.plot(data, countries, pop)