"""add quality metric and training tables

Revision ID: 33107394dcd1
Revises: 003
Create Date: 2025-10-07 17:48:09.877745

"""

import sqlalchemy as sa
import sqlmodel
from alembic import op

# revision identifiers, used by Alembic.
revision = "33107394dcd1"
down_revision = "003"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "qualitymetric",
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("description", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.PrimaryKeyConstraint("name"),
    )
    op.create_table(
        "currentqualitypredictionmodelweight",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("grid_uuid", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("prediction_model_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("metric_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("weight", sa.Float(), nullable=False),
        sa.ForeignKeyConstraint(
            ["grid_uuid"],
            ["grid.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["metric_name"],
            ["qualitymetric.name"],
        ),
        sa.ForeignKeyConstraint(
            ["prediction_model_name"],
            ["qualitypredictionmodel.name"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "qualitymetricstatistics",
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("grid_uuid", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("count", sa.Integer(), nullable=False),
        sa.Column("value_sum", sa.Float(), nullable=False),
        sa.Column("squared_value_sum", sa.Float(), nullable=False),
        sa.ForeignKeyConstraint(
            ["grid_uuid"],
            ["grid.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["name"],
            ["qualitymetric.name"],
        ),
        sa.PrimaryKeyConstraint("name", "grid_uuid"),
    )
    op.create_table(
        "currentqualityprediction",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("grid_uuid", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("value", sa.Float(), nullable=False),
        sa.Column("prediction_model_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("metric_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("foilhole_uuid", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("gridsquare_uuid", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.ForeignKeyConstraint(
            ["foilhole_uuid"],
            ["foilhole.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["grid_uuid"],
            ["grid.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["gridsquare_uuid"],
            ["gridsquare.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["metric_name"],
            ["qualitymetric.name"],
        ),
        sa.ForeignKeyConstraint(
            ["prediction_model_name"],
            ["qualitypredictionmodel.name"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "overallqualityprediction",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("value", sa.Float(), nullable=False),
        sa.Column("foilhole_uuid", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("grid_uuid", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("gridsquare_uuid", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.ForeignKeyConstraint(
            ["foilhole_uuid"],
            ["foilhole.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["grid_uuid"],
            ["grid.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["gridsquare_uuid"],
            ["gridsquare.uuid"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.drop_index(op.f("idx_acquisition_id_hash"), table_name="acquisition", if_exists=True, postgresql_using="hash")
    op.drop_index(op.f("idx_acquisition_id_pattern"), table_name="acquisition", if_exists=True)
    op.drop_index(op.f("idx_acquisition_name"), table_name="acquisition", if_exists=True)
    op.drop_index(op.f("idx_agent_connection_agent_id"), table_name="agentconnection", if_exists=True)
    op.drop_index(op.f("idx_agent_connection_created_at"), table_name="agentconnection", if_exists=True)
    op.drop_index(op.f("idx_agent_connection_heartbeat"), table_name="agentconnection", if_exists=True)
    op.drop_index(op.f("idx_agent_connection_session_id"), table_name="agentconnection", if_exists=True)
    op.drop_index(op.f("idx_agent_connection_status"), table_name="agentconnection", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_acknowledged_at"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_agent_id"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_created_at"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_expires_at"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_priority"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_sent_at"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_sequence"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_session_id"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_status"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_instruction_type"), table_name="agentinstruction", if_exists=True)
    op.drop_index(op.f("idx_agent_ack_agent_id"), table_name="agentinstructionacknowledgement", if_exists=True)
    op.drop_index(op.f("idx_agent_ack_created_at"), table_name="agentinstructionacknowledgement", if_exists=True)
    op.drop_index(op.f("idx_agent_ack_instruction_id"), table_name="agentinstructionacknowledgement", if_exists=True)
    op.drop_index(op.f("idx_agent_ack_processed_at"), table_name="agentinstructionacknowledgement", if_exists=True)
    op.drop_index(op.f("idx_agent_ack_session_id"), table_name="agentinstructionacknowledgement", if_exists=True)
    op.drop_index(op.f("idx_agent_ack_status"), table_name="agentinstructionacknowledgement", if_exists=True)
    op.drop_index(op.f("idx_agent_session_acquisition_uuid"), table_name="agentsession", if_exists=True)
    op.drop_index(op.f("idx_agent_session_agent_id"), table_name="agentsession", if_exists=True)
    op.drop_index(op.f("idx_agent_session_created_at"), table_name="agentsession", if_exists=True)
    op.drop_index(op.f("idx_agent_session_last_activity"), table_name="agentsession", if_exists=True)
    op.drop_index(op.f("idx_agent_session_status"), table_name="agentsession", if_exists=True)
    op.drop_index(op.f("idx_atlas_grid_uuid"), table_name="atlas", if_exists=True)
    op.drop_index(op.f("idx_atlas_id_pattern"), table_name="atlas", if_exists=True)
    op.drop_index(op.f("idx_atlastile_atlas_id"), table_name="atlastile", if_exists=True)
    op.drop_index(op.f("idx_atlastile_id_pattern"), table_name="atlastile", if_exists=True)
    op.drop_index(op.f("idx_foilhole_gridsquare_id"), table_name="foilhole", if_exists=True)
    op.drop_index(op.f("idx_foilhole_id_hash"), table_name="foilhole", if_exists=True, postgresql_using="hash")
    op.drop_index(op.f("idx_foilhole_id_pattern"), table_name="foilhole", if_exists=True)
    op.drop_index(op.f("idx_foilhole_quality"), table_name="foilhole", if_exists=True)
    op.add_column("grid", sa.Column("prediction_updated_time", sa.DateTime(), nullable=False))
    op.drop_index(op.f("idx_grid_acquisition_id"), table_name="grid", if_exists=True)
    op.drop_index(op.f("idx_grid_name"), table_name="grid", if_exists=True)
    op.drop_index(op.f("idx_grid_uuid_hash"), table_name="grid", if_exists=True, postgresql_using="hash")
    op.drop_index(op.f("idx_grid_uuid_pattern"), table_name="grid", if_exists=True)
    op.drop_index(op.f("idx_gridsquare_grid_uuid"), table_name="gridsquare", if_exists=True)
    op.drop_index(op.f("idx_gridsquare_id_hash"), table_name="gridsquare", if_exists=True, postgresql_using="hash")
    op.drop_index(op.f("idx_gridsquare_id_pattern"), table_name="gridsquare", if_exists=True)
    op.drop_index(op.f("idx_micrograph_foilhole_id"), table_name="micrograph", if_exists=True)
    op.drop_index(op.f("idx_micrograph_id_hash"), table_name="micrograph", if_exists=True, postgresql_using="hash")
    op.drop_index(op.f("idx_micrograph_id_pattern"), table_name="micrograph", if_exists=True)
    op.add_column("qualityprediction", sa.Column("metric_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
    op.drop_index(op.f("idx_quality_prediction_foilhole_uuid"), table_name="qualityprediction", if_exists=True)
    op.drop_index(op.f("idx_quality_prediction_gridsquare_uuid"), table_name="qualityprediction", if_exists=True)
    op.drop_index(op.f("idx_quality_prediction_model_foilhole"), table_name="qualityprediction", if_exists=True)
    op.drop_index(op.f("idx_quality_prediction_model_gridsquare"), table_name="qualityprediction", if_exists=True)
    op.drop_index(op.f("idx_quality_prediction_model_name"), table_name="qualityprediction", if_exists=True)
    op.drop_index(op.f("idx_quality_prediction_timestamp"), table_name="qualityprediction", if_exists=True)
    op.drop_index(op.f("idx_quality_prediction_value"), table_name="qualityprediction", if_exists=True)
    op.create_foreign_key(None, "qualityprediction", "qualitymetric", ["metric_name"], ["name"])
    op.add_column("qualitypredictionmodel", sa.Column("level", sa.VARCHAR(), nullable=True))
    op.add_column(
        "qualitypredictionmodelparameter", sa.Column("metric_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True)
    )
    op.drop_index(
        op.f("idx_quality_model_param_composite"), table_name="qualitypredictionmodelparameter", if_exists=True
    )
    op.drop_index(
        op.f("idx_quality_model_param_grid_uuid"), table_name="qualitypredictionmodelparameter", if_exists=True
    )
    op.drop_index(op.f("idx_quality_model_param_key"), table_name="qualitypredictionmodelparameter", if_exists=True)
    op.drop_index(
        op.f("idx_quality_model_param_model_name"), table_name="qualitypredictionmodelparameter", if_exists=True
    )
    op.drop_index(
        op.f("idx_quality_model_param_timestamp"), table_name="qualitypredictionmodelparameter", if_exists=True
    )
    op.create_foreign_key(None, "qualitypredictionmodelparameter", "qualitymetric", ["metric_name"], ["name"])
    op.add_column(
        "qualitypredictionmodelweight", sa.Column("metric_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True)
    )
    op.drop_index(
        op.f("idx_quality_model_weight_grid_model"), table_name="qualitypredictionmodelweight", if_exists=True
    )
    op.drop_index(op.f("idx_quality_model_weight_grid_uuid"), table_name="qualitypredictionmodelweight", if_exists=True)
    op.drop_index(
        op.f("idx_quality_model_weight_micrograph_model"), table_name="qualitypredictionmodelweight", if_exists=True
    )
    op.drop_index(
        op.f("idx_quality_model_weight_micrograph_uuid"), table_name="qualitypredictionmodelweight", if_exists=True
    )
    op.drop_index(
        op.f("idx_quality_model_weight_model_name"), table_name="qualitypredictionmodelweight", if_exists=True
    )
    op.drop_index(op.f("idx_quality_model_weight_quality"), table_name="qualitypredictionmodelweight", if_exists=True)
    op.drop_index(op.f("idx_quality_model_weight_timestamp"), table_name="qualitypredictionmodelweight", if_exists=True)
    op.drop_index(op.f("idx_quality_model_weight_value"), table_name="qualitypredictionmodelweight", if_exists=True)
    op.create_foreign_key(None, "qualitypredictionmodelweight", "qualitymetric", ["metric_name"], ["name"])
    op.drop_column("qualitypredictionmodelweight", "origin")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("qualitypredictionmodelweight", sa.Column("origin", sa.VARCHAR(), autoincrement=False, nullable=True))
    op.drop_constraint(None, "qualitypredictionmodelweight", type_="foreignkey")
    op.create_index(op.f("idx_quality_model_weight_value"), "qualitypredictionmodelweight", ["weight"], unique=False)
    op.create_index(
        op.f("idx_quality_model_weight_timestamp"), "qualitypredictionmodelweight", ["timestamp"], unique=False
    )
    op.create_index(
        op.f("idx_quality_model_weight_quality"), "qualitypredictionmodelweight", ["micrograph_quality"], unique=False
    )
    op.create_index(
        op.f("idx_quality_model_weight_model_name"),
        "qualitypredictionmodelweight",
        ["prediction_model_name"],
        unique=False,
    )
    op.create_index(
        op.f("idx_quality_model_weight_micrograph_uuid"),
        "qualitypredictionmodelweight",
        ["micrograph_uuid"],
        unique=False,
    )
    op.create_index(
        op.f("idx_quality_model_weight_micrograph_model"),
        "qualitypredictionmodelweight",
        ["micrograph_uuid", "prediction_model_name"],
        unique=False,
    )
    op.create_index(
        op.f("idx_quality_model_weight_grid_uuid"), "qualitypredictionmodelweight", ["grid_uuid"], unique=False
    )
    op.create_index(
        op.f("idx_quality_model_weight_grid_model"),
        "qualitypredictionmodelweight",
        ["grid_uuid", "prediction_model_name"],
        unique=False,
    )
    op.drop_column("qualitypredictionmodelweight", "metric_name")
    op.drop_constraint(None, "qualitypredictionmodelparameter", type_="foreignkey")
    op.create_index(
        op.f("idx_quality_model_param_timestamp"), "qualitypredictionmodelparameter", ["timestamp"], unique=False
    )
    op.create_index(
        op.f("idx_quality_model_param_model_name"),
        "qualitypredictionmodelparameter",
        ["prediction_model_name"],
        unique=False,
    )
    op.create_index(op.f("idx_quality_model_param_key"), "qualitypredictionmodelparameter", ["key"], unique=False)
    op.create_index(
        op.f("idx_quality_model_param_grid_uuid"), "qualitypredictionmodelparameter", ["grid_uuid"], unique=False
    )
    op.create_index(
        op.f("idx_quality_model_param_composite"),
        "qualitypredictionmodelparameter",
        ["grid_uuid", "prediction_model_name"],
        unique=False,
    )
    op.drop_column("qualitypredictionmodelparameter", "metric_name")
    op.drop_column("qualitypredictionmodel", "level")
    op.drop_constraint(None, "qualityprediction", type_="foreignkey")
    op.create_index(op.f("idx_quality_prediction_value"), "qualityprediction", ["value"], unique=False)
    op.create_index(op.f("idx_quality_prediction_timestamp"), "qualityprediction", ["timestamp"], unique=False)
    op.create_index(
        op.f("idx_quality_prediction_model_name"), "qualityprediction", ["prediction_model_name"], unique=False
    )
    op.create_index(
        op.f("idx_quality_prediction_model_gridsquare"),
        "qualityprediction",
        ["prediction_model_name", "gridsquare_uuid"],
        unique=False,
    )
    op.create_index(
        op.f("idx_quality_prediction_model_foilhole"),
        "qualityprediction",
        ["prediction_model_name", "foilhole_uuid"],
        unique=False,
    )
    op.create_index(
        op.f("idx_quality_prediction_gridsquare_uuid"), "qualityprediction", ["gridsquare_uuid"], unique=False
    )
    op.create_index(op.f("idx_quality_prediction_foilhole_uuid"), "qualityprediction", ["foilhole_uuid"], unique=False)
    op.drop_column("qualityprediction", "metric_name")
    op.create_index(op.f("idx_micrograph_id_pattern"), "micrograph", ["uuid"], unique=False)
    op.create_index(op.f("idx_micrograph_id_hash"), "micrograph", ["uuid"], unique=False, postgresql_using="hash")
    op.create_index(op.f("idx_micrograph_foilhole_id"), "micrograph", ["foilhole_uuid"], unique=False)
    op.create_index(op.f("idx_gridsquare_id_pattern"), "gridsquare", ["uuid"], unique=False)
    op.create_index(op.f("idx_gridsquare_id_hash"), "gridsquare", ["uuid"], unique=False, postgresql_using="hash")
    op.create_index(op.f("idx_gridsquare_grid_uuid"), "gridsquare", ["grid_uuid"], unique=False)
    op.create_index(op.f("idx_grid_uuid_pattern"), "grid", ["uuid"], unique=False)
    op.create_index(op.f("idx_grid_uuid_hash"), "grid", ["uuid"], unique=False, postgresql_using="hash")
    op.create_index(op.f("idx_grid_name"), "grid", ["name"], unique=False)
    op.create_index(op.f("idx_grid_acquisition_id"), "grid", ["acquisition_uuid"], unique=False)
    op.drop_column("grid", "prediction_updated_time")
    op.create_index(op.f("idx_foilhole_quality"), "foilhole", ["quality"], unique=False)
    op.create_index(op.f("idx_foilhole_id_pattern"), "foilhole", ["uuid"], unique=False)
    op.create_index(op.f("idx_foilhole_id_hash"), "foilhole", ["uuid"], unique=False, postgresql_using="hash")
    op.create_index(op.f("idx_foilhole_gridsquare_id"), "foilhole", ["gridsquare_uuid"], unique=False)
    op.create_index(op.f("idx_atlastile_id_pattern"), "atlastile", ["uuid"], unique=False)
    op.create_index(op.f("idx_atlastile_atlas_id"), "atlastile", ["atlas_uuid"], unique=False)
    op.create_index(op.f("idx_atlas_id_pattern"), "atlas", ["uuid"], unique=False)
    op.create_index(op.f("idx_atlas_grid_uuid"), "atlas", ["grid_uuid"], unique=False)
    op.create_index(op.f("idx_agent_session_status"), "agentsession", ["status"], unique=False)
    op.create_index(op.f("idx_agent_session_last_activity"), "agentsession", ["last_activity_at"], unique=False)
    op.create_index(op.f("idx_agent_session_created_at"), "agentsession", ["created_at"], unique=False)
    op.create_index(op.f("idx_agent_session_agent_id"), "agentsession", ["agent_id"], unique=False)
    op.create_index(op.f("idx_agent_session_acquisition_uuid"), "agentsession", ["acquisition_uuid"], unique=False)
    op.create_index(op.f("idx_agent_ack_status"), "agentinstructionacknowledgement", ["status"], unique=False)
    op.create_index(op.f("idx_agent_ack_session_id"), "agentinstructionacknowledgement", ["session_id"], unique=False)
    op.create_index(
        op.f("idx_agent_ack_processed_at"), "agentinstructionacknowledgement", ["processed_at"], unique=False
    )
    op.create_index(
        op.f("idx_agent_ack_instruction_id"), "agentinstructionacknowledgement", ["instruction_id"], unique=False
    )
    op.create_index(op.f("idx_agent_ack_created_at"), "agentinstructionacknowledgement", ["created_at"], unique=False)
    op.create_index(op.f("idx_agent_ack_agent_id"), "agentinstructionacknowledgement", ["agent_id"], unique=False)
    op.create_index(op.f("idx_agent_instruction_type"), "agentinstruction", ["instruction_type"], unique=False)
    op.create_index(op.f("idx_agent_instruction_status"), "agentinstruction", ["status"], unique=False)
    op.create_index(op.f("idx_agent_instruction_session_id"), "agentinstruction", ["session_id"], unique=False)
    op.create_index(op.f("idx_agent_instruction_sequence"), "agentinstruction", ["sequence_number"], unique=False)
    op.create_index(op.f("idx_agent_instruction_sent_at"), "agentinstruction", ["sent_at"], unique=False)
    op.create_index(op.f("idx_agent_instruction_priority"), "agentinstruction", ["priority"], unique=False)
    op.create_index(op.f("idx_agent_instruction_expires_at"), "agentinstruction", ["expires_at"], unique=False)
    op.create_index(op.f("idx_agent_instruction_created_at"), "agentinstruction", ["created_at"], unique=False)
    op.create_index(op.f("idx_agent_instruction_agent_id"), "agentinstruction", ["agent_id"], unique=False)
    op.create_index(
        op.f("idx_agent_instruction_acknowledged_at"), "agentinstruction", ["acknowledged_at"], unique=False
    )
    op.create_index(op.f("idx_agent_connection_status"), "agentconnection", ["status"], unique=False)
    op.create_index(op.f("idx_agent_connection_session_id"), "agentconnection", ["session_id"], unique=False)
    op.create_index(op.f("idx_agent_connection_heartbeat"), "agentconnection", ["last_heartbeat_at"], unique=False)
    op.create_index(op.f("idx_agent_connection_created_at"), "agentconnection", ["created_at"], unique=False)
    op.create_index(op.f("idx_agent_connection_agent_id"), "agentconnection", ["agent_id"], unique=False)
    op.create_index(op.f("idx_acquisition_name"), "acquisition", ["name"], unique=False)
    op.create_index(op.f("idx_acquisition_id_pattern"), "acquisition", ["uuid"], unique=False)
    op.create_index(op.f("idx_acquisition_id_hash"), "acquisition", ["uuid"], unique=False, postgresql_using="hash")
    op.drop_table("overallqualityprediction")
    op.drop_table("currentqualityprediction")
    op.drop_table("qualitymetricstatistics")
    op.drop_table("currentqualitypredictionmodelweight")
    op.drop_table("qualitymetric")
    # ### end Alembic commands ###
