"""merge score and model_trainings into model_selection_scores

Revision ID: 033e0f7eca4f
Revises: 07e303521594
Create Date: 2025-10-26 17:27:30.400473

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from lecrapaud.config import LECRAPAUD_TABLE_PREFIX

# revision identifiers, used by Alembic.
revision: str = "033e0f7eca4f"
down_revision: Union[str, None] = "07e303521594"
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! ###
    # Check if table exists using inspector
    from sqlalchemy import inspect
    dialect = op.get_bind().dialect.name
    # "precision" is a reserved word — quote differently per dialect
    q = '"' if dialect == "postgresql" else '`'
    inspector = inspect(op.get_bind())
    existing_tables = inspector.get_table_names()
    
    if f"{LECRAPAUD_TABLE_PREFIX}_model_selection_scores" not in existing_tables:
        op.create_table(
            f"{LECRAPAUD_TABLE_PREFIX}_model_selection_scores",
            sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
            sa.Column(
                "created_at",
                sa.TIMESTAMP(timezone=True),
                server_default=sa.text("now()"),
                nullable=False,
            ),
            sa.Column(
                "updated_at",
                sa.TIMESTAMP(timezone=True),
                server_default=sa.text("now()"),
                nullable=False,
            ),
            sa.Column("best_params", sa.JSON(), nullable=True),
            sa.Column("model_path", sa.String(length=255), nullable=True),
            sa.Column("training_time", sa.Integer(), nullable=True),
            sa.Column("model_id", sa.BigInteger(), nullable=False),
            sa.Column("model_selection_id", sa.BigInteger(), nullable=False),
            sa.Column("eval_data_std", sa.Float(), nullable=True),
            sa.Column("rmse", sa.Float(), nullable=True),
            sa.Column("rmse_std_ratio", sa.Float(), nullable=True),
            sa.Column("mae", sa.Float(), nullable=True),
            sa.Column("mape", sa.Float(), nullable=True),
            sa.Column("mam", sa.Float(), nullable=True),
            sa.Column("mad", sa.Float(), nullable=True),
            sa.Column("mae_mam_ratio", sa.Float(), nullable=True),
            sa.Column("mae_mad_ratio", sa.Float(), nullable=True),
            sa.Column("r2", sa.Float(), nullable=True),
            sa.Column("logloss", sa.Float(), nullable=True),
            sa.Column("accuracy", sa.Float(), nullable=True),
            sa.Column("precision", sa.Float(), nullable=True),
            sa.Column("recall", sa.Float(), nullable=True),
            sa.Column("f1", sa.Float(), nullable=True),
            sa.Column("roc_auc", sa.Float(), nullable=True),
            sa.Column("avg_precision", sa.Float(), nullable=True),
            sa.Column("thresholds", sa.JSON(), nullable=True),
            sa.Column("precision_at_threshold", sa.Float(), nullable=True),
            sa.Column("recall_at_threshold", sa.Float(), nullable=True),
            sa.Column("f1_at_threshold", sa.Float(), nullable=True),
            sa.ForeignKeyConstraint(
                ["model_id"],
                [f"{LECRAPAUD_TABLE_PREFIX}_models.id"],
            ),
            sa.ForeignKeyConstraint(
                ["model_selection_id"],
                [f"{LECRAPAUD_TABLE_PREFIX}_model_selections.id"],
                ondelete="CASCADE",
            ),
            sa.PrimaryKeyConstraint("id"),
            sa.UniqueConstraint(
                "model_id",
                "model_selection_id",
                name="uq_model_selection_score_composite",
            ),
        )
        op.create_index(
            op.f("ix_model_selection_scores_id"),
            f"{LECRAPAUD_TABLE_PREFIX}_model_selection_scores",
            ["id"],
            unique=False,
        )

    # Migrate existing data from model_trainings and scores to model_selection_scores
    op.execute(
        f"""
        INSERT INTO {LECRAPAUD_TABLE_PREFIX}_model_selection_scores (
            created_at, updated_at, best_params, model_path, training_time,
            model_id, model_selection_id,
            eval_data_std, rmse, rmse_std_ratio, mae, mape, mam, mad,
            mae_mam_ratio, mae_mad_ratio, r2, logloss, accuracy, {q}precision{q},
            recall, f1, roc_auc, avg_precision, thresholds,
            precision_at_threshold, recall_at_threshold, f1_at_threshold
        )
        SELECT
            mt.created_at,
            mt.updated_at,
            mt.best_params,
            mt.model_path,
            COALESCE(mt.training_time, s.training_time) as training_time,
            mt.model_id,
            mt.model_selection_id,
            s.eval_data_std, s.rmse, s.rmse_std_ratio, s.mae, s.mape,
            s.mam, s.mad, s.mae_mam_ratio, s.mae_mad_ratio, s.r2,
            s.logloss, s.accuracy, s.{q}precision{q}, s.recall, s.f1,
            s.roc_auc, s.avg_precision, s.thresholds,
            s.precision_at_threshold, s.recall_at_threshold, s.f1_at_threshold
        FROM {LECRAPAUD_TABLE_PREFIX}_model_trainings mt
        LEFT JOIN {LECRAPAUD_TABLE_PREFIX}_scores s ON s.model_training_id = mt.id
    """
    )

    # Drop the old tables
    op.drop_table(f"{LECRAPAUD_TABLE_PREFIX}_scores")
    op.drop_table(f"{LECRAPAUD_TABLE_PREFIX}_model_trainings")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    dialect = op.get_bind().dialect.name
    q = '"' if dialect == "postgresql" else '`'
    # Recreate the old tables
    op.create_table(
        f"{LECRAPAUD_TABLE_PREFIX}_model_trainings",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column(
            "created_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("best_params", sa.JSON(), nullable=True),
        sa.Column("model_path", sa.String(length=255), nullable=True),
        sa.Column("training_time", sa.Integer(), nullable=True),
        sa.Column("model_id", sa.BigInteger(), nullable=False),
        sa.Column("model_selection_id", sa.BigInteger(), nullable=False),
        sa.ForeignKeyConstraint(
            ["model_id"],
            [f"{LECRAPAUD_TABLE_PREFIX}_models.id"],
        ),
        sa.ForeignKeyConstraint(
            ["model_selection_id"],
            [f"{LECRAPAUD_TABLE_PREFIX}_model_selections.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "model_id", "model_selection_id", name="uq_model_training_composite"
        ),
    )
    op.create_index(
        op.f("ix_model_trainings_id"),
        f"{LECRAPAUD_TABLE_PREFIX}_model_trainings",
        ["id"],
        unique=False,
    )

    op.create_table(
        f"{LECRAPAUD_TABLE_PREFIX}_scores",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column(
            "created_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("type", sa.String(length=50), nullable=False),
        sa.Column("training_time", sa.Integer(), nullable=True),
        sa.Column("eval_data_std", sa.Float(), nullable=True),
        sa.Column("rmse", sa.Float(), nullable=True),
        sa.Column("rmse_std_ratio", sa.Float(), nullable=True),
        sa.Column("mae", sa.Float(), nullable=True),
        sa.Column("mape", sa.Float(), nullable=True),
        sa.Column("mam", sa.Float(), nullable=True),
        sa.Column("mad", sa.Float(), nullable=True),
        sa.Column("mae_mam_ratio", sa.Float(), nullable=True),
        sa.Column("mae_mad_ratio", sa.Float(), nullable=True),
        sa.Column("r2", sa.Float(), nullable=True),
        sa.Column("logloss", sa.Float(), nullable=True),
        sa.Column("accuracy", sa.Float(), nullable=True),
        sa.Column("precision", sa.Float(), nullable=True),
        sa.Column("recall", sa.Float(), nullable=True),
        sa.Column("f1", sa.Float(), nullable=True),
        sa.Column("roc_auc", sa.Float(), nullable=True),
        sa.Column("avg_precision", sa.Float(), nullable=True),
        sa.Column("thresholds", sa.JSON(), nullable=True),
        sa.Column("precision_at_threshold", sa.Float(), nullable=True),
        sa.Column("recall_at_threshold", sa.Float(), nullable=True),
        sa.Column("f1_at_threshold", sa.Float(), nullable=True),
        sa.Column("model_training_id", sa.BigInteger(), nullable=False),
        sa.ForeignKeyConstraint(
            ["model_training_id"],
            [f"{LECRAPAUD_TABLE_PREFIX}_model_trainings.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "model_training_id", name="unique_score_per_model_training"
        ),
    )
    op.create_index(
        op.f("ix_scores_id"), f"{LECRAPAUD_TABLE_PREFIX}_scores", ["id"], unique=False
    )

    # Migrate data back (note: we'll lose the type column data, defaulting to 'testset')
    op.execute(
        f"""
        INSERT INTO {LECRAPAUD_TABLE_PREFIX}_model_trainings (
            id, created_at, updated_at, best_params, model_path, 
            training_time, model_id, model_selection_id
        )
        SELECT 
            id, created_at, updated_at, best_params, model_path, 
            training_time, model_id, model_selection_id
        FROM {LECRAPAUD_TABLE_PREFIX}_model_selection_scores
    """
    )

    op.execute(
        f"""
        INSERT INTO {LECRAPAUD_TABLE_PREFIX}_scores (
            created_at, updated_at, type, training_time, eval_data_std, 
            rmse, rmse_std_ratio, mae, mape, mam, mad, mae_mam_ratio, 
            mae_mad_ratio, r2, logloss, accuracy, {q}precision{q}, recall,
            f1, roc_auc, avg_precision, thresholds, precision_at_threshold,
            recall_at_threshold, f1_at_threshold, model_training_id
        )
        SELECT
            created_at, updated_at, 'testset', training_time, eval_data_std,
            rmse, rmse_std_ratio, mae, mape, mam, mad, mae_mam_ratio,
            mae_mad_ratio, r2, logloss, accuracy, {q}precision{q}, recall,
            f1, roc_auc, avg_precision, thresholds, precision_at_threshold,
            recall_at_threshold, f1_at_threshold, id
        FROM {LECRAPAUD_TABLE_PREFIX}_model_selection_scores
    """
    )

    op.drop_index(
        op.f("ix_model_selection_scores_id"),
        table_name=f"{LECRAPAUD_TABLE_PREFIX}_model_selection_scores",
    )
    op.drop_table(f"{LECRAPAUD_TABLE_PREFIX}_model_selection_scores")
    # ### end Alembic commands ###
