def test_multinomial_vs_logistic():

    """
    Test that multinomial regression with two categories is the same as logistic regression
    """

    n = 500
    p = 10
    J = 2

    X = np.random.standard_normal(n*p).reshape((n,p))
    counts = np.random.randint(0,10,n*J).reshape((n,J)) + 2

    mult_x = rr.linear_transform(X, input_shape=(p,J-1))
    loss = rr.multinomial_deviance.linear(mult_x, counts=counts)
    problem = rr.container(loss)
    solver = rr.FISTA(problem)
    solver.fit(debug=False, tol=1e-10)
    coefs1 = solver.composite.coefs

    loss = rr.logistic_deviance.linear(X, successes=counts[:,0], trials = np.sum(counts, axis=1))
    problem = rr.container(loss)
    solver = rr.FISTA(problem)
    solver.fit(debug=False, tol=1e-10)
    coefs2 = solver.composite.coefs

    loss = rr.logistic_deviance.linear(X, successes=counts[:,1], trials = np.sum(counts, axis=1))
    problem = rr.container(loss)
    solver = rr.FISTA(problem)
    solver.fit(debug=False, tol=1e-10)
    coefs3 = solver.composite.coefs

    npt.assert_equal(coefs1.shape,(p,J-1))
    npt.assert_array_almost_equal(coefs1.flatten(), coefs2.flatten(), 5)
    npt.assert_array_almost_equal(coefs1.flatten(), -coefs3.flatten(), 5)

def test_lasso_separable():
    """
    This test verifies that the specification of a separable
    penalty yields the same results as having two linear_atoms
    with selector matrices. The penalty here is a lasso, i.e. l1
    penalty.
    """

    X = np.random.standard_normal((100,20))
    Y = np.random.standard_normal((100,)) + np.dot(X, np.random.standard_normal(20))

    penalty1 = rr.l1norm(10, lagrange=1.2)
    penalty2 = rr.l1norm(10, lagrange=1.2)
    penalty = rr.separable((20,), [penalty1, penalty2], [slice(0,10), slice(10,20)], test_for_overlap=True)

    # ensure code is tested

    print(penalty1.latexify())

    print(penalty.latexify())
    print(penalty.conjugate)
    print(penalty.dual)
    print(penalty.seminorm(np.ones(penalty.shape)))
    print(penalty.constraint(np.ones(penalty.shape), bound=2.))

    pencopy = copy(penalty)
    pencopy.set_quadratic(rr.identity_quadratic(1,0,0,0))
    pencopy.conjugate

    # solve using separable
    
    loss = rr.quadratic.affine(X, -Y, coef=0.5)
    problem = rr.separable_problem.fromatom(penalty, loss)
    solver = rr.FISTA(problem)
    solver.fit(min_its=200, tol=1.0e-12)
    coefs = solver.composite.coefs

    # solve using the usual composite

    penalty_all = rr.l1norm(20, lagrange=1.2)
    problem_all = rr.container(loss, penalty_all)
    solver_all = rr.FISTA(problem_all)
    solver_all.fit(min_its=100, tol=1.0e-12)

    coefs_all = solver_all.composite.coefs

    # solve using the selectors

    penalty_s = [rr.linear_atom(p, rr.selector(g, (20,))) for p, g in
                 zip(penalty.atoms, penalty.groups)]
    problem_s = rr.container(loss, *penalty_s)
    solver_s = rr.FISTA(problem_s)
    solver_s.fit(min_its=500, tol=1.0e-12)
    coefs_s = solver_s.composite.coefs

    np.testing.assert_almost_equal(coefs, coefs_all)
    np.testing.assert_almost_equal(coefs, coefs_s)

