#include "family.h"


Family::Family(vector<int *> & fam_member, int fam_id)
:fam_member(fam_member), fam_id(fam_id){

  this->N=fam_member.size();
  this->Affected=0;
  this->UnAffected=0;
  this->Unknown=0;

  for(int i=0;i<N;i++){
    int ind_id = fam_member[i][0];
    member_map[ind_id] = i;

    if(fam_member[i][4]==0)
      Unknown++;
    else if(fam_member[i][4]==1)
      UnAffected++;
    else if(fam_member[i][4]==2)
      Affected++;
  }
}


void Family::set_gcode(int nmark){

  this->nmark = nmark;
  gcode = new int*[N];
  for(int i=0;i<N;i++){
    gcode[i] = new int[nmark];
    memset(gcode[i],-1,nmark*sizeof(int));
  }

  typed = new int[N];
  memset(typed,0,N*sizeof(int));

  weight = new double[N];
  memset(weight,0,N*sizeof(double));
  test_weight = new double[N];
  memset(test_weight,0,N*sizeof(double));
  test_trio_weight1 = new double[N];
  memset(test_trio_weight1,0,N*sizeof(double));
  test_trio_weight2 = new double[N];
  memset(test_trio_weight2,0,N*sizeof(double));
  test_trio_weight3 = new double[N];
  memset(test_trio_weight3,0,N*sizeof(double));

  cat = new int[N];
  memset(cat,0,N*sizeof(int));
  mis_cat = new int[N];
  memset(mis_cat,0,N*sizeof(int));
  trio_cat = new int[N];
  memset(trio_cat,0,N*sizeof(int));
  mis_trio_cat = new int[N];
  memset(mis_trio_cat,0,N*sizeof(int));
  mis_par_order = new int[N];
  memset(mis_par_order,0,N*sizeof(int));
}


void Family::set_gcode(int *input_code, int loc){

  for(int i=0;i<nmark;i++)
    gcode[loc][i] = input_code[i];
}


void Family::set_gcode(){

  for(int i=0;i<N;i++)
    memset(gcode[i],-1,nmark*sizeof(int));
}


void Family::fill_typed(){  // how many individuals are genotyped

  N_typed=0;
  memset(typed,0,N*sizeof(int));

  for(int i=0;i<N;i++){
    int flag = 0;
    for(int j=0;j<nmark;j++){
      if(gcode[i][j]!=-1){
        flag = 1;
        N_typed++;
        break;
      }
    }
    typed[i] = flag;
  }
}


void Family::fill_info_gcode(vector<Genotype *> & geno_list){

  N_typed=0;
  memset(typed,0,N*sizeof(int));

  for(int i=0;i<N;i++){
    int flag = 0;
    if(geno_list[cat[i]]->info==1){
      flag = 1;
      N_typed++;
    }
    typed[i] = flag;
  }
}


void Family::print(){

  cout << "Family " << fam_id << " has " << N << " members.\n";
  cout << "  " << Affected << " affected individuals;\n";
  cout << "  " << UnAffected << " unaffected individuals;\n";
  cout << "  " << Unknown << " of unknown phenotype.\n\n";
}


