'''
Test that maximizing likelihood allows to correctly recover distribution parameters for all
distributions exposed to the user.
'''
# Standard library imports
from typing import Iterable, List, Tuple

# Third-party imports
import mxnet as mx
import numpy as np
import pytest
from pydantic import PositiveFloat, PositiveInt

# First-party imports
from gluonts.model.common import NPArrayLike
from gluonts.distribution.box_cox_tranform import (
    InverseBoxCoxTransform,
    InverseBoxCoxTransformOutput,
)
from gluonts.distribution import (
    DistributionOutput,
    StudentT,
    StudentTOutput,
    MultivariateGaussian,
    MultivariateGaussianOutput,
    LowrankMultivariateGaussian,
    LowrankMultivariateGaussianOutput,
    NegativeBinomial,
    NegativeBinomialOutput,
    Laplace,
    LaplaceOutput,
    Gaussian,
    GaussianOutput,
    PiecewiseLinear,
    PiecewiseLinearOutput,
    Binned,
    BinnedOutput,
)
from gluonts.distribution.transformed_distribution_output import (
    TransformedDistributionOutput
)
from gluonts.distribution.transformed_distribution import (
    TransformedDistribution
)


NUM_SAMPLES = 2000
BATCH_SIZE = 32
TOL = 0.2
START_TOL_MULTIPLE = 1

np.random.seed(1)
mx.random.seed(1)


def inv_softplus(y: NPArrayLike) -> np.ndarray:
    # y = log(1 + exp(x))  ==>  x = log(exp(y) - 1)
    return np.log(np.exp(y) - 1)


def maximum_likelihood_estimate_sgd(
    distr_output: DistributionOutput,
    samples: mx.ndarray,
    init_biases: List[mx.ndarray.NDArray] = None,
    num_epochs: PositiveInt = 5,
    learning_rate: PositiveFloat = 1e-2,
    hybridize: bool = True,
) -> Iterable[float]:
    model_ctx = mx.cpu()

    arg_proj = distr_output.get_args_proj()
    arg_proj.initialize()

    if hybridize:
        arg_proj.hybridize()

    if init_biases is not None:
        for param, bias in zip(arg_proj.proj, init_biases):
            param.params[param.prefix + "bias"].initialize(
                mx.initializer.Constant(bias), force_reinit=True
            )

    trainer = mx.gluon.Trainer(
        arg_proj.collect_params(),
        'sgd',
        {'learning_rate': learning_rate, 'clip_gradient': 10.0},
    )

    dummy_data = mx.nd.array(np.ones((len(samples), 1)))

    train_data = mx.gluon.data.DataLoader(
        mx.gluon.data.ArrayDataset(dummy_data, samples),
        batch_size=BATCH_SIZE,
        shuffle=True,
    )

    dummy_out = mx.nd.array(np.ones((1, 1)))

    for e in range(num_epochs):
        cumulative_loss = 0
        num_batches = 0
        # inner loop
        for i, (data, sample_label) in enumerate(train_data):
            # data: (batch, 1), the "1" defines the dimension of the projection layer and should be kept
            data = data.as_in_context(model_ctx)
            # sample_label: (batch,), for univariate and (batch, d) for multivariate
            if sample_label.shape[-1] == 1:
                sample_label = sample_label.as_in_context(model_ctx).squeeze(
                    axis=-1
                )
            else:
                sample_label = sample_label.as_in_context(model_ctx)

            with mx.autograd.record():
                distr_args = arg_proj(data)
                distr = distr_output.distribution(distr_args)
                loss = distr.loss(sample_label)
            loss.backward()
            trainer.step(BATCH_SIZE)
            num_batches += 1

            cumulative_loss += mx.nd.mean(loss).asscalar()
        print("Epoch %s, loss: %s" % (e, cumulative_loss / num_batches))
    return [param[0].asnumpy() for param in arg_proj(dummy_out)]


