#include <stdio.h>
#include <stdlib.h>
#include <getopt.h>

#include <math.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_multimin.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_cdf.h>
#include <gsl/gsl_linalg.h>

#define maxnd 10000
#define H 0.5

/*
  PROGRAM: simSFS
  AUTHOR: Athanasios Kousathanas
  VERSION: 1.0

  Generate allele frequency vectors (AFVs) using Wright-Fisher transition matrix methods and sample
  to obtain site frequency spectra (SFS). Code implements methods described by Keightley and Eyre-Walker(2007).

  Features:
  Demography: 1-epoch or 2-epoch population changes
  Selection: only single s

  Compilation:
  use makefile or compile as
  gcc -O2 -o simSFS simSFS.v1.0.c -lm -lgsl -lgslcblas -w
*/

////////////////////////////////////////
//calculate allele frequency change, delta(q)
////////////////////////////////////////
double selfn(double q, double s, double h)
{
  double num, denom;
  num = s*q*(1-q)*(q + h*(1-2*q));
  denom = 1 + 2.0*h*s*q*(1-q) + s*q*q;
  if (denom== 0) denom = .00000000000001;
  return( q + (num/denom));
}
////////////////////////////////////////
//initialise rows of matrix (used by setupmatrix)
////////////////////////////////////////
setuprow(double **a, double p, int row, int n){
  double z;         /*probability of zero failures*/
  double x, y, temp1, temp2;
  int j;

  temp1 = 1.0 - p;
  if (temp1 == 0.0)  temp1 = .000000000001;            /*prevent zero divide*/
  temp1 = p/temp1;       /*First prevent negative log*/
  if (temp1 <=0)   temp1 = .000000000001;
  x = log(temp1);    /*used for later multiplication of series*/
  temp2 = 1.0-p;
  if (temp2<=0)   temp2 = .000000000001;
  z = (double)n*log(temp2);
  if (z > 150)   z = 150;
  a[row][0] = exp(z);
  y = a[row][0];        /*add up for later*/
  for (j = 1;j<=n; j++)
    {

      z = z + log(((double)n + 1.0 - (double)j)/(double)j) + x;
      /*             cumulative binomial coeff.*/
      a[row][j] = exp(z);
      y = y + a[row][j];          /*add up*/
    }
  if (y == 0.0)   y = .0000000000001;
  for (j=0; j<=n; j++) a[row][j] = a[row][j]/y;
}
////////////////////////////////////////
//initialise transition matrix
////////////////////////////////////////
setupmatrix(double **a, int n1, int n2, double s, double h){
  int i;
  double p;

  for (i = 1; i<=n1 - 1; i++)             /*do variable rows*/
    {
      p = (double)i/(double)n1;
      p = selfn(p, s,h);
      setuprow(a, p, i, n2);
    }
  for (i = 0; i<=n2; i++)
    {              /*do constant edges*/
      a[0][i] = 0;
      a[n1][i] = 0;
    }
  a[0][0] = 1.0;                          /*do absorbing corners*/
  a[n1][n2] = 1.0;

}
////////////////////////////////////////
//Perform matrix inversion
////////////////////////////////////////
matrixinvert(int N, double **a,double *inverter){

  int i, j, lotkin_signum;

  gsl_matrix  *lotkin_a,*lotkin_inv;
  gsl_vector *x, *lotkin_b, *lotkin_x;

  gsl_permutation *lotkin_perm;

  /* allocate a, x, b */
  lotkin_a = gsl_matrix_alloc(N-1, N-1);
  lotkin_inv = gsl_matrix_alloc(N-1, N-1);
  x = gsl_vector_alloc(N-1);
  lotkin_b = gsl_vector_alloc(N-1);

  lotkin_x = gsl_vector_alloc(N-1);

  gsl_matrix *identity=gsl_matrix_alloc(N-1, N-1);

  for (i = 0; i<=N-2; i++)
    {

      for (j = 0; j<=N-2; j++)
	{
          gsl_matrix_set( identity, i, i,1);
	}

    }

  for (i = 0; i<=N-2; i++)
    {
      for (j = 0; j<=N-2; j++)
	{
	  gsl_matrix_set(lotkin_a,i,j,gsl_matrix_get(identity,i,j)-a[i+1][j+1]);
	}

    }

  /* LU decomposition and forward&backward substition */

  //gsl_blas_dgemv(CblasNoTrans, 1.0, lotkin_a, x, 0.0, lotkin_b);

  lotkin_perm = gsl_permutation_alloc(N-1);
  gsl_linalg_LU_decomp(lotkin_a, lotkin_perm, &lotkin_signum);
  gsl_linalg_LU_solve(lotkin_a, lotkin_perm, lotkin_b, lotkin_x);
  gsl_linalg_LU_invert (lotkin_a,lotkin_perm, lotkin_inv);

  /*apparently the inversion adds +1 to the first element of the vector
    so I substract it*/

  gsl_matrix_set(lotkin_inv,0,0,gsl_matrix_get(lotkin_inv,0,0)-1);

  for (i=0;i<=N-2;i++){

    inverter[i+1]=gsl_matrix_get(lotkin_inv,0,i);

  }

  gsl_matrix_free(lotkin_a);
  gsl_matrix_free(lotkin_inv);
  gsl_vector_free(x);
  gsl_vector_free(lotkin_b);
  gsl_vector_free(lotkin_x);
  gsl_matrix_free(identity);
  gsl_permutation_free(lotkin_perm);

}
////////////////////////////////////////
//perform TM iteration for t generations
////////////////////////////////////////
tmiterate(double **a,double *mut1,
	  int t, int n1, int n2,int decay,double *sumf){

  int k, i, j = 0;
  double z;

  double * steadystatefreq = (double*) calloc (maxnd+1, sizeof(double));
  double * mut2 = (double*) calloc (maxnd+1, sizeof(double));

  for (i=0; i<=n1; i++) steadystatefreq[i] = 0;
  for (k = 1; k<=t; k++)
    {                 /*t iterations*/

      /* Increment steadystatefreq. distribution*/
      for (i = 0; i<=n1; i++)
	{
	  steadystatefreq[i] = steadystatefreq[i] + mut1[i];
	}

      /* Perform matrix multiplication*/
      for (i = 0; i<=n2; i++)
	{
	  z = 0;
	  for (j = 0; j<=n1; j++)
	    {
	      z = z + a[j][i]*mut1[j];
	    }
	  mut2[i] = z;
	}

      /* Copy result of multiplication*/
      for (i=0; i<=n2; i++) {mut1[i] = mut2[i];

	if (decay==0){
	  sumf[i]+=mut1[i];}
	if(decay==1){sumf[i]=mut1[i];}
      }

    }

  free(steadystatefreq);
  free(mut2);
}

eqf_using_matrix_inversion(int n1,double s,double **a,double *egf_out){

  setupmatrix(a, n1, n1, s, H);
  matrixinvert(n1,a,egf_out);

}
////////////////////////////////////////
//scale AFV for mutations eliminated by selection/not mutated
////////////////////////////////////////

