#include <stdio.h>
#include <math.h>
#include "usempi.h"

#define MS_PRINT 0 /* flag to print out debugging information to file */
                   /* Each process gets a unique file */

#ifdef _CRAY  /* this is for the mallopt() command below */
#include <malloc.h>
#endif

#define OFFSET 2 /* offset for tags number for normal tasks */
                 /* 0 quit 
                    1 run mapcouplings */

/* 
   This subroutine initializes the data structures and
   communication groups used by mpi
*/

void init_mpi_types(void){
  int i;
#define SPECIFY_SIZE 11
  int len[SPECIFY_SIZE];
  MPI_Datatype type[SPECIFY_SIZE];
  specify_calculation mx;
  MPI_Aint disp[SPECIFY_SIZE],base;
#if 0
  MPI_Group g1,g2;
  int master[1]={0};
#endif

  /* make MPI datatype for the structure specify_calculation */

  MPI_Address(&mx.type,disp);          len[0]=1;        type[0]=MPI_INT;
  MPI_Address(&mx.ht,disp+1);          len[1]=HINDEX;   type[1]=MPI_INT;
  MPI_Address(&mx.nptrunc,disp+2);     len[2]=1;        type[2]=MPI_INT;
  MPI_Address(&mx.kt,disp+3);          len[3]=1;        type[3]=MPI_UNSIGNED;
  MPI_Address(&mx.charge,disp+4);      len[4]=1;        type[4]=MPI_INT;
  MPI_Address(&mx.o,disp+5);           len[5]=1;        type[5]=MPI_INT;
  MPI_Address(&mx.multi,disp+6);       len[6]=1;        type[6]=MPI_INT;
  MPI_Address(&mx.momflag,disp+7);     len[7]=1;        type[7]=MPI_INTEGER;
  MPI_Address(&mx.angle,disp+8);       len[8]=HINDEX+2; type[8]=MPI_DOUBLE;
  MPI_Address(&mx.nval,disp+9);        len[9]=1;        type[9]=MPI_INTEGER;
  MPI_Address(&mx.vectorflag,disp+10); len[10]=1;       type[10]=MPI_INTEGER;

  /* we find offsets relative to the "base" of the structure,
     wherever that may be (thus we make no assumptions of order. */
  MPI_Address(&mx,&base);
  for(i=0; i<SPECIFY_SIZE; i++)disp[i]-=base;
  MPI_Type_struct(SPECIFY_SIZE,len,disp,type,&mpi_specify_calculation);
  MPI_Type_commit(&mpi_specify_calculation);

  /* check that we did it right */

  MPI_Type_size(mpi_specify_calculation,&i);
  if(sizeof(mx)!=i){
    fprintf(stderr,__FILE__":  construction of mpi_specify_calculation "
	    "failed\n");
  }
  
#if GROUP_SIZE>0

/*
  Initialize the communicator mpiallslaves which
  is a subgroup consisting of just the slaves.
  This is used in file "improved.c"  
*/

  MPI_Comm_group(MPI_COMM_WORLD,&g1);
  MPI_Group_excl(g1,1,master,&g2);
  MPI_Comm_create(MPI_COMM_WORLD,g2,&mpiallslaves);
  MPI_Group_free(&g1);
  MPI_Group_free(&g2);

#endif

}

/* This is the slave process used by MPI 
   it is a "stand-alone" subroutine to calculate the
   lowest eigenvalues for a given basis.  It also
   calculates the basis when necessary. 
   angle contains HINDEX components of transverse momentum,
   the kmax cutoff for nonzero L and L itself.
*/

#define MAXBASIS 150

