/**************************************************************************
 * niche_evolution.cpp
 * Source code for the model described in Ashby et al "Competing species 
 * leave many potential niches unfilled". Each row of the file 
 * "niche_evolution.txt" records the n-dimensional coordinates of each 
 * species in niche space at a given evolutionary time point, padded by 
 * NaNs. For example, in 2-dimensional niche space, column 1 is the 
 * x-coordinate of species 1, column 2 is the y-coordinate of species 1, 
 * column 3 is the x-coordiante of species 2, and so on.
 *
 * The authors accept no responsibility or liability for any loss or damage
 * arising from the use of this code.
 *
 * Author: Ben Ashby
 * 09/10/2016
 * v1.0 
 *************************************************************************/

#include <iostream>
#include <fstream>
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

/* Model parameter values */
#define ADDITIVE 1 /* Model version: 1=additive, 0=multiplicative*/
#define DIMENSIONS 2 /* Dimensionality of niche space */
#define NICHE_DIMENSION_SIZE 7.0 /* Length of each niche axis */
#define MAXSPECIES 100 /* Maximum number of species allowed */
#define R_GRAD 0.0 /* Rate at which r falls away from its maximum */
#define K_GRAD 0.0 /* Rate at which K falls away from its maximum */
#define EPSILON1 0.1 /* Maximum mutation distance in each dimension*/
#define EPSILON2 0.05 /* Extinction threshold */

/* Solver parameter values */
#define NEVOL 1000 /* Number of iterations (evolutionary timesteps) */
#define MAXTIME 1e4 /* Duration for ecological dynamics */
#define MAXSTEPS 1e6 /* Maximum number of steps for ODE solver */
#define INTERVAL 1000 /* Check if the system is close to equilibrium */
#define EQTOL 1e-4 /* Equilibrium tolerance (ODE) */
#define EPS 1e-6 /* ODE solver tolerance */
#define TINY 1e-6 /* Constant value for solver */
#define PI 3.141592653589793 /* Pi */

/*************************************
 * Function prototypes
 *************************************/
void my_rungkut(double *x, double **positions, double *r, double *K, int species);
void rkqs(double *y, double *dydt, double **alpha, double *r, double *K, int species, double *h1, double *hnext, double *yscale);
void rkck(double *y, double *dydt, double **alpha, double *r, double *K, int species, double *yout, double *yerr, double h1);
void ddt(double *x, double *dev, double **alpha, double *r, double *K, int species);
double FMAX(double l, double r);
double FMIN(double l, double r);
double** array_maker(int rows, int cols);
void free_array(double **array, int rows);
double uniform_rng();
double normal_rng();
double my_mod(double x, double y);

/***************************************
 * Main function
 ***************************************/
