/* Software for Case-Control design only */
/* 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 nau;
int **affected, **unaffected;
int *start, *end;
int **case1, **control1;
char **gene;
float *weight, *rs;
double *freq_c;
int naff, nunaff, nsnps, ngenes;
int *nind, *nind1;
double *A1;
double *A2;
double *A3;
double *A4;
double *e_vector;
double *d_vector;
double *diag_e_vector;
double **e;
double **d;


void alloca_data()
{
  int i;
  affected=new int*[nau+100];
  unaffected=new int*[nau+100];
  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];
  case1=new int*[nau+10];
  control1=new int*[nau+10];

  for (i=1; i<=nau; i++)
    affected[i]=new int[nsnps+100];
  for (i=1; i<=nau; i++)
    unaffected[i]=new int[nsnps+100];
  for (i=1; i<=ngenes; i++)
    gene[i]=new char[100];

  for (i=1; i<=nau; i++)
    case1[i]=new int[nsnps+10];
  for (i=1; i<=nau; i++)
    control1[i]=new int[nsnps+10];

  nind=new int[nsnps+10];
  nind1=new int[nsnps+10];
  freq_c=new double[nsnps+10];

  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];
}

main(int argc, char *argv[])
{
  int i, j, j1, k, ig, flag, count, count1, nr, aff, nsnps_gene; 
  double sw_A, sw_B;
  char a[100];
  FILE *fin, *fout;
  double gpval1, gpval2;
  float nm, nmark;
  float wequiv1;
  
  naff=0; nunaff=0;
  nau=atoi(argv[1]);
  nsnps=atoi(argv[2]);
  ngenes=atoi(argv[3]);

  alloca_data();

  // read 1st input file: aff & unaff data
  /***********************************************************************/

  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) { 
        naff=naff+1;
	for (j=1; j<=nsnps; j++)
	  fscanf(fin,"%d",&affected[naff][j]);
      }
      if (aff==1) { 
        nunaff=nunaff+1;
        for (j=1; j<=nsnps; j++)
          fscanf(fin,"%d",&unaffected[nunaff][j]);
      }
    }
  fclose(fin);

  // read 2nd input file: gene names & start/end positions
  /***********************************************************************/
  fin=fopen("genes.txt","r");
  for (i=1; i<=ngenes; i++)
     fscanf(fin,"%s %d %d",gene[i],&start[i],&end[i]);
  fclose(fin);

  // read weights for variants; e.g. 1 for nonsynonymous and 0 otherwise
  /***********************************************************************/
 
  fin=fopen("weights.txt","r");
  for (i=1; i<=nsnps; i++)
    fscanf(fin,"%f %f",&rs[i],&weight[i]);
  fclose(fin);
 
  /***********************************************************************/
  // Compute analytical P-value for the weighted-sum & number of affecteds that carry a novel variant; 
  // Results are written to file: results.txt
  /***********************************************************************/

  fout=fopen("results.txt","w");
  for (ig=1; ig<=ngenes; ig++)
 {
      nm=0; nmark=0;  nr=0;
      nsnps_gene=end[ig]-start[ig]+1;

  for (j=1; j<=nsnps_gene; j++)
	freq_c[j]=0;
 
  for (i=1; i<=naff; i++)
      for (j=start[ig]; j<=end[ig]; j++)
	  case1[i][j-start[ig]+1]=affected[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];
	}
  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;
      for (i=1; i<=naff; i++)
	nind[j]=nind[j]+case1[i][j];
      for (i=1; i<=nunaff; i++)
	nind1[j]=nind1[j]+control1[i][j];

      if ((freq_c[j]<=0.01) && (nind[j]+nind1[j]>0) && (weight[j+start[ig]-1]*rs[j+start[ig]-1]>0)) {
	sw_A=sw_A+nind[j];
	if (nind1[j]<=0) sw_B=sw_B+nind[j];
	nmark=nmark+2*naff*freq_c[j];
	nm=nm+2*naff*((2.0*naff-1)/(2.0*nunaff)+1)*freq_c[j]*(1-freq_c[j]);
	for (k=1; k<=nsnps_gene; k++)
	  {
	    nind[k]=0; nind1[k]=0;
	    for (i=1; i<=naff; i++)
	      nind[k]=nind[k]+case1[i][k];
	    for (i=1; i<=nunaff; i++)
	      nind1[k]=nind1[k]+control1[i][k];

	    for (i=1; i<=naff; i++)
	      if ((freq_c[k]<=0.01) && (nind[k]+nind1[k]>0) && (weight[k+start[ig]-1]*rs[k+start[ig]-1]>0)) 
                          e[j][k]=e[j][k]+(case1[i][j]-2*freq_c[j])*(case1[i][k]-2*freq_c[k]);
	    for (i=1; i<=nunaff; i++)
	      if ((freq_c[k]<=0.01)  && (nind[k]+nind1[k]>0) && (weight[k+start[ig]-1]*rs[k+start[ig]-1]>0)) 
                          e[j][k]=e[j][k]+(control1[i][j]-2*freq_c[j])*(control1[i][k]-2*freq_c[k]);
	  }
	d[j][j]=2*naff*((2.0*naff-1)/(2.0*nunaff)+1)*freq_c[j]*(1-freq_c[j]);
      }
    }
 
  for (j=1; j<=nsnps_gene; j++)
    for (k=1; k<=nsnps_gene; k++)
      e[j][k]=naff*e[j][k]/(naff+nunaff);

  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);
 
  double 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;

  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(fout,"%s %.10f %.10f ",gene[ig],gpval1,gpval2);    
 
  nr=0;
  for (j=1; j<=naff; j++)
    {
      flag=0;
      for (i=start[ig]; i<=end[ig]; i++)
	{
	  count1=0;
	  for (j1=1; j1<=nunaff; j1++)
	    count1=count1+unaffected[j1][i];
     	  if ((weight[i]*rs[i]>0) && (count1<=0) && (affected[j][i]>0)) flag=1;
	}
      nr=nr+flag;
    }
    fprintf(fout,"%d\n",nr);
 }
  fclose(fout); 
}




