#include "genotypetrio.h"


Genotypetrio::Genotypetrio(Genotype * offspring, Genotype * par1, Genotype * par2, int typed_o, int typed_par1, int typed_par2, int size)
{
  if(size==0)
    return;

  this->offspring = offspring;
  this->par1 = par1;
  this->par2 = par2;

  this->typed_o = typed_o;
  this->typed_par1 = typed_par1;
  this->typed_par2 = typed_par2;

  this->size = size;  // # of SNPs
  this->count=0;

  record_trio_hap_structure();
  determin_phase();
}


Genotypetrio::Genotypetrio(Genotype * offspring, int typed_o, int typed_par1, int typed_par2, int size)  // for founder
{
  if(size==0)
    return;
  if(typed_par1==1)
    printf("Constructor of Genotypetrio for founder fails. Please check\n");
  if(typed_par2==1)
    printf("Constructor of Genotypetrio for founder fails. Please check\n");

  this->offspring = offspring;

  this->typed_o = typed_o;
  this->typed_par1 = typed_par1;
  this->typed_par2 = typed_par2;

  this->size = size;  // # of SNPs
  this->count=0;

  record_trio_hap_structure();
  determin_phase();
}
  

void Genotypetrio::show_gstat()  // print trio genotype with its count
{
  printf("\nOffspring has genotype: ");
  int *a=offspring->get_gcode();
  for(int i=0;i<size;i++)
  {
    if(a[i]!=-1)
      printf("%d", a[i]);
    else
      printf("*");
    if(i!=size-1)
      printf("-");
  }

  printf(" Parent 1 has genotype: ");
  if(typed_par1==1)
  {
    a=par1->get_gcode();
    for(int i=0;i<size;i++)
    {
      if(a[i]!=-1)
        printf("%d", a[i]);
      else
        printf("*");
      if(i!=size-1)
        printf("-");
    }
  }
  else
  {
    for(int i=0;i<size;i++)
    {
      printf("*");
      if(i!=size-1)
        printf("-");
    }
  }
  
  printf(" Parent 2 has genotype: ");
  if(typed_par2==1)
  {
    a=par2->get_gcode();
    for(int i=0;i<size;i++)
    {
      if(a[i]!=-1)
        printf("%d", a[i]);
      else
        printf("*");
      if(i!=size-1)
        printf("-");
    }
  }
  else
  {
    for(int i=0;i<size;i++)
    {
      printf("*");
      if(i!=size-1)
        printf("-");
    }
  }

  printf(": %d ", count);
  if(phase==1)
    printf("offspring haplotype phase is determined\n");
  else
    printf("\n");
}


void Genotypetrio::show_content()  // print inh_vec
{
  printf("\nHaplotype:\n");
  printf("%d possibilities.\n", inh_vec.size());
  if(typed_par1==1 && typed_par2==1)
  {
    for(int i=0;i<inh_vec.size();i++)
    {
      printf("Parent 1 has a pair of haplotype: \n");
      display(inh_vec[i].hp_par1);
      printf("Parent 2 has a pair of haplotype: \n");
      display(inh_vec[i].hp_par2);
      printf("Offspring has a pair of haplotype: \n");
      display(inh_vec[i].hp_offspring);

      printf("\n");
    }
  }
  else if(typed_par1==1 && typed_par2==0)
  {
    for(int i=0;i<inh_vec.size();i++)
    {
      printf("Parent 1 has a pair of haplotype: \n");
      display(inh_vec[i].hp_par1);
      printf("Offspring has a pair of haplotype: \n");
      display(inh_vec[i].hp_offspring);

      printf("\n");
    }
  }
  else if(typed_par1==0 && typed_par2==0)
  {
    for(int i=0;i<inh_vec.size();i++)
    {
      printf("Offspring has a pair of haplotype: \n");
      display(inh_vec[i].hp_offspring);

      printf("\n");
    }
  }
}


