# Generated by Django 3.0.4 on 2020-06-07 18:35

from django.db import migrations
import csv
import datetime

GLOBAL_QUARANTINE_SCORE_CVS_PATH = "../results/scores/global_Rt_num.csv"
GLOBAL_QUARANTINE_SCORE_CONF_CVS_PATH = "../results/scores/global_Rt_conf.csv"
US_QUARANTINE_SCORE_CSV_PATH = "../results/scores/us_Rt_num.csv"
US_QUARANTINE_SCORE_CONF_CSV_PATH = "../results/scores/us_Rt_conf.csv"

def load_quarantine_score(apps, schema_editor):
    Area = apps.get_model('model_api', 'Area')
    QuarantineScoreDataPoint = apps.get_model('model_api', 'QuarantineScoreDataPoint')

    # Country level score.
    with open (GLOBAL_QUARANTINE_SCORE_CVS_PATH) as f1, open(GLOBAL_QUARANTINE_SCORE_CONF_CVS_PATH) as f2:
        reader1 = csv.reader(f1)
        reader2 = csv.reader(f2)
        
        header = next(reader1, None)
        next(reader2, None)

        for row in reader1:
            conf_row = next(reader2, None)
            
            state = ""
            country = row[1]
            
            area = Area.objects.get(country=country, state=state)

            for i in range(2, len(header)):
                raw_date = header[i]
                date = datetime.datetime(*[int(item) for item in raw_date.split('-')])
                val = float(row[i]) if float(row[i]) <= 7 else 7 # Set a upper limit of 7.
                if conf_row[i] == "NaN":
                    conf = 0
                else:
                    conf = float(conf_row[i])

                quarantine_score = QuarantineScoreDataPoint(
                    area=area,
                    date=date,
                    val=val,
                    conf=conf
                )
                quarantine_score.save()


    # US state-level score.
    with open(US_QUARANTINE_SCORE_CSV_PATH) as f1, open(US_QUARANTINE_SCORE_CONF_CSV_PATH) as f2:
        reader1 = csv.reader(f1)
        reader2 = csv.reader(f2)

        header = next(reader1, None)
        next(reader2, None)

        for row in reader1:
            conf_row = next(reader2, None)

            state = row[1]
            country = "US"

            area = Area.objects.get(country=country, state=state)

            for i in range(2, len(header)):
                raw_date = header[i]
                date = datetime.datetime(*[int(item) for item in raw_date.split('-')])
                val = float(row[i]) if float(row[i]) <= 7 else 7 # Set a upper limit of 7.
                if conf_row[i] == "NaN":
                    conf = 0
                else:
                    conf = float(conf_row[i])

                quarantine_score = QuarantineScoreDataPoint(
                    area=area,
                    date=date,
                    val=val,
                    conf=conf
                )
                quarantine_score.save()


def delete_quarantine_score(apps, schema_editor):
    Area = apps.get_model('model_api', 'Area')
    QurantineScore = apps.get_model('model_api', 'QuarantineScoreDataPoint')

    Area.objects.all().delete()
    QurantineScore.objects.all().delete()
            

class Migration(migrations.Migration):

    dependencies = [
        ('model_api', '0010_quarantinescore'),
    ]

    operations = [
        migrations.RunPython(load_quarantine_score, delete_quarantine_score),
    ]
