/*  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>
#include <limits>
#include <ctype.h>

using namespace std;


bool custom_isnan(double var){
  // checks if a variable is nan
  volatile double d = var;
  return d != d;
}

class covars {
	// a class for reading and storing covariate data across individuals
 public:
  int Ncovars, Ninds; //the number of covariates and individuals read in
  vector<vector<double> > covariates; //the value of each covariate for each individual
  vector<string> covnames; // the names of the covariances, taken from the text file
  void readCovar(string infile,map<string,int> IDMap,int Ninds, int normalizecovar = 0);
};

void covars::readCovar(string infile,map<string,int> IDMap,int Ninds, int normalizecovar){
  // reads in covariates from a text file, matching the IDs using the map provided
	
  string line;
  double value;
  ifstream covfile;
  covfile.open(infile.c_str());
  int i,j;
  istringstream linestream, tempstream;
  vector<string> splitline;

  //get header
  getline(covfile,line);
  linestream.str(line);
  copy(istream_iterator<string>(linestream), istream_iterator<string>(), back_inserter<vector<string> >(splitline));
  Ncovars = splitline.size() - 2;
  covariates.resize(Ncovars);
  covnames.resize(Ncovars);
  for (i = 0; i < Ncovars; i++) {
    covariates[i].resize(Ninds,numeric_limits<double>::quiet_NaN( ));
    covnames[i] = splitline[i + 2];
  }
  for (i = 0; i < Ncovars; i++){
	for (j = 0; j < Ninds; j++){
		covariates[i][j] = numeric_limits<double>::quiet_NaN( );
	}
  }
  int textwarned=0;

  if (covfile.is_open()){
    cout << "Reading covariants file" << endl;
    while (!covfile.eof()){
      getline(covfile,line);
      if (line.length() == 0) break;
      splitline.clear();
      linestream.clear();
      linestream.str(line);
      copy(istream_iterator<string>(linestream), istream_iterator<string>(), back_inserter<vector<string> >(splitline));
      if (IDMap.find(splitline[0] + ":" + splitline[1]) == IDMap.end()) continue;
      for (i = 2; i < splitline.size(); i++) {

	  tempstream.clear();	    
	  tempstream.str(splitline[i]);
	  tempstream >> value;
	  
	  if (!tempstream) {
		  if (!textwarned & splitline[i] != "NA") {
			  cout << "Warning: Non-numeric input present in covariate file. Treating as missing." << endl;
			  textwarned = 1;
		  }
		  value = numeric_limits<double>::quiet_NaN( );
	  }
	  tempstream.str("");
	  tempstream.clear();
	
	covariates[i - 2][ IDMap[splitline[0] + ":" + splitline[1]] ] = value;
	
      }	
    }
    covfile.close();
  } else {
    cout << "Error: Cannot open covar file" << endl;
  }

  int N;
  double sumx, sumxx, meanx, sdx;

  if (normalizecovar){
    for (i = 0; i < Ncovars; i++) {
      N = 0;
      sumx = 0;
      sumxx = 0;
      // normalize the covariate by its mean and sd
      // we have time, so lets do this in three passes to ensure numeric stability
      for (j = 0; j < Ninds; j++){
	if (custom_isnan(covariates[i][j])) continue;
	N++;
	sumx += (covariates[i][j]);
      }
      meanx = sumx/N;
      for (j = 0; j < Ninds; j++){
	if (custom_isnan(covariates[i][j])) continue;
	sumxx += (covariates[i][j] - meanx)*(covariates[i][j] - meanx);
      }
      sdx = sqrt(sumxx/(N - 1));
      for (j = 0; j < Ninds; j++){
	if (custom_isnan(covariates[i][j])) continue;
	covariates[i][j] = (covariates[i][j] - meanx)/sdx;
      }
    }
  }
}
  
