def test_group_lasso_atom():


    ps = np.array([0]*5 + [3]*3)
    weights = {3:2., 0:2.3}

    lagrange = 1.5
    lipschitz = 0.2
    p = gl.group_lasso(ps, weights=weights, lagrange=lagrange)
    z = 30 * np.random.standard_normal(8)
    q = rr.identity_quadratic(lipschitz, z, 0, 0)

    x = p.solve(q)
    a = ml.mixed_lasso_lagrange_prox(z, lagrange, lipschitz, 
                                     np.array([],np.int), 
                                     np.array([],np.int), 
                                     np.array([], np.int), 
                                     np.array([], np.int), 
                                     np.array([0,0,0,0,0,1,1,1]), np.array([np.sqrt(5), 2]))

    result = np.zeros_like(a)
    result[:5] = z[:5] / np.linalg.norm(z[:5]) * max(np.linalg.norm(z[:5]) - weights[0] * lagrange/lipschitz, 0)
    result[5:] = z[5:] / np.linalg.norm(z[5:]) * max(np.linalg.norm(z[5:]) - weights[3] * lagrange/lipschitz, 0)

    lipschitz = 1.
    q = rr.identity_quadratic(lipschitz, z, 0, 0)
    x2 = p.solve(q)
    pc = p.conjugate
    a2 = pc.solve(q)

    np.testing.assert_allclose(z-a2, x2)

def test_l1prox():
    '''
    this test verifies that the l1 prox in lagrange form can be solved
    by a primal/dual specification 

    obviously, we don't to solve the l1 prox this way,
    but it verifies that specification is working correctly

    '''

    l1 = rr.l1norm(4, lagrange=0.3)
    ww = np.random.standard_normal(4)*3
    ab = l1.proximal(rr.identity_quadratic(0.5, ww, 0,0))

    l1c = copy(l1)
    l1c.quadratic = rr.identity_quadratic(0.5, ww, None, 0.)
    a = rr.simple_problem.nonsmooth(l1c)
    solver = rr.FISTA(a)
    solver.fit(tol=1.e-10)

    ad = a.coefs

    l1c = copy(l1)
    l1c.quadratic = rr.identity_quadratic(0.5, ww, None, 0.)
    a = rr.dual_problem.fromprimal(l1c)
    solver = rr.FISTA(a)
    solver.fit(tol=1.0e-14)

    ac = a.primal

    np.testing.assert_allclose(ac, ab, rtol=1.0e-4)
    np.testing.assert_allclose(ac, ad, rtol=1.0e-4)

def test_l1prox_bound():
    '''
    this test verifies that the l1 prox in bound form can be solved
    by a primal/dual specification 

    obviously, we don't to solve the l1 prox this way,
    but it verifies that specification is working correctly

    '''

    l1 = rr.l1norm(4, bound=2.)
    ww = np.random.standard_normal(4)*2
    ab = l1.proximal(rr.identity_quadratic(0.5, ww, 0, 0))

    l1c = copy(l1)
    l1c.quadratic = rr.identity_quadratic(0.5, ww, None, 0.)
    a = rr.simple_problem.nonsmooth(l1c)
    solver = rr.FISTA(a)
    solver.fit(min_its=100)

    l1c = copy(l1)
    l1c.quadratic = rr.identity_quadratic(0.5, ww, None, 0.)
    a = rr.dual_problem.fromprimal(l1c)
    solver = rr.FISTA(a)
    solver.fit(min_its=100)

    ac = a.primal

    np.testing.assert_allclose(ac + 0.1, ab + 0.1, rtol=1.e-4)

    def test_simple_problem(self):
        tests = []
        atom, q, prox_center, L = self.atom, self.q, self.prox_center, self.L
        loss = self.loss

        problem = rr.simple_problem(loss, atom)
        solver = rr.FISTA(problem)
        solver.fit(tol=1.0e-12, FISTA=self.FISTA, coef_stop=self.coef_stop, min_its=100)

        tests.append((atom.proximal(q), solver.composite.coefs, 'solving prox with simple_problem with monotonicity\n %s' % str(self)))

        # write the loss in terms of a quadratic for the smooth loss and a smooth function...

        q = rr.identity_quadratic(L, prox_center, 0, 0)
        lossq = rr.quadratic.shift(prox_center.copy(), coef=0.6*L)
        lossq.quadratic = rr.identity_quadratic(0.4*L, prox_center.copy(), 0, 0)
        problem = rr.simple_problem(lossq, atom)

        tests.append((atom.proximal(q), 
              problem.solve(coef_stop=self.coef_stop, 
                            FISTA=self.FISTA, 
                            tol=1.0e-12), 
               'solving prox with simple_problem ' +
               'with monotonicity  but loss has identity_quadratic %s\n ' % str(self)))

        problem = rr.simple_problem(loss, atom)
        solver = rr.FISTA(problem)
        solver.fit(tol=1.0e-12, monotonicity_restart=False,
                   coef_stop=self.coef_stop, FISTA=self.FISTA, min_its=100)

        tests.append((atom.proximal(q), solver.composite.coefs, 'solving prox with simple_problem no monotonicity_restart\n %s' % str(self)))

        d = atom.conjugate
        problem = rr.simple_problem(loss, d)
        solver = rr.FISTA(problem)
        solver.fit(tol=1.0e-12, monotonicity_restart=False, 
                   coef_stop=self.coef_stop, FISTA=self.FISTA, min_its=100)
        tests.append((d.proximal(q), problem.solve(tol=1.e-12,
                                                FISTA=self.FISTA,
                                                coef_stop=self.coef_stop,
                                                monotonicity_restart=False), 
               'solving dual prox with simple_problem no monotonocity\n %s ' % str(self)))

        if not self.interactive:
            for test in tests:
                yield (all_close,) + test + (self,)
        else:
            for test in tests:
                yield all_close(*((test + (self,))))

