/*  trinculo, Flexible tools for multinomial association
    Copyright (C) 2014  Luke Jostins

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.*/

#include <iostream>
#include <fstream>
#include <vector>
#include <sstream>
#include <string.h>
#include <stdlib.h>
#include <iterator>
#include <map>

using namespace std;


class bedfile {
  char * bedblock; //a block of binary data read in a binary plink (.bed) file
  int allocated;
 public:
  int Nvars, Ninds; //number of variants and individuals
  vector<string> FID, IID, MID, PID, CHR, RSID, A1, A2; //individual and variant information read from .fam and .bim file
  map<string,int> IDMap; //a map from IIDs to positions in the bedblock, used for matching IDs across different files
  vector<int> pheno, sex; //more information from the fam file
  vector<long int> pos; //more information from the bim file
  vector<double> genpos; //more information from the bim file
  int getDosage(int i, int j);
  void setDosage(int i, int j, int d);
  void readBed(string inprefix);
  void emptyBed(int newNvars, int newNinds);
  void writeBed(string inprefix);
  void makeIDMap();
  bedfile();
  ~bedfile();
};

bedfile::bedfile(){
	allocated=0;
}

bedfile::~bedfile(){
	if (allocated) delete bedblock;
}

void bedfile::makeIDMap() {
	// populate the ID map
  for (int i = 0; i < Ninds; i++) IDMap[FID[i] + ":" + IID[i]] = i;
}

void bedfile::emptyBed(int newNvars, int newNinds){
	//generate an empty bed file with arbitrary individual and variant information, reading for populating (used for simulated data)
	int i,j;
	Nvars = newNvars;
	Ninds = newNinds;
	stringstream sstm;
	
	for (i = 0; i < Ninds; i++){
		sstm << "IND" << i;
		FID.push_back(sstm.str());
		IID.push_back(sstm.str());
		sstm.str("");
		sstm.clear();
		MID.push_back("0");
		PID.push_back("0");
		sex.push_back(-9);
		pheno.push_back(-9);
	}
	for (i = 0; i < Nvars; i++){
		CHR.push_back("1");
		sstm << "SNP" << i;
		RSID.push_back(sstm.str());
		sstm.str("");
		sstm.clear();
		genpos.push_back(i);
		pos.push_back(i);
		A1.push_back("A");
		A2.push_back("C");
	}
    long int size;
    size = (1 + (int)(Ninds - 1)/4);
    size *= Nvars;
    size += 3;
	
	cout << "Allocating empty bed (size=" << size << ")" << endl;
    bedblock = new char [size]; 
	allocated = 1;
	
    bedblock[0] = (char) 108;
	bedblock[1] = (char) 27;
	bedblock[2] = (char) 1;
	for (i = 3; i < size; i++) bedblock[i] = 0; //technically this isn't required (as the below ensures initailization), but try telling that to valgrind
	for (i = 0; i < Nvars; i++)
		for (j = 0; j < Ninds; j++)
			setDosage(i,j,0);
}

void bedfile::readBed(string inprefix){
	// read in data from a .fam, .bim and .bed file (in that order)
	
  string line;
  int value;
  double value2;
  ifstream fam, bim, bed;
  istringstream linestream, tempstream;
  vector<string> splitline;
  fam.open((string(inprefix) + ".fam").c_str());
  
  Ninds=0;
  
  if (fam.is_open()){
    cout << "Reading FAM file" << endl;
    while (!fam.eof()){
      
      getline(fam,line);
      if (line.length() == 0) break;

      splitline.clear(); 
      linestream.clear();
      tempstream.clear();
      
      linestream.str(line);
      copy(istream_iterator<string>(linestream), istream_iterator<string>(), back_inserter<vector<string> >(splitline));
      
      FID.push_back(splitline[0]); 
      IID.push_back(splitline[1]); 
      MID.push_back(splitline[2]); 
      PID.push_back(splitline[3]);
      
      tempstream.str(splitline[4]);
      tempstream >> value;
      sex.push_back(value);

      tempstream.clear();
      tempstream.str(splitline[5]);
      tempstream >> value;
      pheno.push_back(value);
      Ninds++;
    }
    fam.close();
  } else 
    cout << "Error: Cannot open FAM file" << endl;


  bim.open((string(inprefix) + ".bim").c_str());

  Nvars=0;
  if (bim.is_open()){
    cout << "Reading BIM file" << endl;
    while (!bim.eof()){
      getline(bim,line);
      if (line.length() == 0) break;
      splitline.clear();
      linestream.clear();
      tempstream.clear();
      linestream.str(line);
      copy(istream_iterator<string>(linestream), istream_iterator<string>(), back_inserter<vector<string> >(splitline));
      CHR.push_back(splitline[0]);
      RSID.push_back(splitline[1]);
      tempstream.str(splitline[2]);
      tempstream >> value2;
      genpos.push_back(value2);
      tempstream.clear();
      tempstream.str(splitline[3]);
      tempstream >> value;
      pos.push_back(value);
      A1.push_back(splitline[4]);
      A2.push_back(splitline[5]);
      Nvars++;
    }
    bim.close();
  } else {
    cout << "Error: Cannot open BIM file" << endl;
  }
  
  long int size;
  size = (1 + (int)(Ninds - 1)/4);
  size *= Nvars;
  size += 3;
  //cout << "Nvars: " << Nvars << " Ninds: " << Ninds << " size: " << size << endl;
  bedblock = new char [size]; 
  allocated=1;
  bed.open((string(inprefix) + ".bed").c_str(), ios::in|ios::binary|ios::ate);
  
  if (bed.is_open())
    {		
      cout << "Reading BED file" << endl;
      bed.seekg (0, ios::beg);
      bed.read (bedblock, size);
      bed.close();
	  
      if ((unsigned int)bedblock[0] != 108 || (unsigned int) bedblock[1] != 27){
	cout << "Error: BED file misformatted" << endl;
      }
      
      if ((unsigned int) bedblock[2] == 1){
	cout << "BED in SNP-Major format" << endl;
      } else if ((unsigned int) bedblock[2] == 0){
	cout << "Error: BED in Individual-Major format, please reformat to SNP-Major" << endl;
      } else {
	cout << "Error: BED file misformatted" << endl;
      }
    }
  else {
    cout << "Error: Cannot open BED file" << endl;
  }
}

