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

#define all_n n = 0; n < nNodes; n++
#define all_state state = 0; state < nStates; state++
#define all_state1 state1 = 0; state1 < nStates; state1++
#define all_state2 state2 = 0; state2 < nStates; state2++
#define all_features f = 0; f < nFeatureTypes; f++
#define nodePot(a,b) nodePot[a + nNodes*(b)]
#define edgePot(a,b) edgePot[a + nStates*(b)]
#define nodeBel(a,b) nodeBel[a + nNodes*(b)]
#define edgeBel(a,b,c) edgeBel[a + nStates*(b + nStates*(c))]
#define alpha(a,b) alpha[a + nNodes*(b)]
#define beta(a,b) beta[a + nNodes*(b)]
#define tmp(a,b) tmp[a + nStates*(b)]
#define v(a,b) v[a + nStates*(b)]

int maxVal(int x,int y) {
    if(x >= y)
        return x;
    else
        return y;
}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
    int n,s,state,state1,state2,nStates,nNodes,*sentences,nSentences,maxSentenceLength;
    int f,feature,featureInd,*nFeatures,nFeatureTypes,nFeaturesTotal,*featureStart,nParams;
    int *X,*y,nWords;
    int iter,*iVals,maxIter,average;
    double *wv,*nodePot,*nodeBel,*edgePot,*edgeBel,*alpha,*beta,*tmp,*w,*v_start,*v_end,*v,Z,logZ,*kappa;
    double lambda, *stepSizes, rho=1.0, rhoA = 1.0, step, stepA, *wvAvg, *wAvg, *vAvg_start, *vAvg_end, *vAvg, tau = 0;
    mxArray *mxA;
        
    if (nrhs != 6)
        mexErrMsgTxt("6 arguments are needed: {wv,data,lambda,stepSizes,iVals,average}");
    
    
    /* Process optimization arguments */
    lambda = mxGetScalar(prhs[2]);
    stepSizes = mxGetPr(prhs[3]);
    iVals = (int*)mxGetPr(prhs[4]);
    average = (int)mxGetScalar(prhs[5]);
    maxIter = maxVal(mxGetM(prhs[3]),mxGetN(prhs[3]));
    if (!mxIsClass(prhs[4],"int32"))
        mexErrMsgTxt("iVals must be int32");
    if (maxIter != maxVal(mxGetM(prhs[4]),mxGetN(prhs[4])))
        mexErrMsgTxt("iVals and stepSizes must have the same length");
    
    /* ******************************
     * Process data struct 
     ******************************** */
    if (!mxIsStruct(prhs[1]))
        mexErrMsgTxt("data argument needs to be a struct");
    
    mxA = mxGetField(prhs[1],0,"X");
    X = (int*)mxGetPr(mxA);
    nWords = mxGetM(mxA);
    
    y = (int*)mxGetPr(mxGetField(prhs[1],0,"y"));
    
    nStates = (int)mxGetScalar(mxGetField(prhs[1],0,"nStates"));
    
    mxA = mxGetField(prhs[1],0,"sentences");
    sentences = (int*)mxGetPr(mxA);
    nSentences = mxGetM(mxA);
    
    maxSentenceLength = (int)mxGetScalar(mxGetField(prhs[1],0,"maxSentenceLength"));
    
    featureStart = (int*)mxGetPr(mxGetField(prhs[1],0,"featureStart"));
    
    mxA = mxGetField(prhs[1],0,"nFeatures");
    nFeatures = (int*)mxGetPr(mxA);
    nFeatureTypes = maxVal(mxGetM(mxA),mxGetN(mxA));
    nFeaturesTotal = featureStart[nFeatureTypes]-1;
    nParams = nFeaturesTotal*nStates + nStates + nStates + nStates*nStates;
    
    /* ******************************
     * Initialize intermediate values
     ******************************** */
    nodePot = mxCalloc(maxSentenceLength*nStates,sizeof(double));
    edgePot = mxCalloc(nStates*nStates,sizeof(double));
    nodeBel = mxCalloc(maxSentenceLength*nStates,sizeof(double));
    edgeBel = mxCalloc(maxSentenceLength*nStates*nStates,sizeof(double));
    alpha = mxCalloc(maxSentenceLength*nStates,sizeof(double));
    beta = mxCalloc(maxSentenceLength*nStates,sizeof(double));
    tmp = mxCalloc(nStates*nStates,sizeof(double));
    kappa = mxCalloc(maxSentenceLength,sizeof(double));
           
    wv = mxGetPr(prhs[0]);
    w = wv;
    v_start = &wv[nFeaturesTotal*nStates];
    v_end = &wv[nFeaturesTotal*nStates + nStates];
    v = &wv[nFeaturesTotal*nStates + nStates + nStates];
                
    if(average) {
        plhs[0] = mxCreateDoubleMatrix(nParams,1,mxREAL);
        wvAvg = mxGetPr(plhs[0]);
        wAvg = wvAvg;
        vAvg_start = &wvAvg[nFeaturesTotal*nStates];
        vAvg_end = &wvAvg[nFeaturesTotal*nStates + nStates];
        vAvg = &wvAvg[nFeaturesTotal*nStates + nStates + nStates];
    }
    
    for(iter=0; iter < maxIter; iter++) {
        s = iVals[iter]-1;
    
        /* Compute node potentials */
        nNodes = sentences[s + nSentences]-sentences[s]+1;
        /* Compute nodePot */
        for(all_n) { 
            for(all_state)
                nodePot(n,state) = 0;
            for(f = 0; f < nFeatureTypes; f++) {
                feature = X[sentences[s]-1 + n + nWords*f];
                if(feature != 0) {
                    featureInd = featureStart[f]-1 + feature-1;
                    for(all_state) {
                        nodePot(n,state) += w[featureInd + nFeaturesTotal*state];
                    }
                }
            }
        }
        /* Add beginning/end of sentence modification */
        for(all_state) { 
            nodePot(0,state) += v_start[state];
            nodePot(nNodes-1,state) += v_end[state];
        }   
        /* Exponentiate Potentials */
        for(all_n) { 
            for(all_state) {
                nodePot(n,state) = exp(rho*nodePot(n,state));
            }
        }
        /* Edge potentials */
        for(all_state2) {
            for(all_state1) {
                edgePot(state1,state2) = exp(rho*v(state1,state2));
            }
        }
               
        
        /* Forward-Backward to Compute Marginals */
        for(all_n)
            kappa[n] = 0;
        for(all_state) {
            alpha(0,state) = nodePot(0,state);
            kappa[0] += alpha(0,state);
        }
        for(all_state)
            alpha(0,state) /= kappa[0];
        /* Forward Pass */
        for(n = 1; n < nNodes; n++) {
            for(all_state1) {
                for(all_state2)
                    tmp(state1,state2) = alpha(n-1,state1) * edgePot(state1,state2);
            }
            for(all_state2) {
                alpha(n,state2) = 0;
                for(all_state1)
                    alpha(n,state2) += tmp(state1,state2);
                alpha(n,state2) = nodePot(n,state2) * alpha(n,state2);
                kappa[n] += alpha(n,state2);
            }
            for(all_state)
                alpha(n,state) /= kappa[n];
        }
        /* Backward Pass */
       for(all_state)
            beta(nNodes-1,state) = 1;
        for(n = nNodes-2; n >= 0; n--) {
            for(all_state1) {
                for(all_state2)
                    tmp(state1,state2) = nodePot(n+1,state2)*edgePot(state1,state2)*beta(n+1,state2);
            }
            Z = 0;
            for(all_state1) {
                beta(n,state1) = 0;
                for(all_state2)
                    beta(n,state1) += tmp(state1,state2);
                Z += beta(n,state1);
            }
            for(all_state)
                beta(n,state) /= Z;
        }
        
        /* Compute beliefs for use in gradient update */
        /* Node Beliefs */
        for(all_n) {
            Z = 0;
            for(all_state) {
                nodeBel(n,state) = alpha(n,state)*beta(n,state);
                Z += nodeBel(n,state);
            }
            for(all_state)
                nodeBel(n,state) /= Z;
        }         
        /* Edge Beliefs */
        for(n = 0; n < nNodes-1; n++) {
            Z = 0;
            for(all_state1) {
                for(all_state2) {
                    tmp(state1,state2) = alpha(n,state1)*nodePot(n+1,state2)*beta(n+1,state2)*edgePot(state1,state2);
                    Z += tmp(state1,state2);
                }
            }
            for(all_state1) {
                for(all_state2) {
                    edgeBel(state1,state2,n) = tmp(state1,state2)/Z;
                }
            }
        }
                   
        /* Update parameters */
        if (stepSizes[iter]*lambda != 1)
            rho *= 1-stepSizes[iter]*lambda;
        else
            mexErrMsgTxt("I do not like it when stepSize[iter]*lambda = 1!");
        step = stepSizes[iter]/rho;
        
        if(average)
            stepA = -tau*step;

        /* Node parameters */
        for(all_n) {
            for(f = 0; f < nFeatureTypes; f++) {
                feature = X[sentences[s]-1 + n + nWords*f];
                if(feature != 0) {
                    featureInd = featureStart[f]-1 + feature-1;
                    for(all_state) {
                        w[featureInd + nFeaturesTotal*state] -= step*nodeBel(n,state);
                        if (average)
                            wAvg[featureInd + nFeaturesTotal*state] -= stepA*nodeBel(n,state);
                    }
                    state = y[sentences[s]-1 +n]-1;
                    w[featureInd + nFeaturesTotal*state] += step;
                    if (average)
                        wAvg[featureInd + nFeaturesTotal*state] += stepA;
                }
            }
        }
        /* Start/end parameters */
        for(all_state) {
            v_start[state] -= step*nodeBel(0,state);
            v_end[state] -= step*nodeBel(nNodes-1,state);
            if(average) {
                vAvg_start[state] -= stepA*nodeBel(0,state);
                vAvg_end[state] -= stepA*nodeBel(nNodes-1,state);
            }
        }
        v_start[y[sentences[s]-1]-1] += step;
        v_end[y[sentences[s]-1 + nNodes-1]-1] += step;
        if (average) {
            vAvg_start[y[sentences[s]-1]-1] += stepA;
            vAvg_end[y[sentences[s]-1 + nNodes-1]-1] += stepA;
        }
        /* Edge parameters */
        for(n = 0; n < nNodes-1; n++)
        {
            for(all_state1) {
                for(all_state2) {
                    v[state1 + nStates*state2] -= step*edgeBel(state1,state2,n);
                    if (average)
                        vAvg[state1 + nStates*state2] -= stepA*edgeBel(state1,state2,n);
                }
            }
            state1 = y[sentences[s]-1 +n]-1;
            state2 = y[sentences[s]-1 +n+1]-1;
            v[state1 + nStates*state2] += step;
            if (average)
                vAvg[state1 + nStates*state2] += stepA;
        }
        
        if(average) {
            if (iter > 0)
                rhoA *= (double)iter/(double)(iter+1);
            tau += (rho/rhoA)/(double)(iter+1);
        }
    }
    
    /* Scale variables to get final result */
    for(f = 0; f < nParams; f++)
        wv[f] = wv[f]*rho;
    
    if(average) {
        for(f = 0; f < nParams; f++)
            wvAvg[f] = ((tau/rho)*wv[f] + wvAvg[f])*rhoA;
    }
    
    mxFree(nodePot);
    mxFree(edgePot);
    mxFree(alpha);
    mxFree(beta);
    mxFree(tmp);
    mxFree(nodeBel);
    mxFree(edgeBel);
    mxFree(kappa);
}