# -*- coding: utf-8 -*-
# Copyright (c) 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# flake8: noqa
# mypy: ignore-errors
"""v1.13

Revision ID: 6f0d0f71d92e
Revises: 9e89549a08c8
Create Date: 2022-07-14 07:29:41.291905

"""
import sqlalchemy as sa
from alembic import op
from sqlalchemy.orm import sessionmaker

from neural_compressor.ux.components.db_manager.db_manager import DBManager
from neural_compressor.ux.components.db_manager.db_models.domain import Domain
from neural_compressor.ux.components.db_manager.db_models.framework import Framework
from neural_compressor.ux.components.db_manager.db_models.optimization_type import OptimizationType
from neural_compressor.ux.components.db_manager.db_models.precision import (
    Precision,
    precision_optimization_type_association,
)
from neural_compressor.ux.utils.consts import Domains, Frameworks, OptimizationTypes, Precisions

db_manager = DBManager()
Session = sessionmaker(bind=db_manager.engine)


# revision identifiers, used by Alembic.
revision = "6f0d0f71d92e"
down_revision = "9e89549a08c8"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("optimization", schema=None) as batch_op:
        batch_op.add_column(sa.Column("diagnosis_config", sa.String(), nullable=True))

    with Session.begin() as db_session:
        int8_static_quant_id = Precision.add(
            db_session=db_session,
            name=Precisions.INT8_STATIC_QUANTIZATION.value,
        )
        int8_dynamic_quant_id = Precision.add(
            db_session=db_session,
            name=Precisions.INT8_DYNAMIC_QUANTIZATION.value,
        )
        quantization_id = OptimizationType.get_optimization_type_id(
            db_session,
            OptimizationTypes.QUANTIZATION.value,
        )

        for precision_id in [int8_static_quant_id, int8_dynamic_quant_id]:
            query = precision_optimization_type_association.insert().values(
                precision_id=precision_id,
                optimization_type_id=quantization_id,
            )
            db_session.execute(query)

        db_session.add(Framework(name=Frameworks.PT.value))
        db_session.add(Domain(name=Domains.NONE.value))

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("optimization", schema=None) as batch_op:
        batch_op.drop_column("diagnosis_config")

    # ### end Alembic commands ###
