// Declarations of classes that implements a hirearchy
// of models for long-time constant fields from eddy 
// currents induced by diffusion gradients.
//
// LongECModels.h
//
// Jesper Andersson, FMRIB Image Analysis Group
//
// Copyright (C) 2022 University of Oxford
//

#include <cstdlib>
#include <string>
#include <vector>
#include <cmath>
#include <iostream>
#include <fstream>
#include <time.h>
#include "nlohmann/json.hpp"
#include "armawrap/newmat.h"
#ifndef EXPOSE_TREACHEROUS
#define EXPOSE_TREACHEROUS           // To allow us to use .set_sform etc
#endif
#include "newimage/newimageall.h"
#include "miscmaths/miscmaths.h"
#include "EddyHelperClasses.h"
#include "LongECModels.h"
#include "ECScanClasses.h"

using namespace EDDY;

// Initial value for eddy current time constant (half-life in seconds).
// small enough to result in little time dependence, but large enough to have a derivative
double const TC_INITIAL_VALUE = 0.1;
double const INTCPT_INITIAL_VALUE = 0.01;

//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class EC_TimeAndIntercept2Weights
//
// This class translates an EC time constant and intercept to weights
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
arma::colvec EC_TimeAndIntercept2Weights::GetWeightsForPrevious(double Thl,            
								double Intcpt) const EddyTry
{
  arma::colvec weights(_mbg.NGroups());
  for (int i=0; i<weights.n_rows; i++) {
    weights(i) = Intcpt * std::exp(-(i)*(_TR/_mbg.NGroups())*std::log(2.0)/Thl);
  }
  return(weights);
} EddyCatch  

arma::colvec EC_TimeAndIntercept2Weights::GetWeightsForCurrent(double Thl,            
							       double Intcpt) const EddyTry
{
  arma::colvec weights(_mbg.NGroups());
  for (int i=0; i<weights.n_rows; i++) {
    weights(i) = 1.0 - Intcpt * std::exp(-(i)*(_TR/_mbg.NGroups())*std::log(2.0)/Thl);
  }
  return(weights);  
} EddyCatch  

//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class EC_TimeConst2Weights
//
// This class translates an EC time constant to weights
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
arma::colvec EC_TimeConst2Weights::GetWeightsForPrevious(double Thl) const EddyTry
{
  arma::colvec prev(_mbg.NGroups());
  arma::colvec curr(_mbg.NGroups());
  this->calculate_weights(Thl,prev,curr);
  return(prev);
} EddyCatch

arma::colvec EC_TimeConst2Weights::GetWeightsForCurrent(double Thl) const EddyTry
{
  arma::colvec prev(_mbg.NGroups());
  arma::colvec curr(_mbg.NGroups());
  this->calculate_weights(Thl,prev,curr);
  return(curr);  
} EddyCatch

void EC_TimeConst2Weights::calculate_weights(double Thl,
					     arma::colvec& prev,
					     arma::colvec& curr) const EddyTry
{
  arma::colvec gtimes(_mbg.NGroups());       // Times at which diffusion gradients are played out
  arma::colvec otimes(2*_mbg.NGroups());     // Times at which eddy currents are observed
  arma::colvec all(2*_mbg.NGroups(),0.0);    // The effects at all times in order [curr prev]
  for (int i=0; i<gtimes.n_rows; i++) gtimes(i) = i*(_TR/_mbg.NGroups());
  for (int i=0; i<otimes.n_rows; i++) otimes(i) = i*(_TR/_mbg.NGroups());
  for (int i=0; i<otimes.n_rows; i++) {
    for (int j=0; j<gtimes.n_rows; j++) {
      if (gtimes(j)-otimes(i) > 1e-6) { // Only count contribution from gradients in the past
	break;
      }
      all(i) += std::exp(-(otimes(i)-gtimes(j))*std::log(2.0)/Thl);
    }
  }
  // Normalise the weights so that they are 1 for the end of current volume
  auto nf = all(_mbg.NGroups()-1);
  std::transform(all.begin(),all.end(),all.begin(),[nf](auto& c){return(c/nf);});
  // Finally divide it up into prev and curr
  curr = all.head_rows(_mbg.NGroups());
  prev = all.tail_rows(_mbg.NGroups());
  return;
} EddyCatch

//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class IndividualWeightsModel
//
// This class models the weights directly on a volume-by-volume
// basis.
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

IndividualWeightsModel::IndividualWeightsModel(ECScanManager& sm) EddyTry
: _nscans(sm.NScans()), _ngroups(sm.MultiBand().NGroups()), 
_prev(sm.NScans(),arma::colvec(sm.MultiBand().NGroups(),arma::fill::zeros)), 
_curr(sm.NScans(),arma::colvec(sm.MultiBand().NGroups(),arma::fill::ones)), 
_ll(sm.NScans(),1e-6) // Start with very little regularisation
{
  if (!sm.IsDiffusion()) throw EddyException("IndividualWeightsModel::IndividualWeightsModel: sm is not diffusion");

  for (unsigned int s=0; s<sm.NScans(); s++) {
    sm.Scan(s).set_previous(_prev[s]);
    sm.Scan(s).set_current(_curr[s]);
  }
} EddyCatch

