/*
 * Copyright 2009-2020 The VOTCA Development Team (http://www.votca.org)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
#include <libint2/initialize.h>
#define BOOST_TEST_MAIN

#define BOOST_TEST_MODULE ewald_test

// Standard includes
#include <fstream>

// Third party includes
#include <boost/test/unit_test.hpp>

// Local VOTCA includes
#include "votca/xtp/aomatrix.h"
#include "votca/xtp/background.h"
#include "votca/xtp/convergenceacc.h"
#include "votca/xtp/orbitals.h"

// VOTCA includes
#include <votca/tools/eigenio_matrixmarket.h>

using namespace votca::xtp;
using votca::Index;
BOOST_AUTO_TEST_SUITE(ewald_calculations)

BOOST_AUTO_TEST_CASE(writereadcycle_statefile) {
  libint2::initialize();

  // Create a fake polar background
  std::vector<PolarSegment> backgroundSegments;
  PolarSegment seg0("segment0", 0);
  PolarSegment seg1("segment1", 1);
  seg0.push_back(PolarSite(0, "H", Eigen::Vector3d(1, 1, 1)));
  seg0.push_back(PolarSite(1, "H", Eigen::Vector3d(1, 1, 1.5)));
  seg1.push_back(PolarSite(0, "H", Eigen::Vector3d(3, 3, 3)));
  seg1.push_back(PolarSite(1, "H", Eigen::Vector3d(2.7, 3, 3.7)));
  backgroundSegments.push_back(seg0);
  backgroundSegments.push_back(seg1);

  // Other necessary input
  Logger log;
  EwaldOptions options;
  options.alpha = 1;
  options.k_cutoff = 1;
  options.r_cutoff = 1;
  options.shape = Shape::cube;
  options.sharpness = 0.39;

  // Create background
  Background bg(log, 5 * Eigen::Matrix3d::Identity(), options,
                backgroundSegments);

  // Write to state file
  bg.writeToStateFile("testFile.hdf5");

  // Read From state file
  Background bg2(log);
  bg2.readFromStateFile("testFile.hdf5");
  // Compare
  bool background_is_equal = (bg == bg2);
  BOOST_CHECK(background_is_equal);

  libint2::finalize();
}

BOOST_AUTO_TEST_CASE(energy_calc_NaCl){
  Index cry_l = 16;
  Index N = cry_l * cry_l * cry_l;
  Index l = (double)cry_l / 2.0;
  std::vector<PolarSegment> backgroundSegments;
  for (Index i = 0; i < N; ++i) {
    PolarSegment seg("seg", i);
    // note element type doesn't matter for the calcul
    seg.push_back(PolarSite(0, "H", ))

    Index ix = i % cry_l;
    Index iy = (i % (cry_l * cry_l)) / cry_l;
    Index iz = i / (cry_l * cry_l);

    xyz[i][0] = (double)(ix)*0.5;
    xyz[i][1] = (double)(iy)*0.5;
    xyz[i][2] = (double)(iz)*0.5;

    charges[i] = (((ix + iy + iz) % 2) ? 1.0 : -1.0);
  }
  EwaldOptions options;
  options.alpha = 4;
  options.r_cutoff = 1.9;
  options.k_cutoff = 30;
  options.shape = Shape::cube;
  options.sharpness = 0.39;
  Eigen::Matrix3d cellMatrix;
  cellMatrix << l, 0, 0, 0, l, 0, 0, 0, l;

  Logger log;
  Background bg(log, cellMatrix, options, backgroundSegments);

  double energy = bg.interactionEnergy();

}

BOOST_AUTO_TEST_CASE(energy_eval_with_induced_dipoles){
  BOOST_CHECK(true);
}

BOOST_AUTO_TEST_CASE(field_eval_with_induced_dipoles){
BOOST_CHECK(true);
}

BOOST_AUTO_TEST_CASE(background_polarization){
  BOOST_CHECK(true);
}

BOOST_AUTO_TEST_SUITE_END()


#ifndef VBAYES_HPP
#define VBAYES_HPP

#include "parameters.hpp"
#include "data.hpp"
#include "eigen_utils.hpp"
#include "typedefs.hpp"
#include "file_utils.hpp"
#include "variational_parameters.hpp"
#include "vbayes_tracker.hpp"
#include "hyps.hpp"
#include "mpi_utils.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <ctime>
#include <cstdint>
#include <cstdio>
#include <iostream>
#include <iterator>
#include <limits>
#include <queue>
#include <random>
#include <thread>
#include <set>
#include "sys/types.h"
#include "tools/eigen3.3/Dense"
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/device/file.hpp>
#include <boost/math/distributions/fisher_f.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/filesystem.hpp>

namespace boost_m = boost::math;
namespace boost_io = boost::iostreams;

inline double sigmoid(double x){
	return 1.0 / (1.0 + std::exp(-x));
}

class VBayes {
public:
// Constants
	const double PI = 3.1415926535897;
	const double eps = std::numeric_limits<double>::min();
	const double alpha_tol = 1e-4;
	const double logw_tol = 1e-2;
	const double sigma_c = 10000;
	std::vector< std::string > covar_names;
	std::vector< std::string > env_names;

// sizes
	int n_effects;
	long n_samples;
	long n_covar;
	long n_env;
	long n_var;
	long n_var2;
	double Nglobal;
	int world_rank;
	bool first_covar_update;

// Chromosomes in data
	long n_chrs;
	std::vector<long> chrs_present;
	std::vector<long> chrs_index;

	parameters p;
	std::vector<long> fwd_pass;
	std::vector<long> back_pass;
	std::vector< std::vector <long> > main_fwd_pass_chunks, gxe_fwd_pass_chunks;
	std::vector< std::vector <long> > main_back_pass_chunks, gxe_back_pass_chunks;
	std::vector<long> env_fwd_pass, covar_fwd_pass;
	std::vector<long> env_back_pass, covar_back_pass;
	std::map<long, Eigen::MatrixXd> XtX_block_cache, ZtZ_block_cache;

// Data
	GenotypeMatrix&  X;
	EigenDataMatrix& Y;
	EigenDataArrayX Cty;
	EigenDataMatrix& E;
	EigenDataMatrix& C;
	Eigen::MatrixXd CtCRidgeInv;

	Eigen::ArrayXXd& dXtEEX_lowertri;
	std::unordered_map<long, bool> sample_is_invalid;
	std::map<long, int> sample_location;

// Global location of y_m = E[X beta] and y_x = E[X gamma]
	EigenDataMatrix YY, YX, YM, ETA, ETA_SQ;

// genome wide scan computed upstream
	Eigen::ArrayXXd& snpstats;

// Init points
	VariationalParametersLite vp_init;
	std::vector<Hyps> hyps_inits;

// Monitoring
	std::chrono::system_clock::time_point time_check;
	std::chrono::duration<double> elapsed_innerLoop;

	std::vector<Eigen::VectorXd> loco_phenos;

	explicit VBayes(Data& dat) : X(dat.G),
		Y(dat.Y),
		C(dat.C),
		E(dat.E),
		dXtEEX_lowertri(dat.dXtEEX_lowertri),
		snpstats(dat.snpstats),
		p(dat.p),
		hyps_inits(dat.hyps_inits),
		sample_location(dat.sample_location),
		sample_is_invalid(dat.sample_is_invalid),
		vp_init(dat.vp_init){
		MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
#ifdef EIGEN_USE_MKL_ALL
		mkl_set_num_threads_local(p.n_thread);
#endif
		// Data size params
		n_effects      = dat.n_effects;
		n_var          = dat.n_var;
		n_env          = dat.n_env;
		n_var2         = n_effects * dat.n_var;
		n_samples      = dat.n_samples;
		n_covar        = dat.n_covar;
		covar_names    = dat.covar_names;
		env_names      = dat.env_names;
		auto Nlocal    = (double) n_samples;
		Nglobal        = mpiUtils::mpiReduce_inplace(&Nlocal);
		// E = dat.E;
		first_covar_update = true;

		assert(Y.rows() == n_samples);
		assert(X.rows() == n_samples);
		assert(n_covar == C.cols());
		assert(n_covar == covar_names.size());
		assert(n_env == E.cols());
		assert(n_env == env_names.size());

		std::set<int> tmp(X.chromosome.begin(), X.chromosome.end());
		chrs_present.assign(tmp.begin(), tmp.end());
		n_chrs = chrs_present.size();
		for (int cc = 0; cc < n_chrs; cc++) {
			chrs_index.push_back(cc);
		}

		p.main_chunk_size = (unsigned int) std::min((long int) p.main_chunk_size, (long int) n_var);
		p.gxe_chunk_size = (unsigned int) std::min((long int) p.gxe_chunk_size, (long int) n_var);

		// When n_env > 1 this gets set when in updateEnvWeights
		if(n_env == 0) {
			for (auto &hyps : hyps_inits) {
				hyps.s_x << n_var;
			}
		} else {
			for (auto &hyps : hyps_inits) {
				hyps.s_x << n_var, n_var;
			}
		}

		// Allocate memory - fwd/back pass vectors
		for(long kk = 0; kk < n_var * n_effects; kk++) {
			fwd_pass.push_back(kk);
			back_pass.push_back(n_var2 - kk - 1);
		}

		for(long ll = 0; ll < n_env; ll++) {
			env_fwd_pass.push_back(ll);
			env_back_pass.push_back(n_env - ll - 1);
		}

		for(long ll = 0; ll < n_covar; ll++) {
			covar_fwd_pass.push_back(ll);
			covar_back_pass.push_back(n_covar - ll - 1);
		}

		// ceiling of n_var / chunk size
		long n_main_segs, n_gxe_segs;
		n_main_segs = (n_var + p.main_chunk_size - 1) / p.main_chunk_size;
		n_gxe_segs = (n_var + p.gxe_chunk_size - 1) / p.gxe_chunk_size;

		main_fwd_pass_chunks.resize(n_main_segs);
		main_back_pass_chunks.resize(n_main_segs);
		gxe_fwd_pass_chunks.resize(n_gxe_segs);
		gxe_back_pass_chunks.resize(n_gxe_segs);
		for(long kk = 0; kk < n_var; kk++) {
			long main_ch_index = kk / p.main_chunk_size;
			long gxe_ch_index = kk / p.gxe_chunk_size;
			main_fwd_pass_chunks[main_ch_index].push_back(kk);
			main_back_pass_chunks[n_main_segs - 1 - main_ch_index].push_back(kk);
			gxe_fwd_pass_chunks[gxe_ch_index].push_back(kk + n_var);
			gxe_back_pass_chunks[n_gxe_segs - 1 - gxe_ch_index].push_back(kk + n_var);
		}

		for (long ii = 0; ii < n_main_segs; ii++) {
			std::reverse(main_back_pass_chunks[ii].begin(), main_back_pass_chunks[ii].end());
		}
		for (long ii = 0; ii < n_gxe_segs; ii++) {
			std::reverse(gxe_back_pass_chunks[ii].begin(), gxe_back_pass_chunks[ii].end());
		}

		if(n_effects == 1) {
			gxe_back_pass_chunks.clear();
			gxe_fwd_pass_chunks.clear();
		}

		// Generate initial values for each run
		if(p.mode_random_start) {
			std::cout << "Beta and gamma initialised with random draws" << std::endl;
		}

		// Initialise summary vars for vp_init
		if(n_env > 0) {
			vp_init.eta     = E * vp_init.muw.matrix().cast<scalarData>();
			vp_init.eta_sq  = vp_init.eta.array().square().matrix();
			vp_init.eta_sq += E.cwiseProduct(E) * vp_init.sw_sq.matrix().template cast<scalarData>();
		}
		calcPredEffects(vp_init);

		// Cache Cty
		if(n_covar > 0) {
			if (p.debug) std::cout << "Caching Cty" << std::endl;
			EigenDataArrayX Ctylocal;
			Ctylocal = C.transpose() * Y;
			Cty.resize(Ctylocal.rows(), Ctylocal.cols());
			mpiUtils::mpiReduce_double(Ctylocal.data(), Cty.data(), Ctylocal.size());
		}

		// Update main effects
		cache_local_ldblocks(main_fwd_pass_chunks, true);
		cache_local_ldblocks(main_back_pass_chunks, false);
	}

	void cache_local_ldblocks(std::vector<std::vector<long> >iter_chunks, bool is_fwd_pass){
		EigenDataMatrix D;
		for (std::uint32_t ch = 0; ch < iter_chunks.size(); ch++) {
			std::vector<long> chunk = iter_chunks[ch];
			int ee = chunk[0] / n_var;
			long ch_len = chunk.size();
			if (D.cols() != ch_len) {
				D.resize(n_samples, ch_len);
			}
			X.col_block3(chunk, D);

			unsigned long memoize_id = ((is_fwd_pass) ? ch : ch + iter_chunks.size());
			if (XtX_block_cache.count(memoize_id) == 0) {
			    XtX_block_cache[memoize_id] = (D.transpose() * D).template cast<double>();
				XtX_block_cache[memoize_id] = mpiUtils::mpiReduce_inplace(XtX_block_cache[memoize_id]);
			}
		}
	}

	void run(){
		std::cout << std::endl << "Starting variational inference ";
		std::cout << "(will terminate once change in ELBO falls below ";
		std::cout << p.elbo_tol << ")" << std::endl;
#ifndef OSX
		long long kbMax, kbGlobal, kbLocal = fileUtils::getValueRAM();
		MPI_Allreduce(&kbLocal, &kbMax, 1, MPI_LONG_LONG, MPI_MAX, MPI_COMM_WORLD);
		MPI_Allreduce(&kbLocal, &kbGlobal, 1, MPI_LONG_LONG, MPI_SUM, MPI_COMM_WORLD);
		double gbGlobal = kbGlobal / 1000.0 / 1000.0;
		double gbMax = kbMax / 1000.0 / 1000.0;
		if(world_rank == 0) {
			printf(" (RAM usage: %.2f GB in total; max of %.2f GB per rank)", gbGlobal, gbMax);
		}
#endif
		std::cout << std::endl;

		time_check = std::chrono::system_clock::now();
		long n_grid = hyps_inits.size();
		std::vector< VbTracker > trackers(n_grid, p);
		run_inference(hyps_inits, false, 2, trackers);
		write_converged_hyperparams_to_file("", trackers, n_grid);
	}

	void run_inference(const std::vector<Hyps>& hyps_inits,
	                   const bool random_init,
	                   const int round_index,
	                   std::vector<VbTracker>& trackers){
		// Writes results from inference to trackers

		long n_grid = hyps_inits.size();
		// Parrallel starts swapped for multithreaded inference
		int n_thread = 1;

		// Divide grid of hyperparameters into chunks for multithreading
		std::vector< std::vector< int > > chunks(n_thread);
		for (int ii = 0; ii < n_grid; ii++) {
			int ch_index = (ii % n_thread);
			chunks[ch_index].push_back(ii);
		}

		runOuterLoop(round_index, hyps_inits, n_grid, chunks[0], random_init, trackers);

		// Set run with best ELBO to vp_init
		std::vector< double > weights(n_grid);
		if(n_grid > 1) {
			for (int ii = 0; ii <n_grid; ii++) {
				weights[ii] = trackers[ii].logw;
			}
			normaliseLogWeights(weights);
		} else {
			weights[0] = 1;
		}

		long ii_map = std::distance(weights.begin(), std::max_element(weights.begin(), weights.end()));
		vp_init = trackers[ii_map].vp;

		// Compute residual phenotypes
		compute_residuals_per_chr(vp_init, loco_phenos);
	}

	void runOuterLoop(const int round_index,
	                  std::vector<Hyps> all_hyps,
	                  const unsigned long n_grid,
	                  const std::vector<int>& grid_index_list,
	                  const bool random_init,
	                  std::vector<VbTracker>& all_tracker){

		// Run outer loop - don't update trackers
		auto innerLoop_start = std::chrono::system_clock::now();
		runInnerLoop(random_init, round_index, all_hyps, all_tracker);
		auto innerLoop_end = std::chrono::system_clock::now();
		elapsed_innerLoop = innerLoop_end - innerLoop_start;

		long max_count = 0;
		for (const auto& tracker : all_tracker) {
			max_count = std::max(max_count, tracker.count);
		}
		std::cout << "Variational inference finished after " << max_count;
		std::cout << " iterations and " << elapsed_innerLoop.count() << " seconds.";
		std::cout << std::endl << std::endl;
	}

	void runInnerLoop(const bool random_init,
	                  const int round_index,
	                  std::vector<Hyps>& all_hyps,
	                  std::vector<VbTracker>& all_tracker){
		// minimise KL Divergence and assign elbo estimate
		if (p.debug) std::cout << "Starting inner loop" << std::endl;
		int print_interval = 25;
		if(random_init) {
			throw std::logic_error("Random start depreciated.");
		}
		unsigned long n_grid = all_hyps.size();

		std::vector<VariationalParameters> all_vp;
		setup_variational_params(all_hyps, all_vp);

		// Run inner loop until convergence
		std::vector<int> converged(n_grid, 0);
		bool all_converged = false;
		std::vector<Eigen::ArrayXd> w_prev(n_grid), beta_prev(n_grid), gam_prev(n_grid), covar_prev(n_grid);
		std::vector<double> i_logw(n_grid, -1*std::numeric_limits<double>::max());

		for (int nn = 0; nn < n_grid; nn++) {
			if(world_rank == 0) {
				all_tracker[nn].init_interim_output(nn, round_index, n_effects, n_covar, n_env, env_names, all_vp[nn]);
			}
		}

		// SQUAREM objects
		std::vector<Hyps> theta0 = all_hyps;
		std::vector<Hyps> theta1 = all_hyps;
		std::vector<Hyps> theta2 = all_hyps;

		// Allow more flexible start point so that we can resume previous inference run
		long count = p.vb_iter_start;
		while(!all_converged && count < p.vb_iter_max) {
			if (p.debug) std::cout << "Iter count: " << count << std::endl;
			for (int nn = 0; nn < n_grid; nn++) {
				if(n_covar > 0) covar_prev[nn] = all_vp[nn].mean_covars().array();
				beta_prev[nn] = all_vp[nn].mean_beta().array();
				if(n_env > 0) gam_prev[nn] = all_vp[nn].mean_gam().array();
				if(n_env > 1) w_prev[nn] = all_vp[nn].mean_weights().array();
			}
			std::vector<double> logw_prev = i_logw;

			if (p.debug) std::cout << " - update params" << std::endl;
			updateAllParams(count, round_index, all_vp, all_hyps, logw_prev);

			// SQUAREM
			if (p.mode_squarem) {
				if (p.debug) std::cout << " - SQUAREM accelerator" << std::endl;
				if(count % 3 == 0) {
					theta0 = all_hyps;
				} else if (count % 3 == 1) {
					theta1 = all_hyps;
				} else if (count >= p.vb_iter_start + 2) {
					theta2 = all_hyps;
					for (int nn = 0; nn < n_grid; nn++) {
						Hyps rr = theta1[nn] - theta0[nn];
						Hyps vv = (theta2[nn] - theta1[nn]) - rr;
						double step = std::min(-rr.normL2() / vv.normL2(), -1.0);
						Hyps theta = theta0[nn] - 2 * step * rr + step * step * vv;

						// check all hyps in theta remain in valid domain
						while(!theta.domain_is_valid() || std::abs(theta.sigma - theta2[nn].sigma) >= 0.05) {
							step = std::min(step * 0.5, -1.0);
							theta = theta0[nn] - 2 * step * rr + step * step * vv;
						}

						// Copy s_x and pve from theta2
						theta.s_x = theta2[nn].s_x;
						theta.pve = theta2[nn].pve;
						theta.n_effects = theta2[nn].n_effects;
						theta.pve_large = theta2[nn].pve_large;

						all_hyps[nn] = theta;
					}
				}
			}

			// update elbo
			std::vector<double> alpha_diff(n_grid), beta_diff(n_grid), gam_diff(n_grid), covar_diff(n_grid), w_diff(n_grid);;
			for (int nn = 0; nn < n_grid; nn++) {
				i_logw[nn]     = calc_logw(all_hyps[nn], all_vp[nn]);
				if(n_covar > 0) covar_diff[nn] = (covar_prev[nn] - all_vp[nn].mean_covars().array()).abs().maxCoeff();
				beta_diff[nn] = (beta_prev[nn] - all_vp[nn].mean_beta().array()).abs().maxCoeff();
				if(n_env > 0) gam_diff[nn] = (gam_prev[nn] - all_vp[nn].mean_gam().array()).abs().maxCoeff();
				if(n_env > 1) w_diff[nn] = (w_prev[nn] - all_vp[nn].mean_weights().array()).abs().maxCoeff();
			}

			// Interim output
			for (int nn = 0; nn < n_grid; nn++) {
				all_hyps[nn].update_pve();
				if(world_rank == 0) {
					all_tracker[nn].push_interim_hyps(count, all_hyps[nn], i_logw[nn],
					                                  covar_diff[nn], beta_diff[nn], gam_diff[nn], w_diff[nn],
					                                  n_effects,
					                                  n_var, n_covar, n_env, all_vp[nn]);
				}
				if (p.param_dump_interval > 0 && count % p.param_dump_interval == 0) {
					all_tracker[nn].dump_state(std::to_string(count), n_samples, n_covar, n_var,
					                           n_env, n_effects,
					                           all_vp[nn], all_hyps[nn], Y, C,
					                           X, covar_names, env_names, sample_is_invalid,
					                           sample_location);
				}
			}

			// Diagnose convergence
			// Disallow convergence on squarem iter
			if(!p.mode_squarem || count % 3 != 2) {
				for (int nn = 0; nn < n_grid; nn++) {
					double logw_diff = std::abs(i_logw[nn] - logw_prev[nn]);
					if (p.alpha_tol_set_by_user && p.elbo_tol_set_by_user) {
						if (alpha_diff[nn] < p.alpha_tol && logw_diff < p.elbo_tol) {
							converged[nn] = 1;
						}
					} else if (p.alpha_tol_set_by_user) {
						if (alpha_diff[nn] < p.alpha_tol) {
							converged[nn] = 1;
						}
					} else if (p.elbo_tol_set_by_user) {
						if (logw_diff < p.elbo_tol) {
							converged[nn] = 1;
						}
					} else {
						if (alpha_diff[nn] < alpha_tol && logw_diff < logw_tol) {
							converged[nn] = 1;
						}
					}
				}
			}
			if (std::all_of(converged.begin(), converged.end(), [](int i){
				return i == 1;
			})) {
				all_converged = true;
			}
			count++;
			for (int nn = 0; nn < n_grid; nn++) {
				if(!converged[nn]) {
					all_tracker[nn].count_to_convergence++;
				}
			}

			// Report progress to std::cout
			if((count + 1) % print_interval == 0) {
				int n_converged = 0;
				for (auto& n : converged) {
					n_converged += n;
				}

				std::cout << "Completed " << count+1 << " iterations ";
				print_time_check();
			}
		}

		if(any_of(i_logw.begin(), i_logw.end(), [](double x) {
			return !std::isfinite(x);
		})) {
			std::cout << "WARNING: non-finite elbo estimate produced" << std::endl;
		}

		// Dump converged state
		if (p.debug) std::cout << "Dumping converged params" << std::endl;
		for (int nn = 0; nn < n_grid; nn++) {
			all_tracker[nn].dump_state("_converged", n_samples, n_covar, n_var,
			                           n_env, n_effects,
			                           all_vp[nn], all_hyps[nn], Y, C,
			                           X, covar_names, env_names, sample_is_invalid, sample_location);
		}

		// Log all things that we want to track
		for (int nn = 0; nn < n_grid; nn++) {
			all_tracker[nn].logw = i_logw[nn];
			all_tracker[nn].count = count;
			all_tracker[nn].vp = all_vp[nn].convert_to_lite();
			all_tracker[nn].hyps = all_hyps[nn];
		}
	}

	void setup_variational_params(const std::vector<Hyps>& all_hyps,
	                              std::vector<VariationalParameters>& all_vp){
		unsigned long n_grid = all_hyps.size();

		// Init global locations YM YX
		YY.resize(n_samples, n_grid);
		YM.resize(n_samples, n_grid);
		for (int nn = 0; nn < n_grid; nn++) {
			YM.col(nn) = vp_init.ym;
			YY.col(nn) = Y;
		}
		YX.resize(n_samples, n_grid);
		ETA.resize(n_samples, n_grid);
		ETA_SQ.resize(n_samples, n_grid);
		if (n_effects > 1) {
			for (int nn = 0; nn < n_grid; nn++) {
				YX.col(nn) = vp_init.yx;
				ETA.col(nn) = vp_init.eta;
				ETA_SQ.col(nn) = vp_init.eta_sq;
			}
		}

		// Init variational params
		for (int nn = 0; nn < n_grid; nn++) {
			VariationalParameters vp(p, YM.col(nn), YX.col(nn), ETA.col(nn), ETA_SQ.col(nn));
			vp.init_from_lite(vp_init);
			if(n_effects > 1) {
				vp.calcEdZtZ(dXtEEX_lowertri, n_env);
			}
			all_vp.push_back(vp);
		}
	}

/********** VB update functions ************/
	void updateAllParams(const long& count,
	                     const int& round_index,
	                     std::vector<VariationalParameters>& all_vp,
	                     std::vector<Hyps>& all_hyps,
	                     std::vector<double> logw_prev){
		std::vector< std::uint32_t > iter;
		std::vector< std::vector< std::uint32_t > > iter_chunks;
		unsigned long n_grid = all_hyps.size();
		std::vector<double> i_logw(n_grid);

		// Update covar main effects
		for (int nn = 0; nn < n_grid; nn++) {
			if (n_covar > 0) {
				updateCovarEffects(covar_fwd_pass, all_vp[nn], all_hyps[nn]);
				updateCovarEffects(covar_back_pass, all_vp[nn], all_hyps[nn]);
				check_monotonic_elbo(all_hyps[nn], all_vp[nn], count, logw_prev[nn], "updateCovarEffects");
			}
		}

		// Update main & interaction effects
		std::string ms;
		bool is_fwd_pass = (count % 2 == 0);
		if(is_fwd_pass) {
			ms = "updateAlphaMu_fwd_main";
			updateAlphaMu(main_fwd_pass_chunks, all_hyps, all_vp, is_fwd_pass, logw_prev, count);
		} else {
			ms = "updateAlphaMu_back_gxe";
			updateAlphaMu(gxe_back_pass_chunks, all_hyps, all_vp, is_fwd_pass, logw_prev, count);
		}
		for (int nn = 0; nn < n_grid; nn++) {
			check_monotonic_elbo(all_hyps[nn], all_vp[nn], count, logw_prev[nn], ms);
		}

		if(is_fwd_pass) {
			ms = "updateAlphaMu_fwd_gxe";
			updateAlphaMu(gxe_fwd_pass_chunks, all_hyps, all_vp, is_fwd_pass, logw_prev, count);
		} else {
			ms = "updateAlphaMu_back_main";
			updateAlphaMu(main_back_pass_chunks, all_hyps, all_vp, is_fwd_pass, logw_prev, count);
		}
		for (int nn = 0; nn < n_grid; nn++) {
			check_monotonic_elbo(all_hyps[nn], all_vp[nn], count, logw_prev[nn], ms);
		}

		// Update env-weights
		if (p.debug) std::cout << " - update env weights" << std::endl;
		if (n_effects > 1 && n_env > 1) {
			for (int nn = 0; nn < n_grid; nn++) {
				for (int uu = 0; uu < p.env_update_repeats; uu++) {
					updateEnvWeights(env_fwd_pass, all_hyps[nn], all_vp[nn]);
					updateEnvWeights(env_back_pass, all_hyps[nn], all_vp[nn]);
				}
				check_monotonic_elbo(all_hyps[nn], all_vp[nn], count, logw_prev[nn], "updateEnvWeights");
			}
		}

		// Maximise hyps
		if (round_index > 1 && p.mode_empirical_bayes) {
			if (p.debug) std::cout << " - max hyps" << std::endl;
			for (int nn = 0; nn < n_grid; nn++) {
				maximiseHyps(all_hyps[nn], all_vp[nn]);
				check_monotonic_elbo(all_hyps[nn], all_vp[nn], count, logw_prev[nn], "maxHyps");
			}
		}

		// Update PVE
		for (int nn = 0; nn < n_grid; nn++) {
			all_hyps[nn].update_pve();
		}
	}

	void updateCovarEffects(const std::vector<long>& iter,
	                        VariationalParameters& vp,
	                        const Hyps& hyps){
		//
		if(p.joint_covar_update) {
			if(first_covar_update) {
				if(p.debug) {
					std::cout << "Performing first VB update of covar effects" << std::endl;
				}
				Eigen::MatrixXd CtCRidge = C.transpose() * C;
				CtCRidge = mpiUtils::mpiReduce_inplace(CtCRidge);
				CtCRidge += Eigen::MatrixXd::Identity(n_covar, n_covar) / sigma_c;
				CtCRidgeInv = CtCRidge.inverse();
				first_covar_update = false;
			}

			Eigen::VectorXd rr_k = vp.muc;
			auto Calpha = C * vp.muc.matrix();
			Eigen::MatrixXd A = C.transpose() * (Y - (vp.ym + vp.yx.cwiseProduct(vp.eta)) + Calpha);
			A = mpiUtils::mpiReduce_inplace(A);
			Eigen::VectorXd muc = CtCRidgeInv * A;
			for (int cc = 0; cc < n_covar; cc++) {
				vp.sc_sq(cc) = hyps.sigma * CtCRidgeInv(cc, cc);
				vp.muc(cc) = muc(cc);
			}
			vp.ym += C * (vp.muc.matrix() - rr_k);
		} else {
			for (const auto& cc : iter) {
				double rr_k = vp.muc(cc);

				// Update s_sq
				vp.sc_sq(cc) = hyps.sigma * sigma_c / (sigma_c * (Nglobal - 1.0) + 1.0);

				// Update mu
				double Alocal = (vp.ym + vp.yx.cwiseProduct(vp.eta)).dot(C.col(cc));
				auto A = Cty(cc) - mpiUtils::mpiReduce_inplace(&Alocal);
				vp.muc(cc) = vp.sc_sq(cc) * ( (double) A + rr_k * (Nglobal - 1.0)) / hyps.sigma;

				// Update predicted effects
				double rr_k_diff     = vp.muc(cc) - rr_k;
				vp.ym += rr_k_diff * C.col(cc);
			}
		}
	}

	void updateAlphaMu(const std::vector< std::vector<long> >& iter_chunks,
	                   const std::vector<Hyps>& all_hyps,
	                   std::vector<VariationalParameters>& all_vp,
	                   const bool& is_fwd_pass,
	                   std::vector<double> logw_prev,
	                   const long& count){
		// Divide updates into chunks
		// Partition chunks amongst available threads
		unsigned long n_grid = all_hyps.size();
		EigenDataMatrix D;
		// snp_batch x n_grid
		Eigen::MatrixXd AA;
		// snp_batch x n_grid
		Eigen::MatrixXd rr_diff;

		for (std::uint32_t ch = 0; ch < iter_chunks.size(); ch++) {
			std::vector<long> chunk = iter_chunks[ch];
			int ee                 = chunk[0] / n_var;
			long ch_len   = chunk.size();

			// D is n_samples x snp_batch
			if(D.cols() != ch_len) {
				D.resize(n_samples, ch_len);
			}
			if(rr_diff.rows() != ch_len) {
				rr_diff.resize(ch_len, n_grid);
			}
			if(AA.rows() != ch_len) {
				AA.resize(ch_len, n_grid);
			}
			X.col_block3(chunk, D);

			// Most work done here
			// AA is snp_batch x n_grid
			AA = computeGeneResidualCorrelation(D, ee);

			// Update parameters based on AA
			for (int nn = 0; nn < n_grid; nn++) {
				Eigen::Ref<Eigen::VectorXd> A = AA.col(nn);

				unsigned long memoize_id = ((is_fwd_pass) ? ch : ch + iter_chunks.size());
				adjustParams(nn, memoize_id, chunk, D, A, all_hyps, all_vp, rr_diff);
			}

			// Update residuals
			if(ee == 0) {
				YM.noalias() += D * rr_diff.cast<scalarData>();
			} else {
				YX.noalias() += D * rr_diff.cast<scalarData>();
			}

			if(p.debug) {
				if(ee == 0) {
					for (int nn = 0; nn < n_grid; nn++) {
						check_monotonic_elbo(all_hyps[nn], all_vp[nn], count, logw_prev[nn], "updateAlphaMu_main_internal");
					}
				} else {
					for (int nn = 0; nn < n_grid; nn++) {
						check_monotonic_elbo(all_hyps[nn], all_vp[nn], count, logw_prev[nn], "updateAlphaMu_gxe_internal");
					}
				}
			}
		}
	}

	void adjustParams(const int& nn, const unsigned long& memoize_id,
	                  const std::vector<long>& chunk,
	                  const EigenDataMatrix& D,
	                  const Eigen::Ref<const Eigen::VectorXd>& A,
	                  const std::vector<Hyps>& all_hyps,
	                  std::vector<VariationalParameters>& all_vp,
	                  Eigen::Ref<Eigen::MatrixXd> rr_diff){

		int ee                 = chunk[0] / n_var;
		unsigned long ch_len   = chunk.size();
		Eigen::MatrixXd Dlocal(ch_len, ch_len), Dglobal(ch_len, ch_len);
		if (ee == 0) {
			if (XtX_block_cache.count(memoize_id) == 0) {
				if(p.n_thread == 1) {
					Dlocal.triangularView<Eigen::StrictlyUpper>() = (D.transpose() * D).template cast<double>();
				} else {
					Dlocal = (D.transpose() * D).template cast<double>();
				}
				MPI_Allreduce(Dlocal.data(), Dglobal.data(), Dlocal.size(), MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
				XtX_block_cache[memoize_id] = Dglobal;
			}
			_internal_updateAlphaMu_beta(chunk, A, XtX_block_cache[memoize_id], all_hyps[nn], all_vp[nn], rr_diff.col(nn));
		} else {
			if(p.gxe_chunk_size > 1) {
				auto it = ZtZ_block_cache.find(memoize_id);
				if (n_env == 1 && it != ZtZ_block_cache.end()) {
					Dglobal = ZtZ_block_cache[memoize_id];
				} else if(p.n_thread == 1) {
					Dlocal.triangularView<Eigen::StrictlyUpper>() = (D.transpose() * all_vp[nn].eta_sq.asDiagonal() * D).template cast<double>();
					MPI_Allreduce(Dlocal.data(), Dglobal.data(), Dlocal.size(), MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
				} else {
					Dlocal = (D.transpose() * all_vp[nn].eta_sq.asDiagonal() * D).template cast<double>();
					MPI_Allreduce(Dlocal.data(), Dglobal.data(), Dlocal.size(), MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
				}

				if(n_env == 1 && it == ZtZ_block_cache.end()) {
					ZtZ_block_cache[memoize_id] = Dglobal;
				}
			}
			_internal_updateAlphaMu_gam(chunk, A, Dglobal, all_hyps[nn], all_vp[nn], rr_diff.col(nn));
		}
	}

	template <typename EigenMat>
	Eigen::MatrixXd computeGeneResidualCorrelation(const EigenMat& D,
	                                               const int& ee){
		// Most work done here
		// variant correlations with residuals
		EigenMat resLocal;
		if(n_effects == 1) {
			// Main effects update in main effects only model
			resLocal.noalias() = (YY - YM).transpose() * D;
			resLocal.transposeInPlace();
		} else if (ee == 0) {
			// Main effects update in interaction model
			resLocal.noalias() = (YY - YM - YX.cwiseProduct(ETA)).transpose() * D;
			resLocal.transposeInPlace();
		} else {
			// Interaction effects
			resLocal.noalias() = D.transpose() * ((YY - YM).cwiseProduct(ETA) - YX.cwiseProduct(ETA_SQ));
		}
		EigenMat resGlobal(resLocal.rows(), resLocal.cols());
		// long P = resLocal.rows() * resLocal.cols();
		mpiUtils::mpiReduce_double(resLocal.data(), resGlobal.data(), resLocal.size());
		return(resGlobal.template cast<double>());
	}

	void _internal_updateAlphaMu_beta(const std::vector<long>& iter_chunk,
	                                  const Eigen::Ref<const Eigen::VectorXd>& A,
	                                  const Eigen::Ref<const Eigen::MatrixXd>& D_corr,
	                                  const Hyps& hyps,
	                                  VariationalParameters& vp,
	                                  Eigen::Ref<Eigen::MatrixXd> rr_k_diff){

		unsigned long ch_len = iter_chunk.size();
		int ee = 0;

		Eigen::ArrayXd alpha_cnst;
		if(p.mode_mog_prior_beta) {
			alpha_cnst  = (hyps.lambda / (1.0 - hyps.lambda) + eps).log();
			alpha_cnst -= (hyps.slab_var.log() - hyps.spike_var.log()) / 2.0;
		} else {
			alpha_cnst = (hyps.lambda / (1.0 - hyps.lambda) + eps).log() - hyps.slab_var.log() / 2.0;
		}

		// adjust updates within chunk
		Eigen::VectorXd rr_k(ch_len);
		assert(rr_k_diff.rows() == ch_len);
		for (int ii = 0; ii < ch_len; ii++) {
			long jj = iter_chunk[ii];

			// Log prev value
			rr_k(ii) = vp.mean_beta(jj);

			// Update s_sq
			vp.s1_beta_sq(jj)                        = hyps.slab_var(ee);
			vp.s1_beta_sq(jj)                       /= (hyps.slab_relative_var(ee) * (Nglobal-1) + 1);
			if(p.mode_mog_prior_beta) vp.s2_beta_sq(jj)  = hyps.spike_var(ee);
			if(p.mode_mog_prior_beta) vp.s2_beta_sq(jj) /= (hyps.spike_relative_var(ee) * (Nglobal-1) + 1);

			// Update mu
			double offset = rr_k(ii) * (Nglobal-1.0);
			for (int mm = 0; mm < ii; mm++) {
				offset -= rr_k_diff(mm, 0) * D_corr(mm, ii);
			}
			double AA = A(ii) + offset;
			vp.mu1_beta(jj)                            = vp.s1_beta_sq(jj) * AA / hyps.sigma;
			if (p.mode_mog_prior_beta) vp.mu2_beta(jj) = vp.s2_beta_sq(jj) * AA / hyps.sigma;

			// Update alpha
			double ff_k;
			ff_k                        = vp.mu1_beta(jj) * vp.mu1_beta(jj) / vp.s1_beta_sq(jj);
			ff_k                       += std::log(vp.s1_beta_sq(jj));
			if (p.mode_mog_prior_beta) ff_k -= vp.mu2_beta(jj) * vp.mu2_beta(jj) / vp.s2_beta_sq(jj);
			if (p.mode_mog_prior_beta) ff_k -= std::log(vp.s2_beta_sq(jj));
			vp.alpha_beta(jj)           = sigmoid(ff_k / 2.0 + alpha_cnst(ee));

			rr_k_diff(ii, 0) = vp.mean_beta(jj) - rr_k(ii);

			check_nan(vp.alpha_beta(jj), ff_k, offset, hyps, iter_chunk[ii], rr_k_diff, A, D_corr, vp, alpha_cnst);
		}
	}

	void check_nan(const double& alpha,
	               const double& ff_k,
	               const double& offset,
	               const Hyps& hyps,
	               const long& ii,
	               const Eigen::Ref<const Eigen::MatrixXd>& rr_k_diff,
	               const Eigen::Ref<const Eigen::VectorXd>& A,
	               const Eigen::Ref<const Eigen::MatrixXd>& D_corr,
	               VariationalParameters& vp,
	               const Eigen::Ref<const Eigen::ArrayXd>& alpha_cnst){
		// check for NaNs and spit out diagnostics if so.

		if(std::isnan(alpha)) {
			// TODO: print diagnostics to cout
			// TODO: write all snpstats to file
			std::cout << "NaN detected at SNP index: (";
			std::cout << ii % n_var << ", " << ii / n_var << ")" << std::endl;
			std::cout << "alpha_cnst" << std::endl << alpha_cnst << std::endl << std::endl;
			std::cout << "offset" << std::endl << offset << std::endl << std::endl;
			std::cout << "hyps" << std::endl << hyps << std::endl << std::endl;
			std::cout << "rr_k_diff" << std::endl << rr_k_diff << std::endl << std::endl;
			std::cout << "A" << std::endl << A << std::endl << std::endl;
			std::cout << "D_corr" << std::endl << D_corr << std::endl << std::endl;
			throw std::runtime_error("NaN detected");
		}
	}

	void _internal_updateAlphaMu_gam(const std::vector<long>& iter_chunk,
	                                 const Eigen::Ref<const Eigen::VectorXd>& A,
	                                 const Eigen::Ref<const Eigen::MatrixXd>& D_corr,
	                                 const Hyps& hyps,
	                                 VariationalParameters& vp,
	                                 Eigen::Ref<Eigen::MatrixXd> rr_k_diff){

		long ch_len = iter_chunk.size();
		int ee     = 1;

		Eigen::ArrayXd alpha_cnst;
		if(p.mode_mog_prior_gam) {
			alpha_cnst  = (hyps.lambda / (1.0 - hyps.lambda) + eps).log();
			alpha_cnst -= (hyps.slab_var.log() - hyps.spike_var.log()) / 2.0;
		} else {
			alpha_cnst = (hyps.lambda / (1.0 - hyps.lambda) + eps).log() - hyps.slab_var.log() / 2.0;
		}

		// adjust updates within chunk
		// Need to be able to go backwards during a back_pass
		Eigen::VectorXd rr_k(ch_len);
		assert(rr_k_diff.rows() == ch_len);
		for (int ii = 0; ii < ch_len; ii++) {
			// variant index
			long jj = (iter_chunk[ii] % n_var);

			// Log prev value
			rr_k(ii) = vp.mean_gam(jj);

			// Update s_sq
			double tmp = (p.gxe_chunk_size > 1 && p.n_thread == 1) ? vp.EdZtZ(jj) : D_corr(ii, ii);
			vp.s1_gam_sq(jj)                        = hyps.slab_var(ee);
			vp.s1_gam_sq(jj)                       /= (hyps.slab_relative_var(ee) * vp.EdZtZ(jj) + 1);
			if(p.mode_mog_prior_gam) vp.s2_gam_sq(jj)  = hyps.spike_var(ee);
			if(p.mode_mog_prior_gam) vp.s2_gam_sq(jj) /= (hyps.spike_relative_var(ee) * vp.EdZtZ(jj) + 1);

			// Update mu
			double offset = rr_k(ii) * vp.EdZtZ(jj);
			for (int mm = 0; mm < ii; mm++) {
				offset -= rr_k_diff(mm, 0) * D_corr(mm, ii);
			}
			double AA = A(ii) + offset;
			vp.mu1_gam(jj)                       = vp.s1_gam_sq(jj) * AA / hyps.sigma;
			if (p.mode_mog_prior_gam) vp.mu2_gam(jj) = vp.s2_gam_sq(jj) * AA / hyps.sigma;

			// Update alpha
			double ff_k;
			ff_k                        = vp.mu1_gam(jj) * vp.mu1_gam(jj) / vp.s1_gam_sq(jj);
			ff_k                       += std::log(vp.s1_gam_sq(jj));
			if (p.mode_mog_prior_gam) ff_k -= vp.mu2_gam(jj) * vp.mu2_gam(jj) / vp.s2_gam_sq(jj);
			if (p.mode_mog_prior_gam) ff_k -= std::log(vp.s2_gam_sq(jj));
			vp.alpha_gam(jj)           = sigmoid(ff_k / 2.0 + alpha_cnst(ee));

			rr_k_diff(ii, 0) = vp.mean_gam(jj) - rr_k(ii);

			check_nan(vp.alpha_gam(jj), ff_k, offset, hyps, iter_chunk[ii], rr_k_diff, A, D_corr, vp, alpha_cnst);
		}
	}

	void maximiseHyps(Hyps& hyps,
	                  const VariationalParameters& vp){

		// max sigma
		hyps.sigma  = calcExpLinear(hyps, vp);
		if (n_covar > 0) {
			hyps.sigma += (vp.sc_sq + vp.muc.square()).sum() / sigma_c;
			hyps.sigma /= (Nglobal + (double) n_covar);
		} else {
			hyps.sigma /= Nglobal;
		}

		// beta - max lambda
		int ee = 0;
		hyps.lambda[ee] = vp.alpha_beta.sum();

		// beta - max spike & slab variances
		hyps.slab_var[ee]  = (vp.alpha_beta * (vp.s1_beta_sq + vp.mu1_beta.square())).sum();
		hyps.slab_var[ee] /= hyps.lambda[ee];
		hyps.slab_relative_var[ee] = hyps.slab_var[ee] / hyps.sigma;
		if(p.mode_mog_prior_beta) {
			hyps.spike_var[ee]  = ((1.0 - vp.alpha_beta) * (vp.s2_beta_sq + vp.mu2_beta.square())).sum();
			hyps.spike_var[ee] /= ( (double)n_var - hyps.lambda[ee]);
			hyps.spike_relative_var[ee] = hyps.spike_var[ee] / hyps.sigma;
		}

		// beta - finish max lambda
		hyps.lambda[ee] /= n_var;

		// gamma
		if(n_effects > 1) {
			ee = 1;
			hyps.lambda[ee] = vp.alpha_gam.sum();

			// max spike & slab variances
			hyps.slab_var[ee]  = (vp.alpha_gam * (vp.s1_gam_sq + vp.mu1_gam.square())).sum();
			hyps.slab_var[ee] /= hyps.lambda[ee];
			hyps.slab_relative_var[ee] = hyps.slab_var[ee] / hyps.sigma;
			if(p.mode_mog_prior_gam) {
				hyps.spike_var[ee]  = ((1.0 - vp.alpha_gam) * (vp.s2_gam_sq + vp.mu2_gam.square())).sum();
				hyps.spike_var[ee] /= ( (double)n_var - hyps.lambda[ee]);
				hyps.spike_relative_var[ee] = hyps.spike_var[ee] / hyps.sigma;
			}

			// finish max lambda
			hyps.lambda[ee] /= n_var;
		}
	}

	void updateEnvWeights(const std::vector<long>& iter,
	                      Hyps& hyps,
	                      VariationalParameters& vp){

		Eigen::ArrayXd varG = vp.var_gam();
		for (int ll : iter) {
			// Log previous mean weight
			double r_ll = vp.muw(ll);

			// Update s_sq
			double denom;
			denom  = vp.yx.cwiseProduct(E.col(ll)).squaredNorm();
			if(world_rank == 0) {
				denom += (varG * dXtEEX_lowertri.col(dXtEEX_col_ind(ll, ll, n_env))).sum();
			}
			denom  = mpiUtils::mpiReduce_inplace(&denom);
			denom += hyps.sigma;
			vp.sw_sq(ll) = hyps.sigma / denom;

			// Remove dependance on current weight
			vp.eta -= (r_ll * E.col(ll)).matrix();

			// Update mu
			Eigen::ArrayXd env_vars = Eigen::ArrayXd::Zero(n_var);
			if(world_rank == 0) {
				for (int mm = 0; mm < n_env; mm++) {
					if(mm != ll) {
						env_vars += vp.muw(mm) * dXtEEX_lowertri.col(dXtEEX_col_ind(ll, mm, n_env));
					}
				}
				// MPI_Bcast(env_vars.data(), env_vars.size(), MPI_DOUBLE, 0, MPI_COMM_WORLD);
			}
			env_vars  = mpiUtils::mpiReduce_inplace(env_vars);

			double eff = ((Y - vp.ym).array() * E.col(ll).array() * vp.yx.array()).sum();
			eff       -= (vp.yx.array() * E.col(ll).array() * vp.eta.array() * vp.yx.array()).sum();
			eff        = mpiUtils::mpiReduce_inplace(&eff);
			eff       -= (varG * env_vars).sum();
			vp.muw(ll) = vp.sw_sq(ll) * eff / hyps.sigma;

			// Update eta
			vp.eta += (vp.muw(ll) * E.col(ll)).matrix();
		}

		// Recompute eta_sq
		vp.eta_sq  = vp.eta.array().square().matrix();
#ifdef DATA_AS_FLOAT
		vp.eta_sq += E.cwiseProduct(E) * vp.sw_sq.matrix().cast<float>();
#else
		vp.eta_sq += E.cwiseProduct(E) * vp.sw_sq.matrix();
#endif

		// Recompute expected value of diagonal of ZtZ
		vp.calcEdZtZ(dXtEEX_lowertri, n_env);

		double colVarZ = 0;
		if(world_rank == 0) {
			Eigen::ArrayXd muw_sq_combos(n_env * (n_env + 1) / 2);
			for (int ll = 0; ll < n_env; ll++) {
				for (int mm = 0; mm < n_env; mm++) {
					muw_sq_combos(dXtEEX_col_ind(ll, mm, n_env)) = vp.muw(mm) * vp.muw(ll);
				}
			}

			colVarZ = 2 * (dXtEEX_lowertri.rowwise() * muw_sq_combos.transpose()).sum();
			for (int ll = 0; ll < n_env; ll++) {
				colVarZ -= (vp.muw(ll) * vp.muw(ll) * dXtEEX_lowertri.col(dXtEEX_col_ind(ll, ll, n_env))).sum();
			}
			colVarZ /= (Nglobal - 1.0);
		}
		colVarZ = mpiUtils::mpiReduce_inplace(&colVarZ);

		// WARNING: Hard coded index
		// WARNING: Updates S_x in hyps
		hyps.s_x(0) = (double) n_var;
		hyps.s_x(1) = colVarZ;
	}

	double calc_logw(const Hyps& hyps,
	                 const VariationalParameters& vp){

		// Expectation of linear regression log-likelihood
		double int_linear = -1.0 * calcExpLinear(hyps, vp) / 2.0 / hyps.sigma;
		int_linear -= Nglobal * std::log(2.0 * PI * hyps.sigma) / 2.0;

		// kl-beta
		double kl_beta = calcKLBeta(hyps, vp);

		double kl_gamma = 0;
		if(n_effects > 1) {
			kl_gamma += calcKLGamma(hyps, vp);
		}

		// covariates
		double kl_covar = 0.0;
		if(n_covar > 0) {
			kl_covar += (double) n_covar * (1.0 - std::log(hyps.sigma * sigma_c)) / 2.0;
			kl_covar += vp.sc_sq.log().sum() / 2.0;
			kl_covar -= vp.sc_sq.sum() / 2.0 / hyps.sigma / sigma_c;
			kl_covar -= vp.muc.square().sum() / 2.0 / hyps.sigma / sigma_c;
		}

		// weights
		double kl_weights = 0.0;
		if(n_effects > 1 && n_env > 1) {
			kl_weights += (double) n_env / 2.0;
			kl_weights += vp.sw_sq.log().sum() / 2.0;
			kl_weights -= vp.sw_sq.sum() / 2.0;
			kl_weights -= vp.muw.square().sum() / 2.0;
		}

		double res = int_linear + kl_beta + kl_gamma + kl_covar + kl_weights;

		return res;
	}

	/********** Helper functions ************/
	void print_time_check(){
		auto now = std::chrono::system_clock::now();
		std::chrono::duration<double> elapsed_seconds = now-time_check;
		std::cout << " (" << elapsed_seconds.count();
		std::cout << " seconds since last timecheck";
#ifndef OSX
		std::cout << ", estimated RAM usage = ";
		std::cout << fileUtils::getValueRAM() << "KB)" << std::endl;
#else
		std::cout << ")" << std::endl;
#endif
		time_check = now;
	}

	void calcPredEffects(VariationalParametersLite& vp) const {
		Eigen::VectorXd rr_beta = vp.mean_beta();

		vp.ym = X * rr_beta;
		if(n_covar > 0) {
			vp.ym += C * vp.muc.matrix().cast<scalarData>();
		}

		if(n_effects > 1) {
			Eigen::VectorXd rr_gam = vp.mean_gam();
			vp.yx = X * rr_gam;
		}
	}

	void check_monotonic_elbo(const Hyps& hyps,
	                          VariationalParameters& vp,
	                          const long& count,
	                          double& logw_prev,
	                          const std::string& prev_function){
		double i_logw     = calc_logw(hyps, vp);
		if(i_logw < logw_prev) {
			std::cout << count << ": " << prev_function;
			std::cout << " " << logw_prev << " -> " << i_logw;
			std::cout << " (difference of " << i_logw - logw_prev << ")"<< std::endl;
		}
		logw_prev = i_logw;
	}

	void normaliseLogWeights(std::vector< double >& my_weights){
		// Safer to normalise log-weights than niavely convert to weights
		// Skip non-finite values!
		long nn = my_weights.size();
		double max_elem = *std::max_element(my_weights.begin(), my_weights.end());
		for (int ii = 0; ii < nn; ii++) {
			my_weights[ii] = std::exp(my_weights[ii] - max_elem);
		}

		double my_sum = 0.0;
		for (int ii = 0; ii < nn; ii++) {
			if(std::isfinite(my_weights[ii])) {
				my_sum += my_weights[ii];
			}
		}

		for (int ii = 0; ii < nn; ii++) {
			my_weights[ii] /= my_sum;
		}

		int nonfinite_count = 0;
		for (int ii = 0; ii < nn; ii++) {
			if(!std::isfinite(my_weights[ii])) {
				nonfinite_count++;
			}
		}

		if(nonfinite_count > 0) {
			std::cout << "WARNING: " << nonfinite_count << " grid points returned non-finite ELBO.";
			std::cout << "Skipping these when producing posterior estimates.";
		}
	}

	double calcExpLinear(const Hyps& hyps,
	                     const VariationalParameters& vp){
		// Expectation of ||Y - C tau - X beta - Z gamma||^2
		double int_linear = 0, resLocal = 0;

		// Expectation of linear regression log-likelihood
		resLocal  = (Y - vp.ym).squaredNorm();
		if(n_effects > 1) {
			resLocal -= 2.0 * (Y - vp.ym).cwiseProduct(vp.eta).cwiseProduct(vp.yx).sum();
			if (n_env > 1) {
				resLocal += vp.yx.cwiseProduct(vp.eta_sq).dot(vp.yx);
			} else {
				resLocal += vp.yx.cwiseProduct(vp.eta).squaredNorm();
			}
		}
		MPI_Allreduce(&resLocal, &int_linear, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

		// variances
		if(n_covar > 0) {
			// covar main
			int_linear += (Nglobal - 1.0) * vp.sc_sq.sum();
		}
		// beta
		int_linear += (Nglobal - 1.0) * vp.var_beta().sum();
		if(n_effects > 1) {
			// gamma
			int_linear += (vp.EdZtZ * vp.var_gam()).sum();
		}

		return int_linear;
	}

	double calcKLBeta(const Hyps& hyps,
	                  const VariationalParameters& vp){
		// KL Divergence of log[ p(beta | u, theta) / q(u, beta) ]
		double res = 0;
		int ee = 0;

		res += std::log(hyps.lambda(ee) + eps) * vp.alpha_beta.sum();
		res += std::log(1.0 - hyps.lambda(ee) + eps) * ((double) n_var - vp.alpha_beta.sum());

		// Need std::log for eps guard to work
		for (std::uint32_t kk = 0; kk < n_var; kk++) {
			res -= vp.alpha_beta(kk) * std::log(vp.alpha_beta(kk) + eps);
			res -= (1 - vp.alpha_beta(kk)) * std::log(1 - vp.alpha_beta(kk) + eps);
		}

		// beta
		if(p.mode_mog_prior_beta) {
			res += n_var / 2.0;

			res -= vp.mean_beta_sq(1).sum() / 2.0 / hyps.slab_var(ee);
			res -= vp.mean_beta_sq(2).sum() / 2.0 / hyps.spike_var(ee);

			res += (vp.alpha_beta * vp.s1_beta_sq.log()).sum() / 2.0;
			res += ((1.0 - vp.alpha_beta) * vp.s2_beta_sq.log()).sum() / 2.0;

			res -= std::log(hyps.slab_var(ee))  * vp.alpha_beta.sum() / 2.0;
			res -= std::log(hyps.spike_var(ee)) * (n_var - vp.alpha_beta.sum()) / 2.0;
		} else {
			res += (vp.alpha_beta * vp.s1_beta_sq.log()).sum() / 2.0;
			res -= (vp.alpha_beta * (vp.mu1_beta.square() + vp.s1_beta_sq)).sum() / 2.0 / hyps.slab_var(ee);

			res += (1 - std::log(hyps.slab_var(ee))) * vp.alpha_beta.sum() / 2.0;
		}
		return res;
	}

	double calcKLGamma(const Hyps& hyps,
	                   const VariationalParameters& vp){
		// KL Divergence of log[ p(beta | u, theta) / q(u, beta) ]
		double res = 0;
		int ee = 1;

		res += std::log(hyps.lambda(ee) + eps) * vp.alpha_gam.sum();
		res += std::log(1.0 - hyps.lambda(ee) + eps) * ((double) n_var - vp.alpha_gam.sum());

		// Need std::log for eps guard to work
		for (std::uint32_t kk = 0; kk < n_var; kk++) {
			res -= vp.alpha_gam(kk) * std::log(vp.alpha_gam(kk) + eps);
			res -= (1 - vp.alpha_gam(kk)) * std::log(1 - vp.alpha_gam(kk) + eps);
		}

		// beta
		if(p.mode_mog_prior_gam) {
			res += n_var / 2.0;

			res -= vp.mean_gam_sq(1).sum() / 2.0 / hyps.slab_var(ee);
			res -= vp.mean_gam_sq(2).sum() / 2.0 / hyps.spike_var(ee);

			res += (vp.alpha_gam * vp.s1_gam_sq.log()).sum() / 2.0;
			res += ((1.0 - vp.alpha_gam) * vp.s2_gam_sq.log()).sum() / 2.0;

			res -= std::log(hyps.slab_var(ee))  * vp.alpha_gam.sum() / 2.0;
			res -= std::log(hyps.spike_var(ee)) * (n_var - vp.alpha_gam.sum()) / 2.0;
		} else {
			res += (vp.alpha_gam * vp.s1_gam_sq.log()).sum() / 2.0;
			res -= (vp.alpha_gam * (vp.mu1_gam.square() + vp.s1_gam_sq)).sum() / 2.0 / hyps.slab_var(ee);

			res += (1 - std::log(hyps.slab_var(ee))) * vp.alpha_gam.sum() / 2.0;
		}
		return res;
	}

	void
	compute_residuals_per_chr(const VariationalParametersLite &vp,
	                          std::vector<Eigen::VectorXd> &loco_phenos) const {
		std::vector<Eigen::VectorXd> pred_main, pred_int;
		pred_main.resize(n_chrs);
		pred_int.resize(n_chrs);
		loco_phenos.resize(n_chrs);

		// casts used if DATA_AS_FLOAT
		Eigen::VectorXd map_residuals;
		if (n_effects > 1) {
			map_residuals = (Y - vp.ym - vp.yx.cwiseProduct(vp.eta)).cast<double>();
		} else {
			map_residuals = (Y - vp.ym).cast<double>();
		}

		// Compute predicted effects from each chromosome
		Eigen::VectorXd Eq_beta = vp.mean_beta();
		for (auto cc : chrs_index) {
			pred_main[cc] = X.mult_vector_by_chr(chrs_present[cc], Eq_beta);
		}

		if (n_effects > 1) {
			Eigen::VectorXd Eq_gam = vp.mean_gam();
			for (auto cc : chrs_index) {
				pred_int[cc]  = X.mult_vector_by_chr(chrs_present[cc], Eq_gam);
			}
		}

		// Compute mean-centered residuals for each chromosome
		for (auto cc : chrs_index) {
			if (n_effects > 1) {
				loco_phenos[cc] = map_residuals + pred_main[cc] + pred_int[cc].cwiseProduct(vp.eta.cast<double>());
			} else {
				loco_phenos[cc] = map_residuals + pred_main[cc];
			}
			EigenUtils::center_matrix(loco_phenos[cc]);
		}
	}

//	TODO make vp const
	void compute_LOCO_pvals(VariationalParametersLite &vp,
	                        Eigen::MatrixXd &neglogPvals,
	                        Eigen::MatrixXd &testStats) {

		calcPredEffects(vp);
		std::vector<Eigen::VectorXd> resid_pheno;
		compute_residuals_per_chr(vp, resid_pheno);
		assert(resid_pheno.size() == n_chrs);
		neglogPvals.resize(n_var, (n_env > 0 ? 4 : 1));
		testStats.resize(n_var, (n_env > 0 ? 4 : 1));

		std::map<long,long> chr_changes;
		long chr = X.chromosome[0];
		for (long jj = 0; jj < n_var; jj++) {
			if (chr != X.chromosome[jj]) {
				chr_changes[chr] = jj;
				chr = X.chromosome[jj];
			}
		}
		chr_changes[chr] = n_var;

		EigenDataMatrix D;
		long max_chunk_size = 256, start = 0;
		bool append = false;
		Eigen::MatrixXd chunk_neglogPvals, chunk_testStats;
		while (start < n_var) {
			long chr = X.chromosome[start];
			long cc = std::find(chrs_present.begin(), chrs_present.end(), chr) - chrs_present.begin();

			long chunkSize = std::min(max_chunk_size,chr_changes[chr] - start);
			std::vector<long> chunk(chunkSize);
			std::iota(chunk.begin(),chunk.end(),start);

			if (D.cols() != chunkSize) {
				D.resize(n_samples, chunkSize);
			}
			X.col_block3(chunk, D);
			::compute_LOCO_pvals(resid_pheno[cc], D, chunk_neglogPvals, chunk_testStats, vp.eta);

			neglogPvals.block(start,0,chunkSize,(n_env > 0 ? 4 : 1)) = chunk_neglogPvals;
			testStats.block(start,0,chunkSize,(n_env > 0 ? 4 : 1)) = chunk_testStats;
			start += chunkSize;
		}
	}

	void compute_LOSO_pvals(GenotypeMatrix &Xtest,
	                        const VariationalParametersLite &vp,
	                        const long &LOSO_window,
	                        Eigen::Ref<Eigen::VectorXd> neglogp_beta,
	                        Eigen::Ref<Eigen::VectorXd> neglogp_gam_robust,
	                        Eigen::Ref<Eigen::VectorXd> neglogp_joint,
	                        Eigen::Ref<Eigen::VectorXd> test_stat_beta,
	                        Eigen::Ref<Eigen::VectorXd> test_stat_gam_robust,
	                        Eigen::Ref<Eigen::VectorXd> test_stat_joint) const {
		assert(neglogp_beta.rows()  == n_var);
		assert(neglogp_joint.rows() == n_var);
		assert(test_stat_beta.rows()  == n_var);
		assert(test_stat_joint.rows() == n_var);
		assert(LOSO_window >= 0);

		int n_effects = (n_env > 0) ? 2 : 1;
		double N = n_samples;

		Eigen::VectorXd y_resid = (Y - vp.ym).cast<double>();
		if(n_env > 0) {
			y_resid -= vp.yx.cwiseProduct(vp.eta).cast<double>();
		}
		long front = 0, back = 0;

		// Compute p-vals per variant (p=3 as residuals mean centered)
		Eigen::MatrixXd H(n_samples, n_effects);
		boost::math::students_t t_dist(n_samples - n_effects);
		boost::math::fisher_f f_dist(n_effects, n_samples - n_effects);
		for(uint32_t jj = 0; jj < n_var; jj++ ) {
			int chr_test = Xtest.chromosome[jj];
			long pos_test = Xtest.position[jj];

			while (front < n_var && X.position[front] < pos_test + LOSO_window && X.chromosome[front] == chr_test) {
				y_resid += vp.mean_beta(front) * X.col(front);
				if (n_env > 0) y_resid += vp.mean_gam(front) * vp.eta.cwiseProduct(X.col(front));
				front++;
			}

			while (X.position[back] < pos_test - LOSO_window || X.chromosome[back] != chr_test) {
				y_resid -= vp.mean_beta(back) * X.col(back);
				if (n_env > 0) y_resid -= vp.mean_gam(back) * vp.eta.cwiseProduct(X.col(back));
				back++;
			}

			H.col(0) = Xtest.col(jj).cast<double>();

			double rss_alt, rss_null;
			Eigen::MatrixXd HtH(H.cols(), H.cols()), Hty(H.cols(), 1);
			Eigen::MatrixXd HtH_inv(H.cols(), H.cols()), HtVH(H.cols(), H.cols());
			if(n_env == 0) {

				prep_lm(H, y_resid, HtH, HtH_inv, Hty, rss_alt);

				double beta_tstat, beta_pval;
				student_t_test(n_samples, HtH_inv, Hty, rss_alt, 0, beta_tstat, beta_pval);

				neglogp_beta(jj) = -1 * log10(beta_pval);
				test_stat_beta(jj) = beta_tstat;
			} else if (n_env > 0) {
				H.col(1) = H.col(0).cwiseProduct(vp.eta.cast<double>());

				prep_lm(H, y_resid, HtH, HtH_inv, Hty, rss_alt, HtVH);
				rss_null = y_resid.squaredNorm();
				rss_null = mpiUtils::mpiReduce_inplace(&rss_null);

				// Single-var tests
				double beta_tstat, gam_tstat, rgam_stat, beta_pval, gam_pval, rgam_pval;
				hetero_chi_sq(HtH_inv, Hty, HtVH, 1, rgam_stat, rgam_pval);
				student_t_test(n_samples, HtH_inv, Hty, rss_alt, 1, gam_tstat, gam_pval);
				student_t_test(n_samples, HtH_inv, Hty, rss_alt, 0, beta_tstat, beta_pval);

				// F-test over main+int effects of snp_j
				double joint_fstat, joint_pval;
				joint_fstat       = (rss_null - rss_alt) / 2.0;
				joint_fstat      /= rss_alt / (N - 3.0);
				joint_pval        = 1.0 - cdf(f_dist, joint_fstat);

				neglogp_beta[jj]  = -1 * log10(beta_pval);
				neglogp_gam_robust[jj]   = -1 * log10(rgam_pval);
				neglogp_joint[jj] = -1 * log10(joint_pval);
				test_stat_beta[jj]  = beta_tstat;
				test_stat_gam_robust[jj]   = rgam_stat;
				test_stat_joint[jj] = joint_fstat;
			}
		}
	}

	/********** Output functions ************/
	void write_converged_hyperparams_to_file(const std::string& file_prefix,
	                                         const std::vector< VbTracker >& trackers,
	                                         const long& my_n_grid){
		boost_io::filtering_ostream outf;
		std::string ofile       = fileUtils::fstream_init(outf, p.out_file, file_prefix, "");
		std::cout << "Writing converged hyperparameter values to " << ofile << std::endl;

		// Compute normalised weights using finite elbo
		std::vector< double > weights(my_n_grid);
		if(my_n_grid > 1) {
			for (int ii = 0; ii < my_n_grid; ii++) {
				weights[ii] = trackers[ii].logw;
			}
			normaliseLogWeights(weights);
		} else {
			weights[0] = 1;
		}

		/*** Hyps - header ***/
		outf << "count elbo sigma";

		for (int ee = 0; ee < n_effects; ee++) {
			outf << " pve" << ee;
			if((ee == 0 && p.mode_mog_prior_beta) || (ee == 1 && p.mode_mog_prior_gam)) {
				outf << " pve_large" << ee;
			}
			outf << " sigma" << ee;
			if((ee == 0 && p.mode_mog_prior_beta) || (ee == 1 && p.mode_mog_prior_gam)) {
				outf << " sigma_spike" << ee;
				outf << " sigma_spike_dilution" << ee;
			}
			outf << " lambda" << ee;
		}
		outf << std::endl;

		/*** Hyps - converged ***/
		for (int ii = 0; ii < my_n_grid; ii++) {
			outf << std::setprecision(8) << std::fixed;
			outf << " " << trackers[ii].count;
			outf << " " << trackers[ii].logw;
			outf << std::setprecision(6) << std::fixed;
			outf << " " << trackers[ii].hyps.sigma;

			outf << std::setprecision(8) << std::fixed;
			for (int ee = 0; ee < n_effects; ee++) {

				// PVE
				outf << " " << trackers[ii].hyps.pve(ee);
				if((ee == 0 && p.mode_mog_prior_beta) || (ee == 1 && p.mode_mog_prior_gam)) {
					outf << " " << trackers[ii].hyps.pve_large(ee);
				}

				// MoG variances
				outf << std::scientific << std::setprecision(5);
				outf << " " << trackers[ii].hyps.slab_relative_var(ee);
				outf << std::setprecision(8) << std::fixed;

				if((ee == 0 && p.mode_mog_prior_beta) || (ee == 1 && p.mode_mog_prior_gam)) {
					outf << std::scientific << std::setprecision(5);
					outf << " " << trackers[ii].hyps.spike_relative_var(ee);

					outf << std::setprecision(3) << std::fixed;
					outf << " " << trackers[ii].hyps.slab_relative_var(ee) / trackers[ii].hyps.spike_relative_var(ee);
					outf << std::setprecision(8) << std::fixed;
				}

				// Lambda
				outf << " " << trackers[ii].hyps.lambda(ee);
			}
			outf << std::endl;
		}
		boost_io::close(outf);
	}

	void output_vb_results() {
		if(world_rank == 0) {
			std::string format = fileUtils::filepath_format(p.out_file, "", "_converged_vparams_*");
			std::cout << "Writing variational parameters to " << format << std::endl;

			std::string path = fileUtils::filepath_format(p.out_file, "", "_converged_vparams");
			vp_init.dump_to_prefix(path, X, env_names, covar_names);
		}

		if (n_env > 0) {
			std::string path = fileUtils::filepath_format(p.out_file, "", "_converged_eta");
			std::cout << "Writing eta to " << path << std::endl;
			fileUtils::dump_predicted_vec_to_file(vp_init.eta, path, "eta", sample_location);
		}

		/*********** Stats from MAP to file ************/
		// Predicted effects to file
		calcPredEffects(vp_init);
		compute_residuals_per_chr(vp_init, loco_phenos);
		Eigen::VectorXd Ealpha = Eigen::VectorXd::Zero(n_samples);
		if(n_covar > 0) {
			Ealpha += (C * vp_init.muc.matrix().cast<scalarData>()).cast<double>();
		}

		std::string header = "Y";
		long n_cols = 1;
		if (n_covar > 0) {
			header += " Ealpha"; n_cols += 1;
		}
		header += " Xbeta"; n_cols += 1;
		if (n_env > 0) {
			header += " eta Xgamma"; n_cols += 2;
		}
		if(n_chrs > 0) {
			for(auto cc : chrs_present) {
				header += " resid_excl_chr" + std::to_string(cc);
			}
			n_cols += n_chrs;
		}

		Eigen::MatrixXd tmp(n_samples, n_cols);
		int cc = 0;
		tmp.col(cc) = Y; cc++;
		if (n_covar > 0) {
			tmp.col(cc) = Ealpha; cc++;
			tmp.col(cc) = vp_init.ym - Ealpha; cc++;
		} else {
			tmp.col(cc) = vp_init.ym; cc++;
		}
		if (n_env > 0) {
			tmp.col(cc) = vp_init.eta; cc++;
			tmp.col(cc) = vp_init.yx; cc++;
		}
		for(int cc1 = 0; cc1 < n_chrs; cc1++) {
			tmp.col(cc) = loco_phenos[cc1]; cc++;
		}
		assert(cc == n_cols);

		std::string path = fileUtils::filepath_format(p.out_file, "", "_converged_yhat");
		std::cout << "Writing predicted and residualised phenotypes to " << path << std::endl;
		fileUtils::dump_predicted_vec_to_file(tmp, path, header, sample_location);

		// Save eta & residual phenotypes to separate files
		boost_io::filtering_ostream tmp_outf;

		for (long cc = 0; cc < n_chrs; cc++) {
			std::string path = fileUtils::filepath_format(p.out_file, "",
			                                              "_converged_resid_pheno_chr" + std::to_string(chrs_present[cc]));
			std::cout << "Writing residualised pheno to " << path << std::endl;
			fileUtils::dump_predicted_vec_to_file(loco_phenos[cc], path,
			                                      "chr" + std::to_string(chrs_present[cc]),
			                                      sample_location);
		}
		std::cout << std::endl;
	}

	void compute_pvals(VariationalParametersLite vp){
		calcPredEffects(vp);
		compute_residuals_per_chr(vp, loco_phenos);

		// Compute LOCO p-values
		Eigen::MatrixXd neglogPvals, testStats;
		if(!p.drop_loco) {
			if (p.debug) std::cout << "Computing single-SNP hypothesis tests with LOCO strategy" << std::endl;
			compute_LOCO_pvals(vp, neglogPvals, testStats);
		} else {
			if (p.debug) {
				std::cout << "Computing single-snp hypothesis tests while excluding SNPs within ";
				std::cout << p.LOSO_window << " of the test SNP" << std::endl;
			}
			Eigen::VectorXd neglogp_beta(n_var), neglogp_gam(n_var), neglogp_rgam(n_var), neglogp_joint(n_var);
			Eigen::VectorXd test_stat_beta(n_var), test_stat_gam(n_var), test_stat_rgam(n_var), test_stat_joint(n_var);
			compute_LOSO_pvals(X, vp, p.LOSO_window,
			                   neglogp_beta, neglogp_rgam, neglogp_joint,
			                   test_stat_beta, test_stat_rgam, test_stat_joint);
			neglogPvals.resize(neglogp_beta.rows(), 3);
			neglogPvals.col(0) = neglogp_beta;
			neglogPvals.col(1) = Eigen::VectorXd::Constant(neglogp_beta.rows(), -1);
			neglogPvals.col(2) = neglogp_rgam;
			testStats.resize(neglogp_beta.rows(), 3);
			testStats.col(0) = test_stat_beta;
			testStats.col(1) = Eigen::VectorXd::Constant(neglogp_beta.rows(), -1);
			testStats.col(2) = test_stat_rgam;
		}

		if(world_rank == 0) {
			boost_io::filtering_ostream outf;
			std::string ofile = fileUtils::fstream_init(outf, p.out_file, "", "_loco_pvals");
			std::cout << "Writing single-SNP hypothesis test results ";
			if(p.drop_loco) {
				std::cout << "(leave out SNPs within " << p.LOSO_window << "bp)";
			} else {
				std::cout << "(computed with LOCO strategy)";
			}
			std::cout << " to " << ofile << std::endl;
			bool append = false;
			fileUtils::write_snp_stats_to_file(outf, n_effects, X, append, neglogPvals, testStats);
			boost_io::close(outf);
		}
	}
};

#endif


//---------------------------------------------------------------------------//
//!
//! \file   Utility_UnitlessMomentumConversionPolicy.hpp
//! \author Alex Robinson
//! \brief  Policy structs that define unitless momentum conversions
//!
//---------------------------------------------------------------------------//

#ifndef UTILITY_UNITLESS_MOMENTUM_CONVERSION_POLICY_HPP
#define UTILITY_UNITLESS_MOMENTUM_CONVERSION_POLICY_HPP

// Boost Includes
#include <boost/units/base_dimension.hpp>

// FRENSIE Includes
#include "Utility_UnitConversionPolicyDecl.hpp"
#include "Utility_PhysicalConstants.hpp"

//! The dimensionless momentum dimension tag
struct DimensionlessMomentum : boost::units::base_dimension<DimensionlessMomentum,1>{};

namespace Utility{

/*! \brief Policy struct for converting between atomic dimensionless momentum
 * (p/(me*e^2/h_bar)) and the native dimensionless momentum (p/(me*c))
 */
struct AtomicDimensionlessMomentumConversionPolicy : public BaseUnitConversionPolicy<AtomicDimensionlessMomentumConversionPolicy, DimensionlessMomentum>
{
  //! Atomic dimensionless momentum per native dimensionless momentum
  static inline double conversionFactor()
  { return Utility::PhysicalConstants::inverse_fine_structure_constant; }
};

/*! \brief Policy struct for converting between native dimensionless
 * momentum (p/(me*c)) and native dimensionless momentum (p/(me*c))
 */
struct MeCDimensionlessMomentumConversionPolicy : public BaseUnitConversionPolicy<MeCDimensionlessMomentumConversionPolicy, DimensionlessMomentum>
{
  // Native dimensionless momentum per native dimensionless momentum
  static inline double conversionFactor()
  { return 1.0; }
};

} // end Utility namespace

#endif // end UTILITY_UNITLESS_MOMENTUM_CONVERSION_POLICY_HPP

//---------------------------------------------------------------------------//
// end Utility_UnitlessMomentumConversionPolicy.hpp
//---------------------------------------------------------------------------//


#pragma once

#include <string>

#include <boost/format.hpp>

#include "gnuplot-iostream.h"
#include "../util/util.hpp"

#ifdef PLOT
constexpr bool plot = true;
#elif
constexpr bool plot = false;
#endif // PLOT

template <typename Data, typename F>
void plot_with_tube(std::string title, const Data &data, F ref_func,
                    double margin) {
  const std::string tubeColour = "#6699FF55";

  Gnuplot gp;
  gp << "set title '" << title << "'\n"
     << "plot '-' u 1:2:3 title 'Acceptable margin: analytical ±"
     << boost::format("%.3f") % margin << "' w filledcu fs solid fc rgb '"
     << tubeColour
     << "', '-' u 1:2 "
        "title 'Test result' w l\n";

  auto range = util::fmap(
      [&](auto x) {
        // with…
        const double t = x.first;
        const double analyt = ref_func(t);

        return std::make_tuple(t, analyt + margin, analyt - margin);
      },
      data);

  gp.send1d(range);
  gp.send1d(data);
  gp << "pause mouse key\nwhile (MOUSE_CHAR ne 'q') { pause mouse "
        "key; }\n";
}


/*
 * @BEGIN LICENSE
 *
 * libresponse_psi4 by Eric Berquist, a plugin to:
 *
 * Psi4: an open-source quantum chemistry software package
 *
 * Copyright (c) 2007-2019 The Psi4 Developers.
 *
 * The copyrights for code used from other parties are included in
 * the corresponding files.
 *
 * This file is part of Psi4.
 *
 * Psi4 is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * Psi4 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 Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with Psi4; if not, write to the Free Software Foundation, Inc., 51
 * Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * @END LICENSE
 */

#include <armadillo>

#include <libresponse/libresponse.h>
#include <libresponse/linear/interface.h>

#include "psi4/liboptions/liboptions.h"
#include "psi4/libpsio/psio.hpp"
#include "psi4/psi4-dec.h"

#include "psi4/libmints/basisset.h"
#include "psi4/libmints/matrix.h"
#include "psi4/libmints/vector.h"
#include "psi4/libmints/wavefunction.h"
#include "psi4/libpsi4util/PsiOutStream.h"

#include "matvec_psi4.h"
// #include "operatordatatype.h"
#include "parse_operators.h"
#include "wrappers.h"

namespace psi {
namespace libresponse_psi4 {

extern "C" PSI_API int read_options(const std::string &name, Options &options) {
    if (name == "LIBRESPONSE_PSI4" || options.read_globals()) {
        options.add_str("order", "linear");
        options.add_str("solver", "linear");
        options.add_str("hamiltonian", "rpa");
        options.add_str("spin", "singlet");
        options.add_int("maxiter", 60);
        // void add(std::string key, DataType* data);
        // OperatorDataType etc.
        options.add_array("OPERATOR_DIPOLE");
        options.add_array("OPERATOR_QUADRUPOLE");
        // options.add_array("OPERATOR_MULTIPOLE");
        options.add_array("OPERATOR_NABLA");
        options.add_array("OPERATOR_ANGMOM");
    }

    return true;
}

extern "C" PSI_API SharedWavefunction libresponse_psi4(SharedWavefunction ref_wfn,
                                                       Options &options) {
    // This is the Psi4 output file. Replace the cout stream buffer
    // with the one from Psi4.
    std::ostream *ofs = outfile->stream();
    std::streambuf *cout_original_buff = std::cout.rdbuf();
    std::cout.rdbuf(ofs->rdbuf());

    options.print();

    const size_t NOrb = ref_wfn->nmo();
    const size_t NOa = ref_wfn->nalpha();
    const size_t NOb = ref_wfn->nbeta();
    const size_t NVa = NOrb - NOa;
    const size_t NVb = NOrb - NOb;

    const size_t NDen = ref_wfn->same_a_b_dens() ? 1 : 2;

    const std::shared_ptr<BasisSet> bset = ref_wfn->basisset();
    const size_t NBasis = bset->nbf();

    arma::cube C(NBasis, NOrb, NDen);
    arma::mat moene(NOrb, NDen);

    // Psi4 is C++; storage is row major!
    const SharedMatrix mCa = ref_wfn->Ca()->transpose();
    const SharedMatrix mCb = ref_wfn->Cb()->transpose();
    const SharedVector vEa = ref_wfn->epsilon_a();
    const SharedVector vEb = ref_wfn->epsilon_b();
    const arma::mat Ca(mCa->get_pointer(), mCa->rowdim(), mCa->coldim(), true);
    const arma::vec Ea(vEa->pointer(), vEa->dim(), true);
    C.slice(0) = Ca;
    moene.col(0) = Ea;
    if (NDen == 2) {
        const arma::mat Cb(mCb->get_pointer(), mCb->rowdim(), mCb->coldim(), true);
        const arma::vec Eb(vEb->pointer(), vEb->dim(), true);
        C.slice(1) = Cb;
        moene.col(1) = Eb;
    }

    arma::uvec occupations(4);
    occupations(0) = NOa;
    occupations(1) = NVa;
    occupations(2) = NOb;
    occupations(3) = NVb;

    // Configuration.
    libresponse::configurable libresponse_options;
    set_defaults(libresponse_options);
    // TODO input option
    libresponse_options.cfg<int>("print_level", 2);
    // Psi4 must take C_left/C_right, not Dg!
    libresponse_options.cfg<bool>("_do_compute_generalized_density", false);

    // Frequencies.
    std::vector<double> omega;
    // Force the static case for now.
    omega.push_back(0.0);

    // 1-electron integrals for operators.
    std::vector<libresponse::operator_spec> operators;
    parse_operators(ref_wfn, options, operators);

    // 2-electron integral engine.
    MatVec_Psi4 *matvec = new MatVec_Psi4(ref_wfn, options);

    libresponse::SolverIterator_linear *solver_iterator = new libresponse::SolverIterator_linear();

    arma::cube results;

    libresponse::solve_linear_response(results,
                                       matvec,
                                       solver_iterator,
                                       C,
                                       moene,
                                       occupations,
                                       omega,
                                       operators,
                                       libresponse_options);

    delete matvec;
    delete solver_iterator;

    // Final printing.
    std::ostringstream os;
    os.precision(6);
    os.setf(std::ios_base::fixed, std::ios_base::floatfield);

    // os << " " << dashes << std::endl;
    // print_polarizability(os, results.slice(0));
    // os << " " << rcarats << std::endl;

    std::cout << os.str();

    // Put the original cout stream buffer back.
    std::cout.rdbuf(cout_original_buff);

    // Typically you would build a new wavefunction and populate it with data
    return ref_wfn;
}

} // namespace libresponse_psi4
} // namespace psi


#include <boost/make_shared.hpp>

#include <sm/eigen/gtest.hpp>
#include <aslam/backend/test/SampleDvAndError.hpp>

TEST(ErrorTermTestSuite, testMEstimatorGetter) {
  using aslam::backend::FixedWeightMEstimator;
  using aslam::backend::GemanMcClureMEstimator;

  Eigen::Vector2d v;
  Point2d p(v);
  LinearErr error_term(&p);
  boost::shared_ptr<FixedWeightMEstimator> mestimator(
      new FixedWeightMEstimator(4));
  error_term.setMEstimatorPolicy(mestimator);
  EXPECT_EQ(error_term.getCurrentMEstimatorWeight(), 4);
  error_term.getMEstimatorPolicy<FixedWeightMEstimator>()->setWeight(5);
  EXPECT_EQ(error_term.getCurrentMEstimatorWeight(), 5);
  boost::shared_ptr<GemanMcClureMEstimator> null_ptr =
      error_term.getMEstimatorPolicy<GemanMcClureMEstimator>();
  EXPECT_FALSE(null_ptr);
}

TEST(ErrorTermTestSuite, testInvR)
{
  using namespace aslam::backend;
  int D = 4;
  int E = 6;
  std::vector<DesignVariable*> dvs;
  std::vector<ErrorTerm*> errs;
  try {
    buildSystem(D, E, dvs, errs);
    for (size_t i = 0; i < errs.size(); ++i) {
      boost::shared_ptr<GemanMcClureMEstimator>
          me(new GemanMcClureMEstimator(0.321 * i)); // just some sigma2 value
      errs[i]->setMEstimatorPolicy(me);
    }
    for (size_t i = 0; i < errs.size(); ++i) {
      ErrorTerm* e = errs[i];
      Eigen::MatrixXd invR = sm::eigen::randomCovarianceXd(e->dimension());
      e->vsSetInvR(invR);
      e->evaluateError();
      Eigen::VectorXd ee = e->vsError();
      // Check that the raw squared error is e^T invR e
      double rse = e->getRawSquaredError();
      double trueRse = ee.dot(invR * ee);
      ASSERT_NEAR(rse, trueRse, 1e-6);
      // Check that the weighted squared error is w e^T invR e
      double w = e->getMEstimatorWeight(trueRse);
      double trueWRse = w * trueRse;
      double wrse = e->getWeightedSquaredError();
      ASSERT_NEAR(wrse, trueWRse, 1e-6);
      // Check that the weighted squared error, d,  produces  w e^T invR e == d^T d
      Eigen::VectorXd we;
      e->getWeightedError(we, false);
      // No M-estimator.
      rse = we.dot(we);
      ASSERT_NEAR(rse, trueRse, 1e-6);
      e->getWeightedError(we, true);
      // with M-estimator.
      wrse = we.dot(we);
      ASSERT_NEAR(wrse, trueWRse, 1e-6);
      JacobianContainerSparse<> jcRaw(e->dimension());
      e->evaluateJacobians(jcRaw);
      {
        // No M-Estimator
        JacobianContainerSparse<> jc(errs[i]->dimension());
        e->getWeightedJacobians(jc, false);
        e->getWeightedError(we, false);
        Eigen::MatrixXd J = jcRaw.asDenseMatrix();
        Eigen::MatrixXd wJ = jc.asDenseMatrix();
        Eigen::MatrixXd JtInvRJ = J.transpose() * invR * J;
        Eigen::MatrixXd wJtwJ = wJ.transpose() * wJ;
        // Check that wJ^T wJ == w J^T invR J
        ASSERT_DOUBLE_MX_EQ(JtInvRJ, wJtwJ, 1e-6, "Checking the weighted Jacobian");
        // Check that wJ^T we == w J^T invR e
        Eigen::VectorXd JtInvRe = J.transpose() * invR * ee;
        Eigen::VectorXd wJtwe = wJ.transpose() * we;
        ASSERT_DOUBLE_MX_EQ(JtInvRe, wJtwe, 1e-6, "Checking the weighted rhs");
      }
      {
        // with M-Estimator
        JacobianContainerSparse<> jc(errs[i]->dimension());
        e->getWeightedJacobians(jc, true);
        e->getWeightedError(we, true);
        w = e->getMEstimatorWeight(trueRse);
        ASSERT_NE(1.0, w);
        Eigen::MatrixXd J = jcRaw.asDenseMatrix();
        Eigen::MatrixXd wJ = jc.asDenseMatrix();
        Eigen::MatrixXd JtInvRJ =  J.transpose() * invR * J * w;
        Eigen::MatrixXd wJtwJ = wJ.transpose() * wJ;
        // Check that wJ^T wJ == w J^T invR J
        ASSERT_DOUBLE_MX_EQ(JtInvRJ, wJtwJ, 1e-6, "Checking the weighted Jacobian");
        // Check that wJ^T we == w J^T invR e
        Eigen::VectorXd JtInvRe = J.transpose() * invR * ee * w;
        Eigen::VectorXd wJtwe = wJ.transpose() * we;
        ASSERT_DOUBLE_MX_EQ(JtInvRe, wJtwe, 1e-6, "Checking the weighted rhs");
      }
    }
    deleteSystem(dvs, errs);
  } catch (const std::exception& e) {
    deleteSystem(dvs, errs);
    FAIL() << e.what();
  }
}

TEST(ErrorTermTestSuite, testNonSquaredErrorTerm) {
  using namespace aslam::backend;
  try {

    Eigen::Vector2d v(1.0, 2.0);
    Point2d p(v);
    p.setBlockIndex(0);
    TestNonSquaredError::grad_t g(1.0, 2.0);
    TestNonSquaredError e(&p, g);

    e._p = 3.0;
    for (double w : {0.1, 1.0, 10.}){
      e.setWeight(w);
      for (double mEstWeight : {0.1, 1.0, 10.}){
        if(mEstWeight != 1.0){
          e.setMEstimatorPolicy(boost::make_shared<FixedWeightMEstimator>(mEstWeight));
        }
        else {
          e.clearMEstimatorPolicy();
        }

        EXPECT_DOUBLE_EQ(w * 4.0, e.updateRawError());
        EXPECT_DOUBLE_EQ(w * mEstWeight * 4.0, e.getWeightedError());
        EXPECT_DOUBLE_EQ(e.getRawError(), e.getError(false));
        EXPECT_DOUBLE_EQ(e.getWeightedError(), e.getError(true));

        Eigen::MatrixXd grad;
        grad.resize(1,2); grad << 4., 8.;
        grad *= w;

        JacobianContainerSparse<> jc(1);

        // raw Jacobian
        jc.clear();
        e.evaluateJacobians(jc, false);
        Eigen::MatrixXd J = jc.asDenseMatrix();
        EXPECT_TRUE((J.array() == grad.array()).all()) << "J: " << J << std::endl <<
            "Grad: " << grad << std::endl;

        jc.clear();
        e.evaluateRawJacobians(jc);
        J = jc.asDenseMatrix();
        EXPECT_TRUE((J.array() == grad.array()).all()) << "J: " << J << std::endl <<
            "Grad: " << grad << std::endl;

        // weighted Jacobian
        jc.clear();
        e.evaluateJacobians(jc, true);
        J = jc.asDenseMatrix();
        grad *= mEstWeight;
        EXPECT_TRUE((J.array() == grad.array()).all()) << "J: " << J << std::endl <<
            "Grad: " << grad << std::endl;

        jc.clear();
        e.evaluateWeightedJacobians(jc);
        J = jc.asDenseMatrix();
        EXPECT_TRUE((J.array() == grad.array()).all()) << "J: " << J << std::endl <<
            "Grad: " << grad << std::endl;

      }
    }
  } catch (const std::exception& e) {
    FAIL() << e.what();
  }
}


/**
 * \description This is the main for the new version of the mert algorithm developed during the 2nd MT marathon
*/

#include <limits>
#include <unistd.h>
#include <cstdlib>
#include <iostream>
#include <fstream>
#include <cmath>
#include <ctime>

#include <getopt.h>
#include <boost/scoped_ptr.hpp>

#include "Data.h"
#include "Point.h"
#include "Scorer.h"
#include "ScorerFactory.h"
#include "ScoreData.h"
#include "FeatureData.h"
#include "Optimizer.h"
#include "OptimizerFactory.h"
#include "Types.h"
#include "Timer.h"
#include "Util.h"

#include "moses/ThreadPool.h"

using namespace std;
using namespace MosesTuning;

namespace {

const char kDefaultOptimizer[] = "powell";
const char kDefaultScorer[] = "BLEU";
const char kDefaultScorerFile[] = "statscore.data";
const char kDefaultFeatureFile[] = "features.data";
const char kDefaultInitFile[] = "init.opt";
const char kDefaultPositiveString[] = "";
const char kDefaultSparseWeightsFile[] = "";

// Used when saving optimized weights.
const char kOutputFile[] = "weights.txt";

/**
 * Runs an optimisation, or a random restart.
 */
class OptimizationTask : public Moses::Task {
 public:
  OptimizationTask(Optimizer* optimizer, const Point& point)
      : m_optimizer(optimizer), m_point(point) {}

  ~OptimizationTask() {}

  virtual void Run() {
    m_score = m_optimizer->Run(m_point);
  }

  virtual bool DeleteAfterExecution() {
    return false;
  }

  void resetOptimizer() {
    if (m_optimizer) {
      delete m_optimizer;
      m_optimizer = NULL;
    }
  }

  statscore_t getScore() const {
    return m_score;
  }

  const Point& getPoint() const  {
    return m_point;
  }

 private:
  // Do not allow the user to instanciate without arguments.
  OptimizationTask() {}

  Optimizer* m_optimizer;
  Point m_point;
  statscore_t m_score;
};

bool WriteFinalWeights(const char* filename, const Point& point) {
  ofstream ofs(filename);
  if (!ofs) {
    cerr << "Cannot open " << filename << endl;
    return false;
  }

  ofs << point << endl;

  return true;
}

void usage(int ret)
{
  cerr<<"usage: mert -d <dimensions> (mandatory )"<<endl;
  cerr<<"[-n] retry ntimes (default 1)"<<endl;
  cerr<<"[-m] number of random directions in powell (default 0)"<<endl;
  cerr<<"[-o] the indexes to optimize(default all)"<<endl;
  cerr<<"[-t] the optimizer(default powell)"<<endl;
  cerr<<"[-r] the random seed (defaults to system clock)"<<endl;
  cerr<<"[--sctype|-s] the scorer type (default BLEU)"<<endl;
  cerr<<"[--scconfig|-c] configuration string passed to scorer"<<endl;
  cerr<<"[--scfile|-S] comma separated list of scorer data files (default score.data)"<<endl;
  cerr<<"[--ffile|-F] comma separated list of feature data files (default feature.data)"<<endl;
  cerr<<"[--ifile|-i] the starting point data file (default init.opt)"<<endl;
  cerr<<"[--sparse-weights|-p] required for merging sparse features"<<endl;
#ifdef WITH_THREADS
  cerr<<"[--threads|-T] use multiple threads (default 1)"<<endl;
#endif
  cerr<<"[--shard-count] Split data into shards, optimize for each shard and average"<<endl;
  cerr<<"[--shard-size] Shard size as proportion of data. If 0, use non-overlapping shards"<<endl;
  cerr<<"[-v] verbose level"<<endl;
  cerr<<"[--help|-h] print this message and exit"<<endl;
  exit(ret);
}

static struct option long_options[] = {
  {"pdim", 1, 0, 'd'},
  {"ntry",1,0,'n'},
  {"nrandom",1,0,'m'},
  {"rseed",required_argument,0,'r'},
  {"optimize",1,0,'o'},
  {"type",1,0,'t'},
  {"sctype",1,0,'s'},
  {"scconfig",required_argument,0,'c'},
  {"scfile",1,0,'S'},
  {"ffile",1,0,'F'},
  {"ifile",1,0,'i'},
  {"sparse-weights",required_argument,0,'p'},
#ifdef WITH_THREADS
  {"threads", required_argument,0,'T'},
#endif
  {"shard-count", required_argument, 0, 'a'},
  {"shard-size", required_argument, 0, 'b'},
  {"verbose",1,0,'v'},
  {"help",no_argument,0,'h'},
  {0, 0, 0, 0}
};

struct ProgramOption {
  string to_optimize_str;
  int pdim;
  int ntry;
  int nrandom;
  int seed;
  bool has_seed;
  string optimize_type;
  string scorer_type;
  string scorer_config;
  string scorer_file;
  string feature_file;
  string init_file;
  string positive_string;
  string sparse_weights_file;
  size_t num_threads;
  float shard_size;
  size_t shard_count;

  ProgramOption()
      : to_optimize_str(""),
        pdim(-1),
        ntry(1),
        nrandom(0),
        seed(0),
        has_seed(false),
        optimize_type(kDefaultOptimizer),
        scorer_type(kDefaultScorer),
        scorer_config(""),
        scorer_file(kDefaultScorerFile),
        feature_file(kDefaultFeatureFile),
        init_file(kDefaultInitFile),
        positive_string(kDefaultPositiveString),
        sparse_weights_file(kDefaultSparseWeightsFile),
        num_threads(1),
        shard_size(0),
        shard_count(0) { }
};

void ParseCommandOptions(int argc, char** argv, ProgramOption* opt) {
  int c;
  int option_index;

  while ((c = getopt_long(argc, argv, "o:r:d:n:m:t:s:S:F:v:p:P:", long_options, &option_index)) != -1) {
    switch (c) {
      case 'o':
        opt->to_optimize_str = string(optarg);
        break;
      case 'd':
        opt->pdim = strtol(optarg, NULL, 10);
        break;
      case 'n':
        opt->ntry = strtol(optarg, NULL, 10);
        break;
      case 'm':
        opt->nrandom = strtol(optarg, NULL, 10);
        break;
      case 'r':
        opt->seed = strtol(optarg, NULL, 10);
        opt->has_seed = true;
        break;
      case 't':
        opt->optimize_type = string(optarg);
        break;
      case's':
        opt->scorer_type = string(optarg);
        break;
      case 'c':
        opt->scorer_config = string(optarg);
        break;
      case 'S':
        opt->scorer_file = string(optarg);
        break;
      case 'F':
        opt->feature_file = string(optarg);
        break;
      case 'i':
        opt->init_file = string(optarg);
        break;
      case 'p':
        opt->sparse_weights_file=string(optarg);
        break;
      case 'v':
        setverboselevel(strtol(optarg, NULL, 10));
        break;
#ifdef WITH_THREADS
      case 'T':
        opt->num_threads = strtol(optarg, NULL, 10);
        if (opt->num_threads < 1) opt->num_threads = 1;
        break;
#endif
      case 'a':
        opt->shard_count = strtof(optarg, NULL);
        break;
      case 'b':
        opt->shard_size = strtof(optarg, NULL);
        break;
      case 'h':
        usage(0);
        break;
      case 'P':
        opt->positive_string = string(optarg);
        break;
      default:
        usage(1);
    }
  }
}

} // anonymous namespace

int main(int argc, char **argv)
{
  ResetUserTime();

  ProgramOption option;
  ParseCommandOptions(argc, argv, &option);

  vector<unsigned> to_optimize;
  vector<vector<parameter_t> > start_list;
  vector<parameter_t> min;
  vector<parameter_t> max;
  vector<bool> positive;
  // NOTE: those mins and max are the bound for the starting points of the algorithm, not strict bound on the result!

  if (option.pdim < 0)
    usage(1);

  cerr << "shard_size = " << option.shard_size << " shard_count = " << option.shard_count << endl;
  if (option.shard_size && !option.shard_count) {
    cerr << "Error: shard-size provided without shard-count" << endl;
    exit(1);
  }
  if (option.shard_size > 1 || option.shard_size < 0) {
    cerr << "Error: shard-size should be between 0 and 1" << endl;
    exit(1);
  }

  if (option.has_seed) {
    cerr << "Seeding random numbers with " << option.seed << endl;
    srandom(option.seed);
  } else {
    cerr << "Seeding random numbers with system clock " << endl;
    srandom(time(NULL));
  }

  if (option.sparse_weights_file.size()) ++option.pdim;

  // read in starting points
  string onefile;
  while (!option.init_file.empty()) {
    getNextPound(option.init_file, onefile, ",");
    vector<parameter_t> start;
    ifstream opt(onefile.c_str());
    if (opt.fail()) {
      cerr << "could not open initfile: " << option.init_file << endl;
      exit(3);
    }
    start.resize(option.pdim);//to do:read from file
    int j;
    for (j = 0; j < option.pdim && !opt.fail(); j++) {
      opt >> start[j];
    }
    if (j < option.pdim) {
      cerr << option.init_file << ":Too few starting weights." << endl;
      exit(3);
    }
    start_list.push_back(start);
    // for the first time, also read in the min/max values for scores
    if (start_list.size() == 1) {
      min.resize(option.pdim);
      for (j = 0; j < option.pdim && !opt.fail(); j++) {
        opt >> min[j];
      }
      if (j < option.pdim) {
        cerr << option.init_file << ":Too few minimum weights." << endl;
        cerr << "error could not initialize start point with " << option.init_file << endl;
        cerr << "j: " << j << ", pdim: " << option.pdim << endl;
        exit(3);
      }
      max.resize(option.pdim);
      for (j = 0; j < option.pdim && !opt.fail(); j++) {
        opt >> max[j];
      }
      if (j < option.pdim) {
        cerr << option.init_file << ":Too few maximum weights." << endl;
        exit(3);
      }
    }
    opt.close();
  }

  vector<string> ScoreDataFiles;
  if (option.scorer_file.length() > 0) {
    Tokenize(option.scorer_file.c_str(), ',', &ScoreDataFiles);
  }

  vector<string> FeatureDataFiles;
  if (option.feature_file.length() > 0) {
    Tokenize(option.feature_file.c_str(), ',', &FeatureDataFiles);
  }

  if (ScoreDataFiles.size() != FeatureDataFiles.size()) {
    throw runtime_error("Error: there is a different number of previous score and feature files");
  }

  // it make sense to know what parameter set were used to generate the nbest
  boost::scoped_ptr<Scorer> scorer(
      ScorerFactory::getScorer(option.scorer_type, option.scorer_config));

  //load data
  Data data(scorer.get(), option.sparse_weights_file);

  for (size_t i = 0; i < ScoreDataFiles.size(); i++) {
    cerr<<"Loading Data from: "<< ScoreDataFiles.at(i) << " and " << FeatureDataFiles.at(i) << endl;
    data.load(FeatureDataFiles.at(i), ScoreDataFiles.at(i));
  }

  scorer->setScoreData(data.getScoreData().get());

  data.removeDuplicates();

  PrintUserTime("Data loaded");

  // starting point score over latest n-best, accumulative n-best
  //vector<unsigned> bests;
  //compute bests with sparse features needs to be implemented
  //currently sparse weights are not even loaded
  //statscore_t score = TheScorer->score(bests);

  if (option.to_optimize_str.length() > 0) {
    cerr << "Weights to optimize: " << option.to_optimize_str << endl;

    // Parse the string to get weights to optimize, and set them as active.
    vector<string> features;
    Tokenize(option.to_optimize_str.c_str(), ',', &features);

    for (vector<string>::const_iterator it = features.begin();
         it != features.end(); ++it) {
      const int feature_index = data.getFeatureIndex(*it);

      // Note: previous implementaion checked whether
      // feature_index is less than option.pdim.
      // However, it does not make sense when we optimize 'discrete' features,
      // given by '-o' option like -o "d_0,lm_0,tm_2,tm_3,tm_4,w_0".
      if (feature_index < 0) {
        cerr << "Error: invalid feature index = " << feature_index << endl;
        exit(1);
      }
      cerr << "FeatNameIndex: " << feature_index << " to insert" << endl;
      to_optimize.push_back(feature_index);
    }
  } else {
    //set all weights as active
    to_optimize.resize(option.pdim);//We'll optimize on everything
    for (int i = 0; i < option.pdim; i++) {
      to_optimize[i] = 1;
    }
  }

  positive.resize(option.pdim);
  for (int i = 0; i < option.pdim; i++)
    positive[i] = false;
  if (option.positive_string.length() > 0) {
    // Parse string to get weights that need to be positive
    std::string substring;
    int index;
    while (!option.positive_string.empty()) {
      getNextPound(option.positive_string, substring, ",");
      index = data.getFeatureIndex(substring);
      //index = strtol(substring.c_str(), NULL, 10);
      if (index >= 0 && index < option.pdim) {
        positive[index] = true;
      } else {
        cerr << "Index " << index
             << " is out of bounds in positivity list. Allowed indexes are [0,"
             << (option.pdim-1) << "]." << endl;
      }
    }
  }

#ifdef WITH_THREADS
  cerr << "Creating a pool of " << option.num_threads << " threads" << endl;
  Moses::ThreadPool pool(option.num_threads);
#endif

  Point::setpdim(option.pdim);
  Point::setdim(to_optimize.size());
  Point::set_optindices(to_optimize);

  //starting points consist of specified points and random restarts
  vector<Point> startingPoints;

  for (size_t i = 0; i < start_list.size(); ++i) {
    startingPoints.push_back(Point(start_list[i], min, max));
  }

  for (int i = 0; i < option.ntry; ++i) {
    startingPoints.push_back(Point(start_list[0], min, max));
    startingPoints.back().Randomize();
  }

  vector<vector<OptimizationTask*> > allTasks(1);

  //optional sharding
  vector<Data> shards;
  if (option.shard_count) {
    data.createShards(option.shard_count, option.shard_size, option.scorer_config, shards);
    allTasks.resize(option.shard_count);
  }

  // launch tasks
  for (size_t i = 0; i < allTasks.size(); ++i) {
    Data& data_ref = data;
    if (option.shard_count)
      data_ref = shards[i]; //use the sharded data if it exists

    vector<OptimizationTask*>& tasks = allTasks[i];
    Optimizer *optimizer = OptimizerFactory::BuildOptimizer(option.pdim, to_optimize, positive, start_list[0], option.optimize_type, option.nrandom);
    optimizer->SetScorer(data_ref.getScorer());
    optimizer->SetFeatureData(data_ref.getFeatureData());
    // A task for each start point
    for (size_t j = 0; j < startingPoints.size(); ++j) {
      OptimizationTask* task = new OptimizationTask(optimizer, startingPoints[j]);
      tasks.push_back(task);
#ifdef WITH_THREADS
      pool.Submit(task);
#else
      task->Run();
#endif
    }
  }

  // wait for all threads to finish
#ifdef WITH_THREADS
  pool.Stop(true);
#endif

  statscore_t total = 0;
  Point totalP;

  // collect results
  for (size_t i = 0; i < allTasks.size(); ++i) {
    statscore_t best = 0, mean = 0, var = 0;
    Point bestP;
    for (size_t j = 0; j < allTasks[i].size(); ++j) {
      statscore_t score = allTasks[i][j]->getScore();
      mean += score;
      var += score * score;
      if (score > best) {
        bestP = allTasks[i][j]->getPoint();
        best = score;
      }
    }

    mean /= static_cast<float>(option.ntry);
    var /= static_cast<float>(option.ntry);
    var = sqrt(abs(var - mean * mean));

    if (verboselevel() > 1) {
      cerr << "shard " << i << " best score: " << best << " variance of the score (for " << option.ntry << " try): " << var << endl;
    }

    totalP += bestP;
    total += best;
    if (verboselevel() > 1)
      cerr << "bestP " << bestP << endl;
  }

  //cerr << "totalP: " << totalP << endl;
  Point finalP = totalP * (1.0 / allTasks.size());
  statscore_t final = total / allTasks.size();

  if (verboselevel() > 1)
    cerr << "bestP: " << finalP << endl;

  // L1-Normalization of the best Point
  if (static_cast<int>(to_optimize.size()) == option.pdim) {
    finalP.NormalizeL1();
  }

  cerr << "Best point: " << finalP << " => " << final << endl;

  if (!WriteFinalWeights(kOutputFile, finalP)) {
    cerr << "Warning: Failed to write the final point" << endl;
  }

  for (size_t i = 0; i < allTasks.size(); ++i) {
    allTasks[i][0]->resetOptimizer();
    for (size_t j = 0; j < allTasks[i].size(); ++j) {
      delete allTasks[i][j];
    }
  }

  PrintUserTime("Stopping...");

  return 0;
}


#include <boost/test/unit_test.hpp>

#include <giecs/bits.hpp>
#include <giecs/memory/context.hpp>
#include <giecs/memory/accessor.hpp>
#include <giecs/memory/accessors/linear.hpp>

BOOST_AUTO_TEST_SUITE(context);

using namespace giecs;

BOOST_AUTO_TEST_CASE(accessor)
{
    typedef Bits<6> byte;
    typedef Bits<24> word;

    auto c1 = memory::Context<4, word>();

    struct Index
    {
        int x,y;
        Index()
            :x(0),y(0)
        {}

        Index(int x)
            : Index(x, 0)
        {
        }

        Index(int x_, int y_)
            : x(x_), y(y_)
        {
            this->y += this->x / 4;
            this->x %= 4;
        }

        Index operator+(Index const& i) const
        {
            return Index(this->x+i.x, this->y+i.y);
        }

        bool operator<(Index const& i) const
        {
            return ((int)*this  < (int)i);
        }

        Index& operator++()
        {
            *this = Index(this->x+1, this->y);
            return *this;
        }

        operator int() const
        {
            return (this->x + this->y*4);
        }
    };

    struct Array
    {
        byte* ptr;

        Array(void const* ptr_)
            :ptr((byte*)ptr_)
        {}

        byte& operator[] (Index const& i) const
        {
            return this->ptr[(int)i];
        }

        byte& operator* () const
        {
            return *this->ptr;
        }

        Array& operator++()
        {
            this->ptr++;
            return *this;
        }
    };

    auto acc = memory::accessors::Linear<4, word, Index, byte, Array, Index>(c1);

    Index len = Index(0,4);

    int a = 0;
    for(Index i = Index(); i < len; ++i)
    {
        ++a;
        acc[i] = a;
    }

    a = 0;
    for(Index i = Index(); i < len; ++i)
    {
        ++a;
        byte b = acc[i];
        BOOST_CHECK( b == a );
    }
}

BOOST_AUTO_TEST_CASE(synchronization)
{
    auto c1 = memory::Context<8, uint8_t>();

    auto acc1 = memory::accessors::Linear<8, uint8_t, int, uint8_t>(c1);
    auto acc2 = memory::accessors::Linear<8, uint8_t, int, uint16_t>(c1);

    static int const n = 4;

    uint8_t v8 = 0;
    for(int i = 0; i < 2*n; i++)
    {
        ++v8;
        acc1[i] = v8;
    }

    v8 = 0;
    for(int i = 0; i < n; i++)
    {
        uint16_t v16 = 0;

        ++v8;
        uint8_t r8 = acc1[i*2];
        BOOST_CHECK(r8 == v8);
        v16 = v8;

        printf("%x ", v8);

        ++v8;
        r8 = acc1[i*2+1];
        BOOST_CHECK(r8 == v8);
        v16 |= v8 << 8;

        uint16_t r16 = acc2[i];
        printf("%x -> %x == %x\n", v8, v16, r16);
        BOOST_CHECK(r16 == v16);
    }
}

BOOST_AUTO_TEST_SUITE_END();


//==================================================================================================
/*!
  @file

  @copyright 2015 NumScale SAS
  @copyright 2015 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_GENERIC_FUNCTION_FREXP_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_GENERIC_FUNCTION_FREXP_HPP_INCLUDED

#include <boost/simd/meta/cardinal_of.hpp>
#include <boost/dispatch/function/overload.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#include <boost/config.hpp>
#include <boost/fusion/include/std_pair.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;
  BOOST_DISPATCH_OVERLOAD ( frexp_
                          , (typename A0, typename A2)
                          , bd::cpu_
                          , bd::generic_< bd::floating_<A0> >
                          , bd::generic_< bd::integer_<A2> >
                            )
  {
    static_assert
      ( bs::cardinal_of<A0>::value == bs::cardinal_of<A2>::value
      , "boost.simd cardinalities are inconsistent in frexp call"
      );

    BOOST_FORCEINLINE A0 operator() ( A0 const& a0,A2 & a2) const BOOST_NOEXCEPT
    {
      A0 a1;
      bs::frexp(a0, a1, a2);
      return a1;
    }
  };

  BOOST_DISPATCH_OVERLOAD ( frexp_
                          , (typename A0)
                          , bd::cpu_
                          , bd::generic_<bd::floating_<A0> >
                          )
  {
    using exponent_t = bd::as_integer_t<A0, signed>;

    BOOST_FORCEINLINE std::pair<A0,exponent_t>  operator() ( A0 const& a0) const BOOST_NOEXCEPT
    {
      A0 first;
      exponent_t second;
      bs::frexp( a0, first, second );
      return {first, second};
    }
  };

  BOOST_DISPATCH_OVERLOAD ( frexp_
                          , (typename A0, typename A2)
                          , bd::cpu_
                          , bd::generic_< bd::floating_<A0> >
                          , bd::generic_< bd::integer_<A2> >
                          , boost::simd::fast_tag
                            )
  {
    static_assert
      ( bs::cardinal_of<A0>::value == bs::cardinal_of<A2>::value
      , "boost.simd cardinalities are inconsistent in frexp call"
      );

    BOOST_FORCEINLINE A0 operator() ( A0 const& a0,A2 & a2
                                    , bs::fast_tag const& f) const BOOST_NOEXCEPT
    {
      A0 a1;
      bs::frexp(a0, a1, a2, f);
      return a1;
    }
  };

  BOOST_DISPATCH_OVERLOAD ( frexp_
                          , (typename A0)
                          , bd::cpu_
                          , bd::generic_<bd::floating_<A0> >
                          , boost::simd::fast_tag
                          )
  {
    using exponent_t = bd::as_integer_t<A0, signed>;

    BOOST_FORCEINLINE std::pair<A0,exponent_t>  operator() ( A0 const& a0
                                                           , bs::fast_tag const & f    ) const BOOST_NOEXCEPT
    {
      A0 first;
      exponent_t second;
      bs::frexp( a0, first, second ,f);
      return {first, second};
    }
  };
} } }

#endif