#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

/*
 Implementation of various combinatorial enumeration algorithms.
*/


/* Returns the key associated with the permutation in P.  P is
   destroyed in the process. */
int perm_to_key(int* P,
		const int n)
{
  int i, K=0, m=1, j;

  /* Convert to the compressed representation */
  for (i=n-1; i>=0; --i)
    {
      for (j=i-1; j>=0; --j) if (P[j] > P[i]) --P[j]; 
    }

  /* Get the key */
  for (j=1; j<n; ++j)
    {
      K += P[j]*m;
      m *= j+1;
    }

  return K;
}


/* Converts the integer key K into a permutation of length n.  The
 result is stored in the first n positions of P, but P should be
 allocated to 2*n to provide workspace.  The value of K is altered by
 this function. */
void key_to_perm(int K,
		 int* P,
		 const int n)
{
  int i, j, d;

  /* Get the compressed representation. */
  P[0] = 0;
  for (i=1; i<n; ++i)
    {
      K /= i;
      P[i] = K % (i+1);
      K -= P[i];
    }
  
  /* Convert to the standard representation. */
  for (j=0; j<n; ++j) P[j+n] = j;
  for (j=n-1; j>=0; --j) 
    {
      d = P[n + P[j]];
      for (i=P[j]; i<n; ++i) P[n+i] = P[n+i+1];
      P[j] = d;
    }
}


/* Check that key_to_perm and perm_to_key are inverses. */
void test_perm_key(const int n)
{
  int* P = (int*)malloc(n*sizeof(int));
  int K=0, q, k;

  while (log(K) < lgamma(n+1))
    {
      printf("Key:%5d\t", K);
      key_to_perm(K, P, n);
      printf("Permuatation: (");
      for (k=0; k<n; ++k) printf((k<n-1) ? "%3d," : "%3d", P[k]);
      q = perm_to_key(P, n);
      printf(")\tKey:%5d\n", q);
      ++K;
    }

  free(P);
}


/* List all permutations on a set of size n using recursion. */
int* rperm(const int n)
{
  int i, j, k, J, K, p, q, ii;
  int *R;
  int *V;

  /* Exit the recursion. */
  if (n == 1)
    {
      V = (int*)malloc(sizeof(int));
      V[0] = 0;
      return V;
    }

  /* Recursion. */
  V = rperm(n-1);

  /* The number of permutations in V is (n-1)!. */
  J = (int)rint(exp(lgamma(n)));

  /* The number of permutations in the result is n! = (n-1)! * n. */
  K = J*n;

  /* Storage for n! permutations. */
  R = (int*)malloc(n*K*sizeof(int));

  /* Take each permutation on the smaller set, and insert "n-1" at
     every possible point. */
  p=0, q=0;
  for (j=0; j<J; ++j)
    {
      for (i=0; i<n; ++i)
	{
	  ii=0;
	  for (k=0; k<n; ++k) 
	    {
	      if (k != i) R[q+k] = V[p+ii++];
	    } 
	  R[q+i] = n-1;
	  q += n;
	}
      p += n-1;
    }

  free(V);
  return R;
}


/* Test the recursive routine for listing permutations. */
void test_rperm(const int n)
{
  int nf, ii=0, i, j;

  int* P = rperm(n);
  nf = (int)rint(exp(lgamma(n+1)));

  for (i=0; i<nf; ++i)
    {
      for (j=0; j<n; ++j) printf("%3d", P[ii++]);
      printf("\n");
    }
}


/* Return the partition n = A_1 + ... + A_k where A_i <= B[i] in a
 list of all such partitions.  When the list is exhausted, it starts
 again from the beginning.

 Warning: there is no verification that the set of partitions is
 non-empty.  If sum_i B[i] < n then this function will never
 terminate. */
void next_partition(const int n,
		    const int k,
		    const int* const B,
		    int* const A)
{
  int j, P, i;

  while (1)
    {
      /* Cycle from right to left until we find a component that we
	 can increment. */
      for (j=k-2; j>=0; --j) 
	{
	  P = 0;
	  for (i=0; i<j; ++i) P += A[i];
	  
	  if ( (P+A[j] < n) && (A[j] < B[j]) ) 
	    {
	      ++A[j];
	      break;
	    }
	  else A[j] = 0;
	}
      
      /* Set the final component so that the sum is n. */
      P = 0;
      for (i=0; i<k-1; ++i) P += A[i];
      A[k-1] = n-P;

      /* Check if the final component satisfies the bound. */
      if (A[k-1] <= B[k-1]) return;
    }
}


/* Test the next_partition() routine on a simple example. */
void test_next_partition()
{
  int n=5, k=4, j;
  int *B;

  B = (int*)malloc(4*sizeof(int));
  B[0]=5, B[1]=5, B[2]=5, B[3]=3;
  
  /* Initialize the partition. */
  int* A = (int*)malloc(4*sizeof(int));
  A[0]=5, A[1]=0, A[2]=0, A[3]=0;

  while (1)
    {
      for (j=0; j<k; ++j) printf("%3d", A[j]);
      printf("\n");

      next_partition(n, k, B, A);

      if ( (A[0] == 5) && (A[1] == 0) && (A[2] == 0) && (A[3] == 0) ) break;
    }
}


