# Generated by Django 3.0.4 on 2020-03-19 21:06

from django.db import migrations
import csv

DATA_SOURCE_PATHS = {
    "GLOBAL": {
        "INFECTION": "../results/forecasts/global_data.csv",
        "DEATH": "../results/forecasts/global_deaths.csv"
    },
    "US": {
        "INFECTION": "../results/forecasts/us_data.csv",
        "DEATH": "../results/forecasts/us_deaths.csv"
    } 
}

def load_csv(apps, path, scope, type):
    """
    Reads the given CSV and returns a list of Covid19DataPoint/Covid19DeathDataPoint
    objects. The date, area, and val fields on the objects are set. Note
    that the objects have not been saved into the database yet.
    :param apps: Django apps object.
    :param path: Path to CSV with prediction data.
    :param scope: The scope of the historical data, "GLOBAL" or "US"
    :param type: The type of the historical data, "INFECTION" or "DEATH"
    :return: List of Covid19DataPoint/Covid19DeathDataPoint objects (NOT SAVED YET).
    """
    Area = apps.get_model('model_api', 'Area')
    Covid19DataPoint = apps.get_model('model_api', 'Covid19DataPoint')
    Covid19DeathDataPoint = apps.get_model(
        'model_api', 'Covid19DeathDataPoint')

    with open(path, 'r') as f:
        reader = csv.reader(f)
        header = next(reader, None)

        data = []
        using_death_data_point = type == "DEATH"
        using_state_level = scope == "US"

        for row in reader:
            area = None

            # Determine the country / state.
            if not using_state_level:
                country = row[1]
                state = ''
            else:
                country = 'US'
                state = row[1]

            # Try to find the corresponding area.
            try:
                area = Area.objects.get(country=country, state=state)
            except Area.DoesNotExist:
                area = Area(state=state, country=country)
                area.save()
                msg = 'Created area ' + ((state + ' ') if state != '' else '') + country + '.'

            except Area.MultipleObjectsReturned:
                msg = "Found multiple areas for country '{0}'".format(
                    country)
                if state:
                    msg += " and state '{0}'".format(state)
                msg += ' in model_api_area. Skip this area.'
                print(msg)
                continue

            # Load the data.
            for i in range(2, len(header)):
                date = header[i]

                # Skip invalid values.
                raw_val = row[i]
                if raw_val in ['NaN', '-Inf', 'Inf']:
                    continue

                # Skip negative values.
                val = int(float(raw_val))
                if val < 0:
                    continue

                if not using_death_data_point:
                    data.append(Covid19DataPoint(
                        area=area,
                        date=date,
                        val=val,
                    ))
                else:
                    data.append(Covid19DeathDataPoint(
                        area=area,
                        date=date,
                        val=val,
                    ))

    return data


def load_covid19_data(apps, schema_editor):
    print()
    Area = apps.get_model('model_api', 'Area')
    Covid19DataPoint = apps.get_model('model_api', 'Covid19DataPoint')
    Covid19DeathDataPoint = apps.get_model(
        'model_api', 'Covid19DeathDataPoint')

    for scope in DATA_SOURCE_PATHS:
        for type in DATA_SOURCE_PATHS[scope]:
            path = DATA_SOURCE_PATHS[scope][type]
            observed_data = load_csv(apps, path, scope, type)
            for d in observed_data:
                d.save() 


def delete_covid19_data(apps, schema_editor):
    Area = apps.get_model('model_api', 'Area')
    Covid19DataPoint = apps.get_model('model_api', 'Covid19DataPoint')

    # Clear any previously-existing data.
    Area.objects.all().delete()
    Covid19DataPoint.objects.all().delete()


class Migration(migrations.Migration):

    dependencies = [
        ('model_api', '0001_initial'),
    ]

    operations = [
        migrations.RunPython(load_covid19_data, delete_covid19_data),
    ]
