#include "missing.h"


Missing::Missing(const vector<int> & missing_site, const map<int, double> & hap_list, const vector<char *> & content, int size){

  if(size==0)
    return;
  
  this->size = size;  // # of SNPs
  this->count=0;
  this->ind_clus=0;
  this->ind_1df=0;

  gcode = new int[size];
  memset(gcode,0,size*sizeof(int));
  
  for(int i=0;i<missing_site.size();i++)
    gcode[missing_site[i]] = -1;

  vector<int*> *vec;
  if(missing_site.size()==size){
    vec = new vector<int*>();
    vec->push_back(gcode);
  }
  else
    vec = complete_missing_genotype(0);

  int vsize = vec->size();
  for(int i=0;i<vsize;i++){
    Genotype *geno = new Genotype((*vec)[i],hap_list,size);
    const vector<pair<int, int> > hv = geno->get_hap_vec();
    if(hv.size()>0)
      missing_pat_genotype.push_back(geno);
    else
      delete geno;
  }

  if(vsize>1){
    for(int k=0;k<vsize;k++)
      delete[] (*vec)[k];
  }
  vec->clear();
  delete vec;

  keep(missing_site,hap_list,content);
  dim = sites.size();

  int row=hap_list.size();
  if(row==0)
    row=1<<size;
  int col = missing_pat_genotype.size();

  if(dim>0){
    weight = new double*[dim];
    for(int i=0;i<dim;i++){
      weight[i] = new double[col];
      memset(weight[i],0,col*sizeof(double));
    }

    gfreq = new double[col];
    memset(gfreq,0,col*sizeof(double));

    centmatrix = new double*[dim];
    for(int i=0;i<dim;i++){
      centmatrix[i] = new double[dim];
      memset(centmatrix[i],0,dim*sizeof(double));
    }

    derivative = new double*[dim];
    for(int i=0;i<dim;i++){
      derivative[i] = new double[row-1];
      memset(derivative[i],0,(row-1)*sizeof(double));
    }

    for(int i=0;i<col;i++){
      vector<int*> *complete_geno_code = missing_pat_genotype[i]->get_geno_code();
      int complete_geno = complete_geno_code->size();
      for(int j=0;j<complete_geno;j++){
        int code = 0;
        for(int k=0;k<size;k++){
          int gc = (*complete_geno_code)[j][k];
          code = 10*code+gc;
        }
        mapping_function[code] = i;
      }
    }
  }
}
  