def test_adding_quadratic_lasso():

    X, y, beta, active, sigma = instance(n=300, p=200)
    Q = rr.identity_quadratic(0.01, 0, np.random.standard_normal(X.shape[1]), 0)

    L1 = lasso.gaussian(X, y, 20, quadratic=Q)
    beta1 = L1.fit(solve_args={'min_its':500, 'tol':1.e-12})
    G1 = X[:,L1.active].T.dot(X.dot(beta1) - y) + Q.objective(beta1,'grad')[L1.active]
    np.testing.assert_allclose(G1 * np.sign(beta1[L1.active]), -20)

    lin = rr.identity_quadratic(0.0, 0, np.random.standard_normal(X.shape[1]), 0)
    L2 = lasso.gaussian(X, y, 20, quadratic=lin)
    beta2 = L2.fit(solve_args={'min_its':500, 'tol':1.e-12})
    G2 = X[:,L2.active].T.dot(X.dot(beta2) - y) + lin.objective(beta2,'grad')[L2.active]
    np.testing.assert_allclose(G2 * np.sign(beta2[L2.active]), -20)

def test_sqrt_lasso_pvals(n=100,
                          p=200,
                          s=7,
                          sigma=5,
                          rho=0.3,
                          snr=7.):

    counter = 0

    while True:
        counter += 1
        X, y, beta, active, sigma = instance(n=n, 
                                             p=p, 
                                             s=s, 
                                             sigma=sigma, 
                                             rho=rho, 
                                             snr=snr)

        lam_theor = np.mean(np.fabs(np.dot(X.T, np.random.standard_normal((n, 1000)))).max(0)) / np.sqrt(n)
        Q = rr.identity_quadratic(0.01, 0, np.ones(p), 0)

        weights_with_zeros = 0.7*lam_theor * np.ones(p)
        weights_with_zeros[:3] = 0.

        L = lasso.sqrt_lasso(X, y, weights_with_zeros)
        L.fit()
        v = {1:'twosided',
             0:'onesided'}[counter % 2]
        if set(active).issubset(L.active):
            S = L.summary(v)
            return [p for p, v in zip(S['pval'], S['variable']) if v not in active]

