# sputnik.irk.ru/virus/info/
import requests
import numpy as np
from datetime import datetime
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

countries = {
    'South Korea':   (52, 'Ю. Корея', '#77D577'),
    'United States': (328, 'США', '#00DFDF'),
    'Italy':         (60.5, 'Италия', '#D7CD73'),
    'Russia':        (146.7, 'Россия', '#0000FF'),
    'Irkutsk':       (2.4, 'Иркутская область', '#FF0000'),
}

# download world data
csv_url = 'https://raw.githubusercontent.com/owid/covid-19-data/master/public/data/owid-covid-data.csv'
ans = requests.get(csv_url)
data = {}
date0 = datetime(2020,1,1)
for line in ans.text.split('\n')[1:-1]:
    country, date, n = line.split(',')[2:5]
    if n=='':
        continue
    if country not in data:
        data[country] = np.zeros(365, dtype=np.int)
    iday = (datetime.strptime(date, '%Y-%m-%d') - date0).days
    if iday < 0:
        continue
    data[country][iday] = int(float(n))

# load Irkutsk data
data['Irkutsk'] = np.zeros(365, dtype=np.int)
for line in reversed(open('irk.csv').readlines()[1:]):
    date, n = line.strip().split(',')[:2]
    iday = (datetime.strptime(date, '%Y-%m-%d') - date0).days
    data['Irkutsk'][iday] = int(n)

ndays = 5
for country, (population, label, color) in countries.items():
    y = data[country] / population

    # linear
    plt.figure(1)
    ind = np.where(y>1)[0]
    y = y[ind[0]:ind[-1]+1]
    plt.plot(y, label=label, color=color)

    # new
    plt.figure(2)
    dy = y[ndays:] - y[:-ndays]
    plt.plot(dy, label=label, color=color)

    # log
    plt.figure(3)
    plt.plot(y, label=label, color=color)

    # procent
    plt.figure(4)
    y = np.log10(y)
    procent = [np.NaN]*ndays
    for iday in range(ndays, len(y)):
        rate = (10 ** (y[iday] - y[iday - ndays])) ** (1 / ndays)
        procent.append((rate - 1) * 100)
    plt.plot(procent, label=label, color=color)

szx, szy = 960/100, 700/100
plt.figure(1).set_size_inches(szx, szy, forward=True)
plt.title('Линейный, ' + date)
plt.ylim(0, np.max(data['Russia']/countries['Russia'][0])*2)
plt.legend()
plt.xlabel('Дни с начала эпидемии')
plt.ylabel('Кол-во заражений на 1 млн.')
plt.tight_layout()
plt.savefig('linear.png', dpi=100)

plt.figure(2).set_size_inches(szx, szy, forward=True)
plt.title('Кол-во новых заражений (за 5 дней), ' + date)
plt.legend()
plt.xlabel('Дни с начала эпидемии')
plt.ylabel('Кол-во заражений на 1 млн.')
plt.tight_layout()
plt.savefig('new.png', dpi=100)

plt.figure(3).set_size_inches(szx, szy, forward=True)
plt.title('Логарифмический, ' + date)
plt.legend()
plt.yscale('log')
plt.xlabel('Дни с начала эпидемии')
plt.ylabel('Кол-во заражений на 1 млн.')
plt.tight_layout()
plt.savefig('log.png', dpi=100)

plt.figure(4).set_size_inches(szx, szy, forward=True)
plt.title('Процент прироста (за 5 дней), ' + date)
plt.xlabel('Дни с начала эпидемии')
plt.ylabel('% увеличения в день')
axes = plt.gca()
axes.yaxis.grid()
plt.legend()
plt.tight_layout()
plt.savefig('procent.png', dpi=100)