# Generated by Django 3.0.7 on 2020-08-24 20:05

from django.db import migrations, models


def update_units(apps, schema_editor):

    UnitType = apps.get_model('kpis', 'UnitType')
    Unit = apps.get_model('kpis', 'Unit')
    Department = apps.get_model('kpis', 'Department')
    KPI = apps.get_model('kpis', 'KPI')
    db_alias = schema_editor.connection.alias

    # Create UnitTypes
    kinds = {}
    for kind in ['Division', 'Department', 'Beamline', 'Group']:
        kinds[kind.lower()] = UnitType.objects.using(db_alias).create(name=kind)

    # All existing Units are Beamlines
    Unit.objects.using(db_alias).all().update(kind=kinds['beamline'])

    # Create Units for Departments/Divisions
    to_create = []
    for dept in Department.objects.using(db_alias).all():
        info = {
            'name': dept.name,
            'acronym': dept.acronym,
            'parent': dept.division and Unit.objects.using(db_alias).filter(name=dept.division.name).first() or None,
            'kind': dept.division and kinds['department'] or kinds['division']
        }
        to_create.append(Unit(**info))
    Unit.objects.using(db_alias).bulk_create(to_create)

    for unit in Unit.objects.using(db_alias).exclude(department__isnull=True).filter(parent__isnull=True):
        unit.parent = Unit.objects.using(db_alias).filter(name=unit.department.name).first()
        unit.save()

    for kpi in KPI.objects.using(db_alias).all():
        if kpi.beamline:
            kpi.units.add(kpi.beamline)
        if kpi.department:
            kpi.units.add(Unit.objects.using(db_alias).get(name=kpi.department.name))


class Migration(migrations.Migration):

    dependencies = [
        ('kpis', '0021_auto_20200824_1343'),
    ]

    operations = [
        migrations.AddField(
            model_name='kpi',
            name='units',
            field=models.ManyToManyField(blank=True, to='kpis.Unit'),
        ),
        migrations.RunPython(update_units),
    ]