def test_sqrt_lasso(n=100, p=20):

    y = np.random.standard_normal(n)
    X = np.random.standard_normal((n,p))

    lam_theor = np.mean(np.fabs(np.dot(X.T, np.random.standard_normal((n, 1000)))).max(0)) / np.sqrt(n)
    Q = rr.identity_quadratic(0.01, 0, np.random.standard_normal(p) / 5., 0)

    weights_with_zeros = 0.5*lam_theor * np.ones(p)
    weights_with_zeros[:3] = 0.

    huge_weights = weights_with_zeros * 10000

    for q, fw in product([None, Q],
                         [0.5*lam_theor, weights_with_zeros, huge_weights]):

        L = lasso.sqrt_lasso(X, y, fw, quadratic=q, solve_args={'min_its':300, 'tol':1.e-12})
        L.fit(solve_args={'min_its':300, 'tol':1.e-12})
        C = L.constraints

        S = L.summary('onesided', compute_intervals=True)
        S = L.summary('twosided')

        yield (np.testing.assert_array_less,
               np.dot(L.constraints.linear_part, L.onestep_estimator),
               L.constraints.offset)

def test_gaussian(n=100, p=20):

    y = np.random.standard_normal(n)
    X = np.random.standard_normal((n,p))

    lam_theor = np.mean(np.fabs(np.dot(X.T, np.random.standard_normal((n, 1000)))).max(0))
    Q = rr.identity_quadratic(0.01, 0, np.ones(p), 0)

    weights_with_zeros = 0.5*lam_theor * np.ones(p)
    weights_with_zeros[:3] = 0.

    huge_weights = weights_with_zeros * 10000

    for q, fw in product([Q, None],
                         [0.5*lam_theor, weights_with_zeros, huge_weights]):

        L = lasso.gaussian(X, y, fw, 1., quadratic=Q)
        L.fit()
        C = L.constraints

        sandwich = gaussian_sandwich_estimator(X, y)
        L = lasso.gaussian(X, y, fw, 1., quadratic=Q, covariance_estimator=sandwich)
        L.fit()
        C = L.constraints

        S = L.summary('onesided', compute_intervals=True)
        S = L.summary('twosided')

        nt.assert_raises(ValueError, L.summary, 'none')
        print(L.active)
        yield (np.testing.assert_array_less,
               np.dot(L.constraints.linear_part, L.onestep_estimator),
               L.constraints.offset)

    def __iter__(self):
        for offset, FISTA, coef_stop, L, q, w in itertools.product(self.offset_choices,
                                                                   self.FISTA_choices,
                                                                   self.coef_stop_choices,
                                                                   self.L_choices,
                                                                   self.quadratic_choices,
                                                                   self.weight_choices):
            self.FISTA = FISTA
            self.coef_stop = coef_stop
            self.L = L

            if self.mode == 'lagrange':
                atom = self.klass(w, lagrange=self.lagrange)
            else:
                atom = self.klass(w, bound=self.bound)
            atom.use_sklearn = self.use_sklearn and have_sklearn_iso # test out both prox maps if available

            if q: 
                atom.quadratic = rr.identity_quadratic(0, 0, np.random.standard_normal(atom.shape)*0.02)

            if offset:
                atom.offset = 0.02 * np.random.standard_normal(atom.shape)

            solver = Solver(atom, interactive=self.interactive, 
                            coef_stop=coef_stop,
                            FISTA=FISTA,
                            L=L)
            yield solver

