#include "LH.h"
#include "fft.h"

static char answer1;
static char answer;





void init_termmode()
{
  FILE *input;
  char fn[]="term_mode.in";

  input=fopen(fn,"r");
  fscanf(input,"%c",&answer1);
  fclose(input);
  answer='y';
}

void continue_wish()
{
  if (answer1=='n'){
     printf("Do you wish to continue simulation (y or n)\n");
     answer=getchar();        
     while (!((answer=='y') || (answer=='n')))
         answer=getchar();
  } else
      answer='n';
}

int continue_sim()
{
  return answer=='y';
}

int index_orig(int i, int j)
{
  return i*Norig+j;
}

int index(int i, int j)
{
  return i*N+j;
}


void parmread(int *newrun, int *itr, double *scale, double *betascale, double *sigmascale, double *alphascale, int *SUBSAMPLE, int *INTERMED_PRINT, int *startitr, int *no_linpar, double *sigma, double *alpha, double *alpha_low, double *alpha_up, double *Hz, double *Hbeta)
{
  FILE *indfil;
  char fn[40];

  /*read some input parameters from LH.in */ 
  /* alpha_up is a prior parameter for alpha */
  sprintf(fn,"discr_def.in");
  indfil=fopen(fn,"r");
  fscanf(indfil,"%ld%ld%ld%ld",&Morig,&Norig,&M,&N);
  fclose(indfil);
  NMorig=Norig*Morig; NM=N*M;
  
  sprintf(fn,"LH.in");
  indfil=fopen(fn,"r");
  fscanf(indfil,"%d%d%lf%lf%lf%lf%d%d%lf%lf",newrun,itr,scale,betascale,sigmascale,alphascale,SUBSAMPLE,INTERMED_PRINT,Hz,Hbeta);
  fclose(indfil);
  
  if (*newrun){
    sprintf(fn,"parms.in");
    indfil=fopen(fn,"r");
    fscanf(indfil,"%d%lf%lf%lf%lf",no_linpar,sigma,alpha,alpha_low,alpha_up);
    fclose(indfil);  
    *startitr=0;  
  } else {
    sprintf(fn,"parms.out");
    indfil=fopen(fn,"r");
    fscanf(indfil,"%d%lf%lf%lf%lf",no_linpar,sigma,alpha,alpha_low,alpha_up);
    fclose(indfil); 
    indfil=fopen("startitr.out","r");
    fscanf(indfil,"%d",startitr);
    fclose(indfil);
  }

  
  
  if (*no_linpar>maxlinpar){
    printf("error: max no. of lin par is %d\n",maxlinpar);
    exit(1);
  }
  
  if ((M<2*(Morig-1)) || (N<2*(Norig-1))){
    printf("error: insufficient extension of grid\n");
    exit(1);
  }
  
}

void initz(fftstruct *z, fftstruct *S, double *obsdata, double *linpred, int *is_observed, double sigma, double divNM, double *cellareas){
  int l, j;
  char fn[40];
  FILE *outfile;
  
  for (l=0; l<Morig; l++)
    for (j=0; j<Norig; j++)
      if (is_observed[index_orig(l,j)])
	(*z).x[index(l,j)][0]=(log(0.001+obsdata[index_orig(l,j)])-log(cellareas[index_orig(l,j)])-linpred[index_orig(l,j)])/sigma;   
  
  transform(M,N,z,S);
  
  sprintf(fn,"z.ini");
  outfile=fopen(fn,"w");
  for (l=0; l<M; l++){
    for (j=0; j<N; j++)
      fprintf(outfile,"%2.2f ",(*z).x[index(l,j)][0]);
    fprintf(outfile,"\n");
  }
  fclose(outfile);
}

