# -*- 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.14

Revision ID: 6ece06672ed3
Revises: 6f0d0f71d92e
Create Date: 2022-08-31 07:16:24.229939

"""
import json

from sqlalchemy.orm import session, sessionmaker

from neural_compressor.ux.components.db_manager.db_manager import DBManager
from neural_compressor.ux.components.db_manager.db_models.dataloader import Dataloader
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.components.db_manager.db_models.transform import Transform
from neural_compressor.ux.components.db_manager.utils import (
    get_framework_dataloaders_config,
    get_framework_transforms_config,
    update_dataloaders_params,
)
from neural_compressor.ux.utils.consts import Frameworks, OptimizationTypes, Precisions

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


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


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with Session.begin() as db_session:
        # Add Mixed Precision optimization
        mixed_precision_id = OptimizationType.add(
            db_session=db_session,
            name=OptimizationTypes.MIXED_PRECISION.value,
        )
        bf16_precision_id = Precision.get_precision_by_name(
            db_session=db_session,
            precision_name=Precisions.BF16.value,
        )[0]

        query = precision_optimization_type_association.insert().values(
            precision_id=bf16_precision_id,
            optimization_type_id=mixed_precision_id,
        )
        db_session.execute(query)

        update_dataloaders(db_session)

        update_transformations(db_session)
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def update_dataloaders(db_session: session.Session):
    """Update dataloaders in existing database."""
    onnxrt_id = Framework.get_framework_id(db_session, framework_name=Frameworks.ONNX.value)
    tensorflow_id = Framework.get_framework_id(
        db_session,
        framework_name=Frameworks.TF.value,
    )

    onnx_dataloaders_to_update = [
        "CIFAR10",
        "CIFAR100",
        "MNIST",
        "FashionMNIST",
    ]
    tf_dataloaders_to_update = [
        "CIFAR10",
        "CIFAR100",
        "MNIST",
        "FashionMNIST",
        "TFRecordDataset",
    ]

    onnx_dataloaders_config = get_framework_dataloaders_config(Frameworks.ONNX)
    tf_dataloaders_config = get_framework_dataloaders_config(Frameworks.TF)

    # Update ONNX dataloaders' params,
    update_dataloaders_params(
        db_session,
        onnxrt_id,
        onnx_dataloaders_to_update,
        onnx_dataloaders_config,
    )

    # Update TF dataloaders' params
    update_dataloaders_params(
        db_session,
        tensorflow_id,
        tf_dataloaders_to_update,
        tf_dataloaders_config,
    )

    # Add new TF dataloaders
    dataloaders_to_add = ["COCONpy", "VOCRecord"]
    for dataloader_data in tf_dataloaders_config:
        if dataloader_data.get("name") in dataloaders_to_add:
            db_session.add(
                Dataloader(
                    name=dataloader_data.get("name"),
                    help=dataloader_data.get("help"),
                    show_dataset_location=dataloader_data.get("show_dataset_location"),
                    params=json.dumps(dataloader_data.get("params", [])),
                    framework_id=tensorflow_id,
                ),
            )


def update_transformations(db_session: session.Session):
    """Update transformations in existing database."""
    onnxrt_id = Framework.get_framework_id(db_session, framework_name=Frameworks.ONNX.value)
    tensorflow_id = Framework.get_framework_id(
        db_session,
        framework_name=Frameworks.TF.value,
    )

    onnx_transforms_config = get_framework_transforms_config(Frameworks.ONNX)
    tf_transforms_config = get_framework_transforms_config(Frameworks.TF)

    onnx_transforms_to_add = [
        "CropToBoundingBox",
        "Cast",
        "ResizeWithRatio",
    ]
    tf_transforms_to_add = [
        "ResizeWithRatio",
        "ParseDecodeImagenet",
    ]

    for transform_data in tf_transforms_config:
        if transform_data.get("name") in tf_transforms_to_add:
            db_session.add(
                Transform(
                    name=transform_data.get("name"),
                    help=transform_data.get("help"),
                    params=json.dumps(transform_data.get("params", [])),
                    framework_id=tensorflow_id,
                ),
            )

    for transform_data in onnx_transforms_config:
        if transform_data.get("name") in onnx_transforms_to_add:
            db_session.add(
                Transform(
                    name=transform_data.get("name"),
                    help=transform_data.get("help"),
                    params=json.dumps(transform_data.get("params", [])),
                    framework_id=onnxrt_id,
                ),
            )