double IndividualWeightsModel::GetDerivParam(const EDDY::ECScan& sc, 
					     unsigned int        i,
					     unsigned int        s) const EddyTry
{
  if (s >= _nscans) throw EddyException("IndividualWeightsModel::GetDerivParam: s is out of range");
  if (i >= NDerivs()) throw EddyException("IndividualWeightsModel::GetDerivParam: i is out of range");
  if (i < _ngroups) return(sc.GetWeightsForPrevious()(i));
  else return(sc.GetWeightsForCurrent()(i-_ngroups));
} EddyCatch

double IndividualWeightsModel::GetDerivScale(unsigned int i) const EddyTry 
{
  if (i >= NDerivs()) throw EddyException("IndividualWeightsModel::GetDerivScale: i is out of range");
  return(1e-5); 
} EddyCatch

void IndividualWeightsModel::SetDerivParam(EDDY::ECScan&       sc, 
					   unsigned int        i,
					   unsigned int        s,  // Not needed, but checked anyway
					   double              val) const EddyTry
{
  if (i >= NDerivs()) throw EddyException("IndividualWeightsModel::SetDerivParam: i is out of range");
  if (s >= _nscans) throw EddyException("IndividualWeightsModel::SetDerivParam: s is out of range");
  if (i < _ngroups) {
    arma::colvec p = sc.GetWeightsForPrevious();
    p(i) = val;
    sc.set_previous(p);
  }
  else {
    arma::colvec p = sc.GetWeightsForCurrent();
    p(i-_ngroups) = val;
    sc.set_current(p);
  }
  return;
} EddyCatch

double IndividualWeightsModel::GetLevenbergLambda(unsigned int s) const EddyTry {
  if (s >= _nscans) throw EddyException("IndividualWeightsModel::GetLevenbergLambda: s is out of range");
  return(_ll[s]);
} EddyCatch

void IndividualWeightsModel::SetLevenbergLambda(unsigned int s, double lambda) EddyTry {
  if (s >= _nscans) throw EddyException("IndividualWeightsModel::SetLevenbergLambda: s is out of range");
  _ll[s] = lambda;
  return;
} EddyCatch

std::vector<arma::colvec> IndividualWeightsModel::GetParameters(const ECScanManager& sm) const EddyTry
{
  // First verify that the parameters are consistent
  for (unsigned int s=0; s<sm.NScans(); s++) {
    if (!model_and_scan_agree(sm.Scan(s),s)) {
      char message[256]; sprintf(message,"IndividualWeightsModel::GetParameters(sm): mismatch between model and scan %u",s);
      throw EddyException(std::string(message));
    }
  }
  std::vector<arma::colvec> rval(_nscans);
  for (unsigned int s=0; s<_nscans; s++) rval[s] = arma::join_cols(_prev[s],_curr[s]);
  return(rval);
} EddyCatch

arma::colvec IndividualWeightsModel::GetParameters(const ECScan& sc, unsigned int s) const EddyTry
{
  // First verify that the parameters are consistent
  if (s >= _nscans) throw EddyException("IndividualWeightsModel::GetParameters(sc,s): s out of bounds");
  if (!model_and_scan_agree(sc,s)) {
    char message[256]; sprintf(message,"IndividualWeightsModel::GetParameters(sc,s): mismatch between model and scan %u",s);
    throw EddyException(std::string(message));
  }
  return(arma::join_cols(_prev[s],_curr[s]));
} EddyCatch

void IndividualWeightsModel::UpdateParameters(ECScanManager& sm, const std::vector<arma::colvec>& update) EddyTry
{
  // First check agreement between model and update
  if (update.size() != _nscans || update[0].n_rows != NDerivs()) {
    throw EddyException("IndividualWeightsModel::UpdateParameters(sm,update): size mismatch between update and model");
  }
  // Next check agreement between model and scans
  if (sm.NScans() != _nscans) throw EddyException("IndividualWeightsModel::UpdateParameters(sm,update): size mismatch between sm.NScans() and _nscans");
  for (unsigned int s=0; s<sm.NScans(); s++) {
    if (!model_and_scan_agree(sm.Scan(s),s)) {
      char message[256]; sprintf(message,"IndividualWeightsModel::UpdateParameters(sm,update): mismatch between model and scan %u",s);
      throw EddyException(std::string(message));
    }
  }
  // Now update parameters
  for (unsigned int s=0; s<sm.NScans(); s++) {
    _prev[s] += update[s].head(_ngroups);
    _curr[s] += update[s].tail(_ngroups);
    sm.Scan(s).set_previous(_prev[s]);
    sm.Scan(s).set_current(_curr[s]);
  }
  return;
} EddyCatch

void IndividualWeightsModel::UpdateParameters(ECScan& sc, unsigned int s, const arma::colvec& update) EddyTry
{
  // First check agreement between model and update
  if (update.n_rows != NDerivs()) {
    throw EddyException("IndividualWeightsModel::UpdateParameters(sc,s,update): size mismatch between update and model");
  }
  if (s >= _nscans) throw EddyException("IndividualWeightsModel::UpdateParameters(sc,s,update): s out of bounds");
  // Next check agreement between model and scans
  if (!model_and_scan_agree(sc,s)) {
    char message[256]; sprintf(message,"IndividualWeightsModel::UpdateParameters(sc,s,update): mismatch between model and scan %u",s);
    throw EddyException(std::string(message));
  }
  // Now update parameters
  _prev[s] += update.head(_ngroups);
  _curr[s] += update.tail(_ngroups);
  sc.set_previous(_prev[s]);
  sc.set_current(_curr[s]);
  
  return;
} EddyCatch

