#include <iomanip> 

// little function to calculate the CDF of the normal distribution
double pnorm(double Z) { return 0.5*erfc(Z * M_SQRT1_2); }


class MultinomThread {
	// an object to hold data and carry out various multinomial tests on it

public:
	
	// non-shared dynamic objects (i.e. unique to this MultinomThread object)
	MultinomNewton *NR;
	int snp;
	int *constrain;
	vector<double> specLs;
	string assocOutput;
	string selectOutput;
	string waldOutput;
	string phenoLTRsOutput;
	string errorOutput;
	string modelParOutput;
	string modelErrOutput;
	
	
	// shared, static objects (i.e. shared across MultinomThread objects)
	inputs *opts;
	int Npred;
	int Nind;
	int Npheno;
	bedfile *bed;
	dosagefile *dosage;
	
	//functions
	~MultinomThread();
	void LoadInitial(int newNpred, int newNind, int newNpheno, inputs *newOpts, bedfile *newBed, dosagefile *newDosage, covars *cov, phenos *phen, vector<int> conditionsnps, int phenfocus);
	void SetSNP(int newSNP);
	void LoadData();
	void calcModelSel();
	void calcAssoc();
	void calcWald();
	void calcError();
	void calcFullMD();
	void calcPhenoLRTs();
	void RunTests();
};

MultinomThread::~MultinomThread(){
	delete constrain;
	delete NR;
}
void MultinomThread::LoadInitial(int newNpred, int newNind, int newNpheno, inputs *newOpts, bedfile *newBed, dosagefile *newDosage, covars *cov, phenos *phen, vector<int> conditionsnps, int phenfocus){
	//allocate initial parameters

	// allocate the model fit object
	NR = new MultinomNewton;
	NR->AllocateSpace(newNpred, newNind, newNpheno);
	//MultinomNewton NRtemp(newNpred, newNind, newNpheno);
	//NR = &NRtemp;
	

	// set options and genotypes
	opts = newOpts;
	bed = newBed;
	dosage = newDosage;
	Npred=newNpred;
	Nind=newNind;
	Npheno=newNpheno;
	
	// allocate likelihood vector (for model selection)
	constrain = (int *) new int[newNpheno *(1 + newNpred)];
    specLs.resize(pow(2.0,newNpheno)); 
	
    // set confounders, SNPs to condition on and phenotypes
    if (opts->covarfile != "") NR->setCovars(cov);
    if (opts->conditionfile != ""){
		if (opts->dosagefile == "") {
			NR->setConditional(conditionsnps, bed);
		} else {
			NR->setConditional(conditionsnps, dosage);
		}
	}
    NR->setPheno(phen,phenfocus);
  
    // set the default prior, if requested
    int i, j, l;
	if (opts->defaultPrior){
		int Nparam = Npheno*(1 + Npred);
		double *Sigma = new double [Nparam*Nparam];
		for (i = 0; i < Nparam*Nparam; i++) Sigma[i] = 0;
		for (i = 0; i < Nparam; i++) Sigma[i + i*Nparam] = opts->covarsigma;
		for (i = 0; i < Npheno; i++){
			for (j = 0; j < Npheno; j++){
				if (i == j) Sigma[1 + i*(Npred + 1) + (1 + j*(Npred + 1))*Nparam] = opts->priorsigma;
				else Sigma[1 + i*(Npred + 1) + (1 + j*(Npred + 1))*Nparam] = opts->priorcor*opts->priorsigma;
				for (l = 0; l < conditionsnps.size(); l++)
					if (i == j) Sigma[(2 + cov->Ncovars + l) + i*(1 + Npred) + ((2 + cov->Ncovars + l) + j*(Npred + 1))*Nparam] = opts->priorsigma;
					else  Sigma[(2 + cov->Ncovars + l) + i*(1 + Npred) + ((2 + cov->Ncovars + l) + j*(Npred + 1))*Nparam] = opts->priorcor*opts->priorsigma;
			} 
		}
		NR->setPrior(Sigma);
	}
	
    // set the prior from a file, if requested
    if (opts->priorfile != "") {
		priors pr;
		if (pr.readPriors(opts->priorfile, Npheno, 1, Npred - 1 - conditionsnps.size(), conditionsnps.size(), opts->covarsigma)) throw(1);
		NR->setPrior(pr.Sigma);
	}
	
	//TODO - this above could probably be changed to 
}