int main (int argc, char* argv[]) {
    
    double **positions, pop_size[MAXSPECIES], r[MAXSPECIES], K[MAXSPECIES];
    double prog_gap, next_prog, d;
    int i, j, k, species, species_count;
    char filename[100];
    
    /* Random seed */
    srand(time(0));
    prog_gap = 0.05;
    
    /* Create output file */
    sprintf(filename, "niche_evolution.txt");
    std::cout << "filename:" << filename << "\n";
    std::ifstream infile(filename);
    if(infile.good()) {
        std::cout << "File already exists, deleting file...\n";
        remove(filename);
    }
    std::ofstream out(filename, std::ios::app);
    if (!out){
        std::cout << "Cannot create file\n";
        exit(1);
    }
    
    /* Initial conditions */
    positions = array_maker(MAXSPECIES, DIMENSIONS);
    for (i=0; i<DIMENSIONS; i++) {
        positions[0][i] = NICHE_DIMENSION_SIZE*0.5;
    }
    species = 1;
    pop_size[0] = 1;
    d = 0;
    for(k=0;k<DIMENSIONS;k++) {
        d += (positions[0][k]/NICHE_DIMENSION_SIZE-0.5)*(positions[0][k]/NICHE_DIMENSION_SIZE-0.5);
    }
    d = sqrt(d);
    r[0] = FMAX(1e-30, 1 - R_GRAD*d);
    K[0] = FMAX(1e-30, 1 - K_GRAD*d);

    next_prog = prog_gap;
    
    /* Evolution routine */    
    for (i=0; i<NEVOL; i++){
        
        /* Output progress */
        if((double)i/(double)NEVOL>=next_prog){
            std::cout << "Progress = " << 100*next_prog << "%\n";
            next_prog+=prog_gap;
        }
        
        /* Call ODE solver */
        my_rungkut(pop_size, positions, r, K, species);
        
        /* Remove species that are below the extinction threshold */
        species_count = 0;
        for (j=0; j<species; j++){
            if(pop_size[j]>=EPSILON2){
                if(species_count<j){
                    for(k=0;k<DIMENSIONS;k++) {
                        positions[species_count][k] = positions[j][k];
                    }
                    pop_size[species_count] = pop_size[j];
                    r[species_count] = r[j];
                    K[species_count] = K[j];
                }
                species_count++;
            }
            else{
                pop_size[j]=0;
            }
        }
        species = species_count;
        
        /* Record output */
        for (j=0; j<species; j++){
            for (k=0; k<DIMENSIONS; k++) out << positions[j][k] << " ";
        }
        /* Pad output*/
        for(j=species;j<MAXSPECIES;j++){
            for (k=0; k<DIMENSIONS; k++) out << "NaN" << " ";
        }
        out << "\n";
        
        /* Introduce mutants close to each species (up to MAXSPECIES) */
        if(species<=(0.5*MAXSPECIES)){
            for (j=0; j<species; j++){
                d = 0;
                for(k=0;k<DIMENSIONS;k++) {
                    positions[species+j][k] = my_mod(positions[j][k] + EPSILON1*(0.5 - uniform_rng()), NICHE_DIMENSION_SIZE);
                    d += (positions[species+j][k]/NICHE_DIMENSION_SIZE - 0.5)*(positions[species+j][k]/NICHE_DIMENSION_SIZE - 0.5);
                }
                d = sqrt(d);
                
                /* Copy population size and mutate demographic parameters */
                pop_size[species+j] = pop_size[j]/10;
                
                r[species+j] = FMAX(1e-30, 1 - R_GRAD*d);
                K[species+j] = FMAX(1e-30, 1 - K_GRAD*d);
            }
            species*=2;
        }
        else{
            /* Choose MAXSPECIES/2 species to mutate - randomly choose
             beginning or end of list */
            if(uniform_rng()<0.5){
                for (j=0; j<(0.5*MAXSPECIES); j++){
                    if((species+j)>=MAXSPECIES) break;
                    d = 0;
                    for(k=0;k<DIMENSIONS;k++) {
                        positions[species+j][k] = my_mod(positions[j][k] + EPSILON1*(0.5 - uniform_rng()),NICHE_DIMENSION_SIZE);
                        d += (positions[species+j][k]/NICHE_DIMENSION_SIZE - 0.5)*(positions[species+j][k]/NICHE_DIMENSION_SIZE - 0.5);
                    }
                    d = sqrt(d);
                    
                    /* Copy population size and mutate demographic parameters */
                    pop_size[species+j] = pop_size[j]/10;
                    r[species+j] = FMAX(1e-30, 1 - R_GRAD*d);
                    K[species+j] = FMAX(1e-30, 1 - K_GRAD*d);
                }
            }
            else{
                for (j=0; j<(0.5*MAXSPECIES); j++){
                    if((species+j)>=MAXSPECIES) break;
                    for(k=0;k<DIMENSIONS;k++) {
                        positions[species+j][k] = my_mod(positions[species-j][k] + EPSILON1*(0.5 - uniform_rng()),NICHE_DIMENSION_SIZE);
                    }
                    
                    /* Copy population size and mutate demographic parameters */
                    pop_size[species+j] = pop_size[species-j]/10;
                    r[species+j] = FMAX(1e-30, 1 - R_GRAD*d);
                    K[species+j] = FMAX(1e-30, 1 - K_GRAD*d);
                }
            }
            species = MAXSPECIES;
        }
    }
    
    /* Free memory */
    free_array(positions, MAXSPECIES);
}