void slave(void){
  specify_calculation a;
  full_basis basis;
  eigensystem eigen=INITIAL_EIGENSYSTEM;
  int myrank;
  element couplings[NPARAMS];
  size_mem global_reduced_memory;
  size_mem global_els_memory;
#if MS_PRINT 
  FILE *fp;
  char out[10];
  int i,j;
#endif  
  MPI_Status status;

  MPI_Comm_rank(MPI_COMM_WORLD,&myrank);

#if MS_PRINT /* open debugging file with name based on rank */
  sprintf(out,"p%i.tmp",myrank);
  if((fp = fopen(out,"w")) == NULL){
     fprintf(stderr,__FILE__":  Can't open file %s in usempi.c, exiting\n",out);
    goto error;
  }
#endif

  /* on CRAY, be more carefull about memory management.  */

#ifdef _CRAY
      mallopt(M_LOWFIT,1);
#endif

      /* Main loop of program */

  for(;;){

    MPI_Probe(MPI_ANY_SOURCE,MPI_ANY_TAG,MPI_COMM_WORLD,&status);

    /* Stop program */
    
    if(status.MPI_TAG==0)goto exit;

    /* Receive coupling constants. */

    else if(status.MPI_TAG==1)
      MPI_Recv(couplings,NPARAMS,MPI_DOUBLE,0,MPI_ANY_TAG,MPI_COMM_WORLD,&status);
   
    /* Calculate spectrum */

    else {    
#if MS_PRINT /* diagnostic print to debugging file*/
      fprintf(fp,"About to receive Job %i\n",status.MPI_TAG-OFFSET);
      fflush(fp);
#endif
      MPI_Recv(&a,1,mpi_specify_calculation,0,MPI_ANY_TAG,MPI_COMM_WORLD,&status);
#if MS_PRINT /* diagnostic print to debugging file */
      fprintf(fp,"Received job %i, nvals=%i\n",
	      status.MPI_TAG-OFFSET,nvals);
	fprintf(fp,"  ht=(");
	for(i=0; i<HINDEX; i++)fprintf(fp," %i",a.ht[i]);
	fprintf(fp,"), np=%i kt=%i charge=%i o=%i and multi=%i,\n",
		a.npmax,a.kt,a.charge,a.o,a.multi);
      fprintf(fp,"  angles=");
      for(j=0; j<HINDEX+2; j++)fprintf(fp," %g",angle[j]);
      fprintf(fp,"\n");
#endif

/*
    construct one or two bases.  The order for makebasis() 
    and rebasis() is important for memory fragmentation 
    problems on the CRAY. 
    mapcouplings should have already been run.

    Now calculate the actual spectrum 
*/

      spectrum(&a,couplings,&eigen,&basis);

      if(eigen.values==NULL){
	 fprintf(stderr,__FILE__":  slave() %i detecting an error in spectrum()\n",
		myrank);
	goto error;
      }

#if MS_PRINT /* diagnositic print */
      fprintf(fp,"  values: ");
      for(i=0; i<nvals; i++)fprintf(fp,"%.10g ",val[i]);
      fprintf(fp,"\n");
#endif

      MPI_Send(eigen.values,eigen.values_length,MPI_DOUBLE,0,status.MPI_TAG,MPI_COMM_WORLD);

      if(a.vectorflag){
	if(eigen.vectors.r != NULL)
	  MPI_Send(eigen.vectors.r,eigen.vectors.length,
		   MPI_DOUBLE,0,status.MPI_TAG,MPI_COMM_WORLD);
	else if(eigen.vectors.c != NULL)
	  MPI_Send(eigen.vectors.c,2*eigen.vectors.length,
		   MPI_DOUBLE,0,status.MPI_TAG,MPI_COMM_WORLD);
      }
/*
       To prevent memory fragmentation problems on systems
       without virtual memory paging (like the CRAY), we
       deallocate the memory used by spectrum() in reverse order.

       The other two things that can be done  to minimize
       memory fragmentation is to choose
       sufficiently large values for MAX_ELS and MAX_BASE.
*/

#if defined(_CRAY)
      freespectrum();
#endif
    }
  
    /* for normal case, just do next loop iteration */
#if MS_PRINT /* if everything is OK, write over old info */
    rewind(fp);
#endif
  }
    
error:  /* On error, print out memory usage and tell the master. */
/* The master will then send out a command that everyone quit. */
   fprintf(stderr,__FILE__":  Error in slave %i after receiving job %i, exiting.\n",
	  myrank,status.MPI_TAG);
#if MS_PRINT
  printhels(fp);
  printelements(fp);
  printlaso(fp);
#endif
#ifdef _CRAY
#if 0 /* Choose one of the following */
  printf("MM Running malloc_stats(0), slave %i:\n",myrank);
  malloc_stats(0);
#elif 0
  printf("MM Running malloc_stats(1), slave %i:\n",myrank);
  malloc_stats(1);
#elif 1  /* the following is very long */
  printf("MM Running malloc_stats(2), slave %i:\n",myrank);
  malloc_stats(2);
#endif
#endif
  /* Just the message tag is read by the master */
  MPI_Send(&myrank,1,MPI_INT,0,0,MPI_COMM_WORLD);
  
exit:  /* Normal exit */
  free_eigensystem(&eigen);
  free_basis(&basis);
  MPI_Comm_rank(mpiallslaves,&myrank);
  MPI_Reduce(&(max_els.memory),&global_els_memory,1,MPI_SIZE_T,
	     MPI_MAX,0,mpiallslaves);
  MPI_Reduce(&reduced_basis_memory,&global_reduced_memory,1,MPI_SIZE_T,
	     MPI_MAX,0,mpiallslaves);
  if(myrank==0){
    printf("Global maximum memory for Hamiltonian matrix is " SIZE_MEM_PRINT 
	   " bytes\n",global_els_memory);
    printf("Global maximum memory used to construct bases: " SIZE_MEM_PRINT 
	   " bytes\n",global_reduced_memory);
  }
#if MS_PRINT /* debugging file with name based on rank */
  fclose(fp);
#endif
  return;
}

