# 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.
"""Intrinsics of TVM-Python Hybrid Script for Python compilation time
semantic support."""

from tvm.runtime import const, convert
import tvm.te
from tvm.ir.container import Array
from tvm import target as _tgt
from tvm.tir import expr as _expr
from tvm.tir import call_pure_intrin
from tvm.tir.stmt import For

from .util import _internal_assert

# pylint: disable=redefined-builtin

LOOP_INTRIN = {
    'range'       : For.Serial,
    'unroll'      : For.Unrolled,
    'parallel'    : For.Parallel,
    'vectorize'   : For.Vectorized,
    'const_range' : (For.Unrolled, ),
}


def _range(annotation, args):
    """Handling TVM loop types"""
    n = args.__len__()
    if n == 1:
        low, ext = const(0, dtype='int32'), args[0]
    else:
        _internal_assert(n == 2, "A loop intrinsic should only have 1 or 2 arguments!")
        low, ext = args[0], args[1]
    if not tvm.tir.analysis.expr_deep_equal(low, const(0, dtype='int32')):
        ext = ext - low
    for_type = LOOP_INTRIN[annotation]
    iter_var = None
    return iter_var, low, ext, for_type


range = unroll = vectorize = parallel = const_range = _range #pylint: disable=invalid-name


def bind(func_id, args):
    """Handling TVM thread binding"""
    _internal_assert(func_id == "bind", "This function cannot be directly invoked!")
    _internal_assert(args.__len__() == 2, "A loop bind should only have 2 arguments!")
    _internal_assert(isinstance(args[0], str), \
                     "A loop bind's first argument should be a string!")
    low, ext = const(0, "int32"), args[1]
    iter_var = tvm.te.thread_axis((low, ext), args[0])
    for_type = None
    return iter_var, low, ext, for_type


def _math_intrin(func_id, args):
    # pylint: disable=import-outside-toplevel
    from tvm.tir import op
    return getattr(op, func_id)(*args)

sqrt = log = exp = tanh = sigmoid = power = popcount = _math_intrin #pylint: disable=invalid-name


def _min_max(func_id, args):
    _internal_assert(args.__len__() == 2, "Max/Min function should have 2 elements")
    return getattr(_expr, func_id.title())(args[0], args[1])


min = max = _min_max #pylint: disable=invalid-name


def _allocate_tensor(func_id, args):
    """Handling TVM tensor allocation.
    You may refer hybrid.intrin.allocate for more details."""
    n = args.__len__()
    _internal_assert(isinstance(convert(args[0]), Array), \
                     "allocate's first argument should be a tuple of shape!")
    shape = args[0]
    for i in shape:
        _internal_assert(isinstance(i, _expr.PrimExpr), "The shape should be an expression")
    if n > 1:
        _internal_assert(isinstance(args[1], str),
                         "The data type should be an str")
        _internal_assert(args[1].startswith('int') or args[1].startswith('float'), \
                         "The data type should be either int or float!")
        dtype = args[1]
    else:
        dtype = 'float32'
    if n > 2:
        _internal_assert(isinstance(args[2], str), \
                         "The data scope should be an string")
        _internal_assert(func_id != 'output_tensor', "Output tensor cannot specify scope")
        scope = args[2]
    else:
        scope = 'global' if func_id != 'output_tensor' else 'output'
    return (shape, dtype, scope)


output_tensor = allocate = _allocate_tensor #pylint: disable=invalid-name


def len(func_id, args):
    """Iterpret the len function"""
    _internal_assert(args.__len__() == 1, "Only 1 argument is expected!")
    _internal_assert(func_id == "len", "This function cannot be directly invoked!")
    try:
        return convert(args[0].__len__())
    except: #pylint: disable=bare-except
        _internal_assert(args[0].shape.__len__() == 1, "Only one-dimension array can get len")
        return convert(args[0].shape[0])


def _cast(func_id, args):
    _internal_assert(args.__len__() == 1 and isinstance(args[0], _expr.PrimExpr), \
                     "Only one expression can be cast")
    return _expr.Cast(func_id, args[0])

float16 = float32 = float64 = _cast #pylint: disable=invalid-name
int8 = int16 = int32 = int64 = _cast #pylint: disable=invalid-name
uint8 = uint16 = uint32 = uint64 = _cast #pylint: disable=invalid-name


def ceil_div(func_id, args):
    _internal_assert(func_id == "ceil_div", "This function cannot be directly invoked!")
    _internal_assert(args.__len__() == 2, "2 arguments expected for division!")
    _internal_assert(isinstance(args[0], _expr.PrimExpr), "Only expressions can div")
    _internal_assert(isinstance(args[1], _expr.PrimExpr), "Only expressions can div")
    a, b = args[0], args[1]
    return (a + b - 1) // b


def likely(func_id, args):
    _internal_assert(args.__len__() == 1, \
                     "Only one expression can be likely")
    _internal_assert(func_id == "likely", "This function cannot be directly invoked!")
    return call_pure_intrin(args[0].dtype, 'likely', *args)


def max_num_threads(func_id, args):
    _internal_assert(func_id == "max_num_threads", "This function cannot be directly invoked!")
    _internal_assert(args.__len__() <= 1, "At most one argument accepted!")
    if args.__len__() == 0:
        res = _tgt.Target.current().max_num_threads
    else:
        _internal_assert(isinstance(args[0], _expr.IntImm), "In tvm bool should be uint")
        res = _tgt.Target.current(args[0].value).max_num_threads
    return convert(res)