def test_proximal_maps():
    bound = 0.14
    lagrange = 0.13
    shape = 20

    Z = np.random.standard_normal(shape) * 4
    W = 0.02 * np.random.standard_normal(shape)
    U = 0.02 * np.random.standard_normal(shape)
    linq = rr.identity_quadratic(0, 0, W, 0)

    for L, atom, q, offset, FISTA, coef_stop in itertools.product(
        [0.5, 1, 0.1],
        [A.l1norm, A.supnorm, A.l2norm, A.positive_part, A.constrained_max],
        [None, linq],
        [None, U],
        [False, True],
        [True, False],
    ):

        p = atom(shape, lagrange=lagrange, quadratic=q, offset=offset)
        d = p.conjugate
        yield ac, p.lagrange_prox(Z, lipschitz=L), Z - d.bound_prox(
            Z * L, lipschitz=1.0 / L
        ) / L, "testing lagrange_prox and bound_prox starting from atom %s " % atom
        # some arguments of the constructor

        nt.assert_raises(AttributeError, setattr, p, "bound", 4.0)
        nt.assert_raises(AttributeError, setattr, d, "lagrange", 4.0)

        nt.assert_raises(AttributeError, setattr, p, "bound", 4.0)
        nt.assert_raises(AttributeError, setattr, d, "lagrange", 4.0)

        for t in solveit(p, Z, W, U, linq, L, FISTA, coef_stop):
            yield t

        b = atom(shape, bound=bound, quadratic=q, offset=offset)

        for t in solveit(b, Z, W, U, linq, L, FISTA, coef_stop):
            yield t

    lagrange = 0.1
    for L, atom, q, offset, FISTA, coef_stop in itertools.product(
        [0.5, 1, 0.1], sorted(A.nonpaired_atoms), [None, linq], [None, U], [False, True], [False, True]
    ):

        p = atom(shape, lagrange=lagrange, quadratic=q, offset=offset)
        d = p.conjugate
        yield ac, p.lagrange_prox(Z, lipschitz=L), Z - d.bound_prox(
            Z * L, lipschitz=1.0 / L
        ) / L, "testing lagrange_prox and bound_prox starting from atom %s " % atom
        # some arguments of the constructor

        nt.assert_raises(AttributeError, setattr, p, "bound", 4.0)
        nt.assert_raises(AttributeError, setattr, d, "lagrange", 4.0)

        nt.assert_raises(AttributeError, setattr, p, "bound", 4.0)
        nt.assert_raises(AttributeError, setattr, d, "lagrange", 4.0)

        for t in solveit(p, Z, W, U, linq, L, FISTA, coef_stop):
            yield t

    def step_valid(self,
                   max_trials=10):
        """
        Try and move Y_valid
        by accept reject stopping after `max_trials`.
        """

        X, L, mults = self.X, self.L, self.mults
        n, p = X.shape

        count = 0
        Q_old = self.Q_valid

        while True:
            count += 1
            self.Q_valid = self.Q_inter + identity_quadratic(0, 0, self.randomization.rvs(size=self.X.shape[1]) * 
                                                             self.scale_valid, 0) 

            if len(self.mults) > 0:
                proposal_value = self.choose_lambda(self.Y,
                                                    shift_size=0)

                if proposal_value[0] in self.accept_values:
                    break
            else:
                break

            if count >= max_trials:
                self.Q_valid = Q_old
                break

    def __iter__(self):
        for offset, FISTA, coef_stop, L, q in itertools.product(self.offset_choices,
                                                                self.FISTA_choices,
                                                                self.coef_stop_choices,
                                                                self.L_choices,
                                                                self.quadratic_choices):
            self.FISTA = FISTA
            self.coef_stop = coef_stop
            self.L = L

            if self.mode == 'lagrange':
                atom = self.klass(self.shape, lagrange=self.lagrange)
            else:
                atom = self.klass(self.shape, bound=self.bound)

            if q: 
                atom.quadratic = rr.identity_quadratic(0,0,np.random.standard_normal(atom.shape)*0.02)

            if offset:
                atom.offset = 0.02 * np.random.standard_normal(atom.shape)

            solver = Solver(atom, interactive=self.interactive, 
                            coef_stop=coef_stop,
                            FISTA=FISTA,
                            L=L)

            # make sure certain lines of code are tested
            assert(atom == atom)
            atom.latexify(), atom.dual, atom.conjugate

            yield solver

    def __iter__(self):
        for offset, FISTA, coef_stop, L, q, groups in itertools.product(self.offset_choices,
                                                                        self.FISTA_choices,
                                                                        self.coef_stop_choices,
                                                                        self.L_choices,
                                                                        self.quadratic_choices,
                                                                        self.group_choices):
            self.FISTA = FISTA
            self.coef_stop = coef_stop
            self.L = L

            if self.mode == 'lagrange':
                atom = self.klass(groups, lagrange=self.lagrange)
            else:
                atom = self.klass(groups, bound=self.bound)

            if q: 
                atom.quadratic = rr.identity_quadratic(0,0,np.random.standard_normal(atom.shape)*0.02)

            if offset:
                atom.offset = 0.02 * np.random.standard_normal(atom.shape)

            solver = Solver(atom, interactive=self.interactive, 
                            coef_stop=coef_stop,
                            FISTA=FISTA,
                            L=L)
            yield solver

