// Definitions of classes that implements useful
// concepts for the eddy current project.
//
// EddyHelperClasses.cpp
//
// Jesper Andersson, FMRIB Image Analysis Group
//
// Copyright (C) 2011 University of Oxford
//

#include <cstdlib>
#include <string>
#include <sstream>
#include <vector>
#include <cmath>
#include <algorithm>
#include "nlohmann/json.hpp"
#include "armawrap/newmat.h"
#include "newimage/newimageall.h"
#include "miscmaths/miscmaths.h"
#include "EddyHelperClasses.h"
#include "EddyUtils.h"

using namespace std;
using namespace EDDY;



//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class JsonReader
//
// This class manages reading .json files.
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

JsonReader::~JsonReader() { delete _content; }

NEWMAT::Matrix JsonReader::SliceOrdering() const EddyTry
{
  // Check for existence of "SliceTiming" field
  auto SOObj = _content->find("SliceTiming");
  if (SOObj == _content->end()) { // Field not found
    throw EddyException("JsonReader::SliceOrdering: Failed to find object \"SliceTiming\"");
  }
  std::vector<double> slice_times = (*_content)["SliceTiming"];  

  // Sort slice indicies in order of increasing acquisition time
  std::vector<std::pair<unsigned int,double> > indextime(slice_times.size());
  for (unsigned int i=0; i<slice_times.size(); i++) { indextime[i] = std::make_pair(i,slice_times[i]); }
  std::sort(indextime.begin(),indextime.end(),[](std::pair<unsigned int,double> p1, std::pair<unsigned int,double> p2)
	    { return((p1.second == p2.second) ? (p1.first < p2.first) : (p1.second < p2.second)); });

  // Do a little sanity check
  unsigned int mbf = 1;
  for (unsigned int i=1; i<indextime.size(); i++) {
    if (indextime[i].second > indextime[i-1].second) break;
    else mbf++;
  }
  double ctime = indextime[0].second;
  unsigned int cnt = 1;
  for (unsigned int i=1; i<indextime.size(); i++) {
    if (indextime[i].second == ctime) cnt++;
    else {
      if (cnt != mbf) throw EddyException("JsonReader::SliceOrdering: Inconsistent MB groups");
      else {
	ctime = indextime[i].second;
	cnt = 1;
      }
    }
  }
  // Do extra check on MBF if info is available
  auto MBFObj = _content->find("MultibandAccelerationFactor");
  if (MBFObj != _content->end()) {
    int MBFInt = (*_content)["MultibandAccelerationFactor"];
    if (MBFInt != mbf) throw("JsonReader::SliceOrdering: Mismatch between \"SliceTiming\" and \"MultibandAccelerationFactor\"");
  }
  // Finally, repack it in slspec format
  NEWMAT::Matrix M(indextime.size()/mbf,mbf);
  unsigned int ii=0;
  for (int ri=0; ri<M.Nrows(); ri++) {
    for (int ci=0; ci<M.Ncols(); ci++) M(ri+1,ci+1) = std::round(static_cast<double>(indextime[ii++].first));
  }

  return(M);
} EddyCatch

NEWMAT::ColumnVector JsonReader::PEVector() const EddyTry
{
  NEWMAT::ColumnVector pe(3); pe = 0;
  // Check for existence of "PhaseEncodingDirection" field
  auto PEObj = _content->find("PhaseEncodingDirection");
  if (PEObj == _content->end()) { // Field not found
    throw EddyException("JsonReader::PEVector: Failed to find object \"PhaseEncodingDirection\"");
  }
  std::string PEString = (*_content)["PhaseEncodingDirection"];
  if (PEString == std::string("i")) pe(1) = 1;
  else if (PEString == std::string("i-")) pe(1) = -1;
  else if (PEString == std::string("j")) pe(2) = 1;
  else if (PEString == std::string("j-")) pe(2) = -1;
  else throw EddyException("JsonReader::PEVector: Failed to decode \"PhaseEncodingDirection\"");

  return(pe);
} EddyCatch

double JsonReader::TotalReadoutTime() const EddyTry
{
  auto TRTObj = _content->find("TotalReadoutTime");
  if (TRTObj == _content->end()) { // Field not found
    throw EddyException("JsonReader::TotalReadoutTime: Failed to find object \"TotalReadoutTime\"");
  }

  return((*_content)["TotalReadoutTime"]);
} EddyCatch

double JsonReader::EchoTime() const EddyTry
{
  auto ETObj = _content->find("EchoTime");
  if (ETObj == _content->end()) { // Field not found
    throw EddyException("JsonReader::EchoTime: Failed to find object \"EchoTime\"");
  }

  return((*_content)["EchoTime"]);
} EddyCatch

double JsonReader::RepetitionTime() const EddyTry
{
  auto RTObj = _content->find("RepetitionTime");
  if (RTObj == _content->end()) { // Field not found
    throw EddyException("JsonReader::RepetitionTime: Failed to find object \"RepetitionTime\"");
  }

  return((*_content)["RepetitionTime"]);
} EddyCatch

void JsonReader::common_read() EddyTry
{
  try {
    std::ifstream ifs(_fname);
    // nlohmann::json tmp = nlohmann::json::parse(ifs);
    _content = new nlohmann::json;
    (*_content) = nlohmann::json::parse(ifs);
    ifs.close();
  }
  catch (std::exception& e) {
    throw EddyException("JsonReader::common_read(): exception thrown with e.what() = " + std::string(e.what()));
  }
} EddyCatch

//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class DiffPara (Diffusion Parameters)
//
// This class manages the diffusion parameters for a given
// scan.
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

bool DiffPara::operator==(const DiffPara& rhs) const EddyTry
{
  return(EddyUtils::AreInSameShell(*this,rhs) && EddyUtils::HaveSameDirection(*this,rhs));
} EddyCatch

//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class AcqPara (Acquisition Parameters)
//
// This class manages the acquisition parameters for a given
// scan.
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

AcqPara::AcqPara(const NEWMAT::ColumnVector&   pevec,
                 double                        rotime) EddyTry
: _pevec(pevec), _rotime(rotime)
{
  if (pevec.Nrows() != 3) throw EddyException("AcqPara::AcqPara: Wrong dimension pe-vector");
  if (rotime < 0.01 || rotime > 0.2) throw EddyException("AcqPara::AcqPara: Unrealistic read-out time");
  int cc = 0; //Component count
  for (int i=0; i<3; i++) { if (fabs(pevec(i+1)) > 1e-6) cc++; }
  if (!cc) throw EddyException("AcqPara::AcqPara: Zero Phase-encode vector");
  if (cc > 1) throw EddyException("AcqPara::AcqPara: Oblique pe-vectors not yet implemented");
} EddyCatch