egf_scaling_s(int n1,double *v_s_in, double *v_s_out,double f0){
  int n1d=2*n1;
  double sumx,sumy;
  int i;
  sumx=0;
  sumy=0;

  for (i=1;i<n1d;i++){
    sumx+=v_s_in[i];
  }

  for (i=1;i<n1d;i++){
    v_s_out[i]/=sumx;
    sumy+=v_s_out[i];
  }

  v_s_out[0]=(1-sumy)*(1-f0);
}
////////////////////////////////////////
//scale AFV for sites not mutated
////////////////////////////////////////
egf_scaling_f0(int n1,double *fv,double f0){

  int i;
  int n1d=2*n1;
  for (i=1;i<n1d;i++){
    fv[i]*=(1-f0);
  }
  fv[0]+=f0;
}
////////////////////////////////////////
//binomial sampling from AFV to generate SFS
////////////////////////////////////////
binomial_sampling(int n1,int n1b,int sample, double *invy,int *discrete,gsl_rng *rgen){

  int s1,success;
  int i;

  double prob;

  /*
    based on equilibrium frequency vector (use it as a probability vector),
    I randomly generate numbers from 0 to n1d-1 for sample sites.
  */

  gsl_ran_discrete_t *r= gsl_ran_discrete_preproc (n1,invy);
  for (i=0;i<sample;i++){
    s1=gsl_ran_discrete (rgen, r);
    prob=(double)(s1)/n1;

    success=gsl_ran_binomial(rgen,prob,n1b+1);
    discrete[success]++;
  }
  gsl_ran_discrete_free(r);
}

output_sfs_to_file(int n1,int *sfs1,int *sfs2,char *filename,long seed){
  int i;
  FILE *file;

  file = fopen(filename,"w");
 fprintf(file,"#seed used: %d\n",seed);
 fprintf(file,"#neutral\n");
  for (i=0;i<=n1;i++){
    fprintf(file,"%d ",sfs1[i]);
  }
 fprintf(file,"\n#selected\n");

  for (i=0;i<=n1;i++){
    fprintf(file,"%d ",sfs2[i]);
  }

  fclose(file);
}
////////////////////////////////////////
//Calculate weighted average of two vectors
// Normally vectors w(s) x(s) by N1 and N2 respectively
////////////////////////////////////////

vector_average(int n1,int n2,double *fv1,double *fv2, double *fv_out){
  int i;
  int n1d=n1*2;
  int n2d=n2*2;
  for (i=0;i<=n2d;i++){

    fv_out[i]=(n1*fv1[i]+n2*fv2[i])/(n1+n2);

  }
}
//////////////////////////////////////
//Obtain AFV for 1-epoch model
////////////////////////////////////////
calculate_FV_one_epoch(int n1,double s,double *FV){
  int i=0;
  int n1d=2*n1;
  double **FVW = calloc(maxnd+1, sizeof(double *));

  for(i = 0; i < maxnd+1; i++){
    FVW[i] = calloc(maxnd+1,  sizeof(double));}
  eqf_using_matrix_inversion(n1d,s,FVW,FV);
  FV[1]+=1;
  for(i = 0; i < maxnd+1; i++)
    free(FVW[i]);
  free(FVW);
}
////////////////////////////////////////
//Obtain AFV for 2-epoch model
////////////////////////////////////////
calculate_FV_two_epoch(int n1,int n2,int t,double s,double *FV){
  int i=0;
  double **egf_0 = calloc(maxnd+1, sizeof(double *));
  for(i = 0; i < maxnd+1; i++){
    egf_0[i] = calloc(maxnd+1,  sizeof(double));
  }

  double * FVW = (double*) calloc (maxnd+1, sizeof(double));
  double * FVX = (double*) calloc (maxnd+1, sizeof(double));
  double * startingfreq = (double*) calloc (maxnd+1, sizeof(double));

  int n1d=n1*2;
  int n2d=n2*2;

  /*
    Contributions of mutations to AFV before the change in population size.
  */

  //obtain AFV at equilibrium

  eqf_using_matrix_inversion(n1d,s,egf_0,FVW);
  FVW[1]+=1.0;//cumulative vector includes 1st mutant
  //perform change in population size N1->N2
  setupmatrix(egf_0,n1d,n2d,s,H);
  tmiterate(egf_0,FVW,1, n1d, n2d,1,FVW);
  setupmatrix(egf_0,n2d,n2d,s,H);
  tmiterate(egf_0,FVW,t-1, n2d, n2d,1,FVW);

  //dumpvector(FVW,0,n2,"FVW");
  /*
    Contributions of mutations to AFV after the change in population size.
  */
  setupmatrix(egf_0,n2d,n2d,s,H);
  startingfreq[1]=1.0;
  tmiterate(egf_0,startingfreq, t-1, n2d, n2d,0,FVX);
  FVX[1]+=1.0;//cumulative vector includes 1st mutant
  //dumpvector(FVX,0,n2,"FVX");

  /*
    averaging of vectors FVW (N1) and FVX(N2)
  */
  vector_average(n1,n2,FVW,FVX,FV);

  for(i = 0; i < maxnd+1; i++){
    free(egf_0[i]);
  }

  free(egf_0);
  free(FVX);
  free(FVW);
  free(startingfreq);
}
////////////////////////////////////////
//print given matrix
////////////////////////////////////////

dumpmatrix(double **m, int s1,
	   int rows, int s2, int cols, char *s){
  int i, j;
  printf("\n%s\n", s);
  for (i = s1; i<=rows; i++)
    {
      for (j = s2; j<=cols; j++)
	{
	  printf(" %2.3f", m[i][j]);
	}
      printf("\n");
    }
  printf("\n");
}
////////////////////////////////////////
//print given vector
////////////////////////////////////////
dumpvector(double *v, int min, int max, char *s){
  int i, control = 0;
  //   printf("\n");
  printf("%s", s); printf("\n");
  for (i = min; i<=max; i++)
    {
      printf("%3d %lf ", i, v[i]);
      control++;
      if (control==5)
	{
	  control = 0;
	  printf("\n");
	}
    }
  printf("\n");
}

/*
////////////////////////////////////////
MAIN START
////////////////////////////////////////
*/

