<gh_stars>0
import numpy
from plotly.subplots import make_subplots
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"


def test_univariate_gaussian():

    # Question 1 - Draw samples and print fitted model
    samples = np.random.normal(10, 1, 1000)
    univariate_gaussian = UnivariateGaussian()
    univariate_gaussian.fit(samples)
    print((univariate_gaussian.mu_, univariate_gaussian.var_))

    # Question 2 - Empirically showing sample mean is consistent
    axis_x = np.linspace(10, 1000, 100)
    results = numpy.zeros(100)
    counter = 0
    for i in range(10, 1001, 10):
        results[counter] = np.abs(np.mean(samples[:i]) - 10)
        counter = counter + 1

    fig = make_subplots(rows=1, cols=1) \
        .add_traces([go.Scatter(x=axis_x, y=results, mode='lines', marker=dict(color="black"),
                                showlegend=False)]) \
        .update_layout(title_text=r"$\text{Univariate Gaussian Estimation - Question 2}$",
                       height=800)
    fig.update_xaxes(title_text="Sample size", row=1, col=1)
    fig.update_yaxes(title_text="Absolute distance", row=1, col=1)
    fig.show()

    # Question 3 - Plotting Empirical PDF of fitted model
    np.sort(samples)
    pdf_results = univariate_gaussian.pdf(samples)
    fig = make_subplots(rows=1, cols=1) \
        .add_traces([go.Scatter(x=samples, y=pdf_results, mode='markers',
                                marker=dict(color="red"), showlegend=False)])
    fig.update_xaxes(title_text="Sample value", row=1, col=1)
    fig.update_yaxes(title_text="PDF Value", row=1, col=1)
    fig.update_layout(title_text=r"$\text{Empirical PDF Function- Question 3}$", height=800)
    fig.show()


def test_multivariate_gaussian():

    # Question 4 - Draw samples and print fitted model
    multivariate_gaussian = MultivariateGaussian()
    mu = np.array([0, 0, 4, 0])
    sigma = np.array([[1, 0.2, 0, 0.5], [0.2, 2, 0, 0], [0, 0, 1, 0], [0.5, 0, 0, 1]])
    samples = np.random.multivariate_normal(mu, sigma, 1000)
    multivariate_gaussian.fit(samples)
    print(multivariate_gaussian.mu_)
    print(multivariate_gaussian.cov_)

    # Question 5 - Likelihood evaluation
    f1 = np.linspace(-10, 10, 200)
    f3 = np.linspace(-10, 10, 200)
    results = np.zeros([200, 200])
    for i in range(0, 200):
        for j in range(0, 200):
            mu = np.array([f1[i], 0, f3[j], 0])
            results[i][j] = multivariate_gaussian.log_likelihood(mu, sigma, samples)

    fig = make_subplots(rows=1, cols=1).add_traces(go.Heatmap(x=f1, y=f3, z=results))
    fig.update_xaxes(title_text="f1 values", row=1, col=1)
    fig.update_yaxes(title_text="f3 values", row=1, col=1)
    fig.update_layout(title_text=r"$\text{HeatMap - Question 5}$", height=800)
    fig.show()

    # Question 6 - Maximum likelihood

    indices = np.unravel_index(np.argmax(results, axis=None), results.shape)
    print("The model achieved the maximum log-likelihood value is:")
    print("f1 = " + str(f1[indices[0]]))
    print("f3 = " + str(f1[indices[1]]))


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