void MultinomThread::SetSNP(int newSNP) {
	snp = newSNP;
}

void MultinomThread::LoadData(){
	//load up genotypes for a new SNP
	
    if (opts->dosagefile == "")
      NR->setGenos(bed,snp);
    else
      NR->setDosage(dosage,snp);
	
    NR->setMissing();	
    NR->zeroBetas();
	
	assocOutput.clear();
	selectOutput.clear();
	waldOutput.clear();
}


void MultinomThread::calcModelSel(){
	// calculate the marginal likelihood for every possible sharing model and save the output line to selectOutput
    
	stringstream ss;
	int i, j;
	
	double precision;
	int focus, focustemp, freeCount;

	
	for (focus = 0; focus < floor(pow(2.0,Npheno)); focus++){
		NR->zeroBetas();
		for (i = 0; i < Npheno; i++){
			for (j = 0; j < (Npred + 1); j++) constrain[j + i*(Npred + 1)] = 0;
			constrain[1 + i*(Npred + 1)] = 1;
		}
		focustemp = focus;
		freeCount = 0;
		for (i = (Npheno - 1); i >= 0; i--){
			if (focustemp >= floor(pow(2.0,i))){
				constrain[1 + i*(Npred + 1)] = 0;
				focustemp -= pow(2.0,i);
				freeCount++;
			}
		}	 
		precision = NR->fitLogistic(constrain,opts->fititer,opts->fitreltol, opts->fitabstol);
		if (precision > 1e-5){
			if (opts->dosagefile == "")
				cout << "Warning: Model did not coverge for SNP " << bed->RSID[snp] << endl;
			else
				cout << "Warning: Model did not coverge for SNP " << dosage->RSID[snp] << endl;
		}
		specLs[focus] = NR->calcL();
		if (opts->priorfile != "" || opts->defaultPrior) {
			specLs[focus]  += NR->calcPrior(constrain);
			specLs[focus]  += (Npheno*(Npred) + freeCount)*0.5*log(2*PI);
			NR->calcI();
			for (int i = 0; i < (Npred+1)*Npheno; i++){
				if (constrain[i]) {
					for (int j = 0; j < (Npred + 1)*Npheno; j++) { NR->I[i + j*(Npred + 1)*Npheno] = 0; NR->I[j + i*(Npred + 1)*Npheno] = 0;}
					NR->I[i + i*(Npred + 1)*Npheno] = 1;
				}
			}
			specLs[focus]  -= 0.5*log(fabs(inverse(NR->I,(Npred+1)*Npheno)));
		}
	}
	if (opts->dosagefile == "")
		ss << bed->CHR[snp] << "\t" << bed->pos[snp] << "\t" << bed->RSID[snp] << "\t" << bed->A1[snp] << "\t" << bed->A2[snp];
	else
		ss << dosage->RSID[snp];
	for (i = 0; i < specLs.size(); i++) ss << "\t" << specLs[i] ;
	ss << endl;
	selectOutput = ss.str();
	
} 