int main(int argc, char **argv)
{

  //initialise input variables and use default values
  int N1=100,n=10,sampleS=1000,sampleN=1000;
  int N2=N1;
  int t=N1;
  double f0=0.9,S=0;
  char outfile[maxnd+1]="sfs.out";
  long seed=time (NULL) * getpid();

  //other variables
  int option_index = 0;
  int c=0;
  static int verbose_flag;

  const gsl_rng_type * T;
  T = gsl_rng_taus;
  gsl_rng  *rgen =gsl_rng_alloc(T);
  gsl_rng_set(rgen,seed);

  int * discrete0 = (int*) calloc (maxnd+1, sizeof(int));
  int * discrete1 = (int*) calloc (maxnd+1, sizeof(int));

  double * FV0 = (double*) calloc (maxnd+1, sizeof(double));
  double * FVS = (double*) calloc (maxnd+1, sizeof(double));

  //parse input
  while (1)
    {

      static struct option long_options[] =
	{
	  {"h",no_argument,  &verbose_flag, 1},
	  {"N1",     required_argument,0, 'a'},
	  {"N2",  required_argument,0, 'b'},
	  {"n",  required_argument, 0, 'c'},
	  {"f0",  required_argument, 0, 'd'},
	  {"s",    required_argument, 0, 'e'},
	  {"LS",    required_argument, 0, 'f'},
	  {"LN",    required_argument, 0, 'g'},
	  {"t",    required_argument, 0, 'h'},
	  {"o",    required_argument, 0, 'i'},
	  {"seed",    required_argument, 0, 'j'},
	  {NULL, 0, 0, 0}
	};
      /* getopt_long stores the option index here. */
      c = getopt_long_only (argc, argv, "",
			    long_options, &option_index);

      /* Detect the end of the options. */
      if (c == -1)
	break;

      switch (c)
	{
	case 0:
	  /* If this option set a flag, do nothing else now. */
	  if (long_options[option_index].flag != 0)
	    {
	      printf("\
option\tdefault_value\tdescription\n\
N1\t100\tsize N1\n\
N2\t100\tsize N2\n\
t\t100\ttime since size change\n\
n\t10\tallele sample size\n\
f0\t0.9\t1-f0 is proportional to mutation rate\n\
s\t0\tselection coeffient\n\
LS\t1000\tno. neutral sites\n\
LN\t1000\tno. selected sites\n\
o\tsfs.out\toutput file\n\
seed\ttime*pid\tset seed for random generator\n\
");
	    }
	  break;
	  printf ("option %s", long_options[option_index].name);
	  if (optarg)
	    printf (" with arg %s", optarg);
	  printf ("\n");
	  break;

	case 'a':
	  N1=atoi(optarg);
	  break;
	case 'b':
	  N2=atoi(optarg);
	  break;
	case 'c':
	  n=atoi(optarg);
	  break;
	case 'd':
	  f0=atof(optarg);
	  break;
	case 'e':
	  S=atof(optarg);
	  break;
	case 'f':
	  sampleS=atoi(optarg);
	  break;
	case 'g':
	  sampleN=atoi(optarg);
	  break;
	case 'h':
	  t=atoi(optarg);
	  break;
	case 'i':
	  strncpy(outfile,optarg,maxnd+1);
	  break;
	case 'j':
	  seed=atoi(optarg);
	  gsl_rng_set(rgen,seed);
	  break;
	case '?':
	  abort();
	  break;
	default:
	  abort ();
	}
    }

  if (optind < argc)
    {
      printf ("non-option ARGV-elements: ");
      while (optind < argc)
	printf ("%s ", argv[optind++]);
      putchar ('\n');
    }

  if (N1>1000||N2>1000){printf ("Use N1,N2 <1000\n"); abort();}
  /*
    calculate allele frequency vectors for neutral (FV0)
    and selected sites (FVS)
  */
  if (N1==N2){
    calculate_FV_one_epoch(N1,0.0,FV0);
    calculate_FV_one_epoch(N1,S,FVS);
  }else{
    calculate_FV_two_epoch(N1,N2,t,0.0,FV0);
    calculate_FV_two_epoch(N1,N2,t,S,FVS);
  }

  /*
    scale FVS to include the frequency of sites at which mutations
    have been eliminated by selection or have not experienced a mutation
  */
  egf_scaling_s(N2,FV0,FVS, f0);
  egf_scaling_f0(N2,FVS,f0);

  /*
    scale FV0 to account for sites that have never experienced a mutation.
  */
  egf_scaling_s(N2,FV0, FV0,f0);
  egf_scaling_f0(N2,FV0,f0);

  /*sampling
    sample sampleN neutral sites from FV0
    sample sampleS selected sites from FVS
  */

  binomial_sampling(N2,n,sampleN,FV0,discrete0,rgen);
  binomial_sampling(N2,n,sampleS, FVS,discrete1,rgen);

  output_sfs_to_file(n,discrete0,discrete1,outfile,seed);

  //free vectors
  free(FV0);
  free(FVS);

  free(discrete0);
  free(discrete1);

  gsl_rng_free (rgen);

  return 0;
}

/*
////////////////////////////////////////
END OF MAIN
////////////////////////////////////////
*/


/* cheb/eval.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_chebyshev.h>

/* For efficiency there are separate implementations of each of these
   functions */

double
gsl_cheb_eval (const gsl_cheb_series * cs, const double x)
{
  size_t i;
  double d1 = 0.0;
  double d2 = 0.0;

  double y = (2.0 * x - cs->a - cs->b) / (cs->b - cs->a);
  double y2 = 2.0 * y;

  for (i = cs->order; i >= 1; i--)
    {
      double temp = d1;
      d1 = y2 * d1 - d2 + cs->c[i];
      d2 = temp;
    }

  return y * d1 - d2 + 0.5 * cs->c[0];
}

double
gsl_cheb_eval_n (const gsl_cheb_series * cs, const size_t n, const double x)
{
  size_t i;
  double d1 = 0.0;
  double d2 = 0.0;

  size_t eval_order = GSL_MIN (n, cs->order);

  double y = (2.0 * x - cs->a - cs->b) / (cs->b - cs->a);
  double y2 = 2.0 * y;

  for (i = eval_order; i >= 1; i--)
    {
      double temp = d1;
      d1 = y2 * d1 - d2 + cs->c[i];
      d2 = temp;
    }

  return y * d1 - d2 + 0.5 * cs->c[0];
}

int
gsl_cheb_eval_err (const gsl_cheb_series * cs, const double x,
                   double *result, double *abserr)
{
  size_t i;
  double d1 = 0.0;
  double d2 = 0.0;

  double y = (2. * x - cs->a - cs->b) / (cs->b - cs->a);
  double y2 = 2.0 * y;

  double absc = 0.0;

  for (i = cs->order; i >= 1; i--)
    {
      double temp = d1;
      d1 = y2 * d1 - d2 + cs->c[i];
      d2 = temp;
    }

  *result = y * d1 - d2 + 0.5 * cs->c[0];

  /* Estimate cumulative numerical error */

  for (i = 0; i <= cs->order; i++)
    {
      absc += fabs(cs->c[i]);
    }

  /* Combine truncation error and numerical error */

  *abserr = fabs (cs->c[cs->order]) + absc * GSL_DBL_EPSILON;

  return GSL_SUCCESS;
}

int
gsl_cheb_eval_n_err (const gsl_cheb_series * cs,
                     const size_t n, const double x,
                     double *result, double *abserr)
{
  size_t i;
  double d1 = 0.0;
  double d2 = 0.0;

  double y = (2. * x - cs->a - cs->b) / (cs->b - cs->a);
  double y2 = 2.0 * y;

  double absc = 0.0;

  size_t eval_order = GSL_MIN (n, cs->order);

  for (i = eval_order; i >= 1; i--)
    {
      double temp = d1;
      d1 = y2 * d1 - d2 + cs->c[i];
      d2 = temp;
    }

  *result = y * d1 - d2 + 0.5 * cs->c[0];

  /* Estimate cumulative numerical error */

  for (i = 0; i <= eval_order; i++)
    {
      absc += fabs(cs->c[i]);
    }

  /* Combine truncation error and numerical error */

  *abserr = fabs (cs->c[eval_order]) + absc * GSL_DBL_EPSILON;

  return GSL_SUCCESS;
}