def test_conjugate_solver():

    # Solve Lagrange problem 
    Y = np.random.standard_normal(500); Y[100:150] += 7; Y[250:300] += 14
    loss = R.quadratic.shift(-Y, coef=0.5)

    sparsity = R.l1norm(len(Y), lagrange = 1.4)
    D = sparse.csr_matrix((np.identity(500) + np.diag([-1]*499,k=1))[:-1])
    fused = R.l1norm.linear(D, lagrange = 25.5)
    problem = R.container(loss, sparsity, fused)
    
    solver = R.FISTA(problem)
    solver.fit(max_its=500, tol=1e-10)
    solution = solver.composite.coefs

    # Solve constrained version
    delta1 = np.fabs(D * solution).sum()
    delta2 = np.fabs(solution).sum()
    fused_constraint = R.l1norm.linear(D, bound = delta1)
    sparsity_constraint = R.l1norm(500, bound = delta2)
    constrained_problem = R.container(loss, fused_constraint, sparsity_constraint)
    constrained_solver = R.FISTA(constrained_problem)
    vals = constrained_solver.fit(max_its=500, tol=1e-10)
    constrained_solution = constrained_solver.composite.coefs

    npt.assert_almost_equal(np.fabs(constrained_solution).sum(), delta2, 3)
    npt.assert_almost_equal(np.fabs(D * constrained_solution).sum(), delta1, 3)


    # Solve with (shifted) conjugate function

    loss = R.quadratic.shift(-Y, coef=0.5)
    true_conjugate = R.quadratic.shift(Y, coef=0.5)
    problem = R.container(loss, fused_constraint, sparsity_constraint)
    solver = R.FISTA(problem.conjugate_composite(true_conjugate))
    solver.fit(max_its=500, tol=1e-10)
    conjugate_coefs = problem.conjugate_primal_from_dual(solver.composite.coefs)
                      

    # Solve with generic conjugate function

    loss = R.quadratic.shift(-Y, coef=0.5)
    problem = R.container(loss, fused_constraint, sparsity_constraint)
    solver2 = R.FISTA(problem.conjugate_composite(conjugate_tol=1e-12))
    solver2.fit(max_its=500, tol=1e-10)
    conjugate_coefs_gen = problem.conjugate_primal_from_dual(solver2.composite.coefs)



    d1 = np.linalg.norm(solution - constrained_solution) / np.linalg.norm(solution)
    d2 = np.linalg.norm(solution - conjugate_coefs) / np.linalg.norm(solution)
    d3 = np.linalg.norm(solution - conjugate_coefs_gen) / np.linalg.norm(solution)

    npt.assert_array_less(d1, 0.01)
    npt.assert_array_less(d2, 0.01)
    npt.assert_array_less(d3, 0.01)

def test_affine_linear_offset_l1norm():

    """
    Test linear, affine and offset with the l1norm atom
    """
    
    n = 1000
    p = 10
    
    X = np.random.standard_normal((n,p))
    Y = 10*np.random.standard_normal(n)
    
    coefs = []
    
    loss = rr.quadratic.affine(X,-Y, coef=0.5)
    sparsity = rr.l1norm(p, lagrange=5.)
    problem = rr.container(loss, sparsity)
    solver = rr.FISTA(problem)
    solver.fit(debug=False, tol=1e-10)
    coefs.append(1.*solver.composite.coefs)
    
    loss = rr.quadratic.affine(X,-Y, coef=0.5)
    sparsity = rr.l1norm.linear(np.eye(p), lagrange=5.)
    problem = rr.container(loss, sparsity)
    solver = rr.FISTA(problem)
    solver.fit(debug=False, tol=1e-10)
    coefs.append(1.*solver.composite.coefs)
    
    loss = rr.quadratic.affine(X,-Y, coef=0.5)
    sparsity = rr.l1norm.affine(np.eye(p),np.zeros(p), lagrange=5.)
    problem = rr.container(loss, sparsity)
    solver = rr.FISTA(problem)
    solver.fit(debug=False, tol=1e-10)
    coefs.append(1.*solver.composite.coefs)
    
    loss = rr.quadratic.affine(X,-Y, coef=0.5)
    sparsity = rr.l1norm.linear(np.eye(p), lagrange=5., offset=np.zeros(p))
    problem = rr.container(loss, sparsity)
    solver = rr.FISTA(problem)
    solver.fit(debug=False, tol=1e-10)
    coefs.append(1.*solver.composite.coefs)
    
    loss = rr.quadratic.affine(X,-Y, coef=0.5)
    sparsity = rr.l1norm.shift(np.zeros(p), lagrange=5.)
    problem = rr.container(loss, sparsity)
    solver = rr.FISTA(problem)
    solver.fit(debug=False, tol=1e-10)
    coefs.append(1.*solver.composite.coefs)

    for i,j in itertools.combinations(range(len(coefs)), 2):
        npt.assert_almost_equal(coefs[i], coefs[j])