void Family::compute_weight(){  // compute complete-data weight

  memset(weight,0,N*sizeof(double));

  if(N_typed>0){
    double **cholent = new double*[N_typed];
    for(int i=0;i<N_typed;i++){
      cholent[i] = new double[N_typed+1];
      memset(cholent[i],0,(N_typed+1)*sizeof(double));
    }

    double **chol = new double*[N_typed];
    for(int i=0;i<N_typed;i++){
      chol[i] = new double[N_typed];
      memset(chol[i],0,N_typed*sizeof(double));
    }
  
    double **cholaug = new double*[N_typed];
    for(int i=0;i<N_typed;i++){
      cholaug[i] = new double[N_typed+1];
      memset(cholaug[i],0,(N_typed+1)*sizeof(double));
    }

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

    // populate kinship coefficient matrix
    int nb1 = -1;
    int nb2 = -1;
    int ind_1;
    int ind_2;
    double value;
    map<pair<int, int>, double>::iterator iter;

    for(int i=0;i<N;i++){
      if(typed[i]==1){
        nb1++;
        nb2 = nb1;
        ind_1 = fam_member[i][0];
        ind_2 = ind_1;
        iter = kinship_coeff.find(pair<int, int>(ind_1,ind_2));
        if(iter == kinship_coeff.end()){
          printf("No IBD coefficient for individual %d from family %d. Please check...\n\n", ind_1,fam_id);
          exit(1);
        }
        kincoefMatrix[nb1][nb1] = 1;

        for(int j=i+1;j<N;j++){
          if(typed[j]==1){
            nb2++;
            ind_2 = fam_member[j][0];
            iter = kinship_coeff.find(pair<int, int>(ind_1,ind_2));
            if(iter == kinship_coeff.end()){
              printf("No IBD coefficient between individual %d and individual %d from family %d. Please check...\n\n", ind_1,ind_2,fam_id);
              exit(1);
            }
            value = iter->second;
            if(value>1e-8)
            {
              kincoefMatrix[nb1][nb2] = 2*value;
              kincoefMatrix[nb2][nb1] = 2*value;
            }
          }
        }
      }
    }

    for(int i=0;i<N_typed;i++){
      cholent[i][i] = 1;
      cholent[i][N_typed] = 1;
    }

    if(cholesky(kincoefMatrix,N_typed,cholent,N_typed+1,chol,cholaug,0) != 1){
      printf("cholesky decomposition of the cov matrix failed for family %d. Might be due to inconsistent IBD coefficient values...\n", fam_id);
      exit(1);
    }

    double **s = self_colmultiply(cholaug,N_typed,1,N_typed,N_typed+1,N_typed+1);
    int count = -1;
    for(int i=0;i<N;i++){
      if(typed[i] == 1){
        count++;
        weight[i] = s[count][0];
      }
    }

    clear_matrix(cholent,N_typed,N_typed+1);
    clear_matrix(chol,N_typed,N_typed);
    clear_matrix(cholaug,N_typed,N_typed+1);
    clear_matrix(kincoefMatrix,N_typed,N_typed);
    clear_matrix(s,N_typed,1);
  }
}


void Family::compute_test_weight(double mean){

  memset(test_weight,0,N*sizeof(double));

  int ind_1;
  int ind_2;
  double value;
  map<pair<int, int>, double>::iterator iter;

  for(int i=0;i<N;i++){
    ind_1 = fam_member[i][0];
    if(fam_member[i][4]==1 || fam_member[i][4]==2){
      ind_2 = ind_1;
      iter = kinship_coeff.find(pair<int, int>(ind_1,ind_2));      
      if(iter != kinship_coeff.end()){      
        if(fam_member[i][4]==1)
          test_weight[i]+=mean;
        else if(fam_member[i][4]==2)
          test_weight[i]+=1;
      }
    }

    for(int j=i+1;j<N;j++){
      if(fam_member[j][4]==1 || fam_member[j][4]==2 || fam_member[i][4]==1 || fam_member[i][4]==2){
        ind_2 = fam_member[j][0];
        iter = kinship_coeff.find(pair<int, int>(ind_1,ind_2));
        if(iter != kinship_coeff.end()){
          value = iter->second;
          if(value>1e-8){
            if(fam_member[j][4]==1)
              test_weight[i]+=2*value*mean;
            else if(fam_member[j][4]==2)
              test_weight[i]+=2*value;
            if(fam_member[i][4]==1)
              test_weight[j]+=2*value*mean;
            else if(fam_member[i][4]==2)
              test_weight[j]+=2*value;
          }
        }
      }
    }  // complete j
  }  // complete i
}


