#include "mex.h"
#include <stdio.h>
#include <math.h>
#include <string.h>


/*
 * Important Notice: September 20, 2010
 *
 * In this head file, we assume that the features in the tree strucutre
 * are well ordered. That is to say, the indices of the left nodes is always less
 * than the right nodes. Ideally, this can be achieved by reordering the features.
 *
 * The advantage of this ordered features is that, we donot need to use an explicit
 * variable for recording the indices.
 *
 * To deal with the more general case when the features might not be well ordered,
 * we provide the functions in the head file "general_altra.h". Compared with the files in this head file,
 * we need an additional parameter G, which contains the indices of the nodes.
 *
 *
 */

/*
 * -------------------------------------------------------------------
 *                       Functions and parameter
 * -------------------------------------------------------------------
 *
 * altra solves the following problem
 *
 * 1/2 \|x-v\|^2 + \sum \lambda_i \|x_{G_i}\|,
 *
 * where x and v are of dimension n,
 *       \lambda_i >=0, and G_i's follow the tree structure
 *
 * It is implemented in Matlab as follows:
 *
 * x=altra(v, n, ind, nodes);
 *
 * ind is a 3 x nodes matrix.
 *       Each column corresponds to a node.
 *
 *       The first element of each column is the starting index,
 *       the second element of each column is the ending index
 *       the third element of each column corrreponds to \lambbda_i.
 *
 * -------------------------------------------------------------------
 *                       Notices:
 * -------------------------------------------------------------------
 *
 * 1. The nodes in the parameter "ind" should be given in the 
 *    either
 *           the postordering of depth-first traversal
 *    or 
 *           the reverse breadth-first traversal.
 *
 * 2. When each elements of x are penalized via the same L1 
 *    (equivalent to the L2 norm) parameter, one can simplify the input
 *    by specifying 
 *           the "first" column of ind as (-1, -1, lambda)
 *
 *    In this case, we treat it as a single "super" node. Thus in the value
 *    nodes, we only count it once.
 *
 * 3. The values in "ind" are in [1,n].
 *
 * 4. The third element of each column should be positive. The program does
 *    not check the validity of the parameter. 
 *
 *    It is still valid to use the zero regularization parameter.
 *    In this case, the program does not change the values of 
 *    correponding indices.
 *    
 *
 * -------------------------------------------------------------------
 *                       History:
 * -------------------------------------------------------------------
 *
 * Composed by Jun Liu on April 20, 2010
 *
 * For any question or suggestion, please email j.liu@asu.edu.
 *
 */


void altra(double *x, double *v, int n, double *ind, int nodes){
    
    int i, j, m;
    double lambda,twoNorm, ratio;
    
    /*
     * test whether the first node is special
     */
    if ((int) ind[0]==-1){
        
        /*
         *Recheck whether ind[1] equals to zero
         */
        if ((int) ind[1]!=-1){
            printf("\n Error! \n Check ind");
            exit(1);
        }        
        
        lambda=ind[2];
        
        for(j=0;j<n;j++){
            if (v[j]>lambda)
                x[j]=v[j]-lambda;
            else
                if (v[j]<-lambda)
                    x[j]=v[j]+lambda;
                else
                    x[j]=0;
        }
        
        i=1;
    }
    else{
        memcpy(x, v, sizeof(double) * n);
        i=0;
    }
            
    /*
     * sequentially process each node
     *
     */
	for(;i < nodes; i++){
        /*
         * compute the L2 norm of this group         
         */
		twoNorm=0;
		for(j=(int) ind[3*i]-1;j< (int) ind[3*i+1];j++)
			twoNorm += x[j] * x[j];        
        twoNorm=sqrt(twoNorm);
        
        lambda=ind[3*i+2];
        if (twoNorm>lambda){
            ratio=(twoNorm-lambda)/twoNorm;
            
            /*
             * shrinkage this group by ratio
             */
            for(j=(int) ind[3*i]-1;j<(int) ind[3*i+1];j++)
                x[j]*=ratio;            
        }
        else{
            /*
             * threshold this group to zero
             */
            for(j=(int) ind[3*i]-1;j<(int) ind[3*i+1];j++)
                x[j]=0;
        }
	}
}



/*
 * altra_mt is a generalization of altra to the 
 * 
 * multi-task learning scenario (or equivalently the multi-class case)
 *
 * altra_mt(X, V, n, k, ind, nodes);
 *
 * It applies altra for each row (1xk) of X and V
 *
 */