void calclinpred(double *beta, double *X, double *linpred, int no_linpar,int *is_observed, double *popdensity){
  int i, j, k, notleaveout;
  
  notleaveout=1;

  for (i=0; i<Morig; i++)
    for (j=0; j<Norig; j++){ 
      linpred[index_orig(i,j)]=0;
      for (k=0; k<no_linpar; k++){
	linpred[index_orig(i,j)]+=beta[k]*X[i*Norig*maxlinpar+j*maxlinpar+k];
      }
      
      if (popdensity[index_orig(i,j)]<=0){
	printf("nonpositive popdensity= %f\n",popdensity[index_orig(i,j)]);
	exit(1);
      }
      linpred[index_orig(i,j)]+=notleaveout*log(popdensity[index_orig(i,j)]);
    }
}

void initialize(fftstruct **y, fftstruct **z, fftstruct **gradz, fftstruct **gradzprop, fftstruct **zprop, fftstruct **yprop, Doublearray *expyprop, Doublearray *expy,  fftstruct **S, fftstruct **Sprop, Doublearray *meany, Doublearray *meanlambda, Doublearray *ssy, Doublearray *sslambda, Doublearray *obsdata, Doublearray *linpred, Doublearray *linpredprop, Doublearray *meanlinpred, int **is_observed, int newrun, Doublearray *X, Doublearray *beta, Doublearray *betaprop, int no_linpar, Doublearray *residuals, double **cellareas, double **popdensity)
{
  char fn[40];
  FILE *indfil;
  int j;

  /* printf("%d %d %d\n",sizeof(doublecomplex),sizeof(fftw_complex),sizeof(double)); */

  init_fftstruct(M,N,y);
  init_fftstruct(M,N,yprop);
  init_fftstruct(M,N,z); 
  init_fftstruct(M,N,zprop);
  init_fftstruct(M,N,gradz); 
  init_fftstruct(M,N,gradzprop);
  init_fftstruct(M,N,S); 
  init_fftstruct(M,N,Sprop);   

  *expyprop= (double *) calloc(NM,sizeof(double));
  *expy= (double *) calloc(NM,sizeof(double));
  *meany= (double *) calloc(NMorig,sizeof(double)); 
  *meanlambda= (double *) calloc(NMorig,sizeof(double)); 
  *ssy= (double *) calloc(NMorig,sizeof(double));
  *sslambda= (double *) calloc(NMorig,sizeof(double));
  *obsdata=(double *) calloc(NMorig,sizeof(double));
  *linpred=(double *) calloc(NMorig,sizeof(double));
  *linpredprop=(double *) calloc(NMorig,sizeof(double));
  *meanlinpred=(double *) calloc(NMorig,sizeof(double));
  *X=(double *) calloc(NMorig*maxlinpar,sizeof(double));
  *beta=(double *) calloc(maxlinpar,sizeof(double));
  *betaprop=(double *) calloc(maxlinpar,sizeof(double));
  *residuals= (double *) calloc(NMorig,sizeof(double));
  *is_observed= (int *) calloc(NMorig,sizeof(int));
  *cellareas=(double *) calloc(NMorig,sizeof(double));
  *popdensity=(double *) calloc(NMorig,sizeof(double));

  /* printvariables3(*y,*yprop,*z,*zprop,*gradz,*gradzprop,*S,*Sprop,*linpred,*linpredprop); */

  for (j=0; j<NM; j++)
    (**y).x[j][0]=(**y).x[j][1]=(**gradz).x[j][0]=(**gradz).x[j][1]=(**gradzprop).x[j][0]=(**gradzprop).x[j][1]=(**zprop).x[j][0]=(**zprop).x[j][1]=(**yprop).x[j][0]=(**yprop).x[j][1]=(*expyprop)[j]=(*expy)[j]=(**S).x[j][0]=(**S).x[j][1]=0;

  for (j=0; j<NMorig; j++)
    (*obsdata)[j]=(*is_observed)[j]=0;

  for (j=0; j<NMorig; j++)
    (*linpred)[j]=(*linpredprop)[j]=0;
 

  if (newrun){
    sprintf(fn,"beta.in");
    indfil=fopen(fn,"r");
    for (j=0; j<no_linpar; j++){
      fscanf(indfil,"%lf",&((*beta)[j]));
      (*betaprop)[j]=0;
    }
    fclose(indfil);
  } else {
    sprintf(fn,"beta.out");
    indfil=fopen(fn,"r");
    for (j=0; j<no_linpar; j++){
      fscanf(indfil,"%lf",&((*beta)[j]));
      (*betaprop)[j]=0;
    }
    fclose(indfil);
  }

  for (j=0; j<NMorig; j++)
    (*meany)[j]=(*meanlambda)[j]=(*ssy)[j]=(*sslambda)[j]=(*meanlinpred)[j]=(*residuals)[j]=0;
  
  if (newrun){
    for (j=0; j<NM; j++)
      (**z).x[j][0]=(**z).x[j][1]=0;
  } 
  else 
    { /* a previous run exists for which the final state was hidden in z.out */
      sprintf(fn,"z.out");
      indfil=fopen(fn,"r");
      for (j=0; j<NM; j++){
	fscanf(indfil,"%lf",&((**z).x[j][0]));
	(**z).x[j][1]=0;
      }
      fclose(indfil);
      
      /* means and ss from last run are read */
    
      sprintf(fn,"meany.out");
      indfil=fopen(fn,"r");
      for (j=0; j<NMorig; j++)
	fscanf(indfil,"%lf",&((*meany)[j]));
      fclose(indfil);
    
    
      sprintf(fn,"meanlambda.out");
      indfil=fopen(fn,"r");
      for (j=0; j<NMorig; j++)
	fscanf(indfil,"%lf",&((*meanlambda)[j]));
      fclose(indfil);
    
      sprintf(fn,"ssy.out");
      indfil=fopen(fn,"r");
      for (j=0; j<NMorig; j++)
	fscanf(indfil,"%lf",&((*ssy)[j]));
      fclose(indfil);
    
      sprintf(fn,"sslambda.out");
      indfil=fopen(fn,"r");
      for (j=0; j<NMorig; j++)
	fscanf(indfil,"%lf",&((*sslambda)[j]));
      fclose(indfil);        
    
      sprintf(fn,"meanlinpred.out");
      indfil=fopen(fn,"r");
      for (j=0; j<NMorig; j++)
	fscanf(indfil,"%lf",&((*meanlinpred)[j]));
      fclose(indfil);        
    
      sprintf(fn,"residuals.out");
      indfil=fopen(fn,"r");
      for (j=0; j<NMorig; j++)
	fscanf(indfil,"%lf",&((*residuals)[j]));
      fclose(indfil);
      
    } /* end else */
}

