#!/usr/bin/python

import requests
import csv
import datetime
import os
import matplotlib.pyplot as pp
import numpy as np
import sys
import importlib
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 from wikidata and manually adjusted. not complete wrt the ECDC data!
pop = {'Afghanistan': {'wd_key': 'Afghanistan', 'pop': 34940837},
 'Albania': {'wd_key': 'Albania', 'pop': 3020209},
 'Algeria': {'wd_key': 'Algeria', 'pop': 41318142},
 'Andorra': {'wd_key': 'Andorra', 'pop': 76177},
 'Angola': {'wd_key': 'Angola', 'pop': 29784193},
 'Antigua and Barbuda': {'wd_key': 'Antigua and Barbuda', 'pop': 102012},
 'Argentina': {'wd_key': 'Argentina', 'pop': 44938712},
 'Armenia': {'wd_key': 'Armenia', 'pop': 2930450},
 'Australia': {'wd_key': 'Australia', 'pop': 24511800},
 'Austria': {'wd_key': 'Austria', 'pop': 8809212},
 'Azerbaijan': {'wd_key': 'Azerbaijan', 'pop': 10000000},
 'Bahamas': {'wd_key': 'The Bahamas', 'pop': 395361},
 'Bahrain': {'wd_key': 'Bahrain', 'pop': 1492584},
 'Bangladesh': {'wd_key': 'Bangladesh', 'pop': 164669751},
 'Barbados': {'wd_key': 'Barbados', 'pop': 285719},
 'Belarus': {'wd_key': 'Belarus', 'pop': 9408400},
 'Belgium': {'wd_key': 'Belgium', 'pop': 11431406},
 'Belize': {'wd_key': 'Belize', 'pop': 374681},
 'Benin': {'wd_key': 'Benin', 'pop': 11175692},
 'Bhutan': {'wd_key': 'Bhutan', 'pop': 807610},
 'Bolivia': {'wd_key': 'Bolivia', 'pop': 11051600},
 'Bosnia and Herzegovina': {'wd_key': 'Bosnia and Herzegovina', 'pop': 3507017},
 'Botswana': {'wd_key': 'Botswana', 'pop': 2291661},
 'Brazil': {'wd_key': 'Brazil', 'pop': 210147125},
 'Brunei': {'wd_key': 'Brunei', 'pop': 428697},
 'Bulgaria': {'wd_key': 'Bulgaria', 'pop': 7000039},
 'Burkina Faso': {'wd_key': 'Burkina Faso', 'pop': 19193382},
 'Burundi': {'wd_key': 'Burundi', 'pop': 10864245},
 'Cambodia': {'wd_key': 'Cambodia', 'pop': 16005373},
 'Cameroon': {'wd_key': 'Cameroon', 'pop': 24053727},
 'Canada': {'wd_key': 'Canada', 'pop': 37894799},
 'Cape Verde': {'wd_key': 'Cape Verde', 'pop': 546388},
 'Central African Republic': {'wd_key': 'Central African Republic', 'pop': 4659080},
 'Chad': {'wd_key': 'Chad', 'pop': 15477751},
 'Chile': {'wd_key': 'Chile', 'pop': 18054726},
 'China': {'wd_key': "People's Republic of China", 'pop': 1409517397},
 'Colombia': {'wd_key': 'Colombia', 'pop': 49065615},
 'Comoros': {'wd_key': 'Comoros', 'pop': 823652},
 'Costa Rica': {'wd_key': 'Costa Rica', 'pop': 4905769},
 'Croatia': {'wd_key': 'Croatia', 'pop': 4105493},
 'Cuba': {'wd_key': 'Cuba', 'pop': 11484636},
 'Cyprus': {'wd_key': 'Cyprus', 'pop': 1141166},
 'Czech Republic': {'wd_key': 'Czech Republic', 'pop': 10693939},
 'Denmark': {'wd_key': 'Denmark', 'pop': 5827463},
 'Djibouti': {'wd_key': 'Djibouti', 'pop': 956985},
 'Dominica': {'wd_key': 'Dominica', 'pop': 73925},
 'Dominican Republic': {'wd_key': 'Dominican Republic', 'pop': 10403761},
 'Ecuador': {'wd_key': 'Ecuador', 'pop': 16624858},
 'Egypt': {'wd_key': 'Egypt', 'pop': 94798827},
 'El Salvador': {'wd_key': 'El Salvador', 'pop': 6420746},
 'Equatorial Guinea': {'wd_key': 'Equatorial Guinea', 'pop': 1267689},
 'Eritrea': {'wd_key': 'Eritrea', 'pop': 3497000},
 'Estonia': {'wd_key': 'Estonia', 'pop': 1324820},
 'Ethiopia': {'wd_key': 'Ethiopia', 'pop': 104957438},
 'Fiji': {'wd_key': 'Fiji', 'pop': 905502},
 'Finland': {'wd_key': 'Finland', 'pop': 5501043},
 'France': {'wd_key': 'France', 'pop': 66628000},
 'Gabon': {'wd_key': 'Gabon', 'pop': 2025137},
 'Gambia': {'wd_key': 'The Gambia', 'pop': 2100568},
 'Georgia': {'wd_key': 'Georgia', 'pop': 3717100},
 'Germany': {'wd_key': 'Germany', 'pop': 83149300},
 'Ghana': {'wd_key': 'Ghana', 'pop': 26908262},
 'Greece': {'wd_key': 'Greece', 'pop': 10760421},
 'Grenada': {'wd_key': 'Grenada', 'pop': 107825},
 'Guatemala': {'wd_key': 'Guatemala', 'pop': 17263239},
 'Guinea': {'wd_key': 'Equatorial Guinea', 'pop': 1267689},
 'Guinea-Bissau': {'wd_key': 'Guinea-Bissau', 'pop': 1861283},
 'Guyana': {'wd_key': 'Guyana', 'pop': 777859},
 'Haiti': {'wd_key': 'Haiti', 'pop': 10981229},
 'Honduras': {'wd_key': 'Honduras', 'pop': 9265067},
 'Hungary': {'wd_key': 'Hungary', 'pop': 9937628},
 'Iceland': {'wd_key': 'Iceland', 'pop': 357050},
 'India': {'wd_key': 'India', 'pop': 1326093247},
 'Indonesia': {'wd_key': 'Indonesia', 'pop': 263991379},
 'Iran': {'wd_key': 'Iran', 'pop': 79966230},
 'Iraq': {'wd_key': 'Iraq', 'pop': 38274618},
 'Ireland': {'wd_key': 'Ireland', 'pop': 4761865},
 'Israel': {'wd_key': 'Israel', 'pop': 9093000},
 'Italy': {'wd_key': 'Italy', 'pop': 60317000},
 'Jamaica': {'wd_key': 'Jamaica', 'pop': 2890299},
 'Japan': {'wd_key': 'Japan', 'pop': 126785797},
 'Jordan': {'wd_key': 'Jordan', 'pop': 10428241},
 'Kazakhstan': {'wd_key': 'Kazakhstan', 'pop': 18276500},
 'Kenya': {'wd_key': 'Kenya', 'pop': 48468138},
 'Kuwait': {'wd_key': 'Kuwait', 'pop': 4600000},
 'Kyrgyzstan': {'wd_key': 'Kyrgyzstan', 'pop': 6201500},
 'Laos': {'wd_key': 'Laos', 'pop': 6858160},
 'Latvia': {'wd_key': 'Latvia', 'pop': 1909000},
 'Lebanon': {'wd_key': 'Lebanon', 'pop': 6100075},
 'Lesotho': {'wd_key': 'Lesotho', 'pop': 2233339},
 'Liberia': {'wd_key': 'Liberia', 'pop': 4731906},
 'Libya': {'wd_key': 'Libya', 'pop': 6678567},
 'Liechtenstein': {'wd_key': 'Liechtenstein', 'pop': 37922},
 'Lithuania': {'wd_key': 'Lithuania', 'pop': 2790842},
 'Luxembourg': {'wd_key': 'Luxembourg', 'pop': 626108},
 'Macedonia': {'wd_key': 'North Macedonia', 'pop': 2075301},
 'Madagascar': {'wd_key': 'Madagascar', 'pop': 25570895},
 'Malawi': {'wd_key': 'Malawi', 'pop': 18622104},
 'Malaysia': {'wd_key': 'Malaysia', 'pop': 31624264},
 'Maldives': {'wd_key': 'Maldives', 'pop': 436330},
 'Mali': {'wd_key': 'Mali', 'pop': 18541980},
 'Malta': {'wd_key': 'Malta', 'pop': 465292},
 'Mauritania': {'wd_key': 'Mauritania', 'pop': 4420184},
 'Mauritius': {'wd_key': 'Mauritius', 'pop': 1264613},
 'Mexico': {'wd_key': 'Mexico', 'pop': 130526945},
 'Moldova': {'wd_key': 'Moldova', 'pop': 2550900},
 'Monaco': {'wd_key': 'Monaco', 'pop': 38695},
 'Mongolia': {'wd_key': 'Mongolia', 'pop': 3075647},
 'Montenegro': {'wd_key': 'Montenegro', 'pop': 622359},
 'Morocco': {'wd_key': 'Morocco', 'pop': 36029138},
 'Mozambique': {'wd_key': 'Mozambique', 'pop': 29668834},
 'Myanmar': {'wd_key': 'Myanmar', 'pop': 53370609},
 'Namibia': {'wd_key': 'Namibia', 'pop': 2533794},
 'Nepal': {'wd_key': 'Nepal', 'pop': 29400000},
 'Netherlands': {'wd_key': 'Kingdom of the Netherlands', 'pop': 17100715},
 'New Zealand': {'wd_key': 'New Zealand', 'pop': 4942500},
 'Nicaragua': {'wd_key': 'Nicaragua', 'pop': 6217581},
 'Niger': {'wd_key': 'Niger', 'pop': 21477348},
 'Nigeria': {'wd_key': 'Nigeria', 'pop': 190886311},
 'Norway': {'wd_key': 'Norway', 'pop': 5367580},
 'Oman': {'wd_key': 'Oman', 'pop': 4829480},
 'Pakistan': {'wd_key': 'Pakistan', 'pop': 197015955},
 'Panama': {'wd_key': 'Panama', 'pop': 4098587},
 'Papua New Guinea': {'wd_key': 'Papua New Guinea', 'pop': 8251162},
 'Paraguay': {'wd_key': 'Paraguay', 'pop': 6811297},
 'Peru': {'wd_key': 'Peru', 'pop': 29381884},
 'Philippines': {'wd_key': 'Philippines', 'pop': 104918090},
 'Poland': {'wd_key': 'Poland', 'pop': 38433600},
 'Portugal': {'wd_key': 'Portugal', 'pop': 10600000},
 'Qatar': {'wd_key': 'Qatar', 'pop': 2639211},
 'Romania': {'wd_key': 'Romania', 'pop': 19586539},
 'Russia': {'wd_key': 'Russia', 'pop': 146804372},
 'Rwanda': {'wd_key': 'Rwanda', 'pop': 12208407},
 'Saint Kitts and Nevis': {'wd_key': 'Saint Kitts and Nevis', 'pop': 55345},
 'Saint Lucia': {'wd_key': 'Saint Lucia', 'pop': 178844},
 'Saint Vincent and the Grenadines': {'wd_key': 'Saint Vincent and the Grenadines', 'pop': 109897},
 'San Marino': {'wd_key': 'San Marino', 'pop': 33400},
 'Saudi Arabia': {'wd_key': 'Saudi Arabia', 'pop': 33000000},
 'Senegal': {'wd_key': 'Senegal', 'pop': 15850567},
 'Serbia': {'wd_key': 'Serbia', 'pop': 7022268},
 'Seychelles': {'wd_key': 'Seychelles', 'pop': 95843},
 'Sierra Leone': {'wd_key': 'Sierra Leone', 'pop': 7557212},
 'Singapore': {'wd_key': 'Singapore', 'pop': 5888926},
 'Slovakia': {'wd_key': 'Slovakia', 'pop': 5397036},
 'Slovenia': {'wd_key': 'Slovenia', 'pop': 2066880},
 'Somalia': {'wd_key': 'Somalia', 'pop': 11031386},
 'South Africa': {'wd_key': 'South Africa', 'pop': 57725600},
 'South Korea': {'wd_key': 'South Korea', 'pop': 51466201},
 'South Sudan': {'wd_key': 'South Sudan', 'pop': 12575714},
 'Spain': {'wd_key': 'Spain', 'pop': 46733038},
 'Sri Lanka': {'wd_key': 'Sri Lanka', 'pop': 21444000},
 'Sudan': {'wd_key': 'Sudan', 'pop': 40533330},
 'Suriname': {'wd_key': 'Suriname', 'pop': 563402},
 'Sweden': {'wd_key': 'Sweden', 'pop': 10343403},
 'Switzerland': {'wd_key': 'Switzerland', 'pop': 8466017},
 'Syria': {'wd_key': 'Syria', 'pop': 18269868},
 'Tajikistan': {'wd_key': 'Tajikistan', 'pop': 8921343},
 'Tanzania': {'wd_key': 'Tanzania', 'pop': 57310019},
 'Thailand': {'wd_key': 'Thailand', 'pop': 65931550},
 'Timor': {'wd_key': 'East Timor', 'pop': 1296311},
 'Togo': {'wd_key': 'Togo', 'pop': 7797694},
 'Trinidad and Tobago': {'wd_key': 'Trinidad and Tobago', 'pop': 1369125},
 'Tunisia': {'wd_key': 'Tunisia', 'pop': 11565204},
 'Turkey': {'wd_key': 'Turkey', 'pop': 82003882},
 'Uganda': {'wd_key': 'Uganda', 'pop': 42862958},
 'Ukraine': {'wd_key': 'Ukraine', 'pop': 42558328},
 'United Arab Emirates': {'wd_key': 'United Arab Emirates', 'pop': 9400145},
 'United Kingdom': {'wd_key': 'United Kingdom', 'pop': 66022273},
 'United States': {'wd_key': 'United States of America', 'pop': 325145963},
 'Uruguay': {'wd_key': 'Uruguay', 'pop': 3456750},
 'Uzbekistan': {'wd_key': 'Uzbekistan', 'pop': 32387200},
 'Vatican': {'wd_key': 'Vatican City', 'pop': 1000},
 'Venezuela': {'wd_key': 'Venezuela', 'pop': 28515829},
 'Vietnam': {'wd_key': 'Vietnam', 'pop': 94660000},
 'Yemen': {'wd_key': 'Yemen', 'pop': 28250420},
 'Zambia': {'wd_key': 'Zambia', 'pop': 17094130},
 'Zimbabwe': {'wd_key': 'Zimbabwe', 'pop': 16529904}
}

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

    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):
        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:
            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)