"""
Estimating a sparse and low rank matrix
=======================================

"""
print(__doc__)
import numpy as np
import pickle
import os

# .. Generate synthetic data ..
np.random.seed(1)

sigma_2 = 0.6
N = 100
d = 20
blocks = np.array([2 * d / 10, 1 * d / 10, 1 * d / 10, 3 * d / 10, 3 * d / 10]).astype(
    np.int
)
epsilon = 10 ** (-15)

mu = np.zeros(d)
Sigma = np.zeros((d, d))
blck = 0
for k in range(len(blocks)):
    v = 2 * np.random.rand(int(blocks[k]), 1)
    v = v * (abs(v) > 0.9)
    Sigma[blck: blck + blocks[k], blck: blck + blocks[k]] = np.dot(v, v.T)
    blck = blck + blocks[k]
X = np.random.multivariate_normal(
    mu, Sigma + epsilon * np.eye(d), N
) + sigma_2 * np.random.randn(N, d)
Sigma_hat = np.cov(X.T)

threshold = 1e-5
Sigma[np.abs(Sigma) < threshold] = 0
Sigma[np.abs(Sigma) >= threshold] = 1

# .. generate some data ..

n_features = np.multiply(*Sigma.shape)
n_samples = n_features
print("#features", n_features)
A = np.random.randn(n_samples, n_features)

sigma = 1.0
b = A.dot(Sigma.ravel()) + sigma * np.random.randn(n_samples)

# ... save data ...

if not os.path.exists('data'):
    os.makedirs('data')

with open("data/sparse_nuclear_norm.pkl", 'wb') as fh:
    pickle.dump({'A': A, 'b': b, 'Sigma': Sigma, 'n_features': n_features, 'n_samples': n_samples}, fh)