double mytrunc(double mean, double H){
  /* if (mean>H)  fprintf(acceptrates,"truncation: mean=%f\n",mean);*/
  return (mean>H) ? H : mean;
}


void betaproposal(double *beta, double *betaprop, double *X, double betascale, double *logq, double *logqprop, double *linpred, double *linpredprop, int *is_observed, double *logfbetprop, int no_linpar, double *obsdata, double sigma, fftstruct *y, double Hbeta, double *popdensity, double *cellareas)
{

  int i, j, q;
  double gradbeta[maxlinpar], gradbetaprop[maxlinpar];
  double u[maxlinpar];

  for (q=0; q<no_linpar; q++){
    gradbeta[q]=betaprop[q]=0; /* sum_obs=0; sum_expy=0; */
    for (i=0; i<Morig; i++)
      for (j=0; j<Norig; j++)
	if (is_observed[index_orig(i,j)])
	  gradbeta[q]+=trunc_dloglikelihood(obsdata[index_orig(i,j)],(*y).x[index(i,j)][0],sigma,linpred[index_orig(i,j)],Hbeta,cellareas[index_orig(i,j)])*X[i*Norig*maxlinpar+j*maxlinpar+q];
    betaprop[q]+=beta[q]+betascale*gradbeta[q]/2;
  }

  /* for (q=0; q<no_linpar; q++)
     printf("gradbeta %f\n",gradbeta[q]);
     printf("\n"); */
 
  for (i=0; i<(no_linpar+2)/2; i++){
    gauss_sim(&(u[2*i]),&(u[2*i+1])); 
    /* printf("%f %f\n",u[2*i],u[2*i+1]); */
  }

  *logq=0;
  for (q=0; q<no_linpar; q++){
    *logq-=pow(u[q],2)/2;
    betaprop[q]+=u[q]*sqrt(betascale);
  }

  calclinpred(betaprop,X,linpredprop,no_linpar,is_observed,popdensity);
 
  for (q=0; q<no_linpar; q++){
    gradbetaprop[q]=0;
    for (i=0; i<Morig; i++)
      for (j=0; j<Norig; j++)
	if (is_observed[index_orig(i,j)])
	  gradbetaprop[q]+trunc_dloglikelihood(obsdata[index_orig(i,j)],(*y).x[index(i,j)][0],sigma,linpredprop[index_orig(i,j)],Hbeta,cellareas[index_orig(i,j)])*X[i*Norig*maxlinpar+j*maxlinpar+q];
  } 

  *logqprop=0;
  for (q=0; q<no_linpar; q++){
    *logqprop-=pow(beta[q]-betaprop[q]-betascale*gradbetaprop[q]/2,2)/(2*betascale);
  }
 

  *logfbetprop=0;
  for (i=0; i<Morig; i++)
    for (j=0; j<Norig; j++)
      if (is_observed[index_orig(i,j)])
	*logfbetprop+=loglikelihood(obsdata[index_orig(i,j)],(*y).x[index(i,j)][0],sigma,linpredprop[index_orig(i,j)],cellareas[index_orig(i,j)]);
  
  /*     printf("u1 %f betaprop: %f sum_obsyprop: %f sum_expyprop: %f logfzprop: %f logqprop %f logq %f\n",u1,*betaprop,*sum_obsyprop,*sum_expyprop,*logfzprop,*logqprop,*logq); */

}


