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

pio.templates.default = "plotly"


def test_univariate_gaussian():
    # Question 1 - Draw samples and print fitted model
    mu = 10
    var = 1

    samples = np.random.normal(mu, var, 1000)
    uni_gaussian = UnivariateGaussian()
    uni_gaussian.fit(samples)
    print("\nEstimated Expectation and Variance:")
    print((uni_gaussian.mu_, uni_gaussian.var_))

    # Question 2 - Empirically showing sample mean is consistent
    distance = np.empty(100)
    nums = np.arange(10, 1010, 10)

    for i, num in enumerate(nums):
        uni_gaussian.fit(samples[:num])
        distance[i] = abs(mu - uni_gaussian.mu_)

    plot1 = go.Figure(data=[go.Scatter(x=nums, y=distance, mode='markers')])
    plot1.update_layout(title_text='Distance of Empirical Mean from Actual Mean for Different Sample Sizes') \
        .update_yaxes(title_text='Distance of Empirical Mean from Actual Mean') \
        .update_xaxes(title_text='Sample Size')
    plot1.show()

    # Question 3 - Plotting Empirical PDF of fitted model
    pdfs = uni_gaussian.pdf(samples)

    plot2 = go.Figure(data=[go.Scatter(x=samples, y=pdfs, mode='markers')])
    plot2.update_layout(title_text='PDF Values of Random Samples') \
        .update_yaxes(title_text='PDF Value') \
        .update_xaxes(title_text='Sample Value')
    plot2.show()


def test_multivariate_gaussian():
    # Question 4 - Draw samples and print fitted model
    mu = np.array([0, 0, 4, 0])
    cov = 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, cov, 1000)
    multi_gaussian = MultivariateGaussian()
    multi_gaussian.fit(samples)
    print("\nEstimated Expectation:")
    print(multi_gaussian.mu_)
    print("\nEstimated Covariance:")
    print(multi_gaussian.cov_)

    # Question 5 - Likelihood evaluation
    values = np.linspace(-10, 10, 200)

    log_likelihood = np.empty((200, 200))

    max_likelihood, max_pair = None, None
    for i, f1 in enumerate(values):
        for j, f3 in enumerate(values):
            new_mu = np.array([f1, 0, f3, 0])
            log_likelihood[i][j] = multi_gaussian.log_likelihood(new_mu, cov, samples)

            if max_likelihood is None or log_likelihood[i][j] > max_likelihood:
                max_likelihood = log_likelihood[i][j]
                max_pair = (f1, f3)

    fig = go.Figure(data=go.Heatmap(x=values, y=values, z=log_likelihood)) \
        .update_layout(title_text="Likelihood of a Given Sample for Different Expectations") \
        .update_yaxes(title_text="f1 value") \
        .update_xaxes(title_text="f3 value")

    fig.show()

    # Question 6 - Maximum likelihood
    print("\nMax log-likelihood: ", round(max_likelihood, 3))
    print("Max pair:", (round(max_pair[0], 3), round(max_pair[1], 3)))


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