/* Software for the Affected Relative Pair Design */
/* Iuliana Ionita-Laza ii2135@columbia.edu */
/* September 2011 */


#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <gsl/gsl_cdf.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_permutation.h>
#include <gsl/gsl_sort.h>
#include <gsl/gsl_sort_vector.h>
#include <time.h>
#include <gsl/gsl_sf_gamma.h>
#include <unistd.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_eigen.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_eigen.h>
#include <gsl/gsl_blas.h>


#define nall1 3000

int **arp1, **arp2, **unaffected;
int *start, *end;
char **gene;
float *weight, *rs;
int narp, nunaff, nsnps, ngenes, nau;
int kin;
int **sib11, **sib21, **control1;
float *nind, *nind1;
float *freq_c;
float  **W;

double *A1;
double *A2;
double *A3;
double *A4;
double *e_vector;
double *d_vector;
double *diag_e_vector;
double **e;
double **d;
float **p1, **p3, **p6;

void alloca_data()
{
  int i;
  arp1=new int*[nau+100];
  arp2=new int*[nau+100];
  unaffected=new int*[nau+100];
  sib11=new int*[nau+100];
  sib21=new int*[nau+100];
  control1=new int*[nau+100];
  for (i=1; i<=nau; i++)
    control1[i]=new int[nall1+10];

  nind=new float[nsnps+10];
  nind1=new float[nsnps+10];
  freq_c=new float[nsnps+10];

  for (i=1; i<=nau; i++)
    arp1[i]=new int[nsnps+100];
  for (i=1; i<=nau; i++)
    arp2[i]=new int[nsnps+100];
  for (i=1; i<=nau; i++)
    unaffected[i]=new int[nsnps+100];

  for (i=1; i<=nau; i++)
    sib11[i]=new int[nall1+10];
  for (i=1; i<=nau; i++)
    sib21[i]=new int[nall1+10];


  W=new float*[20];
  for (i=1; i<=20; i++)
    W[i]=new float[20];

  p1=new float*[20];
  for (i=0; i<=20; i++)
    p1[i]=new float[20];

  p3=new float*[20];
  for (i=0; i<=20; i++)
    p3[i]=new float[20];

  p6=new float*[20];
  for (i=0; i<=20; i++)
    p6[i]=new float[20];

  A1=new double[nall1*nall1+10];
  A2=new double[nall1*nall1+10];
  A3=new double[nall1*nall1+10];
  A4=new double[nall1*nall1+10];
  e_vector=new double[nall1*nall1+10];
  d_vector=new double[nall1*nall1+10];
  diag_e_vector=new double[nall1*nall1+10];

  d=new double*[nall1];
  e=new double*[nall1];

  for (i=1; i<=nall1; i++)
    d[i]=new double[nall1+10];
  for (i=1; i<=nall1; i++)
    e[i]=new double[nall1+10];

  start=new int[ngenes+100];
  end=new int[ngenes+100];
  gene=new char*[ngenes+100];
  weight=new float[nsnps+100];
  rs=new float[nsnps+100];
  for (i=1; i<=ngenes; i++)
    gene[i]=new char[100];
}