def test_proximal_method():

    X = np.random.standard_normal((100, 50))
    X[:,:7] *= 5

    qX = identity_quadratic(1,X,0,0)
    P = FM.nuclear_norm(X.shape, lagrange=1)
    RP = todense(P.proximal(qX))

    B = FM.nuclear_norm(X.shape, bound=1)
    RB = todense(B.proximal(qX))

    BO = FM.operator_norm(X.shape, bound=1)
    PO = FM.operator_norm(X.shape, lagrange=1)

    RPO = todense(PO.proximal(qX))
    RBO = todense(BO.proximal(qX))

    D = np.linalg.svd(X, full_matrices=0)[1]
    lD = np.linalg.svd(RP, full_matrices=0)[1]
    lagrange_rank = (lD > 1.e-10).sum()
    all_close(lD[:lagrange_rank] + P.lagrange, D[:lagrange_rank], 'proximal method lagrange', None)

    bD = np.linalg.svd(RB, full_matrices=0)[1]
    bound_rank = (bD > 1.e-10).sum()

    all_close(bD[:bound_rank], projl1(D, B.bound)[:bound_rank], 'proximal method bound', None)

    nt.assert_true(np.linalg.norm(RPO+RB-X) / np.linalg.norm(X) < 0.01)
    nt.assert_true(np.linalg.norm(RBO+RP-X) / np.linalg.norm(X) < 0.01)

def test_gaussian(n=100, p=20):

    y = np.random.standard_normal(n)
    X = np.random.standard_normal((n,p))

    lam_theor = np.mean(np.fabs(np.dot(X.T, np.random.standard_normal((n, 1000)))).max(0))
    Q = identity_quadratic(0.01, 0, np.ones(p), 0)

    weights_with_zeros = 0.1 * np.ones(p)
    weights_with_zeros[:3] = 0.

    for q, fw in product([Q, None],
                         [0.5*lam_theor, weights_with_zeros]):

        L = lasso.gaussian(X, y, fw, 1., quadratic=Q)
        L.fit()
        C = L.constraints

        I = L.intervals
        S = L.summary('onesided')
        S = L.summary('twosided')

        yield (np.testing.assert_array_less,
               np.dot(L.constraints.linear_part, L._onestep),
               L.constraints.offset)

def test_conjugate_sqerror():
    """
    This verifies the conjugate class can compute the conjugate
    of a quadratic function.
    """

    ridge_coef = 0.4

    X = np.random.standard_normal((10,4))
    Y = np.random.standard_normal(10)
    l = rr.squared_error(X, Y)

    q = rr.identity_quadratic(ridge_coef,0,0,0)
    atom_conj = rr.conjugate(l, q, tol=1.e-12, min_its=100)
    w = np.random.standard_normal(4)
    u11, u12 = atom_conj.smooth_objective(w)

    # check that objective is half of squared error
    np.testing.assert_allclose(l.smooth_objective(w, mode='func'), 0.5 * np.linalg.norm(Y - np.dot(X, w))**2)
    np.testing.assert_allclose(atom_conj.atom.smooth_objective(w, mode='func'), 0.5 * np.linalg.norm(Y - np.dot(X, w))**2)

    XTX = np.dot(X.T, X) 
    XTXi = np.linalg.pinv(XTX)

    quadratic_term = XTX + ridge_coef * np.identity(4)
    linear_term = np.dot(X.T, Y) + w
    b = u22 = np.linalg.solve(quadratic_term, linear_term)
    u21 = (w*u12).sum() - l.smooth_objective(u12, mode='func') - q.objective(u12, mode='func')
    np.testing.assert_allclose(u12, u22, rtol=1.0e-05)
    np.testing.assert_approx_equal(u11, u21)

