/* ****************************************************************

This is the C version of program 3.4 from page 79 of
"Modeling Infectious Disease in humans and animals"
by Keeling & Rohani.

It is the SEIR model with four different age-groups and 
yearly "movements" between the groups mimicking the school year

This code is written to be simple, transparent and readily compiled.
Far more elegant and efficient code can be written.

This code can be compiled using gnu or intel C compilers:
icc -o Program_3_4  Program_3_4.c  -lm
g++ -o Program_3_4  Program_3_4.c  -lm

If you have any difficulties, we suggest you set DEBUG to 1 or -1

******************************************************************/

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

// Set DEBUG to 1 to output some debugging information
// Set DEBUG to -1 to output more debugging information
#define DEBUG 1

// Set up basic parameters
double beta[4][4]={2.089, 2.089, 2.086, 2.037, 2.089, 9.336, 2.086, 2.037, 2.086, 2.086, 2.086, 2.037, 2.037, 2.037, 2.037, 2.037};
double gamm=1.0/5.0;;
double sigma=1.0/8.0;
double S0[4]={0.05,0.01,0.01,0.008};
double E0[4]={0.0001,0.0001,0.0001,0.0001};
double I0[4]={0.0001,0.0001,0.0001,0.0001};
double MaxTime=100*365;


// Set up fixed parameters
double nu[4]={1.0/(55*365),0,0,0};
double mu[4]={0,0,0,1.0/(55*365)};
double n[4]={6.0/75,4.0/75,10.0/75,55.0/75};


// Set up variables and rates of change
// Note we no-longer explicitly model R.
double t,S[4],E[4],I[4],Pop[12];
double dPop[12];

// Initialise the equations and Runge-Kutta integration
void Diff(double Pop[12])
{
  int i,j;

  // Set up temporary variables to make the equations look neater

  double tmpS[4],tmpE[4],tmpI[4],Inf;

  for(i=0;i<4;i++)
    {
      tmpS[i]=Pop[i];
      tmpE[i]=Pop[i+4];
      tmpI[i]=Pop[i+8];
    }
  
  /* The differential equations */

  for(i=0;i<4;i++)
    {
      Inf=0;
      for(j=0;j<4;j++)
	{
	  Inf+=beta[i][j]*tmpI[j];
	}

      dPop[i] = nu[i]*55/75.0 - Inf*tmpS[i] - mu[i]*tmpS[i];              // dS(i)/dt
      dPop[i+4] = Inf*tmpS[i] - sigma*tmpE[i] - mu[i]*tmpE[i];            // dE(i)/dt
      dPop[i+8] = sigma*tmpE[i] - gamm*tmpI[i] - mu[i]*tmpI[i];  // dI(i)/dt
    }
}

void Runge_Kutta(double);
void Read_in_parameters(FILE *);
void Perform_Checks();
void Output_Data(FILE *);


// The main program

int main(int argc, char** argv)
{
  int i,j;
  char FileName[1000];
  double sum_beta, step, Every, t_next;
  FILE *in, *out;
   
  /* Find the given parameter  file name and open the file 
     otherwise use parameters set at top of program */
  if(argc>1)
    {
      strcpy(FileName, *(argv+1));  
      
      if((in = fopen(FileName,"r"))==NULL)
	{
	  printf("Cannot read file %s\n",FileName);
	  exit(1);
	}

      /* Read in the parameters */
      
      Read_in_parameters(in);

      if(DEBUG) printf("\nReading parameters from file %s\n",FileName);
    }
  else
    {
      if(DEBUG) printf("\nUsing default parameters\n");
    }


  /* Check all parameters are OK & set up intitial conditions */
 
  Perform_Checks();

  /* Find a suitable time-scale for outputs */

  sum_beta=0;
  for(i=0;i<4;i++)
    {
      for(j=0;j<4;j++)
	{
	  sum_beta+=beta[i][j];
	}
      S[i]=S0[i]; E[i]=E0[i]; I[i]=I0[i];
    }

  step=0.01/((sum_beta+gamm+sigma+nu[0]));
  
  Every=pow(10,floor(log10((1.0/(sum_beta+gamm+sigma+nu[0])))));
  while(MaxTime/Every>10000)
    {
      Every*=10.0;
    }
  

  if(DEBUG) printf("Using a time step of %g and outputing data every %g\n\n",step,Every);
    

  if((out = fopen("Output","w"))==NULL)
    {
      printf("Cannot open file Output to write data\n");
      exit(1);
    }

  /* The main iteration routine */

  t=0; 

  Output_Data(out);

  do
    {

      t_next=t+365;

      do
	{
	  
	  Runge_Kutta(step);  t+=step;
	  
	  /* If time has moved on sufficiently, output the current data */
	  if( floor(t/Every) > floor((t-step)/Every) )
	    {
	      Output_Data(out);
	    }
	  else
	    {
	      if(DEBUG<0) Output_Data(stdout);
	    }
	}
      while(t<t_next);

      S[3]=S[3]+S[2]/10;
      S[2]=S[2]+S[1]/4-S[2]/10;
      S[1]=S[1]+S[0]/6-S[1]/4;
      S[0]=S[0]-S[0]/6;

      E[3]=E[3]+E[2]/10;
      E[2]=E[2]+E[1]/4-E[2]/10;
      E[1]=E[1]+E[0]/6-E[1]/4;
      E[0]=E[0]-E[0]/6;

      I[3]=I[3]+I[2]/10;
      I[2]=I[2]+I[1]/4-I[2]/10;
      I[1]=I[1]+I[0]/6-I[1]/4;
      I[0]=I[0]-I[0]/6;

    }
  while(t<MaxTime);
  
  Output_Data(out);

  fclose(out);
}






