#include "function.h"


/******************************************************
  Seven functions to calculate expectation of genotypes
  between two outbred individuals,
  coef2: IBD2; coef1: IBD1; coef0: IBD0
 *****************************************************/
double fun1(double coef2, double coef1, double coef0, double p){
// genotypes aa and aa
  return coef2*p*p+coef1*p*p*p+coef0*p*p*p*p;
}

double fun2(double coef1, double coef0, double p, double q){
// genotypes aa and ab
  return coef1*p*p*q+2*coef0*p*p*p*q;
}

double fun3(double coef0, double p, double q){
// genotypes aa and bb
  return coef0*p*p*q*q;
}

double fun4(double coef0, double p, double q1, double q2){
// genotypes aa and bc
  return 2*coef0*p*p*q1*q2;
}

double fun7(double coef2, double coef1, double coef0, double p, double q){
// genotypes ab and ab
  return 2*coef2*p*q+coef1*p*q*(p+q)+4*coef0*p*p*q*q;
}

double fun8(double coef1, double coef0, double p, double q1, double q2){
// genotypes ab and ac
  return coef1*p*q1*q2+4*coef0*p*p*q1*q2;
}

double fun9(double coef0, double q1, double q2, double q3, double q4){
// genotypes ab and cd
  return 4*coef0*q1*q2*q3*q4;
}


/******************************************************
  Compute expectation for two related individuals
 *****************************************************/
double expfortwo(double IBD_2, double IBD_1, double IBD_0, const pair<int, int> & hp_1, const pair<int, int> & hp_2, map<int, double> & hap_freq_table){
  // i has haplotype pair h1-h2, j has haplotype pair h3-h4 (h1<=h2, h3<=h4)

  double value;

  int h1=hp_1.first;
  int h2=hp_1.second;
  int h3=hp_2.first;
  int h4=hp_2.second;

  if(h1 == h2 && h1 == h3 && h3 == h4)
    value = fun1(IBD_2,IBD_1,IBD_0,hap_freq_table[h1]);
  else if(h1 == h2 && h1 == h3 && h3 != h4)
    value = fun2(IBD_1,IBD_0,hap_freq_table[h1],hap_freq_table[h4]);
  else if(h1 == h2 && h1 == h4 && h3 != h4)
    value = fun2(IBD_1,IBD_0,hap_freq_table[h1],hap_freq_table[h3]);
  else if(h1 == h2 && h1 != h3 && h3 == h4)
    value = fun3(IBD_0,hap_freq_table[h1],hap_freq_table[h3]);
  else if(h1 == h2 && h1 != h3 && h1 != h4 && h3 != h4)
    value = fun4(IBD_0,hap_freq_table[h1],hap_freq_table[h3],hap_freq_table[h4]);
  else if(h1 != h2 && h1 == h3 && h3 == h4)
    value = fun2(IBD_1,IBD_0,hap_freq_table[h3],hap_freq_table[h2]);
  else if(h1 != h2 && h2 == h3 && h3 == h4)
    value = fun2(IBD_1,IBD_0,hap_freq_table[h3],hap_freq_table[h1]);
  else if(h1 != h2 && h1 != h3 && h2 != h3 && h3 == h4)
    value = fun4(IBD_0,hap_freq_table[h3],hap_freq_table[h1],hap_freq_table[h2]);
  else if(h1 != h2 && h1 == h3 && h2 == h4 && h3 != h4)
    value = fun7(IBD_2,IBD_1,IBD_0,hap_freq_table[h1],hap_freq_table[h2]);
  else if(h1 != h2 && h1 == h3 && h2 != h4 && h3 != h4)
    value = fun8(IBD_1,IBD_0,hap_freq_table[h1],hap_freq_table[h2],hap_freq_table[h4]);
  else if(h1 != h2 && h1 == h4 && h2 != h3 && h3 != h4)
    value = fun8(IBD_1,IBD_0,hap_freq_table[h1],hap_freq_table[h2],hap_freq_table[h3]);
  else if(h1 != h2 && h1 != h3 && h2 == h4 && h3 != h4)
    value = fun8(IBD_1,IBD_0,hap_freq_table[h2],hap_freq_table[h1],hap_freq_table[h3]);
  else if(h1 != h2 && h1 != h4 && h2 == h3 && h3 != h4)
    value = fun8(IBD_1,IBD_0,hap_freq_table[h2],hap_freq_table[h1],hap_freq_table[h4]);
  else if(h1 != h2 && h1 != h3 && h1 != h4 && h2 != h3 && h2 != h4 && h3 != h4)
    value = fun9(IBD_0,hap_freq_table[h1],hap_freq_table[h2],hap_freq_table[h3],hap_freq_table[h4]);

  return value;
}


