import time
import numpy as np

from scipy import sparse
from numba import njit
from numpy.linalg import norm
from scipy.sparse.linalg import svds

from andersoncd.lasso import dual_lasso


def primal_grp(R, w, alpha, grp_size):
    return (0.5 * norm(R) ** 2 + alpha *
            norm(w.reshape(-1, grp_size), axis=1).sum())


@njit
def BST(x, u):
    """Block soft-thresholding of vector x at level u."""
    norm_x = norm(x)
    if norm_x < u:
        return np.zeros_like(x)
    else:
        return (1 - u / norm_x) * x


def BST_vec(x, u, grp_size):
    norm_grp = norm(x.reshape(-1, grp_size), axis=1)
    scaling = np.maximum(1 - u / norm_grp, 0)
    return (x.reshape(-1, grp_size) * scaling[:, None]).reshape(x.shape[0])


@njit
def _bcd(X, w, R, alpha, lc, groups):
    grp_size = w.shape[0] // lc.shape[0]
    for g in groups:
        grp = slice(g * grp_size, (g + 1) * grp_size)
        Xg = X[:, grp]
        old_w_g = w[grp].copy()
        w[grp] = BST(old_w_g + Xg.T @ R / lc[g], alpha / lc[g])
        if norm(w[grp] - old_w_g) != 0:
            R += np.sum((old_w_g - w[grp]) * Xg, axis=1)


@njit
def _bcd_sparse(
        X_data, X_indices, X_indptr, w, R, alpha, lc):
    grp_size = w.shape[0] // lc.shape[0]
    grad = np.zeros(grp_size)

    for g in range(lc.shape[0]):
        grad.fill(0)
        grp = slice(g * grp_size, (g + 1) * grp_size)

        for j in range(grp_size * g, grp_size * (g + 1)):
            for ix in range(X_indptr[j], X_indptr[j + 1]):
                grad[j % g] += X_data[ix] * R[X_indices[ix]]
        old_w_g = w[grp].copy()
        w[grp] = BST(old_w_g + grad / lc[g], alpha / lc[g])
        if norm(w[grp] - old_w_g) != 0:
            for j in range(g * grp_size, (g + 1) * grp_size):
                for ix in range(X_indptr[j], X_indptr[j + 1]):
                    R[X_indices[ix]] += (old_w_g[j % grp_size] -
                                         w[j % grp_size]) * X_data[ix]