int
gsl_cheb_eval_mode_e (const gsl_cheb_series * cs,
                      const double x, gsl_mode_t mode,
                      double *result, double *abserr)
{
  size_t i;
  double d1 = 0.0;
  double d2 = 0.0;

  double y = (2. * x - cs->a - cs->b) / (cs->b - cs->a);
  double y2 = 2.0 * y;

  double absc = 0.0;

  size_t eval_order;

  if (GSL_MODE_PREC (mode) == GSL_PREC_DOUBLE)
    eval_order = cs->order;
  else
    eval_order = cs->order_sp;

  for (i = eval_order; i >= 1; i--)
    {
      double temp = d1;
      d1 = y2 * d1 - d2 + cs->c[i];
      d2 = temp;
    }

  *result = y * d1 - d2 + 0.5 * cs->c[0];

  /* Estimate cumulative numerical error */

  for (i = 0; i <= eval_order; i++)
    {
      absc += fabs(cs->c[i]);
    }

  /* Combine truncation error and numerical error */

  *abserr = fabs (cs->c[eval_order]) + absc * GSL_DBL_EPSILON;

  return GSL_SUCCESS;
}

double
gsl_cheb_eval_mode (const gsl_cheb_series * cs,
                    const double x, gsl_mode_t mode)
{
  double result, abserr;
  int status = gsl_cheb_eval_mode_e (cs, x, mode, &result, &abserr);

  if (status != GSL_SUCCESS)
    {
      GSL_ERROR_VAL("gsl_cheb_eval_mode", status, result);
    };

  return result;
}


#define _XOPEN_SOURCE 600

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <math.h>
#include <time.h>

void
setupInitialSeed(
		 int id,
		 gsl_rng *rgen)
{

  struct timespec spec;

  clock_gettime( CLOCK_REALTIME, &spec );

  gsl_rng_set( rgen, (unsigned long int) ( round( spec.tv_nsec / 1.0e6 ) + 20 * id ) );

}

void createInitialSample(
			 double *x,
			 long N,
			 gsl_rng *rgen,
			 int state_dim )
{

  for ( long i = 0; i < N; i++)
    for ( int d = 0; d < state_dim; d++ )
      x[ state_dim * i + d ] = gsl_ran_gaussian_ziggurat( rgen, 1.0 );

}


/*This program evaluates the differential equations for a photon's geodesic in a perturbed spacetime in SPHERICAL coordinates and restricted to RADIAL MOTION with a Runge-Kutta-Fehlberg 45 method.*/

/*This program solves the particular case for flat FRW perturbed spacetime with metric: $g_{ab} = {[g]}_{ab} + h_{ab}$. Where ${[g]}_{ab}$ corresponds to the flat FRW metric and $h_{ab}$ corresponds to the perturbation in the conformal Newtonian gauge. A Plummer potential with adequate parameters have been used to simulate the perturbation.
The equations are written in the form $\frac{d(x or p)^{\alpha}}{d\lambda}=f(x^{\alpha},p^{\beta})$ and the indices $\alpha$ and $\beta$ run from 0 to 1 since the motion is only radial.
Where $p^{\alpha}={\dot{x}}^{\alpha}$.
The coordinates for the photon's geodesics are then: (ct,r) = (x0,x1).*/

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <gsl/gsl_errno.h>          //GSL error management module
#include <gsl/gsl_spline.h>         //GSL interpolation module

/* PHYSICAL CONSTANTS AND PARAMETERS */
#define A 10.0     //Distance parameter of the perturbations
#define G 43007.01     //Gravitational constant
#define M 0.0     //Mass of the perturbation
#define C 299792.458  //Speed of light

/* PROGRAM PARAMETERS (ONLY USING NLINES FRW) */
#define NLINES 100000 //Number of lines in geodesic_solution.dat file
#define NSTEPS 75000000 //Number of steps for solving geodesics
#define NLINESFRW 10000 //Number of lines in frw.dat file
#define DLAMBDA 0.01   //Geodesics parameter step

//We define mydbl as a generic variable type which can be changed when needed to: float, double, long double
typedef long double mydbl;

/* RUNGE-KUTTA-FEHLBERG METHOD CONSTANTS */
const mydbl c2 = 0.25, c3 = 3.0/8.0, c4 = 12.0/13.0, c5 = 1.0, c6 = 0.5; //Multiplying increment in t
const mydbl a21 = 0.25, a31 = 3.0/32.0, a32 = 9.0/32.0, a41 = 1932.0/2197.0, a42 = -7200.0/2197.0, a43 = 7296.0/2197.0, a51 = 439.0/216.0, a52 = -8.0, a53 = 3680.0/513.0, a54 = -845.0/4104.0, a61 = -8.0/27.0, a62 = 2.0, a63 = -3544.0/2565.0, a64 = 1859.0/4104.0, a65 = -11.0/40.0; //This adds to the approximated solution in step i. First script means K value and second which K is multiplying
const mydbl r1 = 1.0/360.0, r3 = -128.0/4275.0, r4 = -2197.0/75240.0, r5 = 1.0/50.0, r6 = 2.0/55.0; //Coefficients to calculate R
const mydbl b1 = 25.0/216.0, b3 = 1408.0/2565.0, b4 = 2197.0/4104.0, b5 = -1.0/5.0;

/* FUNCTIONS OF THE PROGRAM */

/*Find and return the maximum value for objects in a vector*/
mydbl maxValue(mydbl arr[], size_t size)
{
  size_t i;
  mydbl maxValue = arr[0];

  for(i = 1; i < size; i++)
    {
      if(arr[i] > maxValue)
	{
	  maxValue = arr[i];
	}
    }
  return maxValue;
}

/*Interpolation of function inside *spline object evaluated at an abscisa 'x'.
Argument *spline is a pointer to a spline object which stores the type of interpolation to be made. x is the independent variable where the function is evaluated. *acc is a pointer to a lookup object for interpolations.*/
double interpolator(gsl_spline *spline, double x, gsl_interp_accel *acc)
{
  double a = gsl_spline_eval(spline, x, acc);  //Interpolates data to abcisa x using method in spline and acceleration object acc

  return a;  //Return value of interpolated function at x
}

/*Function for the gravitational potential to be used. Potential for Plummer model.*/
mydbl potential(mydbl r)
{
  return -G*M/(sqrtl(A*A + r*r));
}

/*Derivative of potential respecto to radial coordinate.*/
mydbl der_potential(mydbl r)
{
  return G*M*r/(powl(A*A+r*r, 1.5));
}

/*Function of the 0th momentum component differential equation for the geodesics.
${p0}^{dot} = f0(x^{\alpha},p^{\alpha})$.*/
mydbl geodesic_equation_0(gsl_spline *spline1, gsl_interp_accel *acc1, gsl_spline *spline2, gsl_interp_accel *acc2, mydbl p0, mydbl pr, mydbl x0, mydbl r)
{
  double t = (double)(1.0*x0/C);
  mydbl a = (mydbl) 1.0*interpolator(spline1, t, acc1);
  mydbl adot = (mydbl) 1.0*interpolator(spline2, t, acc2);
  mydbl f = -2.0*der_potential(r)*p0*pr/(C*C + 2.0*potential(r)) - (1.0 - 2.0*potential(r)/(C*C))*(a*adot)*(pr*pr)/(C + 2.0*potential(r)/C);
  return f;
}