/******************************************************
  Compute prob that offspring has haplotypes ij given (grand)parents haplotypes kl and st
 *****************************************************/
double transProb(const pair<int, int> & hp_0, const pair<int, int> & hp_1, const pair<int, int> & hp_2, map<int, double> & hap_freq_table, map<pair<int, int>, double> & geno_freq_table, int index){

  double result=0;

  int k=hp_1.first;
  int l=hp_1.second;
  int s=hp_2.first;
  int t=hp_2.second;

  if(index==1){  // index is the generation difference
    int i=hp_0.first;
    int j=hp_0.second;
  
    if((i==k && j==s) || (i==s && j==k))
      result+=0.25;
    if((i==k && j==t) || (i==t && j==k))
      result+=0.25;
    if((i==l && j==s) || (i==s && j==l))
      result+=0.25;
    if((i==l && j==t) || (i==t && j==l))
      result+=0.25;
  }
  else if(index>1){
    pair<int, int> hp_3;
    int min, max;

    // inherit k and s
    min = k<=s ? k : s;
    max = k+s-min;
    hp_3 = pair<int, int>(min,max);
    result+=transProbonep(hp_0,hp_3,hap_freq_table,geno_freq_table,index-1);

    // inherit k and t
    min = k<=t ? k : t;
    max = k+t-min;
    hp_3 = pair<int, int>(min,max);
    result+=transProbonep(hp_0,hp_3,hap_freq_table,geno_freq_table,index-1);

    // inherit l and s
    min = l<=s ? l : s;
    max = l+s-min;
    hp_3 = pair<int, int>(min,max);
    result+=transProbonep(hp_0,hp_3,hap_freq_table,geno_freq_table,index-1);

    // inherit l and t
    min = l<=t ? l : t;
    max = l+t-min;
    hp_3 = pair<int, int>(min,max);
    result+=transProbonep(hp_0,hp_3,hap_freq_table,geno_freq_table,index-1);

    result/=4.0;
  }

  return result;
}


/******************************************************
  Compute prob that offspring has haplotypes ij given one (grand)parent haplotypes kl (i<=j, k<=l)
 *****************************************************/
double transProbonep(const pair<int, int> & hp_0, const pair<int, int> & hp_1, map<int, double> & hap_freq_table, map<pair<int, int>, double> & geno_freq_table, int index){

  double result=0;

  int i=hp_0.first;
  int j=hp_0.second;
  int k=hp_1.first;
  int l=hp_1.second;

  if(k==l){  // parent is homozygote
    if(i==k)
      result = hap_freq_table[j];
    else if(j==k)
      result = hap_freq_table[i];
    if(index>1){  // index is the generation difference
      for(int s=0;s<index-1;s++)
        result*=0.5;
    }
  }
  else{  // parent is heterozygote
    if(i==j){  // offspring is homozygote
      if(i==k || i==l)
        result = hap_freq_table[j];
    }
    else{  // offspring is heterozygote
      if(i==k && j==l)
        result = hap_freq_table[i]+hap_freq_table[j];
      else if(i==k || i==l)
        result = hap_freq_table[j];
      else if(j==k || j==l)
        result = hap_freq_table[i];
    }
    for(int s=0;s<index;s++)
      result*=0.5;
  }

  if(index>1){
    double coef=1.0;
    for(int s=0;s<index-1;s++)
      coef*=0.5;
    result+=geno_freq_table[hp_0]*(1.0-coef);
  }

  return result;
}


