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


void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
	/* Variables */
	int n,n1,n2,e,nNodes,maxState,nEdges,*edgeEnds,s,*y,*swapPositions,sizeNewEdgePot[3];
	double *nodePot,*edgePot,*newNodePot,*newEdgePot;
	
	/* Inputs */
	nodePot = mxGetPr(prhs[0]);
	edgePot = mxGetPr(prhs[1]);
	edgeEnds = (int*)mxGetPr(prhs[2]);
	s = (int)mxGetScalar(prhs[3]);
	y = (int*)mxGetPr(prhs[4]);
	swapPositions = (int*)mxGetPr(prhs[5]);
	
	nNodes = mxGetDimensions(prhs[0])[0];
	maxState = mxGetDimensions(prhs[0])[1];
	nEdges = mxGetDimensions(prhs[2])[0];

	/* Outpus */
	sizeNewEdgePot[0] = 2;
	sizeNewEdgePot[1] = 2;
	sizeNewEdgePot[2] = nEdges;
	plhs[0] = mxCreateDoubleMatrix(nNodes,2,mxREAL);
	plhs[1] = mxCreateNumericArray(3,sizeNewEdgePot,mxDOUBLE_CLASS,mxREAL);
	newNodePot = mxGetPr(plhs[0]);
	newEdgePot = mxGetPr(plhs[1]);
	
	for(n=0;n<nNodes;n++) {
		newNodePot[n] = nodePot[n+nNodes*s];
		newNodePot[n+nNodes] = nodePot[n+nNodes*y[swapPositions[n]]];
	}
	
	for(e=0;e<nEdges;e++) {
		n1 = edgeEnds[e];
		n2 = edgeEnds[e+nEdges];
		newEdgePot[0+2*(0+2*e)] = edgePot[s+maxState*(s+maxState*e)];
		newEdgePot[0+2*(1+2*e)] = edgePot[s+maxState*(y[swapPositions[n2]]+maxState*e)];
		newEdgePot[1+2*(0+2*e)] = edgePot[y[swapPositions[n1]]+maxState*(s+maxState*e)];
		newEdgePot[1+2*(1+2*e)] = edgePot[y[swapPositions[n1]]+maxState*(y[swapPositions[n2]]+maxState*e)];
	}
}