void Read_in_parameters(FILE *in)
{
  int i,j;
  char str[200];
  
  fscanf(in,"%s",str);
  for(i=0;i<4;i++)
    {
      for(j=0;j<4;j++)
	{
	  fscanf(in,"%lf",&beta[i][j]);
	}
    }
    
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&sigma);
    
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&gamm);
  
  fscanf(in,"%s",str);
  for(i=0;i<4;i++)
    {
      fscanf(in,"%lf",&S0[i]);
    }

  fscanf(in,"%s",str);
  for(i=0;i<4;i++)
    {
      fscanf(in,"%lf",&E0[i]);
    }
  
  fscanf(in,"%s",str);
  for(i=0;i<4;i++)
    {
      fscanf(in,"%lf",&I0[i]);
    }
  
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&MaxTime);
  
  fclose(in);
}

void Perform_Checks()
{
  int i,j;

  for(i=0;i<4;i++)
    {
      for(j=0;j<4;j++)
	{
	  if(beta[i][j]<0) {printf("ERROR: Transmission rate beta(%d,%d) = %g is less than zero\n",i,j,beta[i][j]); exit(1);}
	}
    

      if(I0[i]<0) {printf("ERROR: Initial level of infectious in age-group %d (%g) is less than zero\n",i,I0[i]); exit(1);}

      if(E0[i]<0) {printf("ERROR: Initial level of exposed in age-group %d (%g) is less than zero\n",i,E0[i]); exit(1);}

      if(S0[i]<0) {printf("ERROR: Initial level of susceptibles in age-group %d (%g) is less than zero\n",i,S0[i]); exit(1);}

      if(S0[i]+E0[i]+I0[i]>n[i]) {printf("ERROR: Initial level of susceptibles and infecteds in age-group %d (%g) is greater than calculated group size (%g)\n",i,S0[i]+E0[i]+I0[i],n[i]); exit(1);}
    }
  
  if(gamm<=0) {printf("ERROR: Recovery rate gamma (%g) is less than or equal to zero\n",gamm); exit(1);}

  if(sigma<=0) {printf("ERROR: Exposed to Infectious rate sigma (%g) is less than or equal to zero\n",sigma); exit(1);}

  if(MaxTime<=0) {printf("ERROR: Maximum run time (%g) is less than or equal to zero\n",MaxTime); exit(1);}

}



void Output_Data(FILE *out)
{
  int i;
  
  if(out!=stdout) 
    {      
      fprintf(out,"%g \t",t);
      for(i=0;i<4;i++) { fprintf(out," %g",S[i]);}
      for(i=0;i<4;i++) { fprintf(out," %g",E[i]);}
      for(i=0;i<4;i++) { fprintf(out," %g",I[i]);}
      fprintf(out,"\n");
    }
  if(DEBUG)
    {
      printf("%g \t",t);
      for(i=0;i<4;i++) { printf(" %g",S[i]);}
      for(i=0;i<4;i++) { printf(" %g",E[i]);}
      for(i=0;i<4;i++) { printf(" %g",I[i]);}
      printf("\n");
    }
    
}




void Runge_Kutta(double step)
{
  int i;
  double dPop1[12], dPop2[12], dPop3[12], dPop4[12];
  double tmpPop[12], initialPop[12];

  /* Integrates the equations one step, using Runge-Kutta 4
     Note: we work with arrays rather than variables to make the
     coding easier */

  for(i=0;i<4;i++)
    {
      initialPop[i]=S[i];
      initialPop[i+4]=E[i];
      initialPop[i+8]=I[i];
    }

  Diff(initialPop);
  for(i=0;i<12;i++)
    {
      dPop1[i]=dPop[i];
      tmpPop[i]=initialPop[i]+step*dPop1[i]/2;
    }

  Diff(tmpPop);
  for(i=0;i<12;i++)
    {
      dPop2[i]=dPop[i];
      tmpPop[i]=initialPop[i]+step*dPop2[i]/2;  
    }

  Diff(tmpPop);
  for(i=0;i<12;i++)
    {
      dPop3[i]=dPop[i];
      tmpPop[i]=initialPop[i]+step*dPop3[i]; 
    }

  Diff(tmpPop);

  for(i=0;i<12;i++)
    {
      dPop4[i]=dPop[i];

      tmpPop[i]=initialPop[i]+(dPop1[i]/6 + dPop2[i]/3 + dPop3[i]/3 + dPop4[i]/6)*step;
    }


  for(i=0;i<12;i++)
    {
      if(tmpPop[i]<0)
	{
	  for(i=0;i<12;i++)
	    {
	      printf("%d)  %g %g\n",i,initialPop[i],dPop1[i]);
	    }
	  exit(1);
	}
    }


  for(i=0;i<4;i++)
    {
      S[i]=tmpPop[i];
      E[i]=tmpPop[i+4];
      I[i]=tmpPop[i+8];
    }
  

  return;
}