/*Function of the 1th (radial) momentum component differential equation for the geodesics.
${p1}^{dot} = f1(x^{\alpha},p^{\alpha})$.*/
mydbl geodesic_equation_r(gsl_spline *spline1, gsl_interp_accel *acc1, gsl_spline *spline2, gsl_interp_accel *acc2, mydbl p0, mydbl pr, mydbl x0, mydbl r)
{
  double t = (double)(1.0*x0/C);
  mydbl a = (mydbl) 1.0*interpolator(spline1, t, acc1);
  mydbl adot = (mydbl) 1.0*interpolator(spline2, t, acc2);
  mydbl f = - (der_potential(r)*p0*p0)/(a*a*(C*C - 2.0*potential(r))) - (2.0*adot*p0*pr)/(C*a) + der_potential(r)*(pr*pr)/(C*C - 2.0*potential(r));
  return f;
}

/*To set the initial value of pr, it must hold $g_{\mu\nu}p^{\mu}p^{\nu} = 0$.
This factor multiplies p0 to guarantee that p1 fulfill the null geodesic condition.*/
mydbl condition_factor(mydbl r, double a)
{
  return (mydbl)(1.0/a)*sqrtl((1.0+2.0*potential(r)/(C*C))/(1.0 - 2.0*potential(r)/(C*C)));
}

/*$cp^{0}$ multiplied by this factor allows to obtain the energy for a local inertial observer in this spacetime.*/
mydbl energy_factor(mydbl r)
{
  mydbl g = sqrtl(1.0 + 2.0*potential(r)/(C*C));
  return g;
}

/*Violation of null geodesics condition $g_{\mu\nu}p^{\mu}p^{\nu} = 0$.*/
mydbl violation(mydbl r, mydbl p0, mydbl pr, double a)
{
  mydbl f = -(1.0+2.0*potential(r)/(C*C))*p0*p0 + (mydbl)(1.0*a*a)*(1.0-2.0*potential(r)/(C*C))*pr*pr;
  return f;
}

/*Function for solving the geodesic's differential equations using a Runge-Kutta-Fehlber method with adaptative stepsize.
Arguments are pointer so variables in that memory addresses are changed every time this function is called. There is no need to define auxiliary variables to store the results of computations.*/
void runge_kutta_fehlberg(gsl_spline *spline1, gsl_interp_accel *acc1, gsl_spline *spline2, gsl_interp_accel *acc2, mydbl *x0, mydbl *x1, mydbl *p0, mydbl *p1, mydbl *lambda, mydbl *dlambda, mydbl lambda_endpoint, mydbl tol, mydbl dlambda_max, mydbl dlambda_min, FILE *fp, double aem, mydbl energy1)
{
  /*Define the k quantities for all the variables*/
  mydbl k1x0, k1x1, k1p0, k1p1;
  mydbl k2x0, k2x1, k2p0, k2p1;
  mydbl k3x0, k3x1, k3p0, k3p1;
  mydbl k4x0, k4x1, k4p0, k4p1;
  mydbl k5x0, k5x1, k5p0, k5p1;
  mydbl k6x0, k6x1, k6p0, k6p1;

  /*Define a vector that carries the errors for each differential equation and a variable for storing the biggest error*/
  mydbl r[4], rmax;

  /*Define the increments for the values to be solved in the differential equations*/
  mydbl dp0, dp1, dx0, dx1;

  /*Define some relevant quantities to be printed in a file*/
  mydbl energy, difft, t, v, difference;
  double aobs, difftfrw;

  /*Relevant quantity to determine the new stepsize*/
  mydbl delta;

  while(*lambda < lambda_endpoint)
    {
      /*Check if the independent variable plus the actual stepsize is still in the desired interval*/
      if((*lambda + *dlambda) > lambda_endpoint)
	{
	  *dlambda = lambda_endpoint - *lambda;
	}

      /*Calculate the k quantities for all the variables*/
      /*This section calculates the k1 quantities*/
      k1x0 = (*dlambda)*(*p0); k1x1 = (*dlambda)*(*p1);
      k1p0 = *dlambda*geodesic_equation_0(spline1, acc1, spline2, acc2, *p0, *p1, *x0, *x1);
      k1p1 = *dlambda*geodesic_equation_r(spline1, acc1, spline2, acc2, *p0, *p1, *x0, *x1);

      /*This section calculates the k2 quantities*/
      k2x0 = *dlambda*(*p0 + a21*k1p0); k2x1 = *dlambda*(*p1 + a21*k1p1);
      k2p0 = *dlambda*geodesic_equation_0(spline1, acc1, spline2, acc2, *p0 + a21*k1p0, *p1 + a21*k1p1, *x0 + a21*k1x0, *x1 + a21*k1x1);
      k2p1 = *dlambda*geodesic_equation_r(spline1, acc1, spline2, acc2, *p0 + a21*k1p0, *p1 + a21*k1p1, *x0 + a21*k1x0, *x1 + a21*k1x1);

      /*This section calculates the k3 quantities*/
      k3x0 = *dlambda*(*p0 + a31*k1p0 + a32*k2p0); k3x1 = *dlambda*(*p1 + a31*k1p1 + a32*k2p1);
      k3p0 = *dlambda*geodesic_equation_0(spline1, acc1, spline2, acc2, *p0 + a31*k1p0 + a32*k2p0, *p1 + a31*k1p1 + a32*k2p1, *x0 + a31*k1x0 + a32*k2x0, *x1 + a31*k1x1 + a32*k2x1);
      k3p1 = *dlambda*geodesic_equation_r(spline1, acc1, spline2, acc2, *p0 + a31*k1p0 + a32*k2p0, *p1 + a31*k1p1 + a32*k2p1, *x0 + a31*k1x0 + a32*k2x0, *x1 + a31*k1x1 + a32*k2x1);

      /*This section calculates the k4 quantities*/
      k4x0 = *dlambda*(*p0 + a41*k1p0 + a42*k2p0 + a43*k3p0); k4x1 = *dlambda*(*p1 + a41*k1p1 + a42*k2p1 + a43*k3p1);
      k4p0 = *dlambda*geodesic_equation_0(spline1, acc1, spline2, acc2, *p0 + a41*k1p0 + a42*k2p0 + a43*k3p0, *p1 + a41*k1p1 + a42*k2p1 + a43*k3p1, *x0 + a41*k1x0 + a42*k2x0 + a43*k3x0, *x1 + a41*k1x1 + a42*k2x1 + a43*k3x1);
      k4p1 = *dlambda*geodesic_equation_r(spline1, acc1, spline2, acc2, *p0 + a41*k1p0 + a42*k2p0 + a43*k3p0, *p1 + a41*k1p1 + a42*k2p1 + a43*k3p1, *x0 + a41*k1x0 + a42*k2x0 + a43*k3x0, *x1 + a41*k1x1 + a42*k2x1 + a43*k3x1);

      /*This section calculates the k5 quantities*/
      k5x0 = *dlambda*(*p0 + a51*k1p0 + a52*k2p0 + a53*k3p0 + a54*k4p0); k5x1 = *dlambda*(*p1 + a51*k1p1 + a52*k2p1 + a53*k3p1 + a54*k4p1);
      k5p0 = *dlambda*geodesic_equation_0(spline1, acc1, spline2, acc2, *p0 + a51*k1p0 + a52*k2p0 + a53*k3p0 + a54*k4p0, *p1 + a51*k1p1 + a52*k2p1 + a53*k3p1 + a54*k4p1, *x0 + a51*k1x0 + a52*k2x0 + a53*k3x0 + a54*k4x0, *x1 + a51*k1x1 + a52*k2x1 + a53*k3x1 + a54*k4x1);
      k5p1 = *dlambda*geodesic_equation_r(spline1, acc1, spline2, acc2, *p0 + a51*k1p0 + a52*k2p0 + a53*k3p0 + a54*k4p0, *p1 + a51*k1p1 + a52*k2p1 + a53*k3p1 + a54*k4p1, *x0 + a51*k1x0 + a52*k2x0 + a53*k3x0 + a54*k4x0, *x1 + a51*k1x1 + a52*k2x1 + a53*k3x1 + a54*k4x1);

      /*This section calculates the k6 quantities*/
      k6x0 = *dlambda*(*p0 + a61*k1p0 + a62*k2p0 + a63*k3p0 + a64*k4p0 + a65*k5p0); k6x1 = *dlambda*(*p1 + a61*k1p1 + a62*k2p1 + a63*k3p1 + a64*k4p1 + a65*k5p1);
      k6p0 = *dlambda*geodesic_equation_0(spline1, acc1, spline2, acc2, *p0 + a61*k1p0 + a62*k2p0 + a63*k3p0 + a64*k4p0 + a65*k5p0, *p1 + a61*k1p1 + a62*k2p1 + a63*k3p1 + a64*k4p1 + a65*k5p1, *x0 + a61*k1x0 + a62*k2x0 + a63*k3x0 + a64*k4x0 + a65*k5x0, *x1 + a61*k1x1 + a62*k2x1 + a63*k3x1 + a64*k4x1 + a65*k5x1);
      k6p1 = *dlambda*geodesic_equation_r(spline1, acc1, spline2, acc2, *p0 + a61*k1p0 + a62*k2p0 + a63*k3p0 + a64*k4p0 + a65*k5p0, *p1 + a61*k1p1 + a62*k2p1 + a63*k3p1 + a64*k4p1 + a65*k5p1, *x0 + a61*k1x0 + a62*k2x0 + a63*k3x0 + a64*k4x0 + a65*k5x0, *x1 + a61*k1x1 + a62*k2x1 + a63*k3x1 + a64*k4x1 + a65*k5x1);

      /*Determination of the biggest error between the errors for each differential equation*/
      r[0] = fabsl(r1*k1p0 + r3*k3p0 + r4*k4p0 + r5*k5p0 + r6*k6p0)/(*dlambda);
      r[1] = fabsl(r1*k1p1 + r3*k3p1 + r4*k4p1 + r5*k5p1 + r6*k6p1)/(*dlambda);
      r[2] = fabsl(r1*k1x0 + r3*k3x0 + r4*k4x0 + r5*k5x0 + r6*k6x0)/(*dlambda);
      r[3] = fabsl(r1*k1x1 + r3*k3x1 + r4*k4x1 + r5*k5x1 + r6*k6x1)/(*dlambda);
      rmax = maxValue(r, 4);

      /*Check whether the worst approximation is less than the introduced tolerance*/
      if(rmax <= tol)
	{
	  *lambda = *lambda + *dlambda;

	  dp0 = b1*k1p0 + b3*k3p0 + b4*k4p0 + b5*k5p0;
	  dp1 = b1*k1p1 + b3*k3p1 + b4*k4p1 + b5*k5p1;
	  dx0 = b1*k1x0 + b3*k3x0 + b4*k4x0 + b5*k5x0;
	  dx1 = b1*k1x1 + b3*k3x1 + b4*k4x1 + b5*k5x1;
	  *p0 = *p0 + dp0; *p1 = *p1 + dp1; *x0 = *x0 + dx0; *x1 = *x1 + dx1;

	  /*Relevant quantities to save in the .dat file. These only are calculated if the approximation is accepted*/
	  energy = C*energy_factor(*x1)*(*p0);
  	  difft = (energy - energy1)/energy1;
  	  t = *x0/C;
  	  aobs = interpolator(spline1, (double)(1.0*t), acc1);
  	  v = violation(*x1, *p0, *p1, aobs);
  	  difftfrw = (aem/aobs) - 1.0;
  	  difference = difft - (mydbl)(1.0*difftfrw);
  	  fprintf(fp, "%16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8e %16.8Le %16.8Le\n", *lambda, *dlambda, *x0, *x1, *p0, *p1, energy, v, difft, difftfrw, difference, rmax);
	}

      /*Calculate a new stepsize*/
      delta = 0.84*powl(tol/rmax, 0.25);

      if(delta <= 0.1)
	{
	  *dlambda = *dlambda*0.1;
	}
      else if(delta >= 4.0)
	{
	  *dlambda = *dlambda*4.0;
	}
      else
	{
	  *dlambda = *dlambda*delta;
	}

      /*Check whether the stepsize is too big or too small*/
      if(*dlambda > dlambda_max)
	{
	  *dlambda = dlambda_max;
	}
      else if(*dlambda < dlambda_min)
	{
	  printf("Minimum value for stepsize dlambda exceeded.\n");
	  break;
	}
    }
}

