#include "includes.h"
#include "grid.h"
#include "vector.h"

#include <fstream.h>
extern ofstream errorfs;

#include "declare.h"

/*
  This file contains source code for functions used during the mapping 
  routine, mainly mundane things such as adjusting the datasets when
  we change the assumed location of the gene
*/


void augmentdata(const igrid& data, igrid& augdata, 
		 const igrid& freqallele, igrid& augfreqallele,
		 const dgrid& freqfreq, dgrid& augfreqfreq, 
		 const dgrid& jointfreq2, dgrid& augjointfreq2, 
		 //const dvector& y, dvector& augy, 
		 const ivector& allelecounts, ivector& augallelecounts,
		 const dvector& mut_markers, dvector& augmut_markers,
		 const int &C, const int &ctrall, 
		 const int &debug)
{
  int i, j;
  
  //make augdata
  for (i=0; i<data.dimx(); i++) {
    for (j=0; j<C; j++)
      augdata(i, j) = data(i, j);
    augdata(i, C) = ctrall;
    for (j=C; j<data.dimy(); j++)
      augdata(i, j+1) = data(i, j); }
  
  if (debug==1) {
    cout << "augdata" << endl;
    for (i=0;i<augdata.dimx();i++) {
      for (j=0;j<augdata.dimy();j++) {
	cout << augdata(i,j) << " "; }
      cout << endl; }
  }

  //make augfreqallele
  for (i=0; i<freqallele.dimx(); i++) {
    for (j=0; j<C; j++)
      augfreqallele(i, j) = freqallele(i, j);
    augfreqallele(i, C) = 0;
    for (j=C; j<freqallele.dimy(); j++)
      augfreqallele(i, j + 1) = freqallele(i, j);
  }
  augfreqallele(0, C) = ctrall;

  if (debug==1) {
    cout << "augfreqallele" << endl;
    for (i=0;i<augfreqallele.dimx();i++) {
      for (j=0;j<augfreqallele.dimy();j++) {
	cout << augfreqallele(i,j) << " "; }
      cout << endl; }
  }

  //make augfreqfreq
  for (i=0; i<freqfreq.dimx(); i++) {
    for (j = 0; j < C; ++j)
      augfreqfreq(i, j) = freqfreq(i, j);
    augfreqfreq(i, C) = 0;
    for (j = C; j < freqfreq.dimy(); j++)
      augfreqfreq(i, j + 1) = freqfreq(i, j); }
  augfreqfreq(0, C) = 1.0;
  
  if (debug==1) {
    cout << "augfreqfreq" << endl;
    for (i=0;i<augfreqfreq.dimx();i++) {
      for (j=0;j<augfreqfreq.dimy();j++) {
	cout << augfreqfreq(i,j) << " "; }
      cout << endl; }
  }

  //make augjointfreq2
  //make first C-2 columns
  for (j = 0; j < C-1; j++) {
    for (i=0; i< jointfreq2.dimx(); i++) {
      augjointfreq2(i,j)=jointfreq2(i,j); } }
  //make column C-1
  for (i=0; i<jointfreq2.dimx();i++) {
    augjointfreq2(i,C-1)=0; }
  for (i=0;i<allelecounts(C-1);i++) {
    augjointfreq2(i,C-1)=freqfreq(i,C-1); }
  //make column C
  for (i=0; i<jointfreq2.dimx();i++) {
    augjointfreq2(i,C)=0; }
  for (i=0;i<allelecounts(C);i++) {
    augjointfreq2(i,C)=freqfreq(i,C); }
  //make rest of object
  for (j=C; j<jointfreq2.dimy(); j++) {
    for (i=0; i< jointfreq2.dimx(); i++) {
      augjointfreq2(i,j+1)=jointfreq2(i,j); } } 

  if (debug==1) {
    cout << "augjointfreq2" << endl;
    for (i=0;i<augjointfreq2.dimx();i++) {
      for (j=0;j<augjointfreq2.dimy();j++) {
	cout << augjointfreq2(i,j) << " "; }
      cout << endl; }
  }


  //make augy
  //  for (j = 0; j < C; j++)
  //augy(j) = y(j);
  //for (j = C; j < y.dim(); j++)
  //augy(j + 1) = y(j);
  //augy(C) = augy(C - 1) = y(C - 1) / 2.0;
  //
  //if (debug==1) {
  //cout << "augy" << endl;
  //for (i=0;i<augy.dim();i++) cout << augy(i) << " ";
  //cout << endl; }

  //make augallelecounts
  for (j=0;j<C;j++) augallelecounts(j)=allelecounts(j);
  augallelecounts(C)=1;
  for (j=C;j<allelecounts.dim();j++) augallelecounts(j+1)=allelecounts(j);
  
  if (debug==1) {
    cout << "augallelecounts" << endl;
    for (i=0;i<augallelecounts.dim();i++) cout << augallelecounts(i) << " ";
    cout << endl; }

  //make augmut_markers
  for (j=0;j<C;j++) augmut_markers(j)=mut_markers(j);
  augmut_markers(C)=0;
  for (j=C;j<mut_markers.dim();j++) augmut_markers(j+1)=mut_markers(j);

  if (debug==1) {
    cout << "augmut_markers" << endl;
    for (i=0;i<augmut_markers.dim();i++) cout << augmut_markers(i) << " ";
    cout << endl; }

  
} // augmentdata

