/* Affine-invariant MCMC method of Goodman and Weare 2010.  Here we analyze a one-dimensional Gaussian. */

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

#define amult 2.0      /* Parameter for selection of multiplier */
#define Nmult 30    /* Number of walkers = Nmult times number of parameters */
#define Ndat 28 /* Number of data points */
#define dLogl 8.0    /* Threshold for convergence */
#define Nconv 50      /* Spread of steps to test convergence */
#define dlogltol 20.0 /* Require for convergence that walkers be this close */
#define Maxsteps 10000 /* Maximum number of steps allowed */
#define Ncolumn 4     /* Number of columns of data */
#define PI 3.14159265359


void initialize(double **y, double *logl, int Nparam, int Nwalker);
void mcmc(double **y, double *logl, int N, int Nparam, int Nwalker, char *datfile);
double loglikelihood(double *z, int Nparam);
int inbounds(double *z, int Nparam);
void stepprint(double **y, double *logl, int Nparam, int Nwalker, char *datfile);

int converged,Nstep;
int try,success;
double *limmin,*limmax;
double *ybest,loglbest;
double **loglhistory;
double **datax;

int main(int argc, char *argv[])
{
    int j,i,N,M,Nparam,Nwalker;
    char **name;
    double **y,*logl;
    FILE *limits,*data,*acceptfrac;
    
    srand48(getpid());
    
    Nparam=atoi(argv[3]);
    ybest=calloc(Nparam,sizeof(double));
    limmin=calloc(Nparam,sizeof(double));
    limmax=calloc(Nparam,sizeof(double));
    limits=fopen("limits.dat","r");
    for (i=0; i<Nparam; i++)
    {
        fscanf(limits,"%lg %lg",&limmin[i],&limmax[i]);
    }
    fclose(limits);
    datax=calloc(Ndat,sizeof(double *));
    for (i=0; i<Ndat; i++)
       datax[i]=calloc(Ncolumn,sizeof(double));
    data=fopen("data.dat","r");
    for (i=0; i<Ndat; i++)
        for (j=0; j<Ncolumn; j++)
            fscanf(data,"%lg",&datax[i][j]);
    fclose(data);

    
    M=atoi(argv[1]);
    N=atoi(argv[2]);
    Nwalker=Nmult*Nparam;
    y=calloc(Nwalker,sizeof(double *));
    for (i=0; i<Nwalker; i++)
    {
        y[i]=calloc(Nparam,sizeof(double));
    }
    logl=calloc(Nwalker,sizeof(double));
    name=calloc(M,sizeof(char *));
    for (j=0; j<M; j++)
        name[j]=calloc(20,sizeof(char));
    for (j=0; j<M; j++)
        sprintf(name[j],"chain%d.dat",j+1);

    try=0;
    success=0;
    for (j=0; j<M; j++)
    {
        loglbest=-1.0e20;
        initialize(y,logl,Nparam,Nwalker);
        mcmc(y,logl,N,Nparam,Nwalker,name[j]);
        acceptfrac=fopen("acceptfrac.dat","a");
        fprintf(acceptfrac,"Chain %d: %d tried, %d successful\n",j,try,success);
        fclose(acceptfrac);
        try=0;
        success=0;
    }
    
    for (j=0; j<M; j++)
        free(name[j]);
    free(name);
    for (i=0; i<Nwalker; i++)
        free(y[i]);
    free(y);
    free(logl);
    free(limmin);
    free(limmax);
    free(ybest);
    for (i=0; i<Ndat; i++)
        free(datax[i]);
    free(datax);
}

void initialize(double **y, double *logl, int Nparam, int Nwalker)
{
    int i,j;
    double loglike;

    for (i=0; i<Nwalker; i++)
    {
        if (i==0)
        do
        {
            for (j=0; j<Nparam; j++)
                y[i][j]=limmin[j]+(limmax[j]-limmin[j])*drand48();
/* Best results from previous runs */
        } while (!inbounds(y[i],Nparam));
        else
        {
            do
            {
                for (j=0; j<Nparam; j++)
                    y[i][j]=y[0][j]+0.001*(limmax[j]-limmin[j])*(2.0*drand48()-1.0);
            } while (!inbounds(y[i],Nparam));
        }
        logl[i]=loglikelihood(y[i],Nparam);
    }
}

void mcmc(double **y, double *logl, int N, int Nparam, int Nwalker, char *datfile)
{
    int i,j;
    double logl0;
    FILE *dat;
    
    converged=0;           /* Convergence flag */
    Nstep=0;               /* Restart Nstep counter */
    loglhistory=calloc(Nwalker,sizeof(double *));
    for (i=0; i<Nwalker; i++)
        loglhistory[i]=calloc(Maxsteps,sizeof(double));
    
    /* For this code, print out initial walkers, all steps, walkers after burn-in, and walkers after completion. */
    
    dat=fopen(datfile,"a");
    for (i=0; i<Nwalker; i++)
    {
        fprintf(dat,"%3d ",i+1);
        for (j=0; j<Nparam; j++)
        {
            fprintf(dat,"%9.6f ",y[i][j]);
        }
        logl0=loglikelihood(y[i],Nparam);
        fprintf(dat,"%.4f\n",logl0);
        loglhistory[i][Nstep]=logl0;
    }
    fprintf(dat,"\n");
    fclose(dat);
    Nstep++;
    for (i=0; (i<N || !converged) && (Nstep<Maxsteps-2*N); i++)
    {
        stepprint(y,logl,Nparam,Nwalker,datfile);
    }
    dat=fopen(datfile,"a");
    fprintf(dat,"\n");
    for (i=0; i<Nwalker; i++)
    {
        fprintf(dat,"%3d ",i+1);     
        for (j=0; j<Nparam; j++)
        {
            fprintf(dat,"%9.6f ",y[i][j]);
        }
        logl0=loglikelihood(y[i],Nparam);
        fprintf(dat,"%.4f\n",logl0);
    }
    fprintf(dat,"1000 ");
    for (i=0; i<Nparam; i++)
        fprintf(dat,"0 ");
    fprintf(dat,"-1.0e20\n");
    fprintf(dat,"\n");
    fclose(dat);
    for (i=N; i<2*N; i++)
    {
        stepprint(y,logl,Nparam,Nwalker,datfile);
    }
    dat=fopen(datfile,"a");
    fprintf(dat,"\n");
    fprintf(dat,"2000 ");
    for (i=0; i<Nparam; i++)
        fprintf(dat,"0 ");
    fprintf(dat,"-1.0e20\n");
    for (i=0; i<Nwalker; i++)
    {
        fprintf(dat,"%3d ",i+1);     
        for (j=0; j<Nparam; j++)
        {
            fprintf(dat,"%9.6f ",y[i][j]);
        }
        logl0=loglikelihood(y[i],Nparam);
        fprintf(dat,"%.4f\n",logl0);
    }
    fclose(dat);
    for (i=0; i<Nwalker; i++)
        free(loglhistory[i]);
    free(loglhistory);
}