nlohmann::json IndividualWeightsModel::WriteReport(const ECScanManager& sm, const std::string& fname, bool write_old_style_file) const EddyTry
{
  if (_nscans != sm.NScans()) throw EddyException("IndividualWeightsModel::WriteReport: Mismatch between _nscans and sm.NScans()");
  nlohmann::json long_ec_json;
  if (write_old_style_file) {
    try {
      std::ofstream fout;
      fout.open(fname.c_str(),std::ios::out|std::ios::trunc);
      fout << "Long time constant EC report for individual weights model. Two rows per volume." << std::endl; 
      fout << "The first row gives the weights for the EC-field from the previous diffusion gradient." << std::endl;
      fout << "The second row gives the weights for the EC-field from the current diffusion gradient." << std::endl;
      fout << "Each column represents the weight for one MB-group" << std::endl;
      for (unsigned int s=0; s<_nscans; s++) {
	arma::colvec par = this->GetParameters(sm.Scan(s),s);
	arma::colvec previous = par.head_rows(_ngroups);
	for (unsigned int c=0; c<_ngroups; c++) fout << previous(c) << " ";
	fout << std::endl;
	arma::colvec current = par.tail_rows(_ngroups);
	for (unsigned int c=0; c<_ngroups; c++) fout << current(c) << " ";
	fout << std::endl;
      }
      fout.close();
    }
    catch(...) { throw EddyException("IndividualWeightsModel::WriteReport: Error when attempting to write file " + fname); }
  }
  std::vector<nlohmann::json> volumes(_nscans);
  for (unsigned int s=0; s<_nscans; s++) {
    volumes[s]["Volume #"] = s;
    arma::colvec par = this->GetParameters(sm.Scan(s),s);
    volumes[s]["Previous"] = arma::conv_to<std::vector<double> >::from(par.head_rows(_ngroups));
    volumes[s]["Current"] = arma::conv_to<std::vector<double> >::from(par.tail_rows(_ngroups));
  }
  long_ec_json["Model"] = "Individual weights";
  long_ec_json["Format"] = "Long time-constant eddy currents was modelled as two sets of weights per volume,\nthe first is the weights of the EC field from the previous diffusion gradient\n and the second is the weights for the EC field for the current diffusion gradient.\nThe weights pertains to the MB-groups.";
  long_ec_json["Volumes"] = volumes;

  return(long_ec_json);
} EddyCatch

bool IndividualWeightsModel::model_and_scan_agree(const ECScan& sc, unsigned int s) const EddyTry
{
  return(!(arma::any(sc.GetWeightsForPrevious() != _prev[s]) || arma::any(sc.GetWeightsForCurrent() != _curr[s])));
} EddyCatch

//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class IndividualTimeConstantsModel
//
// This class models the time constants (half-life in seconds) 
// on a volume-by-volume basis.
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

IndividualTimeConstantsModel::IndividualTimeConstantsModel(ECScanManager& sm) EddyTry
: _nscans(sm.NScans()), _tai2w(sm.RepetitionTime(),sm.MultiBand()), _par(sm.NScans(),arma::colvec(4,arma::fill::none)),
_ll(sm.NScans(),0.1) // Start with a fair bit of regularisation 
{
  if (!sm.IsDiffusion()) throw EddyException("IndividualTimeConstantsModel::IndividualTimeConstantsModel: sm is not diffusion");

  for (unsigned int s=0; s<_nscans; s++) {
    _par[s](0) = TC_INITIAL_VALUE;
    _par[s](1) = INTCPT_INITIAL_VALUE;
    _par[s](2) = TC_INITIAL_VALUE;
    _par[s](3) = INTCPT_INITIAL_VALUE;
  }
  arma::colvec prev = _tai2w.GetWeightsForPrevious(_par[0](0),_par[0](1));
  arma::colvec curr = _tai2w.GetWeightsForCurrent(_par[0](2),_par[0](3));

  std::cout << "IndividualTimeConstantsModel::IndividualTimeConstantsModel" << std::endl;
  std::cout << "Previous:" << std::endl;
  for (unsigned int i=0; i<prev.n_rows; i++) std::cout << prev(i) << std::endl;
  std::cout << "Current:" << std::endl;
  for (unsigned int i=0; i<prev.n_rows; i++) std::cout << curr(i) << std::endl;

  for (unsigned int s=0; s<_nscans; s++) {
    sm.Scan(s).set_previous(prev);
    sm.Scan(s).set_current(curr);
  }
} EddyCatch

double IndividualTimeConstantsModel::GetDerivParam(const EDDY::ECScan& sc, unsigned int i, unsigned int s) const EddyTry
{
  if (i >= NDerivs()) throw EddyException("IndividualTimeConstantsModel::GetDerivParam: i out of bounds");
  if (s >= _nscans) throw EddyException("IndividualTimeConstantsModel::GetDerivParam: s out of bounds");
  return(_par[s](i));
} EddyCatch

