/***************************************************************************
 *   copyright Pleuni PENNINGS  2009 - 2012  MUNICH GERMANY - CAMBRIDGE MA USA                
 *   pleuni@dds.nl                                
 *   code for the simulations for Pennings 2012 PLOS COMP BIO
 *	 Standing genetic variation and the evolution of drug resistance in HIV 
 ***************************************************************************/

// System wide headers:
#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <sstream> 
#include <iostream> 
#include <fstream>
#include <vector> 
#include <math.h>
#include <string>
#include "gsl/gsl_rng.h"
#include "gsl/gsl_randist.h"
using namespace std;

//Global variables
//Parameters to be set in "GetParameters"
double cost; //cost of mutant at Druglevel = 0
double fitness_max_WT; //fitness of WT is max at Druglevel = 0
double fitness_MT;//fitness of the mutant at Druglevel = 0
double fitness_WT_drug;//fitness of the wt at max drug level
double fitness_MT_drug;//fitness of the mutant at max drug level
double mu;//mutation rate (from wt to resistant) per generation
unsigned int K; // carrying capacity
int Break; //length of treatment interruption (in generations, 200 gen = 1 year)
int I; //number of virus particals from latent cells
double t_wtg; //time it takes for Druglevel to reach the level where F_wt==1 to investigate "tail of treatment" (not used for Pennings 2012 PLOS COMP BIO)
double t_half; //time it takes for Druglevel to half
double ave_druglevel; // for the equilibrium state
double Druglevel; // druglevel at any timepoint
double Druglevel_for_output;


//Parameters for running the program
int N_start; //starting pop size
unsigned int seed, nRuns; // nRuns is the number of replicate runs of the simulation
int CheckProbFix; //  whether the fixation probability should be determined for a specific generation (not used for Pennings 2012 PLOS COMP BIO)
int Tmut; //  generation in which the mutant should be introduced to determine its fixation probability (not used for Pennings 2012 PLOS COMP BIO)
int tspd = 10; //time steps per day
int Ngen; //number of genarations of mutation

//Variables needed in the program
unsigned int repeat; //which of the nRuns is running
ofstream output; // to write output
unsigned int popArray[2]; //  Array that will hold the WT pop size and the Mut pop size (needed in this kind of array for the binomial sampling each generation) 
unsigned int popArrayReplace[2];//the newborns every generation 

//things to be set at timeline
int timeline = 10000; // each simulation runs over 10000 generations
bool Drug[10000]; //whether drug is taken on a given day or not
bool Mut[10000];//whether mutation is going on or not 
bool logN; //whether K should be chosen random from a lognormal distribution

//things to be filled during "popAdapt" 
int PopSizeWT[10000]; // population size of the wild type in each generation
int PopSizeMut[10000]; // pop size of the mutants in each generation
int PopTotal[10000]; // sum of WT and Mut popsizes in each generation
int NewMutants[10000]; // number of new mutants in in each generation

//descriptive statistics over a couple of runs
double MeanPopSizeWT[10000]; // to calculate mean popsize in given generation averaged over all nRuns
double MeanPopSizeMut[10000]; // Mut pop size averaged over all nRuns
double MeanNewMutants[10000]; // New mutants per generation averaged over all nRuns
double MutPresent[10000]; // probability that mutant is present averaged over all nRuns
double MeanTfix = 0; 

//To be determined after all runs are done
double FixProb = 0; //  to calculate the number of runs where the mutant has fixed and thereby the probability that it fixes. 

void getParameters(){ // parameters are entered in the command line or through a shell script
	cerr << "enter seed: "; cin >> seed; // starting seed
	cerr << "enter number of runs: "; cin >> nRuns; //number of replicates
	cerr << "enter N_start: "; cin >> N_start; //The starting pop size 
	cerr << "enter Max fitness_WT ";	cin >> fitness_max_WT; 
	cerr << "enter fitness_WT (at max Druglevel):";cin>>fitness_WT_drug; 
	cerr << "enter fitness_MT (fitness of MUTANT):";cin>>fitness_MT; 
	cerr << "enter fitness_MT (at max Druglevel):";cin>>fitness_MT_drug; 
	cerr << "enter mu ";cin >> mu; 
	cerr << "enter K (carrying capacity) ";	cin >> K;
	cerr << "enter number of viral particals to enter pop from latent cells"; cin >> I; 
	cerr << "enter break length (0 for no break, 8500 for start of treatment):";cin>>Break;
 	cerr << "enter t_wtg, the number of days it takes before druglevels are low enough for wt virus to grow again"; cin >> t_wtg; //always 0 in Pennings PLOS COMP BIO 2012
	cerr <<"enter average drug level: "; cin >> ave_druglevel; //always 1 in Pennings PLOS COMP BIO 2012
 	cerr << "enter Ngen numb of generations of mutation after start of treatment"; cin >> Ngen; //set to 0 if only looking at standing genetic variation
 	cerr << "enter 1 if K should be from logNormal distribution"; cin >> logN; //always 0 in Pennings PLOS COMP BIO 2012
	if (Break==8500){seed =1000*seed;} // this is for historical reasons to make sure i get exactly the results as in the paper Pennings PLOS COMP BIO 2012 
	cost = fitness_max_WT-fitness_MT; // the absolute cost of the resistance mutation in the absence of drugs
	t_half=-t_wtg/log2(1-1/fitness_max_WT); // halflife of the drugs, always 0 in Pennings PLOS COMP BIO 2012
}

