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', (['(0)'], {}), True, 'import numpy as np\n'), (11, 'numpy.random.randn', 'np.random.randn', (['n_samples', 'n_features'], {}), True, 'import numpy as np\n'), (14, 'numpy.zeros', 'np.zeros', (['n_features'], {}), True, 'import numpy as np\n'), (16, 'numpy.random.randint', 'np.random.randint', (['(0)', 'n_features', '(10)'], {}), True, 'import numpy as np\n'), (25, 'sklearn.linear_model.ARDRegression', 'ARDRegression', ([], {'fit_intercept': '(False)', 'n_iter': '(1000)'}), False, 'from sklearn.linear_model import ARDRegression, LinearRegression\n'), (28, 'sklearn.linear_model.LinearRegression', 'LinearRegression', ([], {'fit_intercept': '(False)'}), False, 'from sklearn.linear_model import ARDRegression, LinearRegression\n'), (45, 'copy.deepcopy', 'deepcopy', (['likelihood_precision_prior'], {}), False, 'from copy import deepcopy\n'), (46, 'copy.deepcopy', 'deepcopy', (['parameter_precision_prior'], {}), False, 'from copy import deepcopy\n'), (85, 'numpy.zeros', 'np.zeros', (['(1, n_features)'], {}), True, 'import numpy as np\n'), (90, 'sds.distributions.composite.MatrixNormalGamma', 'MatrixNormalGamma', ([], {'column_dim': 'n_features', 'row_dim': '(1)', 'M': 'M', 'K': 'K', 'alphas': 'alphas', 'betas': 'betas'}), False, 'from sds.distributions.composite import MatrixNormalGamma\n'), (93, 'copy.deepcopy', 'deepcopy', (['prior'], {}), False, 'from copy import deepcopy\n'), (94, 'sds.distributions.lingauss.LinearGaussianWithDiagonalPrecision', 'LinearGaussianWithDiagonalPrecision', ([], {'column_dim': 'n_features', 'row_dim': '(1)', 'affine': '(False)'}), False, 'from sds.distributions.lingauss import LinearGaussianWithDiagonalPrecision\n'), (102, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(6, 5)'}), True, 'import matplotlib.pyplot as plt\n'), (103, 'matplotlib.pyplot.title', 'plt.title', (['"""Weights of the model"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (104, 'matplotlib.pyplot.plot', 'plt.plot', (['w'], {'color': '"""orange"""', 'linestyle': '"""-"""', 'linewidth': '(2)', 'label': '"""Ground truth"""'}), True, 'import matplotlib.pyplot as plt\n'), (105, 'matplotlib.pyplot.plot', 'plt.plot', (['clf.coef_'], {'color': '"""darkblue"""', 'linestyle': '"""-"""', 'linewidth': '(2)', 'label': '"""Sklearn ARD"""'}), True, 'import matplotlib.pyplot as plt\n'), (106, 'matplotlib.pyplot.plot', 'plt.plot', (['our_ard'], {'color': '"""red"""', 'linestyle': '"""-"""', 'linewidth': '(2)', 'label': '"""Our ARD"""'}), True, 'import matplotlib.pyplot as plt\n'), (110, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Features"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (111, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Values of the weights"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (112, 'matplotlib.pyplot.legend', 'plt.legend', ([], {'loc': '(1)'}), 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', (['X', 'w'], {}), True, 'import numpy as np\n'), (55, 'copy.deepcopy', 'deepcopy', (['parameter_prior'], {}), False, 'from copy import deepcopy\n'), (58, 'sds.distributions.lingauss.SingleOutputLinearGaussianWithKnownPrecision', 'SingleOutputLinearGaussianWithKnownPrecision', ([], {'column_dim': 'n_features', 'lmbda': 'beta', 'affine': '(False)'}), False, 'from sds.distributions.lingauss import SingleOutputLinearGaussianWithKnownPrecision\n'), (67, 'sds.distributions.lingauss.SingleOutputLinearGaussianWithKnownMean', 'SingleOutputLinearGaussianWithKnownMean', ([], {'column_dim': 'n_features', 'W': 'param', 'affine': '(False)'}), False, 'from sds.distributions.lingauss import SingleOutputLinearGaussianWithKnownMean\n'), (74, 'sds.distributions.gaussian.GaussianWithKnownMeanAndDiagonalPrecision', 'GaussianWithKnownMeanAndDiagonalPrecision', ([], {'dim': 'n_features'}), False, 'from sds.distributions.gaussian import GaussianWithKnownMeanAndDiagonalPrecision\n'), (86, 'numpy.eye', 'np.eye', (['n_features'], {}), True, 'import numpy as np\n'), (87, 'numpy.ones', 'np.ones', (['(1,)'], {}), True, 'import numpy as np\n'), (88, 'numpy.ones', 'np.ones', (['(1,)'], {}), True, 'import numpy as np\n'), (39, 'numpy.ones', 'np.ones', (['(1,)'], {}), True, 'import numpy as np\n'), (42, 'numpy.ones', 'np.ones', (['(n_features,)'], {}), True, 'import numpy as np\n'), (98, 'numpy.atleast_2d', 'np.atleast_2d', (['y'], {}), True, 'import numpy as np\n'), (21, 'numpy.sqrt', 'np.sqrt', (['alpha_'], {}), True, 'import numpy as np\n'), (40, 'numpy.ones', 'np.ones', (['(1,)'], {}), True, 'import numpy as np\n'), (43, 'numpy.ones', 'np.ones', (['(n_features,)'], {}), True, 'import numpy as np\n'), (53, 'numpy.zeros', 'np.zeros', (['(n_features,)'], {}), True, 'import numpy as np\n'), (54, 'numpy.diag', 'np.diag', (['alphas'], {}), True, 'import numpy as np\n'), (18, 'numpy.sqrt', 'np.sqrt', (['lambda_'], {}), True, 'import numpy as np\n')]