// generate multinomial stuff
#include <iostream>
#include <random>
#include <cmath>
#include <time.h> 
#include "readBed.h"
#include "GenORs.h"
#include <sys/time.h>

using namespace std;

//TODO: split this into functions

int main(int argn, char **argv){
	
	// Stage 0: Set defaults, initalize variables, etc
	timeval t1;
	gettimeofday(&t1, NULL);
	srand(t1.tv_usec * t1.tv_sec);
	
	vector< vector<string> > conflines;
    	vector<string> splitline;
    	istringstream linestream;
	ifstream input;
	string line;
	int i,j, k;
	int Nline = 0;
	
	int Npheno = 2; //number of categories (excluding controls)
	int Nvar = 100; //number of variants to simulate
	int Ncovar = 0; //number of covariates to simulate
	double f = 0.5; //allele frequency
	int makeTag = 0; //flag indicating if a single tag-SNP should be generated
	string covFile = ""; //text file containing a covariance matrix to sample effect sizes from
	int covSize = 0; //size of the covariance matrix
	double * cov; //the covariance matrix
	
	if (argn < 3){
		cout << "  usage: sims configfile.txt outprefix" << endl;
		return(0);
	}
	
	//Parse the config file
    input.open(argv[1]);
    if (input.is_open()){
		while (!input.eof()){
	        splitline.clear(); 
	        linestream.clear();
			getline(input,line);
	        splitline.clear(); 
	        linestream.str(line);
	        copy(istream_iterator<string>(linestream), istream_iterator<string>(), back_inserter<vector<string> >(splitline));
			if (splitline.size() == 0) continue;
			Nline++;
			conflines.push_back(splitline);
			if (splitline[0] == "Npheno:") {Npheno = atoi(splitline[1].c_str()); continue;}
      	  	if (splitline[0] == "Ncovar:") {Ncovar = atoi(splitline[1].c_str()); continue;}
      	  	if (splitline[0] == "Nvar:") {Nvar = atoi(splitline[1].c_str()); continue;}
      	  	if (splitline[0] == "f:") {f = atof(splitline[1].c_str()); continue;}
			if (splitline[0] == "tag:") {makeTag = 1;}
			if (splitline[0] == "CovFile:") {covFile = splitline[1];}
		}
	} else cout << "Error: Cannot open input file" << endl;
	input.close();
	//cout << "Input read" << endl;
	// Read the covariance matrix (if there is one)
	if (covFile != ""){
		covSize = 0;
		cov = new double [1000];
		input.open(covFile.c_str());
		if (input.is_open()){
			while (!input.eof()){
	        	splitline.clear(); 
	        	linestream.clear();
				getline(input,line);
	        	splitline.clear(); 
	        	linestream.str(line);
	        	copy(istream_iterator<string>(linestream), istream_iterator<string>(), back_inserter<vector<string> >(splitline));
				for (i = 0; i < splitline.size(); i++){
					cov[i + covSize*splitline.size()] = atof(splitline[i].c_str());
				}
				if (splitline.size() == 0) continue;
				covSize++;
			}	
		} else cout << "Error: Cannot open input file" << endl;
		if (covSize != Npheno){
			cout << "Error: Size of covariance matrix is different to number of phenotypes" << endl;
			throw(3);
		}

	}
	GenORs ORMaker(cov, covSize);
	
	//Set the number of individuals to simulate
	int * Ncounts = new int [Npheno + 1];
	int found=0;
	for (i = 0; i < Nline; i++){
		if (conflines[i][0] == "Ncounts:"){
			for (j = 0; j <= Npheno; j++) {
				Ncounts[j] = atoi(conflines[i][1 + j].c_str());
				found=1;
			}
		}
	}
	
	//Calclulate the total number of individuals across all categories
	int Ntotal=0;
	for (i = 0; i <= Npheno; i++) {
		if (!found) Ncounts[i] = 1000; //set to 1000 per category if not specified in config file
		Ntotal += Ncounts[i];
	}

	// set the underlied frequencies of each category (before ascertainment)
	double * phenoK = new double [Npheno + 1];
	found=0;
	for (i = 0; i < Nline; i++){
		if (conflines[i][0] == "prevs:"){
			for (j = 0; j <= Npheno; j++) {
				phenoK[j] = atof(conflines[i][1 + j].c_str());
				found=1;
			}
		}
	}

	// if not set in config file, set to a default value of 0.01 for non-control categories
	if (!found){
		phenoK[0] = 1;
		for (i = 1; i <= Npheno; i++) {
			phenoK[i] = 0.01;
			phenoK[0] -= 0.01;
		}
	}
	
	//set the effect size
	double * beta = new double [Npheno];
	found=0;
	for (i = 0; i < Nline; i++){
		//cout << "go:" << conflines[i][0] << endl;
		if (conflines[i][0] == "ORs:"){
			for (j = 0; j < Npheno; j++) {
				beta[j] = log(atof(conflines[i][1 + j].c_str()));
				found=1;
			}
		}
	}

	// if not set in the config file, set the first half of phenotypes to OR=1.22, and the rest to OR=1
	if (!found){
		for (i = 0; i < Npheno; i++) {
			beta[i] = 0;
			if (((i+1)*2)/Npheno > 1) beta[i] = 0.2;
		}
	}

	//set covariates
	double * confound_covar, * beta_covar;	
	if (Ncovar != 0){
		confound_covar = new double [Ncovar];
		beta_covar  = new double [Npheno * Ncovar];

		// set the effect od covariates on the allele frequency
		found=0;
		for (j = 0; j < Ncovar; j++) confound_covar[j] = 0; //default is no confounding
		for (i = 0; i < Nline; i++){
			if (conflines[i][0] == "FreqConfound:"){
				for (j = 0; j < Ncovar; j++) {
					confound_covar[j] = (atof(conflines[i][1 + j].c_str()));
					found=1;
				}
			}
		}

		// set the effect of covariates on phenotype categories
		found=0;
		for (i = 0; i < Nline; i++) if (conflines[i][0] == "DiseaseConfound:") {found=1;break;}
		for (j = 0; j < Ncovar; j++){
			for (k = 0; k < Npheno; k++){
				beta_covar[j*Npheno + k] = 0; //default to no confounding
				if (found) beta_covar[j*Npheno + k] = atof(conflines[1 + (i*found) + j][k].c_str());
			}
		}
	}

	//simulate a tagSNP (if set in the config file)
	double r2;
	double f_tag;
	if (makeTag){
		for (i = 0; i < Nline; i++){
			if (conflines[i][0] == "tag:") {
				r2 = atof(conflines[i][1].c_str());
				f_tag = atof(conflines[i][2].c_str());
				break;
			}
		}
	}	
	double f_AB = f*f_tag + sqrt(r2*f*f_tag*(1 - f)*(1 - f_tag));
	double f_BgA = f_AB / f;
	double f_Bga = (f_tag - f_AB) / (1 - f);
	
	
	// Output text describing the simulations
	
	cout << endl << "##Trinculo simulations##" << endl;
	
	cout << Npheno << " phenotypes with prevalences (" << phenoK[1];
	for (i = 1; i < Npheno; i++) cout << ", " << phenoK[1 + i];
	cout << ")" << endl;
	
	if (covFile == ""){
		cout << "Allele frequency is " << f << " and odds ratios are (" << exp(beta[0]);
		for (i = 1; i < Npheno; i++) cout << ", " << exp(beta[i]);
		cout << ")" << endl;
	} else {
		cout << "Allele frequency is " << f << " and odds ratios are drawn from covariance matrix in " << covFile << "." << endl;
	}
	cout << "Sampling " << Ntotal << " individuals, case numbers are (" << (Ncounts[1]);
	for (i = 1; i < Npheno; i++) cout << ", " << (Ncounts[i+1]);
	cout << ")" << endl;
	if (Ncovar != 0){
		cout << "Including " << Ncovar << " covariates:" << endl;
  		cout << "Effect of PCs on allele frequency: (" << confound_covar[0];
		for (i = 1; i <= Npheno; i++) cout << ", " << confound_covar[i];
		cout << ")" << endl;
		for (i = 0; i < Npheno; i++){
			cout << "Effect of PCs on phenotype " << i << " (" << beta_covar[0*Npheno + i];
			for (j = 1; j < Ncovar; j++) cout << ", " << beta_covar[j*Npheno + i];
			cout << ")" << endl;
		}
	}
	
	// define random number generators
	std::default_random_engine generator (rand());
  	std::normal_distribution<double> rnorm(0,1);
  	std::binomial_distribution<int> rbinom(2,f);
	std::uniform_real_distribution<double> runif(0.0,1.0);
  	
	// temporary variables (who are you, the variable name police?)
  	double temp1 = 0;
	double temp2 = 0;
	
	// set the intercept on the basis of the prevalences
  	double * mu = new double [Npheno];
	for (i = 0; i < Npheno; i++){
		temp1 = 2*f*beta[i];
	  	temp2 = 2*f*(1-f)*beta[i]*beta[i];
	  	if (Ncovar > 0) for (j = 0; j < Ncovar; j++) temp2 += beta_covar[j*Npheno + i]*beta_covar[j*Npheno + i];
	  	mu[i] = log(phenoK[i + 1]) - log(phenoK[0]) - temp1 - temp2/2;
 	 }
	
	//allocate some space to put the results
 	bedfile bed;
	double * covar;
	
	if (Ncovar > 0)  covar = new double [Ncovar*Ntotal];
	int * y;
	y = (int *) malloc(Ntotal * sizeof(int));
	if (makeTag) bed.emptyBed(2*Nvar,Ntotal); //allocate some space
	else bed.emptyBed(Nvar,Ntotal);
	
	// more temporary variables
	double * thisCov;
	if (Ncovar > 0) thisCov = new double [Ncovar];
	vector<double> ps(Npheno+1,0.2);
	vector<double> ps2(3,0.2);
	
	int ind, thisG, thisG2, thisY = 0;
	double pTot;
	
	//Carry out the simulations (the stdout describes what happens from this point)
	
	cout << endl << "##Running simulations##" << endl;
	
	cout << "Generating covariates and phenotypes...";
	for (ind = 0; ind < Ntotal; ind++){
		for (j = 0; j < Ncovar; j++) covar[ind*Ncovar + j] = rnorm(generator);
		y[ind] = thisY;
		Ncounts[thisY]--;
		if (Ncounts[thisY] == 0) thisY++;
	}
	cout << "done" << endl;
	
	cout << "Generating genotypes..." << endl;
	
	int g, var;
	double fadj;
	for (var = 0; var < Nvar; var++){
		if (covFile != ""){
			ORMaker.sampleORs();
		  	for (i = 0; i < Npheno; i++){
				beta[i] = log(ORMaker.ORs[i]);
				temp1 = 2*f*beta[i];
			  	temp2 = 2*f*(1-f)*beta[i]*beta[i];
			  	for (j = 0; j < Ncovar; j++) temp2 += beta_covar[j*Npheno + i]*beta_covar[j*Npheno + i];
			  	mu[i] = log(phenoK[i + 1]) - log(phenoK[0]) - temp1 - temp2/2;
		 	 }
		}
		if (var % 100 == 0) cout << "..." << var << " completed..." << endl;
		for (ind = 0; ind < Ntotal; ind++){
			
			fadj = f;
			for (j = 0; j < Ncovar; j++) fadj +=  covar[ind*Ncovar + j]*confound_covar[j];
			for (g = 0; g <= 2; g++){
				ps[0] = 1;
				pTot=1;
				for (i = 0; i < Npheno; i++) {
					temp1 = mu[i] + g*beta[i];
					for (j = 0; j < Ncovar; j++) temp1 += covar[ind*Ncovar + j]*beta_covar[j*Npheno + i];
					ps[1 + i] = exp(temp1);
					pTot += ps[1 + i];
				}
				ps2[g] = ps[y[ind]]/pTot;
			}
			ps2[0] *= (1 - fadj)*(1 - fadj);
			ps2[1] *= 2*(1 - fadj)*fadj;
			ps2[2] *= fadj*fadj;
	  		std::discrete_distribution<int> rmultinom(ps2.begin(),ps2.end());
			thisG = rmultinom(generator);
			if (makeTag){
				thisG2 = 0;
				if (thisG == 0){
					if (runif(generator) < f_Bga) thisG2++;
					if (runif(generator) < f_Bga) thisG2++;
				} else if (thisG == 1){
					if (runif(generator) < f_BgA) thisG2++;
					if (runif(generator) < f_Bga) thisG2++;
				}else if (thisG == 2){
					if (runif(generator) < f_BgA) thisG2++;
					if (runif(generator) < f_BgA) thisG2++;
				}
				bed.setDosage(2*var,ind,thisG);
				bed.setDosage(2*var + 1,ind,thisG2);
			} else bed.setDosage(var,ind,thisG);

			
			
		}
	}
	cout << "done" << endl;
	
	cout << "Writing data..." << endl;
	
    ofstream covarFILE, phenoFILE;
    covarFILE.open((string(argv[2]) + ".covar").c_str());
    phenoFILE.open((string(argv[2]) + ".pheno").c_str());
	
	covarFILE << "FID IID"; 
	for (i = 0; i < Ncovar; i++) covarFILE << " PC" << i;
	covarFILE << endl;
	phenoFILE << "FID IID PHENO" << endl;
	
	for (ind = 0; ind < Ntotal; ind++){
		covarFILE << bed.FID[ind] << " " << bed.IID[ind];
		for (j = 0; j < Ncovar; j++){
			covarFILE << " ";
			covarFILE << covar[ind*Ncovar + j];
		}
		covarFILE << endl;
		phenoFILE << bed.FID[ind] << " " << bed.IID[ind] << " " << y[ind] << endl;
	}
	
	bed.writeBed(string(argv[2]));
	cout << "...done" << endl;
	
  	return 0;
}