// Relative-to-Relative Transition Matrix in Lange (1976) 
inline float transition_matrix(float f1, float f2, float a11, float keff, float keff_k)
{
  int i, j, k, l, m, n;
  float i33_vector[15];
  float t33a_vector[15];
  float t33b_vector[15];
  float u33a_vector[15];
  float u33b_vector[15];

  float  v;
  W=new float*[20];
  for (i=1; i<=20; i++)
    W[i]=new float[20];


  for (i=0; i<3; i++)
    for (j=0; j<3; j++)
      {
      if (i==j) i33_vector[3*i+j]=1;
      else i33_vector[3*i+j]=0;
      }
  for (i=0; i<3; i++)
    for (j=0; j<3; j++)
      { 
	if ((i==0) && (j==0)) t33a_vector[3*i+j]=1-f1;
	if ((i==0) && (j==1)) t33a_vector[3*i+j]=f1;
	if ((i==0) && (j==2)) t33a_vector[3*i+j]=0;
	if ((i==1) && (j==0)) t33a_vector[3*i+j]=(1-f1)/2.0;
	if ((i==1) && (j==1)) t33a_vector[3*i+j]=0.5;
	if ((i==1) && (j==2)) t33a_vector[3*i+j]=f1/2.0;
	if ((i==2) && (j==0)) t33a_vector[3*i+j]=0;
	if ((i==2) && (j==1)) t33a_vector[3*i+j]=1-f1;
	if ((i==2) && (j==2)) t33a_vector[3*i+j]=f1;
      } 
  for (i=0; i<3; i++)
    for (j=0; j<3; j++)
      {
        if ((i==0) && (j==0)) t33b_vector[3*i+j]=1-f2;
	if ((i==0) && (j==1)) t33b_vector[3*i+j]=f2;
	if ((i==0) && (j==2)) t33b_vector[3*i+j]=0;
	if ((i==1) && (j==0)) t33b_vector[3*i+j]=(1-f2)/2.0;
	if ((i==1) && (j==1)) t33b_vector[3*i+j]=0.5;
	if ((i==1) && (j==2)) t33b_vector[3*i+j]=f2/2.0;
        if ((i==2) && (j==0)) t33b_vector[3*i+j]=0;
        if ((i==2) && (j==1)) t33b_vector[3*i+j]=1-f2;
        if ((i==2) && (j==2)) t33b_vector[3*i+j]=f2;
      }

  for (i=0; i<3; i++)
    for (j=0; j<3; j++)
      {
        if ((i==0) && (j==0)) u33a_vector[3*i+j]=(1-f1)*(1-f1);
	if ((i==0) && (j==1)) u33a_vector[3*i+j]=2*f1*(1-f1);
	if ((i==0) && (j==2)) u33a_vector[3*i+j]=f1*f1;
	if ((i==1) && (j==0)) u33a_vector[3*i+j]=(1-f1)*(1-f1);
	if ((i==1) && (j==1)) u33a_vector[3*i+j]=2*f1*(1-f1);
	if ((i==1) && (j==2)) u33a_vector[3*i+j]=f1*f1;
	if ((i==2) && (j==0)) u33a_vector[3*i+j]=(1-f1)*(1-f1);
	if ((i==2) && (j==1)) u33a_vector[3*i+j]=2*f1*(1-f1);
        if ((i==2) && (j==2)) u33a_vector[3*i+j]=f1*f1;
      }

  for (i=0; i<3; i++)
    for (j=0; j<3; j++)
      {
        if ((i==0) && (j==0)) u33b_vector[3*i+j]=(1-f2)*(1-f2);
	if ((i==0) && (j==1)) u33b_vector[3*i+j]=2*f2*(1-f2);
	if ((i==0) && (j==2)) u33b_vector[3*i+j]=f2*f2;
	if ((i==1) && (j==0)) u33b_vector[3*i+j]=(1-f2)*(1-f2);
	if ((i==1) && (j==1)) u33b_vector[3*i+j]=2*f2*(1-f2);
        if ((i==1) && (j==2)) u33b_vector[3*i+j]=f2*f2;
        if ((i==2) && (j==0)) u33b_vector[3*i+j]=(1-f2)*(1-f2);
        if ((i==2) && (j==1)) u33b_vector[3*i+j]=2*f2*(1-f2);
        if ((i==2) && (j==2)) u33b_vector[3*i+j]=f2*f2;
      }
  
  for (i=0; i<3; i++)
    for (j=0; j<3; j++)
      for (k=0; k<3; k++)
	for (l=0; l<3; l++)
	  {
	    n=3*i+k;
            m=3*j+l;
            p1[n][m]=0.25*i33_vector[3*i+j]*i33_vector[3*k+l];
	  }
  for (i=0; i<3; i++)
    for (j=0; j<3; j++)
      for (k=0; k<3; k++)
	for (l=0; l<3; l++)
          {
            n=3*i+k;
	    m=3*j+l;
	    p3[n][m]=0.5*t33a_vector[3*i+j]*t33b_vector[3*k+l];
          }
  for (i=0; i<3; i++)
    for (j=0; j<3; j++)
      for (k=0; k<3; k++)
        for (l=0; l<3; l++)
          {
            n=3*i+k;
            m=3*j+l;
            p6[n][m]=0.25*u33a_vector[3*i+j]*u33b_vector[3*k+l];
          }
  for (i=0; i<9; i++)
    for (j=0; j<9; j++)
      W[i+1][j+1]=p1[i][j]+p3[i][j]+p6[i][j];
      v=keff*keff_k*a11*W[5][5];
  return (v);
}