int inbounds(double *z, int Nparam)
{
    int bound;
    int j;
    
    bound=1;
    for (j=0; j<Nparam; j++)  /* Generally, would be Nparam instead of 8 */
    {
        if (z[j]<limmin[j] || z[j]>limmax[j])
            bound=0;
    }
    
    return bound;
}

void stepprint(double **y, double *logl, int Nparam, int Nwalker, char *datfile)
{
    int i,istretch,j,*accept;
    double **ycand,Z,x,*loglcand;
    FILE *dat;
    
    dat=fopen(datfile,"a");
    accept=calloc(Nwalker,sizeof(int));
    ycand=calloc(Nwalker,sizeof(double *));
    for (i=0; i<Nwalker; i++)
        ycand[i]=calloc(Nparam,sizeof(double));
    loglcand=calloc(Nwalker,sizeof(double));
    
    /* Select candidate moves for each of the walkers */
    for (i=0; i<Nwalker; i++)
    {
        try+=1;
        do
        {
            x=drand48();
            Z=sqrt(1.0/amult)+x*(sqrt(amult)-sqrt(1.0/amult));
            Z=Z*Z;
            do
            {
                istretch=((int) (Nwalker*drand48()));
            } while (istretch==i);
            for (j=0; j<Nparam; j++)
                ycand[i][j]=y[istretch][j]+Z*(y[i][j]-y[istretch][j]);
        } while (!inbounds(ycand[i],Nparam));
        if (inbounds(ycand[i],Nparam))
        {
            loglcand[i]=loglikelihood(ycand[i],Nparam);
            x=log(drand48());
            if (x<((Nparam-1.0)*log(Z)+loglcand[i]-logl[i]))
                accept[i]=1;
            else
                accept[i]=0;
            /* If the step has been accepted and the log likelihood is the highest yet, update ybest and loglbest. */
            if (accept[i] && loglcand[i]>loglbest)
            {
                loglbest=loglcand[i];
                for (j=0; j<Nparam; j++)
                    ybest[j]=ycand[i][j];
            }
        }
    }
    converged=1;
    for (i=0; i<Nwalker; i++)
    {
        if (accept[i])
        {
            success+=1;
            for (j=0; j<Nparam; j++)
            {
                y[i][j]=ycand[i][j];
            }
            logl[i]=loglcand[i];
        }
        fprintf(dat," %3d ",i+1);
        for (j=0; j<Nparam; j++)
        {
            fprintf(dat,"%9.6f ",y[i][j]);
        }
        fprintf(dat,"%.4f\n",logl[i]);
        loglhistory[i][Nstep]=logl[i];
        if (Nstep>=Nconv)
        {
            if (loglhistory[i][Nstep]>loglhistory[i][Nstep-Nconv]+dLogl)
                converged=0;
        }
    }
    
    /* If we think we might have converged, next step is to see whether all of the walkers have log likelihoods within some tolerance of the best logl. */
    if (converged==1 && Nstep>=Nconv)   /* Possible convergence; let's check */
    {
        for (i=0; i<Nwalker; i++)
        {
            if (loglhistory[i][Nstep]<loglbest-dlogltol) /* Not close enough */
            {
                for (j=0; j<Nparam; j++)
                {
                    y[i][j]=ybest[j];
                }
                loglhistory[i][Nstep]=loglbest;
                converged=0;
            }
        }
    }
    
    Nstep++;
    
    for (i=0; i<Nwalker; i++)
        free(ycand[i]);
    free(ycand);
    free(accept);
    free(loglcand);
    fclose(dat);
}

double loglikelihood(double *z, int Nparam)
{

/* Probability density is a normalized Gaussian times Ndat.  Normalization depends on the standard deviation but not on the mean. */

    int i;
    double Phi0,gamma,model;
    double Eavg,flux,dflux;
    double logl;

    Phi0=pow(10.0,z[0]);  /* Normalization at 100 GeV */
    gamma=z[1];           /* Power law slope: E^{-gamma} */
    
    logl=0.0;
    for (i=0; i<Ndat; i++)
    {
        Eavg=datax[i][0];
        flux=datax[i][1];
        dflux=sqrt(datax[i][2]*datax[i][2]+datax[i][3]*datax[i][3]);
        model=Phi0*pow((Eavg/100.0),-gamma);
        logl-=0.5*(model-flux)*(model-flux)/(dflux*dflux);
    }
    
    return logl;
}