def test_group_lasso_sparse(n=100):

    def selector(p, slice):
        return np.identity(p)[slice]

    def selector_sparse(p, slice):
        return sparse.csr_matrix(np.identity(p)[slice])

    X = np.random.standard_normal((1000,500))
    Y = np.random.standard_normal((1000,))
    loss = R.quadratic.affine(X, -Y, coef=0.5)

    penalties = [R.l2norm.linear(selector(500, slice(i*100,(i+1)*100)), lagrange=.1) for i in range(5)]
    penalties[0].lagrange = 250.
    penalties[1].lagrange = 225.
    penalties[2].lagrange = 150.
    penalties[3].lagrange = 100.
    group_lasso = R.container(loss, *penalties)

    solver=R.FISTA(group_lasso)
    solver.debug = True
    t1 = time.time()
    vals = solver.fit(max_its=2000, min_its=20,tol=1e-8)
    soln1 = solver.composite.coefs
    t2 = time.time()
    dt1 = t2 - t1



    print soln1[range(10)]

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.0)

    l11 = rr.l1norm(4, lagrange=1.0)
    l12 = rr.l1norm(4, lagrange=1.0)

    X = np.random.standard_normal((10, 4))
    Y = np.random.standard_normal(10) + 3

    loss = rr.quadratic.affine(X, -Y)
    p1 = rr.container(l11, loss, l12)

    solver1 = rr.FISTA(p1)
    solver1.fit(tol=1.0e-12, min_its=500)

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

    f = p2.objective
    ans = scipy.optimize.fmin_powell(f, np.zeros(4), ftol=1.0e-12)
    print f(solver2.composite.coefs), f(ans)
    print f(solver1.composite.coefs), f(ans)

    yield ac, ans, solver2.composite.coefs, "singleton solver"
    yield ac, solver1.composite.coefs, solver2.composite.coefs, "container solver"

def example4(lambda1=10):

    #Example with an initial value for backtracking

    # In the previous examples you'll see a lot of "Increasing inv_step" iterations - these are trying to find an approximate Lipschitz constant in a backtracking loop.
    # For your problem the Lipschitz constant is just the largest eigenvalue of X^TX, so you can precompute this with a few power iterations.

    n = 100
    p = 1000

    X = np.random.standard_normal(n*p).reshape((n,p))
    Y = 10*np.random.standard_normal(n)

    v = np.random.standard_normal(p)
    for i in range(10):
        v = np.dot(X.T, np.dot(X,v))
        norm = np.linalg.norm(v)
        v /= norm
    print "Approximate Lipschitz constant is", norm

    loss = rr.l2normsq.affine(X,-Y,coef=1.)
    sparsity = rr.l1norm(p, lagrange = lambda1)
    nonnegative = rr.nonnegative(p)

    problem = rr.container(loss, sparsity, nonnegative)
    solver = rr.FISTA(problem)

    #Give approximate Lipschitz constant to solver
    solver.fit(debug=True, start_inv_step=norm)

    solution = solver.composite.coefs

