#!/usr/bin/env python

"""
Download data from OWID and store it in a local sqlite3 database.

The data can then be used for further analysis.
Use `NULL` for missing values.
Only insert values for a selection of countries.
"""

import codecs
import csv
import sqlite3
import urllib.request


URL = ('https://raw.githubusercontent.com/owid/covid-19-data/master/'
       'public/data/excess_mortality/excess_mortality.csv')
COUNTRIES = ('Austria', 'Sweden')

con = sqlite3.connect('mortality_database.db')
cur = con.cursor()
cur.execute('''
create table if not exists countries (
    country_id integer primary key,
    name text not null
);
''')
cur.execute('''
create table if not exists mortality (
    mortality_id integer primary key,
    country_id integer,
    week_start_iso_date text,
    week_excess real,
    cumulative_excess real,
    week_excess_per_million real,
    cumulative_excess_per_million real,
    foreign key (country_id) references countries(country_id)
);
''')

country_id = 0
country = ''
with urllib.request.urlopen(URL) as stream:
    reader = csv.DictReader(codecs.iterdecode(stream, 'utf-8'))
    for row in reader:
        if row['location'] not in COUNTRIES:
            continue
        if country != row['location']:
            country = row['location']
            country_id += 1
            cur.execute(f'''
INSERT INTO countries (country_id, name) VALUES ({country_id}, '{country}');
''')
        # TODO: insert data into mortality table


con.commit()
cur.close()
con.close()