void sigmaproposal_rw(double sigma, double *sigmaprop, double sigmascale, double u2,  double *logq, double *logqprop,  double *linpred, int *is_observed, double *logfbetprop, fftstruct *y, double *obsdata, double *cellareas)
{
  int l, j;
  double sum_obsyprop, sum_expyprop;

  *sigmaprop=sigma*exp(u2*sqrt(sigmascale));
     
  if (*sigmaprop>0){
    sum_expyprop=sum_obsyprop=0;
    for (l=0; l<Morig; l++) 
      for (j=0; j<Norig; j++)
	if (is_observed[index_orig(l,j)]){
	  sum_obsyprop+=obsdata[index_orig(l,j)]*(log(cellareas[index_orig(l,j)])+linpred[index_orig(l,j)]+*sigmaprop*(*y).x[index(l,j)][0]);
	  sum_expyprop+=cellareas[index_orig(l,j)]*exp(linpred[index_orig(l,j)]+*sigmaprop*(*y).x[index(l,j)][0]);
	}
    *logqprop=-log(sigma);
    *logq=-log(*sigmaprop);
    *logfbetprop=sum_obsyprop-sum_expyprop;
  } 
  /*     printf("u2 %f sigmaprop: %f sum_obsyprop: %f sum_expyprop: %f logfzprop: %f logqprop %f logq %f sum_yexpyprop %f beta %f \n",u2,*sigmaprop,*sum_obsyprop,*sum_expyprop,*logfzprop,*logqprop,*logq,*sum_yexpyprop,beta); */
} 




void update_grad(fftstruct *y, fftstruct *gradz, fftstruct *S, fftstruct *z, double divNM, double *obsdata, int *is_observed, double sigma, double *linpred, double Hz, double *cellareas)
{
  int l, j;

  for (j=0; j<NM; j++){
    (*gradz).x[j][0]=(*gradz).x[j][1]=0; 
  }
  for (l=0; l<Morig; l++) 
    for (j=0; j<Norig; j++)
      if (is_observed[index_orig(l,j)])
	(*gradz).x[index(l,j)][0]=trunc_dloglikelihood(obsdata[index_orig(l,j)],(*y).x[index(l,j)][0],sigma,linpred[index_orig(l,j)],Hz,cellareas[index_orig(l,j)]); 
  
  transformstar(M,N,gradz,S);

  for (j=0; j<NM; j++)
    (*gradz).x[j][0]=sigma*(*gradz).x[j][0]-(*z).x[j][0];     

  /* (the gradient is now calculated) */
  /* printvariables(0,1,2,3,y,gradz,z);*/
}


