# Generated by Django 4.2.5 on 2023-10-11 13:34

from pathlib import Path

import lamindb_setup as ln_setup
from django.db import migrations

import wetlab.models

CORE_MODELS = {
    "ExperimentType": False,
    "Experiment": False,
    "Well": False,
    "TreatmentTarget": False,
    "Treatment": False,
    "Biosample": False,
    "Techsample": False,
}


def import_registry(registry, directory):
    import pandas as pd

    table_name = registry._meta.db_table
    df = pd.read_parquet(directory / f"{table_name}.parquet")
    old_foreign_key_columns = [column for column in df.columns if column.endswith("_old")]
    for column in old_foreign_key_columns:
        df.drop(column, axis=1, inplace=True)
    df.to_sql(table_name, ln_setup.settings.instance.db, if_exists="append", index=False)


def import_db(apps, schema_editor):
    # import data from parquet files
    directory = Path(f"./lamindb_export/{ln_setup.settings.instance.slug}/")
    if directory.exists():
        for model_name in CORE_MODELS.keys():
            registry = getattr(wetlab.models, model_name)
            import_registry(registry, directory)
            many_to_many_names = [field.name for field in registry._meta.many_to_many]
            for many_to_many_name in many_to_many_names:
                link_orm = getattr(registry, many_to_many_name).through
                import_registry(link_orm, directory)


class Migration(migrations.Migration):
    dependencies = [
        ("wetlab", "0001_initial_squashed_0012"),
        ("lnschema_core", "0024_import_legacy_data"),
        ("lnschema_bionty", "0017_import_legacy_data"),
    ]

    operations = [migrations.RunPython(import_db, reverse_code=migrations.RunPython.noop)]