double IndividualTimeConstantsModel::GetDerivScale(unsigned int i) const EddyTry 
{
  if (i >= NDerivs()) throw EddyException("IndividualTimeConstantsModel::GetDerivScale: i out of bounds");
  if (i==0 || i==2) return(1e-3); // If it is half life
  return(1e-4);                   // If it is intercept 
  // return(1e-3);
} EddyCatch

void IndividualTimeConstantsModel::SetDerivParam(ECScan& sc, unsigned int i, unsigned int s, double val) const
{
  if (i >= NDerivs()) throw EddyException("IndividualTimeConstantsModel::SetDerivParam: i out of bounds");
  if (s >= _nscans) throw EddyException("IndividualTimeConstantsModel::SetDerivParam: s out of bounds");
  if (i<2) {
    if (i==0) sc.set_previous(_tai2w.GetWeightsForPrevious(val,_par[s](1)));
    else sc.set_previous(_tai2w.GetWeightsForPrevious(_par[s](0),val));
  }
  else {
    if (i==2) sc.set_current(_tai2w.GetWeightsForCurrent(val,_par[s](3)));
    else sc.set_current(_tai2w.GetWeightsForCurrent(_par[s](2),val));
  }
}

double IndividualTimeConstantsModel::GetLevenbergLambda(unsigned int s) const EddyTry {
  if (s >= _nscans) throw EddyException("IndividualTimeConstantsModel::GetLevenbergLambda: s is out of range");
  return(_ll[s]);
} EddyCatch

void IndividualTimeConstantsModel::SetLevenbergLambda(unsigned int s, double lambda) EddyTry {
  if (s >= _nscans) throw EddyException("IndividualTimeConstantsModel::SetLevenbergLambda: s is out of range");
  _ll[s] = lambda;
  return;
} EddyCatch

std::vector<arma::colvec> IndividualTimeConstantsModel::GetParameters(const ECScanManager& sm) const EddyTry
{
  // First verify that the parameters are consistent
  for (unsigned int s=0; s<sm.NScans(); s++) {
    if (!model_and_scan_agree(sm.Scan(s),s)) {
      char message[256]; sprintf(message,"IndividualTimeConstantsModel::GetParameters(sm): mismatch between model and scan %u",s);
      throw EddyException(std::string(message));
    }
  }
  return(_par);
} EddyCatch

arma::colvec IndividualTimeConstantsModel::GetParameters(const ECScan& sc, unsigned int s) const EddyTry
{
  // First verify that the parameters are consistent
  if (s >= _nscans) throw EddyException("IndividalTimeConstantsModel::GetParameters(sc,s): s out of bounds");
  if (!model_and_scan_agree(sc,s)) {
    char message[256]; sprintf(message,"IndividualTimeConstantsModel::GetParameters(sc,s): mismatch between model and scan %u",s);
    throw EddyException(std::string(message));
  }
  return(_par[s]);
} EddyCatch

void IndividualTimeConstantsModel::UpdateParameters(ECScanManager& sm, const std::vector<arma::colvec>& update) EddyTry
{
  // First check agreement between model and update
  if (update.size() != _nscans || update[0].n_rows != NDerivs()) {
    throw EddyException("IndividualTimeConstantsModel::UpdateParameters(sm,update): size mismatch between update and model");
  }
  // Next check agreement between model and scans
  if (sm.NScans() != _nscans) throw EddyException("IndividualTimeConstantsModel::UpdateParameters(sm,update): size mismatch between sm.NScans() and _nscans");
  for (unsigned int s=0; s<sm.NScans(); s++) {
    if (!model_and_scan_agree(sm.Scan(s),s)) {
      char message[256]; sprintf(message,"IndividalTimeConstantsModel::UpdateParameters(sm,update): mismatch between model and scan %u",s);
      throw EddyException(std::string(message));
    }
  }
  // Now update parameters
  for (unsigned int s=0; s<sm.NScans(); s++) {
    for (unsigned int i=0; i<NDerivs(); i++) _par[s](i) += update[s](i);
    sm.Scan(s).set_previous(_tai2w.GetWeightsForPrevious(_par[s](0),_par[s](1)));
    sm.Scan(s).set_current(_tai2w.GetWeightsForCurrent(_par[s](2),_par[s](3)));    
  }
  return;
} EddyCatch

void IndividualTimeConstantsModel::UpdateParameters(ECScan& sc, unsigned int s, const arma::colvec& update) EddyTry
{
  // First check agreement between model and update
  if (update.n_rows != NDerivs()) {
    throw EddyException("IndividalTimeConstantsModel::UpdateParameters(sc,s,update): size mismatch between update and model");
  }
  if (s >= _nscans) throw EddyException("IndividalTimeConstantsModel::UpdateParameters(sc,s,update): s out of bounds");
  // Next check agreement between model and scans
  if (!model_and_scan_agree(sc,s)) {
    char message[256]; sprintf(message,"IndividalTimeConstantsModel::UpdateParameters(sc,s,update): mismatch between model and scan %u",s);
    throw EddyException(std::string(message));
  }
  // Now update parameters
  for (unsigned int i=0; i<NDerivs(); i++) _par[s](i) += update(i);
  sc.set_previous(_tai2w.GetWeightsForPrevious(_par[s](0),_par[s](1)));
  sc.set_current(_tai2w.GetWeightsForCurrent(_par[s](2),_par[s](3)));    
  
  return;
} EddyCatch

