#include "hfestimator.h"


HFEstimator::HFEstimator(vector<Family *> & fam_list, vector<Genotype *> & geno_list, map<int, double> & hap_list, int size)
:fam_list(fam_list), geno_list(geno_list), hap_freq_table(hap_list), loci_num(size){

  sample_size=0;
  effect_size=0;

  for(int i=0;i<fam_list.size();i++){
    for(int j=0;j<fam_list[i]->N;j++)
      effect_size+=fam_list[i]->weight[j];
    sample_size+=fam_list[i]->N_typed;
  }

  init_freq_table();
}


void HFEstimator::init_freq_table(){

  int tot;
  if(hap_freq_table.size()==0){
    tot = 1<<loci_num;
    for(int i=0;i<tot;i++)
      hap_freq_table[i] = 1.0/double(tot);
  }
  else{
    tot = hap_freq_table.size();
    map<int, double>::iterator iter=hap_freq_table.begin();
    while(iter!=hap_freq_table.end()){
      iter->second = 1.0/double(tot);
      iter++;
    }
  }
  
  geno_freq_table.clear();
  map<int, double>::iterator iter1=hap_freq_table.begin();
  map<int, double>::iterator iter2;
  while(iter1!=hap_freq_table.end()){
    iter2 = iter1;
    while(iter2!=hap_freq_table.end()){
      int h1 = iter1->first;
      int h2 = iter2->first;

      if(h1 == h2)
        geno_freq_table[pair<int, int>(h1,h2)] = iter1->second*iter2->second;
      else
        geno_freq_table[pair<int, int>(h1,h2)] = 2*iter1->second*iter2->second;

      iter2++;
    }
    iter1++;
  }
}


void HFEstimator::reset_freq_table(map<int,double> & hap_freq){

  hap_freq_table.clear();
  hap_freq_table = hap_freq;
  
  geno_freq_table.clear();
  map<int, double>::iterator iter1=hap_freq_table.begin();
  map<int, double>::iterator iter2;
  while(iter1!=hap_freq_table.end()){
    iter2 = iter1;
    while(iter2!=hap_freq_table.end()){
      int h1 = iter1->first;
      int h2 = iter2->first;

      if(h1 == h2)
        geno_freq_table[pair<int, int>(h1,h2)] = iter1->second*iter2->second;
      else
        geno_freq_table[pair<int, int>(h1,h2)] = 2*iter1->second*iter2->second;

      iter2++;
    }
    iter1++;
  }
}


void HFEstimator::show_incom_geno_freq(){

  for(int i=0;i<geno_list.size();i++){
    int *gcode = geno_list[i]->get_gcode();
    printf(" ");
    for(int j=0;j<loci_num;j++){
      if(gcode[j]==-1)
	    printf("*");
      else
	    printf("%d", gcode[j]);
      if(j!=loci_num-1)
	    printf("-");
    }

    double freq = 0;
    const vector<pair<int, int> > hv = geno_list[i]->get_hap_vec();
    for(int j=0;j<hv.size();j++)
      freq+=geno_freq_table[hv[j]];
    
    printf(" %6d ", geno_list[i]->get_num());
    printf(" %.3f \n", freq);
  }
  printf("\n");
}
    

void HFEstimator::show_hap_freq(){

  map<int, double>::iterator iter = hap_freq_table.begin();
  while(iter!=hap_freq_table.end()){
    if(iter->second<0.00005){
      iter++;
      continue;
    }
    printf("%2d ", iter->first);
    printf("   %.4f\n", iter->second);
    iter++;
  }
  printf("\n");
}


void HFEstimator::show_geno_freq(){

  map<pair<int, int>, double>::iterator giter = geno_freq_table.begin();
  while(giter!=geno_freq_table.end()){
    if(giter->second<0.00005){
      giter++;
      continue;
    }
    printf("(%2d,%2d) ", (giter->first).first,(giter->first).second);
    printf("   %.4f\n", giter->second);
    giter++;
  }
  printf("\n");
}
  

