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

/* NOTE! Here we explicitly use the fact that matlab allocates all of these with the initial value 0 */

#define pow(a,b) (exp((b)*log(a)))

void vextract(TENSOR2i A,TENSOR2i I,TENSOR2i B)
{
  /* matlab's A(I) = B */
  int i;
  for(i=0; i<d2(I); i++)
    _(A,0,_(I,0,i)) = _(B,0,i);
}

void complement(TENSOR2i *J,TENSOR2i I,int n)
{
  /* matlab's J = setdiff(1:n,I) */
  /* assuming I is given in incr order */
  int i=0;
  int j=0;
  int k;
  setdim2((*J),1,n-d2(I));
  for(k=0; k<n; k++) {
    if(_(I,0,i)==k)
      i++;
    else {
      _((*J),0,j) = k;
      j++;
    }
  }
}



void mexFunction( int nlhs, mxArray *plhs[],
                  int nrhs, const mxArray *prhs[] )
{
  int i,m,n,mn,n1,mn1;
  TENSOR2 P;
  TENSOR2i X,ii,jj,dim,dim1,u,z;
  double p,*pptr;

/* ---------------- I/O bureaucracy -------------- */

  /* p = getPX(P,m,n,X,ii)
     p = P(X(ii) = x)   */

  makeT2ad(P,IN(1));
  m  = (int) *mxGetPr(IN(2));
  n  = (int) *mxGetPr(IN(3));

  allocTENSOR2i(&X,1,n);
  allocTENSOR2i(&ii,1,n);
  allocTENSOR2i(&jj,1,n);
  allocTENSOR2i(&dim,1,n);
  allocTENSOR2i(&u,1,n);

  for(i=0;i<n;i++)
    _(dim,0,i) = m;

  mn = (int) (.5+pow(1.0*m,1.0*n));
  /* convert X to int */
  copyarr2tensi(IN(4),&X);
  /* convert ii to int */
  copyarr2tensi(IN(5),&ii);

  /* adjust X,ii */
  for(i=0;i<d2(X);i++) {
    _(X,0,i)--;
    _(ii,0,i)--;
  }

  n1 = n - d2(ii);
  mn1 = (int) (.5+pow(1.0*m,1.0*n1));
  complement(&jj,ii,n);
  allocTENSOR2i(&dim1,1,n1);
  allocTENSOR2i(&z,1,n1);
  for(i=0;i<n1;i++)
    _(dim1,0,i) = m;

  int zi,ui;
  p = 0.0;
  for(zi=0;zi<mn1;zi++) {
    /* z is the "complement" of X */
    /*printf("zi = %d \n",zi+1);*/
    ind2coord(z,zi,dim1);
    /*printTENSOR2i(z);*/
    vextract(u,ii,X);  /* u(ii) = X */
    vextract(u,jj,z);  /* u(jj) = z */
    /*printTENSOR2i(u);*/
    ui = coord2ind(u,dim);
    p += _(P,ui,0);
    /*printf("zi = %d ui = %d p = %1.6f\n",zi+1,ui+1,_(P,ui,0));/**/
  }
  OU(1) = mxCreateDoubleMatrix(1,1,mxREAL);
  pptr = mxGetPr(OU(1));
  *pptr = p;

  freeTENSOR(X);
  freeTENSOR(ii);
  freeTENSOR(jj);
  freeTENSOR(dim);
  freeTENSOR(u);
  freeTENSOR(dim1);
  freeTENSOR(z);
}