import os
import pytest
import scipy.special as sc
import shutil
import tempfile

from uuid import uuid4

from scipy.special._testutils import check_version
from scipy.special._testutils import MissingModule

try:
    import cupy  # type: ignore
except (ImportError, AttributeError):
    cupy = MissingModule('cupy')


def get_test_cases():
    cases_source = [
        (sc.beta, "cephes/beta.h", "out0 = xsf::cephes::beta(in0, in1)"),
        (sc.binom, "binom.h", "out0 = xsf::binom(in0, in1)"),
        (sc.digamma, "digamma.h", "xsf::digamma(in0)"),
        (sc.expn, "cephes/expn.h", "out0 = xsf::cephes::expn(in0, in1)"),
        (sc.hyp2f1, "hyp2f1.h", "out0 = xsf::hyp2f1(in0, in1, in2, in3)"),
        (sc._ufuncs._lambertw, "lambertw.h", "out0 = xsf::lambertw(in0, in1, in2)"),
        (sc.ellipkinc, "cephes/ellik.h", "out0 = xsf::cephes::ellik(in0, in1)"),
        (sc.ellipeinc, "cephes/ellie.h", "out0 = xsf::cephes::ellie(in0, in1)"),
        (sc.gdtrib, "cdflib.h", "out0 = xsf::gdtrib(in0, in1, in2)"),
        (sc.sici, "sici.h", "xsf::sici(in0, &out0, &out1)"),
        (sc.shichi, "sici.h", "xsf::shichi(in0, &out0, &out1)"),
    ]

    cases = []
    for ufunc, header, routine in cases_source:
        preamble = f"#include <xsf/{header}>"
        for signature in ufunc.types:
            cases.append((signature, preamble, routine))
    return cases


dtype_map = {
    "f": "float32",
    "d": "float64",
    "F": "complex64",
    "D": "complex128",
    "i": "int32",
    "l": "int64",
}


def get_params(signature):
    in_, out = signature.split("->")
    in_params = []
    out_params = []
    for i, typecode in enumerate(in_):
        in_params.append(f"{dtype_map[typecode]} in{i}")
    for i, typecode in enumerate(out):
        out_params.append(f"{dtype_map[typecode]} out{i}")
    in_params = ", ".join(in_params)
    out_params = ", ".join(out_params)
    return in_params, out_params


def get_sample_input(signature, xp):
    dtype_map = {
        "f": xp.float32,
        "d": xp.float64,
        "F": xp.complex64,
        "D": xp.complex128,
        "i": xp.int32,
        "l": xp.int64,
    }

    in_, _ = signature.split("->")
    args = []
    for typecode in in_:
        args.append(xp.zeros(2, dtype=dtype_map[typecode]))
    return args


@pytest.fixture(scope="module", autouse=True)
def manage_cupy_cache():
    # Temporarily change cupy kernel cache location so kernel cache will not be polluted
    # by these tests. Remove temporary cache in teardown.
    temp_cache_dir = tempfile.mkdtemp()
    original_cache_dir = os.environ.get('CUPY_CACHE_DIR', None)
    os.environ['CUPY_CACHE_DIR'] = temp_cache_dir

    yield

    if original_cache_dir is not None:
        os.environ['CUPY_CACHE_DIR'] = original_cache_dir
    else:
        del os.environ['CUPY_CACHE_DIR']
    shutil.rmtree(temp_cache_dir)


@check_version(cupy, "13.0.0")
@pytest.mark.parametrize("signature,preamble,routine", get_test_cases())
@pytest.mark.xslow
def test_compiles_in_cupy(signature, preamble, routine, manage_cupy_cache):
    name = f"x{uuid4().hex}"
    in_params, out_params = get_params(signature)

    func = cupy.ElementwiseKernel(
        in_params,
        out_params,
        routine,
        name,
        preamble=preamble,
        options=(f"--include-path={sc._get_include()}", "-std=c++17")
    )

    _ = func(*get_sample_input(signature, cupy))