import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

from sklearn.linear_model import ARDRegression, LinearRegression

# Parameters of the example
np.random.seed(0)
n_samples, n_features = 100, 100
# Create Gaussian data
X = np.random.randn(n_samples, n_features)
# Create weights with a precision lambda_ of 4.
lambda_ = 4.
w = np.zeros(n_features)
# Only keep 10 weights of interest
relevant_features = np.random.randint(0, n_features, 10)
for i in relevant_features:
    w[i] = stats.norm.rvs(loc=0, scale=1. / np.sqrt(lambda_))
# Create noise with a precision alpha of 50.
alpha_ = 50.
noise = stats.norm.rvs(loc=0, scale=1. / np.sqrt(alpha_), size=n_samples)
# Create the target<
y = np.dot(X, w) + noise

clf = ARDRegression(fit_intercept=False, n_iter=1000)
clf.fit(X, y)

ols = LinearRegression(fit_intercept=False)
ols.fit(X, y)

from copy import deepcopy

from sds.distributions.lingauss import SingleOutputLinearGaussianWithKnownPrecision
from sds.distributions.lingauss import SingleOutputLinearGaussianWithKnownMean
from sds.distributions.gaussian import GaussianWithPrecision
from sds.distributions.gaussian import GaussianWithKnownMeanAndDiagonalPrecision
from sds.distributions.gamma import Gamma

likelihood_precision_prior = Gamma(dim=1, alphas=np.ones((1, )),
                                   betas=1e-6 * np.ones((1, )))

parameter_precision_prior = Gamma(dim=n_features, alphas=np.ones((n_features, )),
                                  betas=1e-6 * np.ones((n_features, )))

likelihood_precision_posterior = deepcopy(likelihood_precision_prior)
parameter_precision_posterior = deepcopy(parameter_precision_prior)
parameter_posterior = None

for i in range(100):
    # parameter posterior
    alphas = parameter_precision_posterior.mean()
    parameter_prior = GaussianWithPrecision(dim=n_features,
                                            mu=np.zeros((n_features, )),
                                            lmbda=np.diag(alphas))
    parameter_posterior = deepcopy(parameter_prior)

    beta = likelihood_precision_posterior.mean()
    likelihood_known_precision = SingleOutputLinearGaussianWithKnownPrecision(column_dim=n_features,
                                                                              lmbda=beta,
                                                                              affine=False)

    stats = likelihood_known_precision.statistics(X, y)
    parameter_posterior.nat_param = parameter_prior.nat_param + stats

    # likelihood precision posterior
    param = parameter_posterior.mean()
    likelihood_known_mean = SingleOutputLinearGaussianWithKnownMean(column_dim=n_features,
                                                                    W=param, affine=False)

    stats = likelihood_known_mean.statistics(X, y)
    likelihood_precision_posterior.nat_param = likelihood_precision_prior.nat_param + stats

    # parameter precision posterior
    parameter_likelihood = GaussianWithKnownMeanAndDiagonalPrecision(dim=n_features)

    param = parameter_posterior.mean()
    stats = parameter_likelihood.statistics(param)
    parameter_precision_posterior.nat_param = parameter_precision_prior.nat_param + stats

our_ard = parameter_posterior.mode()

from sds.distributions.composite import MatrixNormalGamma
from sds.distributions.lingauss import LinearGaussianWithDiagonalPrecision

M = np.zeros((1, n_features))
K = 1e-16 * np.eye(n_features)
alphas = 1e-16 * np.ones((1, ))
betas = 1e-16 * np.ones((1, ))

prior = MatrixNormalGamma(column_dim=n_features, row_dim=1,
                          M=M, K=K, alphas=alphas, betas=betas)

posterior = deepcopy(prior)
likelihood = LinearGaussianWithDiagonalPrecision(column_dim=n_features,
                                                 row_dim=1,
                                                 affine=False)

stats = likelihood.statistics(X, np.atleast_2d(y).T)
posterior.nat_param = prior.nat_param + stats
our_ols = posterior.mode()[0]