def test_conjugate_l1norm():
    '''
    this test verifies that numerically computing the conjugate
    is essentially the same as using the smooth_conjugate
    of the atom
    '''

    q = rr.identity_quadratic(1.2,0,0,0)
    l1 = rr.l1norm(4, lagrange=0.3)
    pen2 = copy(l1)
    pen2.set_quadratic(q)

    v1 = rr.smooth_conjugate(l1, q)
    v2 = rr.conjugate(l1, q, tol=1.e-12, min_its=100)
    v3 = rr.conjugate(pen2, None, tol=1.e-12, min_its=100)
    w = np.random.standard_normal(4)

    u11, u12 = v1.smooth_objective(w)
    u21, u22 = v2.smooth_objective(w)
    u31, u32 = v3.smooth_objective(w)
    np.testing.assert_approx_equal(u11, u21)
    np.testing.assert_allclose(u12, u22, rtol=1.0e-05)
    np.testing.assert_approx_equal(u11, u31)
    np.testing.assert_allclose(u12, u32, rtol=1.0e-05)

    v2.smooth_objective(w, mode='func')
    v2.smooth_objective(w, mode='grad')
    nt.assert_raises(ValueError, v2.smooth_objective, w, 'blah')

    def __init__(self, loss, 
                 linear_randomization,
                 quadratic_coef,
                 randomization, 
                 penalty,
                 solve_args={'tol':1.e-10, 'min_its':100, 'max_its':500}):

        (self.loss,
         self.linear_randomization,
         self.randomization,
         self.quadratic_coef) = (loss,
                                 linear_randomization,
                                 randomization,
                                 quadratic_coef)

        # initialize optimization problem

        self.penalty = penalty
        self.problem = rr.simple_problem(loss, penalty)

        random_term = rr.identity_quadratic(
                                quadratic_coef, 0, 
                                self.linear_randomization, 0)

        self.initial_soln = self.problem.solve(random_term,
                                               **solve_args)
        self.initial_grad = self.loss.smooth_objective(self.initial_soln, 
                                                       mode='grad')
        self.opt_vars = self.penalty.setup_sampling( \
            self.initial_grad,
            self.initial_soln,
            self.linear_randomization,
            self.quadratic_coef)

def test_lasso():
    '''
    this test verifies that the l1 prox can be solved
    by a primal/dual specification 

    obviously, we don't to solve the l1 prox this way,
    but it verifies that specification is working correctly

    '''

    l1 = rr.l1norm(4, lagrange=2.)
    l1.quadratic = rr.identity_quadratic(0.5, 0, None, 0.)

    X = np.random.standard_normal((10,4))
    Y = np.random.standard_normal(10) + 3
    
    loss = rr.quadratic.affine(X, -Y, coef=0.5)

    p2 = rr.separable_problem.singleton(l1, loss)
    solver2 = rr.FISTA(p2)
    solver2.fit(tol=1.0e-14, min_its=100)


    f = p2.objective
    ans = scipy.optimize.fmin_powell(f, np.zeros(4), ftol=1.0e-12, xtol=1.e-10)

    print(f(solver2.composite.coefs), f(ans))
    np.testing.assert_allclose(ans + 0.1, solver2.composite.coefs + 0.1, rtol=1.e-3)

def test_proximal_maps():

    X = np.random.standard_normal((100, 50))
    X[:,:7] *= 5

    P = FM.nuclear_norm(X.shape, lagrange=1)
    RP = todense(P.lagrange_prox(X))

    B = FM.nuclear_norm(X.shape, bound=1)
    RB = todense(B.bound_prox(X))

    BO = FM.operator_norm(X.shape, bound=1)
    PO = FM.operator_norm(X.shape, lagrange=1)

    RPO = todense(PO.lagrange_prox(X))
    RBO = todense(BO.bound_prox(X))

    D = np.linalg.svd(X, full_matrices=0)[1]
    lD = np.linalg.svd(RP, full_matrices=0)[1]
    lagrange_rank = (lD > 1.e-10).sum()
    all_close(lD[:lagrange_rank] + P.lagrange, D[:lagrange_rank], 'proximal lagrange', None)

    bD = np.linalg.svd(RB, full_matrices=0)[1]
    bound_rank = (bD > 1.e-10).sum()

    all_close(bD[:bound_rank], projl1(D, B.bound)[:bound_rank], 'proximal bound', None)

    nt.assert_true(np.linalg.norm(RPO+RB-X) / np.linalg.norm(X) < 0.01)
    nt.assert_true(np.linalg.norm(RBO+RP-X) / np.linalg.norm(X) < 0.01)

    # running code to ensure it is tested

    P.conjugate
    P.quadratic = identity_quadratic(1, 0, 0, 0)
    P.conjugate

    BO.conjugate
    BO.quadratic = identity_quadratic(1, 0, 0, 0)
    BO.conjugate

    B.conjugate
    B.quadratic = identity_quadratic(1, 0, 0, 0)
    B.conjugate

    PO.conjugate
    PO.quadratic = identity_quadratic(1, 0, 0, 0)
    PO.conjugate