void altra_mt(double *X, double *V, int n, int k, double *ind, int nodes){
    int i, j;
    
    double *x=(double *)malloc(sizeof(double)*k);
    double *v=(double *)malloc(sizeof(double)*k);
    
    for (i=0;i<n;i++){
        /*
         * copy a row of V to v
         *         
         */
        for(j=0;j<k;j++)
            v[j]=V[j*n + i];
        
        altra(x, v, k, ind, nodes);
        
        /*
         * copy the solution to X         
         */        
        for(j=0;j<k;j++)
            X[j*n+i]=x[j];
    }
    
    free(x);
    free(v);
}




/*
 * compute
 *  lambda2_max=computeLambda2Max(x,n,ind,nodes);
 *
 * compute the 2 norm of each group, which is divided by the ind(3,:),
 * then the maximum value is returned
 */
    /*
     *This function does not consider the case ind={[-1, -1, 100]',...}
     *
     *This functions is not used currently.
     */

void computeLambda2Max(double *lambda2_max, double *x, int n, double *ind, int nodes){
    int i, j, m;
    double lambda,twoNorm;
    
    *lambda2_max=0;
    
    for(i=0;i < nodes; i++){
        /*
         * compute the L2 norm of this group         
         */
		twoNorm=0;
		for(j=(int) ind[3*i]-1;j< (int) ind[3*i+1];j++)
			twoNorm += x[j] * x[j];        
        twoNorm=sqrt(twoNorm);
        
        twoNorm=twoNorm/ind[3*i+2];
        
        if (twoNorm >*lambda2_max )
            *lambda2_max=twoNorm;        
	}
}

/*
 * -------------------------------------------------------------------
 *                       Function and parameter
 * -------------------------------------------------------------------
 *
 * treeNorm compute
 *
 *        \sum \lambda_i \|x_{G_i}\|,
 *
 * where x is of dimension n,
 *       \lambda_i >=0, and G_i's follow the tree structure
 *
 * The file is implemented in the following in Matlab:
 *
 * tree_norm=treeNorm(x, n, ind,nodes);
 */


void treeNorm(double *tree_norm, double *x, int n, double *ind, int nodes){
    
    int i, j, m;
    double twoNorm, lambda;
    
    *tree_norm=0;
    
    /*
     * test whether the first node is special
     */
    if ((int) ind[0]==-1){
        
        /*
         *Recheck whether ind[1] equals to zero
         */
        if ((int) ind[1]!=-1){
            printf("\n Error! \n Check ind");
            exit(1);
        }        
        
        lambda=ind[2];
        
        for(j=0;j<n;j++){
            *tree_norm+=fabs(x[j]);
        }
        
        *tree_norm=*tree_norm * lambda;
        
        i=1;
    }
    else{
        i=0;
    }
            
    /*
     * sequentially process each node
     *
     */
	for(;i < nodes; i++){
        /*
         * compute the L2 norm of this group         
         */
		twoNorm=0;
		for(j=(int) ind[3*i]-1;j< (int) ind[3*i+1];j++)
			twoNorm += x[j] * x[j];        
        twoNorm=sqrt(twoNorm);
        
        lambda=ind[3*i+2];
        
        *tree_norm=*tree_norm + lambda*twoNorm;
	}
}


/*
 * -------------------------------------------------------------------
 *                       Function and parameter
 * -------------------------------------------------------------------
 *
 * findLambdaMax compute
 * 
 * the lambda_{max} that achieves a zero solution for
 *
 *     min  1/2 \|x-v\|^2 +  \lambda_{\max} * \sum  w_i \|x_{G_i}\|,
 *
 * where x is of dimension n,
 *       w_i >=0, and G_i's follow the tree structure
 *
 * The file is implemented in the following in Matlab:
 *
 * lambdaMax=findLambdaMax(v, n, ind,nodes);
 */

