#include "genotype.h"


Genotype::Genotype(int *input_code, const map<int, double> & hap_list, int size){

  if(input_code==0 || size==0)
    return;
  
  this->size = size;  // # of SNPs

  gcode = new int[size];
  memset(gcode,0,size*sizeof(int));
  
  int flag = 0;  // if there is missing genotype at each locus
  for(int i=0;i<size;i++){
    gcode[i]=input_code[i];
    if(input_code[i]==-1)
      flag = 1;  // missing
  }

  if(flag==1)
    geno_code = complete_genotype(0); // enumerate all possible complete genotype that are compatible with the observed genotype
  else{
    geno_code = new vector<int*>();
    geno_code->push_back(gcode);
  }

  this->weighted_count=0;
  this->count=0;
  this->case_num=0;
  this->control_num=0;

  record_hap_structure(hap_list);
}
  

vector<int*> * Genotype::complete_genotype(int index){

  int code = gcode[index];

  if(index==size-1){
    vector<int*> *vec = new vector<int*>();
    
    if(code==0 || code==1 || code==2){
      int *a1 = new int[size];
      memset(a1,0,size*sizeof(int));
       
      a1[size-1] = code;
      vec->push_back(a1);
    }
    else if(code==-1){
      int *a1 = new int[size];
      int *a2 = new int[size];
      int *a3 = new int[size];
      memset(a1,0,size*sizeof(int));
      memset(a2,0,size*sizeof(int));
      memset(a3,0,size*sizeof(int));
      a1[size-1]=0;
      a2[size-1]=1;
      a3[size-1]=2;
      vec->push_back(a1);
      vec->push_back(a2);
      vec->push_back(a3);
    }

    return vec;
  }
  
  vector<int*> *vec = complete_genotype(index+1);

  int pos = index;
  
  if(code==0 || code==1 || code==2){
    for(int i=0;i<vec->size();i++)
      (*vec)[i][pos] = code;
    
    return vec;  
  }
  
  // else if code == -1
  int vsize = vec->size();
  for(int i=0;i<vsize;i++){
    int *a1 = new int[size];
    int *a2 = new int[size];
    memset(a1,0,size*sizeof(int));
    memset(a2,0,size*sizeof(int));
    
    for(int j=pos+1;j<size;j++)
      a2[j]=a1[j]=(*vec)[i][j];
    
    a1[pos]=1;
    a2[pos]=2;
    (*vec)[i][pos]=0;
    vec->push_back(a1);
    vec->push_back(a2);
  }
  
  return vec;
}


void Genotype::show_gstat(){
  // print gcode with its count
  for(int i=0;i<size;i++){
    if(gcode[i]!=-1)
      printf("%d", gcode[i]);
    else
      printf("*");
    if(i!=size-1)
      printf("-");
  }
  
  printf(": %d \n", count);
}


void Genotype::show_content(){
  // print geno_code and hap_vec
  printf("\nGenotype:\n");
  for(int i=0;i<geno_code->size();i++){
    for(int j=0;j<size;j++){
      printf("%d", (*geno_code)[i][j]);
      if(j!=size-1)
	    printf("-");
      else 
	    printf("\n");
    }
  }
  
  printf("\nHaplotype:\n");
  for(int i=0;i<hap_vec.size();i++){
    int h1 = hap_vec[i].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_vec[i].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");
    }
    printf("\n");
  }
}


vector<int*> * Genotype::hap_decompose(int *geno, int index, int &bit){
  // enumerate all possible haplotype pairs
  int code = geno[index];

  if(index==size-1){
    vector<int*> *vec = new vector<int*>();

    int *a1 = new int[size];
    int *a2 = new int[size];
    memset(a1,0,size*sizeof(int));
    memset(a2,0,size*sizeof(int));

    if(code==0)
      a1[size-1]=a2[size-1]=0;
    else if(code==1){
      a1[size-1]=0;
      a2[size-1]=1;
      bit = 1;  // record heterozygous sites
    }
    else if(code==2)
      a1[size-1]=a2[size-1]=1;
    
    vec->push_back(a1);
    vec->push_back(a2);
    return vec;
  }
  
  vector<int*> *vec = hap_decompose(geno,index+1,bit);
  
  int pos = index;
  
  // first heterozygous genotype, no permutation, reset bit flag
  if(code==1 && bit==0){
    for(int i=0;i<vec->size();i++)
      (*vec)[i][pos] = i%2;
    bit = 1;
    return vec;
  }
  if(code==0 || code==2){
    for(int i=0;i<vec->size();i++)
      (*vec)[i][pos] = code/2;
    return vec;
  }
  
  // else if code == 1 and not the first heterozygous site
  int vsize = vec->size();
  for(int i=0;i<vsize;i+=2){
    int *a1 = new int[size];
    int *a2 = new int[size];
    memset(a1,0,size*sizeof(int));
    memset(a2,0,size*sizeof(int));

    for(int j=pos+1;j<size;j++){
      a1[j]=(*vec)[i][j];
      a2[j]=(*vec)[i+1][j];
    }
    
    a1[pos]=1;
    a2[pos]=0;
    (*vec)[i][pos]=0;
    (*vec)[i+1][pos]=1;
    vec->push_back(a1);
    vec->push_back(a2);
  }

  return vec;
}