bool AcqPara::operator==(const AcqPara& rh) const EddyTry
{
  if (fabs(this->_rotime-rh._rotime) > 1e-6) return(false);
  for (int i=0; i<3; i++) {
    if (fabs(this->_pevec(i+1)-rh._pevec(i+1)) > 1e-6) return(false);
  }
  return(true);
} EddyCatch

std::vector<unsigned int> AcqPara::BinarisedPhaseEncodeVector() const EddyTry
{
  std::vector<unsigned int> rval(3,0);
  for (unsigned int i=0; i<3; i++) rval[i] = (_pevec(i+1) == 0) ? 0 : 1;
  return(rval);
} EddyCatch

DiffStats::DiffStats(const NEWIMAGE::volume<float>& diff, const NEWIMAGE::volume<float>& mask) EddyTry
: _md(diff.zsize(),0.0), _msd(diff.zsize(),0.0), _n(mask.zsize(),0)
{
#pragma omp parallel for // shared(diff,mask,_md,_msd,_n) private(j,k)
  for (int k=0; k<diff.zsize(); k++) {
    for (int j=0; j<diff.ysize(); j++) {
      for (int i=0; i<diff.xsize(); i++) {
	if (mask(i,j,k)) {
          _md[k] += diff(i,j,k);
          _msd[k] += diff(i,j,k)*diff(i,j,k);
	  _n[k] += 1;
	}
      }
    }
    if (_n[k]) { _md[k] /= double(_n[k]); _msd[k] /= double(_n[k]); }
  }
} EddyCatch

MultiBandGroups::MultiBandGroups(unsigned int nsl,
				 unsigned int mb,
				 int          offs) EddyTry
: _nsl(nsl), _mb(mb), _offs(offs)
{
  if (std::abs(_offs) > 1) throw EddyException("MultiBandGroups::MultiBandGroups: offs out of range");
  if (((int(_nsl)+std::abs(_offs)) % int(_mb)) || (_mb==1 && _offs!=0)) throw EddyException("MultiBandGroups::MultiBandGroups: Incompatible nsl, mb and offs");
  unsigned int ngrp = (_nsl+static_cast<unsigned int>(std::abs(_offs))) / _mb;
  _grps.resize(ngrp);
  for (unsigned int grp=0; grp<ngrp; grp++) {
    int sindx = (offs == 1) ? -1 : 0;
    for (int i=sindx+grp; i<int(_nsl); i+=ngrp) {
      if (i >= 0 && i < int(_nsl)) _grps[grp].push_back(static_cast<unsigned int>(i));
    }
  }
  _to.resize(ngrp); for (unsigned int grp=0; grp<ngrp; grp++) _to[grp]=grp; // Set temporal order to slice order
} EddyCatch

MultiBandGroups::MultiBandGroups(const std::string& fname) EddyTry
{
  // Read text file specifying what slices acquired when
  std::string line;
  std::ifstream ifs(fname.c_str());
  if (ifs.is_open()) {
    while (std::getline(ifs,line)) {
      std::vector<unsigned int> tmp_vec; // Empty vector of slice numbers
      unsigned int tmp_ui;
      std::stringstream ss(line);
      while (ss >> tmp_ui) tmp_vec.push_back(tmp_ui);
      _grps.push_back(tmp_vec); // Add vector of slice numbers to the end of vector of groups
    }
    if (ifs.eof()) ifs.close();
    else throw EddyException("MultiBandGroups::MultiBandGroups: Problem reading file");
  }
  else throw EddyException("MultiBandGroups::MultiBandGroups: Unable to open file");
  assert_grps(); // Check that _grps is kosher
} EddyCatch

MultiBandGroups::MultiBandGroups(const NEWMAT::Matrix& slices) EddyTry
{
  _grps.resize(slices.Nrows());
  for (unsigned int i=0; i<_grps.size(); i++) {
    _grps[i].resize(slices.Ncols());
    for (unsigned int j=0; j<_grps[i].size(); j++) _grps[i][j] = static_cast<unsigned int>(std::round(slices(i+1,j+1)));
  }
  assert_grps(); // Check that _grps is kosher
} EddyCatch

void MultiBandGroups::assert_grps() EddyTry
{
  // First deduce number of slices and MB factor.
  _nsl = 0;
  _mb = 0;
  for (unsigned int grp=0; grp<_grps.size(); grp++) {
    _mb = std::max(static_cast<unsigned int>(_grps.size()),_mb);
    for (unsigned int sl=0; sl<_grps[grp].size(); sl++) {
      _nsl = std::max(_grps[grp][sl],_nsl);
    }
  }
  _nsl++;
  _offs = 0; // Arbitrary, not used.
  // Next make sure that each slices is specified once and only once.
  std::vector<unsigned int> check_vec(_nsl,0);
  for (unsigned int grp=0; grp<_grps.size(); grp++) {
    for (unsigned int sl=0; sl<_grps[grp].size(); sl++) {
      check_vec[_grps[grp][sl]] += 1;
    }
  }
  for (unsigned int i=0; i<check_vec.size(); i++) {
    if (check_vec[i] != 1) throw EddyException("MultiBandGroups::MultiBandGroups: Logical error in file");
  }
  // Set time order to the same as storage order
  _to.resize(_grps.size());
  for (unsigned int i=0; i<_grps.size(); i++) _to[i] = i;
} EddyCatch

/*!
 * Will write three files containing the mean differences, the
 * mean squared differences and the number of valid voxels. The
 * organisation of the (text) files is such that the nth column
 * of the mth row corresponds to the nth slice for the mth scan.
 * \param bfname Base file name from which will be created 'bfname'.MeanDifference,
 * 'bfname'.MeanSquaredDifference and 'bfname'.NoOfVoxels.
 */
void DiffStatsVector::Write(const std::string& bfname) const EddyTry
{
  std::string fname = bfname + std::string(".MeanDifference");
  std::ofstream file;
  file.open(fname.c_str(),ios::out|ios::trunc);
  for (unsigned int i=0; i<_n; i++) file << _ds[i].MeanDifferenceVector() << endl;
  file.close();

  fname = bfname + std::string(".MeanSquaredDifference");
  file.open(fname.c_str(),ios::out|ios::trunc);
  for (unsigned int i=0; i<_n; i++) file << _ds[i].MeanSqrDiffVector() << endl;
  file.close();

  fname = bfname + std::string(".NoOfVoxels");
  file.open(fname.c_str(),ios::out|ios::trunc);
  for (unsigned int i=0; i<_n; i++) file << _ds[i].NVoxVector() << endl;
  file.close();
} EddyCatch

ReplacementManager::ReplacementManager(const std::vector<std::vector<unsigned int> >& shi,    // Shell indicies
				       const std::vector<double>&                     bvs,    // b-values of the shells
				       unsigned int                                   nsl,    // # of slices
				       const OutlierDefinition&                       old,    // Class defining an outlier
				       unsigned int                                   etc,    // =1->const (across slices) type 1 error, =2->const type 2 error
				       OLType                                         olt,    // Slice-wise, group-wise or both.
				       const MultiBandGroups&                         mbg)    // multi-band structure