void prepareTimeline(){
	int m=0; 
	for(m=0; m < 10000; m++) {Drug[m]=1;Mut[m]=0;}//to reach pop dyn equilibirum 
	for(m=(8500-Break); m < 8500 ; m++) {Drug[m]=0;}// no drugs during break 
	// to simulate the start of treatment, I use a very long window of no treatment, by setting break =8500
	if (Break == 8500){for(m=8500-200; m < 8500+Ngen ; m++) {Mut[m]=1;}}//turn on mutation only during 200 generations before treatment and Ngen during treatment 
	//to simulate a real treatment interruption, I set break to a small number 
	if (Break != 8500){for(m=(8500-Break); m < 8500 ; m++) {Mut[m]=1;}}//during the break there is mutation
	}

void popAdapt(gsl_rng *rng){
	//popAdapt has two parts: initialize (to make a monomorphic population and seed the random number generator) and evolve (a loop that loops until the ancestral allele is lost). 
	
	//initialize
	gsl_rng_set (rng, (unsigned long)repeat+seed+1); 
	//set everything to zero for a new run
	if (logN==1){N_start=K=pow(10,3.6+gsl_ran_gaussian(rng,0.5)); I=K*(1-fitness_max_WT)/(34);} //choose the population size (not used for Pennings PLOS COMP BIO 2012)
	for(int m=0; m < timeline ; m++){PopSizeWT[m]=0; PopSizeMut[m]=0; NewMutants[m]=0;PopTotal[m]=0;} //set the pop sizes to 0
	PopSizeWT[0]=N_start; PopTotal[0]=N_start; //set the popsize in the first generation to N_start
	MeanPopSizeWT[0]=N_start; MeanPopSizeMut[0]=0; MeanNewMutants[0]=0; MutPresent[0]=0; 
	popArray[0]=PopSizeWT[0]; popArray[1]=PopSizeMut[0]; 
	popArrayReplace[0]=0; popArrayReplace[1]=0; 
	Druglevel=0;
	
	//evolve
	for (int t=1; t<timeline;t++)  { //each loop is one generation, continues until end of timeline
		for (int ts=0; ts<(tspd);ts++)  { // each generation is split in tspd steps (10 in this code)
			//druglevel
			if (Drug[t]==1){Druglevel=ave_druglevel;}
			if (Drug[t]==0){Druglevel=Druglevel*pow(2,-1/(tspd*t_half));} // it takes t_half days for the drug to half, and there are tspd steps in one day
			if (t==8499+Break&&ts==(tspd-1)) Druglevel_for_output=Druglevel; //not relevant for current code
			double deathrate = 1./tspd; 
		//how many should die & have the chance to reproduce? 
			int deathWT=gsl_ran_binomial(rng, deathrate,popArray[0]);
			int deathMT=gsl_ran_binomial(rng, deathrate,popArray[1]);
		//how many should replace them? 
			double birthrateWT=(fitness_max_WT-(fitness_max_WT-fitness_WT_drug)*Druglevel);  
			double birthrateMT=(fitness_MT-(fitness_MT-fitness_MT_drug)*Druglevel);
			popArrayReplace[0]=gsl_ran_poisson(rng, birthrateWT*deathWT); //independent reproduction for WTs and Muts, no competition
			popArrayReplace[1]=gsl_ran_poisson(rng, birthrateMT*deathMT); //independent reproduction for WTs and Muts, no competition
		//make sure the pop doesnt grow too big (i.e. it should not be >K)
			int add_deaths=0; 
			if (popArray[0]+popArray[1]>K){add_deaths=popArray[0]+popArray[1]-K;}
			double ratio = 1.*popArray[0]/(1.*popArray[0]+1.*popArray[1]);
			int add_deathsWT=gsl_ran_binomial(rng,ratio,add_deaths);
			deathWT+=add_deathsWT; deathMT+=(add_deaths-add_deathsWT);
			if (deathMT>popArray[1])deathMT=popArray[1];
		//mutation
			int mut_from_wt =0; int mut_from_res=0; int wt_from_res=0; //mutations from replication and from the reservoir 
			if (Mut[t]==1){				
				mut_from_wt = gsl_ran_binomial(rng, mu, popArrayReplace[0]); mut_from_res = gsl_ran_binomial(rng, mu, I*1./tspd); }//determine number new muts
			wt_from_res = I*1./tspd - mut_from_wt; //make sure wt doesnt die out, wt virus comes from reservoir
			NewMutants[t]+=mut_from_wt; //bookkeeping
		//do it all 
			popArray[0]=popArray[0]+popArrayReplace[0]-deathWT-mut_from_wt;
			popArray[1]=popArray[1]+popArrayReplace[1]-deathMT+mut_from_wt;
		//add wt and resistant particles from reservoir
			popArray[0]=popArray[0]+wt_from_res;popArray[1]=popArray[1]+mut_from_res;
				NewMutants[t]+=mut_from_res;
			}
		PopTotal[t]=popArray[0]+popArray[1]; 
		PopSizeWT[t]=popArray[0]; PopSizeMut[t]=popArray[1]; //bookkeeping
		MeanPopSizeWT[t]+=PopSizeWT[t]; MeanPopSizeMut[t]+=PopSizeMut[t]; MeanNewMutants[t]+=NewMutants[t]; if (PopSizeMut[t]>0) MutPresent[t]+=1; // bookkeeping 
		}
	//at the end of a simulation run, I need to know whether adaptation happened
	if (PopSizeMut[9999]>PopSizeWT[9999]){
		int gen; for (gen=1; PopSizeMut[gen]<PopSizeWT[gen]; gen++); MeanTfix = MeanTfix+gen; }//get time of fixation after start of treatment
	if ((double)PopSizeMut[timeline-1]>(double)PopSizeWT[timeline-1]) {FixProb+=1; } //I assume a mutation will fix  popsizemut is larger than popsize WT
}


