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

pio.templates.default = "simple_white"


def test_univariate_gaussian():
    # Question 1 - Draw samples and print fitted model
    actual_mean = 10
    actual_var = 1
    num_of_samples = 1000
    normal_arr_10_1 = np.random.normal(actual_mean, actual_var, num_of_samples)
    estimator = UnivariateGaussian()
    estimator.fit(normal_arr_10_1)
    print(estimator.mu_, estimator.var_)

    # Question 2 - Empirically showing sample mean is consistent
    errors_from_mean = list()
    step = 10
    start = 10
    for i in range(start, num_of_samples + 1, step):
        estimator.fit(normal_arr_10_1[0: i])
        error = abs(estimator.mu_ - actual_mean)
        errors_from_mean.append(error)

    fig = px.scatter(x=[i for i in range(start, num_of_samples + 1, step)], y=errors_from_mean)
    fig.update_layout(title_text="Error of estimated mean by number of samples", title_x=0.5,
                      xaxis_title="Number of samples",
                      yaxis_title="Error of the estimated mean")
    fig.show()

    # Question 3 - Plotting Empirical PDF of fitted model
    pdf_results = estimator.pdf(normal_arr_10_1)
    fig2 = px.scatter(x=[i for i in range(1, 1001)], y=pdf_results)
    fig2.update_layout(title_text="PDF of the randomly drawn samples", title_x=0.5,
                       xaxis_title="Sample number",
                       yaxis_title="PDF value of the sample")
    fig2.show()


def test_multivariate_gaussian():
    # Question 4 - Draw samples and print fitted model
    num_of_samples = 1000
    mu_4d = np.array([0, 0, 4, 0])
    cov_4d = np.array([[1, 0.2, 0, 0.5],
                       [0.2, 2, 0, 0],
                       [0, 0, 1, 0],
                       [0.5, 0, 0, 1]])
    mult_gaus_samples_10_1 = np.random.multivariate_normal(mu_4d, cov_4d, num_of_samples)

    multi_gaus = MultivariateGaussian()
    multi_gaus.fit(mult_gaus_samples_10_1)
    print(multi_gaus.mu_)
    print(multi_gaus.cov_)

    # Question 5 - Likelihood evaluation
    num_of_linspace_samples = 200
    linspace_res = np.linspace(-10, 10, num_of_linspace_samples)
    print("linspace_res is: ", linspace_res) #todo - remove, for testing
    for i in range(num_of_linspace_samples):
        pass

    # Question 6 - Maximum likelihood
    # raise NotImplementedError()


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