import functools
import itertools

import numpy as np
import pytest
import torch
from scipy import integrate, special

from ricciardi import ricciardi


@pytest.fixture(params=itertools.product([0.01], [0.01, 0.02], [0.002], [0.01], [0.02]))
def params(request):
    p = request.param
    return dict(sigma=p[0], tau=p[1], tau_rp=p[2], V_r=p[3], theta=p[4])


@np.vectorize
def ierfcx_exact(x, y):
    return integrate.quad(special.erfcx, x, y)[0]


def ricciardi_exact(mu, sigma=0.01, tau=0.02, tau_rp=0.002, V_r=0.01, theta=0.02):
    min_u = (V_r - mu) / sigma
    max_u = (theta - mu) / sigma
    return 1.0 / (tau_rp + tau * torch.pi**0.5 * ierfcx_exact(-max_u, -min_u))


@pytest.mark.parametrize(
    "x",
    [
        torch.linspace(-0.01, 0.1, 1001),
        torch.linspace(-10.0, 50.0, 1001),
    ],
)
@pytest.mark.parametrize(
    "dtype", [torch.long, torch.bfloat16, torch.half, torch.float, torch.double]
)
@pytest.mark.parametrize("vmap", [False, True])
def test_ricciardi(x, params, dtype, vmap):
    x = x.to(dtype)
    if vmap:
        out = torch.vmap(ricciardi)(x, **params)
    else:
        out = ricciardi(x, **params)
    expected = ricciardi_exact(x.double().numpy(), **params)
    expected = torch.from_numpy(expected).nan_to_num()
    if torch.is_floating_point(x):
        expected = expected.to(dtype)
    else:
        expected = expected.float()
    torch.testing.assert_close(out, expected)


@pytest.mark.parametrize(
    "x",
    [
        torch.linspace(-0.01, 0.1, 1001, requires_grad=True).double(),
        torch.linspace(-10.0, 50.0, 1001, requires_grad=True).double(),
    ],
)
def test_ricciardi_grad(x, params):
    torch.autograd.gradcheck(functools.partial(ricciardi, **params), x)


def test_ricciardi_raises():
    with pytest.raises(TypeError):
        ricciardi(0.0)

    with pytest.raises(TypeError):
        ricciardi(torch.tensor(1.0j))
