/*--
function z = barqpqz(x,P,A,B)
    used in the eigensolver for constrained ncuts
Input:
    x = N x 1 real vector
    P = N x N real sparse matrix
    A = N x C real sparse matrix
    B = C x C real sparse matrix
Output:
    z = (I - A * B * A') * P * (I - A * B * A') * x    
See also: 
    cncut.m    
Reference:
    Chapter 5, page 101 of my thesis at:
        www.cs.berkeley.edu/~stellayu/thesis.html
    
Stella X. Yu, July 11 2003.

--*/

# include "math.h"
# include "mex.h"

/* define subroutines */
/* y = P x, y_i = \sum_j P_ij x_j */
/* warning: y is assumed to be initialized zero */
void Px(double *y, int *ir, int *jc, double *pr, double *x, int nc)
{
    int i, j, k;
    double s;         
    for (j=0; j<nc; j++) {
        s = x[j];
        for (k=jc[j]; k<jc[j+1]; k++) {
            i = ir[k];
            y[i] = y[i] + pr[k] * s;
        }
    }
}

/* y = y - P x, y_i = y_i - \sum_j P_ij x_j */
void IdPx(double *y, int *ir, int *jc, double *pr, double *x, int nc)
{
    int i, j, k;
    double s;

    for (j=0; j<nc; j++) {
        s = x[j];
        for (k=jc[j]; k<jc[j+1]; k++) {
            i = ir[k];
            y[i] = y[i] - pr[k] * s;
        }
    }
}

/* y = P' x, y_i = \sum_j  P_ji x_j */
void Ptx(double *y, int *ir, int *jc, double *pr, double *x, int nc)
{
    int i, j, k;
    double s;
    for (i=0; i<nc; i++) {
        s = 0;
        for (k=jc[i]; k<jc[i+1]; k++) {
            j = ir[k];
            s = s + pr[k] * x[j];
        }
        y[i] = s;
    }
}    

/* main program */    
void mexFunction(
    int nargout,
    mxArray *out[],
    int nargin,
    const mxArray *in[]
)
{    
    /* declare variables */
    int m[3], n[3], np, nc, i, k, *ir[3], *jc[3];
    double *x, *y, *z, *t, *u, *pr[3];
    
    /* check argument */
    if (nargin < 4) {
        mexErrMsgTxt("Four input arguments required");
    }
    if (nargout>1) {
        mexErrMsgTxt("Too many output arguments.");
    }
    
    if (!(mxIsDouble(in[0]))) {
        mexErrMsgTxt("First input argument must be of type double");
    }

    for (k=1; k<4; k++) {
        if (!(mxIsSparse(in[k]))) {
            mexErrMsgTxt("Input arguments must be of type sparse");
        }
    }
      
    /* read in data */
    x = mxGetPr(in[0]);
    for (k=1; k<4; k++) {
        i = k - 1;
        m[i] = mxGetM(in[k]);
        n[i] = mxGetN(in[k]);
        pr[i] = mxGetPr(in[k]);
        ir[i] = mxGetIr(in[k]);
        jc[i] = mxGetJc(in[k]);
    }    
    np = mxGetM(in[0]);
    if ( np != m[0] ||  np != n[0] || np != m[1] || n[1] != m[2] || m[2] != n[2] ) {
        mexErrMsgTxt("barqpqz: dimension mismatch.");
    }
    nc = n[1];
    
    /* create output */
    out[0] = mxCreateDoubleMatrix(np,1,mxREAL);
    z = mxCalloc(nc, sizeof(double));
    t = mxCalloc(nc, sizeof(double));
    u = mxCalloc(np, sizeof(double));
    if (out[0]==NULL || z==NULL || t==NULL || u == NULL) {
        mexErrMsgTxt("Not enough space for the output matrix.");
    }   
    y = mxGetPr(out[0]);
    
    /* computation */
    
    /* u = x - A B A' x = y - A t, t = B z, z = A' x */
    memcpy(u,x,np*sizeof(double));
    Ptx(z, ir[1],jc[1],pr[1],x,nc);
    Px(t,  ir[2],jc[2],pr[2],z,nc);    
    IdPx(u,ir[1],jc[1],pr[1],t,nc);
        
    /* y = P (I - A B A' ) x = P u */
    Px(y,  ir[0],jc[0],pr[0],u,np);
    
    /* y = (I - A B A') y = y - A t, t = B z, z = A' y */
    memset(t,0,nc*sizeof(double));
    Ptx(z, ir[1],jc[1],pr[1],y,nc);
    Px(t,  ir[2],jc[2],pr[2],z,nc);
    IdPx(y,ir[1],jc[1],pr[1],t,nc);
    
    mxFree(z);
    mxFree(t);
    mxFree(u);
}