plt.figure(figsize=(6, 5))
plt.title("Weights of the model")
plt.plot(w, color='orange', linestyle='-', linewidth=2, label="Ground truth")
plt.plot(clf.coef_, color='darkblue', linestyle='-', linewidth=2, label="Sklearn ARD")
plt.plot(our_ard, color='red', linestyle='-', linewidth=2, label="Our ARD")
# plt.plot(ols.coef_, color='yellowgreen', linestyle=':', linewidth=2, label="Sklearn OLS")
# plt.plot(our_ols.flatten(), color='cyan', linestyle='-', linewidth=2, label="Our OLS")

plt.xlabel("Features")
plt.ylabel("Values of the weights")
plt.legend(loc=1)

plt.show()

[(8, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (11, 'numpy.random.randn', 'np.random.randn', True, 'import numpy as np\n'), (14, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (16, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (25, 'sklearn.linear_model.ARDRegression', 'ARDRegression', False, 'from sklearn.linear_model import ARDRegression, LinearRegression\n'), (28, 'sklearn.linear_model.LinearRegression', 'LinearRegression', False, 'from sklearn.linear_model import ARDRegression, LinearRegression\n'), (45, 'copy.deepcopy', 'deepcopy', False, 'from copy import deepcopy\n'), (46, 'copy.deepcopy', 'deepcopy', False, 'from copy import deepcopy\n'), (85, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (90, 'sds.distributions.composite.MatrixNormalGamma', 'MatrixNormalGamma', False, 'from sds.distributions.composite import MatrixNormalGamma\n'), (93, 'copy.deepcopy', 'deepcopy', False, 'from copy import deepcopy\n'), (94, 'sds.distributions.lingauss.LinearGaussianWithDiagonalPrecision', 'LinearGaussianWithDiagonalPrecision', False, 'from sds.distributions.lingauss import LinearGaussianWithDiagonalPrecision\n'), (102, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (103, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (104, 'matplotlib.pyplot.plot', 'plt.plot', True, 'import matplotlib.pyplot as plt\n'), (105, 'matplotlib.pyplot.plot', 'plt.plot', True, 'import matplotlib.pyplot as plt\n'), (106, 'matplotlib.pyplot.plot', 'plt.plot', True, 'import matplotlib.pyplot as plt\n'), (110, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (111, 'matplotlib.pyplot.ylabel', 'plt.ylabel', True, 'import matplotlib.pyplot as plt\n'), (112, 'matplotlib.pyplot.legend', 'plt.legend', True, 'import matplotlib.pyplot as plt\n'), (114, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (23, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (55, 'copy.deepcopy', 'deepcopy', False, 'from copy import deepcopy\n'), (58, 'sds.distributions.lingauss.SingleOutputLinearGaussianWithKnownPrecision', 'SingleOutputLinearGaussianWithKnownPrecision', False, 'from sds.distributions.lingauss import SingleOutputLinearGaussianWithKnownPrecision\n'), (67, 'sds.distributions.lingauss.SingleOutputLinearGaussianWithKnownMean', 'SingleOutputLinearGaussianWithKnownMean', False, 'from sds.distributions.lingauss import SingleOutputLinearGaussianWithKnownMean\n'), (74, 'sds.distributions.gaussian.GaussianWithKnownMeanAndDiagonalPrecision', 'GaussianWithKnownMeanAndDiagonalPrecision', False, 'from sds.distributions.gaussian import GaussianWithKnownMeanAndDiagonalPrecision\n'), (86, 'numpy.eye', 'np.eye', True, 'import numpy as np\n'), (87, 'numpy.ones', 'np.ones', True, 'import numpy as np\n'), (88, 'numpy.ones', 'np.ones', True, 'import numpy as np\n'), (39, 'numpy.ones', 'np.ones', True, 'import numpy as np\n'), (42, 'numpy.ones', 'np.ones', True, 'import numpy as np\n'), (98, 'numpy.atleast_2d', 'np.atleast_2d', True, 'import numpy as np\n'), (21, 'numpy.sqrt', 'np.sqrt', True, 'import numpy as np\n'), (40, 'numpy.ones', 'np.ones', True, 'import numpy as np\n'), (43, 'numpy.ones', 'np.ones', True, 'import numpy as np\n'), (53, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (54, 'numpy.diag', 'np.diag', True, 'import numpy as np\n'), (18, 'numpy.sqrt', 'np.sqrt', True, 'import numpy as np\n')]