// a single EM iteration, index=0, EW; index=1, CDW
void HFEstimator::EM_Iteration(int index){
  // maximazation step

  // init the updated genotype frequency table
  map<pair<int, int>, double> ngf;
  map<pair<int, int>, double>::iterator giter = geno_freq_table.begin();

  while(giter!=geno_freq_table.end()){
    ngf[giter->first] = 0;
    giter++;
  }

  // compute phenotype frequency with current genotype frequency
  if(index==0){
    // EW
    for(int i=0;i<geno_list.size();i++){
      if(geno_list[i]->info==1){
        const vector<pair<int, int> > hv = geno_list[i]->get_hap_vec();
        double Pi = 0;  // record observed genotype frequency
        int ni = geno_list[i]->get_num();
        
        if(ni>0){
		  for(int j=0;j<hv.size();j++){
            pair<int, int> conf = hv[j];
            Pi+=geno_freq_table[conf];
          }
    
          for(int j=0;j<hv.size();j++){
            pair<int, int> conf = hv[j];
            ngf[conf] += double(ni)*geno_freq_table[conf]/Pi;
          }
        }
      }
    }
  }
  else if(index==1){
    // CDW
    for(int i=0;i<geno_list.size();i++){
      if(geno_list[i]->info==1){
        const vector<pair<int, int> > hv = geno_list[i]->get_hap_vec();
	    double Pi = 0;  // record observed genotype frequency
	    double ni = geno_list[i]->get_count();
      
	    if(ni>1e-8 || ni<-1e-8){
          for(int j=0;j<hv.size();j++){
            pair<int, int> conf = hv[j];
            Pi+=geno_freq_table[conf];
          }
    
          for(int j=0;j<hv.size();j++){
            pair<int, int> conf = hv[j];
            ngf[conf] += ni*geno_freq_table[conf]/Pi;
          }
		}
      }
    }
  }
  
  double sum = 0;
  giter = geno_freq_table.begin();
  while(giter!=geno_freq_table.end()){
    sum+=ngf[giter->first];
    giter++;
  }

  giter = geno_freq_table.begin();
  while(giter!=geno_freq_table.end()){
    giter->second = ngf[giter->first]/sum;
    giter++;
  }
  
  update_hap_freq();
  
  // Expection step  
  update_geno_freq();

  ngf.clear();
}


void HFEstimator::update_hap_freq(){

  map<int, double>::iterator iter = hap_freq_table.begin();
  while(iter!=hap_freq_table.end()){
    iter->second = 0;
    iter++;
  }
    
  map<pair<int, int>, double>::iterator giter = geno_freq_table.begin();
  
  while(giter!= geno_freq_table.end()){
    pair<int, int> conf = giter->first;
    if(conf.first==conf.second)
      hap_freq_table[conf.first] += geno_freq_table[conf];
    else{
      hap_freq_table[conf.first] += .5*geno_freq_table[conf];
      hap_freq_table[conf.second] += .5*geno_freq_table[conf];
    }
    giter++;
  }
}


void HFEstimator::update_geno_freq(){

  map<pair<int, int>, double>::iterator giter = geno_freq_table.begin();
  while(giter!=geno_freq_table.end()){
    giter->second = 0;
    giter++;
  }

  giter = geno_freq_table.begin();
  while(giter!=geno_freq_table.end()){
    pair<int, int> conf = giter->first;
    int k1 = conf.first;
    int k2 = conf.second;
    if(k1==k2)
      giter->second = hap_freq_table[k1]*hap_freq_table[k2];
    else
      giter->second = 2*hap_freq_table[k1]*hap_freq_table[k2];
    
    giter++;
  }
}


HFEstimator::~HFEstimator(){

  hap_freq_table.clear();
  geno_freq_table.clear();
}