@pytest.mark.parametrize("mu_sigma_nu", [(2.3, 0.7, 6.0)])
@pytest.mark.parametrize("hybridize", [True, False])
def test_studentT_likelihood(
    mu_sigma_nu: Tuple[float, float, float], hybridize: bool
) -> None:
    '''
    Test to check that maximizing the likelihood recovers the parameters
    '''
    # test instance
    mu, sigma, nu = mu_sigma_nu

    # generate samples
    mus = mx.nd.zeros((NUM_SAMPLES, 1)) + mu
    sigmas = mx.nd.zeros((NUM_SAMPLES, 1)) + sigma
    nus = mx.nd.zeros((NUM_SAMPLES, 1)) + nu

    distr = StudentT(mus, sigmas, nus)
    samples = distr.sample()

    # nu takes very long to learn, so we initialize it at the true value.
    # transform used is softplus(x) + 2
    init_bias = [
        mu - START_TOL_MULTIPLE * TOL * mu,
        inv_softplus(sigma - START_TOL_MULTIPLE * TOL * sigma),
        inv_softplus(nu - 2),
    ]

    mu_hat, sigma_hat, nu_hat = maximum_likelihood_estimate_sgd(
        StudentTOutput(),
        samples,
        init_biases=init_bias,
        hybridize=hybridize,
        num_epochs=PositiveInt(10),
        learning_rate=1e-2,
    )

    assert (
        np.abs(mu_hat - mu) < TOL * mu
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
    assert (
        np.abs(sigma_hat - sigma) < TOL * sigma
    ), f"sigma did not match: sigma = {sigma}, sigma_hat = {sigma_hat}"
    assert (
        np.abs(nu_hat - nu) < TOL * nu
    ), "nu0 did not match: nu0 = %s, nu_hat = %s" % (nu, nu_hat)


@pytest.mark.parametrize("mu_sigma", [(1.0, 0.1)])
@pytest.mark.parametrize("hybridize", [True, False])
def test_gaussian_likelihood(mu_sigma: Tuple[float, float], hybridize: bool):
    '''
    Test to check that maximizing the likelihood recovers the parameters
    '''
    # test instance
    mu, sigma = mu_sigma

    # generate samples
    mus = mx.nd.zeros((NUM_SAMPLES, 1)) + mu
    sigmas = mx.nd.zeros((NUM_SAMPLES, 1)) + sigma

    distr = Gaussian(mus, sigmas)
    samples = distr.sample()

    init_biases = [
        mu - START_TOL_MULTIPLE * TOL * mu,
        inv_softplus(sigma - START_TOL_MULTIPLE * TOL * sigma),
    ]

    mu_hat, sigma_hat = maximum_likelihood_estimate_sgd(
        GaussianOutput(),
        samples,
        init_biases=init_biases,
        hybridize=hybridize,
        learning_rate=PositiveFloat(0.001),
        num_epochs=PositiveInt(5),
    )

    assert (
        np.abs(mu_hat - mu) < TOL * mu
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
    assert (
        np.abs(sigma_hat - sigma) < TOL * sigma
    ), f"alpha did not match: sigma = {sigma}, sigma_hat = {sigma_hat}"


@pytest.mark.timeout(10)
def test_multivariate_gaussian() -> None:
    num_samples = 2000
    dim = 2

    mu = np.arange(0, dim) / float(dim)

    L_diag = np.ones((dim,))
    L_low = 0.1 * np.ones((dim, dim)) * np.tri(dim, k=-1)
    L = np.diag(L_diag) + L_low
    Sigma = L.dot(L.transpose())

    distr = MultivariateGaussian(mu=mx.nd.array(mu), L=mx.nd.array(L))

    samples = distr.sample(num_samples)

    mu_hat, L_hat = maximum_likelihood_estimate_sgd(
        MultivariateGaussianOutput(dim=dim),
        samples,
        init_biases=None,  # todo we would need to rework biases a bit to use it in the multivariate case
        hybridize=False,
        learning_rate=PositiveFloat(0.01),
        num_epochs=PositiveInt(10),
    )

    distr = MultivariateGaussian(
        mu=mx.nd.array([mu_hat]), L=mx.nd.array([L_hat])
    )

    Sigma_hat = distr.variance[0].asnumpy()

    assert np.allclose(
        mu_hat, mu, atol=0.1, rtol=0.1
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
    assert np.allclose(
        Sigma_hat, Sigma, atol=0.1, rtol=0.1
    ), f"Sigma did not match: sigma = {Sigma}, sigma_hat = {Sigma_hat}"


@pytest.mark.timeout(10)
def test_lowrank_multivariate_gaussian() -> None:
    num_samples = 2000
    dim = 2
    rank = 1

    mu = np.arange(0, dim) / float(dim)
    D = np.eye(dim) * (np.arange(dim) / dim + 0.5)
    W = np.sqrt(np.ones((dim, rank)) * 0.2)
    Sigma = D + W.dot(W.transpose())

    distr = LowrankMultivariateGaussian(
        mu=mx.nd.array([mu]),
        D=mx.nd.array([np.diag(D)]),
        W=mx.nd.array([W]),
        dim=dim,
        rank=rank,
    )

    assert np.allclose(
        distr.variance[0].asnumpy(), Sigma, atol=0.1, rtol=0.1
    ), f"did not match: sigma = {Sigma}, sigma_hat = {distr.variance[0]}"

    samples = distr.sample(num_samples).squeeze().asnumpy()

    mu_hat, D_hat, W_hat = maximum_likelihood_estimate_sgd(
        LowrankMultivariateGaussianOutput(dim=dim, rank=rank),
        samples,
        learning_rate=PositiveFloat(0.01),
        num_epochs=PositiveInt(10),
        init_biases=None,  # todo we would need to rework biases a bit to use it in the multivariate case
        hybridize=False,
    )

    distr = LowrankMultivariateGaussian(
        dim=dim,
        rank=rank,
        mu=mx.nd.array([mu_hat]),
        D=mx.nd.array([D_hat]),
        W=mx.nd.array([W_hat]),
    )

    Sigma_hat = distr.variance.asnumpy()

    assert np.allclose(
        mu_hat, mu, atol=0.2, rtol=0.1
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"

    assert np.allclose(
        Sigma_hat, Sigma, atol=0.1, rtol=0.1
    ), f"alpha did not match: sigma = {Sigma}, sigma_hat = {Sigma_hat}"


@pytest.mark.parametrize("mu", [6.0])
@pytest.mark.parametrize("hybridize", [True, False])
def test_deterministic_l2(mu: float, hybridize: bool) -> None:
    '''
    Test to check that maximizing the likelihood recovers the parameters.
    This tests uses the Gaussian distribution with fixed variance and sample mean.
    This essentially reduces to determistic L2.
    '''
    # generate samples
    mu = mu
    mus = mx.nd.zeros(NUM_SAMPLES) + mu

    deterministic_distr = Gaussian(mu=mus, sigma=0.1 * mx.nd.ones_like(mus))
    samples = deterministic_distr.sample()

    class GaussianFixedVarianceOutput(GaussianOutput):
        @classmethod
        def domain_map(cls, F, mu, sigma):
            sigma = 0.1 * F.ones_like(sigma)
            return mu.squeeze(axis=-1), sigma.squeeze(axis=-1)

    mu_hat, _ = maximum_likelihood_estimate_sgd(
        GaussianFixedVarianceOutput(),
        samples,
        init_biases=[3 * mu, 0.1],
        hybridize=hybridize,
        num_epochs=PositiveInt(1),
    )

    assert (
        np.abs(mu_hat - mu) < TOL * mu
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"


@pytest.mark.parametrize("mu", [1.0])
@pytest.mark.parametrize("hybridize", [True, False])
def test_deterministic_l1(mu: float, hybridize: bool) -> None:
    '''
    Test to check that maximizing the likelihood recovers the parameters.
    This tests uses the Laplace distribution with fixed variance and sample mean.
    This essentially reduces to determistic L1.
    '''
    # generate samples
    mu = mu
    mus = mx.nd.zeros(NUM_SAMPLES) + mu

    class LaplaceFixedVarianceOutput(LaplaceOutput):
        @classmethod
        def domain_map(cls, F, mu, b):
            b = 0.1 * F.ones_like(b)
            return mu.squeeze(axis=-1), b.squeeze(axis=-1)

    deterministic_distr = Laplace(mu=mus, b=0.1 * mx.nd.ones_like(mus))
    samples = deterministic_distr.sample()

    mu_hat, _ = maximum_likelihood_estimate_sgd(
        LaplaceFixedVarianceOutput(),
        samples,
        init_biases=[3 * mu, 0.1],
        learning_rate=1e-3,
        hybridize=hybridize,
    )

    assert (
        np.abs(mu_hat - mu) < TOL * mu
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"


@pytest.mark.parametrize("mu_alpha", [(2.5, 0.7)])
@pytest.mark.parametrize("hybridize", [True, False])
def test_neg_binomial(mu_alpha: Tuple[float, float], hybridize: bool) -> None:
    '''
    Test to check that maximizing the likelihood recovers the parameters
    '''
    # test instance
    mu, alpha = mu_alpha

    # generate samples
    mus = mx.nd.zeros((NUM_SAMPLES, 1)) + mu
    alphas = mx.nd.zeros((NUM_SAMPLES, 1)) + alpha

    neg_bin_distr = NegativeBinomial(mu=mus, alpha=alphas)
    samples = neg_bin_distr.sample()

    init_biases = [
        inv_softplus(mu - START_TOL_MULTIPLE * TOL * mu),
        inv_softplus(alpha + START_TOL_MULTIPLE * TOL * alpha),
    ]

    mu_hat, alpha_hat = maximum_likelihood_estimate_sgd(
        NegativeBinomialOutput(),
        samples,
        hybridize=hybridize,
        init_biases=init_biases,
        num_epochs=PositiveInt(15),
    )

    assert (
        np.abs(mu_hat - mu) < TOL * mu
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
    assert (
        np.abs(alpha_hat - alpha) < TOL * alpha
    ), f"alpha did not match: alpha = {alpha}, alpha_hat = {alpha_hat}"


@pytest.mark.timeout(10)
@pytest.mark.parametrize("mu_b", [(3.3, 0.7)])
@pytest.mark.parametrize("hybridize", [True, False])
def test_laplace(mu_b: Tuple[float, float], hybridize: bool) -> None:
    '''
    Test to check that maximizing the likelihood recovers the parameters
    '''
    # test instance
    mu, b = mu_b

    # generate samples
    mus = mx.nd.zeros((NUM_SAMPLES, 1)) + mu
    bs = mx.nd.zeros((NUM_SAMPLES, 1)) + b

    laplace_distr = Laplace(mu=mus, b=bs)
    samples = laplace_distr.sample()

    init_biases = [
        mu - START_TOL_MULTIPLE * TOL * mu,
        inv_softplus(b + START_TOL_MULTIPLE * TOL * b),
    ]

    mu_hat, b_hat = maximum_likelihood_estimate_sgd(
        LaplaceOutput(), samples, hybridize=hybridize, init_biases=init_biases
    )

    assert (
        np.abs(mu_hat - mu) < TOL * mu
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
    assert (
        np.abs(b_hat - b) < TOL * b
    ), f"b did not match: b = {b}, b_hat = {b_hat}"


@pytest.mark.parametrize(
    "gamma, slopes, knot_spacings",
    [(2.0, np.array([3, 1, 3, 4]), np.array([0.3, 0.2, 0.35, 0.15]))],
)
@pytest.mark.parametrize("hybridize", [True, False])
def test_piecewise_linear(
    gamma: float,
    slopes: np.ndarray,
    knot_spacings: np.ndarray,
    hybridize: bool,
) -> None:
    '''
    Test to check that minimizing the CRPS recovers the quantile function
    '''
    num_samples = 500  # use a few samples for timeout failure

    gammas = mx.nd.zeros((num_samples,)) + gamma
    slopess = mx.nd.zeros((num_samples, len(slopes))) + mx.nd.array(slopes)
    knot_spacingss = mx.nd.zeros(
        (num_samples, len(knot_spacings))
    ) + mx.nd.array(knot_spacings)

    pwl_sqf = PiecewiseLinear(gammas, slopess, knot_spacingss)

    samples = pwl_sqf.sample()

    # Parameter initialization
    gamma_init = gamma - START_TOL_MULTIPLE * TOL * gamma
    slopes_init = slopes - START_TOL_MULTIPLE * TOL * slopes
    knot_spacings_init = knot_spacings
    # We perturb knot spacings such that even after the perturbation they sum to 1.
    mid = len(slopes) // 2
    knot_spacings_init[:mid] = (
        knot_spacings[:mid] - START_TOL_MULTIPLE * TOL * knot_spacings[:mid]
    )
    knot_spacings_init[mid:] = (
        knot_spacings[mid:] + START_TOL_MULTIPLE * TOL * knot_spacings[mid:]
    )

    init_biases = [gamma_init, slopes_init, knot_spacings_init]

    # check if it returns original parameters of mapped
    gamma_hat, slopes_hat, knot_spacings_hat = maximum_likelihood_estimate_sgd(
        PiecewiseLinearOutput(len(slopes)),
        samples,
        init_biases=init_biases,
        hybridize=hybridize,
        learning_rate=PositiveFloat(0.01),
        num_epochs=PositiveInt(20),
    )

    # Since the problem is highly non-convex we may not be able to recover the exact parameters
    # Here we check if the estimated parameters yield similar function evaluations at different quantile levels.
    quantile_levels = np.arange(0.1, 1.0, 0.1)

    # create a LinearSplines instance with the estimated parameters to have access to .quantile
    pwl_sqf_hat = PiecewiseLinear(
        mx.nd.array(gamma_hat),
        mx.nd.array(slopes_hat).expand_dims(axis=0),
        mx.nd.array(knot_spacings_hat).expand_dims(axis=0),
    )

    # Compute quantiles with the estimated parameters
    quantiles_hat = np.squeeze(
        pwl_sqf_hat.quantile(
            mx.nd.array(quantile_levels).expand_dims(axis=0), axis=1
        ).asnumpy()
    )

    # Compute quantiles with the original parameters
    # Since params is replicated across samples we take only the first entry
    quantiles = np.squeeze(
        pwl_sqf.quantile(
            mx.nd.array(quantile_levels)
            .expand_dims(axis=0)
            .repeat(axis=0, repeats=num_samples),
            axis=1,
        ).asnumpy()[0, :]
    )

    for ix, (quantile, quantile_hat) in enumerate(
        zip(quantiles, quantiles_hat)
    ):
        assert np.abs(quantile_hat - quantile) < TOL * quantile, (
            f"quantile level {quantile_levels[ix]} didn't match:"
            f" "
            f"q = {quantile}, q_hat = {quantile_hat}"
        )


@pytest.mark.skip("this test fails when run locally")
@pytest.mark.timeout(20)
@pytest.mark.parametrize("lambdas", [(0.1, 0.01)])
@pytest.mark.parametrize("mu_sigma", [(-1.5, 0.5)])
@pytest.mark.parametrize("hybridize", [True])
def test_box_cox_tranform(
    lambdas: Tuple[float, float],
    mu_sigma: Tuple[float, float],
    hybridize: bool,
):
    '''
    Test to check that maximizing the likelihood recovers the parameters
    '''
    # test instance
    lam_1, lam_2 = lambdas
    mu, sigma = mu_sigma

    # generate samples
    lamdas_1 = mx.nd.zeros((NUM_SAMPLES,)) + lam_1
    lamdas_2 = mx.nd.zeros((NUM_SAMPLES,)) + lam_2
    transform = InverseBoxCoxTransform(lamdas_1, lamdas_2)

    mus = mx.nd.zeros((NUM_SAMPLES,)) + mu
    sigmas = mx.nd.zeros((NUM_SAMPLES,)) + sigma
    gausian_distr = Gaussian(mus, sigmas)

    # Here the base distribution is Guassian which is transformed to
    # non-Gaussian via the inverse Box-Cox transform.
    # Sampling from `trans_distr` gives non-Gaussian samples
    trans_distr = TransformedDistribution(gausian_distr, transform)

    # Given the non-Gaussian samples find the true parameters
    # of the Box-Cox transformation as well as the underlying Gaussian distribution.
    samples = trans_distr.sample()

    init_biases = [
        mu - START_TOL_MULTIPLE * TOL * mu,
        inv_softplus(sigma - START_TOL_MULTIPLE * TOL * sigma),
        lam_1 - START_TOL_MULTIPLE * TOL * lam_1,
        inv_softplus(lam_2 - START_TOL_MULTIPLE * TOL * lam_2),
    ]

    mu_hat, sigma_hat, lam_1_hat, lam_2_hat = maximum_likelihood_estimate_sgd(
        TransformedDistributionOutput(
            GaussianOutput(),
            InverseBoxCoxTransformOutput(lb_obs=lam_2, fix_lambda_2=True),
        ),
        samples,
        init_biases=init_biases,
        hybridize=hybridize,
        learning_rate=PositiveFloat(0.01),
        num_epochs=PositiveInt(18),
    )

    assert (
        np.abs(lam_1_hat - lam_1) < TOL * lam_1
    ), f"lam_1 did not match: lam_1 = {lam_1}, lam_1_hat = {lam_1_hat}"
    # assert (
    #     np.abs(lam_2_hat - lam_2) < TOL * lam_2
    # ), f"lam_2 did not match: lam_2 = {lam_2}, lam_2_hat = {lam_2_hat}"

    assert np.abs(mu_hat - mu) < TOL * np.abs(
        mu
    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
    assert (
        np.abs(sigma_hat - sigma) < TOL * sigma
    ), f"sigma did not match: sigma = {sigma}, sigma_hat = {sigma_hat}"


@pytest.mark.parametrize("num_bins", [6])
@pytest.mark.parametrize(
    "bin_probabilites", [np.array([0.3, 0.1, 0.05, 0.2, 0.1, 0.25])]
)
@pytest.mark.parametrize(
    "hybridize", [False]
)  # TODO: hybridize does not work in the test
def test_binned_likelihood(
    num_bins: float, bin_probabilites: np.ndarray, hybridize: bool
):
    '''
    Test to check that maximizing the likelihood recovers the parameters
    '''

    bin_prob = mx.nd.array(bin_probabilites)
    bin_center = mx.nd.array(np.logspace(-1, 1, num_bins))

    # generate samples
    bin_probs = mx.nd.zeros((NUM_SAMPLES, num_bins)) + bin_prob
    bin_centers = mx.nd.zeros((NUM_SAMPLES, num_bins)) + bin_center

    distr = Binned(bin_probs, bin_centers)
    samples = distr.sample()

    # add some jitter to the uniform initialization and normalize
    bin_prob_init = mx.nd.random_uniform(1 - TOL, 1 + TOL, num_bins) * bin_prob
    bin_prob_init = bin_prob_init / bin_prob_init.sum()

    init_biases = [bin_prob_init]

    bin_prob_hat, _ = maximum_likelihood_estimate_sgd(
        BinnedOutput(list(bin_center.asnumpy())),
        samples,
        init_biases=init_biases,
        hybridize=hybridize,
        learning_rate=PositiveFloat(0.05),
        num_epochs=PositiveInt(25),
    )

    assert all(
        mx.nd.abs(mx.nd.array(bin_prob_hat) - bin_prob) < TOL * bin_prob
    ), f"bin_prob did not match: bin_prob = {bin_prob}, bin_prob_hat = {bin_prob_hat}"