void Genotypetrio::display(const pair<int, int> & hap)
{
  int h1 = hap.first;
  for(int j=0;j<size;j++)
  {
    int mask = 1<<j;
    int b1 = (h1&mask)>>j;  // get the jth last bit
    printf("%d", b1);
    if(j!=size-1)
      printf("-");
    else
      printf("\n");
  }
  int h2 = hap.second;
  for(int j=0;j<size;j++)
  {
    int mask = 1<<j;
    int b2 = (h2&mask)>>j;
    printf("%d", b2);
    if(j!=size-1)
      printf("-");
    else
      printf("\n");
  }
}


void Genotypetrio::push_vec(const vector<pair<int, int> > & hap_vec, int f0, int f1, int m0, int m1, int a, int b)
{
  int min = a<=b ? a : b;
  int max = a+b-min;

  if(is_in(hap_vec,min,max))
  {
    inheritance m;
    m.hp_par1 = pair<int, int>(f0,f1);
    m.hp_par2 = pair<int, int>(m0,m1);
    m.hp_offspring = pair<int, int>(min,max);

    inh_vec.push_back(m);
  }
}


void Genotypetrio::push_vec(int f0, int f1, int m0, int m1, int a, int b)
{
  int min = a<=b ? a : b;
  int max = a+b-min;

  inheritance m;
  m.hp_par1 = pair<int, int>(f0,f1);
  m.hp_par2 = pair<int, int>(m0,m1);
  m.hp_offspring = pair<int, int>(min,max);

  inh_vec.push_back(m);
}


void Genotypetrio::push_vec_complete(const vector<pair<int, int> > & hap_vec, int f0, int f1, int m0, int m1, int a, int b)
{
  int min = a<=b ? a : b;
  int max = a+b-min;

  if(is_in(hap_vec,min,max))
  {
    inheritance m;
    m.hp_par1 = pair<int, int>(f0,f1);
    m.hp_par2 = pair<int, int>(m0,m1);
    m.hp_offspring = pair<int, int>(min,max);

    inh_vec_complete.push_back(m);
  }
}


void Genotypetrio::push_vec_complete(int f0, int f1, int m0, int m1, int a, int b)
{
  int min = a<=b ? a : b;
  int max = a+b-min;

  inheritance m;
  m.hp_par1 = pair<int, int>(f0,f1);
  m.hp_par2 = pair<int, int>(m0,m1);
  m.hp_offspring = pair<int, int>(min,max);

  inh_vec_complete.push_back(m);
}