def solver_group(
        X, y, alpha, grp_size, max_iter=10000, tol=1e-4, f_gap=10, K=5,
        use_acc=False, algo='bcd', compute_time=False, tmax=np.infty,
        verbose=True):
    """Solve the GroupLasso with BCD/ISTA/FISTA, eventually with extrapolation.

    Groups are contiguous, of size grp_size.
    Objective:
    norm(y - Xw, ord=2)**2 / 2 + alpha * sum_g ||w_{[g]}||_2

    TODO: filled docstring

    Parameters:
    algo: string
        'bcd', 'pgd', 'fista'

    compute_time : bool, default=False
        If you want to compute timings or not

    tmax : float, default=1000
        Maximum time (in seconds) the algorithm is allowed to run

    alpha: strength of the group penalty
    """

    is_sparse = sparse.issparse(X)
    n_features = X.shape[1]
    if n_features % grp_size != 0:
        raise ValueError("n_features is not a multiple of group size")
    n_groups = n_features // grp_size

    _range = np.arange(n_groups)
    groups = dict(
        bcd=lambda: _range,
        bcdshuf=lambda: np.random.choice(n_groups, n_groups, replace=False),
        rbcd=lambda: np.random.choice(n_groups, n_groups, replace=True))

    if not is_sparse and not np.isfortran(X):
        X = np.asfortranarray(X)

    last_K_w = np.zeros([K + 1, n_features])
    U = np.zeros([K, n_features])

    if algo in ('pgd', 'fista'):
        if is_sparse:
            L = svds(X, k=1)[1][0] ** 2
        else:
            L = norm(X, ord=2) ** 2

    lc = np.zeros(n_groups)
    for g in range(n_groups):
        X_g = X[:, g * grp_size: (g + 1) * grp_size]
        if is_sparse:
            gram = (X_g.T @ X_g).todense()
            lc[g] = norm(gram, ord=2)
        else:
            lc[g] = norm(X_g, ord=2) ** 2

    w = np.zeros(n_features)
    if algo == 'fista':
        z = np.zeros(n_features)
        t_new = 1

    R = y.copy()
    E = []
    gaps = np.zeros(max_iter // f_gap)

    if compute_time:
        times = []
        t_start = time.time()

    for it in range(max_iter):
        if it % f_gap == 0:
            if algo == 'fista':
                R = y - X @ w
            p_obj = primal_grp(R, w, alpha, grp_size)
            E.append(p_obj)
            theta = R / alpha

            if compute_time:
                elapsed_times = time.time() - t_start
                times.append(elapsed_times)
                if verbose:
                    print("elapsed time: %f " % elapsed_times)
                if elapsed_times > tmax:
                    break

            d_norm_theta = np.max(
                norm((X.T @ theta).reshape(-1, grp_size), axis=1))
            if d_norm_theta > 1.:
                theta /= d_norm_theta
            d_obj = dual_lasso(y, theta, alpha)

            gap = p_obj - d_obj

            if verbose:
                print("Iteration %d, p_obj::%.5f, d_obj::%.5f, gap::%.2e" %
                      (it, p_obj, d_obj, gap))
            gaps[it // f_gap] = gap
            if gap < tol:
                print("Early exit")
                break

        if algo.endswith('bcd'):
            if is_sparse:
                _bcd_sparse(
                    X.data, X.indices, X.indptr, w, R, alpha, lc)
            else:
                _bcd(X, w, R, alpha, lc, groups[algo]())
        elif algo == 'pgd':
            w[:] = BST_vec(w + X.T @ R / L, alpha / L, grp_size)
            R[:] = y - X @ w
        elif algo == 'fista':
            w_old = w.copy()
            w[:] = BST_vec(z - X.T @ (X @ z - y) / L, alpha / L, grp_size)
            t_old = t_new
            t_new = (1. + np.sqrt(1 + 4 * t_old ** 2)) / 2.
            z[:] = w + (t_old - 1.) / t_new * (w - w_old)
        else:
            raise ValueError("Unknown algo %s" % algo)

        if use_acc:
            if it < K + 1:
                last_K_w[it] = w
            else:
                for k in range(K):
                    last_K_w[k] = last_K_w[k + 1]
                last_K_w[K - 1] = w

                for k in range(K):
                    U[k] = last_K_w[k + 1] - last_K_w[k]
                C = np.dot(U, U.T)

                try:
                    z = np.linalg.solve(C, np.ones(K))
                    c = z / z.sum()
                    w_acc = np.sum(last_K_w[:-1] * c[:, None],
                                   axis=0)
                    p_obj = primal_grp(R, w, alpha, grp_size)
                    R_acc = y - X @ w_acc
                    p_obj_acc = primal_grp(R_acc, w_acc, alpha, grp_size)
                    if p_obj_acc < p_obj:
                        w = w_acc
                        R = R_acc
                except np.linalg.LinAlgError:
                    if verbose:
                        print("----------Linalg error")

    if compute_time:
        return w, np.array(E), gaps[:it // f_gap + 1], times
    return w, np.array(E), gaps[:it // f_gap + 1]

[((421, 428), 'numpy.linalg.norm', 'norm', (['x'], {}), '(x)\n', (425, 428), False, 'from numpy.linalg import norm\n'), ((624, 655), 'numpy.maximum', 'np.maximum', (['(1 - u / norm_grp)', '(0)'], {}), '(1 - u / norm_grp, 0)\n', (634, 655), True, 'import numpy as np\n'), ((1242, 1260), 'numpy.zeros', 'np.zeros', (['grp_size'], {}), '(grp_size)\n', (1250, 1260), True, 'import numpy as np\n'), ((2672, 2690), 'scipy.sparse.issparse', 'sparse.issparse', (['X'], {}), '(X)\n', (2687, 2690), False, 'from scipy import sparse\n'), ((2877, 2896), 'numpy.arange', 'np.arange', (['n_groups'], {}), '(n_groups)\n', (2886, 2896), True, 'import numpy as np\n'), ((3190, 3219), 'numpy.zeros', 'np.zeros', (['[K + 1, n_features]'], {}), '([K + 1, n_features])\n', (3198, 3219), True, 'import numpy as np\n'), ((3228, 3253), 'numpy.zeros', 'np.zeros', (['[K, n_features]'], {}), '([K, n_features])\n', (3236, 3253), True, 'import numpy as np\n'), ((3410, 3428), 'numpy.zeros', 'np.zeros', (['n_groups'], {}), '(n_groups)\n', (3418, 3428), True, 'import numpy as np\n'), ((3680, 3700), 'numpy.zeros', 'np.zeros', (['n_features'], {}), '(n_features)\n', (3688, 3700), True, 'import numpy as np\n'), ((3816, 3843), 'numpy.zeros', 'np.zeros', (['(max_iter // f_gap)'], {}), '(max_iter // f_gap)\n', (3824, 3843), True, 'import numpy as np\n'), ((463, 479), 'numpy.zeros_like', 'np.zeros_like', (['x'], {}), '(x)\n', (476, 479), True, 'import numpy as np\n'), ((3153, 3173), 'numpy.asfortranarray', 'np.asfortranarray', (['X'], {}), '(X)\n', (3170, 3173), True, 'import numpy as np\n'), ((3737, 3757), 'numpy.zeros', 'np.zeros', (['n_features'], {}), '(n_features)\n', (3745, 3757), True, 'import numpy as np\n'), ((3903, 3914), 'time.time', 'time.time', ([], {}), '()\n', (3912, 3914), False, 'import time\n'), ((6744, 6755), 'numpy.array', 'np.array', (['E'], {}), '(E)\n', (6752, 6755), True, 'import numpy as np\n'), ((1024, 1046), 'numpy.linalg.norm', 'norm', (['(w[grp] - old_w_g)'], {}), '(w[grp] - old_w_g)\n', (1028, 1046), False, 'from numpy.linalg import norm\n'), ((1070, 1109), 'numpy.sum', 'np.sum', (['((old_w_g - w[grp]) * Xg)'], {'axis': '(1)'}), '((old_w_g - w[grp]) * Xg, axis=1)\n', (1076, 1109), True, 'import numpy as np\n'), ((1652, 1674), 'numpy.linalg.norm', 'norm', (['(w[grp] - old_w_g)'], {}), '(w[grp] - old_w_g)\n', (1656, 1674), False, 'from numpy.linalg import norm\n'), ((3124, 3139), 'numpy.isfortran', 'np.isfortran', (['X'], {}), '(X)\n', (3136, 3139), True, 'import numpy as np\n'), ((3597, 3614), 'numpy.linalg.norm', 'norm', (['gram'], {'ord': '(2)'}), '(gram, ord=2)\n', (3601, 3614), False, 'from numpy.linalg import norm\n'), ((4630, 4657), 'andersoncd.lasso.dual_lasso', 'dual_lasso', (['y', 'theta', 'alpha'], {}), '(y, theta, alpha)\n', (4640, 4657), False, 'from andersoncd.lasso import dual_lasso\n'), ((6687, 6698), 'numpy.array', 'np.array', (['E'], {}), '(E)\n', (6695, 6698), True, 'import numpy as np\n'), ((247, 254), 'numpy.linalg.norm', 'norm', (['R'], {}), '(R)\n', (251, 254), False, 'from numpy.linalg import norm\n'), ((2968, 3019), 'numpy.random.choice', 'np.random.choice', (['n_groups', 'n_groups'], {'replace': '(False)'}), '(n_groups, n_groups, replace=False)\n', (2984, 3019), True, 'import numpy as np\n'), ((3042, 3092), 'numpy.random.choice', 'np.random.choice', (['n_groups', 'n_groups'], {'replace': '(True)'}), '(n_groups, n_groups, replace=True)\n', (3058, 3092), True, 'import numpy as np\n'), ((3380, 3394), 'numpy.linalg.norm', 'norm', (['X'], {'ord': '(2)'}), '(X, ord=2)\n', (3384, 3394), False, 'from numpy.linalg import norm\n'), ((3649, 3665), 'numpy.linalg.norm', 'norm', (['X_g'], {'ord': '(2)'}), '(X_g, ord=2)\n', (3653, 3665), False, 'from numpy.linalg import norm\n'), ((5991, 6005), 'numpy.dot', 'np.dot', (['U', 'U.T'], {}), '(U, U.T)\n', (5997, 6005), True, 'import numpy as np\n'), ((4211, 4222), 'time.time', 'time.time', ([], {}), '()\n', (4220, 4222), False, 'import time\n'), ((6147, 6191), 'numpy.sum', 'np.sum', (['(last_K_w[:-1] * c[:, (None)])'], {'axis': '(0)'}), '(last_K_w[:-1] * c[:, (None)], axis=0)\n', (6153, 6191), True, 'import numpy as np\n'), ((3326, 3338), 'scipy.sparse.linalg.svds', 'svds', (['X'], {'k': '(1)'}), '(X, k=1)\n', (3330, 3338), False, 'from scipy.sparse.linalg import svds\n'), ((6071, 6081), 'numpy.ones', 'np.ones', (['K'], {}), '(K)\n', (6078, 6081), True, 'import numpy as np\n'), ((5497, 5524), 'numpy.sqrt', 'np.sqrt', (['(1 + 4 * t_old ** 2)'], {}), '(1 + 4 * t_old ** 2)\n', (5504, 5524), True, 'import numpy as np\n')]