"""rename and refactor optimization

Revision ID: cdd9f3dc10e3
Revises: 160352419195
Create Date: 2021-12-08 14:53:27.286030

"""

import os
import sys

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

sys.path.insert(1, os.path.dirname(os.path.abspath(__file__)))
from migration_helpers.v0_50_helpers import get_empty_keywords_id

# revision identifiers, used by Alembic.
revision = "cdd9f3dc10e3"
down_revision = "160352419195"
branch_labels = None
depends_on = None


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

    # The optimization specification table
    op.create_table(
        "optimization_specification",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("program", sa.String(length=100), nullable=False),
        sa.Column("qc_specification_id", sa.Integer(), nullable=False),
        sa.Column("keywords", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
        sa.Column("protocols", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
        sa.CheckConstraint("program = LOWER(program)", name="ck_optimization_specification_program_lower"),
        sa.ForeignKeyConstraint(
            ["qc_specification_id"],
            ["qc_specification.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "program",
            "qc_specification_id",
            "keywords",
            "protocols",
            name="ux_optimization_specification_keys",
        ),
    )
    op.create_index("ix_optimization_specification_keywords", "optimization_specification", ["keywords"], unique=False)
    op.create_index("ix_optimization_specification_program", "optimization_specification", ["program"], unique=False)
    op.create_index(
        "ix_optimization_specification_protocols", "optimization_specification", ["protocols"], unique=False
    )
    op.create_index(
        "ix_optimization_specification_qc_specification_id",
        "optimization_specification",
        ["qc_specification_id"],
        unique=False,
    )

    # Optimization <-> Singlepoint association (for Trajectory)
    op.drop_constraint("opt_result_association_result_id_fkey", "opt_result_association", type_="foreignkey")
    op.drop_constraint("opt_result_association_opt_id_fkey", "opt_result_association", type_="foreignkey")
    op.alter_column("opt_result_association", "result_id", new_column_name="singlepoint_id")
    op.alter_column("opt_result_association", "opt_id", new_column_name="optimization_id")
    op.create_foreign_key(
        "optimization_trajectory_optimization_id_fkey",
        "opt_result_association",
        "optimization_procedure",
        ["optimization_id"],
        ["id"],
        ondelete="cascade",
    )
    op.create_foreign_key(
        "optimization_trajectory_singlepoint_id_fkey",
        "opt_result_association",
        "singlepoint_record",
        ["singlepoint_id"],
        ["id"],
    )

    # Now the optimization table
    op.execute(
        sa.text(
            "ALTER TABLE optimization_procedure RENAME CONSTRAINT optimization_procedure_id_fkey TO optimization_record_id_fkey"
        )
    )
    op.execute(sa.text("ALTER INDEX optimization_procedure_pkey RENAME TO optimization_record_pkey"))

    op.add_column("optimization_procedure", sa.Column("specification_id", sa.Integer(), nullable=True))
    op.create_foreign_key(
        "optimization_record_specification_id_fkey",
        "optimization_procedure",
        "optimization_specification",
        ["specification_id"],
        ["id"],
    )

    op.drop_constraint("optimization_procedure_final_molecule_fkey", "optimization_procedure", type_="foreignkey")
    op.drop_constraint("optimization_procedure_initial_molecule_fkey", "optimization_procedure", type_="foreignkey")
    op.alter_column("optimization_procedure", "initial_molecule", new_column_name="initial_molecule_id")
    op.alter_column("optimization_procedure", "final_molecule", new_column_name="final_molecule_id")
    op.create_foreign_key(
        "optimization_record_initial_molecule_id_fkey",
        "optimization_procedure",
        "molecule",
        ["initial_molecule_id"],
        ["id"],
    )
    op.create_foreign_key(
        "optimization_record_final_molecule_id_fkey",
        "optimization_procedure",
        "molecule",
        ["final_molecule_id"],
        ["id"],
    )

    op.create_index(
        "ix_optimization_record_initial_molecule_id", "optimization_procedure", ["initial_molecule_id"], unique=False
    )
    op.create_index(
        "ix_optimization_record_final_molecule_id", "optimization_procedure", ["final_molecule_id"], unique=False
    )
    op.create_index(
        "ix_optimization_record_specification_id", "optimization_procedure", ["specification_id"], unique=False
    )

    # Migrate data to the specification tables
    # First, qcspec -> singlepoint specification

    # Now we can go with qc_spec
    empty_kw_id = get_empty_keywords_id(op.get_bind())

    # remove empty strings an json null
    op.execute(
        sa.text(
            r"UPDATE optimization_procedure SET qc_spec = (qc_spec::jsonb - 'keywords')::json WHERE qc_spec::jsonb->>'keywords' = ''"
        )
    )
    op.execute(
        sa.text(
            r"UPDATE optimization_procedure SET qc_spec = (qc_spec::jsonb - 'keywords')::json WHERE qc_spec::jsonb->>'keywords' = 'null'"
        )
    )
    op.execute(
        sa.text(
            r"UPDATE optimization_procedure SET qc_spec = (qc_spec::jsonb - 'keywords')::json WHERE qc_spec::jsonb->>'keywords' IS NULL"
        )
    )

    # Remove program from keywords
    op.execute(sa.text(r"UPDATE optimization_procedure SET keywords = (keywords::jsonb - 'program')::json"))

    # Protocols for qc_spec were always ignored. So set them with the default
    op.execute(
        sa.text(
            f"""
               INSERT INTO qc_specification (program, driver, method, basis, keywords_id, protocols)
               SELECT DISTINCT o.qc_spec->>'program',
                               'deferred'::singlepointdriver,
                               o.qc_spec->>'method',
                               COALESCE(o.qc_spec->>'basis', ''),
                               COALESCE((o.qc_spec->>'keywords')::int, {empty_kw_id}),
                               '{{}}'::jsonb
               FROM optimization_procedure o
               ON CONFLICT DO NOTHING
               """
        )
    )

    # Populate the optimization specification table
    op.execute(
        sa.text(
            r"""UPDATE base_record SET protocols = '{}'::json WHERE protocols::jsonb = '{"trajectory": "all"}'::jsonb"""
        )
    )

    op.execute(
        sa.text(
            f"""
               INSERT INTO optimization_specification (program, keywords, protocols, qc_specification_id)
               SELECT DISTINCT o.program,
                               o.keywords::jsonb,
                               br.protocols::jsonb,
                               (
                               SELECT id from qc_specification sp
                               WHERE sp.program = o.qc_spec->>'program'
                               AND sp.driver = 'deferred'::singlepointdriver
                               AND sp.method = o.qc_spec->>'method'
                               AND sp.basis = COALESCE(o.qc_spec->>'basis', '')
                               AND sp.keywords_id = COALESCE((o.qc_spec->>'keywords')::int, {empty_kw_id}) 
                               AND sp.protocols = '{{}}'::jsonb
                               )
               FROM optimization_procedure o INNER JOIN base_record br on br.id = o.id
               """
        )
    )

    # Now this mess
    op.execute(
        sa.text(
            f"""
               UPDATE optimization_procedure o
               SET specification_id = ss.id
               FROM optimization_specification ss, base_record as br
               WHERE o.id = br.id
               AND o.program = ss.program
               AND o.keywords::jsonb = ss.keywords
               AND br.protocols::jsonb = ss.protocols
               AND ss.qc_specification_id = (
                               SELECT id from qc_specification sp
                               WHERE sp.program = o.qc_spec->>'program'
                               AND sp.driver = 'deferred'::singlepointdriver
                               AND sp.method = o.qc_spec->>'method'
                               AND sp.basis = COALESCE(o.qc_spec->>'basis', '')
                               AND sp.keywords_id = COALESCE((o.qc_spec->>'keywords')::int, {empty_kw_id}) 
                               AND sp.protocols = '{{}}'::jsonb
                               )
               """
        )
    )

    op.drop_index("ix_optimization_program", table_name="optimization_procedure")
    op.drop_column("optimization_procedure", "qc_spec")
    op.drop_column("optimization_procedure", "keywords")
    op.drop_column("optimization_procedure", "schema_version")
    op.drop_column("optimization_procedure", "program")

    # Columns not nullable
    op.alter_column("optimization_procedure", "specification_id", nullable=False)

    # Rename the tables
    op.rename_table("opt_result_association", "optimization_trajectory")

    # Drop and recreate the pkey, since column names changed
    op.drop_constraint("opt_result_association_pkey", "optimization_trajectory", type_="primary")
    op.execute(
        sa.text("ALTER TABLE optimization_trajectory ADD PRIMARY KEY (optimization_id, singlepoint_id, position)")
    )

    op.rename_table("optimization_procedure", "optimization_record")

    # We can now finally remove the protocols from base_record
    op.drop_column("base_record", "protocols")

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    raise RuntimeError("Cannot downgrade")
