# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
# pylint: disable=invalid-name, no-member, too-many-locals, too-many-arguments, undefined-variable
"""SSD multibox operators"""
from __future__ import absolute_import as _abs
import tvm

from tvm import hybrid
from tvm.intrin import exp, sqrt

import topi

from ..nms import non_max_suppression

@hybrid.script
def hybrid_multibox_prior(data, sizes, ratios, steps, offsets):
    """Hybrid routing for multibox_prior operator.

    Parameters
    ----------
    data : tvm.Tensor or numpy NDArray
        4-D tensor with shape [batch, channel, height, width]]

    sizes : tvm ConsExpr
        Sizes for anchor boxes.

    ratios : tvm ConsExpr
        Ratios for anchor boxes.

    steps : tvm ConsExpr
        Priorbox step across y and x, -1 for auto calculation.

    offsets : tvm ConsExpr
        Priorbox center offsets, y and x respectively.

    Returns
    -------
    output : tvm.Tensor or numpy NDArray
        3-D tensor with shape [1, h_in * w_in * (num_sizes + num_ratios - 1), 4]
    """
    in_height = data.shape[2]
    in_width = data.shape[3]
    num_sizes = len(sizes)
    num_ratios = len(ratios)
    num_boxes = in_height * in_width * (num_sizes + num_ratios - 1)
    output = output_tensor((1, num_boxes, 4), "float32")
    steps_h = steps[0] * 1.0 if steps[0] > 0 else 1.0 / in_height
    steps_w = steps[1] * 1.0 if steps[1] > 0 else 1.0 / in_width
    offset_h = offsets[0]
    offset_w = offsets[1]

    # Need to define var out of const_range + if
    w = 0.0
    h = 0.0

    for i in parallel(in_height):
        center_h = (i + offset_h) * steps_h
        for j in range(in_width):
            center_w = (j + offset_w) * steps_w
            for k in const_range(num_sizes + num_ratios - 1):
                if k < num_sizes:
                    w = sizes[k] * in_height / in_width / 2.0
                    h = sizes[k] / 2.0
                else:
                    w = sizes[0] * in_height / in_width \
                        * sqrt(ratios[k - num_sizes + 1] * 1.0) / 2.0
                    h = sizes[0] / sqrt(ratios[k - num_sizes + 1] * 1.0) / 2.0
                count = i * in_width * (num_sizes + num_ratios - 1) \
                        + j * (num_sizes + num_ratios - 1) + k
                output[0, count, 0] = center_w - w
                output[0, count, 1] = center_h - h
                output[0, count, 2] = center_w + w
                output[0, count, 3] = center_h + h

    return output


@tvm.target.generic_func
def multibox_prior(data, sizes=(1,), ratios=(1,), steps=(-1, -1), offsets=(0.5, 0.5), clip=False):
    """Generate prior(anchor) boxes from data, sizes and ratios.

    Parameters
    ----------
    data : tvm.Tensor
        4-D with shape [batch, c_in, h_in, w_in]]

    sizes : tuple of float
        Tuple of sizes for anchor boxes.

    ratios : tuple of float
        Tuple of ratios for anchor boxes.

    steps : Tuple of float
        Priorbox step across y and x, -1 for auto calculation.

    offsets : tuple of int
        Priorbox center offsets, y and x respectively.

    clip : boolean
        Whether to clip out-of-boundary boxes.

    Returns
    -------
    out : tvm.Tensor
        3-D tensor with shape [1, h_in * w_in * (num_sizes + num_ratios - 1), 4]
    """
    out = hybrid_multibox_prior(data, tvm.convert(sizes), tvm.convert(ratios),
                                tvm.convert(steps), tvm.convert(offsets))
    if clip:
        out = topi.clip(out, 0, 1)
    return out


@hybrid.script
def _hybridy_transform_loc(box, pred_loc, variance, clip):
    """Transform prior anchor box to output box through location predictions.
    """
    al = box[0]
    at = box[1]
    ar = box[2]
    ab = box[3]

    px = pred_loc[0]
    py = pred_loc[1]
    pw = pred_loc[2]
    ph = pred_loc[3]

    vx = variance[0]