import numpy as np
from scipy.stats import chi2, norm, invgamma
from scipy.special import gamma
import math
import unittest
import matplotlib.pyplot as plt


def inv_chi2(x, nu):
    """
    Inverse chi-squared probability distribution function.

    http://artax.karlin.mff.cuni.cz/r-help/library/geoR/html/InvChisquare.html
    f(x) = 1 / (2^(n/2) Gamma(n/2)) (1/x)^(n/2-1) e^(-1/(2x))

    :param x: value
    :param nu:
    :return:
    PDF value
    """
    a = 2.0 ** (-nu/2.0) / gamma(nu/2.0)
    # b = x ** (-nu/2.0 + 1.0)
    b = (1.0/x) ** (nu/2.0 - 1.0)
    p = a * b * math.exp(-1.0/(2.0*x))
    return p


class TestDistr(unittest.TestCase):
    def test_inv_chi2_pdf(self):
        """
        Test to the Octave function.
        >> format long
        >> chi2pdf(1.5^-1, 4.0)
           ans =  0.119421885095632

        Inverse-chi-square distribution Calculator
        http://keisan.casio.com/exec/system/1304908316 gives
        gives (1.5, 4) -> 0.0530763933


        R
        > library(geoR)
        > dinvchisq(1.5, df=4)
        > ans=0.05307639

        """
        octave_chi2pdf = 0.119421885095632
        online_calculator = 0.0530763933
        r_value = 0.05307639
        v1 = 1.5
        v2 = 4.0
        my_func = inv_chi2(v1, v2)
        using_inv_gamma = invgamma.pdf(v1, v2/2.0, 0.5)
        print("... Value obtained using inv-gamma :", using_inv_gamma)
        scipy_chi2pdf = chi2.pdf(v1**-1, v2)
        print("... Value obtained using chi2 :", scipy_chi2pdf)
        self.assertAlmostEqual(octave_chi2pdf, scipy_chi2pdf, places=5)
        self.assertAlmostEqual(my_func, octave_chi2pdf, places=5)
        # self.assertAlmostEqual(my_func, online_calculator, places=5)


def prior_sigma(x, nu0, sigma0=0):
    # p = chi2.pdf(nu0**-1, sigma0**2)
    p = invgamma.pdf(x, nu0, sigma0**2)
    return p


def test_plot_prior_sigma():
    xs = np.linspace(0, 5.0, 100)
    ps = prior_sigma(xs, 0.1, 0.1)
    plt.plot(xs, ps)
    plt.show()


def prior_mu(x, mu0, sigma_sq, kappa0):
    p = norm.pdf(x, mu0, sigma_sq**2/kappa0)
    return p


def test_plot_prior_mu():
    xs = np.linspace(0, 5.0, 100)
    ps = prior_mu(xs, 3.0, 1.1, 2.0)
    plt.plot(xs, ps)
    plt.show()


def posterior_sigma(x, mu0, nu0, sigma0, y, kappa0):
    n = len(y)
    bar_y = np.mean(y)
    nu_n = nu0 + n
    s = np.std(y)
    sigma_n_squared = (nu0 * sigma0**2 + (n-1.0) * s**2 + (kappa0 * n) * (bar_y - mu0) / (kappa0 + n)) / nu_n
    p_squared = chi2.pdf(nu_n ** (-1.0), sigma_n_squared)
    p = np.sqrt(p_squared)
    return p


def test_plot_posteriro_sigma():
    xs = np.linspace(0, 5.0, 100)
    y = []


if __name__ == "__main__":
    # test_plot_prior_mu()
    # test_plot_prior_sigma()
    unittest.main()
    # test_prior()


# def mu_prior():
#     """
#     Prior PDF for \mu.
#     :return:
#     """
#     pass
#
#
# def sigma_prior(nu0, sigma0):
#     """
#     Prior PDF for \sigma is given by the inverse-chi-squared distribution.
#     See Gelman et.al., Section 3.3 (p.78)
#     Example output: https://nl.mathworks.com/help/stats/chi2pdf.html
#
#     :return:
#     """
#     pass
