#!/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", "India", "Japan", "Brazil", "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/owid-covid-data.csv"
    date = datetime.date.today()

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

    if not os.path.isfile(datafile):

        for n in range(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
            elif len(row) == 50:
                iso_code,continent,location,date,total_cases,new_cases,new_cases_smoothed,total_deaths,new_deaths,new_deaths_smoothed,total_cases_per_million,new_cases_per_million,new_cases_smoothed_per_million,total_deaths_per_million,new_deaths_per_million,new_deaths_smoothed_per_million,reproduction_rate,icu_patients,icu_patients_per_million,hosp_patients,hosp_patients_per_million,weekly_icu_admissions,weekly_icu_admissions_per_million,weekly_hosp_admissions,weekly_hosp_admissions_per_million,new_tests,total_tests,total_tests_per_thousand,new_tests_per_thousand,new_tests_smoothed,new_tests_smoothed_per_thousand,positive_rate,tests_per_case,tests_units,stringency_index,population,population_density,median_age,aged_65_older,aged_70_older,gdp_per_capita,extreme_poverty,cardiovasc_death_rate,diabetes_prevalence,female_smokers,male_smokers,handwashing_facilities,hospital_beds_per_thousand,life_expectancy,human_development_index = row
            else:
                print(f"WARNING! Table format changed, new header:\n{row})")
                exit(1)
            # break loop if header
            if location=="location":
                # table header
                continue
            # cast to num type
            if isinstance(total_cases, str):
                if total_cases in ('', ' ', '0.0'):
                    total_cases = 0
                try:
                    total_cases = int(total_cases)
                except:
                    try:
                        total_cases = float(total_cases)
                    except:
                        print("tc failed")
            if isinstance(new_cases, str):
                if new_cases in ('', ' ', '0.0'):
                    new_cases = 0
                try:
                    new_cases = int(new_cases)
                except:
                    try:
                        new_cases = float(new_cases)
                    except:
                        print("nc failed")
            if isinstance(new_deaths, str):
                if new_deaths in ('', ' ', '0.0'):
                    new_deaths = 0
                try:
                    new_deaths = int(new_deaths)
                except:
                    try:
                        new_deaths = float(new_deaths)
                    except:
                        print("nd failed")
            if isinstance(total_deaths, str):
                if total_deaths in ('', ' ', '0.0'):
                    total_deaths = 0
                try:
                    total_deaths = int(total_deaths)
                except:
                    try:
                        total_deaths = float(total_deaths)
                    except:
                        print("td failed")
            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)