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]

[((20, 13, 20, 20), 'numpy.linalg.norm', 'norm', ({(20, 18, 20, 19): 'x'}, {}), '(x)', False, 'from numpy.linalg import norm\n'), ((29, 14, 29, 45), 'numpy.maximum', 'np.maximum', ({(29, 25, 29, 41): '1 - u / norm_grp', (29, 43, 29, 44): '0'}, {}), '(1 - u / norm_grp, 0)', True, 'import numpy as np\n'), ((49, 11, 49, 29), 'numpy.zeros', 'np.zeros', ({(49, 20, 49, 28): 'grp_size'}, {}), '(grp_size)', True, 'import numpy as np\n'), ((92, 16, 92, 34), 'scipy.sparse.issparse', 'sparse.issparse', ({(92, 32, 92, 33): 'X'}, {}), '(X)', False, 'from scipy import sparse\n'), ((98, 13, 98, 32), 'numpy.arange', 'np.arange', ({(98, 23, 98, 31): 'n_groups'}, {}), '(n_groups)', True, 'import numpy as np\n'), ((107, 15, 107, 44), 'numpy.zeros', 'np.zeros', ({(107, 24, 107, 43): '[K + 1, n_features]'}, {}), '([K + 1, n_features])', True, 'import numpy as np\n'), ((108, 8, 108, 33), 'numpy.zeros', 'np.zeros', ({(108, 17, 108, 32): '[K, n_features]'}, {}), '([K, n_features])', True, 'import numpy as np\n'), ((116, 9, 116, 27), 'numpy.zeros', 'np.zeros', ({(116, 18, 116, 26): 'n_groups'}, {}), '(n_groups)', True, 'import numpy as np\n'), ((125, 8, 125, 28), 'numpy.zeros', 'np.zeros', ({(125, 17, 125, 27): 'n_features'}, {}), '(n_features)', True, 'import numpy as np\n'), ((132, 11, 132, 38), 'numpy.zeros', 'np.zeros', ({(132, 20, 132, 37): 'max_iter // f_gap'}, {}), '(max_iter // f_gap)', True, 'import numpy as np\n'), ((22, 15, 22, 31), 'numpy.zeros_like', 'np.zeros_like', ({(22, 29, 22, 30): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((105, 12, 105, 32), 'numpy.asfortranarray', 'np.asfortranarray', ({(105, 30, 105, 31): 'X'}, {}), '(X)', True, 'import numpy as np\n'), ((127, 12, 127, 32), 'numpy.zeros', 'np.zeros', ({(127, 21, 127, 31): 'n_features'}, {}), '(n_features)', True, 'import numpy as np\n'), ((136, 18, 136, 29), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((217, 14, 217, 25), 'numpy.array', 'np.array', ({(217, 23, 217, 24): 'E'}, {}), '(E)', True, 'import numpy as np\n'), ((41, 11, 41, 33), 'numpy.linalg.norm', 'norm', ({(41, 16, 41, 32): '(w[grp] - old_w_g)'}, {}), '(w[grp] - old_w_g)', False, 'from numpy.linalg import norm\n'), ((42, 17, 42, 56), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((60, 11, 60, 33), 'numpy.linalg.norm', 'norm', ({(60, 16, 60, 32): '(w[grp] - old_w_g)'}, {}), '(w[grp] - old_w_g)', False, 'from numpy.linalg import norm\n'), ((104, 29, 104, 44), 'numpy.isfortran', 'np.isfortran', ({(104, 42, 104, 43): 'X'}, {}), '(X)', True, 'import numpy as np\n'), ((121, 20, 121, 37), 'numpy.linalg.norm', 'norm', (), '', False, 'from numpy.linalg import norm\n'), ((158, 20, 158, 47), 'andersoncd.lasso.dual_lasso', 'dual_lasso', ({(158, 31, 158, 32): 'y', (158, 34, 158, 39): 'theta', (158, 41, 158, 46): 'alpha'}, {}), '(y, theta, alpha)', False, 'from andersoncd.lasso import dual_lasso\n'), ((216, 18, 216, 29), 'numpy.array', 'np.array', ({(216, 27, 216, 28): 'E'}, {}), '(E)', True, 'import numpy as np\n'), ((13, 18, 13, 25), 'numpy.linalg.norm', 'norm', ({(13, 23, 13, 24): 'R'}, {}), '(R)', False, 'from numpy.linalg import norm\n'), ((101, 24, 101, 75), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n'), ((102, 21, 102, 71), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n'), ((114, 16, 114, 30), 'numpy.linalg.norm', 'norm', (), '', False, 'from numpy.linalg import norm\n'), ((123, 20, 123, 36), 'numpy.linalg.norm', 'norm', (), '', False, 'from numpy.linalg import norm\n'), ((198, 20, 198, 34), 'numpy.dot', 'np.dot', ({(198, 27, 198, 28): 'U', (198, 30, 198, 33): 'U.T'}, {}), '(U, U.T)', True, 'import numpy as np\n'), ((147, 32, 147, 43), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((203, 28, 204, 42), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((112, 16, 112, 28), 'scipy.sparse.linalg.svds', 'svds', (), '', False, 'from scipy.sparse.linalg import svds\n'), ((201, 43, 201, 53), 'numpy.ones', 'np.ones', ({(201, 51, 201, 52): 'K'}, {}), '(K)', True, 'import numpy as np\n'), ((183, 26, 183, 53), 'numpy.sqrt', 'np.sqrt', ({(183, 34, 183, 52): '(1 + 4 * t_old ** 2)'}, {}), '(1 + 4 * t_old ** 2)', True, 'import numpy as np\n')]