void makejointfreq3(dvector &jointfreq3, const dgrid &jointfreq2, 
		    const int &C, const int &debug)
{
  //this function makes the vector of 3-locus haps freqs. around C, from
  //the grid of 2-marker hap frequencies, when mapping

  int i;
  for (i=0;i<jointfreq2.dimx();i++) {
    jointfreq3(i)=jointfreq2(i,C-1); }

  if (debug==1) {
    cout << "jointfreq3" << endl;
    for (i=0;i<jointfreq3.dim();i++) cout << jointfreq3(i) << " ";
    cout << endl; }

}

void make_auganc_hap(const int &C, const int &ctrall, 
		     const int &anc_hap_known,
		     const ivectors &anc_hap, ivector &auganc_hap,
		     int &max_counter, const int &debug)
{
  //this function adds the center allele to the ancestral haplotype

  int i,j;

  //make auganc_hap
  if (anc_hap_known==1 || anc_hap_known==2) {
    if (anc_hap_known==1) {
      for (j=0;j<C;j++) auganc_hap(j)=anc_hap[0](j);
      auganc_hap(C)=ctrall;
      for (j=C;j<anc_hap[0].dim();j++) auganc_hap(j+1)=anc_hap[0](j);
    }
    else {
      for (j=0;j<C;j++) auganc_hap(j)=anc_hap[max_counter](j);
      auganc_hap(C)=ctrall;
      for (j=C;j<anc_hap[0].dim();j++) auganc_hap(j+1)=anc_hap[max_counter](j);
    }
    
    //    max_counter++;

    if (debug==1) {
      cout << "auganc_hap" << endl;
      for (i=0;i<auganc_hap.dim();i++) cout << auganc_hap(i) << " ";
      cout << endl; }
  }
  
}

void updateanchapstore(ivectors &anchapstore, int &anchapcounter,
		       const int &C_loc, const igrid &candanc,
		       const int &debug)
{
  //this function takes the ancestral hap. after every time it is estimated
  //and checks to see if it has been estimated before. If not, it adds the
  //candidate to a list of haps to maximize over in the final generation
  //of the likelihood surface

  int i=0,j;
  int match=0, submatch=1;

  //does the estimated ancestral haplotype match any previously stored
  //candidates?
  for (i=0;i<anchapstore.dim() && anchapstore[i][0]!=0;i++) {
    submatch=1;
    for (j=0;j<anchapstore[i].dim();j++) {
      if (j< C_loc) { 
	if (anchapstore[i][j] != candanc(0,j)){submatch=0; break;}}
      if (j>=C_loc) { 
	if (anchapstore[i][j] != candanc(0,j+1)){submatch=0; break;}} }
    if (submatch==1) { match=1; break; }
  }

  if (match==0) {
    for (j=0;j<anchapstore[anchapcounter].dim();j++) {
      if (j< C_loc) anchapstore[anchapcounter][j] = candanc(0,j);
      else anchapstore[anchapcounter][j] = candanc(0,j+1);
    }
    anchapcounter++;
  }

  if (debug) {
    cout << "anchapstore\n";
    for (i=0;i<anchapstore.dim() && anchapstore[i][0]!=0;i++) {
      for (j=0;j<anchapstore[i].dim();j++) 
	cout << anchapstore[i][j] << " ";
      cout << endl;
    }
  }

}//updateanchapstore

int assignE(const ivector &maxlikCint, const dvector &maxlikCloc)
{
  //this function searches over the results of the calls to max() and
  //picks the interval which as heretofore generated the highest likelihood.
  //Estimates of the ancestral hap. are then reconstructed around 
  //this interval.

  int i, loc=maxlikCint[0];
  double max=maxlikCloc[0];

  for (i=1;i<maxlikCint.dim();i++) {
    if (maxlikCloc[i]>max) { max=maxlikCloc[i]; loc=maxlikCint[i]; }
  }

  return loc;
}//assignE

int pickmaxlikanc(const dvectors &resonetempstore)
{
  //for each point on the likelihood surface, this function picks the
  //ancestral haplotype giving the highest lik. at the assumed location of 
  //the gene (these points are joined together to form the lik. surf.)

  int i,maxind=0;
  double maxlik=resonetempstore[0][4];

  for (i=1;i<resonetempstore.dim();i++) {
    if (maxlik<resonetempstore[i][4]) {
      maxlik=resonetempstore[i][4];
      maxind=i; }
  }
  
  return maxind;
	
}//pickmaxlikanc