EddyTry : _old(old), _etc(etc), _olt(olt), _mbg(mbg), _shi(shi), _bvs(bvs) 
{
  if (_etc != 1 && _etc != 2) throw  EddyException("ReplacementManager::ReplacementManager: etc must be 1 or 2");

  unsigned int nscan=0;
  for (unsigned int i=0; i<_shi.size(); i++) nscan += _shi[i].size();
  _sws = ReplacementManager::StatsInfo(nsl,nscan);
  _gws = ReplacementManager::StatsInfo(_mbg.NGroups(),nscan);
  _swo = ReplacementManager::OutlierInfo(nsl,nscan);
  _gwo = ReplacementManager::OutlierInfo(_mbg.NGroups(),nscan);
} EddyCatch

void ReplacementManager::Update(const DiffStatsVector& dsv) EddyTry
{
  if (dsv.NSlice() != this->NSlice() || dsv.NScan() != this->NScan()) {
    throw EddyException("ReplacementManager::Update: Mismatched DiffStatsVector object");
  }
  // Populate slice-wise stats matrix
  for (unsigned int scan=0; scan<NScan(); scan++) {
    for (unsigned int sl=0; sl<NSlice(); sl++) {
      _sws._nvox[sl][scan] = dsv.NVox(scan,sl);
      _sws._mdiff[sl][scan] = dsv.MeanDiff(scan,sl);
      _sws._msqrd[sl][scan] = dsv.MeanSqrDiff(scan,sl);
    }
  }
  // Populate group-wise stats matrix
  for (unsigned int scan=0; scan<NScan(); scan++) {
    for (unsigned int grp=0; grp<NGroup(); grp++) {
      std::vector<unsigned int> sl_in_grp = _mbg.SlicesInGroup(grp);
      _gws._nvox[grp][scan] = 0; _gws._mdiff[grp][scan] = 0.0; _gws._msqrd[grp][scan] = 0.0;
      for (unsigned int i=0; i<sl_in_grp.size(); i++) {
	_gws._nvox[grp][scan] += _sws._nvox[sl_in_grp[i]][scan];
	_gws._mdiff[grp][scan] += _sws._nvox[sl_in_grp[i]][scan] * _sws._mdiff[sl_in_grp[i]][scan];
	_gws._msqrd[grp][scan] += _sws._nvox[sl_in_grp[i]][scan] * _sws._msqrd[sl_in_grp[i]][scan];
      }
      _gws._mdiff[grp][scan] /= static_cast<double>(_gws._nvox[grp][scan]);
      _gws._msqrd[grp][scan] /= static_cast<double>(_gws._nvox[grp][scan]);
    }
  }
  // Calculate population statistics
  _sss = mean_and_std(_sws,_shi,_old.MinVoxels(),_etc,_swo._ovv); // Slice-wise summary stats
  _gss = mean_and_std(_gws,_shi,_old.MinVoxels(),_etc,_gwo._ovv); // Group-wise summary stats
  // First group-wise map
  for (unsigned int sh=0; sh<_shi.size(); sh++) {
    for (const unsigned int& scan : _shi[sh]) {
      for (unsigned int grp=0; grp<NGroup(); grp++) {
	if (_gws._nvox[grp][scan] >= _old.MinVoxels()) { // Only consider groups with more than minimum valid voxels
	  double sf = (_etc == 1) ? 1.0 / std::sqrt(double(_gws._nvox[grp][scan])) : 1.0;
	  _gwo._nsv[grp][scan] = (_gws._mdiff[grp][scan] - _gss[sh]._diff_mean) / (sf * _gss[sh]._diff_std);
	  _gwo._nsq[grp][scan] = (_gws._msqrd[grp][scan] - _gss[sh]._sqr_mean) / (sf * _gss[sh]._sqr_std);
	  _gwo._ovv[grp][scan] = (-_gwo._nsv[grp][scan] > _old.NStdDev()) ? true : false;
	  if (_old.ConsiderPosOL()) _gwo._ovv[grp][scan] = (_gwo._nsv[grp][scan] > _old.NStdDev()) ? true : _gwo._ovv[grp][scan];
	  if (_old.ConsiderSqrOL()) _gwo._ovv[grp][scan] = (_gwo._nsq[grp][scan] > _old.NStdDev()) ? true : _gwo._ovv[grp][scan];
	}
      }
    }
  }
  // Then slice-wise map
  if (_olt==OLType::GroupWise || _olt==OLType::Both) { // If based (wholy or partially) on mb-groups
    // First pass to identify candidates
    for (unsigned int grp=0; grp<NGroup(); grp++) {
      std::vector<unsigned int> sl_in_grp = _mbg.SlicesInGroup(grp);
      for (unsigned int scan=0; scan<NScan(); scan++) {
	for (unsigned int i=0; i<sl_in_grp.size(); i++) {
	  _swo._nsv[sl_in_grp[i]][scan] = _gwo._nsv[grp][scan];
	  _swo._nsq[sl_in_grp[i]][scan] = _gwo._nsq[grp][scan];
	  _swo._ovv[sl_in_grp[i]][scan] = _gwo._ovv[grp][scan];
	}
      }
    }
    if (_olt==OLType::Both) { // If we should also consider slice-based stats
      // First do a second pass of the group-wise stats to weed out groups driven by single slice
      for (unsigned int grp=0; grp<NGroup(); grp++) {
	std::vector<unsigned int> sl_in_grp = _mbg.SlicesInGroup(grp);
	for (unsigned int sh=0; sh<_shi.size(); sh++) {
	  for (const unsigned int& scan : _shi[sh]) {
	    bool group_kosher = true;
	    for (unsigned int i=0; i<sl_in_grp.size(); i++) {
	      if (_sws._nvox[sl_in_grp[i]][scan] >= _old.MinVoxels()) { // Only check slices with more than minimum valid voxels
		double sf = (_etc == 1) ? 1.0 / std::sqrt(double(_sws._nvox[sl_in_grp[i]][scan])) : 1.0;
		double tmp_ns = (_sws._mdiff[sl_in_grp[i]][scan] - _sss[sh]._diff_mean) / (sf * _sss[sh]._diff_std);
		group_kosher = (-tmp_ns > (_old.NStdDev()-1.0)) ? true : false;
		if (_old.ConsiderPosOL()) group_kosher = (tmp_ns > (_old.NStdDev()-1.0) && _gwo._nsv[grp][scan] > 0.0) ? true : group_kosher;
		if (_old.ConsiderSqrOL()) {
		  tmp_ns = (_sws._msqrd[sl_in_grp[i]][scan] - _sss[sh]._sqr_mean) / (sf * _sss[sh]._sqr_std);
		  group_kosher = (tmp_ns > (_old.NStdDev()-1.0)) ? true : group_kosher;
		}
	      }
	      if (!group_kosher) break;
	    }
	    if (!group_kosher) { // Reset group
	      for (unsigned int i=0; i<sl_in_grp.size(); i++) {
		_swo._nsv[sl_in_grp[i]][scan] = 0.0;    // Means that it will be set by slice
		_swo._nsq[sl_in_grp[i]][scan] = 0.0;    // Means that it will be set by slice
		_swo._ovv[sl_in_grp[i]][scan] = false;
	      }
	    }
	  }
	}
      }
      // Then add slices that are outliers in their own right
      for (unsigned int sh=0; sh<_shi.size(); sh++) {
	for (const unsigned int& scan : _shi[sh]) {
	  for (unsigned int sl=0; sl<NSlice(); sl++) {
	    if (_sws._nvox[sl][scan] >= _old.MinVoxels()) { // Only consider slices with more than minimum valid voxels
	      double sf = (_etc == 1) ? 1.0 / std::sqrt(double(_sws._nvox[sl][scan])) : 1.0;
	      double tmp_ns = (_sws._mdiff[sl][scan] - _sss[sh]._diff_mean) / (sf * _sss[sh]._diff_std);
	      _swo._nsv[sl][scan] = (std::abs(tmp_ns) > std::abs(_swo._nsv[sl][scan])) ? tmp_ns : _swo._nsv[sl][scan];
	      tmp_ns = (_sws._msqrd[sl][scan] - _sss[sh]._sqr_mean) / (sf * _sss[sh]._sqr_std);
	      _swo._nsq[sl][scan] = (std::abs(tmp_ns) > std::abs(_swo._nsq[sl][scan])) ? tmp_ns : _swo._nsq[sl][scan];
	      _swo._ovv[sl][scan] = (-_swo._nsv[sl][scan] > _old.NStdDev()) ? true : false;
	      if (_old.ConsiderPosOL()) _swo._ovv[sl][scan] = (_swo._nsv[sl][scan] > _old.NStdDev()) ? true : _swo._ovv[sl][scan];
	      if (_old.ConsiderSqrOL()) _swo._ovv[sl][scan] = (_swo._nsq[sl][scan] > _old.NStdDev()) ? true : _swo._ovv[sl][scan];
	    }
	  }
	}
      }
    }
  }
  else { // If slice-based
    for (unsigned int sh=0; sh<_shi.size(); sh++) {
      for (const unsigned int& scan : _shi[sh]) {
	for (unsigned int sl=0; sl<NSlice(); sl++) {
	  if (_sws._nvox[sl][scan] >= _old.MinVoxels()) { // Only consider slices with more than minimum valid voxels
	    double sf = (_etc == 1) ? 1.0 / std::sqrt(double(_sws._nvox[sl][scan])) : 1.0;
	    _swo._nsv[sl][scan] = (_sws._mdiff[sl][scan] - _sss[sh]._diff_mean) / (sf * _sss[sh]._diff_std);
	    _swo._nsq[sl][scan] = (_sws._msqrd[sl][scan] - _sss[sh]._sqr_mean) / (sf * _sss[sh]._sqr_std);
	    _swo._ovv[sl][scan] = (-_swo._nsv[sl][scan] > _old.NStdDev()) ? true : false;
	    if (_old.ConsiderPosOL()) _swo._ovv[sl][scan] = (_swo._nsv[sl][scan] > _old.NStdDev()) ? true : _swo._ovv[sl][scan];
	    if (_old.ConsiderSqrOL()) _swo._ovv[sl][scan] = (_swo._nsq[sl][scan] > _old.NStdDev()) ? true : _swo._ovv[sl][scan];
	  }
	}
      }
    }
  }
} EddyCatch

