/*
Function to partition a given number of processors nprocs into
a processor grid nprocs=px*py*pz ..., such that:

	px*weights[0] + py*weights[1] + ...

is minimized.

Use -DTEST to compile the test code.

The function can be called from a Fortran program, e.g.:
        PARAMETER (NDIM=3)
	INTEGER NP(NDIM)
        DOUBLE PRECISION  WEIGHTS(NDIM)
        DATA WEIGHTS/1.,1.,1./
        ... ...
        CALL PARTITION(32,NP,WEIGHTS,NDIM)
        (or IERR=PARTITION(32,NP,WEIGHTS,NDIM) )
        ... ...

By ZLB.
*/

#ifdef VERBOSE
#  define TEST
#endif

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

#define MAXFACTORS	16

int  partition_(int *n, int *np, double *weights, int *ndim);

static int compare(double sum, int *npwork, double best_sum, int *np);
/* the functions with a "stack" argument are recursive */
static void next_factor(int facno, int prod, int dimen);
static void next_dimension(int dimen);

#ifdef TEST
void main(void)
{
  int i, n, dim, *np;
  double *weights;

  fprintf(stderr,"Enter dimension: "); scanf("%d",&dim);
  np=malloc(dim*sizeof(*np));
  weights=malloc(dim*sizeof(*weights));
  if (np==NULL || weights==NULL) {
    fprintf(stderr,"Memory allocation error!\n");
    exit(1);
  }
  	for (n=0; n<dim; n++) weights[n]=1;
  	weights[0]=1; weights[1]=2; weights[2]=3;
  fprintf(stderr,"Enter weights: ");
  for (n=0; n<dim; n++) scanf("%lf",weights+n);

  while (1) {
    fprintf(stderr,"Enter N (0 to quit): ");
    scanf("%d",&n);
    if (!n) break;
    if (i=partition_(&n, np, weights, &dim)) {
      fprintf(stderr,"Error code %d\n",i);
    } else {
      fprintf(stderr,"Best choice=");
      for (i=0; i<dim; i++) fprintf(stderr,"%d ",np[i]);
      fprintf(stderr,"\n");
    }
  }
  exit(0);
}
#endif

static int n_save, ndim, *np, *npwork;
static double *weights, best_sum;
static int numfac;
static int factors[MAXFACTORS], powers[MAXFACTORS];

#ifdef TEST
  static int counter;
#endif

int partition_(int *n0, int *np0, double *weights0, int *ndim0)
/* find np[0], np[1],..., np[*ndim] such that:
	np[0]*np[1]*...*np[*ndim]=*n0
	(np[0]-1)*weights[0] +
	(np[1]-1)*weights[1] +
	... +
	(np[*ndim]-1)*weights[*ndim] = minimum,
  returns !0 on errors */
{
  int n, id, iq, p;

  n=n_save=abs(*n0);
  ndim=*ndim0;
  np=np0;
  weights=weights0;

  if (n<=1) {
    for (id=0; id<ndim; id++) np[id]=n;
    return 0;
  }

  np[0]=n; best_sum=(np[0]-1)*weights[0];
  for (id=1; id<ndim; id++) np[id]=1;

  /* First, factorize n into prime factors */
  p=numfac=0;
  id=2;
  while (id<=n) {
    iq=n/id;
    if (n == iq*id) {
      if (!p) p=1; else p++;
      n=iq;
    } else {
      if (p) {
        if (numfac==MAXFACTORS) return 2;
        factors[numfac]=id;
        powers[numfac++]=p;
        p=0;
      }
      id= id==2 ? 3:id+2;
    }
  }
  if (p) {
    if (numfac==MAXFACTORS) return 2;
    factors[numfac]=id;
    powers[numfac++]=p;
  }
#ifdef TEST
    fprintf(stderr,"Number of factors = %d: ",numfac);
    for (id=0; id<numfac; id++)
      fprintf(stderr,"%d^%d ",factors[id],powers[id]);
    if (numfac) fprintf(stderr,"\n");
#endif

  /* Loop on all possible combinations to find the minimum */

#ifdef TEST
    counter=0;
#endif

  if ((npwork=malloc(ndim*sizeof(*npwork)))==NULL) return 3;
  next_dimension(0);
  free(npwork);


#ifdef TEST
    fprintf(stderr,"Number of cases searched=%d\n",counter);
#endif

  return 0;
}

static void next_factor(int facno, int prod, int dimen)
{
  int i,p,f;

  if (facno>=numfac) {
    npwork[dimen++]=prod;
    next_dimension(dimen);
    return;
  }

  p=powers[facno]; f=factors[facno];
  for (i=0; i<=p; i++, powers[facno]--) {
    next_factor(facno+1,prod,dimen);
    prod*=f;
  }
  powers[facno]=p;

  return;
}

static void next_dimension(int dimen)
{
  static int i, prod;
  static double sum;

  if (dimen>=ndim-1) {
    prod=1;
    for (i=0; i<dimen; i++) prod*=npwork[i];
    npwork[dimen]=n_save/prod;

    /* evaluate this case */
    sum=0.0;
    for (i=0; i<ndim; i++) sum+=(npwork[i]-1)*weights[i];
    if (compare(sum, npwork, best_sum, np)) {
      best_sum=sum;
      memcpy(np, npwork, ndim*sizeof(*np));
    }
#ifdef TEST
      counter++;
#endif
#ifdef VERBOSE
      for (i=0; i<ndim; i++) fprintf(stderr,"%-8d ", npwork[i]);
      printf("\tsum=%-08.2lf\n",sum);
#endif
  } else {
    next_factor(0, 1, dimen);
  }

  return;
}

static int compare(double sum, int *npwork, double best_sum, int *np)
/* returns true if (sum,npwork) < (best_sum,np) */
{
  int i;

  if (sum<best_sum) return 1;
  if (sum>best_sum) return 0;
  for (i=0; i<ndim; i++) {
    if (npwork[i]<np[i]) return 1;
    if (npwork[i]>np[i]) return 0;
  }
  return 0;
}