void conddensity_alpha(fftstruct *y, fftstruct *S, double sigma, double *logfbetprop, double *obsdata, int *is_observed, fftstruct *z, double divNM, double *linpred, double *cellareas)
{
  int l, j;
  double sum_obsy, sum_expy;
    
  for (j=0; j<NM; j++){
    (*y).x[j][0]=(*z).x[j][0]; (*y).x[j][1]=0;
  }


  transformstar(M,N,y,S);


  /* ubet. Y har middelvaerdi 0 og varians 1 */

  sum_obsy=sum_expy=0;    
  /* check in Gaussian case 
     sums_obs=sums_y=0;    */
  for (l=0; l<Morig; l++) 
    for (j=0; j<Norig; j++)
      if (is_observed[index_orig(l,j)]){
	sum_obsy+=obsdata[index_orig(l,j)]*(log(cellareas[index_orig(l,j)])+linpred[index_orig(l,j)]+sigma*(*y).x[index(l,j)][0]);
	sum_expy+=cellareas[index_orig(l,j)]*exp(linpred[index_orig(l,j)]+sigma*(*y).x[index(l,j)][0]);
	/* sums_obs+=pow(obsdata[index_orig(l,j)],2); 
	   sums_y+=pow(linpred[index_orig(l,j)]+sigma*(*y).x[index(l,j)][0],2); */
      }

  *logfbetprop=sum_obsy-sum_expy;
 
 /* check in Gaussian case 
   *logfbetprop=-0.5*(sums_obs+sums_y-2**sum_obsy)/0.3;  */ 
  /* (the conditional density is now calculated) */
  
  /* printvariables(sigma,*logfz,*sum_obsy,*sum_expy,y,y,y); */
}
                        

void conddensity_grad(fftstruct *y, fftstruct *S, double sigma,  fftstruct *gradz, double *logfbet, double *ssz, double *obsdata, int *is_observed, fftstruct *z, double divNM, double *linpred, double Hz, double *cellareas)
{
  int l, j;
  double  sum_obsy, sum_expy;

  for (j=0; j<NM; j++){
    (*y).x[j][0]=(*z).x[j][0]; (*y).x[j][1]=0;
  }

  transformstar(M,N,y,S);

  for (j=0; j<NM; j++)
    (*gradz).x[j][0]=(*gradz).x[j][1]=0;
 
  /* ubet. Y har middelvaerdi 0 og varians 1 */
  sum_obsy=sum_expy=0;    
  /* check in Gaussian case 
     sums_obs=sums_y=0; */
  for (l=0; l<Morig; l++) 
    for (j=0; j<Norig; j++)
      if (is_observed[index_orig(l,j)]){
	sum_obsy+=obsdata[index_orig(l,j)]*(log(cellareas[index_orig(l,j)])+linpred[index_orig(l,j)]+sigma*(*y).x[index(l,j)][0]);
	sum_expy+=cellareas[index_orig(l,j)]*exp(linpred[index_orig(l,j)]+sigma*(*y).x[index(l,j)][0]);
	/* check in Gaussian case 
	   sums_obs+=pow(obsdata[index_orig(l,j)],2);
	   sums_y+=pow(linpred[index_orig(l,j)]+sigma*y[index(l,j)][0],2); */
      }

  for (l=0; l<Morig; l++) 
    for (j=0; j<Norig; j++)
      if (is_observed[index_orig(l,j)])
	(*gradz).x[index(l,j)][0]=trunc_dloglikelihood(obsdata[index_orig(l,j)],(*y).x[index(l,j)][0],sigma,linpred[index_orig(l,j)],Hz,cellareas[index_orig(l,j)]); 

  transformstar(M,N,gradz,S);

  for (j=0; j<NM; j++)
    (*gradz).x[j][0]=sigma*(*gradz).x[j][0]-(*z).x[j][0];   
  /* (the gradient is now calculated) */
  
  *ssz=0;
  for (j=0; j<NM; j++)
    *ssz-=(0.5*pow((*z).x[j][0],2));
  
  *logfbet=sum_obsy-sum_expy; 

  /* check in Gaussian case 
   *logfbet=-0.5*(sums_obs+sums_y-2**sum_obsy)/0.3; */
  
  /* (the conditional density is now calculated) */
  
  /* printvariables(sigma,*logfz,*sum_obsy,*sum_expy,y,gradz,z); */
}

