"""rankify_predictions

Revision ID: 609c7cc51794
Revises: 50e1f1bc2cac
Create Date: 2019-04-03 16:23:22.903979

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '609c7cc51794'
down_revision = '97cf99b7348f'
branch_labels = None
depends_on = None


def add_rank_abs_with_ties_to_predictions(schema, op):
    """Add the new column rank_abs_with_ties directly after rank_abs

    Altering a table to add a column in Postgres can only put columns at the end of the table.
    This is not ideal. So we recreate the table and copy everything over.
    """
    op.execute(f"alter table {schema}.predictions rename to predictions_old")
    op.create_table(
        "predictions",
        sa.Column("model_id", sa.Integer(), nullable=False),
        sa.Column("entity_id", sa.BigInteger(), nullable=False),
        sa.Column("as_of_date", sa.DateTime(), nullable=False),
        sa.Column("score", sa.Numeric(6, 5), nullable=True),
        sa.Column("label_value", sa.Integer(), nullable=True),
        sa.Column("rank_abs_no_ties", sa.Integer(), nullable=True),
        sa.Column("rank_abs_with_ties", sa.Integer(), nullable=True),
        sa.Column("rank_pct_no_ties", sa.Numeric(6, 5), nullable=True),
        sa.Column("rank_pct_with_ties", sa.Numeric(6, 5), nullable=True),
        sa.Column("matrix_uuid", sa.Text(), nullable=True),
        sa.Column("test_label_timespan", sa.Interval(), nullable=True),
        schema=schema,
    )
    op.execute(f"""
insert into {schema}.predictions
select model_id, entity_id, as_of_date, score, label_value, rank_abs, null, rank_pct, null, matrix_uuid, test_label_timespan from {schema}.predictions_old
""")
    op.drop_table("predictions_old", schema=schema)
    op.create_primary_key(
        f"pk_{schema}_predictions",
        "predictions",
        schema=schema,
        columns=['model_id', 'entity_id', 'as_of_date']
    )
    op.create_foreign_key(
        constraint_name=f"{schema}_predictions_matrix_uuid_fkey",
        source_table="predictions",
        source_schema=schema,
        referent_table="matrices",
        referent_schema="model_metadata",
        local_cols=["matrix_uuid"],
        remote_cols=["matrix_uuid"],
    )
    op.create_foreign_key(
        constraint_name=f"{schema}_predictions_model_id_fkey",
        source_table="predictions",
        source_schema=schema,
        referent_table="models",
        referent_schema="model_metadata",
        local_cols=["model_id"],
        remote_cols=["model_id"],
    )

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('prediction_metadata',
    sa.Column('model_id', sa.Integer(), nullable=False),
    sa.Column('matrix_uuid', sa.Text(), nullable=False),
    sa.Column('tiebreaker_ordering', sa.Text(), nullable=True),
    sa.Column('random_seed', sa.Integer(), nullable=True),
    sa.Column('predictions_saved', sa.Boolean(), nullable=True),
    sa.ForeignKeyConstraint(['matrix_uuid'], ['model_metadata.matrices.matrix_uuid'], ),
    sa.ForeignKeyConstraint(['model_id'], ['model_metadata.models.model_id'], ),
    sa.PrimaryKeyConstraint('model_id', 'matrix_uuid'),
    schema='test_results'
    )
    op.create_table('prediction_metadata',
    sa.Column('model_id', sa.Integer(), nullable=False),
    sa.Column('matrix_uuid', sa.Text(), nullable=False),
    sa.Column('tiebreaker_ordering', sa.Text(), nullable=True),
    sa.Column('random_seed', sa.Integer(), nullable=True),
    sa.Column('predictions_saved', sa.Boolean(), nullable=True),
    sa.ForeignKeyConstraint(['matrix_uuid'], ['model_metadata.matrices.matrix_uuid'], ),
    sa.ForeignKeyConstraint(['model_id'], ['model_metadata.models.model_id'], ),
    sa.PrimaryKeyConstraint('model_id', 'matrix_uuid'),
    schema='train_results'
    )
    add_rank_abs_with_ties_to_predictions("train_results", op)
    add_rank_abs_with_ties_to_predictions("test_results", op)
    op.add_column('models', sa.Column('random_seed', sa.Integer(), nullable=True), schema='model_metadata')
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('models', 'random_seed', schema='model_metadata')
    op.drop_column('predictions', 'rank_abs_with_ties', schema='train_results')
    op.drop_column('predictions', 'rank_abs_with_ties', schema='test_results')
    op.drop_column('predictions', 'rank_pct_with_ties', schema='train_results')
    op.drop_column('predictions', 'rank_pct_with_ties', schema='test_results')
    op.alter_column('predictions', 'rank_abs_no_ties', new_column_name='rank_abs', schema='train_results')
    op.alter_column('predictions', 'rank_pct_no_ties', new_column_name='rank_pct', schema='train_results')
    op.alter_column('predictions', 'rank_abs_no_ties', new_column_name='rank_abs', schema='test_results')
    op.alter_column('predictions', 'rank_pct_no_ties', new_column_name='rank_pct', schema='test_results')
    op.drop_table('prediction_metadata', schema='train_results')
    op.drop_table('prediction_metadata', schema='test_results')
    # ### end Alembic commands ###