std::vector<unsigned int> ReplacementManager::OutliersInScan(unsigned int scan) const EddyTry
{
  throw_if_oor(scan);
  std::vector<unsigned int> ol;
  for (unsigned int sl=0; sl<NSlice(); sl++) if (_swo._ovv[sl][scan]) ol.push_back(sl);
  return(ol);
} EddyCatch

bool ReplacementManager::ScanHasOutliers(unsigned int scan) const EddyTry
{
  throw_if_oor(scan);
  for (unsigned int sl=0; sl<NSlice(); sl++) if (_swo._ovv[sl][scan]) return(true);
  return(false);
} EddyCatch

nlohmann::json ReplacementManager::WriteReport(const std::vector<unsigned int>& i2i,
					       const std::string&               fname,
					       bool                             write_old_style_file) const EddyTry
{
  if (write_old_style_file) { 
    std::ofstream fout;
    fout.open(fname.c_str(),ios::out|ios::trunc);
    if (fout.fail()) throw EddyException("ReplacementManager::WriteReport:Failed to open outlier report file " + fname);
    for (unsigned int sl=0; sl<NSlice(); sl++) {
      for (unsigned int s=0; s<NScan(); s++) {
	if (_swo._ovv[sl][s]) {
	  fout << "Slice " << sl << " in scan " << i2i[s] << " is an outlier with mean " << _swo._nsv[sl][s] << " standard deviations off, and mean squared " << _swo._nsq[sl][s] << " standard deviations off." << endl;
	}
      }
    }
    fout.close();
  }

  // Next make json-object
  std::vector<nlohmann::json> ol_object_vector;
  for (unsigned int sl=0; sl<NSlice(); sl++) {
    for (unsigned int s=0; s<NScan(); s++) {
      if (_swo._ovv[sl][s]) {
	nlohmann::json ol_object;
	ol_object["Slice"] = sl;
	ol_object["Volume"] = i2i[s];
	ol_object["Mean: # of standard deviations off"] = _swo._nsv[sl][s];
	ol_object["Mean-squared: # of standard deviations off"] = _swo._nsq[sl][s];
	ol_object_vector.push_back(ol_object);
      }
    }
  }
  nlohmann::json ol_report_object;
  ol_report_object["Outlier list"] = ol_object_vector;
  
  return(ol_report_object);
} EddyCatch

