# 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.
# ruff: noqa: E722
"""Configure pytest"""

# pylint: disable=invalid-name
import threading

import numpy as np
import pytest

import tvm
import tvm.testing
from tvm import rpc, te
from tvm.contrib import random


def test_randint():
    """Tests randint function"""
    m = 10240
    n = 10240
    A = random.randint(-127, 128, size=(m, n), dtype="int32")

    def verify(target="llvm"):
        if not tvm.testing.device_enabled(target):
            print(f"skip because {target} is not enabled...")
            return
        if not tvm.get_global_func("tvm.contrib.random.randint", True):
            print("skip because extern function is not available")
            return
        dev = tvm.cpu(0)
        f = tvm.compile(te.create_prim_func([A]), target=target)
        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
        f(a)
        na = a.numpy()
        assert abs(np.mean(na)) < 0.3
        assert np.min(na) == -127
        assert np.max(na) == 127

    verify()


def test_uniform():
    """Tests uniform function"""
    m = 10240
    n = 10240
    A = random.uniform(0, 1, size=(m, n))

    def verify(target="llvm"):
        if not tvm.testing.device_enabled(target):
            print(f"skip because {target} is not enabled...")
            return
        if not tvm.get_global_func("tvm.contrib.random.uniform", True):
            print("skip because extern function is not available")
            return
        dev = tvm.cpu(0)
        f = tvm.compile(te.create_prim_func([A]), target=target)
        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
        f(a)
        na = a.numpy()
        assert abs(np.mean(na) - 0.5) < 1e-1
        assert abs(np.min(na) - 0.0) < 1e-3
        assert abs(np.max(na) - 1.0) < 1e-3

    verify()


def test_normal():
    """Tests normal function"""
    m = 10240
    n = 10240
    A = random.normal(3, 4, size=(m, n))

    def verify(target="llvm"):
        if not tvm.testing.device_enabled(target):
            print(f"skip because {target} is not enabled...")
            return
        if not tvm.get_global_func("tvm.contrib.random.normal", True):
            print("skip because extern function is not available")
            return
        dev = tvm.cpu(0)
        f = tvm.compile(te.create_prim_func([A]), target=target)
        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
        f(a)
        na = a.numpy()
        assert abs(np.mean(na) - 3) < 1e-1
        assert abs(np.std(na) - 4) < 1e-2

    verify()


@pytest.mark.gpu
def test_random_fill():
    """Tests random_fill function"""

    def test_local(dev, dtype):
        if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
            print("skip because extern function is not available")
            return
        value = tvm.runtime.empty((512, 512), dtype, dev)
        random_fill = tvm.get_global_func("tvm.contrib.random.random_fill")
        random_fill(value)

        assert np.count_nonzero(value.numpy()) == 512 * 512

        # make sure arithmentic doesn't overflow too
        np_values = value.numpy()
        assert np.isfinite(np_values * np_values + np_values).any()

    def test_rpc(dtype):
        if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
            print("skip because extern function is not available")
            return
        if not tvm.testing.device_enabled("rpc") or not tvm.runtime.enabled("llvm"):
            return

        def check_remote(server):
            remote = rpc.connect(server.host, server.port)
            value = tvm.runtime.empty((512, 512), dtype, remote.cpu())
            random_fill = remote.get_function("tvm.contrib.random.random_fill")
            random_fill(value)

            assert np.count_nonzero(value.numpy()) == 512 * 512

            # make sure arithmentic doesn't overflow too
            np_values = value.numpy()
            assert np.isfinite(np_values * np_values + np_values).any()

        check_remote(rpc.Server("127.0.0.1"))

    # Packed sub-byte dtypes (e.g. int4) are intentionally unsupported by
    # random_fill since #19714 and raise an error instead.
    for dtype in [
        "bool",
        "int8",
        "uint8",
        "int16",
        "uint16",
        "int32",
        "int32",
        "int64",
        "uint64",
        "float16",
        "float32",
        "float64",
    ]:
        for target, dev in tvm.testing.enabled_targets():
            if tvm.target.Target(target).kind.name == "llvm":
                test_local(dev, dtype)
            else:
                tvm.testing.run_with_gpu_lock(test_local, dev, dtype)
        test_rpc(dtype)


def test_random_fill_mt():
    """Check random filler applicability in case of nontrivial thread pool configuration.
    Particularly when MaxConcurrency != num_workers_used_ which is actual for big-little systems.
    """
    no_exception_happened = True

    def test_body():
        try:
            num_thread_used = 1
            configure_threads = tvm.get_global_func("runtime.config_threadpool")
            configure_threads(1, num_thread_used)

            test_input = tvm.runtime.empty((10, 10))
            random_fill = tvm.get_global_func("tvm.contrib.random.random_fill_for_measure")
            random_fill(test_input)
        except:  # pylint: disable=bare-except
            nonlocal no_exception_happened
            no_exception_happened = False

    # ThreadPool object is thread local. To eliminate effect on other test cases put it into thread
    x = threading.Thread(target=test_body)
    x.start()
    x.join()
    assert no_exception_happened


if __name__ == "__main__":
    test_randint()
    test_uniform()
    test_normal()
    test_random_fill()
    test_random_fill_mt()