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

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

It is the SIR model with a probability of mortality and
unequal births and deaths. This code assumes Density-
Dependent Transmission

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_2_3  Program_2_3.c  -lm
g++ -o Program_2_3  Program_2_3.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 0

// Set up basic parameters
double beta=520.0/365.0;
double gamm=1.0/7.0;
double nu=1.0/(70.0*365.0);
double mu=1.0/(70.0*365.0);
double rho=0.5;
double X0=0.2;
double Y0=1e-6;
double N0=1;

double MaxTime=1e5;

// Set up variables and rates of change
double t,X,Y,Z,Pop[3];
double dPop[3];

// Initialise the equations and Runge-Kutta integration
void Diff(double Pop[3])
{
  // Set up temporary variables to make the equations look neater

  double tmpX, tmpY, tmpZ;

  tmpX=Pop[0]; tmpY=Pop[1]; tmpZ=Pop[2];

  /* The differential equations */

  dPop[0] = nu - beta*tmpX*tmpY - mu*tmpX;                      // dX/dt
  dPop[1] = beta*tmpX*tmpY - ((gamm+mu)/(1.0-rho))*tmpY;        // dY/dt
  dPop[2] = gamm*tmpY - mu*tmpZ;                                // dZ/dt

  return;
}

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)
{
  char FileName[1000],str[200];
  double step, Every;
  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");
    }

  if(DEBUG) printf("beta=%g\ngamma=%g\nnu=%g\nmu=%g\nrho=%g\nInitial X=%g\nInitial Y=%g\nInitial N=%g\nRuns until time %g\n\n",beta,gamm,nu,mu,rho,X0,Y0,N0,MaxTime);

  /* Check all parameters are OK & set up intitial conditions */

  Perform_Checks();

  X=X0; Y=Y0; Z=N0-X-Y;

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

  step=(N0)*0.01/((beta+(gamm+mu)/(1.0-rho))*X0);
  
  Every=pow(10,floor(log10((N0/((beta+(gamm+mu)/(1.0-rho))*X0)))));
  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
    {
      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<MaxTime);

  Output_Data(out);

  fclose(out);
}




void Read_in_parameters(FILE *in)
{
  char str[200];

  fscanf(in,"%s",str);
  fscanf(in,"%lf",&beta);
  
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&gamm);
  
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&nu);
  
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&mu);
  
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&rho);
  
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&X0);
  
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&Y0);

  fscanf(in,"%s",str);
  fscanf(in,"%lf",&N0);
  
  fscanf(in,"%s",str);
  fscanf(in,"%lf",&MaxTime);
  
  fclose(in);
}

void Perform_Checks()
{
  if(X0<=0) {printf("ERROR: Initial level of susceptibles (%g) is less than or equal to zero\n",X0); exit(1);}

  if(Y0<=0) {printf("ERROR: Initial level of infecteds (%g) is less than or equal to zero\n",Y0); exit(1);}

 if(N0<=0) {printf("ERROR: Initial population level (%g) is less than or equal to zero\n",N0); exit(1);}

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

  if(gamm<=0) {printf("ERROR: Recovery rate gamma (%g) is less than or equal to zero\n",gamm); exit(1);}
  
  if(nu<0) {printf("ERROR: Birth rate mu (%g) is less than zero\n",mu); exit(1);}
  
  if(mu<0) {printf("ERROR: Death rate mu (%g) is less than zero\n",mu); exit(1);}
  
  if(rho<0 || rho>1) {printf("ERROR: Mortality probability, rho (%g) is not a valid probability\n",mu); exit(1);}

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

  if(X0+Y0>N0) {printf("WARNING: Initial level of susceptibles+infecteds (%g+%g=%g) is greater than N0=%g\n",X0,Y0,X0+Y0,N0);}

  if(beta*(1-rho)*nu<(gamm+mu)*mu) {printf("WARNING: Basic reproductive ratio (R_0=%g) is less than one\n",beta*(1-rho)*nu/((gamm+mu)*mu));}

}


void Output_Data(FILE *out)
{
  if(out!=stdout) fprintf(out,"%g \t %g %g %g\n",t,X,Y,Z);

  if(DEBUG) printf("%g \t %g %g %g\n",t,X,Y,Z);
}



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

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

  initialPop[0]=X; initialPop[1]=Y; initialPop[2]=Z;

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

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

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

  Diff(tmpPop);

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

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


  X=tmpPop[0]; Y=tmpPop[1]; Z=tmpPop[2];

  return;
}