/******************************************************
  Compute expectation for three related individuals, j is i's one parent side sib, and k is i's the other parent side sib
 *****************************************************/
double uncleaunt(const pair<int, int> & hp_0, const pair<int, int> & hp_1, const pair<int, int> & hp_2, map<int, double> & hap_freq_table, map<pair<int, int>, double> & geno_freq_table){

  // ip is one parent of i that relates to jf, ip' is the other parent of i that relates to jm
  double result=0;

  map<pair<int, int>, double>::iterator giter1;
  map<pair<int, int>, double>::iterator giter2;

  giter1 = geno_freq_table.begin();
  while(giter1!=geno_freq_table.end()){
    double coef1 = expfortwo(0.25,0.5,0.25,giter1->first,hp_1,hap_freq_table);
    if(coef1>0){
      double sum = 0;
      giter2 = geno_freq_table.begin();
      while(giter2!=geno_freq_table.end()){
        double coef = transProb(hp_0,giter1->first,giter2->first,hap_freq_table,geno_freq_table,1);
        if(coef>0)
          sum+=expfortwo(0.25,0.5,0.25,giter2->first,hp_2,hap_freq_table)*coef;
        giter2++;
      }
      result+=coef1*sum;
    }
    giter1++;
  }

  return result;
}


double poz(double z){

  double y, x, w;
       
  if (z == 0.0)
    x = 0.0;
  else{
    y = 0.5 * fabs (z);
    if (y >= (Z_MAX * 0.5))
      x = 1.0;
    else if (y < 1.0){
      w = y*y;
      x = ((((((((0.000124818987 * w
           -0.001075204047) * w +0.005198775019) * w
           -0.019198292004) * w +0.059054035642) * w
           -0.151968751364) * w +0.319152932694) * w
           -0.531923007300) * w +0.797884560593) * y * 2.0;
    }
    else{
      y -= 2.0;
      x = (((((((((((((-0.000045255659 * y
           +0.000152529290) * y -0.000019538132) * y
           -0.000676904986) * y +0.001390604284) * y
           -0.000794620820) * y -0.002034254874) * y
           +0.006549791214) * y -0.010557625006) * y
           +0.011630447319) * y -0.009279453341) * y
           +0.005353579108) * y -0.002141268741) * y
           +0.000535310849) * y +0.999936657524;
    }
  }

  return (z > 0.0 ? ((x + 1.0) * 0.5) : ((1.0 - x) * 0.5));
}


double critz(double p)
{
  double minz = -Z_MAX;    /* minimum of range of z */
  double maxz = Z_MAX;     /* maximum of range of z */
  double zval = 0.0;       /* computed/returned z value */
  double pval;     /* prob (z) function, pval := poz (zval) */

  if (p <= 0.0 || p >= 1.0)
    return (0.0);
       
  while (maxz - minz > Z_EPSILON){
    pval = poz(zval);
    if (pval > p)
      maxz = zval;
    else
      minz = zval;
    zval = (maxz + minz) * 0.5;
  }

  return (zval);
}


double pochisq(double x, int df)
{
  double a, y, s;
  double e, c, z;
         /* computes probability of normal z score */
  int even;     /* true if df is an even number */
       
  if (x <= 0.0 || df < 1)
    return (1.0);
       
  a = 0.5 * x;
  even = (2*(df/2)) == df;
  if (df > 1)
    y = ex (-a);
  s = (even ? y : (2.0 * poz(-sqrt (x))));
  if (df > 2){
    x = 0.5 * (df - 1.0);
    z = (even ? 1.0 : 0.5);
    if (a > BIGX){
      e = (even ? 0.0 : LOG_SQRT_PI);
      c = log (a);
      while (z <= x){
        e = log (z) + e;
        s += ex (c*z-a-e);
        z += 1.0;
      }
      return (s);
    }
    else{
      e = (even ? 1.0 : (I_SQRT_PI / sqrt (a)));
      c = 0.0;
      while (z <= x){
        e = e * (a / z);
        c = c + e;
        z += 1.0;
      }
      return (c * y + s);
    }
  }
  else
    return (s);
}