nlohmann::json IndividualTimeConstantsModel::WriteReport(const ECScanManager& sm, const std::string& fname, bool write_old_style_file) const EddyTry
{
  if (_nscans != sm.NScans()) throw EddyException("IndividualTimeConstantsModel::WriteReport: Mismatch between _nscans and sm.NScans()");
  if (write_old_style_file) {
    try {
      std::ofstream fout;
      fout.open(fname.c_str(),std::ios::out|std::ios::trunc);
      fout << "Long time constant EC report for individual time constants model. Three rows per volume." << std::endl; 
      fout << "The first row gives [TC Intercept] for previous, followed by [TC Intercept] for current." << std::endl;
      fout << "The second row gives the weights for the EC-field from the previous diffusion gradient." << std::endl;
      fout << "The third row gives the weights for the EC-field from the current diffusion gradient." << std::endl;
      fout << "Each column of rows two and three represents the weight for one MB-group" << std::endl;
      std::vector<arma::colvec> tc = this->GetParameters(sm);
      for (unsigned int s=0; s<_nscans; s++) {
	fout << tc[s](0) << " " << tc[s](1) << " " << tc[s](2) << " " << tc[s](3) << std::endl; 
	arma::colvec previous = sm.Scan(s).GetWeightsForPrevious();
	for (unsigned int c=0; c<_tai2w.NGroups(); c++) fout << previous(c) << " ";
	fout << std::endl;
	arma::colvec current = sm.Scan(s).GetWeightsForCurrent();
	for (unsigned int c=0; c<_tai2w.NGroups(); c++) fout << current(c) << " ";
	fout << std::endl;
      }
      fout.close();      
    }
    catch(...) { throw EddyException("IndividualTimeConstantsModel::WriteReport: Error when attempting to write file " + fname); }
  }
  nlohmann::json long_ec_json;
  std::vector<nlohmann::json> volumes(_nscans);
  for (unsigned int s=0; s<_nscans; s++) {
    volumes[s]["Volume #"] = s;
    volumes[s]["Half-life (s) Intercept for previous"] = arma::conv_to<std::vector<double> >::from(this->GetParameters(sm.Scan(s),s).rows(0,1));
    volumes[s]["Half-life (s) Intercept for current"] = arma::conv_to<std::vector<double> >::from(this->GetParameters(sm.Scan(s),s).rows(2,3));
    volumes[s]["Resulting weights for previous"] = arma::conv_to<std::vector<double> >::from(sm.Scan(s).GetWeightsForPrevious());
    volumes[s]["Resulting weights for current"] = arma::conv_to<std::vector<double> >::from(sm.Scan(s).GetWeightsForCurrent());	
  }
  long_ec_json["Model"] = "Individual time-constants";
  long_ec_json["Format"] = "Long time-constant eddy currents was modelled as one time-constant (half-life in seconds) \nand one intercept for previous and current, resulting in four parameters per volume.\nThe resulting weights for the MB-groups are also given.";
  long_ec_json["Volumes"] = volumes;
  
  return(long_ec_json);
} EddyCatch

bool IndividualTimeConstantsModel::model_and_scan_agree(const ECScan& sc, unsigned int s) const EddyTry
{
  return(!(arma::any(sc.GetWeightsForPrevious() != _tai2w.GetWeightsForPrevious(_par[s](0),_par[s](1))) || 
	   arma::any(sc.GetWeightsForCurrent() != _tai2w.GetWeightsForCurrent(_par[s](2),_par[s](3)))));
} EddyCatch

//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class JointWeightsModel
//
// This class models the weights jointly across all volumes.
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

JointWeightsModel::JointWeightsModel(ECScanManager& sm) EddyTry
: _nscans(sm.NScans()), _ngroups(sm.MultiBand().NGroups()), 
_prev(arma::colvec(sm.MultiBand().NGroups(),arma::fill::zeros)), _curr(arma::colvec(sm.MultiBand().NGroups(),arma::fill::ones)),
_ll(1e-6) // Start with very little regularisation
{
  if (!sm.IsDiffusion()) throw EddyException("JointWeightsModel::JointWeightsModel: sm is not diffusion");

  for (unsigned int s=0; s<sm.NScans(); s++) {
    sm.Scan(s).set_previous(_prev);
    sm.Scan(s).set_current(_curr);
  }
} EddyCatch

double JointWeightsModel::GetDerivParam(const EDDY::ECScan& sc, 
					unsigned int        i,
					unsigned int        s) const EddyTry
{
  if (s >= _nscans) throw EddyException("JointWeightsModel::GetDerivParam: s is out of range");
  if (i >= NDerivs()) throw EddyException("JointWeightsModel::GetDerivParam: i is out of range");
  if (i < _ngroups) return(sc.GetWeightsForPrevious()(i));
  else return(sc.GetWeightsForCurrent()(i-_ngroups));
} EddyCatch

double JointWeightsModel::GetDerivScale(unsigned int i) const EddyTry 
{
  if (i >= NDerivs()) throw EddyException("JointWeightsModel::GetDerivScale: i is out of range");
  return(1e-5); 
} EddyCatch

