"""
Group lasso with overlap
========================

Comparison of solvers for a least squares with
overlapping group lasso regularization.

References
----------
This example is modeled after the experiments in `Adaptive Three Operator Splitting <https://arxiv.org/pdf/1804.02339.pdf>`_, Appendix E.3.
"""

import os
import numpy as np
from sklearn import preprocessing
import pickle


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

n_samples, n_features = 100, 1002

# .. generate some data ..
# .. the first set of blocks is
groups = [np.arange(8 * i, 8 * i + 10) for i in range(125)]
ground_truth = np.zeros(n_features)
g = np.random.randint(0, len(groups), 10)
for i in g:
    ground_truth[groups[i]] = np.random.randn()

A = np.random.randn(n_samples, n_features)
p = 0.95  # create a matrix with correlations between features
for i in range(1, n_features):
    A[:, i] = p * A[:, i] + (1 - p) * A[:, i-1]
A[:, 0] /= np.sqrt(1 - p ** 2)
A = preprocessing.StandardScaler().fit_transform(A)
b = A.dot(ground_truth) + np.random.randn(n_samples)

# make labels in {0, 1}
b = np.sign(b)
b = (b + 1) // 2

# ... save data ...

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

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