std::vector<nlohmann::json> ReplacementManager::WriteMatrixReport(const std::vector<unsigned int>& i2i,
								  unsigned int                     nscan,
								  const std::string&               om_fname,
								  const std::string&               nstdev_fname,
								  const std::string&               n_sqr_stdev_fname,
								  bool                             write_old_style_file) const EddyTry
{
  if (write_old_style_file) {
    std::ofstream fout;
    try {
      if (!om_fname.empty()) {
	fout.open(om_fname.c_str(),ios::out|ios::trunc);
	fout << "One row per scan, one column per slice. Outlier: 1, Non-outlier: 0" << endl;
	// Repack into matrix with b0's in place
	NEWMAT::Matrix rpovv(nscan,NSlice()); rpovv = 0.0;
	for (unsigned int scan=0; scan<NScan(); scan++) {
	  for (unsigned int slice=0; slice<NSlice(); slice++) {
	    if (_swo._ovv[slice][scan]) rpovv(i2i[scan]+1,slice+1) = 1.0;
	  }
	}
	// Write repacked version
	for (unsigned int scan=0; scan<nscan; scan++) {
	  for (unsigned int slice=0; slice<NSlice(); slice++) {
	    if (rpovv(scan+1,slice+1) > 0) fout << "1 ";
	    else fout << "0 ";
	  }
	  fout << endl;
	}
	fout.close();
      }
    }
    catch (...) { throw EddyException("ReplacementManager::WriteMatrixReport: Error when writing file " + om_fname); }
    // Now start on nstdev of mean file
    if (!nstdev_fname.empty()) {
      try {
	fout.open(nstdev_fname.c_str(),ios::out|ios::trunc);
	fout << "One row per scan, one column per slice. b0s set to zero" << endl;
	// Repack into matrix with b0's in place
	NEWMAT::Matrix rpnsv(nscan,NSlice()); rpnsv = 0.0;
	for (unsigned int scan=0; scan<NScan(); scan++) {
	  for (unsigned int slice=0; slice<NSlice(); slice++) {
	    rpnsv(i2i[scan]+1,slice+1) = _swo._nsv[slice][scan];
	  }
	}
	// Write repacked version
	for (unsigned int scan=0; scan<nscan; scan++) {
	  for (unsigned int slice=0; slice<NSlice(); slice++) {
	    double tmp = rpnsv(scan+1,slice+1);
	    fout << tmp << " ";
	  }
	  fout << endl;
	}
	fout.close();
      }
      catch (...) { throw EddyException("ReplacementManager::WriteMatrixReport: Error when writing file " + nstdev_fname); }
    }
    // Finally do nstdev of squared mean file
    if (!n_sqr_stdev_fname.empty()) {
      try {
	fout.open(n_sqr_stdev_fname.c_str(),ios::out|ios::trunc);
	fout << "One row per scan, one column per slice. b0s set to zero" << endl;
	// Repack into matrix with b0's in place
	NEWMAT::Matrix rpnsv(nscan,NSlice()); rpnsv = 0.0;
	for (unsigned int scan=0; scan<NScan(); scan++) {
	  for (unsigned int slice=0; slice<NSlice(); slice++) {
	    rpnsv(i2i[scan]+1,slice+1) = _swo._nsq[slice][scan];
	  }
	}
	// Write repacked version
	for (unsigned int scan=0; scan<nscan; scan++) {
	  for (unsigned int slice=0; slice<NSlice(); slice++) {
	    double tmp = rpnsv(scan+1,slice+1);
	    fout << tmp << " ";
	  }
	  fout << endl;
	}
	fout.close();
      }
      catch (...) { throw EddyException("ReplacementManager::WriteMatrixReport: Error when writing file " + n_sqr_stdev_fname); }
    }
  }
  // Next make json-objects
  std::vector<std::vector<bool> > ol_map(nscan,std::vector<bool>(NSlice(),false));  
  std::vector<std::vector<double> > stdev_map(nscan,std::vector<double>(NSlice(),0.0));  
  std::vector<std::vector<double> > sqr_stdev_map(nscan,std::vector<double>(NSlice(),0.0)); 
  for (unsigned int scan=0; scan<NScan(); scan++) {
    for (unsigned int slice=0; slice<NSlice(); slice++) {
      ol_map[i2i[scan]][slice] = _swo._ovv[slice][scan];
      stdev_map[i2i[scan]][slice] = _swo._nsv[slice][scan];
      sqr_stdev_map[i2i[scan]][slice] = _swo._nsq[slice][scan];
    }
  }
  nlohmann::json ol_map_json = {"Outlier map", ol_map};
  nlohmann::json stdev_map_json = {"Standard deviation of mean difference map", stdev_map};
  nlohmann::json sqr_stdev_map_json = {"Standard deviation of mean-square-difference map", sqr_stdev_map};
  std::vector<nlohmann::json> all_maps = {ol_map_json, stdev_map_json, sqr_stdev_map_json};

  return(all_maps);
} EddyCatch

void ReplacementManager::WriteDebugInfo(const std::string&               fname, 
					const std::vector<unsigned int>& i2i,
					unsigned int                     nscan) const EddyTry
{
  nlohmann::json top;
  std::vector<std::vector<bool> > outlier_map(nscan,std::vector<bool>(NSlice(),false));
  std::vector<std::vector<double> > nstdev_map(nscan,std::vector<double>(NSlice(),0.0));  
  std::vector<std::vector<double> > sqr_nstdev_map(nscan,std::vector<double>(NSlice(),0.0));

  top["NumberOfScans"] = nscan;
  top["NumberOfSlices"] = NSlice();
  top["NumberOfMBGroups"] = NGroup();
  top["OutlierDefinition"]["NoOfStdDev"] = _old.NStdDev();
  top["OutlierDefinition"]["MinNoOfVoxels"] = _old.MinVoxels();
  top["OutlierDefinition"]["ConsiderPositive"] = (_old.ConsiderPosOL()) ? "Yes" : "No";
  top["OutlierDefinition"]["ConsiderSquaredDiff"] = (_old.ConsiderSqrOL()) ? "Yes" : "No";
  top["OutlierDefinition"]["ConstantErrorType"] = _etc;
  top["OutlierDefinition"]["OutlierType"] = (_olt==OLType::SliceWise) ? "Slice-wise" : ((_olt==OLType::GroupWise) ? "Group-wise" : "Both");

  for (unsigned int sh=0; sh<_shi.size(); sh++) {
    if (_bvs.size()) top["ShellInfo"][sh]["BValue"] = _bvs[sh];
    else top["ShellInfo"][sh]["BValue"] = "Unknown b-value";
    top["ShellInfo"][sh]["SliceWiseDiffMean"] = _sss[sh]._diff_mean;
    top["ShellInfo"][sh]["SliceWiseSqrMean"] = _sss[sh]._sqr_mean;
    top["ShellInfo"][sh]["GroupWiseDiffMean"] = _gss[sh]._diff_mean;
    top["ShellInfo"][sh]["GroupWiseSqrMean"] = _gss[sh]._sqr_mean;
    double sf = (_etc == 1) ? 1.0 / std::sqrt(_sws._nvox[static_cast<int>(NSlice()/2)][static_cast<int>(NScan()/2)]) : 1.0; // Arbitrarily mid-slice of mid-scan
    top["ShellInfo"][sh]["SliceWiseDiffStdDev"] = sf * _sss[sh]._diff_std;
    top["ShellInfo"][sh]["SliceWiseSqrStdDev"] = sf * _sss[sh]._sqr_std;
    sf = (_etc == 1) ? 1.0 / std::sqrt(_gws._nvox[static_cast<int>(NGroup()/2)][static_cast<int>(NScan()/2)]) : 1.0;       // Arbitrarily mid-group of mid-scan
    top["ShellInfo"][sh]["GroupWiseDiffStdDev"] = sf * _gss[sh]._diff_std;
    top["ShellInfo"][sh]["GroupWiseSqrStdDev"] = sf * _gss[sh]._sqr_std;
    std::vector<unsigned int> global_index(_shi[sh].size());
    for (unsigned int i=0; i<_shi[sh].size(); i++) global_index[i] = i2i[_shi[sh][i]];
    top["ShellInfo"][sh]["Index"] = global_index;
  }

  top["SliceWiseInfo"]["RawStats"]["NoOfVoxels"] = unpack(_sws._nvox,nscan,i2i);
  top["SliceWiseInfo"]["RawStats"]["DiffMean"] = unpack(_sws._mdiff,nscan,i2i);
  top["SliceWiseInfo"]["RawStats"]["SqrMean"] = unpack(_sws._msqrd,nscan,i2i);
  top["SliceWiseInfo"]["Transformed"]["OutlierMap"] = unpack(_swo._ovv,nscan,i2i);
  top["SliceWiseInfo"]["Transformed"]["NDiffStDev"] = unpack(_swo._nsv,nscan,i2i);
  top["SliceWiseInfo"]["Transformed"]["NSqrStDev"] = unpack(_swo._nsq,nscan,i2i);

  top["GroupWiseInfo"]["RawStats"]["NoOfVoxels"] = unpack(_gws._nvox,nscan,i2i);
  top["GroupWiseInfo"]["RawStats"]["DiffMean"] = unpack(_gws._mdiff,nscan,i2i);
  top["GroupWiseInfo"]["RawStats"]["SqrMean"] = unpack(_gws._msqrd,nscan,i2i);
  top["GroupWiseInfo"]["Transformed"]["OutlierMap"] = unpack(_gwo._ovv,nscan,i2i);
  top["GroupWiseInfo"]["Transformed"]["NDiffStDev"] = unpack(_gwo._nsv,nscan,i2i);
  top["GroupWiseInfo"]["Transformed"]["NSqrStDev"] = unpack(_gwo._nsq,nscan,i2i);

  try {
    std::ofstream ofile(fname);
    ofile << std::setw(4) << top << std::endl;
    ofile.close();
  }
  catch (const std::exception& e) {
    throw EddyException("ReplacementManager::WriteDebugInfo: Exception thrown with message " + std::string(e.what()));
  }
} EddyCatch