/**************************************************************************
 * ODE solver
 *************************************************************************/
void my_rungkut(double *x, double **positions, double *r, double *K, int species){
    
    int i, j, k, exitflag, count;
    double t, nextcheck, d, hnext[1], h[1];
    double **alpha, *dxdt, *xscale, *xmax, *xmin;
    
    /* Allocate memory */
    alpha = array_maker(species, species);
    dxdt = (double *)malloc(species*sizeof(double));
    xscale = (double *)malloc(species*sizeof(double));
    xmax = (double *)malloc(species*sizeof(double));
    xmin = (double *)malloc(species*sizeof(double));
    
    /* Other parameters */
    exitflag = 0;
    count=0;
    h[0] = 1e-3;
    hnext[0] = 1e-3;
    t=0;
    nextcheck = INTERVAL;
    
    /* Calculate competition coefficients */
    for (i=0;i<species;i++) {
        for (j=0;j<species;j++) {
            if(i==j){
                alpha[i][j]=1;
            }
            else if(i<j){
                if(ADDITIVE>0){ /* Additive version */
                    alpha[i][j] = 0;
                    for (k=0;k<DIMENSIONS;k++) {
                        d = FMIN((positions[i][k]-positions[j][k])*(positions[i][k]-positions[j][k]),FMIN((NICHE_DIMENSION_SIZE+positions[i][k]-positions[j][k])*(NICHE_DIMENSION_SIZE+positions[i][k]-positions[j][k]),(positions[i][k]-positions[j][k]-NICHE_DIMENSION_SIZE)*(positions[i][k]-positions[j][k]-NICHE_DIMENSION_SIZE)));
                        alpha[i][j]+=exp(-0.5*d);
                    }
                    alpha[i][j]/=((double)(DIMENSIONS));
                }
                else{ /* Multiplicative version */
                    d=0;
                    for (k=0;k<DIMENSIONS;k++) {
                        d += FMIN((positions[i][k]-positions[j][k])*(positions[i][k]-positions[j][k]),FMIN((NICHE_DIMENSION_SIZE+positions[i][k]-positions[j][k])*(NICHE_DIMENSION_SIZE+positions[i][k]-positions[j][k]),(positions[i][k]-positions[j][k]-NICHE_DIMENSION_SIZE)*(positions[i][k]-positions[j][k]-NICHE_DIMENSION_SIZE)));
                    }
                    alpha[i][j] = exp(-0.5*d);
                }
            }
            else{
                alpha[i][j] = alpha[j][i];
            }
        }
    }
    
    /* Set up initial population */
    for (i=0;i<species;i++) {
         xmax[i] = x[i];
         xmin[i] = x[i];
    }
    
    /* Main loop */
    do{
        /* This ensures the final step lands us on the final time point */
        if(1.1*hnext[0]>(MAXTIME-t)){
            hnext[0] = MAXTIME-t;
            h[0] = MAXTIME-t;
            t=MAXTIME;
        }
        else{
            h[0] = hnext[0];
            t+=h[0];
        }
        if(t==MAXTIME) {
            exitflag=1;
        }
        
         /* This is where your equations are first solved */
        ddt(x, dxdt, alpha, r, K, species);
        
        /* Adjust the step size to maintain accuracy */        
        for (i=0;i<species;i++) xscale[i]=fabs(x[i])+fabs(dxdt[i]*(*h))+TINY;
        rkqs(x, dxdt, alpha, r, K, species, h, hnext, xscale);
        count++;
        
        /* Check if we're close to equilibrium */
        for (i=0; i<species; i++){
            x[i] = FMAX(0, x[i]);
            xmax[i]=FMAX(xmax[i], x[i]);
            xmin[i]=FMIN(xmin[i], x[i]);
        }
        if(t>nextcheck){
            exitflag = 1;
            for (i=0; i<species; i++){
                if(fabs(xmax[i]-xmin[i])>EQTOL){
                    exitflag = 0;
                    break;
                }
            }
            if(exitflag==1){
                break;
            }
            nextcheck+=INTERVAL;
            for(i=0;i<species;i++) {
                xmax[i] = x[i];
                xmin[i] = x[i];
            }
        } 
    }while(count<MAXSTEPS && t<=MAXTIME && !exitflag);
    
    /* Free memory */
    free_array(alpha, species);
    free(dxdt);
    free(xscale);
    free(xmax);
    free(xmin);
}