void JointWeightsModel::SetDerivParam(EDDY::ECScan&       sc, 
				      unsigned int        i,
				      unsigned int        s,  // Not needed, but checked anyway
				      double              val) const EddyTry
{
  if (i >= NDerivs()) throw EddyException("JointWeightsModel::GetDerivParam: i is out of range");
  if (s >= _nscans) throw EddyException("JointWeightsModel::GetDerivParam: s is out of range");
  if (i < _ngroups) {
    arma::colvec p = sc.GetWeightsForPrevious();
    p(i) = val;
    sc.set_previous(p);
  }
  else {
    arma::colvec p = sc.GetWeightsForCurrent();
    p(i-_ngroups) = val;
    sc.set_current(p);
  }
  return;
} EddyCatch

double JointWeightsModel::GetLevenbergLambda(unsigned int s) const EddyTry {
  if (s >= _nscans) throw EddyException("JointWeightsModel::GetLevenbergLambda: s is out of range"); // Not needed, but do anyway
  return(_ll);
} EddyCatch

void JointWeightsModel::SetLevenbergLambda(unsigned int s, double lambda) EddyTry {
  if (s >= _nscans) throw EddyException("JointWeightsModel::SetLevenbergLambda: s is out of range"); // Not needed, but do anyway
  _ll = lambda;
  return;
} EddyCatch

std::vector<arma::colvec> JointWeightsModel::GetParameters(const ECScanManager& sm) const EddyTry
{
  // First verify that the parameters are consistent
  for (unsigned int s=0; s<sm.NScans(); s++) {
    if (!model_and_scan_agree(sm.Scan(s),s)) {
      char message[256]; sprintf(message,"JointWeightsModel::GetParameters(sm): mismatch between model and scan %u",s);
      throw EddyException(std::string(message));
    }
  }
  std::vector<arma::colvec> rval(1);
  rval[0] = arma::join_cols(_prev,_curr);
  return(rval);
} EddyCatch

arma::colvec JointWeightsModel::GetParameters(const ECScan& sc, unsigned int s) const EddyTry
{
  // First verify that the parameters are consistent
  if (s >= _nscans) throw EddyException("JointWeightsModel::GetParameters(sc,s): s out of bounds");
  if (!model_and_scan_agree(sc,s)) {
    char message[256]; sprintf(message,"JointWeightsModel::GetParameters(sc,s): mismatch between model and scan %u",s);
    throw EddyException(std::string(message));
  }
  return(arma::join_cols(_prev,_curr));
} EddyCatch

void JointWeightsModel::UpdateParameters(ECScanManager& sm, const std::vector<arma::colvec>& update) EddyTry
{
  // First check agreement between model and update
  if (update.size() != 1 || update[0].n_rows != NDerivs()) {
    throw EddyException("JointWeightsModel::UpdateParameters(sm,update): size mismatch between update and model");
  }
  // Next check agreement between model and scans
  if (sm.NScans() != _nscans) throw EddyException("JointWeightsModel::UpdateParameters(sm,update): size mismatch between sm.NScans() and _nscans");
  for (unsigned int s=0; s<sm.NScans(); s++) {
    if (!model_and_scan_agree(sm.Scan(s),s)) {
      char message[256]; sprintf(message,"JointWeightsModel::UpdateParameters(sm,update): mismatch between model and scan %u",s);
      throw EddyException(std::string(message));
    }
  }
  // Now update parameters
  _prev += update[0].head(_ngroups);
  _curr += update[0].tail(_ngroups);
  for (unsigned int s=0; s<sm.NScans(); s++) {
    sm.Scan(s).set_previous(_prev);
    sm.Scan(s).set_current(_curr);
  }
  return;
} EddyCatch

void JointWeightsModel::UpdateParameters(ECScan& sc, unsigned int s, const arma::colvec& update) EddyTry
{
  throw EddyException("JointWeightsModel::UpdateParameters(sc,s,update): not allowed to update individual scan");
  return;
} EddyCatch

nlohmann::json JointWeightsModel::WriteReport(const ECScanManager& sm, const std::string& fname, bool write_old_style_file) const EddyTry
{
  if (_nscans != sm.NScans()) throw EddyException("JointWeightsModel::WriteReport: Mismatch between _nscans and sm.NScans()");
  if (write_old_style_file) {
    try {
      std::ofstream fout;
      fout.open(fname.c_str(),std::ios::out|std::ios::trunc);
      fout << "Long time constant EC report for joint weights model. Two rows." << std::endl; 
      fout << "The first row gives the weights for the EC-field from the previous diffusion gradient." << std::endl;
      fout << "The second row gives the weights for the EC-field from the current diffusion gradient." << std::endl;
      fout << "Each column represents the weight for one MB-group" << std::endl;
      std::vector<arma::colvec> par = this->GetParameters(sm);
      arma::colvec previous = par[0].head_rows(_ngroups);
      for (unsigned int c=0; c<_ngroups; c++) fout << previous(c) << " ";
      fout << std::endl;
      arma::colvec current = par[0].tail_rows(_ngroups);
      for (unsigned int c=0; c<_ngroups; c++) fout << current(c) << " ";
      fout << std::endl;
      fout.close();
    }
    catch(...) { throw EddyException("JointWeightsModel::WriteReport: Error when attempting to write file " + fname); }
  }
  nlohmann::json long_ec_json;
  long_ec_json["Model"] = "Joint weights";
  long_ec_json["Format"] = "Long time-constant eddy currents was modelled as two sets of weights,\nthe first is the weights of the EC field from the previous diffusion gradient\n and the second is the weights for the EC field for the current diffusion gradient.\nThe weights pertains to the MB-groups and are pooled across all volumes.";
  std::vector<arma::colvec> par = this->GetParameters(sm);
  long_ec_json["Previous"] = arma::conv_to<std::vector<double> >::from(par[0].head_rows(_ngroups));
  long_ec_json["Current"] = arma::conv_to<std::vector<double> >::from(par[0].tail_rows(_ngroups));

  return(long_ec_json);
} EddyCatch

