import numpy as np
import cython
# We now need to fix a datatype for our arrays. I've used the variable
# DTYPE for this, which is assigned to the usual NumPy runtime
# type info object.
DTYPE = np.intc

# @cython.cfunc means here that this function is a plain C function (so faster).
# To get all the benefits, we type the arguments and the return value.
@cython.exceptval(check=False)
@cython.cfunc
def clip(a: cython.int, min_value: cython.int, max_value: cython.int) -> cython.int:
    return min(max(a, min_value), max_value)


def compute(array_1, array_2, a: cython.int, b: cython.int, c: cython.int):

    # Annotation is also used within functions to type variables. It
    # can only be used at the top indentation level (there are non-trivial
    # problems with allowing them in other places, though we'd love to see
    # good and thought out proposals for it).
    x_max: cython.Py_ssize_t  = array_1.shape[0]
    y_max: cython.Py_ssize_t  = array_1.shape[1]

    assert array_1.shape == array_2.shape
    assert array_1.dtype == DTYPE
    assert array_2.dtype == DTYPE

    result = np.zeros((x_max, y_max), dtype=DTYPE)

    # It is very important to type ALL your variables. You do not get any
    # warnings if not, only much slower code (they are implicitly typed as
    # Python objects).
    # For the "tmp" variable, we want to use the same data type as is
    # stored in the array, so we use int because it correspond to np.intc.
    # NB! An important side-effect of this is that if "tmp" overflows its
    # datatype size, it will simply wrap around like in C, rather than raise
    # an error like in Python.

    tmp: cython.int

    # cython.Py_ssize_t is the proper C type for Python array indices.
    x: cython.Py_ssize_t
    y:  cython.Py_ssize_t

    for x in range(x_max):
        for y in range(y_max):

            tmp = clip(array_1[x, y], 2, 10)
            tmp = tmp * a + array_2[x, y] * b
            result[x, y] = tmp + c

    return result