int main(void)
{
  /***READ SCALE FACTOR DATA AND PREPARE OBJECTS FOR INTERPOLATION ***/
  int i;            //For array manipulation

  /*Pointer to scale_factor.data file*/
  FILE *frw;
  frw = fopen("../scale_factor.dat","r");

  /*Variables and arrays to read the data*/
  double cosmictime[NLINESFRW], conftime, scale_factor[NLINESFRW], der_scale_factor[NLINESFRW];

  /*Reading the data*/
  for(i=0; i<NLINESFRW; i++)
    {
      fscanf(frw,"%lf %lf %lf %lf", &cosmictime[i], &conftime, &scale_factor[i], &der_scale_factor[i]);
    }

  /*Free space in memory*/
  fclose(frw);

  /*** Initializes objects for interpolation. 1 is for interpolation of scale factor, 2 is for interpolation of derivative of scale factor ***/

  /*Allocate space in memory*/
  gsl_interp_accel *acc1 = gsl_interp_accel_alloc();  //Acceleration type object (for index lookup)
  gsl_interp_accel *acc2 = gsl_interp_accel_alloc();
  gsl_spline *spline1 = gsl_spline_alloc(gsl_interp_cspline, NLINESFRW);  //Spline type object (define interpolation type and space in memory, works for both)
  gsl_spline *spline2 = gsl_spline_alloc(gsl_interp_cspline, NLINESFRW);

  /*Initializes objects for interpolation*/
  gsl_spline_init(spline1, cosmictime, scale_factor, NLINESFRW);  //Initializes spline object for data cosmictime, scale_factor of size NLINES
  gsl_spline_init(spline2, cosmictime, der_scale_factor, NLINESFRW);  //Initializes spline object for data cosmictime, der_scale_factor of size NLINES

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

  /***SOLVES GEODESIC EQUATIONS FOR PERTURBED FRW UNIVERSE WITH STATIC PLUMMER POTENTIAL ***/

  /*Parameters for RKF method*/
  mydbl dlambda_max = 1.0e-01, dlambda_min = 1.0e-05, tol = 1.0e-10, lambdaEndPoint = 1.0e+05;

  /*Initial conditions*/
  mydbl ti = 7.0 ,x0, r = -500.0, p0 = 1.0e-3, pr, lambda = 0.0, energy1, energy, v, difft, difference, dlambda = dlambda_max;
  double difftfrw, aem, aobs;
  x0 = C*ti;
  aem = interpolator(spline1, (double)(1.0*ti), acc1);
  pr = condition_factor(r, aem)*p0;
  energy1 = C*energy_factor(r)*p0;
  v = violation(r, p0, pr, aem);
  difft = (energy1 - energy1)/energy1;
  difftfrw = (aem/aem) - 1.0;
  difference = difft - (mydbl)(1.0*difftfrw);

  /*Pointer to file where solution of differential equation will be saved.*/
  FILE *geodesic;
  geodesic = fopen("geodesic_solution.dat","w");

  /*Write line of initial values in file*/
  fprintf(geodesic, "%16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8Le %16.8e %16.8Le %16.8Le\n", lambda, dlambda, x0, r, p0, pr, energy1, v, difft, difftfrw, difference, (mydbl)0.0);

  /*Solution of the differential equation*/
  runge_kutta_fehlberg(spline1, acc1, spline2, acc2, &x0, &r, &p0, &pr, &lambda, &dlambda, lambdaEndPoint, tol, dlambda_max, dlambda_min, geodesic, aem, energy1);

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

  /***RELEASING ALL SPACE USED IN MEMORY***/
  fclose(geodesic); //Close file storing the results
  gsl_spline_free(spline1);  //Free memory of spline object
  gsl_spline_free(spline2);
  gsl_interp_accel_free(acc1);  //Free memory of accel object
  gsl_interp_accel_free(acc2);
}


