#include <fstream>
#include <iostream>
#include <iomanip>
#include <string>
#include <sstream>
#include <stdarg.h>
#include <stdlib.h>
#include <math.h>
#include <map>
#include <matvec/doublematrix.h>
#include <matvec/vector.h>
#include <matvec/session.h>
#include "util.h"
#include <conio.h>

// Single trait, fixed effects models

using namespace std; 

class TermData{
public:
  double value;
  unsigned level;
};

class DataNode{
public:
  vector<TermData > trmVec;
  double depVar;
};

class MME;

class ModelTerm{
public:
	unsigned start;
	string name;
	static MME *myMMEPtr;
	Recoder<string> myRecoder;
	vector<unsigned> factors;
	
	unsigned code(string str){return myRecoder.code(str);}
	unsigned nLevels(){return myRecoder.size();}
	string   getTermString();
	unsigned getTermLevel (){
		return code(getTermString());
	}
	double   getTermValue ();
};

class MME {
public:
	string fileName;
	Tokenizer colType;
	Tokenizer colName;
	Tokenizer colData;
	unsigned numCols;
	unsigned depCol;
	vector <ModelTerm> modelTrmVec;
	vector <DataNode> dataVec;
	unsigned numTerms;
	unsigned mmeSize;
	matvec::doubleMatrix lhs;
	matvec::Vector<double> rhs, sol;
	
	void putColNames(string str);
	void putColTypes(string str);
	void inputData();
    void displayData();
	static double getDouble(string& Str);
	void calcStarts();
	void getSolution();
	void calcWPW();
	void display();
};

string ModelTerm::getTermString(){
	unsigned numFactors = factors.size();
	string trmStr;
	unsigned factorIndex = factors[0];
	if(myMMEPtr->colType[factorIndex]=="COV"){
		trmStr = myMMEPtr->colName[factorIndex];
	}
	else {
		trmStr = myMMEPtr->colData[factorIndex];
	}
	for (unsigned i=1;i<numFactors;i++){
		factorIndex = factors[i];
		if(myMMEPtr->colType[factorIndex]=="COV"){
			trmStr += "*" + myMMEPtr->colName[factorIndex];
		}
		else{
			trmStr += "*" + myMMEPtr->colData[factorIndex];
		}
	}
	return trmStr;
}

double ModelTerm::getTermValue(){
	unsigned numFactors = factors.size();
	double value = 1.0;  
	for (unsigned i=0;i<numFactors;i++){
		unsigned factorIndex = factors[i];
		if (myMMEPtr->colType[factorIndex]=="COV"){
			string covStr = myMMEPtr->colData[factorIndex];
			value *= MME::getDouble(covStr);
		}
	}
	return value;
}

void MME::putColNames(string str){
  string sep(" ");
  colName.getTokens(str,sep);
  numCols = colName.size();
}

void MME::putColTypes(string str){
  string sep(" ");
  colType.getTokens(str,sep);
  if (numCols!=colType.size()){
    cerr <<"number of column names and column types do not match\n";
    exit (-1);
  }
  for (unsigned i=0;i<numCols;i++){
	if (colType[i] == "DEP") {
		depCol = i;
		return;
	}
  }
  cout << "Could not find dependent variable \n";
  exit(1);
}

void MME::inputData(){
	DataNode dataNode;
	numTerms = modelTrmVec.size();
	dataNode.trmVec.resize(numTerms);
	ifstream datafile;
	datafile.open(fileName.c_str());
	if(!datafile) {
		cerr << "Couldn't open data file: " << fileName << endl;
		exit (-1);
	}
	unsigned linewidth = 1024;
	char *line = new char [linewidth];
	string sep(" ");
	while (datafile.getline(line,linewidth)){
		string inputStr(line);
		colData.getTokens(inputStr,sep);
		dataNode.depVar = getDouble(colData[depCol]);
		for (unsigned i=0;i<numTerms;i++){ 
			dataNode.trmVec[i].level = modelTrmVec[i].getTermLevel();
			dataNode.trmVec[i].value = modelTrmVec[i].getTermValue();
		}
		dataVec.push_back(dataNode);
	}
}
	
double MME::getDouble(string& Str) {
  istringstream inputStrStream(Str.c_str());
  double val;
  inputStrStream >> val; 
  return val;
}

void MME::calcStarts(){
	modelTrmVec[0].start = 0; 
	for (unsigned i=1;i<numTerms;i++){
		modelTrmVec[i].start = modelTrmVec[i-1].start + modelTrmVec[i-1].nLevels();
	}
	mmeSize = modelTrmVec[numTerms-1].start + modelTrmVec[numTerms-1].nLevels();
}

void MME::calcWPW(){
	unsigned ii,jj;
	double vi,vj,tr_value;
	rhs.resize(mmeSize,0.0);
	lhs.resize(mmeSize,mmeSize,0.0);
	for (unsigned i=0;i<dataVec.size();i++){
		for (unsigned mi=0;mi<numTerms;mi++){ 
			ii = modelTrmVec[mi].start + dataVec[i].trmVec[mi].level - 1;
			vi = dataVec[i].trmVec[mi].value;
			tr_value = dataVec[i].depVar;
			rhs[ii] += vi*tr_value;
			for (unsigned mj=0;mj<numTerms;mj++){
				jj = modelTrmVec[mj].start + dataVec[i].trmVec[mj].level - 1; 
				vj = dataVec[i].trmVec[mj].value;
				lhs[ii][jj] += vi*vj; 
			}
		}
	}
}

void MME::getSolution(){
	inputData();
	calcStarts();
	calcWPW();
	sol = lhs.ginv0()*rhs;
}

void MME::display(){
	cout << "LHS " << endl;
	for (unsigned i = 0;i<mmeSize;i++){
		for (unsigned j = 0;j<mmeSize;j++){
			cout << setw(5) << lhs[i][j] <<" ";
		}
		cout << endl;
	}
	cout << "RHS " << endl;
	cout << rhs << endl;
	for (unsigned i=0;i<modelTrmVec.size();i++){
		cout << "Solutions for " << modelTrmVec[i].name << endl;
		Recoder<string>::iterator it;
		for (it=modelTrmVec[i].myRecoder.begin(); it!=modelTrmVec[i].myRecoder.end();it++){
			unsigned ii = modelTrmVec[i].start + it->second - 1;
			cout << setw(10) << it->first << " " << sol[ii] << endl;
		}
	}
}
		

MME* ModelTerm::myMMEPtr;

int main() {
	try{
		matvec::SESSION.initialize("matvec_trash");
		MME mme;
		ModelTerm::myMMEPtr = &mme;
		mme.fileName = "Data/twoway.dat";
		mme.putColNames("A B y");
		mme.putColTypes("CLASS CLASS DEP");
		
		ModelTerm mterm;
	 	mterm.name = "A";
		mterm.factors.resize(1);
		mterm.factors[0] = 0;
		mme.modelTrmVec.push_back(mterm);
		
		mterm.name = "B";
		mterm.factors[0] = 1;
		mme.modelTrmVec.push_back(mterm);	
		
		mterm.name = "A*B";
		mterm.factors.resize(2);
		mterm.factors[0] = 0;
		mterm.factors[1] = 1;
		mme.modelTrmVec.push_back(mterm);
		
		mme.getSolution();
		mme.display();
		getch();
	}
	catch (matvec::exception &ex) {
		cerr << ex.what() << "\n";
		exit(1);
	}
	catch (...) {
		cerr << "other exceptions were caught\n";
		exit(1);
	}
}