import numpy as np
import numpy.random as rd
from .compact_func import Classo,pathlasso


'''
Here is the function that does stability selection. It returns the distribution as an d-array. 

There is three different stability selection methods implemented here : 'first' ; 'max' ; 'lam'

    - 'first' will compute the whole path until q parameters pop.  
               It will then look at those paremeters, and repeat it for each subset of sample. 
               It this case it will also return distr_path wich is an n_lam x d - array . It is usefull to have it is one want to plot it.
               
   - 'max' will do the same but it will stop at a certain lamin that is set at 1e-2 * lambdamax here, 
              then will look at the q parameters for which the  max_lam (|beta_i(lam)|) is the highest. 
              
  - 'lam' will, for each subset of sample, compute the classo solution at a fixed lambda. That it will look at the q highest value of |beta_i(lam)|.
              
'''



def stability(matrix,StabSelmethod = 'first',numerical_method = "Path-Alg", Nlam = 100, lamin=1e-2,
              lam = 0.1, q = 10 ,B = 50, percent_nS = 0.5 ,
              formulation = 'LS', seed = 1, rho=1.345,
              rho_classification=-1.,
              true_lam = False, e=1., w=None, intercept=False):
    

    rd.seed(seed)
    n, d = len(matrix[2]), len(matrix[0][0])
    if intercept : d+=1
    nS = int(percent_nS*n)
    distribution=np.zeros(d)
    
    lambdas= np.linspace(1.,lamin,Nlam)



    if (StabSelmethod == 'first') :
    
        
        distr_path = np.zeros((Nlam,d))
        for i in range(B):
            subset = build_subset(n,nS)
            submatrix = build_submatrix(matrix,subset)
            # compute the path until n_active = q.
            BETA = np.array(pathlasso(submatrix,lambdas=lambdas,n_active=q+1,lamin=0,
                             typ=formulation, meth = numerical_method,
                             rho = rho, rho_classification=rho_classification,e=e*percent_nS, w=w, intercept=intercept)[0])

            distr_path = distr_path + (abs(BETA) >= 1e-1)
        distribution = distr_path[-1]
        return(distribution * 1./B, distr_path * 1./B,lambdas)
    
    elif (StabSelmethod == 'lam') :
    
        for i in range(B):
            subset = build_subset(n,nS)
            submatrix = build_submatrix(matrix,subset)
            regress = Classo(submatrix,lam,typ = formulation,
                             meth=numerical_method, rho = rho, rho_classification=rho_classification, e=e*percent_nS, true_lam = true_lam, w=w, intercept=intercept)
            if (type(regress)==tuple): beta =regress[0]
            else : beta = regress
            qbiggest = biggest_indexes(abs(beta),q)
            for i in qbiggest:
                distribution[i]+=1
                
    
    elif (StabSelmethod == 'max') :
        
       
        
        for i in range(B):
            subset = build_subset(n,nS)
            submatrix = build_submatrix(matrix,subset)
            # compute the path until n_active = q, and only take the last Beta
            BETA = pathlasso(submatrix,n_active=0,lambdas=lambdas,
                             typ=formulation,meth = numerical_method,
                             rho = rho, rho_classification=rho_classification, e=e*percent_nS , w=w, intercept=intercept)[0]
            betamax = np.amax( abs(np.array(BETA)), axis = 0 )
            qmax = biggest_indexes(betamax,q)
            for i in qmax:
                distribution[i]+=1           
                

    return(distribution * 1./B)


'''
Auxilaries functions that are used in the main function which is stability

'''


# returns the list of the q highest componants of an array, using the fact that it is probably sparse.
def biggest_indexes(array,q):
    qbiggest = []
    nonnul = non_nul_indices(array)
    reduc_array = array[nonnul]
    for i1 in range(q):
        if not np.any(nonnul): break
        reduc_index = np.argmax(reduc_array)
        index = nonnul[reduc_index]
        if (reduc_array[reduc_index]==0.): break
        reduc_array[reduc_index]=0.
        qbiggest.append(index)
    return(qbiggest)

# return the list of indices where the componant of the array is null
def non_nul_indices(array):
    L = []
    for i in range(len(array)):
        if not (array[i]==0.):L.append(i)
    return(L)

# for a certain threshold, it returns the features that should be selected
def selected_param(distribution,threshold,threshold_label):
    selected, to_label  = [False]*len(distribution), [False]*len(distribution)
    for i in range(len(distribution)):
        if (distribution[i] > threshold): selected[i]=True
        if (distribution[i] > threshold_label): to_label[i]=True
    return(np.array(selected), np.array(to_label))


# submatrices associated to this subset
def build_submatrix(matrix,subset):
    (A,C,y) = matrix
    subA,suby = A[subset],y[subset]
    return((subA,C,suby))

# random subset of [1,n] of size nS
def build_subset(n,nS): return(rd.permutation(n)[:nS])