template<typename T>
std::vector<std::vector<T> > ReplacementManager::unpack(const std::vector<std::vector<T> >& mat,
							unsigned int                        nscan,
							const std::vector<unsigned int>&    i2i) const EddyTry
{
  unsigned int nsl = mat.size(); // Number of slices/MB-groups
  std::vector<std::vector<T> > rval(nscan,std::vector<T>(nsl,static_cast<T>(0)));
  for (unsigned int scan=0; scan<NScan(); scan++) {
    for (unsigned int sl=0; sl<nsl; sl++) {
      rval[i2i[scan]][sl] = mat[sl][scan];
    }
  }
  return(rval);
} EddyCatch

void ReplacementManager::DumpOutlierMaps(const std::string& bfname) const EddyTry
{
  std::string fname = bfname + ".SliceWiseOutlierMap";
  std::ofstream fout;
  fout.open(fname.c_str(),ios::out|ios::trunc);
  fout << "One row per scan, one column per slice. Outlier: 1, Non-outlier: 0" << endl;
  for (unsigned scan=0; scan<NScan(); scan++) {
    for (unsigned int slice=0; slice<NSlice(); slice++) {
      fout << _swo._ovv[slice][scan] << " ";
    }
    fout << endl;
  }
  fout.close();

  fname = bfname + ".SliceWiseNoOfStdevMap";
  fout.open(fname.c_str(),ios::out|ios::trunc);
  fout << "One row per scan, one column per slice." << endl;
  for (unsigned scan=0; scan<NScan(); scan++) {
    for (unsigned int slice=0; slice<NSlice(); slice++) {
      fout << _gwo._nsv[slice][scan] << " ";
    }
    fout << endl;
  }
  fout.close();

  fname = bfname + ".GroupWiseOutlierMap";
  fout.open(fname.c_str(),ios::out|ios::trunc);
  fout << "One row per scan, one column per slice. Outlier: 1, Non-outlier: 0" << endl;
  for (unsigned scan=0; scan<NScan(); scan++) {
    for (unsigned int slice=0; slice<NSlice(); slice++) {
      fout << _gwo._ovv[slice][scan] << " ";
    }
    fout << endl;
  }
  fout.close();

  fname = bfname + ".GroupWiseNoOfStdevMap";
  fout.open(fname.c_str(),ios::out|ios::trunc);
  fout << "One row per scan, one column per slice." << endl;
  for (unsigned scan=0; scan<NScan(); scan++) {
    for (unsigned int slice=0; slice<NSlice(); slice++) {
      fout << _swo._nsv[slice][scan] << " ";
    }
    fout << endl;
  }
  fout.close();
} EddyCatch

