from django.db import migrations


def copy_annotation(apps, schema_editor):
    Category = apps.get_model('api', 'Category')
    Span = apps.get_model('api', 'Span')
    for model in [Category, Span]:
        for annotation in model.objects.all():
            if model == Category:
                LabelModel = apps.get_model('api', 'CategoryType')
            else:
                LabelModel = apps.get_model('api', 'SpanType')
            label = LabelModel.objects.get(pk=annotation.label.id)
            annotation.new_label = label
            annotation.save()


def delete_annotation(apps, schema_editor):
    Category = apps.get_model('api', 'Category')
    Span = apps.get_model('api', 'Span')
    for model in [Category, Span]:
        for annotation in model.objects.all():
            annotation.new_label = None
            annotation.save()


class Migration(migrations.Migration):

    dependencies = [
        ('api', '0022_auto_20211221_1430'),
    ]

    operations = [
        migrations.RunPython(
            code=copy_annotation,
            reverse_code=delete_annotation
        ),
    ]