#  Copyright (c) 2020, Apple Inc. All rights reserved.
#
#  Use of this source code is governed by a BSD-3-clause license that can be
#  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause

import numpy as np

from coremltools.converters.mil.input_types import ColorLayout, ImageType
from coremltools.converters.mil.mil import Builder as mb
# import mil internal ops to add it to the builder
from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass
from coremltools.converters.mil.mil.passes.helper import block_context_manager
from coremltools.converters.mil.mil.passes.pass_registry import register_pass
from coremltools.converters.mil.mil.types import nptype_from_builtin


@register_pass(namespace="mil_backend")
class insert_image_preprocessing_ops(AbstractGraphPass):
    """
    Insert preprocessing ops, right after the input if its of type Image
    """
    def apply(self, prog):
        for f_name, f in prog.functions.items():
            if f_name == 'main':
                if len(f.operations) == 0:
                    continue
                # The new image processing ops will be inserted in front of the first op in the function
                with mb.set_before_op(f.operations[0]):
                    _insert_image_preprocessing_ops(f)

@block_context_manager
def _insert_image_preprocessing_ops(block):
    input_types = list(block.input_types)

    for input_type in input_types:
        if isinstance(input_type, ImageType):
            if input_type.name not in block.inputs:
                continue

            input_var = block.inputs[input_type.name]
            placeholder_op = block.placeholder_inputs[input_type.name]
            first_op = block.operations[0]
            old_var = placeholder_op.outputs[0]
            has_bias = np.any(np.array(input_type.bias) != 0)
            last_output = input_var
            input_nptype = nptype_from_builtin(type(last_output.dtype()))
            if input_type.scale != 1:
                last_output = mb.mul(
                    x=last_output,
                    y=np.array(input_type.scale, dtype=input_nptype),
                    name=input_var.name + "__scaled__",
                )
            if has_bias:
                if input_type.color_layout in (
                    ColorLayout.GRAYSCALE,
                    ColorLayout.GRAYSCALE_FLOAT16,
                ):
                    last_output = mb.add(
                        x=last_output,
                        y=np.array(input_type.bias, dtype=input_nptype),
                        name=input_var.name + "__biased__",
                    )
                else:
                    if len(last_output.shape) == 3:
                        last_output = mb.add(
                            x=last_output,
                            y=np.array(input_type.bias, dtype=input_nptype).reshape([3, 1, 1]),
                            name=input_var.name + "__biased__",
                        )
                    elif len(last_output.shape) == 4:
                        last_output = mb.add(
                            x=last_output,
                            y=np.array(input_type.bias, dtype=input_nptype).reshape([1, 3, 1, 1]),
                            name=input_var.name + "__biased__",
                        )
                    else:
                        raise TypeError("Unsupported rank for image input type.")

            if last_output != input_var:
                block.replace_uses_of_var_after_op(anchor_op=last_output.op,
                                                   old_var=old_var,
                                                   new_var=last_output)
