#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 i,n,s,state,state1,state2,nStates,nNodes,*sentences,nSentences,maxSentenceLength;
    int f,feature,featureInd,*nFeatures,nFeatureTypes,nFeaturesTotal,*featureStart;
    int *X,*y,nWords;
    double *func,*wv,*nodePot,*nodeBel,*edgePot,*edgeBel,*alpha,*beta,*tmp,*w,*v_start,*v_end,*v,Z,logZ,*kappa;
    double *grad, *gw, *gv_start, *gv_end, *gv;
    mxArray *mxA;
        
    /* Process data struct */
    if (!mxIsStruct(prhs[1]))
        mexErrMsgTxt("data argument needs to be a struct: crfChain_lossC(wv,data)");
    
    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;
    
    /* 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(nlhs > 0) {
        plhs[0] = mxCreateDoubleMatrix(1,1,mxREAL);
        func = mxGetPr(plhs[0]);
        if(nlhs > 1) {
            plhs[1] = mxCreateDoubleMatrix(nFeaturesTotal*nStates + nStates + nStates + nStates*nStates,1,mxREAL);
            grad = mxGetPr(plhs[1]);
            gw = grad;
            gv_start = &grad[nFeaturesTotal*nStates];
            gv_end = &grad[nFeaturesTotal*nStates + nStates];
            gv = &grad[nFeaturesTotal*nStates + nStates + nStates];
        }
    }
    
    /* Edge Potentials are constant */
    for(all_state2) {
        for(all_state1) {
            edgePot(state1,state2) = exp(v(state1,state2));
        }
    }
        
    for(s = 0;s < nSentences; s++) {
        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];
        }        
        
        /* Update Objective Function based on true labels */
        for(all_n) {
            state = y[sentences[s]-1 + n]-1;
            *func -= nodePot(n,state);
        }
        for(n = 0; n < nNodes-1; n++) {
            state1 = y[sentences[s]-1 + n]-1;
            state2 = y[sentences[s]-1 + n+1]-1;
            *func -= v(state1,state2);
        }
        
        /* Exponentiate Potentials */
        for(all_n) {
            for(all_state)
                nodePot(n,state) = exp(nodePot(n,state));
        }
               
        
        /* Forward-Backward to Compute Marginals */
                
        /* Forward Pass */
        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];
        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;
        }
        
        /* 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 objective function based log-partition */
        logZ = 0;
        for(all_n) {
            logZ += log(kappa[n]);
        }
        *func += logZ;
        
        
        /* Update Gradient */
        if (nlhs > 1) {
            /* Node Features */
            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) {
                            gw[featureInd + nFeaturesTotal*state] += nodeBel(n,state);
                        }
                        state = y[sentences[s]-1 +n]-1;
                        gw[featureInd + nFeaturesTotal*state] -= 1;
                    }
                }
            }
            /* Start/end features */
            for(all_state) {
                gv_start[state] += nodeBel(0,state);
                gv_end[state] += nodeBel(nNodes-1,state);
            }
            state = y[sentences[s]-1]-1;
            gv_start[state] -= 1;
            state = y[sentences[s]-1 + nNodes-1]-1;
            gv_end[state] -= 1;
            /* Edge features */
            for(n = 0; n < nNodes-1; n++)
            {
                for(all_state1) {
                    for(all_state2) {
                        gv[state1 + nStates*state2] += edgeBel(state1,state2,n);
                    }
                }
                state1 = y[sentences[s]-1 +n]-1;
                state2 = y[sentences[s]-1 +n+1]-1;
                gv[state1 + nStates*state2] -= 1;
            }
        }
    }
    mxFree(nodePot);
    mxFree(edgePot);
    mxFree(alpha);
    mxFree(beta);
    mxFree(tmp);
    mxFree(nodeBel);
    mxFree(edgeBel);
    mxFree(kappa);
}