"""Add contributed values table

Revision ID: 159ba85908fd
Revises: d5d88ac1d291
Create Date: 2019-11-01 15:39:50.970246

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy import orm
from sqlalchemy.dialects import postgresql
from qcfractal.storage_sockets.models.sql_base import MsgpackExt
from qcfractal.storage_sockets.models.collections_models import ContributedValuesORM

import numpy as np

# revision identifiers, used by Alembic.
revision = "159ba85908fd"
down_revision = "d5d88ac1d291"
branch_labels = None
depends_on = None


def migrate_contributed_values_data():

    bind = op.get_bind()
    session = orm.Session(bind=bind)

    # Dataset and reaction datasets tables
    ds_ids_data = session.execute("select id, contributed_values_data from dataset;").fetchall()
    print(f"Migrating datasets with ids: {[ds[0] for ds in ds_ids_data]}")

    rds_ids_data = session.execute("select id, contributed_values_data from reaction_dataset;").fetchall()
    print(f"Migrating reaction datasets with ids: {[ds[0] for ds in rds_ids_data]}")

    ds_ids_data.extend(rds_ids_data)

    for ds in ds_ids_data:
        (ds_id, ds_contrib) = ds
        if ds_contrib is None:
            continue

        for key, dict_values in ds_contrib.items():

            idx, vals = [], []
            for key, value in dict_values["values"].items():
                idx.append(key)
                vals.append(value)

            dict_values["values"] = np.array(vals)
            dict_values["index"] = np.array(idx)

            cv = ContributedValuesORM(**dict_values)
            cv.collection_id = ds_id

            session.add(cv)

    session.commit()


def upgrade():

    # rename old column with data
    op.alter_column("dataset", "contributed_values", new_column_name="contributed_values_data")
    op.alter_column("reaction_dataset", "contributed_values", new_column_name="contributed_values_data")

    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "contributed_values",
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("collection_id", sa.Integer(), nullable=False),
        sa.Column("citations", sa.JSON(), nullable=True),
        sa.Column("theory_level", sa.JSON(), nullable=False),
        sa.Column("theory_level_details", sa.JSON(), nullable=True),
        sa.Column("comments", sa.String(), nullable=True),
        sa.Column("values", MsgpackExt(), nullable=False),
        sa.Column("index", MsgpackExt(), nullable=False),
        sa.Column("external_url", sa.String(), nullable=True),
        sa.Column("doi", sa.String(), nullable=True),
        sa.Column("units", sa.String(), nullable=False),
        sa.Column("values_structure", sa.JSON(), nullable=True, default=lambda: {}),
        sa.ForeignKeyConstraint(["collection_id"], ["collection.id"], ondelete="cascade"),
        sa.PrimaryKeyConstraint("name", "collection_id"),
    )

    op.alter_column("contributed_values", "values_structure", server_default=None, nullable=False)

    migrate_contributed_values_data()

    op.drop_column("dataset", "contributed_values_data")
    op.drop_column("reaction_dataset", "contributed_values_data")


def downgrade():
    # ### Won't work on production data because data will be lost ###

    op.add_column(
        "reaction_dataset",
        sa.Column("contributed_values", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True),
    )
    op.add_column(
        "dataset",
        sa.Column("contributed_values", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True),
    )
    op.drop_table("contributed_values")
