# Generated by Django 5.1.8 on 2026-01-07 20:13

from pathlib import Path

import django.db.models.deletion
import lamindb.base.fields
from django.db import migrations, models
from lamin_utils import logger


def export_tables_to_parquet(apps, schema_editor):
    """Export data from tables being dropped to parquet files."""
    from datetime import datetime

    import pandas as pd

    # Create export directory with timestamp
    export_dir = Path("data_exports") / datetime.now().strftime("%Y%m%d_%H%M%S")
    export_dir.mkdir(parents=True, exist_ok=True)

    # Models to export
    models_to_export = [
        "ArtifactBiosample",
        "ArtifactDonor",
        "ArtifactExperiment",
        "ArtifactTechsample",
        "ArtifactWell",
        "Biosample",
        "Donor",
        "Experiment",
        "Techsample",
        "Well",
    ]

    exported = False

    for model_name in models_to_export:
        try:
            Model = apps.get_model("wetlab", model_name)

            # Check if there's any data
            if not Model.objects.exists():
                continue

            exported = True

            # Get field names, excluding many-to-many fields
            regular_fields = []
            m2m_fields = []

            for field in Model._meta.get_fields():
                if field.concrete:
                    if field.many_to_many:
                        m2m_fields.append(field.name)
                    else:
                        # For FK fields, get the _id field name
                        if hasattr(field, "related_model") and field.related_model:
                            regular_fields.append(f"{field.name}_id")
                        else:
                            regular_fields.append(field.name)

            # Fetch all data at once using .values()
            data = list(Model.objects.values(*regular_fields))

            # Handle M2M fields separately if any exist
            if m2m_fields:
                {row["id"]: idx for idx, row in enumerate(data)}

                for m2m_field_name in m2m_fields:
                    # Get the through table and fetch all relationships at once
                    m2m_field = Model._meta.get_field(m2m_field_name)
                    through_model = m2m_field.remote_field.through

                    # Get the field names for the source and target
                    source_field = None
                    target_field = None
                    for f in through_model._meta.get_fields():
                        if hasattr(f, "related_model"):
                            if f.related_model == Model:
                                source_field = f.name
                            elif f.related_model == m2m_field.related_model:
                                target_field = f.name

                    if source_field and target_field:
                        # Fetch all M2M relationships at once
                        m2m_data = through_model.objects.values_list(
                            f"{source_field}_id", f"{target_field}_id"
                        )

                        # Group by source ID
                        m2m_dict = {}
                        for source_id, target_id in m2m_data:
                            if source_id not in m2m_dict:
                                m2m_dict[source_id] = []
                            m2m_dict[source_id].append(target_id)

                        # Add M2M data to rows
                        for row in data:
                            row[m2m_field_name] = m2m_dict.get(row["id"], [])

            # Convert to DataFrame and export
            df = pd.DataFrame(data)
            output_path = export_dir / f"{model_name}.parquet"
            df.to_parquet(output_path, index=False)
            logger.important(
                f"Exported {len(df)} rows from {model_name} to {output_path}"
            )

        except Exception as e:
            logger.error(f"Error exporting {model_name}: {e}")
            raise e

    if exported:
        logger.important(
            f"All exports completed. Files saved to: {export_dir}\n  To keep using these records, you can create them in the Record registry."
        )