/*  This subroutine gives tasks to the various slaves. 
    Eigenvalues are output in the vector vals whose
    memory should be already allocated.  vals[whichvals[i]+j]
    contains the jth eigenvalue, j=0,1,2... for job i,
    i=1,2,...njob. 

    Since all the jobs are assumed to have the same couplings,
    the generation of improved matrix elements is parallelized.

    master returns 0 on successful completion.
*/


int master(specify_calculation *base, element *couplings, 
	   element *vals, int *whichvals, int *nvals, int njob){
  int job,receive,slave,nslaves,j;
#if MS_PRINT /* diagnostic print to debugging file */
  FILE *fp;
  int i,k;
  char out[10]="p0.tmp";
#endif  
  MPI_Status status,status2;

  MPI_Comm_size(MPI_COMM_WORLD,&nslaves);
  nslaves--; /* don't include master amongst slaves */
  

#if MS_PRINT /* open debugging file*/
  if((fp = fopen(out,"w")) == NULL){
     fprintf(stderr,__FILE__":  Can't open file %s in usempi.c, exiting\n",out);
    MPI_Type_free(&xt2);
    return 1;
  }
#endif

  for(j=0; j<nslaves; j++)
    MPI_Send(couplings,NPARAMS,MPI_DOUBLE,j+1,1,MPI_COMM_WORLD);
  
  /* loop through all jobs that need completion */
  for(job=0, receive=0; receive<njob;){

    /* Receive completed jobs once all slaves have been seeded. */
    if(job<nslaves)
      slave=job+1;  /* slaves are numbered 1,2,3,... */
    else{
      MPI_Probe(MPI_ANY_SOURCE,MPI_ANY_TAG,MPI_COMM_WORLD,&status);
      j=status.MPI_TAG-OFFSET; /* remove offset in tag number */
      if(status.MPI_TAG<OFFSET||j>=njob){
	 fprintf(stderr,__FILE__":  Master receiving tag=%i from slave %i, exiting\n",
		status.MPI_TAG,status.MPI_SOURCE);
#if MS_PRINT /* debugging file with name based on rank */
	fclose(fp);
#endif
	return 1;
      }
      MPI_Recv(vals+whichvals[j],nvals[j],
	       MPI_DOUBLE,MPI_ANY_SOURCE,
	       MPI_ANY_TAG,MPI_COMM_WORLD,&status2);
      if(base[j].vectorflag){
	fprintf(stderr,__FILE__":  master(...) not ready to receive vectors!\n");
	return 1;
      }

#if 0 /* debugging print: print jobs received */
      if(status.MPI_TAG<10){
	printf("Master receiving job %i from slave %i offset %i;\n ",
	       status.MPI_TAG,status.MPI_SOURCE,whichvals[status.MPI_TAG]);
	for(i=0; i<nvals[j]; i++)
	  printf(" %f",vals[whichvals[j]+i]);
	printf("\n");
      }
#endif
      slave=status.MPI_SOURCE;
      receive++;
    }

    /* send jobs to the slaves */
    if(job<njob){
#if MS_PRINT /* debugging print to file: print jobs sent to slaves. */
      fprintf(fp,"Master sending job=%i to slave %i, nvals=%i\n",
	      job,slave,nvals[job]);
      for(k=0; k<2; k++){
	fprintf(fp,"  ht=(");
	for(i=0; i<HINDEX; i++)fprintf(fp," %i",base[job][k]->ht[i]);
	fprintf(fp,"), np=%i kt=%i cyclic=%i o=%i and multi=%i,\n",
		base[job][k]->npmax,base[job][k]->kt,base[job][k]->cyclic,
		base[job][k]->o,base[job][k]->multi);
      }
      fprintf(fp,"  angle=");
      for(i=0; i<HINDEX+2; i++)fprintf(fp," %f",angle[job][i]);
      fprintf(fp," nvals=%i\n",nvals[job]);
      fflush(fp);
#endif
      /* include offset of one in tag number */
      MPI_Send(base+job,1,mpi_specify_calculation,slave,job+OFFSET,MPI_COMM_WORLD);
      job++;
    }
  }

  /* normal return, no error */
#if MS_PRINT /* open debugging file with name based on rank */
  fclose(fp);
#endif
  return 0;  
}