from __future__ import print_function
import numpy as np
import pandas as pd
import regreg.api as rr

from ...tests.flags import SET_SEED, SMALL_SAMPLES
from ...tests.instance import logistic_instance
from ...tests.decorators import (wait_for_return_value, 
                                 set_seed_iftrue, 
                                 set_sampling_params_iftrue,
                                 register_report)
import selection.tests.reports as reports

from ...api import (randomization, 
                    glm_group_lasso, 
                    pairs_bootstrap_glm, 
                    multiple_queries, 
                    discrete_family, 
                    projected_langevin, 
                    glm_group_lasso_parametric, 
                    glm_target)
from ..glm import (glm_parametric_covariance, 
                   glm_nonparametric_bootstrap, 
                   restricted_Mest, 
                   set_alpha_matrix)

@register_report(['truth', 'active'])
@set_sampling_params_iftrue(SMALL_SAMPLES, ndraw=10, burnin=10)
@set_seed_iftrue(SET_SEED)
@wait_for_return_value()
def test_multiple_queries(s=3,
                          n=200,
                          p=20,
                          signal=7,
                          rho=0.1,
                          lam_frac=0.7,
                          nview=4,
                          ndraw=100, burnin=0,
                          bootstrap=True,
                          test = 'global'):

    #randomizer = randomization.laplace((p,), scale=1)
    randomizer = randomization.logistic((p,), scale=1)
    X, y, beta, _ = logistic_instance(n=n, p=p, s=s, rho=rho, signal=signal)

    nonzero = np.where(beta)[0]

    loss = rr.glm.logistic(X, y)
    epsilon = 1.

    lam = lam_frac * np.mean(np.fabs(np.dot(X.T, np.random.binomial(1, 1. / 2, (n, 10000)))).max(0))
    W = np.ones(p)*lam
    #W[0] = 0 # use at least some unpenalized
    penalty = rr.group_lasso(np.arange(p),
                             weights=dict(zip(np.arange(p), W)), lagrange=1.)

    view = []
    for i in range(nview):
        view.append(glm_group_lasso(loss, epsilon, penalty, randomizer))


    mv = multiple_queries(view)
    mv.solve()

    active_union = np.zeros(p, np.bool)
    for i in range(nview):
        active_union += view[i].selection_variable['variables']

    nactive = np.sum(active_union)
    #print("nactive", nactive)

    if set(nonzero).issubset(np.nonzero(active_union)[0]):
        if nactive==s:
            return None

        active_set = np.nonzero(active_union)[0]

        if test == 'selected zeros':