std::vector<ReplacementManager::SummaryStats> ReplacementManager::mean_and_std(// Input
									       const EDDY::ReplacementManager::StatsInfo&     stats, // Slice/group-wise stats
									       const std::vector<std::vector<unsigned int> >& shi,   // Shell indicies
									       unsigned int                                   minvox,// Smallest allowed # of voxels in a slice
									       unsigned int                                   etc,   // ErrorTypeControl (type 1 or 2)
									       const std::vector<std::vector<bool> >&         ovv) const EddyTry   // Slices/groups currently labeled as outliers
{
  std::vector<ReplacementManager::SummaryStats> rval(shi.size());

  // First pass to get means
  for (unsigned int sh=0; sh<shi.size(); sh++) {
    unsigned int ntot = 0;
    for (unsigned int sl_gr=0; sl_gr<stats._mdiff.size();  sl_gr++) {
      for (const unsigned int& s : shi[sh]) {
	if (!ovv[sl_gr][s] && stats._nvox[sl_gr][s] >= minvox) {
	  if (etc == 1) {
	    rval[sh]._diff_mean += stats._nvox[sl_gr][s] * stats._mdiff[sl_gr][s];
	    rval[sh]._sqr_mean += stats._nvox[sl_gr][s] * stats._msqrd[sl_gr][s];
	    ntot += stats._nvox[sl_gr][s];
	  }
	  else {
	    rval[sh]._diff_mean += stats._mdiff[sl_gr][s];
	    rval[sh]._sqr_mean += stats._msqrd[sl_gr][s];
	    ntot += 1;
	  }
	}
      }
    }
    rval[sh]._diff_mean /= double(ntot);
    rval[sh]._sqr_mean /= double(ntot);
  }
  // Second pass to get standard deviations
  // If etc==1, i.e. if we want to keep the false-positive
  // rate constant we, "guesstimate" the underlying
  // voxel-wise standard deviation.
  for (unsigned int sh=0; sh<shi.size(); sh++) {
    unsigned int ntot = 0;
    for (unsigned int sl_gr=0; sl_gr<stats._mdiff.size();  sl_gr++) {
      for (const unsigned int& s : shi[sh]) { 
	// for (unsigned int scan=0; scan<stats._mdiff[sl_gr].size(); scan++) {
	if (!ovv[sl_gr][s] && stats._nvox[sl_gr][s] >= minvox) {
	  if (etc == 1) {
	    rval[sh]._diff_std += stats._nvox[sl_gr][s] * this->sqr(stats._mdiff[sl_gr][s] - rval[sh]._diff_mean);
	    rval[sh]._sqr_std += stats._nvox[sl_gr][s] * this->sqr(stats._msqrd[sl_gr][s] - rval[sh]._sqr_mean);
	  }
	  else {
	    rval[sh]._diff_std += this->sqr(stats._mdiff[sl_gr][s] - rval[sh]._diff_mean);
	    rval[sh]._sqr_std += this->sqr(stats._msqrd[sl_gr][s] - rval[sh]._sqr_mean);
	  }
	  ntot += 1;
	}
      }
    }
    rval[sh]._diff_std /= double(ntot - 1); rval[sh]._diff_std = std::sqrt(rval[sh]._diff_std); 
    rval[sh]._sqr_std /= double(ntot - 1); rval[sh]._sqr_std = std::sqrt(rval[sh]._sqr_std);
  }

  return(rval);
} EddyCatch

double MutualInfoHelper::MI(const NEWIMAGE::volume<float>& ima1,
			    const NEWIMAGE::volume<float>& ima2,
			    const NEWIMAGE::volume<float>& mask) const EddyTry
{
  // Make joint histograms
  memset(_mhist1,0,_nbins*sizeof(double));
  memset(_mhist2,0,_nbins*sizeof(double));
  memset(_jhist,0,_nbins*_nbins*sizeof(double));
  float min1, max1, min2, max2;
  if (!_lset) { min1 = ima1.min(); max1 = ima1.max(); min2 = ima2.min(); max2 = ima2.max(); }
  else { min1 = _min1; max1 = _max1; min2 = _min2; max2 = _max2; }
  unsigned int nvox=0;
  for (int z=0; z<ima1.zsize(); z++) {
    for (int y=0; y<ima1.ysize(); y++) {
      for (int x=0; x<ima1.xsize(); x++) {
	if (mask(x,y,z)) {
	  unsigned int i1 = this->val_to_indx(ima1(x,y,z),min1,max1,_nbins);
	  unsigned int i2 = this->val_to_indx(ima2(x,y,z),min2,max2,_nbins);
	  _mhist1[i1] += 1.0;
	  _mhist2[i2] += 1.0;
	  _jhist[i2*_nbins + i1] += 1.0;
	  nvox++;
	}
      }
    }
  }
  // Calculate entropies
  double je=0.0; double me1=0.0; double me2=0.0;
  for (unsigned int i1=0; i1<_nbins; i1++) {
    me1 += this->plogp(_mhist1[i1]/static_cast<double>(nvox));
    me2 += this->plogp(_mhist2[i1]/static_cast<double>(nvox));
    for (unsigned int i2=0; i2<_nbins; i2++) {
      je += this->plogp(_jhist[i2*_nbins + i1]/static_cast<double>(nvox));
    }
  }
  // return mutual information
  return(me1+me2-je);
} EddyCatch

double MutualInfoHelper::SoftMI(const NEWIMAGE::volume<float>& ima1,
				const NEWIMAGE::volume<float>& ima2,
				const NEWIMAGE::volume<float>& mask) const EddyTry
{
  // Make joint histograms
  memset(_mhist1,0,_nbins*sizeof(double));
  memset(_mhist2,0,_nbins*sizeof(double));
  memset(_jhist,0,_nbins*_nbins*sizeof(double));
  float min1, max1, min2, max2;
  if (!_lset) { min1 = ima1.min(); max1 = ima1.max(); min2 = ima2.min(); max2 = ima2.max(); }
  else { min1 = _min1; max1 = _max1; min2 = _min2; max2 = _max2; }
  double nvox=0.0;
  for (int z=0; z<ima1.zsize(); z++) {
    for (int y=0; y<ima1.ysize(); y++) {
      for (int x=0; x<ima1.xsize(); x++) {
	if (mask(x,y,z)) {
	  float mv = mask(x,y,z);
	  float r1, r2;
	  unsigned int i1 = this->val_to_floor_indx(ima1(x,y,z),min1,max1,_nbins,&r1);
	  unsigned int i2 = this->val_to_floor_indx(ima2(x,y,z),min2,max2,_nbins,&r2);
	  _mhist1[i1] += mv*(1.0 - r1);
	  if (r1) _mhist1[i1+1] += mv*r1;
	  _mhist2[i2] += mv*(1.0 - r2);
	  if (r2) _mhist2[i2+1] += mv*r2;
	  _jhist[i2*_nbins + i1] += mv*(1.0 - r1)*(1.0 - r2);
	  if (r1) _jhist[i2*_nbins + i1+1] += mv*r1*(1.0 - r2);
	  if (r2) _jhist[(i2+1)*_nbins + i1] += mv*(1.0 - r1)*r2;
	  if (r1 && r2) _jhist[(i2+1)*_nbins + i1+1] += mv*r1*r2;
	  nvox += mv;
	}
      }
    }
  }
  // Calculate entropies
  double je=0.0; double me1=0.0; double me2=0.0;
  for (unsigned int i1=0; i1<_nbins; i1++) {
    me1 += this->plogp(_mhist1[i1]/nvox);
    me2 += this->plogp(_mhist2[i1]/nvox);
    for (unsigned int i2=0; i2<_nbins; i2++) {
      je += this->plogp(_jhist[i2*_nbins + i1]/nvox);
    }
  }
  // return mutual information
  return(me1+me2-je);
} EddyCatch