void MultinomThread::calcAssoc(){
	//do an association test (either Bayesian or frequentist, as specified by opts) and save the output line to assocOutput
	
	stringstream ss;
	double precision, p, Lnull, Lalt;
	int i,j;
	
    //set the constraints for the null model
    for (i = 0; i < Npheno; i++){
      for (j = 0; j < (Npred + 1); j++) constrain[j + i*(Npred + 1)] = 0;
      constrain[1 + i*(Npred + 1)] = 1;
    }
    
	//fit the null model
	NR->zeroBetas();
	precision = NR->fitLogistic(constrain,opts->fititer,opts->fitreltol, opts->fitabstol);
	if (precision > 1e-5){
		if (opts->dosagefile == "")
			cout << "Warning: Null model did not coverge for SNP " << bed->RSID[snp] << endl;
		else
			cout << "Warning: Null model did not coverge for SNP " << dosage->RSID[snp] << endl;
	}
	
	//get the null maximum or marginal likelihood (as requested)
	Lnull = NR->calcL();
	if (opts->priorfile != "" || opts->defaultPrior) {
		Lnull += NR->calcPrior(constrain);
		Lnull += (Npheno*(Npred))*0.5*log(2*PI);
		NR->calcI();
		for (int i = 0; i < (Npred+1)*Npheno; i++){
			if (constrain[i]) {
				for (int j = 0; j < (Npred + 1)*Npheno; j++) { NR->I[i + j*(Npred + 1)*Npheno] = 0; NR->I[j + i*(Npred + 1)*Npheno] = 0;}
				NR->I[i + i*(Npred + 1)*Npheno] = 1;
			}
		}
		Lnull -= 0.5*log(fabs(inverse(NR->I,(Npred+1)*Npheno)));
	}
	
	// fit the alt model
	for (i = 0; i < Npheno*(Npred + 1); i++) constrain[i] = 0;
	NR->zeroBetas();
	precision = NR->fitLogistic(constrain,opts->fititer,opts->fitreltol, opts->fitabstol);
	if (precision > 1e-5){
		if (opts->dosagefile == "")
			cout << "Warning: Alt model did not coverge for SNP " << bed->RSID[snp] << endl;
		else
			cout << "Warning: Alt model did not coverge for SNP " << dosage->RSID[snp] << endl;
	}
	
	//calculate the alt maximum or marginal likelihood (as requested)
	Lalt = NR->calcL();
	if (opts->priorfile != "" || opts->defaultPrior) {
		Lalt += NR->calcPrior(constrain);
		Lalt += (Npheno*(Npred + 1))*0.5*log(2*PI);
		Lalt -= 0.5*log(fabs(inverse(NR->I,(Npred+1)*Npheno)));
	}
	

    if (opts->dosagefile == "")
		ss << bed->CHR[snp] << "\t" << bed->pos[snp] << "\t" << bed->RSID[snp] << "\t" << bed->A1[snp] << "\t" << bed->A2[snp];
    else
		ss << dosage->RSID[snp];
    for (i = 0 ; i < Npheno; i++) ss << "\t" << setprecision(opts->precision) << fixed << exp(NR->beta[1 + i*(Npred + 1)]);
    ss << "\t" << Lnull << "\t" << Lalt;
	
	//output frequentist or Bayesian results, as requested
    if (opts->priorfile == "" && !opts->defaultPrior){
		p = pchisq(2*(Lalt - Lnull),Npheno); //p-value from LRT
		if (p < 1e-3){
			ss << "\t" << setprecision(opts->precision) << scientific << p << endl;
		} else {
			ss << "\t" << setprecision(opts->precision) << fixed << p << endl;
		}
	} else ss << "\t" << Lalt - Lnull << endl; //Bayes factor
	
	assocOutput = ss.str();
}

void MultinomThread::calcWald(){
 	//do a wald test for each phenotype, and save the output line to waldOutput
	
	double SE, Z;
    int i;
	
	stringstream ss;
	if (opts->dosagefile == "") 
		ss << bed->CHR[snp] << "\t" << bed->pos[snp] << "\t" << bed->RSID[snp] << "\t" << bed->A1[snp] << "\t" << bed->A2[snp];
	else
	ss << dosage->RSID[snp];
	for (i = 0; i < Npheno; i++){
		SE = sqrt(NR->I[1 + i*(Npred + 1) + Npheno*(Npred + 1)*(1 + i*(Npred + 1))]);
		Z = NR->beta[1 + i*(1 + Npred)]/SE;
		ss << "\t" <<  SE << "\t" << Z << "\t" << pnorm(fabs(Z))*2;
	}
	ss << endl;
	
	waldOutput = ss.str();

}