/**************************************************************************
 * ODE equations
 *************************************************************************/
void ddt(double *x, double *dev, double **alpha, double *r, double *K, int species) {
    
    int i,j;
    double comp;
    
    /* Calculate total effect of competition on species i */
    for(i=0;i<species;i++){
        comp = 0;
        for(j=0;j<species;j++){
            comp+=alpha[i][j]*x[j];
        }
        dev[i] = r[i]*x[i]*(1 - comp/K[i]); /* ODEs */
    }
}

/**************************************************************************
 * Part of the Runge-Kutta algorithm
 *************************************************************************/
void rkqs(double *y, double *dydt, double **alpha, double *r, double *K, int species, double *h1, double *hnext, double *yscale) {
    double *ytemp, *yerr;
    double htemp, errmax;
    int i, n, count;
    
    n= species;
    ytemp= (double *)malloc(n * sizeof(double));
    yerr=  (double *)malloc(n * sizeof(double));
    
    htemp= *h1;
    
    rkck(y, dydt, alpha, r, K, species, ytemp, yerr, *h1);
    *hnext= *h1;
    count = 0;
    
    for(;;) {
        rkck(y, dydt, alpha, r, K, species, ytemp, yerr, *h1);
        errmax= 0.0;
        for(i=0;i<n;i++) errmax= FMAX(errmax, fabs(yerr[i]/yscale[i]));
        errmax/= EPS;
        if(errmax<=1.0) break;
        count++;
        if(count>1e4) {
            printf("Unknown error with adaptive step size!\n");
            exit(EXIT_FAILURE);
        }
        htemp= 0.9*(*h1)*pow(errmax, -0.25);
        *h1= (*h1>=0.0 ? FMAX(htemp, 0.1*(*h1)) : FMIN(htemp, 0.1*(*h1)));
    }
    if(errmax > 1.89E-4) {
        *hnext= 0.9*(*h1)*pow(errmax, -0.2);
    }
    else {
        *hnext= 5.0*(*h1);
    }
    
    for(i=0;i<n;i++) y[i]= ytemp[i];
    
    free(ytemp);
    free(yerr);
}

/**************************************************************************
 * Part of the Runge-Kutta algorithm
 *************************************************************************/