void Stacks2YVecsAndWgts::MakeVectors(const NEWIMAGE::volume4D<float>& stacks,
	                              const NEWIMAGE::volume4D<float>& masks,
	                              const NEWIMAGE::volume4D<float>& zcoord,
	                              unsigned int                     i,
	                              unsigned int                     j) EddyTry
{
  int tsz = stacks.tsize();
  int zsz = stacks.zsize();
  std::fill(_n.begin(),_n.end(),0);
  for (int t=0; t<tsz; t++) {
    for (int k=0; k<zsz; k++) {
      if (masks(i,j,k,t) && zcoord(i,j,k,t)>-1.0 && zcoord(i,j,k,t)<zsz) {
	if (zcoord(i,j,k,t) < 0.0) {  // If before first element
 	  _wgt[0][_n[0]] = 1.0 + zcoord(i,j,k,t);
	  _sqrtwgt[0][_n[0]] = std::sqrt(_wgt[0][_n[0]]);
	  _y[0][_n[0]] = stacks(i,j,k,t);
	  _bvi[0][_n[0]].first = t;
	  _bvi[0][_n[0]].second = k;
	  _n[0]++;
	}
	else if (zcoord(i,j,k,t) > zsz-1) { // If beyond last element
 	  _wgt[zsz-1][_n[zsz-1]] = static_cast<float>(zsz) - zcoord(i,j,k,t);
	  _sqrtwgt[zsz-1][_n[zsz-1]] = std::sqrt(_wgt[zsz-1][_n[zsz-1]]);
	  _y[zsz-1][_n[zsz-1]] = stacks(i,j,k,t);
	  _bvi[zsz-1][_n[zsz-1]].first = t;
	  _bvi[zsz-1][_n[zsz-1]].second = k;
	  _n[zsz-1]++;
	}
	else { // If somewhere in the middle
	  int li = static_cast<int>(std::floor(zcoord(i,j,k,t)));
 	  _wgt[li][_n[li]] = 1.0 + static_cast<float>(li) - zcoord(i,j,k,t);
	  _sqrtwgt[li][_n[li]] = std::sqrt(_wgt[li][_n[li]]);
	  _y[li][_n[li]] = stacks(i,j,k,t);
	  _bvi[li][_n[li]].first = t;
	  _bvi[li][_n[li]].second = k;
	  _n[li]++;
	  int ui = static_cast<int>(std::ceil(zcoord(i,j,k,t)));
 	  _wgt[ui][_n[ui]] = 1.0 - static_cast<float>(ui) + zcoord(i,j,k,t);
	  _sqrtwgt[ui][_n[ui]] = std::sqrt(_wgt[ui][_n[ui]]);
	  _y[ui][_n[ui]] = stacks(i,j,k,t);
	  _bvi[ui][_n[ui]].first = t;
	  _bvi[ui][_n[ui]].second = k;
	  _n[ui]++;
	}
      }
    }
  }
} EddyCatch

NEWMAT::RowVector Indicies2KMatrix::GetkVector(const NEWMAT::ColumnVector& bvec,
					       unsigned int                grp) const EddyTry
{
  NEWMAT::RowVector rval(_K.Nrows());
  for (int i=0; i<rval.Ncols(); i++) {
    double th = std::acos(std::min(1.0,std::abs(NEWMAT::DotProduct(bvec,_bvecs[i])))); // theta
    double a = _thpar[1];
    if (a>th) {
      rval(i+1) = _thpar[0] * (1.0 - 1.5*th/a + 0.5*(th*th*th)/(a*a*a));
      if (_grpb.size() > 1 && _grpi[i] != grp) {
	double bvdiff = _log_grpb[_grpi[i]] - _log_grpb[grp];
        rval(i+1) *= std::exp(-(bvdiff*bvdiff) / (2*_thpar[2]*_thpar[2]));
      }
      if (_wgt.size() != 0) rval(i+1) *= _wgt[i];
    }
    else rval(i+1) = 0.0;
  }
  return(rval);
} EddyCatch

void Indicies2KMatrix::common_construction(const std::vector<std::vector<NEWMAT::ColumnVector> >& bvecs,
					   const std::vector<unsigned int>&                       grpi,
					   const std::vector<double>&                             grpb,
					   const std::vector<std::pair<int,int> >&                indx,
					   unsigned int                                           nval,
					   const std::vector<double>&                             hpar,
					   const std::vector<double>                              *wgt) EddyTry
{

  // Set and transform hyper-parameters and group b-values
  _grpb = grpb;
  _log_grpb = _grpb; for (unsigned int i=0; i<_log_grpb.size(); i++) _log_grpb[i] = std::log(_grpb[i]);
  _hpar = hpar;
  _thpar = _hpar; for (unsigned int i=0; i<_thpar.size(); i++) _thpar[i] = std::exp(_hpar[i]);
  if (wgt != nullptr) _wgt = *wgt;

  // First extract relevant vectors of bvecs and bval-groups
  _bvecs.resize(nval); _grpi.resize(nval);
  for (unsigned int i=0; i<nval; i++) {
    _bvecs[i] = bvecs[indx[i].first][indx[i].second];
    _grpi[i] = grpi[indx[i].first];
  }

  // Next make angle matrix
  _K.resize(nval,nval);
  for (unsigned int j=0; j<nval; j++) {
    for (unsigned int i=j; i<nval; i++) {
      if (i==j) _K(i+1,j+1) = 0.0;
      else {
	_K(i+1,j+1) = std::acos(std::min(1.0,std::abs(NEWMAT::DotProduct(_bvecs[i],_bvecs[j]))));
      }
    }
  }

  // Next do first pass for angular covariance
  double sm = _thpar[0]; double a = _thpar[1];
  for (int j=0; j<_K.Ncols(); j++) {
    for (int i=j; i<_K.Nrows(); i++) {
      double th = _K(i+1,j+1); // theta
      if (a>th) _K(i+1,j+1) = sm * (1.0 - 1.5*th/a + 0.5*(th*th*th)/(a*a*a));
      else _K(i+1,j+1) = 0.0;
      if (i!=j) _K(j+1,i+1) = _K(i+1,j+1);
    }
  }

  // Second pass for b-value covariance
  if (_grpb.size() > 1) {
    double l = _thpar[2];
    for (int j=0; j<_K.Ncols(); j++) {
      for (int i=j+1; i<_K.Nrows(); i++) {
	if (_K(i+1,j+1) != 0.0) {
	  double bvdiff = _log_grpb[_grpi[i]] - _log_grpb[_grpi[j]];
	  if (bvdiff) {
	    _K(i+1,j+1) *= std::exp(-(bvdiff*bvdiff) / (2*l*l));
	    _K(j+1,i+1) = _K(i+1,j+1);
	  }
	}
      }
    }
  }

  // (Optional) third pass for weights
  if (_wgt.size() != 0) {
    for (int j=0; j<_K.Ncols(); j++) {
      for (int i=j; i<_K.Nrows(); i++) {
	_K(i+1,j+1) *= _wgt[i]*_wgt[j];
	if (i!=j) _K(j+1,i+1) = _K(i+1,j+1);
      }
    }
  }

  // Fourth pass for error variances
  const double *ev = nullptr;
  ev = (_grpb.size() > 1) ? &(_thpar[3]) : &(_thpar[2]);
  for (int i=0; i<_K.Ncols(); i++) {
    _K(i+1,i+1) += ev[_grpi[i]];
  }

} EddyCatch