int main (int argc, char *argv[]){
	getParameters();
	gsl_rng *rng = gsl_rng_alloc(gsl_rng_mt19937);//choose random number generator
	prepareTimeline();
	
	//do nRuns simulation runs
	for (repeat = 0; repeat < nRuns; repeat++){	//each of these loops is an independent run
		popAdapt(rng);}
	
	//calculate & output means over all generations
	output<<"t\t" <<"PopSizeWT\t"<<"PopSizeMut\t"<<"MutPresent\t"<<"MeanNewMuts\tMut\tDrug\n";
	for (int t=1; t<timeline;t++)  {
		MeanPopSizeWT[t]/=(double)nRuns; 
		MeanPopSizeMut[t]/=(double)nRuns; 
		MeanNewMutants[t]/=(double)nRuns;
		MutPresent[t]/=(double)nRuns;
			output<<t<<"\t"<<MeanPopSizeWT[t]<<"\t"<<MeanPopSizeMut[t]<<"\t"<<MutPresent[t]<<"\t"<<MeanNewMutants[t]<<"\t"<<Mut[t]<<"\t"<<Drug[t]<<"\n";}
	MeanTfix/=(double)nRuns;
	FixProb/=(double)nRuns; //to calculate the probability that adaptation has happened
	
	double totalnewmuts; 
	for (int m = 0; m<10000; m++){totalnewmuts+=MeanNewMutants[m];}
	cout << "averagedruglevel\t"<<ave_druglevel<< "\tFitnessWT\t"<< fitness_max_WT-(fitness_max_WT-fitness_WT_drug)*ave_druglevel;
	cout << "\tFitnessWTmax\t"<< fitness_max_WT<<"\tFitnessMT\t"<< fitness_MT; 
	cout <<"\tFitnessMTDrug\t" << fitness_MT_drug<<"\tNumberNewMutants\t"<<totalnewmuts; 
	cout <<"\tFixprob\t" <<FixProb<<"\tMaxPopsize\t"<<MeanPopSizeWT[8499]; 
	cout <<"\tPopMT\t"<<MeanPopSizeMut[8499]<<"\tLengthBreak\t"<<Break; 
	cout <<"\tTwtg\t"<<t_wtg<<"\tmu\t"<<mu<<"\tI\t"<<I; 
	cout <<"\tdrug_end_break\t"<<Druglevel_for_output<<"\tMeanTfix\t"<<MeanTfix; 
	cout <<"\tcost\t"<<fitness_max_WT-fitness_MT<<"\n";
	
	return 5;
}