"""evaluations

Revision ID: f6528d04bbbd
Revises: 4524fb60d23e
Create Date: 2024-11-19 19:31:38.381105+00:00

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import sqlmodel
import ell.stores.models.core


# revision identifiers, used by Alembic.
revision: str = 'f6528d04bbbd'
down_revision: Union[str, None] = '4524fb60d23e'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('serializedevaluation',
    sa.Column('id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('name', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('created_at', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=False),
    sa.Column('dataset_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('n_evals', sa.Integer(), nullable=False),
    sa.Column('version_number', sa.Integer(), nullable=False),
    sa.Column('commit_message', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('evaluationlabeler',
    sa.Column('id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('name', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('type', sa.Enum('METRIC', 'ANNOTATION', 'CRITERION', name='evaluationlabelertype'), nullable=False),
    sa.Column('labeling_lmp_id', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
    sa.Column('evaluation_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('labeling_rubric', sa.JSON(), nullable=True),
    sa.ForeignKeyConstraint(['evaluation_id'], ['serializedevaluation.id'], ),
    sa.ForeignKeyConstraint(['labeling_lmp_id'], ['serializedlmp.lmp_id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_index(op.f('ix_evaluationlabeler_labeling_lmp_id'), 'evaluationlabeler', ['labeling_lmp_id'], unique=False)
    op.create_table('serializedevaluationrun',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('evaluation_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('evaluated_lmp_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('api_params', sa.JSON(), nullable=True),
    sa.Column('start_time', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
    sa.Column('end_time', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
    sa.Column('success', sa.Boolean(), nullable=True),
    sa.Column('error', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
    sa.ForeignKeyConstraint(['evaluated_lmp_id'], ['serializedlmp.lmp_id'], ),
    sa.ForeignKeyConstraint(['evaluation_id'], ['serializedevaluation.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_index(op.f('ix_serializedevaluationrun_evaluated_lmp_id'), 'serializedevaluationrun', ['evaluated_lmp_id'], unique=False)
    op.create_index(op.f('ix_serializedevaluationrun_evaluation_id'), 'serializedevaluationrun', ['evaluation_id'], unique=False)
    op.create_table('evaluationresultdatapoint',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('invocation_being_labeled_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('evaluation_run_id', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['evaluation_run_id'], ['serializedevaluationrun.id'], ),
    sa.ForeignKeyConstraint(['invocation_being_labeled_id'], ['invocation.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('evaluationrunlabelersummary',
    sa.Column('evaluation_labeler_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('evaluation_run_id', sa.Integer(), nullable=False),
    sa.Column('created_at', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
    sa.Column('updated_at', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
    sa.Column('finalized_at', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
    sa.Column('is_scalar', sa.Boolean(), nullable=False),
    sa.Column('data', sa.JSON(), nullable=True),
    sa.Column('count', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['evaluation_labeler_id'], ['evaluationlabeler.id'], ),
    sa.ForeignKeyConstraint(['evaluation_run_id'], ['serializedevaluationrun.id'], ),
    sa.PrimaryKeyConstraint('evaluation_labeler_id', 'evaluation_run_id')
    )
    op.create_table('evaluationlabel',
    sa.Column('labeler_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.Column('labeled_datapoint_id', sa.Integer(), nullable=False),
    sa.Column('label_invocation_id', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
    sa.Column('manual_label', sa.JSON(), nullable=True),
    sa.ForeignKeyConstraint(['label_invocation_id'], ['invocation.id'], ),
    sa.ForeignKeyConstraint(['labeled_datapoint_id'], ['evaluationresultdatapoint.id'], ),
    sa.ForeignKeyConstraint(['labeler_id'], ['evaluationlabeler.id'], ),
    sa.PrimaryKeyConstraint('labeler_id', 'labeled_datapoint_id')
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('evaluationlabel')
    op.drop_table('evaluationrunlabelersummary')
    op.drop_table('evaluationresultdatapoint')
    op.drop_index(op.f('ix_serializedevaluationrun_evaluation_id'), table_name='serializedevaluationrun')
    op.drop_index(op.f('ix_serializedevaluationrun_evaluated_lmp_id'), table_name='serializedevaluationrun')
    op.drop_table('serializedevaluationrun')
    op.drop_index(op.f('ix_evaluationlabeler_labeling_lmp_id'), table_name='evaluationlabeler')
    op.drop_table('evaluationlabeler')
    op.drop_table('serializedevaluation')
    # ### end Alembic commands ###