double critchi(double p, int df)
{
  double minchisq = 0.0;
  double maxchisq = CHI_MAX;
  double chisqval;
       
  if (p <= 0.0)
    return (maxchisq);
  else if (p >= 1.0)
    return (0.0);
       
  chisqval = df / sqrt (p);    /* fair first value */
  while (maxchisq - minchisq > CHI_EPSILON){
    if (pochisq (chisqval, df) < p)
      maxchisq = chisqval;
    else
      minchisq = chisqval;
    chisqval = (maxchisq + minchisq) * 0.5;
  }
  return (chisqval);
}


double fun1(double p){
// genotypes aa and aa
  return p*p*p;
}

double fun2(double p, double q){
// genotypes aa and ab
  return p*p*q;
}

double fun7(double p, double q){
// genotypes ab and ab
  return p*q*(p+q);
}

double fun8(double p, double q1, double q2){
// genotypes ab and ac
  return p*q1*q2;
}


/******************************************************
  Compute expectation for two related individuals given that they share 1 IBD
 *****************************************************/
double expfortwo(const pair<int, int> & hp, map<int, double> & hap_freq_table){
  // i has haplotype pair h1-h2, j has haplotype pair h1-h2 (h1<=h2)

  double value = 0;

  int h1=hp.first;
  int h2=hp.second;

  if(h1 == h2)
    value = fun1(hap_freq_table[h1]);
  else
    value = fun7(hap_freq_table[h1],hap_freq_table[h2]);

  return value;
}


double expfortwo(const pair<int, int> & hp_1, const pair<int, int> & hp_2, map<int, double> & hap_freq_table){
  // i has haplotype pair h1-h2, j has haplotype pair h3-h4 (h1<=h2, h3<=h4)

  double value = 0;

  int h1=hp_1.first;
  int h2=hp_1.second;
  int h3=hp_2.first;
  int h4=hp_2.second;

  if(h1 == h2 && h1 == h3 && h3 == h4)
    value = fun1(hap_freq_table[h1]);
  else if(h1 == h2 && h1 == h3 && h3 != h4)
    value = fun2(hap_freq_table[h1],hap_freq_table[h4]);
  else if(h1 == h2 && h1 == h4 && h3 != h4)
    value = fun2(hap_freq_table[h1],hap_freq_table[h3]);
  else if(h1 != h2 && h1 == h3 && h3 == h4)
    value = fun2(hap_freq_table[h3],hap_freq_table[h2]);
  else if(h1 != h2 && h2 == h3 && h3 == h4)
    value = fun2(hap_freq_table[h3],hap_freq_table[h1]);
  else if(h1 != h2 && h1 == h3 && h2 == h4 && h3 != h4)
    value = fun7(hap_freq_table[h1],hap_freq_table[h2]);
  else if(h1 != h2 && h1 == h3 && h2 != h4 && h3 != h4)
    value = fun8(hap_freq_table[h1],hap_freq_table[h2],hap_freq_table[h4]);
  else if(h1 != h2 && h1 == h4 && h2 != h3 && h3 != h4)
    value = fun8(hap_freq_table[h1],hap_freq_table[h2],hap_freq_table[h3]);
  else if(h1 != h2 && h1 != h3 && h2 == h4 && h3 != h4)
    value = fun8(hap_freq_table[h2],hap_freq_table[h1],hap_freq_table[h3]);
  else if(h1 != h2 && h1 != h4 && h2 == h3 && h3 != h4)
    value = fun8(hap_freq_table[h2],hap_freq_table[h1],hap_freq_table[h4]);

  return value;
}