/* Use recursion to generate all partitions of n as the sum of k
 non-negative values, with the j^th summand not exceeding B[j].  The
 partitions are returned in a p x k array. */
int* list_partitions(const int n,
		     const int k,
		     const int* const B,
		     int* p)
{
  int M, i, j, m;
  int *Q, *R, **P, *q;

  /* No recursion in this case. */
  if (k == 1)
    {
      if (n > B[0]) { p = 0; return NULL; }
      else
	{
	  Q = (int*)malloc(sizeof(int));
	  Q[0] = n;
	  *p = 1;
	  return Q;
	}
    }

  /* The number of distinct values for the first position. */
  M = (B[0] < n) ? B[0] : n;

  /* P[j] will hold the list for positions 1..k-1 when position 0 has
     value j. */
  P = (int**)malloc((M+1)*sizeof(int*));

  /* q[j] holds the number of rows in P[j]. */
  q = (int*)malloc((M+1)*sizeof(int));

  /* Keep track of the total number of rows. */
  *p = 0;

  /* Loop over the possible values of the first component and recur to
     get the rest. */
  for (j=0; j<=M; ++j)
    {
      m = n-j;
      P[j] = list_partitions(m, k-1, B+1, &q[j]);
      *p += q[j];
    }

  /* Merge P into a single array. */
  Q = (int*)malloc((*p)*k*sizeof(int));
  R = Q;
  for (j=0; j<=M; ++j)
    {
      for (i=0; i<q[j]; ++i)
	{
	  R[0] = j;
	  if (P[j] != NULL) memcpy(R+1, P[j] + i*(k-1), (k-1)*sizeof(int));
	  R += k;
	}
    }

  /* Clean up. */
  for (j=0; j<=M; ++j) free(P[j]);

  free(P);
  free(q);

  return Q;
}



/* A utility function that returns the sum of a double precision
   vector of length n. */
int sum(const int* V,
	const int n)
{
  int S=V[0], j;

  for (j=1; j<n; ++j) S += V[j];
  return S;
}


/* Return an array containing all mxn contingency tables with row and
 column marginals given in R and C.  The number of tables is given in
 the reference argument p.  The tables are packed row-wise into the
 return array. */
int* enumerate_tables(const int m,
		      const int n,
		      const int* R,
		      const int* C,
		      int* p)
{
  int i, ii, q, j, s;
  int *W, *S, *f, **T, *P, *Q;

  /* If there is just 1 row there are either 0 or 1 solutions. */
  if (m == 1)
    {
      s = sum(C, n);
      if (s == R[0])
	{
	  P = (int*)malloc(n*sizeof(int));
	  memcpy(P, C, n*sizeof(int));
	  *p = 1;
	  return P;
	}
      else { *p = 0; return NULL; }
    }

  /* Otherwise do recursion. */
  
  /* Workspace. */
  W = (int*)malloc(n*sizeof(int));
  
  /* Get all possible values for the first row. */
  P = list_partitions(R[0], n, C, &q);
  
  /* Consider each possible value for the first row. */
  T = (int**)malloc(q*sizeof(int*));
  f = (int*)malloc(q*sizeof(int));
  *p = 0;
  for (i=0; i<q; ++i)
    {
      /* Construct the new column totals. */
      for (j=0; j<n; ++j) W[j] = C[j] - P[i*n + j];

      /* Get all possible values for the rest of the table. */
      T[i] = enumerate_tables(m-1, n, R+1, W, &f[i]);
      *p += f[i];
    }

  /* Merge everything into a single array. */
  Q = (int*)malloc(m*n*(*p)*sizeof(int));
  S = Q;
  for (i=0; i<q; ++i)
    {
      for (ii=0; ii<f[i]; ++ii)
	{
	  memcpy(S, P + i*n, n * sizeof(int));
	  S += n;
	  memcpy(S, T[i] + (m-1)*n*ii, (m-1) * n * sizeof(int));
	  S += (m - 1) * n;
	}
    }
  
  /* Clean up. */
  for (i=0; i<q; ++i) free(T[i]);

  free(T);
  free(P);
  free(f);
  free(W);
  
  return Q;
}



/* Test the enumerate_tables(...) using a simple example. */
void test_enumerate_tables()
{
  int m=3, n=2, p, ii, i, i1, i2;
  int *R, *C, *T;

  R = (int*)malloc(m*sizeof(int));
  C = (int*)malloc(n*sizeof(int));

  R[0]=2, R[1] = 3, R[2] = 2;
  C[0]=3, C[1] = 4;

  T = enumerate_tables(m, n, R, C, &p);

  ii = 0;
  for (i=0; i<p; ++i)
    {
      for (i1=0; i1<m; ++i1)
	{
	  for (i2=0; i2<n; ++i2) 
	    {
	      printf("%4d", T[ii]);
	      ++ii;
	    }
	  printf("\n");
	}
      printf("\n");
    }
}



int main(int argc, char** argv)
{
  printf("Generating permutations by key:\n");
  test_perm_key(5);
  printf("\n\n\n");

  printf("Generating permutations recursively:\n");
  test_rperm(5);
  printf("\n\n\n");

  printf("Generating partitions:\n");
  test_next_partition();
  printf("\n\n\n");

  printf("Generating contingency tables:\n");
  test_enumerate_tables();

  return 0;
}