bool JointWeightsModel::model_and_scan_agree(const ECScan& sc, unsigned int s) const EddyTry
{
  return(!(arma::any(sc.GetWeightsForPrevious() != _prev) || arma::any(sc.GetWeightsForCurrent() != _curr)));
} EddyCatch

//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
//
// Class JointTimeConstantModel
//
// This class models the time-constant (half-life in seconds) 
// jointly across all volumes.
//
//@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

JointTimeConstantModel::JointTimeConstantModel(ECScanManager& sm) EddyTry
: _nscans(sm.NScans()), _tai2w(sm.RepetitionTime(),sm.MultiBand()), _par(arma::colvec(4,arma::fill::none)),
_ll(0.1) // Start with a fair bit of regularisation
 
{
  if (!sm.IsDiffusion()) throw EddyException("JointTimeConstantsModel::JointTimeConstantsModel: sm is not diffusion");

  _par(0) = TC_INITIAL_VALUE;
  _par(1) = INTCPT_INITIAL_VALUE;
  _par(2) = TC_INITIAL_VALUE;
  _par(3) = INTCPT_INITIAL_VALUE;
  arma::colvec prev = _tai2w.GetWeightsForPrevious(_par(0),_par(1));
  arma::colvec curr = _tai2w.GetWeightsForCurrent(_par(2),_par(3));

  for (unsigned int s=0; s<sm.NScans(); s++) {
    sm.Scan(s).set_previous(prev);
    sm.Scan(s).set_current(curr);
  }
} EddyCatch

double JointTimeConstantModel::GetDerivParam(const EDDY::ECScan& sc, unsigned int i, unsigned int s) const EddyTry
{
  if (i >= NDerivs()) throw EddyException("JointTimeConstantModel::GetDerivParam: i out of bounds");
  if (s >= _nscans) throw EddyException("JointTimeConstantModel::GetDerivParam: s out of bounds");
  return(_par(i));
} EddyCatch

double JointTimeConstantModel::GetDerivScale(unsigned int i) const EddyTry 
{
  if (i >= NDerivs()) throw EddyException("JointTimeConstantModel::GetDerivScale: i out of bounds");
  return(1e-5); 
} EddyCatch

void JointTimeConstantModel::SetDerivParam(ECScan& sc, unsigned int i, unsigned int s, double val) const
{
  if (i >= NDerivs()) throw EddyException("JointTimeConstantModel::SetDerivParam: i out of bounds");
  if (s >= _nscans) throw EddyException("JointTimeConstantModel::SetDerivParam: s out of bounds");
  if (i < 2) {
    if (i==0) sc.set_previous(_tai2w.GetWeightsForPrevious(val,_par(1)));
    else sc.set_previous(_tai2w.GetWeightsForPrevious(_par(0),val));
  }
  else {
    if (i==2) sc.set_current(_tai2w.GetWeightsForCurrent(val,_par(3)));
    else sc.set_current(_tai2w.GetWeightsForCurrent(_par(2),val));
  }
}

double JointTimeConstantModel::GetLevenbergLambda(unsigned int s) const EddyTry {
  if (s >= _nscans) throw EddyException("JointTimeConstantModel::GetLevenbergLambda: s is out of range"); // Not needed, but do anyway
  return(_ll);
} EddyCatch

void JointTimeConstantModel::SetLevenbergLambda(unsigned int s, double lambda) EddyTry {
  if (s >= _nscans) throw EddyException("JointTimeConstantModel::SetLevenbergLambda: s is out of range"); // Not needed, but do anyway
  _ll = lambda;
  return;
} EddyCatch

std::vector<arma::colvec> JointTimeConstantModel::GetParameters(const ECScanManager& sm) const EddyTry
{
  // First verify that the parameters are consistent
  for (unsigned int s=0; s<sm.NScans(); s++) {
    if (!model_and_scan_agree(sm.Scan(s),s)) {
      char message[256]; sprintf(message,"JointTimeConstantModel::GetParameters(sm): mismatch between model and scan %u",s);
      throw EddyException(std::string(message));
    }
  }
  return(std::vector<arma::colvec>(1,_par));
} EddyCatch

arma::colvec JointTimeConstantModel::GetParameters(const ECScan& sc, unsigned int s) const EddyTry
{
  // First verify that the parameters are consistent
  if (s >= _nscans) throw EddyException("JointTimeConstantModel::GetParameters(sc,s): s out of bounds");
  if (!model_and_scan_agree(sc,s)) {
    char message[256]; sprintf(message,"JointTimeConstantModel::GetParameters(sc,s): mismatch between model and scan %u",s);
    throw EddyException(std::string(message));
  }
  return(_par);
} EddyCatch

