from IMLearn.learners import UnivariateGaussian, MultivariateGaussian
import numpy as np
import plotly.graph_objects as go
import plotly.io as pio

pio.templates.default = "simple_white"
MU = 10
VAR = 1
NUM_SAMPLES = 1000


def test_univariate_gaussian():
    # Question 1 - Draw samples and print fitted model

    gauss = UnivariateGaussian()
    samples = np.random.normal(loc=MU, scale=VAR, size=NUM_SAMPLES)
    gauss = gauss.fit(samples)
    expectation, variance = gauss.mu_, gauss.var_
    print(f"({expectation}, {variance})")

    # Question 2 - Empirically showing sample mean is consistent

    xs = list(range(10, NUM_SAMPLES + 1, 10))
    fitted_samples, abs_dist = [], []
    for i in range(10, NUM_SAMPLES + 1, 10):
        abs_dist.append(np.abs(gauss.fit(samples[:i]).mu_ - expectation))
    fig = go.Figure([go.Scatter(x=xs, y=abs_dist, mode='markers+lines', name=r'$Distance$'),
                     go.Scatter(x=xs, y=[0] * len(xs), mode='lines', name=r'$SmallestDistance$')],
                    layout=go.Layout(title=r"$\text{Distance between the estimated and true value of the expectation:}$",
                                     xaxis_title="$\\text{Number Of Samples}$",
                                     yaxis_title="$\\text{Distance From True Expectation}$"))
    fig.show()

    # Question 3 - Plotting Empirical PDF of fitted model

    pdfs = gauss.pdf(samples)
    fig = go.Figure([go.Scatter(x=samples, y=pdfs, mode='markers', name=r'$PDF values'),
                     go.Scatter(x=samples, y=[0] * len(samples), mode='lines', name=r'$Zero$')],
                    layout=go.Layout(title=r"$\text{PDF Function Values:}$",
                                     xaxis_title="$\\text{Samples}$",
                                     yaxis_title="$\\text{pdf value}$"))
    fig.show()


def test_multivariate_gaussian():
    # Question 4 - Draw samples and print fitted model

    in_cov = [[1, 0.2, 0, 0.5],
              [0.2, 2, 0, 0],
              [0, 0, 1, 0],
              [0.5, 0, 0, 1]]

    in_mu = [0,
             0,
             4,
             0]

    samples = np.random.multivariate_normal(mean=in_mu, cov=in_cov, size=NUM_SAMPLES)
    gauss = MultivariateGaussian()
    gauss = gauss.fit(samples)
    print(gauss.mu_)
    print(gauss.cov_)

    # Question 5 - Likelihood evaluation

    xs = np.linspace(-10, 10, 200)
    likelihoods = []
    maximum = -np.inf
    max_mu = None
    for i in range(len(xs)):
        f1 = xs[i]
        for j in range(len(xs)):
            f3 = xs[j]
            mu = np.array([f1, 0, f3, 0])
            cur = MultivariateGaussian.log_likelihood(mu, np.array(in_cov), samples)
            if cur > maximum:
                maximum = cur
                max_mu = mu
            likelihoods.append(cur)

    xaxis, yaxis = np.tile(xs, 200), np.repeat(xs, 200)
    fig = go.Figure(data=go.Heatmap(z=likelihoods, x=xaxis, y=yaxis),
                    layout=go.Layout(title="$\\text{Multivariate Gaussian Log-Likelihood}$",
                                     xaxis_title="r$\\text{values of feature}\f_1$",
                                     yaxis_title="r$\\text{values of feature}\f_2$"))
    fig.show()

    # Question 6 - Maximum likelihood

    maximum_index = np.argmax(likelihoods)
    print(f"max log_likelihood: {likelihoods[maximum_index]:.3f}")
    print("max mu: " + str(max_mu))


if __name__ == '__main__':
    np.random.seed(0)
    test_univariate_gaussian()
    test_multivariate_gaussian()

    # X = np.array([1, 5, 2, 3, 8, -4, -2, 5, 1, 10, -10,
    #               4, 5, 2, 7, 1, 1, 3, 2, -1, -3, 1, -4,
    #               1, 2, 1, -4, -4, 1, 3, 2, 6, -6, 8, 3,
    #               -6, 4, 1, -2, 3, 1, 4, 1, 4, -2, 3, -1,
    #               0, 3, 5, 0, -2])
    #
    # print(UnivariateGaussian.log_likelihood(1, 1, X))
    # print(UnivariateGaussian.log_likelihood(10, 1, X))