#include <stdio.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_randist.h>

int
dofit(const gsl_multifit_robust_type *T,
      const gsl_matrix *X, const gsl_vector *y,
      gsl_vector *c, gsl_matrix *cov)
{
  int s;
  gsl_multifit_robust_workspace * work
    = gsl_multifit_robust_alloc (T, X->size1, X->size2);

  s = gsl_multifit_robust (X, y, c, cov, work);
  gsl_multifit_robust_free (work);

  return s;
}

int
main (int argc, char **argv)
{
  int i;
  size_t n;
  const size_t p = 2; /* linear fit */
  gsl_matrix *X, *cov;
  gsl_vector *x, *y, *c, *c_ols;
  const double a = 1.45; /* slope */
  const double b = 3.88; /* intercept */
  gsl_rng *r;

  if (argc != 2)
    {
      fprintf (stderr,"usage: robfit n\n");
      exit (-1);
    }

  n = atoi (argv[1]);

  X = gsl_matrix_alloc (n, p);
  x = gsl_vector_alloc (n);
  y = gsl_vector_alloc (n);

  c = gsl_vector_alloc (p);
  c_ols = gsl_vector_alloc (p);
  cov = gsl_matrix_alloc (p, p);

  r = gsl_rng_alloc(gsl_rng_default);

  /* generate linear dataset */
  for (i = 0; i < n - 3; i++)
    {
      double dx = 10.0 / (n - 1.0);
      double ei = gsl_rng_uniform(r);
      double xi = -5.0 + i * dx;
      double yi = a * xi + b;

      gsl_vector_set (x, i, xi);
      gsl_vector_set (y, i, yi + ei);
    }

  /* add a few outliers */
  gsl_vector_set(x, n - 3, 4.7);
  gsl_vector_set(y, n - 3, -8.3);

  gsl_vector_set(x, n - 2, 3.5);
  gsl_vector_set(y, n - 2, -6.7);

  gsl_vector_set(x, n - 1, 4.1);
  gsl_vector_set(y, n - 1, -6.0);

  /* construct design matrix X for linear fit */
  for (i = 0; i < n; ++i)
    {
      double xi = gsl_vector_get(x, i);

      gsl_matrix_set (X, i, 0, 1.0);
      gsl_matrix_set (X, i, 1, xi);
    }

  /* perform robust and OLS fit */
  dofit(gsl_multifit_robust_ols, X, y, c_ols, cov);
  dofit(gsl_multifit_robust_bisquare, X, y, c, cov);

  /* output data and model */
  for (i = 0; i < n; ++i)
    {
      double xi = gsl_vector_get(x, i);
      double yi = gsl_vector_get(y, i);
      gsl_vector_view v = gsl_matrix_row(X, i);
      double y_ols, y_rob, y_err;

      gsl_multifit_robust_est(&v.vector, c, cov, &y_rob, &y_err);
      gsl_multifit_robust_est(&v.vector, c_ols, cov, &y_ols, &y_err);

      printf("%g %g %g %g\n", xi, yi, y_rob, y_ols);
    }

#define C(i) (gsl_vector_get(c,(i)))
#define COV(i,j) (gsl_matrix_get(cov,(i),(j)))

  {
    printf ("# best fit: Y = %g + %g X\n",
            C(0), C(1));

    printf ("# covariance matrix:\n");
    printf ("# [ %+.5e, %+.5e\n",
               COV(0,0), COV(0,1));
    printf ("#   %+.5e, %+.5e\n",
               COV(1,0), COV(1,1));
  }

  gsl_matrix_free (X);
  gsl_vector_free (x);
  gsl_vector_free (y);
  gsl_vector_free (c);
  gsl_vector_free (c_ols);
  gsl_matrix_free (cov);
  gsl_rng_free(r);

  return 0;
}


//
//  her2k.h
//  Linear Algebra Template Library
//
//  Created by Rodney James on 1/6/12.
//  Copyright (c) 2012 University of Colorado Denver. All rights reserved.
//

#ifndef _her2k_h
#define _her2k_h

/// @file her2k.h Performs multiplication of two complex matrices resulting in a Hermitian matrix.

#include <cctype>
#include "latl.h"

namespace LATL
{
   /// @brief Performs multiplcation of complex matrices resulting in a Hermitian matrix.
   ///
   /// For complex matrices A and B, complex Hermitian matrix C, and complex scalar alpha and real
   /// scalar beta
   ///
   ///        C := alpha*A*B'+conj(alpha)*B*A'+beta*C
   /// or
   ///
   ///        C := alpha*A'*B+conj(alpha)*B'*A+beta*C
   /// is computed.
   /// @return 0 if success.
   /// @return -i if the ith argument is invalid.
   /// @tparam real_t Floating point type.
   /// @param uplo Specifies whether the upper or lower triangular part of the Hermitian matrix C
   /// is to be referenced:
   ///
   ///        if uplo = 'U' or 'u' then C is upper triangular,
   ///        if uplo = 'L' or 'l' then C is lower triangular.
   /// @param trans Specifies the operation to be perfomed as follows:
   ///
   ///        if trans = 'N' or 'n' then C := alpha*A*B'+conj(alpha)*B*A'+beta*C
   ///        if trans = 'C' or 'c' then C := alpha*A'*B+conj(alpha)*B'*A+beta*C
   /// @param n Specifies the order of the complex Hermitian matrix C.  n>=0
   /// @param k Specifies the other dimension of the complex matrices A and B.
   ///
   ///        if trans = 'N' or 'n' then A and B are n-by-k
   ///        if trans = 'C' or 'c' then A and B are k-by-n
   /// @param alpha Complex scalar.
   /// @param A Pointer to complex matrix.
   /// @param ldA Column length of the matrix A. If trans = 'N' or 'n' ldA>=n, otherwise ldA>=k.
   /// @param B Pointer to complex matrix.
   /// @param ldB Column length of the matrix B. If trans = 'N' or 'n' ldB>=n, otherwise ldB>=k.
   /// @param beta Real scalar.
   /// @param C Pointer to complex Hermitian n-by-n matrix C.
   /// Only the upper or lower triangular part of C is referenced, depending on the value of uplo above.
   /// @param ldC Column length of the matrix C.  ldC>=n
   /// @ingroup BLAS