void Genotypetrio::record_trio_hap_structure()
{
  if(typed_o==1 && typed_par1==1 && typed_par2==1)
  {
    vector<pair<int, int> > hap_vec_o = offspring->get_hap_vec();
    vector<pair<int, int> > hap_vec_f = par1->get_hap_vec();
    vector<pair<int, int> > hap_vec_m = par2->get_hap_vec();

    int f0, f1, m0, m1;
    for(int i=0;i<hap_vec_f.size();i++)
    {
      f0 = hap_vec_f[i].first;
      f1 = hap_vec_f[i].second;
      for(int j=0;j<hap_vec_m.size();j++)
      {
        m0 = hap_vec_m[j].first;
        m1 = hap_vec_m[j].second;
        if(f0==f1 && m0==m1)  // both parents are homozygotes
          push_vec(hap_vec_o,f0,f1,m0,m1,f0,m0);
        else if(f0==f1 && m0!=m1)  // father is homozygote
        {
          push_vec(hap_vec_o,f0,f1,m0,m1,f0,m0);
          push_vec(hap_vec_o,f0,f1,m0,m1,f0,m1);
        }
        else if(f0!=f1 && m0==m1)  // mother is homozygote
        {
          push_vec(hap_vec_o,f0,f1,m0,m1,f0,m0);
          push_vec(hap_vec_o,f0,f1,m0,m1,f1,m0);
        }
        else if(f0!=f1 && m0!=m1)  // both parents are heterozygotes
        {
          if(f0==m0 && f1==m1)
          {
            push_vec(hap_vec_o,f0,f1,m0,m1,f0,m0);
            push_vec(hap_vec_o,f0,f1,m0,m1,f0,m1);
            push_vec(hap_vec_o,f0,f1,m0,m1,f1,m1);
          }
          else
          {
            push_vec(hap_vec_o,f0,f1,m0,m1,f0,m0);
            push_vec(hap_vec_o,f0,f1,m0,m1,f0,m1);
            push_vec(hap_vec_o,f0,f1,m0,m1,f1,m0);
            push_vec(hap_vec_o,f0,f1,m0,m1,f1,m1);
          }
        }
      }  // complete j
    }  // complete i
  }  // both parents are typed
  else if(typed_o==0 && typed_par1==1 && typed_par2==1)
  {
    vector<pair<int, int> > hap_vec_f = par1->get_hap_vec();
    vector<pair<int, int> > hap_vec_m = par2->get_hap_vec();

    int f0, f1, m0, m1;
    for(int i=0;i<hap_vec_f.size();i++)
    {
      f0 = hap_vec_f[i].first;
      f1 = hap_vec_f[i].second;
      for(int j=0;j<hap_vec_m.size();j++)
      {
        m0 = hap_vec_m[j].first;
        m1 = hap_vec_m[j].second;
        if(f0==f1 && m0==m1)  // both parents are homozygotes
          push_vec(f0,f1,m0,m1,f0,m0);
        else if(f0==f1 && m0!=m1)  // father is homozygote
        {
          push_vec(f0,f1,m0,m1,f0,m0);
          push_vec(f0,f1,m0,m1,f0,m1);
        }
        else if(f0!=f1 && m0==m1)  // mother is homozygote
        {
          push_vec(f0,f1,m0,m1,f0,m0);
          push_vec(f0,f1,m0,m1,f1,m0);
        }
        else if(f0!=f1 && m0!=m1)  // both parents are heterozygotes
        {
          if(f0==m0 && f1==m1)
          {
            push_vec(f0,f1,m0,m1,f0,m0);
            push_vec(f0,f1,m0,m1,f0,m1);
            push_vec(f0,f1,m0,m1,f1,m1);
          }
          else
          {
            push_vec(f0,f1,m0,m1,f0,m0);
            push_vec(f0,f1,m0,m1,f0,m1);
            push_vec(f0,f1,m0,m1,f1,m0);
            push_vec(f0,f1,m0,m1,f1,m1);
          }
        }
      }  // complete j
    }  // complete i
  }  // both parents are typed
  else if(typed_par1==1 && typed_par2==0)
  {
    vector<pair<int, int> > hap_vec_o = offspring->get_hap_vec();
    vector<pair<int, int> > hap_vec_f = par1->get_hap_vec();

    int f0, f1, o0, o1;
    for(int i=0;i<hap_vec_f.size();i++)
    {
      f0 = hap_vec_f[i].first;
      f1 = hap_vec_f[i].second;
      for(int j=0;j<hap_vec_o.size();j++)
      {
        o0 = hap_vec_o[j].first;
        o1 = hap_vec_o[j].second;
        if(o0==f0 || o0==f1 || o1==f0 || o1==f1)
        {
          inheritance m;
          m.hp_par1 = pair<int, int>(f0,f1);
          m.hp_par2 = pair<int, int>(-1,-1);
          m.hp_offspring = pair<int, int>(o0,o1);

          inh_vec.push_back(m);
        }
      }  // complete j
    }  // complete i
  }  // parent 1 is typed
  else if(typed_par1==0 && typed_par2==0)
  {
    vector<pair<int, int> > hap_vec_o = offspring->get_hap_vec();

    int o0, o1;
    for(int i=0;i<hap_vec_o.size();i++)
    {
      o0 = hap_vec_o[i].first;
      o1 = hap_vec_o[i].second;
      inheritance m;
      m.hp_par1 = pair<int, int>(-1,-1);
      m.hp_par2 = pair<int, int>(-1,-1);
      m.hp_offspring = pair<int, int>(o0,o1);

      inh_vec.push_back(m);
    }
  }  // both parents are untyped

  if(inh_vec.size()==0)
    info = 0;
  else
    info = 1;

  if(typed_o==1 && typed_par1==1 && typed_par2==0)
  {
    vector<pair<int, int> > hap_vec_o = offspring->get_hap_vec();
    vector<pair<int, int> > hap_vec_f = par1->get_hap_vec();
    vector<pair<int, int> > hap_vec_m = par2->get_hap_vec();

    int f0, f1, m0, m1;
    for(int i=0;i<hap_vec_f.size();i++)
    {
      f0 = hap_vec_f[i].first;
      f1 = hap_vec_f[i].second;
      for(int j=0;j<hap_vec_m.size();j++)
      {
        m0 = hap_vec_m[j].first;
        m1 = hap_vec_m[j].second;
        if(f0==f1 && m0==m1)  // both parents are homozygotes
          push_vec_complete(hap_vec_o,f0,f1,m0,m1,f0,m0);
        else if(f0==f1 && m0!=m1)  // father is homozygote
        {
          push_vec_complete(hap_vec_o,f0,f1,m0,m1,f0,m0);
          push_vec_complete(hap_vec_o,f0,f1,m0,m1,f0,m1);
        }
        else if(f0!=f1 && m0==m1)  // mother is homozygote
        {
          push_vec_complete(hap_vec_o,f0,f1,m0,m1,f0,m0);
          push_vec_complete(hap_vec_o,f0,f1,m0,m1,f1,m0);
        }
        else if(f0!=f1 && m0!=m1)  // both parents are heterozygotes
        {
          if(f0==m0 && f1==m1)
          {
            push_vec_complete(hap_vec_o,f0,f1,m0,m1,f0,m0);
            push_vec_complete(hap_vec_o,f0,f1,m0,m1,f0,m1);
            push_vec_complete(hap_vec_o,f0,f1,m0,m1,f1,m1);
          }
          else
          {
            push_vec_complete(hap_vec_o,f0,f1,m0,m1,f0,m0);
            push_vec_complete(hap_vec_o,f0,f1,m0,m1,f0,m1);
            push_vec_complete(hap_vec_o,f0,f1,m0,m1,f1,m0);
            push_vec_complete(hap_vec_o,f0,f1,m0,m1,f1,m1);
          }
        }
      }  // complete j
    }  // complete i
  }  // one parent is typed
  else if(typed_o==0 && typed_par1==1 && typed_par2==0)
  {
    vector<pair<int, int> > hap_vec_f = par1->get_hap_vec();
    vector<pair<int, int> > hap_vec_m = par2->get_hap_vec();

    int f0, f1, m0, m1;
    for(int i=0;i<hap_vec_f.size();i++)
    {
      f0 = hap_vec_f[i].first;
      f1 = hap_vec_f[i].second;
      for(int j=0;j<hap_vec_m.size();j++)
      {
        m0 = hap_vec_m[j].first;
        m1 = hap_vec_m[j].second;
        if(f0==f1 && m0==m1)  // both parents are homozygotes
          push_vec_complete(f0,f1,m0,m1,f0,m0);
        else if(f0==f1 && m0!=m1)  // father is homozygote
        {
          push_vec_complete(f0,f1,m0,m1,f0,m0);
          push_vec_complete(f0,f1,m0,m1,f0,m1);
        }
        else if(f0!=f1 && m0==m1)  // mother is homozygote
        {
          push_vec_complete(f0,f1,m0,m1,f0,m0);
          push_vec_complete(f0,f1,m0,m1,f1,m0);
        }
        else if(f0!=f1 && m0!=m1)  // both parents are heterozygotes
        {
          if(f0==m0 && f1==m1)
          {
            push_vec_complete(f0,f1,m0,m1,f0,m0);
            push_vec_complete(f0,f1,m0,m1,f0,m1);
            push_vec_complete(f0,f1,m0,m1,f1,m1);
          }
          else
          {
            push_vec_complete(f0,f1,m0,m1,f0,m0);
            push_vec_complete(f0,f1,m0,m1,f0,m1);
            push_vec_complete(f0,f1,m0,m1,f1,m0);
            push_vec_complete(f0,f1,m0,m1,f1,m1);
          }
        }
      }  // complete j
    }  // complete i
  }  // one parent is typed
}