class Migration(migrations.Migration):
    dependencies = [
        ("bionty", "0063_remove_experimentalfactor_instrument_and_more"),
        ("wetlab", "0048_remove_artifactbiologic_feature_ref_is_name_and_more"),
    ]

    operations = [
        # Export data before dropping models
        migrations.RunPython(
            export_tables_to_parquet,
            reverse_code=migrations.RunPython.noop,
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="artifacts",
        ),
        migrations.RemoveField(
            model_name="artifactdonor",
            name="artifact",
        ),
        migrations.RemoveField(
            model_name="artifactdonor",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="artifactdonor",
            name="donor",
        ),
        migrations.RemoveField(
            model_name="artifactdonor",
            name="feature",
        ),
        migrations.RemoveField(
            model_name="artifactdonor",
            name="run",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="artifacts",
        ),
        migrations.RemoveField(
            model_name="artifactexperiment",
            name="artifact",
        ),
        migrations.RemoveField(
            model_name="artifactexperiment",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="artifactexperiment",
            name="experiment",
        ),
        migrations.RemoveField(
            model_name="artifactexperiment",
            name="feature",
        ),
        migrations.RemoveField(
            model_name="artifactexperiment",
            name="run",
        ),
        migrations.RemoveField(
            model_name="experiment",
            name="artifacts",
        ),
        migrations.RemoveField(
            model_name="artifacttechsample",
            name="artifact",
        ),
        migrations.RemoveField(
            model_name="artifacttechsample",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="artifacttechsample",
            name="feature",
        ),
        migrations.RemoveField(
            model_name="artifacttechsample",
            name="run",
        ),
        migrations.RemoveField(
            model_name="artifacttechsample",
            name="techsample",
        ),
        migrations.RemoveField(
            model_name="techsample",
            name="artifacts",
        ),
        migrations.RemoveField(
            model_name="artifactwell",
            name="artifact",
        ),
        migrations.RemoveField(
            model_name="artifactwell",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="artifactwell",
            name="feature",
        ),
        migrations.RemoveField(
            model_name="artifactwell",
            name="run",
        ),
        migrations.RemoveField(
            model_name="artifactwell",
            name="well",
        ),
        migrations.RemoveField(
            model_name="well",
            name="artifacts",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="branch",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="cell_lines",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="cell_types",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="diseases",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="organism",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="run",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="space",
        ),
        migrations.RemoveField(
            model_name="biosample",
            name="tissues",
        ),
        migrations.RemoveField(
            model_name="techsample",
            name="biosamples",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="branch",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="diseases",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="ethnicity",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="organism",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="run",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="sex",
        ),
        migrations.RemoveField(
            model_name="donor",
            name="space",
        ),
        migrations.RemoveField(
            model_name="experiment",
            name="branch",
        ),
        migrations.RemoveField(
            model_name="experiment",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="experiment",
            name="run",
        ),
        migrations.RemoveField(
            model_name="experiment",
            name="space",
        ),
        migrations.RemoveField(
            model_name="techsample",
            name="branch",
        ),
        migrations.RemoveField(
            model_name="techsample",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="techsample",
            name="run",
        ),
        migrations.RemoveField(
            model_name="techsample",
            name="space",
        ),
        migrations.AlterUniqueTogether(
            name="well",
            unique_together=None,
        ),
        migrations.RemoveField(
            model_name="well",
            name="branch",
        ),
        migrations.RemoveField(
            model_name="well",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="well",
            name="run",
        ),
        migrations.RemoveField(
            model_name="well",
            name="space",
        ),
        migrations.RemoveField(
            model_name="combinationperturbation",
            name="ontology_id",
        ),
        migrations.RemoveField(
            model_name="compoundperturbation",
            name="targets",
        ),
        migrations.AddField(
            model_name="biologic",
            name="source",
            field=lamindb.base.fields.ForeignKey(
                blank=True,
                null=True,
                on_delete=django.db.models.deletion.PROTECT,
                related_name="+",
                to="bionty.source",
            ),
        ),
        migrations.AddField(
            model_name="combinationperturbation",
            name="abbr",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=None,
                max_length=32,
                null=True,
                unique=True,
            ),
        ),
        migrations.AddField(
            model_name="combinationperturbation",
            name="source",
            field=lamindb.base.fields.ForeignKey(
                blank=True,
                null=True,
                on_delete=django.db.models.deletion.PROTECT,
                related_name="+",
                to="bionty.source",
            ),
        ),
        migrations.AddField(
            model_name="combinationperturbation",
            name="synonyms",
            field=lamindb.base.fields.TextField(
                blank=True, db_index=True, default=None, null=True
            ),
        ),
        migrations.AddField(
            model_name="compound",
            name="moa",
            field=lamindb.base.fields.TextField(
                blank=True, db_index=True, default=None, null=True
            ),
        ),
        migrations.AddField(
            model_name="compound",
            name="targets",
            field=models.ManyToManyField(
                related_name="compounds", to="wetlab.perturbationtarget"
            ),
        ),
        migrations.AddField(
            model_name="compound",
            name="type",
            field=lamindb.base.fields.CharField(
                blank=True, db_index=True, default=None, max_length=32, null=True
            ),
        ),
        migrations.AddField(
            model_name="environmentalperturbation",
            name="abbr",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=None,
                max_length=32,
                null=True,
                unique=True,
            ),
        ),
        migrations.AddField(
            model_name="environmentalperturbation",
            name="source",
            field=lamindb.base.fields.ForeignKey(
                blank=True,
                null=True,
                on_delete=django.db.models.deletion.PROTECT,
                related_name="+",
                to="bionty.source",
            ),
        ),
        migrations.AddField(
            model_name="environmentalperturbation",
            name="synonyms",
            field=lamindb.base.fields.TextField(
                blank=True, db_index=True, default=None, null=True
            ),
        ),
        migrations.AddField(
            model_name="geneticperturbation",
            name="abbr",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=None,
                max_length=32,
                null=True,
                unique=True,
            ),
        ),
        migrations.AddField(
            model_name="geneticperturbation",
            name="source",
            field=lamindb.base.fields.ForeignKey(
                blank=True,
                null=True,
                on_delete=django.db.models.deletion.PROTECT,
                related_name="+",
                to="bionty.source",
            ),
        ),
        migrations.AddField(
            model_name="geneticperturbation",
            name="synonyms",
            field=lamindb.base.fields.TextField(
                blank=True, db_index=True, default=None, null=True
            ),
        ),
        migrations.AddField(
            model_name="perturbationtarget",
            name="abbr",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=None,
                max_length=32,
                null=True,
                unique=True,
            ),
        ),
        migrations.AddField(
            model_name="perturbationtarget",
            name="source",
            field=lamindb.base.fields.ForeignKey(
                blank=True,
                null=True,
                on_delete=django.db.models.deletion.PROTECT,
                related_name="+",
                to="bionty.source",
            ),
        ),
        migrations.AddField(
            model_name="perturbationtarget",
            name="synonyms",
            field=lamindb.base.fields.TextField(
                blank=True, db_index=True, default=None, null=True
            ),
        ),
        migrations.AlterField(
            model_name="biologic",
            name="name",
            field=lamindb.base.fields.CharField(
                blank=True, db_index=True, default=None, max_length=255
            ),
        ),
        migrations.AlterField(
            model_name="compound",
            name="ontology_id",
            field=lamindb.base.fields.CharField(
                blank=True, db_index=True, default=None, max_length=255, null=True
            ),
        ),
        migrations.AlterField(
            model_name="geneticperturbation",
            name="sequence",
            field=models.TextField(db_index=True, null=True),
        ),
        migrations.AlterField(
            model_name="geneticperturbation",
            name="system",
            field=models.CharField(db_index=True, max_length=32, null=True),
        ),
        migrations.AlterField(
            model_name="perturbationtarget",
            name="name",
            field=lamindb.base.fields.CharField(
                blank=True, db_index=True, default=None, max_length=255
            ),
        ),
        migrations.RenameField(
            model_name="geneticperturbation",
            old_name="system",
            new_name="type",
        ),
        migrations.AlterUniqueTogether(
            name="compound",
            unique_together={("name", "ontology_id")},
        ),
        migrations.DeleteModel(
            name="ArtifactBiosample",
        ),
        migrations.DeleteModel(
            name="ArtifactDonor",
        ),
        migrations.DeleteModel(
            name="ArtifactExperiment",
        ),
        migrations.DeleteModel(
            name="ArtifactTechsample",
        ),
        migrations.DeleteModel(
            name="ArtifactWell",
        ),
        migrations.DeleteModel(
            name="Biosample",
        ),
        migrations.DeleteModel(
            name="Donor",
        ),
        migrations.DeleteModel(
            name="Experiment",
        ),
        migrations.DeleteModel(
            name="Techsample",
        ),
        migrations.DeleteModel(
            name="Well",
        ),
    ]