void rkck(double *y, double *dydt, double **alpha, double *r, double *K, int species, double *yout, double *yerr, double h1) {
    int i, n;
    double *k2, *k3, *k4, *k5, *k6, *ytemp;
    static double b_m1=0.2, b31=3.0/40.0, b32=9.0/40.0, b41=0.3, b42=-0.9, b43=1.2,
            b51=-11.0/54.0, b52=2.5, b53=-70.0/27.0, b54=35.0/27.0, b61=1631.0/55296,
            b62=175.0/512.0, b63=575.0/13824.0, b64=44275.0/110592, b65=253.0/4096.0,
            c1=37.0/378.0, c3=250.0/621.0, c4=125.0/594.0, c6=512.0/1771.0, dc5=-277.00/14336;
    double dc1=c1-2825.0/27648.0, dc3=c3-18575.0/48384.0, dc4=c4-13525.0/55296.0,
            dc6=c6-0.25;
    
    n=species;
    
    k2= (double *)malloc(n * sizeof(double));
    k3= (double *)malloc(n * sizeof(double));
    k4= (double *)malloc(n * sizeof(double));
    k5= (double *)malloc(n * sizeof(double));
    k6= (double *)malloc(n * sizeof(double));
    ytemp= (double *)malloc(n * sizeof(double));
    
    for(i=0;i<n;i++)
        ytemp[i]= y[i]+b_m1*h1*dydt[i];
    ddt(ytemp, k2, alpha, r, K, species);
    
    for(i=0;i<n;i++)
        ytemp[i]= y[i]+h1*(b31*dydt[i]+b32*k2[i]);
    ddt(ytemp, k3, alpha, r, K, species);
    
    for(i=0;i<n;i++)
        ytemp[i]= y[i]+h1*(b41*dydt[i]+b42*k2[i]+b43*k3[i]);
    ddt(ytemp, k4, alpha, r, K, species);
    
    for(i=0;i<n;i++)
        ytemp[i]= y[i]+h1*(b51*dydt[i]+b52*k2[i]+b53*k3[i]+b54*k4[i]);
    ddt(ytemp, k5, alpha, r, K, species);
    
    for(i=0;i<n;i++)
        ytemp[i]= y[i]+h1*(b61*dydt[i]+b62*k2[i]+b63*k3[i]+b64*k4[i]+b65*k5[i]);
    ddt(ytemp, k6, alpha, r, K, species);
    
    for(i=0;i<n;i++)
        yout[i]= y[i]+h1*(c1*dydt[i]+c3*k3[i]+c4*k4[i]+c6*k6[i]);
    for(i=0;i<n;i++)
        yerr[i]= h1*(dc1*dydt[i]+dc3*k3[i]+dc4*k4[i]+dc5*k5[i]+dc6*k6[i]);
    
    free(k2);
    free(k3);
    free(k4);
    free(k5);
    free(k6);
    free(ytemp);
}

/**************************************************************************
 * Max and Min functions
 **************************************************************************/
double FMAX(double l, double r)
{
    if(l>r)return l;
    else   return r;
}
double FMIN(double l, double r)
{
    if(l<r)return l;
    else   return r;
}

/***************************************
 * Make 2D double array
 ***************************************/
double** array_maker(int rows, int cols) {
    
    double** new_array;
    new_array = (double**) malloc(rows*sizeof(double*));
    for (int i = 0; i < rows; i++)
        new_array[i] = (double*) malloc(cols*sizeof(double));
    
    return new_array;
}

/***************************************
 * Free 2D double array
 ***************************************/
void free_array(double **array, int rows) {
    
    for (int i = 0; i < rows; i++) free(array[i]);
    free(array);
}

/**************************************************************************
 * Uniform random number generator
 *************************************************************************/
double uniform_rng()
{
    return ((double)(rand())+1.)/((double)(RAND_MAX)+1.);
}

/**************************************************************************
 * Returns a normally distributed real number, using the Box-Muller
 * algorithm
 * ***********************************************************************/
double normal_rng()
{
    double u1, r, theta;
    static double u2;
    static int first_run = 1;
    
    if(first_run){
        u1 = uniform_rng();
        first_run=0;
    }
    else{
        u1=u2;
    }
    u2 = uniform_rng();
    r = sqrt(-2.0*log(u1));
    theta = 2.0*PI*u2;
    
    return r*sin(theta);
}

/**************************************************************************
 * Remove boundary effects
 * ***********************************************************************/
double my_mod(double x, double y){
    
    if(x>=y){
        return x-NICHE_DIMENSION_SIZE;
    }
    else if(x<0){
        return x+NICHE_DIMENSION_SIZE;
    }
    else {
        return x;
    }
    
}