void Family::compute_test_trio_weight(double mean){

  memset(test_trio_weight1,0,N*sizeof(double));
  memset(test_trio_weight2,0,N*sizeof(double));
  memset(test_trio_weight3,0,N*sizeof(double));

  int ind_1, ind_1f, ind_1m;
  int ind_2;
  double value_o, value_f, value_m;
  map<pair<int, int>, double>::iterator iter;

  for(int i=0;i<N;i++){
    ind_1 = fam_member[i][0];
    ind_1f = fam_member[i][1];
    ind_1m = fam_member[i][2];
    if(ind_1f!=0 && ind_1m!=0){
      // i is not a founder, no need to compute for founders
      for(int j=0;j<N;j++){
        if(fam_member[j][4]==1 || fam_member[j][4]==2){
          ind_2 = fam_member[j][0];
          if(ind_2==ind_1){
            // j is i
            iter = kinship_coeff.find(pair<int, int>(ind_1,ind_2));
            if(iter != kinship_coeff.end()){              
              if(fam_member[j][4]==1)
                test_trio_weight1[i]+=mean;
              else if(fam_member[j][4]==2)
                test_trio_weight1[i]+=1;
            }
          }
          else if(ind_2==ind_1f){
            // j is if
            iter = kinship_coeff.find(pair<int, int>(ind_1f,ind_2));
            if(iter != kinship_coeff.end()){
              if(fam_member[j][4]==1)
                test_trio_weight2[i]+=mean;
              else if(fam_member[j][4]==2)
                test_trio_weight2[i]+=1;
            }
          }
          else if(ind_2==ind_1m){
            // j is im
            iter = kinship_coeff.find(pair<int, int>(ind_1m,ind_2));
            if(iter != kinship_coeff.end()){
              if(fam_member[j][4]==1)
                test_trio_weight3[i]+=mean;
              else if(fam_member[j][4]==2)
                test_trio_weight3[i]+=1;
            }
          }
          else{
            iter = kinship_coeff.find(pair<int, int>(ind_1,ind_2));
            if(iter != kinship_coeff.end()){
              value_o = iter->second;
              if(value_o>1e-8){
                if(fam_member[j][4]==1){
                  iter = kinship_coeff.find(pair<int, int>(ind_1f,ind_2));
                  if(iter != kinship_coeff.end()){
                    value_f = iter->second;
                    iter = kinship_coeff.find(pair<int, int>(ind_1m,ind_2));
                    if(iter != kinship_coeff.end()){
                      value_m = iter->second;

                      if(value_o>value_f && value_o>value_m)  // group 1
                        test_trio_weight1[i]+=2*value_o*mean;
                      else if(value_f>value_o && value_m<1e-8)  // group 2
                        test_trio_weight2[i]+=2*value_f*mean;
                      else if(value_m>value_o && value_f<1e-8)  // group 3
                        test_trio_weight3[i]+=2*value_m*mean;
                      else{
                        // group 4
                        test_trio_weight2[i]+=2*value_f*mean;
                        test_trio_weight3[i]+=2*value_m*mean;
                      }
                    }
                  }
                }
                else if(fam_member[j][4]==2){
                  iter = kinship_coeff.find(pair<int, int>(ind_1f,ind_2));
                  if(iter != kinship_coeff.end()){
                    value_f = iter->second;
                    iter = kinship_coeff.find(pair<int, int>(ind_1m,ind_2));
                    if(iter != kinship_coeff.end()){
                      value_m = iter->second;

                      if(value_o>value_f && value_o>value_m)  // group 1
                        test_trio_weight1[i]+=2*value_o;
                      else if(value_f>value_o && value_m<1e-8)  // group 2
                        test_trio_weight2[i]+=2*value_f;
                      else if(value_m>value_o && value_f<1e-8)  // group 3
                        test_trio_weight3[i]+=2*value_m;
                      else{
                        // group 4
                        test_trio_weight2[i]+=2*value_f;
                        test_trio_weight3[i]+=2*value_m;
                      }
                    }
                  }
                }
              }  // i and j related
            }
          }
        }  // j's phenotype is 1 or 2
      }  // complete j
    }  // i is not a founder
  }  // complete i
}


Family::~Family(){

  for(int i=0;i<N;i++)
    delete[] fam_member[i];
  fam_member.clear();

  delete[] typed;
  delete[] weight;
  delete[] test_weight;
  delete[] test_trio_weight1;
  delete[] test_trio_weight2;
  delete[] test_trio_weight3;

  kinship_coeff.clear();
  IBD2_coeff.clear();
  IBD1_coeff.clear();
  IBD0_coeff.clear();
  member_map.clear();

  delete[] cat;
  delete[] mis_cat;
  delete[] trio_cat;
  delete[] mis_trio_cat;
  delete[] mis_par_order;

  for(int i=0;i<N;i++)
    delete[] gcode[i];
  delete[] gcode;
}