import pandas as pd
import numpy as np
import numpy.testing as nt
import torch
import scipy
from sincei.ExponentialFamily import Bernoulli, Gaussian, Poisson, Beta, SigmoidBeta, Gamma


def testGaussianDistribution():
    # Verifies compliance of Gaussian definition using randomly selected points.
    X = torch.linspace(-5, 5, 500)
    for _ in range(25):
        # random parameter (mean)
        theta = torch.normal(0.0, 1.0, (1, 1))[0]

        # Gaussian PDF generated by our implementation
        gaussian_dist = Gaussian()
        dist_gauss_values = np.array([gaussian_dist.distribution(x, theta).detach().numpy() for x in X]).flatten()

        # Scipy Gaussian distribution
        scipy_gauss_val = scipy.stats.norm.pdf(X, loc=theta)

        # Assert equality
        nt.assert_array_almost_equal(dist_gauss_values, scipy_gauss_val, decimal=3)


def testBernoulliDistribution():
    # Verifies compliance of Bernoulli definition using randomly selected points.
    X = torch.Tensor([0, 1])

    # Span possible values from -30 to 30
    theta_val = torch.linspace(-30, 30, 1000)
    for theta in theta_val:
        bernoulli_dist = Bernoulli()
        dist_bernoulli_values = np.array([bernoulli_dist.distribution(x, theta).detach().numpy() for x in X]).flatten()

        # Compare to theoretical value
        assert np.abs(dist_bernoulli_values[-1] - 1 / (1 + torch.exp(-theta))) < 1e-6


def testPoissonDistribution():
    # Verifies compliance of Poisson definition using randomly selected points.
    X = torch.linspace(0, 100, 101)

    # Span possible values from -30 to 30
    theta_val = torch.linspace(-50, 5, 10)
    for theta in theta_val:
        poisson_dist = Poisson()
        log_dist_poisson_values = np.array(
            [poisson_dist.log_distribution(x, theta).detach().numpy() for x in X]
        ).flatten()

        # Compare to Scipy values
        np.testing.assert_array_almost_equal(
            log_dist_poisson_values, scipy.stats.poisson.logpmf(X, np.exp(theta)), decimal=3
        )


def testBetaDistribution():
    # Verifies compliance of Beta definition using randomly selected points.
    X = torch.linspace(0 + 1e-4, 1 - 1e-4, 100)

    # Possible parameters: scan through them
    theta_val = torch.linspace(0.01, 0.99, 10)
    nu_val = torch.rand(20) * 10
    for nu in nu_val:
        for theta in theta_val:
            # Our definition
            beta_dist = Beta()
            beta_dist.family_params["nu"] = torch.Tensor([nu])
            log_dist_beta_values = np.array(
                [beta_dist.log_distribution(x, theta).detach().numpy() for x in X]
            ).flatten()
            dist_beta_values = np.array([beta_dist.distribution(x, theta).detach().numpy() for x in X]).flatten()

            # Compare to Scipy.
            np.testing.assert_array_almost_equal(
                log_dist_beta_values, scipy.stats.beta.logpdf(X, theta * nu, nu * (1 - theta)), decimal=2
            )


def testSigmoidBetaDistribution():
    # Verifies compliance of SigmoidBeta definition using randomly selected points.
    X = torch.linspace(0 + 1e-4, 1 - 1e-4, 100)
    # Possible parameters
    theta_val = torch.logit(torch.linspace(0.01, 0.99, 10))
    nu_val = torch.rand(20) * 10

    for nu in nu_val:
        for theta in theta_val:
            # Our definition
            beta_dist = SigmoidBeta()
            beta_dist.family_params["nu"] = torch.Tensor([nu])
            log_dist_beta_values = np.array(
                [beta_dist.log_distribution(x, theta).detach().numpy() for x in X]
            ).flatten()
            dist_beta_values = np.array([beta_dist.distribution(x, theta).detach().numpy() for x in X]).flatten()

            # Comparison to Scipy.
            np.testing.assert_array_almost_equal(
                log_dist_beta_values,
                scipy.stats.beta.logpdf(X, torch.sigmoid(theta) * nu, nu * (1 - torch.sigmoid(theta))),
                decimal=2,
            )


def testGammaDistribution():
    # Verifies compliance of Gamma definition using randomly selected points.
    X = torch.logspace(-8, 5, 1000)
    theta_val = torch.logspace(-5, 3, 50)

    # Possible parameters
    nu_val = torch.rand(20) * 10

    for nu in nu_val:
        for theta in theta_val:
            # Our definition
            gamma_dist = Gamma()
            gamma_dist.family_params["nu"] = torch.Tensor([nu])
            log_dist_gamma_values = np.array(
                [gamma_dist.distribution(torch.Tensor(x), torch.Tensor([theta])).detach().numpy() for x in X]
            ).flatten()

            # Comparison to Scipy.
            np.testing.assert_array_almost_equal(
                scipy.stats.gamma.pdf(X, a=theta + 1, loc=0, scale=1 / nu), log_dist_gamma_values, decimal=2
            )