int bedfile::getDosage (int i,int j){
	// get the dosage for a particular variant for a particular individual
  int loc1, loc2, b1, b2, d;
  i += 1; j += 1; // I worked this out for a 1-index...
  loc1 = 3 + (i-1)*(1 + (int)(Ninds - 1)/4) + int((j-1)/4);
  loc2 = (int) ((j - 4*(int)((j-1)/4)) - 1);
  b1 = (unsigned int) (bedblock[loc1] >> (2*loc2) & (char) 1);
  b2 = (unsigned int) (bedblock[loc1] >> (2*loc2 + 1) & (char) 1);
  
  if (b1 == 0) if (b2 == 0) d=0; else d=1;
  else if (b2 == 0) d=-1; else d = 2;
  
  return(d);
}


void bedfile::setDosage (int i,int j, int d){
	// set the dosage for a particular variant for a particular individual to the value d
  int loc1, loc2, b1, b2;
  i += 1; j += 1; // I worked this out for a 1-index...
  loc1 = 3 + (i-1)*(1 + (int)(Ninds - 1)/4) + int((j-1)/4);
  loc2 = (int) ((j - 4*(int)((j-1)/4)) - 1);
  
  if (d == 2) {b1 = 1; b2 = 1;}
  else if (d == 1) {b1 = 0; b2 = 1;}
  else if (d == 0) {b1 = 0; b2 = 0;}
  else {b1 = 1; b2 = 0;} // if d isn't a real dosage (i.e. 0, 1 or 2), set it to missing
  
  bedblock[loc1] ^= (-b1 ^ bedblock[loc1]) & (1 << (2*loc2));
  bedblock[loc1] ^= (-b2 ^ bedblock[loc1]) & (1 << (2*loc2 + 1));
}



void bedfile::writeBed(string inprefix){
	// write data out to a .fam, .bim and .bed file (in that order)
	
  string line;
  int i;
  int value;
  double value2;
  ofstream fam, bim, bed;
  istringstream linestream, tempstream;
  vector<string> splitline;
  fam.open((string(inprefix) + ".fam").c_str());
  
  if (fam.is_open()){
	  cout << "Writing FAM file" << endl;
	  for (i = 0; i < Ninds; i++) fam << FID[i] << " " << IID[i] << " " << MID[i] << " " << PID[i] << " " << sex[i] << " " << pheno[i] << endl;
	  fam.close();
  }
  else cout << "Error: Cannot open FAM file" << endl;

  bim.open((string(inprefix) + ".bim").c_str());

  if (bim.is_open()){
    cout << "Writing BIM file" << endl;
    for (i = 0; i < Nvars; i++) bim << CHR[i] << " " << RSID[i] << " " << genpos[i] << " " << pos[i] << " " << A1[i] << " " << A2[i] << endl;
    bim.close();
  } else {
    cout << "Error: Cannot open BIM file" << endl;
  }
  
  long int size;
  size = (1 + (int)(Ninds - 1)/4);
  size *= Nvars;
  size += 3;
  bed.open((string(inprefix) + ".bed").c_str(), ios::out|ios::binary);

  if (bed.is_open())
    {
      cout << "Writing BED file (size=" << size << ")" << endl;
      bed.write (bedblock, size);
      bed.close();
  }
  else {
    cout << "Error: Cannot write BED file" << endl;
  }
  
}