void Genotypetrio::determin_phase()
{
  phase = 0;
  vector<pair<int, int> > hap_vec_o = offspring->get_hap_vec();
  if(hap_vec_o.size()==1)
    phase = 1;
  else if(hap_vec_o.size()>1)
  {
    if(inh_vec.size()==1)
      phase = 1;
    else if(inh_vec.size()>1)
    {
      int flag = 0;
      int a0 = (inh_vec[0].hp_offspring).first;
      int a1 = (inh_vec[0].hp_offspring).second;
      for(int i=1;i<inh_vec.size();i++)
      {
        if(a0!=(inh_vec[i].hp_offspring).first || a1!=(inh_vec[i].hp_offspring).second)
        {
          flag = 1;
          break;
        }
      }
      if(flag==0)
        phase = 1;
    }
  }
}


void Genotypetrio::expected_hap_num(map<int, double> & hap_freq_table, double *freq, map<pair<int, int>, double> & geno_freq_table, map<pair<int, int>, int> & hap_pair_map, int **count, const vector<int> & sites)
{
  hap_num.clear();

  int loc;

  if(inh_vec.size()>0)
  {
    if(phase==1)
    {
      pair<int, int> hp_offspring = inh_vec[0].hp_offspring;
      loc = hap_pair_map[hp_offspring];
      for(int j=0;j<sites.size();j++)
        hap_num.push_back(double(count[sites[j]][loc])/2.0-freq[sites[j]]);
    }
    else
    {
      for(int j=0;j<sites.size();j++)
        hap_num.push_back(0);

      double value;
      double Pi = 0;
      if(typed_par1==1 && typed_par2==1)
      {
        for(int j=0;j<inh_vec.size();j++)
        {
          pair<int, int> hp_offspring = inh_vec[j].hp_offspring;
          pair<int, int> hp_father = inh_vec[j].hp_par1;
          pair<int, int> hp_mother = inh_vec[j].hp_par2;

          value = geno_freq_table[hp_father]*geno_freq_table[hp_mother]*transProb(hp_offspring,hp_father,hp_mother,hap_freq_table,geno_freq_table,1);
          Pi+=value;

          loc = hap_pair_map[hp_offspring];
          for(int k=0;k<sites.size();k++)
          {
            if(count[sites[k]][loc]>0)
              hap_num[k]+=count[sites[k]][loc]*value;
          }
        }
      }
      else if(typed_par1==1 && typed_par2==0)
      {
        for(int j=0;j<inh_vec.size();j++)
        {
          pair<int, int> hp_offspring = inh_vec[j].hp_offspring;
          pair<int, int> hp_father = inh_vec[j].hp_par1;

          value = geno_freq_table[hp_father]*transProbonep(hp_offspring,hp_father,hap_freq_table,geno_freq_table,1);
          Pi+=value;

          loc = hap_pair_map[hp_offspring];
          for(int k=0;k<sites.size();k++)
          {
            if(count[sites[k]][loc]>0)
              hap_num[k]+=count[sites[k]][loc]*value;
          }
        }
      }
      else if(typed_par1==0 && typed_par2==0)
      {
        for(int j=0;j<inh_vec.size();j++)
        {
          pair<int, int> hp_offspring = inh_vec[j].hp_offspring;

          Pi+=geno_freq_table[hp_offspring];

          loc = hap_pair_map[hp_offspring];
          for(int k=0;k<sites.size();k++)
          {
            if(count[sites[k]][loc]>0)
              hap_num[k]+=count[sites[k]][loc]*geno_freq_table[hp_offspring];
          }
        }
      }

      for(int j=0;j<sites.size();j++)
        hap_num[j]=hap_num[j]/(2.0*Pi)-freq[sites[j]];
    }
  }
}


Genotypetrio::~Genotypetrio()
{
  inh_vec.clear();
  inh_vec_complete.clear();
  hap_num.clear();
}


int is_in(const vector<pair<int, int> > & hap_vec, int o1, int o2)
{
  for(int i=0;i<hap_vec.size();i++)
  {
    if(hap_vec[i].first==o1 && hap_vec[i].second==o2)
      return 1;
  }
  return 0;
}