void JointTimeConstantModel::UpdateParameters(ECScanManager& sm, const std::vector<arma::colvec>& update) EddyTry
{
  // First check agreement between model and update
  if (update.size() != 1 || update[0].n_rows != NDerivs()) {
    throw EddyException("JointTimeConstantModel::UpdateParameters(sm,update): size mismatch between update and model");
  }
  // Next check agreement between model and scans
  if (sm.NScans() != _nscans) throw EddyException("JointTimeConstantModel::UpdateParameters(sm,update): size mismatch between sm.NScans() and _nscans");
  for (unsigned int s=0; s<sm.NScans(); s++) {
    if (!model_and_scan_agree(sm.Scan(s),s)) {
      char message[256]; sprintf(message,"JointTimeConstantModel::UpdateParameters(sm,update): mismatch between model and scan %u",s);
      throw EddyException(std::string(message));
    }
  }
  // Now update parameters
  for (unsigned int i=0; i<NDerivs(); i++) _par(i) += update[0](i);
  for (unsigned int s=0; s<sm.NScans(); s++) {
    sm.Scan(s).set_previous(_tai2w.GetWeightsForPrevious(_par(0),_par(1)));
    sm.Scan(s).set_current(_tai2w.GetWeightsForCurrent(_par(2),_par(3)));
  }
  return;
} EddyCatch

void JointTimeConstantModel::UpdateParameters(ECScan& sc, unsigned int s, const arma::colvec& update) EddyTry
{
  throw EddyException("JointTimeConstantModel::UpdateParameters(sc,s,update): not allowed to update individual scan");
  return;
} EddyCatch

nlohmann::json JointTimeConstantModel::WriteReport(const ECScanManager& sm, const std::string& fname, bool write_old_style_file) const EddyTry
{
  if (_nscans != sm.NScans()) throw EddyException("JointTimeConstantModel::WriteReport: Mismatch between _nscans and sm.NScans()");
  if (write_old_style_file) {
    try {
      std::ofstream fout;
      fout.open(fname.c_str(),std::ios::out|std::ios::trunc);
      fout << "Long time constant EC report for joint time constant model. Three rows." << std::endl; 
      fout << "The first row gives [TC Intercept] for previous, followed by [TC Intercept] for current." << std::endl;
      fout << "The second row gives the resulting weights for the EC-field from the previous diffusion gradient." << std::endl;
      fout << "The third row gives the resulting weights for the EC-field from the current diffusion gradient." << std::endl;
      fout << "Each column represents the weight for one MB-group" << std::endl;
      std::vector<arma::colvec> tc = this->GetParameters(sm);
      fout << tc[0](0) << " " << tc[0](1) << " " << tc[0](2) << " " << tc[0](3) << std::endl; 
      arma::colvec previous = sm.Scan(0).GetWeightsForPrevious();
      for (unsigned int c=0; c<_tai2w.NGroups(); c++) fout << previous(c) << " ";
      fout << std::endl;
      arma::colvec current = sm.Scan(0).GetWeightsForCurrent();
      for (unsigned int c=0; c<_tai2w.NGroups(); c++) fout << current(c) << " ";
      fout << std::endl;
      fout.close();
    }
    catch(...) { throw EddyException("JointTimeConstantModel::WriteReport: Error when attempting to write file " + fname); }
  }
  nlohmann::json long_ec_json;
  long_ec_json["Model"] = "Joint time-constant";
  long_ec_json["Format"] = "Long time-constant eddy currents was modelled as a single time-constant (half-life in seconds) across all volumes.\nThe resulting weights for the MB-groups are also given.";      
  long_ec_json["Half-life (s) Intercept for previous"] = arma::conv_to<std::vector<double> >::from(this->GetParameters(sm)[0].rows(0,1));
  long_ec_json["Half-life (s) Intercept for current"] = arma::conv_to<std::vector<double> >::from(this->GetParameters(sm)[0].rows(2,3));
  long_ec_json["Resulting weights for previous"] = arma::conv_to<std::vector<double> >::from(sm.Scan(0).GetWeightsForPrevious());
  long_ec_json["Resulting weights for current"] = arma::conv_to<std::vector<double> >::from(sm.Scan(0).GetWeightsForCurrent());	

  return(long_ec_json);
} EddyCatch

bool JointTimeConstantModel::model_and_scan_agree(const ECScan& sc, unsigned int s) const EddyTry
{
  return(!(arma::any(sc.GetWeightsForPrevious() != _tai2w.GetWeightsForPrevious(_par(0),_par(1))) || 
	   arma::any(sc.GetWeightsForCurrent() != _tai2w.GetWeightsForCurrent(_par(2),_par(3)))));
} EddyCatch

nlohmann::json NoLongECModel::WriteReport(const ECScanManager& sm, const std::string& fname, bool write_old_style_file) const EddyTry {
  nlohmann::json long_ec_json;
  long_ec_json["Model"] = "None";
  long_ec_json["Format"] = "No modelling of long time-constant eddy currents was performed.";
  return(long_ec_json);
} EddyCatch