void findLambdaMax(double *lambdaMax, double *v, int n, double *ind, int nodes){
 
    int i, j;
    double lambda=0,squaredWeight=0, lambda1,lambda2;
    double *x=(double *)malloc(sizeof(double)*n);
    double *ind2=(double *)malloc(sizeof(double)*nodes*3);
    int num=0;
       
    for(i=0;i<n;i++){
        lambda+=v[i]*v[i];
    }
    
    if ( (int)ind[0]==-1 )
        squaredWeight=n*ind[2]*ind[2];
    else
        squaredWeight=ind[2]*ind[2];
    
    for (i=1;i<nodes;i++){
        squaredWeight+=ind[3*i+2]*ind[3*i+2];
    }
    
    /* set lambda to an initial guess
     */
    lambda=sqrt(lambda/squaredWeight);
    
    /*
    printf("\n\n   lambda=%2.5f",lambda);
    */
    
    /*
     *copy ind to ind2,
     *and scale the weight 3*i+2
     */
    for(i=0;i<nodes;i++){
        ind2[3*i]=ind[3*i];
        ind2[3*i+1]=ind[3*i+1];
        ind2[3*i+2]=ind[3*i+2]*lambda;
    }
    
    /* test whether the solution is zero or not
     */
    altra(x, v, n, ind2, nodes);    
    for(i=0;i<n;i++){
        if (x[i]!=0)
            break;
    }
    
    if (i>=n) {
        /*x is a zero vector*/
        lambda2=lambda;
        lambda1=lambda;
        
        num=0;
        
        while(1){
            num++;
            
            lambda2=lambda;
            lambda1=lambda1/2;
            /* update ind2
             */
            for(i=0;i<nodes;i++){
                ind2[3*i+2]=ind[3*i+2]*lambda1;
            }
            
            /* compute and test whether x is zero
             */
            altra(x, v, n, ind2, nodes);
            for(i=0;i<n;i++){
                if (x[i]!=0)
                    break;
            }
            
            if (i<n){
                break;
                /*x is not zero
                 *we have found lambda1
                 */
            }
        }

    }
    else{
        /*x is a non-zero vector*/
        lambda2=lambda;
        lambda1=lambda;
        
        num=0;
        while(1){
            num++;            
            
            lambda1=lambda2;
            lambda2=lambda2*2;
            /* update ind2
             */
            for(i=0;i<nodes;i++){
                ind2[3*i+2]=ind[3*i+2]*lambda2;
            }
            
            /* compute and test whether x is zero
             */
            altra(x, v, n, ind2, nodes);
            for(i=0;i<n;i++){
                if (x[i]!=0)
                    break;
            }
            
            if (i>=n){
                break;
                /*x is a zero vector
                 *we have found lambda2
                 */
            }
        }
    }    
    
    /*
    printf("\n num=%d, lambda1=%2.5f, lambda2=%2.5f",num, lambda1,lambda2);
    */
    
    while ( fabs(lambda2-lambda1) > lambda2 * 1e-10 ){
        
        num++;
        
        lambda=(lambda1+lambda2)/2;
        
        /* update ind2
         */
        for(i=0;i<nodes;i++){
            ind2[3*i+2]=ind[3*i+2]*lambda;
        }
        
        /* compute and test whether x is zero
         */
        altra(x, v, n, ind2, nodes);
        for(i=0;i<n;i++){
            if (x[i]!=0)
                break;
        }
        
        if (i>=n){
            lambda2=lambda;
        }
        else{
            lambda1=lambda;
        }
        
       /*
        printf("\n lambda1=%2.5f, lambda2=%2.5f",lambda1,lambda2);
        */
    }
    
   /*
    printf("\n num=%d",num);
    
    printf("   lambda1=%2.5f, lambda2=%2.5f",lambda1,lambda2);
     
    */
    
    *lambdaMax=lambda2;
    
    free(x);
    free(ind2);
}


/*
 * findLambdaMax_mt is a generalization of findLambdaMax to the 
 * 
 * multi-task learning scenario (or equivalently the multi-class case)
 *
 * lambdaMax=findLambdaMax_mt(X, V, n, k, ind, nodes);
 *
 * It applies findLambdaMax for each row (1xk) of X and V
 *
 */


void findLambdaMax_mt(double *lambdaMax, double *V, int n, int k, double *ind, int nodes){
    int i, j;
    
    double *v=(double *)malloc(sizeof(double)*k);
    double lambda;
    
    *lambdaMax=0;
    
    for (i=0;i<n;i++){
        /*
         * copy a row of V to v
         *         
         */
        for(j=0;j<k;j++)
            v[j]=V[j*n + i];
        
        findLambdaMax(&lambda, v, k, ind, nodes);
        
        /*
        printf("\n   lambda=%5.2f",lambda);        
         */
        
        if (lambda>*lambdaMax)
            *lambdaMax=lambda;
    }
    
    /*
    printf("\n *lambdaMax=%5.2f",*lambdaMax);
     */
    
    free(v);
}