   template <typename real_t>
   int HER2K(char uplo, char trans, int_t n, int_t k, complex<real_t> alpha, complex<real_t> *A, int_t ldA, complex<real_t> *B, int_t ldB, real_t beta, complex<real_t> *C, int_t ldC)
   {
      using std::toupper;
      uplo=toupper(uplo);
      trans=toupper(trans);
      using std::real;
      using std::conj;
      const complex<real_t> zero(0.0,0.0);
      const complex<real_t> one(1.0,0.0);
      int_t i,j,l;
      complex<real_t> *a,*b,*c,*at,*bt;
      complex<real_t> s,t;

      if((uplo!='U')&&(uplo!='L'))
         return -1;
      else if((trans!='N')&&(trans!='C'))
         return -2;
      else if(n<0)
         return -3;
      else if(k<0)
         return -4;
      else if(ldA<((trans=='N')?n:k))
         return -7;
      else if(ldB<((trans=='N')?n:k))
         return -9;
      else if(ldC<n)
         return -12;
      else if((n==0)||(((alpha==zero)||(k==0))&&(beta==one)))
         return 0;

      if(alpha==zero)
      {
         if(uplo=='U')
         {
            c=C;
            for(j=0;j<n;j++)
            {
               for(i=0;i<j;i++)
                  c[i]*=beta;
               c[j]=beta*real(c[j]);
               c+=ldC;
            }

         }
         else
         {
            c=C;
            for(j=0;j<n;j++)
            {
               c[j]=beta*real(c[j]);
               for(i=j+1;i<n;i++)
                  c[i]*=beta;
               c+=ldC;
            }
         }
      }
      else if(trans=='N')
      {
         if(uplo=='U')
         {
            c=C;
            for(j=0;j<n;j++)
            {
               for(i=0;i<j;i++)
                  c[i]*=beta;
               c[j]=beta*real(c[j]);
               a=A;
               b=B;
               for(l=0;l<k;l++)
               {
                  s=alpha*conj(b[j]);
                  t=conj(alpha*a[j]);
                  for(i=0;i<j;i++)
                     c[i]+=s*a[i]+t*b[i];
                  c[j]=real(c[j])+real(a[j]*s+b[j]*t);
                  a+=ldA;
                  b+=ldB;
               }
               c+=ldC;
            }
         }
         else
         {
            c=C;
            for(j=0;j<n;j++)
            {
               c[j]=beta*real(c[j]);
               for(i=j+1;i<n;i++)
                  c[i]*=beta;
               a=A;
               b=B;
               for(l=0;l<k;l++)
               {
                  s=alpha*conj(b[j]);
                  t=conj(alpha*a[j]);
                  c[j]=real(c[j])+real(s*a[j]+t*b[j]);
                  for(i=j+1;i<n;i++)
                     c[i]+=s*a[i]+t*b[i];
                  a+=ldA;
                  b+=ldB;
               }
               c+=ldC;
            }
         }
      }
      else
      {
         if(uplo=='U')
         {
            c=C;
            at=A;
            bt=B;
            for(j=0;j<n;j++)
            {
               a=A;
               b=B;
               for(i=0;i<=j;i++)
               {
                  s=zero;
                  t=zero;
                  for(l=0;l<k;l++)
                  {
                     s+=conj(b[l])*at[l];
                     t+=conj(a[l])*bt[l];
                  }
                  if(i<j)
                     c[i]=alpha*t+conj(alpha)*s+beta*c[i];
                  else
                     c[j]=real(alpha*t+conj(alpha)*s)+beta*real(c[j]);
                  a+=ldA;
                  b+=ldB;
               }
               at+=ldA;
               bt+=ldB;
               c+=ldC;
            }
         }
         else
         {
            c=C;
            at=A;
            bt=B;
            for(j=0;j<n;j++)
            {
               a=A+j*ldA;
               b=B+j*ldB;
               for(i=j;i<n;i++)
               {
                  s=zero;
                  t=zero;
                  for(l=0;l<k;l++)
                  {
                     s+=conj(b[l])*at[l];
                     t+=conj(a[l])*bt[l];
                  }
                  if(i>j)
                     c[i]=alpha*t+conj(alpha)*s+beta*c[i];
                  else
                     c[j]=real(alpha*t+conj(alpha)*s)+beta*real(c[j]);
                  a+=ldA;
                  b+=ldB;
               }
               at+=ldA;
               bt+=ldB;
               c+=ldC;
            }
         }
      }
      return 0;
   }

#ifdef __latl_cblas
#include <cblas.h>

   template <> int HER2K<float>(char uplo, char trans, int_t n, int_t k, complex<float> alpha, complex<float> *A, int_t ldA, complex<float> *B, int_t ldB, float beta, complex<float> *C, int_t ldC)
   {
      using std::toupper;
      uplo=toupper(uplo);
      trans=toupper(trans);
      if((uplo!='U')&&(uplo!='L'))
         return -1;
      else if((trans!='N')&&(trans!='C'))
         return -2;
      else if(n<0)
         return -3;
      else if(k<0)
         return -4;
      else if(ldA<((trans=='N')?n:k))
         return -7;
      else if(ldB<((trans=='N')?n:k))
         return -9;
      else if(ldC<n)
         return -12;

      const CBLAS_UPLO Uplo=(uplo=='U')?CblasUpper:CblasLower;
      const CBLAS_TRANSPOSE Trans=(trans=='N')?CblasNoTrans:((trans=='T')?CblasTrans:CblasConjTrans);

      cblas_cher2k(CblasColMajor,Uplo,Trans,n,k,&alpha,A,ldA,B,ldB,beta,C,ldC);

      return 0;
   }

   template <> int HER2K<double>(char uplo, char trans, int_t n, int_t k, complex<double> alpha, complex<double> *A, int_t ldA, complex<double> *B, int_t ldB, double beta, complex<double> *C, int_t ldC)
   {
      using std::toupper;
      uplo=toupper(uplo);
      trans=toupper(trans);
      if((uplo!='U')&&(uplo!='L'))
         return -1;
      else if((trans!='N')&&(trans!='C'))
         return -2;
      else if(n<0)
         return -3;
      else if(k<0)
         return -4;
      else if(ldA<((trans=='N')?n:k))
         return -7;
      else if(ldB<((trans=='N')?n:k))
         return -9;
      else if(ldC<n)
         return -12;

      const CBLAS_UPLO Uplo=(uplo=='U')?CblasUpper:CblasLower;
      const CBLAS_TRANSPOSE Trans=(trans=='N')?CblasNoTrans:((trans=='T')?CblasTrans:CblasConjTrans);

      cblas_zher2k(CblasColMajor,Uplo,Trans,n,k,&alpha,A,ldA,B,ldB,beta,C,ldC);

      return 0;
   }
#endif

}
#endif