#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
#include <time.h> 


using namespace std;


class EstCovMat{
	// Class for storing summary statistics 
	public:
		int k; // number of phenotypes
		int Nloci; // number of loci
		vector<double> betas; // Effect sizes (k x Nloci)
		vector<double> errs; // standard errors on effect sizes (k x k x Nloci)
		EstCovMat(int newK);
		EstCovMat();
		void AddLocus(double *newBetas, double *newErrs);
		double *InferCovariance();
};


EstCovMat::EstCovMat(){
	//initialize the object, starting with no loci or phenotypes
	k = 0;
	Nloci = 0;
}

EstCovMat::EstCovMat(int newK){
	//initialize the object, starting with no loci but newK phenotypes
	k = newK;
	Nloci = 0;
}

void EstCovMat::AddLocus(double *newBetas, double *newErrs){
	// add summary statistics for one new locus
	int i;
	for (i = 0; i < k; i++) betas.push_back(newBetas[i]);
	for (i = 0; i < k*k; i++) errs.push_back(newErrs[i]);
	Nloci++;
}

double *EstCovMat::InferCovariance(){
	// estimate the underying covariance matrix in effect sizes from summary statistics across all loci
	 
	 	
	int k2 = k*(k+1)/2; // the number of parameters in (lower half of) the covariance matrix 
	
	// parameter values:
	double *s = new double [k2]; // the lower triangular of the covariance matrix at this iteration (this will be what will be output at the end)
	double *grad = new double [k2]; //the gradiant of the log-likelihood at this iteration
	double *hess = new double [k2*k2]; //the hessian of the log-likelihood at this iteration
	
	// full matrix representations of above
	double * Sigma = new double [k*k]; //the full representation of the covariance matrix
	double *gradBig = new double [k*k]; //the full representation of the gradiant matrix

	//itermediate and temporary variables
	double *W = new double [k*k]; 
	double *B = new double [k*k];
	double *WB = new double [k*k];
	double *tempMat = new double [k*k];
	double *tempMat2 = new double [k*k];
	int *test = new int [k*k]; 
	double temp;
	int i, j;
	int x,y,z;
	int a,b;
	int iters ;
	int loc ;
	
	// variables required by LAPACK functions
	int * ipiv = new int[3];
	int lwork = k*k*k*k;
	double *work = new double [lwork];
	double *DebugMat;
	int info;
	
	// initialize the parameter matrix to zero
	for (i = 0; i < k2; i++) s[i] = 0;
	
	// start the NR algorithm, running up to 100 iterations
	for (iters = 0; iters < 100; iters++){
		
		// make a full representation of the covariance matrix (for matrix operations)
		for (i = 0; i < k; i++){
			for (j = i; j < k; j++){
				Sigma[i + j*k] = s[j + i*k - i*(i+1)/2];
				Sigma[j + i*k] = s[j + i*k - i*(i+1)/2];
			}
		}
		
		// initialize gradiant and hessian to zero
		for (i = 0; i < k*k; i++) gradBig[i] = 0;
		for (i = 0; i < k2; i++) grad[i] = 0;
		for (i = 0; i < k2*k2; i++) hess[i] = 0;
		
		// calculate the gradiant and the hessian at this iteration
		for (loc = 0; loc < Nloci; loc++){
			for (i = 0; i < k; i++){
				for (j = 0; j < k; j++){
					B[i + k*j] = betas[i + loc*k]*betas[j + loc*k];
				}
			}

			for (i = 0; i < k*k; i++) W[i] = Sigma[i] + errs[i + loc*k*k];
			int N2 = k;
		    dgetrf_(&N2,&N2,W,&N2,ipiv,&info);	
		    dgetri_(&N2,W,&N2,ipiv,work,&lwork,&info);

			for (i = 0; i < k; i++){
				for (j = 0; j < k; j++){
					WB[i + k*j] = 0;
					for (a = 0; a < k; a++){
						WB[i + k*j] += W[i + a*k]*B[a + j*k];
					}
				}
			}
			
			for (i = 0; i < k; i++){
				for (j = 0; j < k; j++){
					temp = 0;
					for (a = 0; a < k; a++){
						temp += WB[i + a*k]*W[a + j*k];
					}
					if (i == j){
						gradBig[i + k*j] += (W[i + j*k] - temp)*0.5;
					} else {
						gradBig[i + k*j] += (W[i + j*k] - temp);
					} 
				}
			}
			
			for (x = 0; x < k2; x++){
				a = floor( ( 2*k + 1 - sqrt( (2*k+1)*(2*k+1) - 8*x ) ) / 2 ) ;
				b = x - k*a + a*(a+1)/2;
				grad[x] = gradBig[a + b*k];
				
				for (i = 0; i < k*k; i++) {test[i] = 0; tempMat[i] = 0;}
				test[a + k*b] = 1;
				test[b + k*a] = 1;
				for (i = 0; i < k*k; i++) tempMat[i] = 0;
				for (i = 0; i < k; i++){
					for (j = 0; j < k; j++){
						tempMat[i + j*k] -= test[i + j*k];
						for (z = 0; z < k; z++){
							tempMat[i + j*k] += test[i + z*k]*WB[z + j*k];
							tempMat[j + i*k] += test[i + z*k]*WB[z + j*k];
						}
					
					}
				}
				for (i = 0; i < k; i++){
					for (j = 0; j < k; j++){
						tempMat2[i + j*k] = 0;
						for (z = 0; z < k; z++){
							tempMat2[i + j*k] += W[i + z*k] * tempMat[z + j*k];
						}
				
					}
				}
				for (i = 0; i < k; i++){
					for (j = 0; j < k; j++){
						tempMat[i + j*k] = 0;
						for (z = 0; z < k; z++){
							if (i == j){
								tempMat[i + j*k] += tempMat2[i + z*k]*W[z + j*k]*0.5;
							} else {
								tempMat[i + j*k] += tempMat2[i + z*k]*W[z + j*k];
							}
						}
				
					}
				}

				for (z = 0; z < k2; z++){
					a = floor( ( 2*k + 1 - sqrt( (2*k+1)*(2*k+1) - 8*z ) ) / 2 ) ;
					b = z - k*a + a*(a+1)/2;
					hess[z + x*k2] += tempMat[a + b*k];
				}
			}
		
		}
		//end calculation of hessian and gradiant
		
		//get the step change at this iteration
		int N = k2;
		int nrhs = 1;
		char trans = 'N';
		solve(hess, grad, N, ipiv);

		//apply the step change
		for (i = 0; i < k2; i++) s[i] = s[i] - grad[i];

		// check if the algorithm has converged, and if it has break
		double maxChange = 0;
		for (i = 0; i < k2; i++) {
			if (fabs(grad[i]/s[i]) > maxChange){
				maxChange = fabs(grad[i]/s[i]);
			}
		}
		if (maxChange < 1e-6) break;
	}
	// end NR algorithm
	
	//return the parameter estimates at termination
	return(s);
}