void  MultinomThread::calcError(){
	// extract the error matrix for the genotypic effect sizes and save the output line to errorOutput
	
	stringstream ss;
	int i, j;
	
	if (opts->dosagefile == "") 
		ss << bed->CHR[snp] << "\t" << bed->pos[snp] << "\t" << bed->RSID[snp] << "\t" << bed->A1[snp] << "\t" << bed->A2[snp];
	else
		ss << dosage->RSID[snp];
	for (i = 0; i < Npheno; i++){
		for (j = 0; j < Npheno; j++){
			ss << "\t" << NR->I[1 + i*(Npred + 1) + Npheno*(Npred + 1)*(1 + j*(Npred + 1))];
		}
	}
	ss << endl;
	errorOutput = ss.str();
}

void MultinomThread::calcFullMD(){
	// extract the full model fit (genotypes and intercepts), and the full error matrix on these parameters, and save them to modelParOutput and modelErrOutput respectively
	
	stringstream ss1, ss2;
	int i, j;
	
	if (opts->dosagefile == ""){ 
		ss1 << bed->CHR[snp] << "\t" << bed->pos[snp] << "\t" << bed->RSID[snp] << "\t" << bed->A1[snp] << "\t" << bed->A2[snp];
 		ss2 << bed->CHR[snp] << "\t" << bed->pos[snp] << "\t" << bed->RSID[snp] << "\t" << bed->A1[snp] << "\t" << bed->A2[snp];
	} else {
		ss1 << dosage->RSID[snp];
		ss2 << dosage->RSID[snp];
	}
	
	for (i = 0; i < Npheno*(Npred + 1); i++){
		ss1 << "\t" << NR->beta[i];
		for (j = 0; j < Npheno*(Npred + 1); j++){
			ss2 << "\t" << NR->I[i + Npheno*(Npred + 1)*j];
		}
	}
	ss1 << endl;
	ss2 << endl;
	
	modelParOutput = ss1.str();
	modelErrOutput = ss2.str();
}

void MultinomThread::calcPhenoLRTs(){
	// do a likelihood ratio test for each individual disease and store the results line to phenoLTRsOutput 
	
	stringstream ss;
	double p;
	int i;
	double Lalt, Lnull;
	
	if (opts->dosagefile == "")
		ss << bed->CHR[snp] << "\t" << bed->pos[snp] << "\t" << bed->RSID[snp] << "\t" << bed->A1[snp] << "\t" << bed->A2[snp];
    else
		ss << dosage->RSID[snp];
	
    for (i = 0; i < Npheno*(Npred + 1); i++) constrain[i] = 0;
	
	for (i = 0; i < Npheno; i++) {
		constrain[1 + i*(Npred + 1)] = 1;
   		NR->zeroBetas();
   		NR->fitLogistic(constrain,opts->fititer,opts->fitreltol, opts->fitabstol);
   		Lnull = NR->calcL();
  		constrain[1 + i*(Npred + 1)] = 0; 
  		p = pchisq(2*(Lalt - Lnull),1);
        if (p < 1e-3){
			ss << "\t" << setprecision(opts->precision) << scientific << p;
        } else {
  			ss << "\t" << setprecision(opts->precision) << fixed << p;
        }
    }
 	ss << endl;
	phenoLTRsOutput = ss.str();
}

void MultinomThread::RunTests(){
	// run all the tests (as specified in opts) for SNP i	
	
	// load up the new SNP
	LoadData();
	
	// calculate model selection statistics, if requested
	if (opts->modelSelect) calcModelSel();

	// calculate association statistics
	calcAssoc();

	// calculate the error matrix, if we will need it
	if (opts->outputErrs || opts->outputFullMd || opts->ses){
		NR->calcI();
		inverse(NR->I,(Npred+1)*Npheno);
	}

	// calculate Wald statistics, if requested
	if (opts->ses) calcWald();

	// extract the error matrix, if requested 
	if (opts->outputErrs) calcError();

	// extract the full model statistics, if requested
    if (opts->outputFullMd) calcFullMD();

	// do per-phenotype LRTs, if requested
    if (opts->PhenoPs) calcPhenoLRTs();
}