def fused_example():

    x=np.random.standard_normal(500); x[100:150] += 7

    sparsity = R.l1norm(500, lagrange=1.3)
    D = (np.identity(500) + np.diag([-1]*499,k=1))[:-1]
    fused = R.l1norm.linear(D, lagrange=10.5)

    loss = R.quadratic.shift(-x, coef=0.5)
    pen = R.container(loss, sparsity,fused)
    solver = R.FISTA(pen)
    vals = solver.fit()
    soln = solver.composite.coefs
    
    # solution

    pylab.figure(num=1)
    pylab.clf()
    pylab.plot(soln, c='g')
    pylab.scatter(np.arange(x.shape[0]), x)

    # objective values

    pylab.figure(num=2)
    pylab.clf()
    pylab.plot(vals)

def test_lasso(n=100):

    l1 = 1.
    sparsity = R.l1norm(n, lagrange=l1)
    
    X = np.random.standard_normal((5000,n))
    Y = np.random.standard_normal((5000,))
    regloss = R.quadratic.affine(-X,Y)

    p=R.container(regloss, sparsity)
    solver=R.FISTA(p)
    solver.debug = True
    t1 = time.time()
    vals1 = solver.fit(max_its=800)
    t2 = time.time()
    dt1 = t2 - t1
    soln = solver.composite.coefs

    time.sleep(5)


    print soln[range(10)]

    print solver.composite.objective(soln)
    print "Times", dt1

def group_lasso_example():

    def selector(p, slice):
        return np.identity(p)[slice]
    penalties = [R.l2norm(selector(500, slice(i*100,(i+1)*100)), lagrange=.1) for i in range(5)]
    penalties[0].lagrange = 250.
    penalties[1].lagrange = 225.
    penalties[2].lagrange = 150.
    penalties[3].lagrange = 100.

    X = np.random.standard_normal((1000,500))
    Y = np.random.standard_normal((1000,))
    loss = R.quadratic.affine(X, -Y, coef=0.5)
    group_lasso = R.container(loss, *penalties)

    solver=R.FISTA(group_lasso)
    solver.debug = True
    vals = solver.fit(max_its=2000, min_its=20,tol=1e-10)
    soln = solver.composite.coefs

    # solution

    pylab.figure(num=1)
    pylab.clf()
    pylab.plot(soln, c='g')

    # objective values

    pylab.figure(num=2)
    pylab.clf()
    pylab.plot(vals)

def test_group_lasso_separable():
    """
    This test verifies that the specification of a separable
    penalty yields the same results as having two linear_atoms
    with selector matrices. The penalty here is a group_lasso, i.e. l2
    penalty.
    """

    X = np.random.standard_normal((100,20))
    Y = np.random.standard_normal((100,)) + np.dot(X, np.random.standard_normal(20))

    penalty1 = rr.l2norm(10, lagrange=.2)
    penalty2 = rr.l2norm(10, lagrange=.2)
    penalty = rr.separable((20,), [penalty1, penalty2], [slice(0,10), slice(10,20)])

    # solve using separable
    
    loss = rr.quadratic.affine(X, -Y, coef=0.5)
    problem = rr.separable_problem.fromatom(penalty, loss)
    solver = rr.FISTA(problem)
    solver.fit(min_its=200, tol=1.0e-12)
    coefs = solver.composite.coefs

    # solve using the selectors

    penalty_s = [rr.linear_atom(p, rr.selector(g, (20,))) for p, g in
                 zip(penalty.atoms, penalty.groups)]
    problem_s = rr.container(loss, *penalty_s)
    solver_s = rr.FISTA(problem_s)
    solver_s.fit(min_its=200, tol=1.0e-12)
    coefs_s = solver_s.composite.coefs

    np.testing.assert_almost_equal(coefs, coefs_s)