void Genotype::record_hap_structure(const map<int, double> & hap_list){

  int tot = 1<<size;

  if(hap_list.size()==0 || hap_list.size()==tot){
    for(int k=0;k<geno_code->size();k++){
      int bit=0;
      vector<int*> *rst = hap_decompose((*geno_code)[k],0,bit);
      for(int i=0;i<rst->size();i+=2){
        int h0=0;
        int h1=0;
        for(int j=0;j<size;j++){
	      if((*rst)[i][j]==1)
	        h0+=(1<<j);
	      if((*rst)[i+1][j]==1)
	        h1+=(1<<j);
        }

        // release heap memory here
        delete[] (*rst)[i];
        delete[] (*rst)[i+1];
        int min = h0<=h1 ? h0 : h1;
        int max = h0+h1-min;  // order the two haplotype with a pair
        hap_vec.push_back(pair<int, int>(min,max));  // record into hap_vec
      }
      // clean up the container
      rst->clear();
      delete rst;
    }

    if(hap_vec.size()==0 || hap_vec.size()==tot*(tot+1)/2)
      info = 0;
    else
      info = 1;
  }
  else{
    for(int k=0;k<geno_code->size();k++){
      int bit=0;
      vector<int*> *rst = hap_decompose((*geno_code)[k],0,bit);
      for(int i=0;i<rst->size();i+=2){
        int h0=0;
        int h1=0;
        for(int j=0;j<size;j++){
	      if((*rst)[i][j]==1)
	        h0+=(1<<j);
	      if((*rst)[i+1][j]==1)
	        h1+=(1<<j);
        }

        // release heap memory here
        delete[] (*rst)[i];
        delete[] (*rst)[i+1];
        int min = h0<=h1 ? h0 : h1;
        int max = h0+h1-min;  // order the two haplotype with a pair

        if(hap_list.find(min)!=hap_list.end() && hap_list.find(max)!=hap_list.end())
          hap_vec.push_back(pair<int, int>(min,max));  // record into hap_vec
      }
      // clean up the container
      rst->clear();
      delete rst;
    }

    if(hap_vec.size()==0 || hap_vec.size()==hap_list.size()*(hap_list.size()+1)/2)
      info = 0;
    else
      info = 1;
  }
}


void Genotype::expected_hap_num(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(hap_vec.size()==1){
    loc = hap_pair_map[hap_vec[0]];
    for(int j=0;j<sites.size();j++)
      hap_num.push_back(double(count[sites[j]][loc])/2.0-freq[sites[j]]);
  }
  else if(hap_vec.size()>1){
    for(int j=0;j<sites.size();j++)
      hap_num.push_back(0);

    double Pi = 0;
    for(int j=0;j<hap_vec.size();j++){
      Pi+=geno_freq_table[hap_vec[j]];

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

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


void Genotype::expectation(map<pair<int, int>, double> & geno_freq_table){

  prob = 0;
  if(hap_vec.size()==1)
    prob = geno_freq_table[hap_vec[0]];
  else if(hap_vec.size()>1){
    for(int j=0;j<hap_vec.size();j++)
      prob+=geno_freq_table[hap_vec[j]];
  }
}


void Genotype::expectation_IBD1(map<int, double> & hap_freq_table){

  prob_IBD1 = 0;
  if(hap_vec.size()==1)
    prob_IBD1 = expfortwo(hap_vec[0],hap_freq_table);
  else if(hap_vec.size()>1){
    for(int i=0;i<hap_vec.size();i++){
      prob_IBD1+=expfortwo(hap_vec[i],hap_freq_table);
      for(int j=i+1;j<hap_vec.size();j++)
        prob_IBD1+=2*expfortwo(hap_vec[i],hap_vec[j],hap_freq_table);
    }
  }
}


void Genotype::compute_gname(){

  int code = 0;
  for(int i=0;i<size;i++)
    code = 10*code+gcode[i];
  gname = code;
}


Genotype::~Genotype(){

  if(geno_code->size()>1)
    delete[] gcode;
  
  for(int k=0;k<geno_code->size();k++)
    delete[] (*geno_code)[k];
  geno_code->clear();
  delete geno_code;

  hap_vec.clear();
  hap_num.clear();
}


void translate_hap(FILE *fp, int code, int size, vector<char> & a0v, vector<char> & a1v){

  for(int j=0;j<size;j++){
    int mask = 1<<j;
    int b1 = (code&mask)>>j;
    if(b1==0)
      fprintf(fp,"%c", a0v[j]);
    else if(b1==1)
      fprintf(fp,"%c", a1v[j]);
    if(j!=size-1)
      fprintf(fp,"-");
  }
}


void translate_hap(FILE *fp, int code, int size){

  for(int j=0;j<size;j++){
    int mask = 1<<j;
    int b1 = (code&mask)>>j;
    fprintf(fp, "%d", b1);
    if(j!=size-1)
      fprintf(fp,"-");
  }
}