/* Gaussian  
double corrfct(double alpha, double h){
    return exp(-pow(h/exp(alpha),2));
}
*/

 
double corrfct(double alpha, double h){
    return exp(-h/exp(alpha));
}


double logprior(double alpha, double alpha_up){
  return 1;
}
 
double logprior_sigma(double sigma){
  /* return -log(sigma)-0.00001/sigma;*/
  /* uniform on sigma  */
  return 0;
}


double trunc_dloglikelihood(double obsdata, double y, double sigma, double linpred, double H, double cellarea){
/* differentieret poisson likelihood,  log-link */
    return obsdata-mytrunc(cellarea*exp(linpred+sigma*y),H);
}

double loglikelihood(double obsdata, double y, double sigma, double linpred, double cellarea){
/* poisson likelihood  log-link */
    return obsdata*(log(cellarea)+linpred+sigma*y)-cellarea*exp(linpred+sigma*y);
}


int sqrteigen_values(fftstruct *S, double alpha)
{
  int i, j;
  int pos;
  
  pos=1;
  for (i=0; i<M; i++)
    for (j=0; j<N; j++){
      (*S).x[index(i,j)][0]=corrfct(alpha,geodist(i,j));
      (*S).x[index(i,j)][1]=0;
    }
  

  transformb(M,N,S);

  /* tager kvadratrod og udskriver sqrt(egenvaerdier) */
  for (i=0; i<M; i++){
    for (j=0; j<N; j++){
      if ((*S).x[index(i,j)][0]<0){
	/* printf("neg. eigenvalue= %f alpha=%f\n",(*S).x[index(i,j)][0],alpha);        
	if (S[index(i,j)][0]<-0.00000001) */
        if ((*S).x[index(i,j)][0]< -0.000001) 
	  pos=0;
	(*S).x[index(i,j)][0]=0;
      }
      (*S).x[index(i,j)][0]=sqrt((*S).x[index(i,j)][0]);
    }
  }
  return pos;
}

double min(int i, int j){
  if (i<j)
    return (double) i;
  else 
    return (double) j;
}



void gauss_sim(double *u1, double *u2)
{
  /* Gaussian proposal */
  double theta, E, R;
  
  theta=2*pi*ran();
  E=-log(ran());
  R=sqrt(2*E);
  
  *u1=(double) R*sin(theta);
  *u2=(double) R*cos(theta);
}

void addGauss(fftstruct *surf, double scale)
{
  /* Gaussian proposal */
  int i;
  double u1, u2;
  double sqrtscale;

  sqrtscale=sqrt(scale);

  if (NM%2==0)
    for (i=0; i<NM/2; i++){
      gauss_sim(&u1,&u2);
      (*surf).x[2*i][0]+=sqrtscale*u1;
      (*surf).x[2*i+1][0]+=sqrtscale*u2;
    }
  else 
    {    
      for (i=0; i<NM/2; i++){
	gauss_sim(&u1,&u2);
	(*surf).x[2*i][0]+=sqrtscale*u1;
	(*surf).x[2*i+1][0]+=sqrtscale*u2;
      }
      gauss_sim(&u1,&u2);
      (*surf).x[NM-1][0]+=sqrtscale*u1;
    }
}

void alphaproposal(double alpha, double *alphaprop, double alphascale, double *logq, double *logqprop){
  double u1, u2;
  
  gauss_sim(&u1,&u2);
  *alphaprop=alpha+u1*alphascale;
  *logq=*logqprop=0;
}