main(int argc, char *argv[])
{
  int i, j, k, d1, w, h, b, ig; 
  float sw_A, sw_B;
  int aff, nsnps_gene;
  char a[100]; 
  float a00, a01, a10, a11;
  float nmark, nm;
  float nr;
 
  float gpval1, gpval2;
  float x30;
  float fi, f1, keff, e_teff,var_teff, keff_k, e_teff_k; 
  
  FILE *fin, *fc; 
  const gsl_rng_type * T;
  gsl_rng * r;

  nau=atoi(argv[1]);
  nsnps=atoi(argv[2]);
  ngenes=atoi(argv[3]);
  kin=atoi(argv[4]);

  alloca_data();
  gsl_rng_env_setup(); 
  T = gsl_rng_default;
  r = gsl_rng_alloc (T);
  gsl_rng_set(r, time(0));


  narp=0; nunaff=0;
  fin=fopen("data.ped","r");
  for (i=1; i<=nau; i++)
    {
      fscanf(fin,"%s %s %s %s %s %d",a,a,a,a,a,&aff);
      if (aff==2) {
        narp=narp+1;
        for (j=1; j<=nsnps; j++)
          fscanf(fin,"%d",&arp1[narp][j]);
        fscanf(fin,"%s %s %s %s %s %d",a,a,a,a,a,&aff);
        for (j=1; j<=nsnps; j++)
          fscanf(fin,"%d",&arp2[narp][j]);
      }
      if (aff==1) {
        nunaff=nunaff+1;
        for (j=1; j<=nsnps; j++)
          fscanf(fin,"%d",&unaffected[nunaff][j]);
      }
    }
  fclose(fin);

  

  fin=fopen("genes.txt","r");
  for (i=1; i<=ngenes; i++)
    fscanf(fin,"%s %d %d",gene[i],&start[i],&end[i]);
  fclose(fin);
  printf("%s ",gene[ngenes]);

  fin=fopen("weights.txt","r");
  for (i=1; i<=nsnps; i++)
    fscanf(fin,"%f %f",&rs[i],&weight[i]);
  fclose(fin);

  if (kin==1) { fi=0.25; d1=1;} //sibs 
  if (kin==2) { fi=1.0/16; d1=0;} //uncle-nephew  
  if (kin==3) { fi=1.0/32; d1=0;} //first cousins      
  if (kin==4) { fi=1.0/64; d1=0;} //second cousins            
  
  fc=fopen("results.txt","w");   
  for (ig=1; ig<=ngenes; ig++)
    {
      nsnps_gene=end[ig]-start[ig]+1;

      for (j=1; j<=nsnps_gene; j++)
	  freq_c[j]=0;  
      
      for (i=1; i<=narp; i++)
        for (j=start[ig]; j<=end[ig]; j++)
	  {
	    sib11[i][j-start[ig]+1]=arp1[i][j];
	    sib21[i][j-start[ig]+1]=arp2[i][j];
	  }
      
      for (i=1; i<=nunaff; i++)
	for (j=start[ig]; j<=end[ig]; j++)
            {
              control1[i][j-start[ig]+1]=unaffected[i][j];
              freq_c[j-start[ig]+1]=freq_c[j-start[ig]+1]+control1[i][j-start[ig]+1];
            }
      
      float flag1=0;
      nr=0;
      for (j=1; j<=narp; j++)
        {
          flag1=0;
          for (i=1; i<=nsnps_gene; i++)
            {
              f1=1.0/(2.0*nunaff+2.0);
              keff=log(4*f1*fi+4*f1*f1*(1-fi))/log(2*f1);
              if (freq_c[i]<=0) {
                if ((sib11[j][i]>=1) && (sib21[j][i]>=1)) flag1=keff;
              }
            }
          nr=nr+flag1;
        }
    
	 for (j=1; j<=nsnps_gene; j++)
           {
             if (freq_c[j]<=0) freq_c[j]=(freq_c[j]+1)/(2.0*nunaff+2);
             else freq_c[j]=(freq_c[j]+0.1)/(2.0*nunaff+10.1);
           }

	 
	 for (j=1; j<=nsnps_gene; j++)
     	     for (k=1; k<=nsnps_gene; k++)
	       { e[j][k]=0; d[j][k]=0;}
    	 

  	 sw_A=0; sw_B=0; nmark=0; nm=0;
         for (j=1; j<=nsnps_gene; j++)
	   {
	     nind[j]=0; nind1[j]=0;
             f1=freq_c[j];
             keff=log(4*f1*fi+4*f1*f1*(1-fi))/log(2*f1);
             e_teff=keff*4*f1*fi;
             var_teff=narp*keff*keff*(narp*16*fi*fi*f1*(1-f1)/(2.0*nunaff)+4*f1*fi);
	     for (i=1; i<=narp; i++)
	      	 if ((sib11[i][j]>=1) && (sib21[i][j]>=1)) nind[j]=nind[j]+keff;
	   

	     for (i=1; i<=nunaff; i++)
               nind1[j]=nind1[j]+control1[i][j];

	     if ((freq_c[j]<=0.01) && (nind[j]+nind1[j]>0) && (rs[i]*weight[i]>0)) {
	       sw_A=sw_A+nind[j];
               if (nind1[j]<=0) sw_B=sw_B+nind[j];
	       nmark=nmark+narp*e_teff;
               nm=nm+var_teff;
               for (k=1; k<=nsnps_gene; k++)
		 {
		   nind[k]=0; nind1[k]=0;
		   f1=freq_c[k];
		   keff_k=log(4*f1*fi+4*f1*f1*(1-fi))/log(2*f1);
		   e_teff_k=keff_k*4*f1*fi;
		   for (i=1; i<=narp; i++)
		       if ((sib11[i][k]>=1) && (sib21[i][k]>=1)) nind[k]=nind[k]+keff_k;
                   for (i=1; i<=nunaff;i++)
                     nind1[k]=nind1[k]+control1[i][k];

		   a00=0; a01=0; a10=0; a11=0;
		   for (i=1; i<=nunaff; i++)
		     {
                       if ((control1[i][j]==1) && (control1[i][k]==1)) a11=a11+1;
		     }
                 
		   a11=a11/nunaff;
		   if ((freq_c[k]<=0.01) && (nind[k]+nind1[k]>0)  && (rs[k]*weight[k]>0)) {
                   x30=transition_matrix(freq_c[j],freq_c[k],a11,keff,keff_k);
                   e[j][k]=x30-e_teff*e_teff_k;}
		 }
	       d[j][j]=var_teff;
             }
           }

         for (j=1; j<=nsnps_gene; j++)
           for (k=1; k<=nsnps_gene; k++)
             e[j][k]=narp*e[j][k];

         for (i=0; i<nsnps_gene; i++)
           for (j=0; j<nsnps_gene; j++)
             A1[i*nsnps_gene+j]=0;
         for (i=0; i<nsnps_gene; i++)
           for (j=0; j<nsnps_gene; j++)
             A2[i*nsnps_gene+j]=0;
         for (i=0; i<nsnps_gene; i++)
           for (j=0; j<nsnps_gene; j++)
             A3[i*nsnps_gene+j]=0;
         for (i=0; i<nsnps_gene; i++)
           for (j=0; j<nsnps_gene; j++)
             A4[i*nsnps_gene+j]=0;

         for (i=0; i<nsnps_gene; i++)
           for (j=0; j<nsnps_gene; j++)
             d_vector[i*nsnps_gene+j]=sqrt(d[i+1][j+1]);
         for (i=0; i<nsnps_gene; i++)
           for (j=0; j<nsnps_gene; j++)
                  e_vector[i*nsnps_gene+j]=e[i+1][j+1];
            
	 for (i=0; i<nsnps_gene; i++)
           for (j=0; j<nsnps_gene; j++)
             {
               if ((i==j) && (e[i+1][j+1]>0)) diag_e_vector[i*nsnps_gene+j]=pow(e[i+1][j+1],-0.5);
	       else diag_e_vector[i*nsnps_gene+j]=0;
             }

         gsl_matrix_view E_matrix = gsl_matrix_view_array(e_vector,nsnps_gene,nsnps_gene);
         gsl_matrix_view D_matrix = gsl_matrix_view_array(d_vector,nsnps_gene,nsnps_gene);
         gsl_matrix_view diagE_matrix = gsl_matrix_view_array(diag_e_vector,nsnps_gene,nsnps_gene);
         gsl_matrix_view K1 = gsl_matrix_view_array(A1, nsnps_gene,nsnps_gene);
         gsl_matrix_view K2 = gsl_matrix_view_array(A2, nsnps_gene,nsnps_gene);
         gsl_matrix_view K3 = gsl_matrix_view_array(A3, nsnps_gene,nsnps_gene);
         gsl_matrix_view K4 = gsl_matrix_view_array(A4, nsnps_gene,nsnps_gene);
	 gsl_blas_dgemm (CblasNoTrans, CblasNoTrans,
                         1.0, &diagE_matrix.matrix, &E_matrix.matrix,
                         0.0, &K1.matrix);
         gsl_blas_dgemm (CblasNoTrans, CblasNoTrans,
                         1.0, &K1.matrix, &diagE_matrix.matrix,
                         0.0, &K2.matrix);
         gsl_blas_dgemm (CblasNoTrans, CblasNoTrans,
                         1.0, &D_matrix.matrix, &K2.matrix,
                         0.0, &K3.matrix);
         gsl_blas_dgemm (CblasNoTrans, CblasNoTrans,
                         1.0, &K3.matrix, &D_matrix.matrix,
                         0.0, &K4.matrix);

         float evar=0;
         for (i=0; i<nsnps_gene; i++)
           for (j=0; j<nsnps_gene; j++)
	     if (i!=j)  evar=evar+gsl_matrix_get(&K4.matrix,i,j);
	 evar=evar+nm;
         
         float wequiv1;
         if (nmark>0) { wequiv1=evar/nmark; if (sw_A>0) gpval1=gsl_cdf_gamma_P(nmark/wequiv1,sw_A/wequiv1,1.0); else gpval1=1;
           if (sw_B>0)  gpval2=gsl_cdf_gamma_P(nmark/wequiv1,sw_B/wequiv1,1.0); else gpval2=1;}
         if (nmark==0) {gpval1=1; gpval2=1;}
         fprintf(fc,"%s %.10f %.10f %f\n",gene[ig],gpval1,gpval2,nr);
    }
    fclose(fc); 
}


