"""Make molecules unique

Revision ID: c0fba229aaa8
Revises: 48d1c43c27e0
Create Date: 2022-12-12 09:39:10.950668

"""

import sqlalchemy as sa
from alembic import op
from sqlalchemy.orm.session import Session

# revision identifiers, used by Alembic.
revision = "c0fba229aaa8"
down_revision = "48d1c43c27e0"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###

    # Find all duplicate molecules
    session = Session(op.get_bind())

    duplicates = session.execute(
        sa.text(
            """
        SELECT molecule_hash, array_agg(id), count(*)
        FROM molecule
        GROUP BY molecule_hash
        HAVING count(*) > 1
    """
        )
    )

    for mol_hash, mol_ids, _ in duplicates:
        mol_ids = sorted(mol_ids)

        # Keep the first id, merge the rest
        keep_id = mol_ids[0]

        dup_ids = f"({','.join(str(x) for x in mol_ids[1:])})"
        op.execute(
            f"UPDATE gridoptimization_dataset_entry  SET initial_molecule_id = {keep_id}  WHERE initial_molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE gridoptimization_record         SET initial_molecule_id = {keep_id}  WHERE initial_molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE gridoptimization_record         SET starting_molecule_id = {keep_id} WHERE starting_molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE manybody_cluster                SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE manybody_dataset_entry          SET initial_molecule_id = {keep_id}  WHERE initial_molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE manybody_record                 SET initial_molecule_id = {keep_id}  WHERE initial_molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE neb_dataset_molecule            SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE neb_initialchain                SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE optimization_dataset_entry      SET initial_molecule_id = {keep_id}  WHERE initial_molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE optimization_record             SET final_molecule_id = {keep_id}    WHERE final_molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE optimization_record             SET initial_molecule_id = {keep_id}  WHERE initial_molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE reaction_component              SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE reaction_dataset_stoichiometry  SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE singlepoint_dataset_entry       SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE singlepoint_record              SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE torsiondrive_dataset_molecule   SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )
        op.execute(
            f"UPDATE torsiondrive_initial_molecule   SET molecule_id = {keep_id}          WHERE molecule_id IN {dup_ids}"
        )

        op.execute(f"DELETE FROM molecule WHERE id IN {dup_ids}")

    # Officially no more duplicate molecules
    op.drop_index("ix_molecule_molecule_hash", table_name="molecule")
    op.create_unique_constraint("ux_molecule_molecule_hash", "molecule", ["molecule_hash"])
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint("ux_molecule_molecule_hash", "molecule", type_="unique")
    op.create_index("ix_molecule_molecule_hash", "molecule", ["molecule_hash"], unique=False)
    # ### end Alembic commands ###