vector<int*> *Missing::complete_missing_genotype(int index){

  int code = gcode[index];

  if(index==size-1){
    vector<int*> *vec = new vector<int*>();
    
    if(code==-1){
      int *a1 = new int[size];
      memset(a1,0,size*sizeof(int));
       
      a1[size-1] = code;
      vec->push_back(a1);
    }
    else if(code==0){
      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_missing_genotype(index+1);

  int pos = index;
  
  if(code==-1){
    for(int i=0;i<vec->size();i++)
      (*vec)[i][pos] = code;
    
    return vec;  
  }
  
  // else if code == 0
  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 Missing::set_weight(map<pair<int, int>, double> & geno_freq_table, map<pair<int, int>, int> & hap_pair_map, int **count){

  int col = missing_pat_genotype.size();

  for(int i=0;i<dim;i++)
    memset(weight[i],0,col*sizeof(double));
  memset(gfreq,0,col*sizeof(double));

  int loc;

  for(int i=0;i<missing_pat_genotype.size();i++){
    const vector<pair<int, int> > hv = missing_pat_genotype[i]->get_hap_vec();
    if(hv.size()==1){
      gfreq[i] = geno_freq_table[hv[0]];

      loc = hap_pair_map[hv[0]];
      for(int j=0;j<dim;j++){
        if(count[sites[j]][loc]>0)
          weight[j][i] = double(count[sites[j]][loc])/2.0;
      }
    }
    else if(hv.size()>1){
      for(int j=0;j<hv.size();j++){
        gfreq[i]+=geno_freq_table[hv[j]];

        loc = hap_pair_map[hv[j]];
        for(int k=0;k<dim;k++){
          if(count[sites[k]][loc]>0)
            weight[k][i]+=count[sites[k]][loc]*geno_freq_table[hv[j]];
        }
      }

      for(int j=0;j<dim;j++)
        weight[j][i]/=(2.0*gfreq[i]);
    }
  }
}


void Missing::compute_centmatrix(double *freq, map<pair<int, int>, double> & geno_freq_table){

  for(int i=0;i<dim;i++)
    memset(centmatrix[i],0,dim*sizeof(double));

  nonzero.clear();
  for(int i=0;i<missing_pat_genotype.size();i++){
    for(int j=0;j<dim;j++){
      if(weight[j][i]>0){
        nonzero.push_back(i);
        break;
      }
    }
  }

  for(int i=0;i<nonzero.size();i++){
    for(int j=0;j<dim;j++){
      centmatrix[j][j]+=weight[j][nonzero[i]]*weight[j][nonzero[i]]*gfreq[nonzero[i]];
      for(int k=j+1;k<dim;k++)
        centmatrix[j][k]+=weight[j][nonzero[i]]*weight[k][nonzero[i]]*gfreq[nonzero[i]];
    }
  }

  for(int i=0;i<dim;i++){
    centmatrix[i][i]-=freq[sites[i]]*freq[sites[i]];
    for(int j=i+1;j<dim;j++){
      centmatrix[i][j]-=freq[sites[i]]*freq[sites[j]];
      centmatrix[j][i] = centmatrix[i][j];
    }
  }
}


void Missing::compute_derivative(map<int, double> & hap_freq_table, map<pair<int, int>, int> & hap_pair_map, int **count, double **rst){

  int row = hap_freq_table.size();

  for(int i=0;i<dim;i++)
    memset(derivative[i],0,(row-1)*sizeof(double));

  int loc;

  double *matrix1 = new double[row-1];
  memset(matrix1,0,(row-1)*sizeof(double));
  double **matrix2 = new double*[dim];
  for(int j=0;j<dim;j++){
    matrix2[j] = new double[row-1];
    memset(matrix2[j],0,(row-1)*sizeof(double));
  }

  for(int i=0;i<missing_pat_genotype.size();i++){
    const vector<pair<int, int> > hv = missing_pat_genotype[i]->get_hap_vec();
    if(hv.size()>1){
      memset(matrix1,0,(row-1)*sizeof(double));
      for(int j=0;j<dim;j++)
        memset(matrix2[j],0,(row-1)*sizeof(double));

      for(int j=0;j<hv.size();j++){
        loc = hap_pair_map[hv[j]];
        for(int k=0;k<row-1;k++)
          matrix1[k]+=rst[k][loc];

        for(int k=0;k<dim;k++){
          if(count[sites[k]][loc]>0){
            for(int l=0;l<row-1;l++)
              matrix2[k][l]+=count[sites[k]][loc]*rst[l][loc];
          }
        }
      }

      for(int j=0;j<dim;j++)
        for(int k=0;k<row-1;k++)
          derivative[j][k]-=(matrix2[j][k]/2.0-weight[j][i]*matrix1[k]);
    }
  }

  for(int i=0;i<dim;i++)
    derivative[i][sites[i]]+=1;

  delete[] matrix1;
  clear_matrix(matrix2,dim,row-1);
}


void Missing::compute_derivative_cluster(map<int, double> & hap_freq_table, double **mul_clus, int col_clus){

  if(dim>0){
    ind_clus = 1;
    int row = hap_freq_table.size();
    derivative_clus = multiply(derivative,dim,row-1,mul_clus,row-1,col_clus);
  }
}


void Missing::compute_derivative_1df(map<int, double> & hap_freq_table, double **mul_1df, int col_1df){

  if(dim>0){
    ind_1df = 1;
    int row = hap_freq_table.size();
    derivative_1df = multiply(derivative,dim,row-1,mul_1df,row-1,col_1df);
  }
}


void Missing::keep(const vector<int> & missing_site, const map<int, double> & hap_list, const vector<char *> & content){

  sites.clear();

  int tot = 1<<size;

  if(missing_site.size()==0){
    int full_dim = hap_list.size();
    if(full_dim==0)
      full_dim = tot;

    for(int i=0;i<full_dim-1;i++)
      sites.push_back(i);
  }
  else if(missing_site.size()>0 && missing_site.size()<size){
    map<int, int> hcom_hinc;

    if(hap_list.size()==0 || hap_list.size()==tot){
      int *hcode = new int[size];
      memset(hcode,0,size*sizeof(int));

      vector<int*> *vec=complete_missing_haplotype(hcode,0);

      for(int i=0;i<tot;i++){
        int h_full = 0;
        int h_part = 0;
        int *a = (*vec)[i];
        for(int j=0;j<size;j++){
          if(a[j]==1){
            h_full+=(1<<j);
            if(!find_element(missing_site,j))
              h_part+=(1<<j);
          }
        }

        hcom_hinc[h_full] = h_part;
      }

      delete[] hcode;

      for(int k=0;k<tot;k++)
        delete[] (*vec)[k];

      vec->clear();
      delete vec;
    }
    else{
      int hap_num = content.size();

      for(int i=0;i<hap_num;i++){
        int h_full = 0;
        int h_part = 0;
        for(int j=0;j<size;j++){
          if(content[i][j]=='1'){
            h_full+=(1<<j);
            if(!find_element(missing_site,j))
              h_part+=(1<<j);
          }
        }

        hcom_hinc[h_full] = h_part;
      }
    }

    map<int, int> e_map;
    int loc=0;
    map<int, int>::iterator iter = hcom_hinc.begin();
    while(iter!=hcom_hinc.end()){
      int h_part = iter->second;
      if(e_map.find(h_part)==e_map.end())
        e_map[h_part] = loc;

      loc++;
      iter++;
    }

    loc=0;
    iter = e_map.begin();
    while(loc!=e_map.size()-1){
      sites.push_back(iter->second);
      iter++;
      loc++;
    }
  }
}


vector<int*> *Missing::complete_missing_haplotype(int *hcode, int index){

  int code = hcode[index];

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

    return vec;
  }
  
  vector<int*> *vec = complete_missing_haplotype(hcode,index+1);

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

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


Missing::~Missing(){

  delete[] gcode;
  for(int i=0;i<missing_pat_genotype.size();i++)
    delete missing_pat_genotype[i];
  missing_pat_genotype.clear();

  if(dim>0){
    for(int i=0;i<dim;i++){
      delete[] weight[i];
      delete[] centmatrix[i];
      delete[] derivative[i];
    }

    delete[] weight;
    delete[] gfreq;
    delete[] centmatrix;
    delete[] derivative;

    sites.clear();
    nonzero.clear();
    mapping_function.clear();

    if(ind_clus==1){
      for(int i=0;i<dim;i++)
        delete[] derivative_clus[i];
      delete[] derivative_clus;
    }

    if(ind_1df==1){
      for(int i=0;i<dim;i++)
        delete[] derivative_1df[i];
      delete[] derivative_1df;
    }
  }
}


int find_element(const vector<int> & s_vec, int m){

  for(int i=0;i<s_vec.size();i++){
    if(s_vec[i]==m)
      return 1;
  }
  return 0;
}