def test_lasso_dual():

    """
    Check that the solution of the lasso signal approximator dual composite is soft-thresholding
    """

    l1 = .1
    sparsity = R.l1norm(10, lagrange=l1)
    x = np.arange(10) - 5
    loss = R.quadratic.shift(-x, coef=0.5)

    pen = R.simple_problem(loss, sparsity)
    solver = R.FISTA(pen)
    pen.lipschitz = 1
    solver.fit(backtrack=False)
    soln = solver.composite.coefs
    st = np.maximum(np.fabs(x)-l1,0) * np.sign(x) 

    np.testing.assert_almost_equal(soln,st, decimal=3)

    pen = R.simple_problem(loss, sparsity)
    solver = R.FISTA(pen)
    solver.fit(monotonicity_restart=False)
    soln = solver.composite.coefs
    st = np.maximum(np.fabs(x)-l1,0) * np.sign(x) 

    np.testing.assert_almost_equal(soln,st, decimal=3)


    pen = R.container(loss, sparsity)
    solver = R.FISTA(pen)
    solver.fit()
    soln = solver.composite.coefs

    np.testing.assert_almost_equal(soln,st, decimal=3)

def lasso_via_dual_split():

    def selector(p, slice):
        return np.identity(p)[slice]
    penalties = [R.l1norm(selector(500, slice(i*100,(i+1)*100)), lagrange=0.2) for i in range(5)]
    x = np.random.standard_normal(500)
    loss = R.quadratic.shift(-x, coef=0.5)
    lasso = R.container(loss,*penalties)
    solver = R.FISTA(lasso)
    np.testing.assert_almost_equal(np.maximum(np.fabs(x)-0.2, 0) * np.sign(x), solver.composite.coefs, decimal=3)

def example3(lambda1=10):

    #Example using a smooth approximation to the non-negativity constraint
    # On large problems this might be faster than using the actual constraint

    n = 100
    p = 1000


    X = np.random.standard_normal(n*p).reshape((n,p))
    Y = 10*np.random.standard_normal(n)

    loss = rr.l2normsq.affine(X,-Y,coef=1.)
    sparsity = rr.l1norm(p, lagrange = lambda1)
    nonnegative = rr.nonnegative(p)
    smooth_nonnegative = rr.smoothed_atom(nonnegative, epsilon = 1e-4)

    problem = rr.container(loss, sparsity, smooth_nonnegative)
    solver = rr.FISTA(problem)
    solver.fit(debug=True)

    solution1 = solver.composite.coefs



    loss = rr.l2normsq.affine(X,-Y,coef=1.)
    sparsity = rr.l1norm(p, lagrange = lambda1)
    nonnegative = rr.nonnegative(p)

    problem = rr.container(loss, sparsity, nonnegative)
    solver = rr.FISTA(problem)
    solver.fit(debug=True)

    solution2 = solver.composite.coefs


    pl.subplot(1,2,1)
    pl.hist(solution1, bins=40)

    pl.subplot(1,2,2)
    pl.scatter(solution2,solution1)
    pl.xlabel("Constraint")
    pl.ylabel("Smooth constraint")

def test_logistic_counts():
    """
    Test the equivalence of binary/count specification in logistic_deviance
    """

    #Form the count version of the problem
    trials = np.random.binomial(5,0.5,100)+1
    successes = np.random.binomial(trials,0.5,len(trials)) 
    n = len(successes)
    p = 2*n
    X = np.random.normal(0,1,n*p).reshape((n,p))

    loss = rr.logistic_deviance.linear(X, successes=successes, trials=trials)
    penalty = rr.quadratic(p, coef=1.)

    prob1 = rr.container(loss, penalty)
    solver1 = rr.FISTA(prob1)
    solver1.fit()
    solution1 = solver1.composite.coefs
    
    #Form the binary version of the problem
    Ynew = []
    Xnew = []

    for i, (s,n) in enumerate(zip(successes,trials)):
        Ynew.append([1]*s + [0]*(n-s))
        for j in range(n):
            Xnew.append(X[i,:])
    Ynew = np.hstack(Ynew)
    Xnew =  np.vstack(Xnew)


    loss = rr.logistic_deviance.linear(Xnew, successes=Ynew)
    penalty = rr.quadratic(p, coef=1.)

    prob2 = rr.container(loss, penalty)
    solver2 = rr.FISTA(prob2)
    solver2.fit()
    solution2 = solver2.composite.coefs

   
    npt.assert_array_almost_equal(solution1, solution2, 3)

def group_lasso_signal_approx():

    def selector(p, slice):
        return np.identity(p)[slice]
    penalties = [R.l2norm(selector(500, slice(i*100,(i+1)*100)), lagrange=10.) for i in range(5)]
    loss = R.quadratic.shift(-x, coef=0.5)
    group_lasso = R.container(loss, **penalties)
    x = np.random.standard_normal(500)
    solver = R.FISTA(group_lasso)
    solver.fit()
    a = solver.composite.coefs

    def __init__(self, X, initial=None, lagrange=1, rho=1):
        self.X = R.affine_transform(X, None)
        self.atom = R.l1norm(X.shape[1], l)
        self.rho = rho
        self.loss = R.quadratic.affine(X, -np.zeros(X.shape[0]), lagrange=rho/2.)
        self.lasso = R.container(self.loss, self.atom)
        self.solver = R.FISTA(self.lasso.problem())

        if initial is None:
            self.beta[:] = np.random.standard_normal(self.atom.primal_shape)
        else:
            self.beta[:] = initial

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

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

        tests.append((atom.proximal(q), solver.composite.coefs, 'solving atom prox with container\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.container(lossq, atom)
        solver = rr.FISTA(problem)
        solver.fit(tol=1.0e-12, FISTA=self.FISTA, coef_stop=self.coef_stop)

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

        d = atom.conjugate
        problem = rr.container(d, loss)
        solver = rr.FISTA(problem)
        solver.fit(tol=1.0e-12, 
                   coef_stop=self.coef_stop, FISTA=self.FISTA, min_its=100)
        tests.append((d.proximal(q), solver.composite.coefs, 'solving dual prox with container\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_1d_fused_lasso():

    """
    Check the 1d fused lasso solution using an equivalent lasso formulation
    """

    n = 100
    l1 = 1.
    
    D = (np.identity(n) - np.diag(np.ones(n-1),-1))[1:]
    extra = np.zeros(n)
    extra[0] = 1.
    D = np.vstack([D,extra])
    D = sparse.csr_matrix(D)

    fused = R.l1norm.linear(D, lagrange=l1)
    
    X = np.random.standard_normal((2*n,n))
    Y = np.random.standard_normal((2*n,))
    loss = R.quadratic.affine(X, -Y, coef=0.5)
    fused_lasso = R.container(loss, fused)
    solver=R.FISTA(fused_lasso)
    vals1 = solver.fit(max_its=25000, tol=1e-10)
    soln1 = solver.composite.coefs
    
    B = np.array(sparse.tril(np.ones((n,n))).todense())
    X2 = np.dot(X,B)
    
    loss = R.quadratic.affine(X2, -Y, coef=0.5)
    sparsity = R.l1norm(n, lagrange=l1)
    lasso = R.container(loss, sparsity)
    solver = R.FISTA(lasso)
    solver.fit(tol=1e-10)

    soln2 = np.dot(B, solver.composite.coefs)

    npt.assert_array_almost_equal(soln1, soln2, 3)

def test_lasso_via_dual_split():
    """
    Test the lasso by breaking it up into multiple l1 atoms over the range of beta
    """

    def selector(p, slice):
        return np.identity(p)[slice]
    
    penalties = [R.l1norm.linear(selector(500, slice(i*100,(i+1)*100)), lagrange=0.2) for i in range(5)]
    x = np.random.standard_normal(500)
    loss = R.quadratic.shift(-x, coef=0.5)
    lasso = R.container(loss,*penalties)
    solver = R.FISTA(lasso)
    solver.fit(tol=1e-8)

    npt.assert_array_almost_equal(np.maximum(np.fabs(x)-0.2, 0) * np.sign(x), solver.composite.coefs, 3)