//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SCALAR_REMQUO_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SCALAR_REMQUO_HPP_INCLUDED
#include <boost/simd/toolbox/arithmetic/functions/remquo.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#include <boost/fusion/tuple.hpp>
#include <boost/mpl/vector.hpp>
#include <boost/simd/include/functions/scalar/remainder.hpp>
#include <boost/simd/include/functions/scalar/idivround.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::remquo_, tag::cpu_,
                       (A0),
                       (scalar_ < floating_<A0> > )
                       (scalar_ < floating_<A0> > )
                       )
  {
    typedef typename dispatch::meta::as_integer<A0,signed>::type           quo;
    typedef boost::fusion::tuple<A0,quo>                           result_type;

    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      result_type res;
      boost::simd::remquo(a0, a1, boost::fusion::at_c<0>(res), boost::fusion::at_c<1>(res));
      return res;
    }

  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::remquo_, tag::cpu_,
                       (A0)(A1),
                       (scalar_ < floating_<A0> > )
                       (scalar_ < floating_<A0> > )
                       (scalar_ < floating_<A0> > )
                       (scalar_ < integer_<A1> > )
                       )
  {
    typedef boost::simd::int32_t result_type;
    inline result_type operator()(A0 const& a0,const A0 & a1,A0 & a2,A1 & a3) const
    {
      a2 = boost::simd::remainder(a0, a1);
      a3 = boost::simd::idivround(a0, a1);
      return 0;
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::remquo_, tag::cpu_,
                       (A0)(A1),
                       (scalar_ < integer_<A0> > )
                       (scalar_ < integer_<A0> > )
                       (scalar_ < floating_<A1> > )
                       (scalar_ < integer_<A0> > )
                       )
  {
    typedef boost::simd::int32_t result_type;
    inline result_type operator()(A0 const& a0,const A0 & a1,A1 & a2,A0 & a3) const
    {
      return remquo(A1(a0), A1(a1), a2, a3);
    }
  };
} } }

#endif


#include "path_guiding.hxx"
#include "scene.hxx"

#ifdef HAVE_JSON
#include "json.hxx"
#include "rapidjson/document.h"
#include "rapidjson/ostreamwrapper.h"
#include "rapidjson/writer.h"
#endif
#include <fstream>

#include <tbb/parallel_for_each.h>
#include <tbb/combinable.h>
#include <tbb/enumerable_thread_specific.h>
#include <Eigen/Eigenvalues>

//#include <range/v3/view/enumerate.hpp>

#ifdef _MSC_VER
static const char* DEBUG_FILE_PREFIX = "D:\\tmp2\\";
#else
static const char* DEBUG_FILE_PREFIX = "/tmp/";
#endif

#ifdef HAVE_JSON
namespace rapidjson_util {

rapidjson::Value ToJSON(const vmf_fitting::VonMisesFischerMixture<> &mixture, Alloc &a)
{
  using namespace rapidjson_util;
  rj::Value jgmm(rj::kObjectType);
  jgmm.AddMember("weights", ToJSON(mixture.weights, a), a);
  jgmm.AddMember("means", ToJSON(mixture.means, a), a);
  jgmm.AddMember("concentrations", ToJSON(mixture.concentrations, a), a);
  return jgmm;
}

} // rapidjson_util namespace
#endif

namespace guiding
{

void AddToPointStatistics(CellData &cell, Span<IncidentRadiance> buffer)
{
  for (const auto &in : buffer)
  {
    if (in.is_original)
      cell.learned.leaf_stats += in.pos;
  }
}

namespace MoVmfRadianceDistribution
{

RadianceDistributionLearned::Parameters::Parameters(const CellData &cell)
{
  using namespace vmf_fitting;
  // TODO: should be the max over all previous samples counts to prevent
  // overfitting in case the sample count decreases.
  double prior_strength = std::max(1., 0.01*cell.max_num_samples);
  params.prior_alpha = prior_strength;
  params.prior_nu = prior_strength;
  params.prior_tau = prior_strength;
  params.maximization_step_every = std::max(1, static_cast<int>(prior_strength) * 10); //param_em_every;
  //params.prior_mode = &cell.current_estimate.radiance_distribution;

  // Below, I provide a fixed, more or less uniform prior. This seems to yield
  // higher variance than taking the previouly learned distribution as prior.
  vmf_fitting::InitializeForUnitSphere(prior_mode);
  params.prior_mode = &prior_mode;
}

Float3 RadianceDistributionSampled::ComputeStochasticFilteredDirection(const IncidentRadiance & rec, RandGen &sampler) const
{
  const float prob = 0.1f;
  const float pdf = Pdf(rec.reverse_incident_dir.cast<double>().eval());
  const float a = -prob/(pdf*float(2.*Pi) + 1.e-9f) + 1.f;
  const float cos_scatter_open_angle = std::min(std::max(a, -1.f), 1.f);
  Double3 scatter_offset = SampleTrafo::ToUniformSphereSection(cos_scatter_open_angle, sampler.UniformUnitSquare());
  return OrthogonalSystemZAligned(rec.reverse_incident_dir) * scatter_offset.cast<float>();
}

rapidjson::Value RadianceDistributionLearned::ToJSON(rapidjson_util::Alloc &a) const
{
  using namespace rapidjson_util;
  rj::Value jtree;
  return jtree;
}

rapidjson::Value RadianceDistributionSampled::ToJSON(rapidjson_util::Alloc &a) const
{
  using namespace rapidjson_util;
  rj::Value jtree;
  return jtree;
}

} // namespace MoVmfRadianceDistribution

namespace quadtree_radiance_distribution
{

Eigen::Vector2f MapSphereToTree(const Eigen::Vector3f &dir)
{
  const double v = (dir[2] + 1.) * 0.5;
  const double phi = std::atan2(dir[1], dir[0]); // in [-pi,+pi]
  const double u = (phi / Pi + 1.) * 0.5;
  return { (float)u, (float)v };
}

Eigen::Vector3f MapTreeToSphere(const Eigen::Vector2f &uv)
{
  const double z = uv[1] * 2. - 1.;
  const double rho = std::sqrt(std::max(0., 1. - z*z));
  const double phi = (uv[0] * 2. - 1.) * Pi;
  const double x = rho*std::cos(phi);
  const double y = rho*std::sin(phi);
  return { (float)x, (float)y, (float)z };
}

Float3 RadianceDistributionSampled::ComputeStochasticFilteredDirection(const IncidentRadiance & rec, RandGen &sampler) const
{
  using namespace guiding::quadtree;

  auto pt = MapSphereToTree(rec.reverse_incident_dir.cast<float>());
  DecentHelper d{pt, this->tree};

  float footprint = .5f; // As in sort of a angular arc.
  for (;!d.IsLeaf();)
  {
    d.TraverseInplace();
    footprint *= 0.5;
  }
  static constexpr float SQRT_4_PI = 3.5449077018110318f;
  footprint *= SQRT_4_PI;
  // Value of pi would mean to rotate the scattered sampled 180 deg.
  // The footprint value will always be smaller. Slightly larger than pi/2 if there is only one root leaf.

  // TODO: float version of UniformSphereSection.
  Double3 scatter_offset = SampleTrafo::ToUniformSphereSection(std::cos(footprint), sampler.UniformUnitSquare());
  return OrthogonalSystemZAligned(rec.reverse_incident_dir) * scatter_offset.cast<float>();
}

inline void PushWeight(
  const quadtree::Tree & tree,
  Accumulators::SoaOnlineVariance<double, long> &node_weights,
  const Eigen::Vector2f &p,
  float w)
{
  quadtree::detail::DecentHelper d{ p, tree };
  for (;;)
  {
    node_weights.Add(d.GetNode().idx, w);
    if (d.IsLeaf())
      break;
    d.TraverseInplace();
  }
}

void RadianceDistributionLearned::IncrementalFit(
  Span<const IncidentRadiance> buffer,
  const Parameters &params,
  RadianceDistributionSampled &dst)
{
  assert(params.round > 0);

  for (auto in : buffer)
  {
    if (in.is_original)
      continue;
    const auto uv = MapSphereToTree(in.reverse_incident_dir);
    PushWeight(tree, node_weights, uv, in.weight);
  }

  Eigen::ArrayXd relative_counts = this->CalcRelativeCounts();
  Eigen::ArrayXd node_means = relative_counts*this->node_weights.Mean();
  const double total_flux = node_means[tree.GetRoot().idx];
  Eigen::ArrayXd node_stddev = relative_counts*this->node_weights.MeanErr((double)LargeFloat, min_sample_count);
  Eigen::ArrayXd mix_factor = good_sample_count / (good_sample_count + (node_weights.Counts().cast<double>()-min_sample_count).max(0.));
  // This stuff is inspired by the UCB algorithm for Multi-Armed-Bandits. It should ensure that, as the number of training passes goes up, we
  // ocassionally re-visit each node and sample from it. It is an answer to the exporation-exploitation problem in reinforcement learning type problems.
  // The difference to the regular UCB algorithm is that the sampling probability is proportional to the upper confidence bounds, instead of sampling the
  // arm with the largest bound.
  dst.node_sample_probs = (dst.node_sample_prior.cast<double>()*mix_factor + (node_means + node_stddev)*(1.-mix_factor)).cast<float>();
  dst.node_sample_probs += total_flux*(std::log((double)params.round) * (node_weights.Counts()+1).cast<double>().cwiseInverse()).sqrt().cast<float>();
  quadtree::PropagateLeafWeightsToParents(tree, AsSpan(dst.node_sample_probs), tree.GetRoot());
  assert(dst.node_sample_probs.allFinite());
}

void RadianceDistributionLearned::InitialFit(Span<const IncidentRadiance> buffer, const Parameters &params)
{
  auto points = TransformVector(buffer, [](const IncidentRadiance &in) -> Eigen::Vector2f { return MapSphereToTree(in.reverse_incident_dir); });
  auto weights = TransformVector(buffer, [](const IncidentRadiance &in) -> float { return in.weight; });

  quadtree::Builder builder{AsSpan(points), AsSpan(weights), 0.1f};
  tree = builder.ExtractTree();

  node_weights = decltype(node_weights)(tree.NumNodes());
  for (int i=0; i<points.size(); ++i)
  {
    PushWeight(tree, node_weights, points[i], weights[i]);
  }

  //fmt::print("initial quadtree fit with {} nodes, root flux {}\n", tree.NumNodes(), node_weights[tree.GetRoot().idx]);
}

RadianceDistributionSampled RadianceDistributionLearned::IterationUpdateAndBake(
  const RadianceDistributionSampled &previous_radiance_dist)
{
  // The adaptation part
  if (node_weights.Counts()[tree.GetRoot().idx]>10)
  {
    const auto old_relative_counts = this->CalcRelativeCounts();
    auto old_flux = (old_relative_counts*node_weights.Mean()).cast<float>().eval();
    quadtree::TreeAdaptor adaptor{tree, AsSpan(old_flux), 0.05};

    tree = adaptor.ExtractTree();

    Accumulators::SoaOnlineVariance<double, long> adapted_stats(tree.NumNodes());
    Eigen::ArrayXf prior(tree.NumNodes());
    Eigen::ArrayXd relative_counts(tree.NumNodes());

    int i=0;
    for (auto nto: adaptor.GetNewToOldMapping())
    {
      auto s = this->node_weights.GetStats(nto.idx);
      float scale = nto.fractional_size*nto.fractional_size;
      s = s.CountsMultiplied(scale);
      adapted_stats.SetStats(i, s);
      prior[i] = scale*previous_radiance_dist.node_sample_probs[nto.idx];
      relative_counts[i] = scale*old_relative_counts[nto.idx];
      ++i;
    }

    this->node_weights = adapted_stats;

    Eigen::ArrayXd mix_factor = good_sample_count / (good_sample_count + (node_weights.Counts().cast<double>()-min_sample_count).max(0.));

    RadianceDistributionSampled dst;
    dst.tree = this->tree;
    dst.node_means = relative_counts*this->node_weights.Mean();
    dst.node_stddev = relative_counts*this->node_weights.MeanErr(safe_error_number, min_sample_count);
    dst.node_sample_probs = (prior.cast<double>()*mix_factor + (dst.node_means + dst.node_stddev)*(1.-mix_factor)).cast<float>();
    quadtree::PropagateLeafWeightsToParents(tree, AsSpan(dst.node_sample_probs), tree.GetRoot());
    dst.node_sample_prior = dst.node_sample_probs;

    this->node_weights = decltype(node_weights)(tree.NumNodes());

    assert(dst.node_means.allFinite());
    assert(dst.node_stddev.allFinite());
    assert(dst.node_sample_probs.allFinite());
    return dst;
  }
  else // Don't bother ...
  {
    return previous_radiance_dist;
  }
}

RadianceDistributionSampled RadianceDistributionLearned::Bake() const
{
  RadianceDistributionSampled dst;
  dst.tree = this->tree;

  auto relative_counts = this->CalcRelativeCounts();
  dst.node_means = relative_counts*this->node_weights.Mean();
  dst.node_stddev = relative_counts*this->node_weights.MeanErr(safe_error_number, 10);
  dst.node_sample_probs = dst.node_means.cast<float>();
  dst.node_sample_prior = dst.node_sample_probs;
  quadtree::PropagateLeafWeightsToParents(tree, AsSpan(dst.node_sample_probs), tree.GetRoot());
  assert(dst.node_means.allFinite());
  assert(dst.node_stddev.allFinite());
  assert(dst.node_sample_probs.allFinite());
  return dst;
}

rapidjson::Value RadianceDistributionLearned::ToJSON(rapidjson_util::Alloc &a) const
{
  using namespace rapidjson_util;
  rj::Value jtree = tree.ToJSON(a);
  //jtree.AddMember("node_weights", rapidjson_util::ToJSON(node_sample_probs, a), a);
  jtree.AddMember("node_means", rapidjson_util::ToJSON(node_weights.Mean(), a), a);
  jtree.AddMember("node_stddev", rapidjson_util::ToJSON(node_weights.MeanErr(), a), a);
  return jtree;
}

rapidjson::Value RadianceDistributionSampled::ToJSON(rapidjson_util::Alloc &a) const
{
  using namespace rapidjson_util;
  rj::Value jtree = tree.ToJSON(a);
  jtree.AddMember("node_weights", rapidjson_util::ToJSON(node_sample_probs, a), a);
  return jtree;
}

} //namespace quadtree_radiance_distribution

boost::filesystem::path GetDebugFilePrefix()
{
  return DEBUG_FILE_PREFIX;
}

PathGuiding::PathGuiding(const Box &region, double cellwidth, const RenderingParameters &params, tbb::task_arena &the_task_arena, const char* name) :
    region{ region },
    recording_tree{},
    name{name},
    param_num_initial_samples{ params.guiding_tree_subdivision_factor },
    param_em_every{ params.guiding_em_every },
    param_prior_strength{ params.guiding_prior_strength },
    the_task_arena{ &the_task_arena }
{
  auto& cell = cell_data.emplace_back();
  cell.index = 0;
  cell.current_estimate.cell_bbox = region;
}

CellData& PathGuiding::LookupCellData(const Double3 &p)
{
  return cell_data[recording_tree.Lookup(p)];
}

const PathGuiding::RadianceEstimate& PathGuiding::FindRadianceEstimate(const Double3 &p) const
{
    return cell_data[recording_tree.Lookup(p)].current_estimate;
}

void PathGuiding::AddSample(
  ThreadLocal& tl, const Double3 &pos,
  Sampler &sampler, const Double3 &reverse_incident_dir, const Spectral3 &radiance)
{
    auto rec = IncidentRadiance{
            pos,
            reverse_incident_dir.cast<float>(),
            radiance.cast<float>().mean(),
    };

    tl.samples.push_back(rec);
}

IncidentRadiance PathGuiding::ComputeStochasticFilterPosition(const IncidentRadiance & rec, const CellData &cd, RandGen &sampler)
{
  IncidentRadiance new_rec{rec};

  const Double3 rv{ sampler.Uniform01() , sampler.Uniform01() , sampler.Uniform01() };

  // Note: Stddev of uniform distribution over [0, 1] is ca 0.3. To scatter samples across this range
  //       I multiply the size of the axes by 1./stddev * 1/2 which is the 3/2rd factor here. The 1/2
  //       is there because the frame is centered in the middle, so the axes lengths need to be the "radius"
  //       of the point cloud, not the diameter.
  //       Other factors are simply ad hoc tuning parameters
  static constexpr double MAGIC = 2.;
  new_rec.pos += cd.current_estimate.points_cov_frame * (rv*2. - Double3::Ones()) * (MAGIC * 3. / 2.);
  new_rec.reverse_incident_dir = cd.current_estimate.radiance_distribution.ComputeStochasticFilteredDirection(rec, sampler);
  new_rec.is_original = false;
  return new_rec;
}

ToyVector<int> PathGuiding::ComputeCellIndices(Span<const IncidentRadiance> samples) const
{
  ToyVector<int> cell_indices(samples.size(), -1);
  tbb::parallel_for(tbb::blocked_range<long>(0, samples.size(), 1000), [&,this](tbb::blocked_range<long> r)
  {
    for (long i=r.begin(); i<r.end(); ++i)
    {
      const int cell = recording_tree.Lookup(samples[i].pos);
      cell_indices[i] = cell;
    }
  });
  return cell_indices;
}

ToyVector<ToyVector<IncidentRadiance>> PathGuiding::SortSamplesIntoCells(Span<const int> cell_indices, Span<const IncidentRadiance> samples) const
{
  ToyVector<long> cell_sample_count(cell_data.size(), 0);
  for (int i : cell_indices)
      ++cell_sample_count[i];

  ToyVector<ToyVector<IncidentRadiance>> sorted_samples(cell_data.size());
  for (int i=0; i<cell_data.size(); ++i)
    sorted_samples[i].reserve(cell_sample_count[i]);

  for (long i=0; i<samples.size(); ++i)
  {
    sorted_samples[cell_indices[i]].push_back(samples[i]);
  }

  return sorted_samples;
}

void PathGuiding::GenerateStochasticFilteredSamplesInplace(Span<int> cell_indices, Span<IncidentRadiance> samples) const
{
  tbb::enumerable_thread_specific<RandGen> tls_samplers;
  tbb::parallel_for(tbb::blocked_range<long>(0, samples.size(), 1000), [&, this](tbb::blocked_range<long> r)
  {
    auto sampler = tls_samplers.local();
    for (long i = r.begin(); i<r.end(); ++i)
    {
      const auto s = ComputeStochasticFilterPosition(samples[i], cell_data[cell_indices[i]], sampler);
      int new_cell = recording_tree.Lookup(s.pos);
      samples[i] = s;
      cell_indices[i] = new_cell;
    }
  });
}

// Also frees the tls sample buffer
// Updates the fit data in the CellData structs.
void PathGuiding::FitTheSamples(Span<ThreadLocal*> thread_locals)
{
  ToyVector<IncidentRadiance> samples;
  for (auto *tl : thread_locals)
  {
    samples.insert(samples.end(), tl->samples.begin(), tl->samples.end());
    ToyVector<IncidentRadiance>{}.swap(tl->samples);
  }

  auto cell_indices = ComputeCellIndices(AsSpan(samples));

#if 1 // Stochastic smoothing filter samples
  const auto noriginal = samples.size();
  samples.reserve(noriginal*2);
  cell_indices.reserve(noriginal*2);

  // This is only safe because the memory was reserved in advance!
  std::copy(samples.begin(), samples.begin()+noriginal, std::back_inserter(samples));
  std::copy(cell_indices.begin(), cell_indices.begin()+noriginal, std::back_inserter(cell_indices));

  GenerateStochasticFilteredSamplesInplace(
    Subspan(AsSpan(cell_indices), noriginal, noriginal),
    Subspan(AsSpan(samples), noriginal, noriginal));
#endif

  auto sorted_samples = SortSamplesIntoCells(AsSpan(cell_indices), AsSpan(samples));

  tbb::enumerable_thread_specific<RandGen> samplers;

  tbb::parallel_for<int>(0, isize(cell_data), [&, this](int cell_idx)
  {
    auto samples = AsSpan(sorted_samples[cell_idx]);
    if (samples.size())
    {
      RandomShuffle(samples.begin(), samples.end(), samplers.local());
      FitTheSamples(cell_data[cell_idx], samples);
    }
  });

  ++sub_round;
}

void PathGuiding::FitTheSamples(CellData &cell, Span<IncidentRadiance> buffer) const
{
  RadianceDistributionLearned::Parameters params{ cell, sub_round };
  cell.learned.radiance_distribution.IncrementalFit(buffer, params, cell.current_estimate.radiance_distribution);
  AddToPointStatistics(cell, buffer);

#ifdef PATH_GUIDING_WRITE_SAMPLES_ACTUALLY_ENABLED
  for (const auto &in : buffer)
  {
    if (in.weight <= 0.)
      continue;
    cell_data_debug[cell.index].Write(in.pos, in.reverse_incident_dir, in.weight);
  }
  //cell_data_debug[cell.index].params = params;
#endif
}

void PathGuiding::BeginRound(Span<ThreadLocal*> thread_locals)
{
  round++;
  sub_round = 1;

  const auto n = recording_tree.NumLeafs();
  assert(n == cell_data.size());

  for (auto* tl : thread_locals)
  {
    tl->samples.reserve(1024*1024);
  }

#ifdef PATH_GUIDING_WRITE_SAMPLES_ACTUALLY_ENABLED
  cell_data_debug.reset(new CellDebug[n]);
  for (std::ptrdiff_t i = 0; i < n; ++i)
  {
    cell_data_debug[i].Open(
      fmt::format("{}{}_samples_round_{}_cell_{}.csv", DEBUG_FILE_PREFIX, name, std::to_string(round), std::to_string(i))
    );
  }
#endif
}

void PathGuiding::FinalizeRound(Span<ThreadLocal*> thread_locals)
{
  the_task_arena->execute([&]() { the_task_group.wait(); } );

  long mem_cell_data = static_cast<long>(cell_data.size()*sizeof(CellData));
  long mem_tls_buffers = 0;
  for (auto* tl : thread_locals)
    mem_tls_buffers += lsize(tl->samples) + static_cast<long>(sizeof(decltype(*tl)));

  if (round <= 1)
  {
    the_task_arena->execute([&, this]() {
      the_task_group.run_and_wait([&, this]() {
        AdaptInitial(thread_locals);
      });
    });
  }
  else
  {
    the_task_arena->execute([&, this]() {
      the_task_group.run_and_wait([&, this]() {
        FitTheSamples(thread_locals);
      });
    });
  }

  // std::cout << "--- expected mem use [MB] ----" << std::endl;
  // std::cout << "Thread local buffer " << mem_cell_data/(1024*1024) << std::endl;
  // std::cout << "Cell data " << mem_cell_data/(1024*1024) << std::endl;
}

namespace
{

using namespace ::guiding::kdtree;

std::pair<Box, Box> ChildrenBoxes(const Tree &tree, Handle node, const Box &node_box)
{
  auto [axis, pos] = tree.Split(node);

  Box leftbox{ node_box };
  leftbox.max[axis] = pos;

  Box rightbox{ node_box };
  rightbox.min[axis] = pos;

  return std::make_pair(leftbox, rightbox);
}

void ComputeLeafBoxes(const Tree &tree, Handle node, const Box &node_box, Span<CellData> out)
{
  if (node.is_leaf)
  {
    // Assign lateral extents.
    //out[node.idx].current_estimate.cell_size = (node_box.max - node_box.min);
    out[node.idx].current_estimate.cell_bbox = node_box;
  }
  else
  {
    auto [left, right] = tree.Children(node);
    auto [bl, br] = ChildrenBoxes(tree, node, node_box);
    ComputeLeafBoxes(tree, left, bl, out);
    ComputeLeafBoxes(tree, right, br, out);
  }
}

} // namespace

namespace
{

void InitializePcaFrame(CellData::CurrentEstimate &ce, const CellData::Learned &learned)
{
  Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> eigensolver{ learned.leaf_stats.Cov() };
  // Note: Eigenvalues are the variances w.r.t. the eigenvector frame.
  //       Take sqrt to obtain stddev.
  //       Then  scale axes of the frame so that computation of points for the stochastic filtering only needs the matrix multiplication.
  const Eigen::Vector3d evs = eigensolver.eigenvalues().cwiseMax(Epsilon);
  ce.points_cov_frame = eigensolver.eigenvectors() * evs.cwiseSqrt().asDiagonal();
  ce.points_mean = learned.leaf_stats.Mean();
  ce.points_stddev = learned.leaf_stats.Var().cwiseSqrt();

  assert(learned.leaf_stats.Count() > 1);
  assert(learned.leaf_stats.Cov().allFinite());
  assert(ce.points_cov_frame.allFinite());
  assert(ce.points_mean.allFinite());
}

}

void PathGuiding::PrepareAdaptedStructures()
{
  if (round > 1)
  {
    WriteDebugData();
    AdaptIncremental();
  }
}

void PathGuiding::AdaptInitial(Span<ThreadLocal*> thread_locals)
{
  // There is only one cell, and all the samples are stored in its CellDataTemporary struct.
  assert(cell_data.size() == 1);

  // For easier access. Also freeing the memory, which I have to eventually.
  ToyVector<IncidentRadiance> samples; samples.reserve(1024*1024);
  for (auto* tls : thread_locals)
  {
    samples.insert(samples.end(), tls->samples.begin(), tls->samples.end());
    ToyVector<IncidentRadiance>{}.swap(tls->samples);
  }

  // Probably this is taking samples from media interaction, in a scene without media.
  if (samples.empty())
    return;

  auto builder = kdtree::MakeBuilder<IncidentRadiance>(/*max_depth*/ MAX_DEPTH, /*min_num_points*/ param_num_initial_samples, [](const IncidentRadiance &s) { return s.pos; });
  recording_tree = builder.Build(AsSpan(samples));

  cell_data.resize(recording_tree.NumLeafs());

  std::cout << "Fitting to " << samples.size() << " initial samples in " << cell_data.size() << " cells ..." << std::endl;

  tbb::parallel_for(0, isize(cell_data), [this, &builder](int i)
  {
      CellData& cd = cell_data[i];
      Span<IncidentRadiance> cell_samples = builder.DataRangeOfLeaf(i);
      assert(cell_samples.size() > 2);

      cd.index = i;
      cd.last_num_samples = cell_samples.size();
      cd.max_num_samples = cell_samples.size();

      AddToPointStatistics(cd, cell_samples);
      InitializePcaFrame(cd.current_estimate, cd.learned);

      RadianceDistributionLearned::Parameters params{ cd, 1 };
      cd.learned.radiance_distribution.InitialFit(cell_samples, params);
      cd.current_estimate.radiance_distribution = cd.learned.radiance_distribution.Bake();
  });

  ComputeLeafBoxes(recording_tree, recording_tree.GetRoot(), region, AsSpan(cell_data));

  previous_max_samples_per_cell = param_num_initial_samples;
  previous_total_samples = samples.size();

  std::cout << "round " << round << ", num cells = " << cell_data.size() << std::endl;
  const double avg_sample_count = static_cast<double>(std::accumulate(cell_data.begin(), cell_data.end(), 0l, [](const long &a, const CellData &b) { return a+b.last_num_samples; })) / cell_data.size();
  std::cout << "  avg sample count = " << avg_sample_count << std::endl;
}

void PathGuiding::AdaptIncremental()
{
  const std::int64_t num_fit_samples = std::accumulate(cell_data.begin(), cell_data.end(), 0l, [](std::int64_t n, const CellData &cd) {
    return n + cd.learned.leaf_stats.Count();
  });
  // Start off with 100 samples per cell = nc.
  // mc = sqrt(N)*c
  // mc_0 = sqrt(N0)*c = nc
  // => c = nc / sqrt(N0)

  // mc = sqrt(M)*c
  // nc = sqrt(N)*c
  // => nc = mc/sqrt(M)*sqrt(N)

  const std::int64_t max_samples_per_cell = num_fit_samples
    ? std::sqrt(num_fit_samples)*previous_max_samples_per_cell/std::sqrt(previous_total_samples)
    : previous_max_samples_per_cell;

  fmt::print(
    "round {}: num_samples={},\n current num_cells={},\n target samples per cell={}\n",
     round, num_fit_samples, recording_tree.NumLeafs(), max_samples_per_cell);

  { // Start tree adaptation
    auto DetermineSplit = [this, max_samples_per_cell](int cell_idx) -> std::pair<int, double>
    {
      // If too many samples fell into this cell -> split it at the sample's centroid.
      const auto& stats = cell_data[cell_idx].learned.leaf_stats;
      if (stats.Count() > max_samples_per_cell)
      {
        const auto mean = stats.Mean();
        const auto var = stats.Var();
        int axis = -1;
        var.maxCoeff(&axis);
        double pos = mean[axis];
        return { axis, pos };
      }
      else
      {
        return { -1, NaN };
      }
    };

    kdtree::TreeAdaptor adaptor(DetermineSplit);
    recording_tree = adaptor.Adapt(recording_tree);

    decltype(cell_data) new_data(recording_tree.NumLeafs());

    int num_cell_over_2x_limit = 0;
    int num_cell_over_1_5x_limit = 0;
    int num_cell_over1x_limit = 0;
    int num_cell_else = 0;
    int num_cell_sample_count_regressed = 0;
    double kl_divergence = 0.;
    const double logq = std::log2(1. / isize(cell_data));

    // This should initialize all sampling mixtures with
    // the previously learned ones. If node is split, assignment
    // is done to both children.
    int i = 0;
    for (const auto& m : adaptor.GetNodeMappings())
    {
      const bool is_split = (m.new_first >= 0) && (m.new_second >= 0);

      auto baked_distribution = cell_data[i].learned.radiance_distribution.IterationUpdateAndBake(cell_data[i].current_estimate.radiance_distribution);

      auto CopyCell = [&src = cell_data[i], new_data = AsSpan(new_data), &baked_distribution](int dst_idx, bool is_split) mutable
      {
        CellData &dst = new_data[dst_idx];
        dst.index = dst_idx;

        dst.current_estimate.radiance_distribution = baked_distribution;
        dst.learned.radiance_distribution = src.learned.radiance_distribution;

        if (src.learned.leaf_stats.Count() >= 10)
        {
          InitializePcaFrame(dst.current_estimate, src.learned);
        }
        else
        {
          dst.current_estimate.points_cov_frame = src.current_estimate.points_cov_frame;
          dst.current_estimate.points_mean = src.current_estimate.points_mean;
          dst.current_estimate.points_stddev = src.current_estimate.points_stddev;
        }

        dst.last_num_samples = src.learned.leaf_stats.Count();
        dst.max_num_samples = std::max(dst.last_num_samples, dst.max_num_samples);
      };

      if (m.new_first >= 0)
      {
        CopyCell(m.new_first, is_split);
      }
      if (m.new_second >= 0)
      {
        CopyCell(m.new_second, is_split);
      }

      const int num_samples = cell_data[i].learned.leaf_stats.Count();
      if (num_samples > 2 * max_samples_per_cell)
        ++num_cell_over_2x_limit;
      else if (2 * num_samples > 3 * max_samples_per_cell)
        ++num_cell_over_1_5x_limit;
      else if (num_samples > max_samples_per_cell)
        ++num_cell_over1x_limit;
      else if (!is_split && num_samples * 3 < cell_data[i].last_num_samples * 4)
        ++num_cell_sample_count_regressed;
      else
        ++num_cell_else;
      const double p = static_cast<double>(num_samples) / num_fit_samples;
      kl_divergence += p > 0. ? (p*std::log2(p) - p * logq) : 0.;
      ++i;
    }

    cell_data = std::move(new_data);

    previous_max_samples_per_cell = max_samples_per_cell;
    previous_total_samples = num_fit_samples;

    fmt::print(
      "num_cell_over_2x_limit = {}\n"
      "num_cell_over_1_5x_limit = {}\n"
      "num_cell_over1x_limit = {}\n"
      "num_cell_else = {}\n"
      "num_cell_regressed = {}\n"
      "kl_divergence = ",
      num_cell_over_2x_limit,
      num_cell_over_1_5x_limit,
      num_cell_over1x_limit,
      num_cell_else,
      num_cell_sample_count_regressed,
      kl_divergence);
    ++i;
  } // End tree adaption

  ComputeLeafBoxes(recording_tree, recording_tree.GetRoot(), region, AsSpan(cell_data));
}

#ifdef PATH_GUIDING_WRITE_SAMPLES_ACTUALLY_ENABLED
CellDebug::~CellDebug()
{
  Close();
}

void CellDebug::Open(std::string filename_)
{
  this->filename = filename_;
  std::cout << "opening " << this->filename << std::endl;
  assert(!file.is_open());
  file.open(filename);
  assert(file.is_open());
}

void CellDebug::Write(const Double3 &pos, const Float3 &dir, float weight)
{
  assert(file.is_open());
  file << pos[0] << "," << pos[1] << "," << pos[2] << ","<< dir[0] << ","<< dir[1] << ","<< dir[2] << ","<< weight << "\n";
}

void CellDebug::Close()
{
  std::cout << "closing " << this->filename << std::endl;
  file.close();
}

#endif

void PathGuiding::WriteDebugData()
{
#if (defined WRITE_DEBUG_OUT & defined HAVE_JSON & !defined NDEBUG)
  const auto filename = fmt::format("{}{}_radiance_records_{}.json",DEBUG_FILE_PREFIX, name, round);
  std::cout << "Writing " << filename << std::endl;

  using namespace rapidjson_util;
  rj::Document doc;
  doc.SetObject();

  recording_tree.DumpTo(doc, doc);

  rj::Value records(rj::kArrayType);

  auto &a = doc.GetAllocator();

  for (const auto &cd : cell_data)
  {
    //std::cout << "cell " << idx << " contains" << celldata.incident_radiance.size() << " records " << std::endl;
    rj::Value jcell(rj::kObjectType);
    jcell.AddMember("id", cd.index, a);

    if (cd.learned.leaf_stats.Count()>0)
    {
      assert(cd.learned.leaf_stats.Cov().allFinite());
      assert(cd.current_estimate.points_cov_frame.allFinite());
      assert(cd.current_estimate.points_mean.allFinite());
      jcell.AddMember("point_distribution_mean", ToJSON(cd.current_estimate.points_mean, a), a);
      jcell.AddMember("point_distribution_frame", ToJSON(cd.current_estimate.points_cov_frame, a), a);
      jcell.AddMember("point_distribution_stddev", ToJSON(cd.current_estimate.points_stddev, a), a);
    }
    else
    {
      jcell.AddMember("point_distribution_mean", ToJSON(Eigen::Vector3d::Zero().eval(), a), a);
      jcell.AddMember("point_distribution_frame", ToJSON(Eigen::Matrix3d::Zero().eval(), a), a);
      jcell.AddMember("point_distribution_stddev", ToJSON(Eigen::Vector3d::Zero().eval(), a), a);
    }
    jcell.AddMember("bbox_min", ToJSON(cd.current_estimate.cell_bbox.min, a), a);
    jcell.AddMember("bbox_max", ToJSON(cd.current_estimate.cell_bbox.max, a), a);
    jcell.AddMember("num_points", cd.learned.leaf_stats.Count(), a);
    //jcell.AddMember("average_weight", cd.learned.fitdata.avg_weights(), a);
    //jcell.AddMember("incident_flux_learned", ToJSON(cd.learned.incident_flux_density_accum.Mean(), a), a);
    //jcell.AddMember("incident_flux_sampled", ToJSON(cd.current_estimate.incident_flux_density, a), a);
    jcell.AddMember("radiance_learned", cd.learned.radiance_distribution.ToJSON(a), a);
    jcell.AddMember("radiance_sampled", cd.current_estimate.radiance_distribution.ToJSON(a), a);

    // Fit parameters
#ifdef PATH_GUIDING_WRITE_SAMPLES_ACTUALLY_ENABLED
    jcell.AddMember("fitparam_prior_nu", cell_data_debug[cd.index].params.prior_nu, a);
    jcell.AddMember("fitparam_prior_tau", cell_data_debug[cd.index].params.prior_tau, a);
    jcell.AddMember("fitparam_prior_alpha", cell_data_debug[cd.index].params.prior_alpha, a);
    jcell.AddMember("fitparam_maximization_step_every", cell_data_debug[cd.index].params.maximization_step_every, a);
#endif

    records.PushBack(jcell.Move(), a);
  }

  assert(records.Size() == cell_data.size());

  doc.AddMember("records", records.Move(), a);

  {
    std::ofstream file(filename);
    rapidjson::OStreamWrapper osw(file);
    rapidjson::Writer<rapidjson::OStreamWrapper> writer(osw);
    doc.Accept(writer);
  }

  //auto datastr = ToString(doc);
  //file << datastr;

#ifdef PATH_GUIDING_WRITE_SAMPLES_ACTUALLY_ENABLED
  cell_data_debug.reset();
#endif
#endif
}

} // namespace guiding


#include <iostream>
#include <vector>

#include <dlib/dnn.h>
#include <sampml/data.hpp>
#include <sampml/random_forest.hpp>
#include <sampml/svm_classifier.hpp>

#include "common.hpp"
#include "dnn.hpp"

std::vector<output_vector> samples;
void transform (std::string_view path) {
    sampml::data::reader<input_vector> data(path);
    transformer t;
    t.submit(data);

    samples = std::move(t.pool);
}

void test() {
    sampml::trainer::svm_classifier<sample_type> svm;
    svm.deserialize(std::string(svm_model.begin(), svm_model.end()));

    sampml::trainer::random_forest<sample_type> rf;
    rf.deserialize(std::string(rf_model.begin(), rf_model.end()));

    test_network_type dnn;
    dlib::deserialize(std::string(dnn_model.begin(), dnn_model.end())) >> dnn.subnet();

    int aimbot_samples = 0,
        clean_samples = 0;
    dlib::running_stats<double> values;
    for(auto& sample : samples) {
        double svm_prob = svm.test(sample),
               rf_prob = std::clamp(rf.test(sample), 0.0, 1.0),
               dnn_prob = dlib::mat(dnn(sample))(1);
        double prob = (svm_prob + rf_prob + dnn_prob)/3;

        const float cutoff = 0.5;
        if(prob > cutoff)
            aimbot_samples++;
        else
            clean_samples++;
        values.add(prob);
    }
    std::cout << "test statistics: " << '\n';
    std::cout << "average: " << values.mean() << ", stddev: " << values.stddev() << '\n'
              << "min: " << values.min() << ", max: " << values.max()  << '\n'
              << "skewness: " << values.skewness() << ", excess kurtosis: " << values.ex_kurtosis() << '\n';

    std::cout << "samples classified as aimbot: " << aimbot_samples << '\n'
              << "samples classified as clean: " << clean_samples << '\n'
              << "possible use of aimbot: " << float(aimbot_samples)/(samples.size()) * 100 << "%\n";
}

int main (int argc, char* argv[]) {
    if(argc == 1) {
        std::cout << "Usage: " << argv[0] << " [raw data file]" << std::endl;
        return 0;
    }

    transform(argv[1]);
    test();
}


//==================================================================================================
/*!
  @file
  @copyright 2016 NumScale SAS
  @copyright 2016 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_SIMD_FUNCTION_GAMMA_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_GAMMA_HPP_INCLUDED

#include <boost/simd/function/std.hpp>
#include <boost/simd/arch/common/detail/generic/gamma_kernel.hpp>

#include <boost/simd/constant/half.hpp>
#include <boost/simd/constant/inf.hpp>
#include <boost/simd/constant/logpi.hpp>
#include <boost/simd/constant/logsqrt2pi.hpp>
#include <boost/simd/constant/minf.hpp>
#include <boost/simd/constant/mone.hpp>
#include <boost/simd/constant/nan.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/constant/two.hpp>
#include <boost/simd/constant/three.hpp>
#include <boost/simd/constant/zero.hpp>

#include <boost/simd/constant/pi.hpp>
#include <boost/simd/constant/two.hpp>

#include <boost/simd/function/abs.hpp>
#include <boost/simd/function/any.hpp>
#include <boost/simd/function/bitwise_or.hpp>
#include <boost/simd/function/floor.hpp>
#include <boost/simd/function/is_even.hpp>

#include <boost/simd/function/fma.hpp>
#include <boost/simd/function/if_dec.hpp>
#include <boost/simd/function/if_else.hpp>
#include <boost/simd/function/if_inc.hpp>
#include <boost/simd/function/if_minus.hpp>
#include <boost/simd/function/if_nan_else.hpp>
#include <boost/simd/function/if_plus.hpp>
#include <boost/simd/function/is_equal.hpp>
#include <boost/simd/function/is_eqz.hpp>
#include <boost/simd/function/is_flint.hpp>
#include <boost/simd/function/is_less.hpp>
#include <boost/simd/function/is_lez.hpp>
#include <boost/simd/function/log.hpp>
#include <boost/simd/function/logical_and.hpp>
#include <boost/simd/function/logical_andnot.hpp>
#include <boost/simd/function/logical_or.hpp>
#include <boost/simd/function/nbtrue.hpp>
#include <boost/simd/function/negif.hpp>
#include <boost/simd/function/sinpi.hpp>
#include <boost/simd/function/sqr.hpp>

#include <boost/simd/meta/as_logical.hpp>
#include <boost/simd/detail/dispatch/function/overload.hpp>
#include <boost/simd/detail/dispatch/meta/scalar_of.hpp>
#include <boost/config.hpp>
#include <cmath>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;
  BOOST_DISPATCH_OVERLOAD_IF ( gamma_
                             , (typename A0, typename X)
                             , (detail::is_native<X>)
                             , bd::cpu_
                             , bs::pack_< bd::floating_<A0>, X>
                             )
  {
    BOOST_FORCEINLINE A0 operator() (const A0& a0) const BOOST_NOEXCEPT
    {
      auto nan_result = logical_and(is_ltz(a0), is_flint(a0));
      #ifndef BOOST_SIMD_NO_INVALIDS
      nan_result = logical_or(bs::is_nan(a0), nan_result);
      #endif
      A0 q = bs::abs(a0);
      auto test = is_less(a0, A0(-33.0));
      std::size_t nb = bs::nbtrue(test);
      A0 r =  Nan<A0>();
      if(nb > 0)
      {
        //treat negative large with reflection
        r = large_negative(q);
        if (nb >= A0::static_size)
          return bs::if_nan_else(nan_result, r);
      }
      A0 r1 = other(a0, test);
      A0 r2 = if_else(test, r, r1);
      return bs::if_else(is_eqz(a0), copysign(Inf<A0>(), a0), bs::if_nan_else(nan_result, r2));
    }
  private :
    static  A0 large_negative(const A0& q) BOOST_NOEXCEPT
    {
      A0 st =  bs::stirling(q);
      A0 p = bs::floor(q);
      A0 sgngam = bs::negif(bs::is_even(p), One<A0>());
      A0 z = q - p;
      auto test2 = is_less(z, bs::Half<A0>() );
      z = bs::if_minus(test2, z, bs::One<A0>());
      z = q*bs::sinpi(z);
      z =  bs::abs(z);
      return sgngam*bs::Pi<A0>()/(z*st);
    }

    template < typename BA0>
    static A0 other(const A0& q, const BA0& test) BOOST_NOEXCEPT
    {
      A0 x =  bs::if_else(test, Two<A0>(), q);
      #ifndef BOOST_SIMD_NO_INFINITIES
      auto inf_result = is_equal(q, Inf<A0>());
      x = if_else(inf_result, Two<A0>(), x);
      #endif
      A0 z = bs::One<A0>();
      auto test1 = is_greater_equal(x,Three<A0>());
      while( bs::any(test1) )
      {
        x = bs::if_plus(test1, x, bs::Mone<A0>());
        z = bs::if_else(   test1, z*x, z);
        test1 = is_greater_equal(x,Three<A0>());
      }
      //all x are less than 3
      test1 = bs::is_ltz(x);
      while( bs::any(test1) )
      {
        z = bs::if_else(test1, z/x, z);
        x = bs::if_plus(test1, x, bs::One<A0>());
        test1 = bs::is_ltz(x);
      }
      //all x are greater than 0 and less than 3
      auto test2 = is_less(x,bs::Two<A0>());
      while( bs::any(test2))
      {
        z = bs::if_else(test2, z/x, z);
        x = bs::if_plus(test2, x, bs::One<A0>());
        test2 = is_less(x,bs::Two<A0>());
      }
      //all x are greater equal 2 and less than 3
      x = z*detail::gamma_kernel<A0>::gamma1(x-bs::Two<A0>());
      #ifndef BOOST_SIMD_NO_INFINITIES
      return if_else(inf_result, q, x);
      #else
      return x;
      #endif
    }
  };

} } }

#endif


#include <iostream>

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/timer.hpp>

#include <vector>
#include <math.h>

#include "NRGclasses.hpp"
#include "NRGfunctions.hpp"
#include "TwoChQS.hpp"

void TwoChQS_UpdateQm1fQ(CNRGbasisarray* pSingleSite,CNRGarray* pAeig,
			  CNRGbasisarray* pAbasis,CNRGmatrix* Qm1fNQ){

  // Boost matrices
  boost::numeric::ublas::matrix<double> Zibl;
  boost::numeric::ublas::matrix<double> Zjbl;
  boost::numeric::ublas::matrix<double> fnbasis;

  boost::numeric::ublas::matrix<double> fnw;
  // Check time
  boost::timer t;
  double time_elapsed;

  //Clear all
  for (int ich=1;ich<=2;ich++)
    {
      Qm1fNQ[ich-1].ClearAll();
      Qm1fNQ[ich-1].SyncNRGarray(*pAeig);
    }

  // Note: assumes  Abasis and Aeig have
  // EXACTLY the same block structure... should check.

  double Qi;
  double Qj;
  double Si,Sj;

  // Loop over blocks (icount counts for each matrix)
  int icount[2]={0,0};
  for (int ibl=0;ibl<pAeig->NumBlocks();ibl++)
    {
      Qi=pAeig->GetQNumber(ibl,0);
      Si=pAeig->GetQNumber(ibl,1);
      int Nstibl=pAeig->GetBlockSize(ibl);

      for (int jbl=ibl+1;jbl<pAeig->NumBlocks();jbl++)
	{
	  Qj=pAeig->GetQNumber(jbl,0);
	  Sj=pAeig->GetQNumber(jbl,1);
	  int Nstjbl=pAeig->GetBlockSize(jbl);

	  //Loop over channels: check if the matrix elements are non-zero
	  for (int ich=1;ich<=2;ich++)
	    {
	      if (  dEqual(Qi+1.0,Qj)&&
		   ( dEqual(Si+0.5,Sj)||dEqual(Si-0.5,Sj) )  )
		{
		  // Set Zi,Zj
		  cout << " BC: Nshell = " << pAeig->Nshell << endl;
		  cout << "Matrix in channel: " << ich << endl;
		  cout << "  Setting up Block i : " << ibl
		       << " (size " <<  Nstibl
		       << ") x Block j " << jbl
		       << " (size " << Nstjbl << ") of " << pAeig->NumBlocks() << endl;
		  Qm1fNQ[ich-1].MatBlockMap.push_back(ibl);
		  Qm1fNQ[ich-1].MatBlockMap.push_back(jbl);
		  Qm1fNQ[ich-1].MatBlockBegEnd.push_back(icount[ich-1]);
		  icount[ich-1]+=Nstibl*Nstjbl;
		  Qm1fNQ[ich-1].MatBlockBegEnd.push_back(icount[ich-1]-1);

		  cout << "    Setting up BLAS matrices..." << endl;
		  //boost::numeric::ublas::matrix<double> Zibl(Nstibl,Nstibl);
		  //boost::numeric::ublas::matrix<double> Zjbl(Nstjbl,Nstjbl);
		  Zibl.resize(Nstibl,Nstibl);
		  Zjbl.resize(Nstjbl,Nstjbl);

		  Zibl=pAeig->EigVec2BLAS(ibl);
		  Zjbl=pAeig->EigVec2BLAS(jbl);
		  cout << "    ...Zs done ..." << endl;

		  //boost::numeric::ublas::matrix<double> fnbasis (Nstibl,Nstjbl);
		  fnbasis.resize(Nstibl,Nstjbl);
		  // Set-up fn basis: Loop in the basis states
		  int istbl=0;
		  for (int ist=pAbasis->GetBlockLimit(ibl,0);
		       ist<=pAbasis->GetBlockLimit(ibl,1);ist++)
		    {
		      int typei=pAbasis->iType[ist];
		      int stcfi=pAbasis->StCameFrom[ist];

		      int jstbl=0;
		      for (int jst=pAbasis->GetBlockLimit(jbl,0);
			   jst<=pAbasis->GetBlockLimit(jbl,1);jst++)
			{
			  int typej=pAbasis->iType[jst];
			  int stcfj=pAbasis->StCameFrom[jst];

			  fnbasis(istbl,jstbl)=0.0;

			  if (stcfi==stcfj)
			    {
			      // New approach

			      // Get SingleSite QNumbers
			      int pos=0;
			      int iblssi=pSingleSite->GetBlockFromSt(typei,pos);
			      int iblssj=pSingleSite->GetBlockFromSt(typej,pos);

			      double Qtildei=pSingleSite->GetQNumber(iblssi,0);
			      double Stildei=pSingleSite->GetQNumber(iblssi,1);
			      double Sztildei=pSingleSite->GetQNumber(iblssi,2);

			      double Qtildej=pSingleSite->GetQNumber(iblssj,0);
			      double Stildej=pSingleSite->GetQNumber(iblssj,1);
			      double Sztildej=pSingleSite->GetQNumber(iblssj,2);

			      double siteqnumsi[]={Qtildei,Stildei,Sztildei};
			      double siteqnumsj[]={Qtildej,Stildej,Sztildej};

			      double Scfi=Si-Sztildei;
			      double Scfj=Sj-Sztildej; // should be the same!!

			      double auxBasis[2]={0.0,0.0};
			      double auxCG[2]={0.0,0.0};
			      double CGnorm=1.0;

			      int sigma=0;

			      // Choose which sigma
			      auxCG[0]=CGordan(Si,Si,0.5,-0.5,Sj,Sj);
			      auxCG[1]=CGordan(Si,Si,0.5,0.5,Sj,Sj);
			      if ( dNEqual(auxCG[0],0.0) )
				{sigma=-1;CGnorm=auxCG[0];}
			      else
				{
				  if ( dNEqual(auxCG[1],0.0) )
				    {sigma=1;CGnorm=auxCG[1];}
				  else fnbasis(istbl,jstbl)=0.0;
				}

// 			      cout << " UpMat: "
// 				   << " ist = " << ist
// 				   << " jst = " << jst
// 				   << " CG0 = " << auxCG[0]
// 				   << " CG1 = " << auxCG[1]
// 				   << " sigma = " << sigma
// 				   << endl;

			      double dSigma=0.5*sigma;
			      // Now Sj=Si+dSigma necessarily
			      // Calculate
			      if (sigma!=0)
				{
				  for (double Szcfj=Scfj;
				       Szcfj>=-Scfj;Szcfj-=1.0)
				    {
				      Sztildei=Si-Szcfj; //Szcfi=Szcfj
				      Sztildej=Sj-Szcfj;
				      siteqnumsi[2]=Sztildei;
				      int blockstatei=pSingleSite->GetBlockFromQNumbers(siteqnumsi);
				      siteqnumsj[2]=Sztildej;
				      int blockstatej=pSingleSite->GetBlockFromQNumbers(siteqnumsj);

				      auxCG[0]=CGordan(Scfi,Szcfj,
						       Stildei,Sztildei,
						       Si,Si);
				      auxCG[1]=CGordan(Scfj,Szcfj,
						       Stildej,Sztildej,
						       Sj,Sj);

				      for (
			  int sitestatei=pSingleSite->GetBlockLimit(blockstatei,0);
			  sitestatei<=pSingleSite->GetBlockLimit(blockstatei,1);
			  sitestatei++)
					{
					  for (
			  int sitestatej=pSingleSite->GetBlockLimit(blockstatej,0);
			  sitestatej<=pSingleSite->GetBlockLimit(blockstatej,1);
			  sitestatej++)
					    {
					      auxBasis[0]+=TwoChQS_fd_table(ich,sigma,sitestatej,sitestatei)*auxCG[0]*auxCG[1]*TwoChQS_SpSm_table(sitestatei,typei)*TwoChQS_SpSm_table(sitestatej,typej);

// 					      cout << " Szcfj = " << Szcfj
// 						   << " Stildei = " << Stildei
// 						   << " Sztildei = " << Sztildei
// 						   << " Stildej = " << Stildej
// 						   << " Sztildej = " << Sztildej
// 						   << " typei = " << typei
// 						   << " sitestatei = " << sitestatei
// 						   << " typej = " << typej
// 						   << " sitestatej = " << sitestatej
// 						   << " CG0 = " << auxCG[0]
// 						   << " CG1 = " << auxCG[1]
// 						   << endl;

					    }
					  // loop in blockstatej
					}
				      // loop in blockstatei
				    }
				  // end Loop in Szcf

				  //Calculate fbasis
				  fnbasis(istbl,jstbl)=auxBasis[0]/CGnorm;

//  				  cout << " auxBasis = " << auxBasis[0]
// 				       << " CGnorm = " << CGnorm
//  				       << endl;

				}
			      // end if (sigma!=0)

			    }
			  // end istcf=jstcf
			  jstbl++;
			}
		      // end jst loop
		      istbl++;
		    }
		  // end ist loop
		  cout << "    ...fnbasis done." << endl;

		  cout << "    Multiplying BLAS matrices... " << endl;

		  fnw.resize(Nstibl,Nstjbl);
		  t.restart();
		  noalias(fnw)=prod (Zibl,
				     boost::numeric::ublas::matrix<double>(prod(fnbasis,trans(Zjbl))) );
		  time_elapsed=t.elapsed();
		  cout << "    ...done. Elapsed time:" << time_elapsed << endl;

//   		  cout << "Z(i="<<ibl<<")  : " <<  Zibl << endl;
//   		  cout << "Z(j="<<jbl<<")  : " <<  Zjbl << endl;
//   		  cout << "fbasis   :" <<  fnbasis << endl;
//   		  cout << "Zi.fbasis.ZjT : " <<  fnw << endl;

		  // Add to Qm1fNQ[ich-1]
		  for (int ii=0;ii<fnw.size1();ii++)
		    for (int jj=0;jj<fnw.size2();jj++)
		      Qm1fNQ[ich-1].MatEl.push_back(fnw(ii,jj));

		}
	      //end if Q=Q'+1 etc

	    }
	  // end channel loop

	}
      // end jbl loop

    }
  // end ibl loop

}
// END subroutine

//////////////////////////////////
//////////////////////////////////
//////////////////////////////////
//////////////////////////////////

void TwoChQS_UpdateMatrixAfterCutting(CNRGbasisarray* pSingleSite,
				       CNRGbasisarray* pAeigCut,
				       CNRGbasisarray* pAbasis,
				       CNRGmatrix* Qm1fNQ,
				       CNRGmatrix* pMQQp1){

  // Boost matrices
  boost::numeric::ublas::matrix<double> Zibl;
  boost::numeric::ublas::matrix<double> Zjbl;
  boost::numeric::ublas::matrix<double> fnbasis;

  boost::numeric::ublas::matrix<double> fnw;
  // Check time
  boost::timer t;
  double time_elapsed;

  //Clear all
  for (int ich=1;ich<=2;ich++)
    {
      Qm1fNQ[ich-1].ClearAll();
      Qm1fNQ[ich-1].SyncNRGarray(*pAeigCut);
    }

  // Note: assumes  Abasis and Aeig have
  // EXACTLY the same block structure... should check.

  double Qi;
  double Qj;
  double Si,Sj;

  double qnums[2];

  // Loop over blocks (icount counts for each matrix)
  int icount[2]={0,0};
  for (int ibl=0;ibl<pAeigCut->NumBlocks();ibl++)
    {
      Qi=pAeigCut->GetQNumber(ibl,0);
      Si=pAeigCut->GetQNumber(ibl,1);
      int Nstibl=pAeigCut->GetBlockSize(ibl);

      // Get the corresponding block in pAbasis!
      qnums[0]=Qi;
      qnums[1]=Si;
      int ii_corr=pAbasis->GetBlockFromQNumbers(qnums);
      int NstiblBC=pAbasis->GetBlockSize(ii_corr);

      for (int jbl=ibl+1;jbl<pAeigCut->NumBlocks();jbl++)
	{
	  Qj=pAeigCut->GetQNumber(jbl,0);
	  Sj=pAeigCut->GetQNumber(jbl,1);
	  int Nstjbl=pAeigCut->GetBlockSize(jbl);

	  // Get the corresponding block in pAbasis!
	  qnums[0]=Qj;
	  qnums[1]=Sj;
	  int jj_corr=pAbasis->GetBlockFromQNumbers(qnums);
	  int NstjblBC=pAbasis->GetBlockSize(jj_corr);

	  if ( (!dEqual(Qj,pAbasis->GetQNumber(jj_corr,0)))||
	       (!dEqual(Sj,pAbasis->GetQNumber(jj_corr,1))) )
	    {
	      cout << "Ops, problems with GetBlockFromQNumbers" << endl;
	      pAbasis->PrintQNumbers();
	      cout << qnums[0] << " " << qnums[1] << endl;
	      cout << pAbasis->GetBlockFromQNumbers(qnums) << endl;
	    }

	  //Loop over channels: check if the matrix elements are non-zero
	  if (  dEqual(Qi+1.0,Qj)&&
		( dEqual(Si+0.5,Sj)||dEqual(Si-0.5,Sj) )  )
	    {
	      for (int ich=1;ich<=2;ich++)
		{
		  // Set Zi,Zj
		  cout << " AC : Nshell = " << pAeigCut->Nshell << endl;
		  cout << "Matrix in channel: " << ich << endl;
		  cout << "  Setting up Block i : " << ibl
		       << " (size " <<  Nstibl << " was " << NstiblBC
		       << ") x Block j " << jbl
		       << " (size " << Nstjbl << " was " << NstjblBC
		       << ") of " << pAeigCut->NumBlocks() << endl;
		  Qm1fNQ[ich-1].MatBlockMap.push_back(ibl);
		  Qm1fNQ[ich-1].MatBlockMap.push_back(jbl);
		  Qm1fNQ[ich-1].MatBlockBegEnd.push_back(icount[ich-1]);
		  icount[ich-1]+=Nstibl*Nstjbl;
		  Qm1fNQ[ich-1].MatBlockBegEnd.push_back(icount[ich-1]-1);

		  cout << "    Setting up BLAS matrices..." << endl;
		  //boost::numeric::ublas::matrix<double> Zibl(Nstibl,Nstibl);
		  //boost::numeric::ublas::matrix<double> Zjbl(Nstjbl,Nstjbl);
		  Zibl.resize(Nstibl,NstiblBC);
		  Zjbl.resize(Nstjbl,NstjblBC);

		  //Zibl=pAeigCut->EigVec2BLAS(ibl);
		  //Zjbl=pAeigCut->EigVec2BLAS(jbl);
		  Zibl=pAeigCut->EigVecCut2BLAS(ibl);
		  Zjbl=pAeigCut->EigVecCut2BLAS(jbl);
		  cout << "    ...Zs done ..." << endl;

		  fnbasis.resize(NstiblBC,NstjblBC);
		  // Set-up fn basis: Loop in the basis states
		  // Watch out here... ii_corr, not ibl
		  int istbl=0;
		  for (int ist=pAbasis->GetBlockLimit(ii_corr,0);
		       ist<=pAbasis->GetBlockLimit(ii_corr,1);ist++)
		    {
		      int typei=pAbasis->iType[ist];
		      int stcfi=pAbasis->StCameFrom[ist];

		      int jstbl=0;
		      for (int jst=pAbasis->GetBlockLimit(jj_corr,0);
			   jst<=pAbasis->GetBlockLimit(jj_corr,1);jst++)
			{
			  int typej=pAbasis->iType[jst];
			  int stcfj=pAbasis->StCameFrom[jst];

			  fnbasis(istbl,jstbl)=0.0;

			  if (stcfi==stcfj)
			    {
			      // New approach

			      // Get SingleSite QNumbers
			      int pos=0;
			      int iblssi=pSingleSite->GetBlockFromSt(typei,pos);
			      int iblssj=pSingleSite->GetBlockFromSt(typej,pos);

			      double Qtildei=pSingleSite->GetQNumber(iblssi,0);
			      double Stildei=pSingleSite->GetQNumber(iblssi,1);
			      double Sztildei=pSingleSite->GetQNumber(iblssi,2);

			      double Qtildej=pSingleSite->GetQNumber(iblssj,0);
			      double Stildej=pSingleSite->GetQNumber(iblssj,1);
			      double Sztildej=pSingleSite->GetQNumber(iblssj,2);

			      double siteqnumsi[]={Qtildei,Stildei,Sztildei};
			      double siteqnumsj[]={Qtildej,Stildej,Sztildej};

			      double Scfi=Si-Sztildei;
			      double Scfj=Sj-Sztildej;

			      double auxBasis[2]={0.0,0.0};
			      double auxCG[2]={0.0,0.0};

			      double CGnorm=1.0;

			      int sigma=0;

			      // Choose which sigma
			      auxCG[0]=CGordan(Si,Si,0.5,-0.5,Sj,Sj);
			      auxCG[1]=CGordan(Si,Si,0.5,0.5,Sj,Sj);
			      if ( dNEqual(auxCG[0],0.0) )
				{sigma=-1;CGnorm=auxCG[0];}
			      else
				{
				  if ( dNEqual(auxCG[1],0.0) )
				    {sigma=1;CGnorm=auxCG[1];}
				  else fnbasis(istbl,jstbl)=0.0;
				}

			      double dSigma=0.5*sigma;
			      // Now Sj=Si+dSigma necessarily
			      // Calculate
			      if (sigma!=0)
				{
				  for (double Szcfj=Scfj;
				       Szcfj>=-Scfj;Szcfj-=1.0)
				    {
				      Sztildei=Si-Szcfj; //Szcfi=Szcfj
				      Sztildej=Sj-Szcfj;
				      siteqnumsi[2]=Sztildei;
				      int blockstatei=pSingleSite->GetBlockFromQNumbers(siteqnumsi);
				      siteqnumsj[2]=Sztildej;
				      int blockstatej=pSingleSite->GetBlockFromQNumbers(siteqnumsj);

				      auxCG[0]=CGordan(Scfi,Szcfj,
						       Stildei,Sztildei,
						       Si,Si);
				      auxCG[1]=CGordan(Scfj,Szcfj,
						       Stildej,Sztildej,
						       Sj,Sj);

				      for (
			  int sitestatei=pSingleSite->GetBlockLimit(blockstatei,0);
			  sitestatei<=pSingleSite->GetBlockLimit(blockstatei,1);
			  sitestatei++)
					{
					  for (
			  int sitestatej=pSingleSite->GetBlockLimit(blockstatej,0);
			  sitestatej<=pSingleSite->GetBlockLimit(blockstatej,1);
			  sitestatej++)
					    {
					      auxBasis[0]+=TwoChQS_fd_table(ich,sigma,sitestatej,sitestatei)*auxCG[0]*auxCG[1]*TwoChQS_SpSm_table(sitestatei,typei)*TwoChQS_SpSm_table(sitestatej,typej);

					    }
					  // loop in blockstatej
					}
				      // loop in blockstatei
				    }
				  // end Loop in Szcf

				  //Calculate fbasis
				  fnbasis(istbl,jstbl)=auxBasis[0]/CGnorm;

				}
			      // end if (sigma!=0)

			    }
			  // end if icf=jcf
			  jstbl++;
			}
		      // end jst loop
		      istbl++;
		    }
		  // end ist loop
		  cout << "    ...fnbasis done." << endl;

		  cout << "    Multiplying BLAS matrices... " << endl;

		  fnw.resize(Nstibl,Nstjbl);
		  t.restart();
		  noalias(fnw)=prod (Zibl,
				     boost::numeric::ublas::matrix<double>(prod(fnbasis,trans(Zjbl))) );
		  time_elapsed=t.elapsed();
		  cout << "    ...done. Elapsed time:" << time_elapsed << endl;

// 		  cout << "Z(i="<<ibl<<")  : " <<  Zibl << endl;
// 		  cout << "Z(j="<<jbl<<")  : " <<  Zjbl << endl;
// 		  cout << "fbasis   :" <<  fnbasis << endl;
// 		  cout << "Zi.fbasis.ZjT : " <<  fnw << endl;

		  // Add to Qm1fNQ[ich-1]
		  for (int ii=0;ii<fnw.size1();ii++)
		    for (int jj=0;jj<fnw.size2();jj++)
		      Qm1fNQ[ich-1].MatEl.push_back(fnw(ii,jj));

		}
	      // end channel loop

	    }
	  //end if Q=Q'+1 etc

	}
      // end jbl loop

    }
  // end ibl loop

}
//////////////////////////////////
//////////////////////////////////
//////////////////////////////////
//////////////////////////////////

// 			      for (Sztildej=-Stildej;
// 				   Sztildej<=Stildej;Sztildej+=1.0)
// 				{
// 				  siteqnumsj[2]=Sztildej;
// 				  int sitestatej=pSingleSite->GetBlockFromQNumbers(siteqnumsj);
// 				  int i1=0;
// 				  for (double dSigma=-0.5;dSigma<=0.5;
// 				       dSigma+=1.0)
// 				    {
// 				      siteqnumsi[2]=Sztildej-dSigma;
// 				      int sitestatei=pSingleSite->GetBlockFromQNumbers(siteqnumsi);

// 				      auxCG[0]=CGordan(Scfi,Si-Sztildej+dSigma,
// 						       Stildei,Sztildej-dSigma,
// 						       Si,Si);
// 				      auxCG[1]=CGordan(Scfj,Sj-Sztildej,
// 						       Stildej,Sztildej,
// 						       Sj,Sj);
// 	  auxBasis[i1]+=TwoChQS_fd_table(ich,(int)(dSigma/0.5),sitestatej,sitestatei)*auxCG[0]*auxCG[1];

// 				      i1++;
// 				    }
// 				  // END loop in dSigma

// 				}
// 			      //END loop in Sztildej

			      // Check to see which of them is non-zero
// 			      auxCG[0]=CGordan(Si,Si,0.5,-0.5,Sj,Sj);
// 			      auxCG[1]=CGordan(Si,Si,0.5,0.5,Sj,Sj);
// 			      if ( !(dEqual(auxCG[0],0.0)) )
// 				fnbasis(istbl,jstbl)=auxBasis[0]/auxCG[0];
// 			      else
// 				{
// 				  if ( !(dEqual(auxCG[1],0.0)) )
// 				    fnbasis(istbl,jstbl)=auxBasis[1]/auxCG[1];
// 				  else fnbasis(istbl,jstbl)=0.0;
// 				}

			      // Loop in Sztildej,sigma:
			      //Sztildei=Sztildej-sigma
// 			      for (Sztildej=-Stildej;
// 				   Sztildej<=Stildej;Sztildej+=1.0)
// 				{
// 				  siteqnumsj[2]=Sztildej;
// 				  int sitestatej=pSingleSite->GetBlockFromQNumbers(siteqnumsj);
// 				  int i1=0;
// 				  for (double dSigma=-0.5;dSigma<=0.5;
// 				       dSigma+=1.0)
// 				    {
// 				      siteqnumsi[2]=Sztildej-dSigma;
// 				      int sitestatei=pSingleSite->GetBlockFromQNumbers(siteqnumsi);

// 				      auxCG[0]=CGordan(Scfi,Si-Sztildej+dSigma,
// 						       Stildei,Sztildej-dSigma,
// 						       Si,Si);
// 				      auxCG[1]=CGordan(Scfj,Sj-Sztildej,
// 						       Stildej,Sztildej,
// 						       Sj,Sj);
// 				      auxBasis[i1]+=TwoChQS_fd_table(ich,(int)(dSigma/0.5),sitestatej,sitestatei)*auxCG[0]*auxCG[1];

// 				      i1++;
// 				    }
// 				  // END loop in dSigma

// 				}
// 			      //END loop in Sztildej
// 			      // Check to see which of them is non-zero
// 			      auxCG[0]=CGordan(Si,Si,0.5,-0.5,Sj,Sj);
// 			      auxCG[1]=CGordan(Si,Si,0.5,0.5,Sj,Sj);
// 			      if ( !(dEqual(auxCG[0],0.0)) )
// 				fnbasis(istbl,jstbl)=auxBasis[0]/auxCG[0];
// 			      else
// 				{
// 				  if ( !(dEqual(auxCG[1],0.0)) )
// 				    fnbasis(istbl,jstbl)=auxBasis[1]/auxCG[1];
// 				  else fnbasis(istbl,jstbl)=0.0;
// 				}

// //fnbasis(istbl,jstbl)=TwoChQS_fd_table(ich,-1,typej,typei)+TwoChQS_fd_table(ich,1,typej,typei);


#ifndef BNI_LEARNING_K2_ALGORITHM_HPP
#define BNI_LEARNING_K2_ALGORITHM_HPP

#include <string>
#include <random>
#include <boost/range/algorithm_ext/erase.hpp>
#include <bayesian/graph.hpp>
#include <bayesian/sampler.hpp>
#include <bayesian/utility.hpp>

namespace bn {
namespace learning {

template<class Eval>
class k2_algorithm {
public:
    k2_algorithm(bn::sampler const& sampling)
        : sampling_(sampling), eval_(sampling_), engine_(make_engine<std::mt19937>())
    {
    }

    double operator()(graph_t& graph, std::unordered_map<vertex_type, std::vector<vertex_type>> preconditon)
    {
        // bestな構造を保持しておく
        sampling_.make_cpt(graph);
        double eval_best = eval_(graph);

        // 頂点の順番をランダムに
        auto vertexes = graph.vertex_list();
        std::shuffle(vertexes.begin(), vertexes.end(), engine_);

        for(auto const& target : vertexes)
        {
            // 自分自身もしくは与えられた除外条件以外のノードを親候補とする
            auto candidature = graph.vertex_list();
            boost::remove_erase_if(candidature,
                [&target, &preconditon](vertex_type const& node) -> bool
                {
                    // C言語でいう副作用完了点を悪用
                    auto ignore_nodes = preconditon.find(target);
                    return
                        node == target ||
                        (ignore_nodes != preconditon.end() &&
                        std::find(ignore_nodes->second.begin(), ignore_nodes->second.end(), node) != ignore_nodes->second.end());
                });

            for(auto const& parent : candidature)
            {
                if(auto edge = graph.add_edge(parent, target))
                {
                    sampling_.make_cpt(graph);
                    auto const eval_now = eval_(graph);

                    if(eval_now < eval_best)
                    {
                        eval_best = eval_now;
                        preconditon[parent].push_back(target);
                    }
                    else graph.erase_edge(edge);
                }
            }
        }

        return eval_best;
    }

private:
    sampler const& sampling_;
    Eval const eval_;
    std::mt19937 engine_;
};

} // namespace learning
} // namespace bn

#endif // BNI_LEARNING_K2_ALGORITHM_HPP


#ifndef ALGORITHM_SAUVOLA_HPP_41GTBVKV4
#define ALGORITHM_SAUVOLA_HPP_41GTBVKV4

#include <log4cxx/logger.h>
#include <log4cxx/basicconfigurator.h>
#include <log4cxx/helpers/exception.h>
#include <boost/foreach.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/unordered_map.hpp>
#include <boost/multi_array.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
//#include "grey_level_histogram.hpp"

//namespace std {using namespace __gnu_cxx;}
using namespace std;

namespace prhlt {
    class Algorithm_SAUVOLA{
        public:
            Algorithm_SAUVOLA(cv::Mat &ex_image);
            ~Algorithm_SAUVOLA();
            cv::Mat run(float ex_standard_deviation_constant, float ex_dynamic_range_constant, int ex_window_height, int ex_window_width);
        private:
            float standard_deviation_constant;
            float dynamic_range_constant;
            int window_height;
            int window_width;
            log4cxx::LoggerPtr logger;
            cv::Mat image;
            cv::Mat binarized_image;
            cv::Mat integral_image;
            cv::Mat integral_image_squared;
            void update_pixel(int x, int y);
            void initialize_integral_images();
            double average_grey_of_region(int from_x, int from_y);
            double standard_deviation_grey_of_region(double mean, int from_x, int from_y);
	};
}

#endif /* end of include guard*/


/* Hector -- A Simple Climate Model
   Copyright (C) 2014-2015  Battelle Memorial Institute

   Please see the accompanying file LICENSE.md for additional licensing
   information.
*/
/*
 *  temperature_component.cpp
 *  hector
 *
 *  Created by Ben Vega-Westhoff on 11/1/16.
 *
 */

#include <boost/lexical_cast.hpp>
#include <cmath>
#include <limits>
#define _USE_MATH_DEFINES

#include "temperature_component.hpp"
#include "core.hpp"
#include "h_util.hpp"
#include "simpleNbox.hpp"
#include "avisitor.hpp"

namespace Hector {

using namespace std;

//------------------------------------------------------------------------------
/*! \brief Constructor
 */
TemperatureComponent::TemperatureComponent() {
}

//------------------------------------------------------------------------------
/*! \brief Destructor
 */
TemperatureComponent::~TemperatureComponent(){
}

//------------------------------------------------------------------------------
// documentation is inherited
string TemperatureComponent::getComponentName() const {
    const string name = TEMPERATURE_COMPONENT_NAME;

    return name;
}

//------------------------------------------------------------------------------
/*! \brief              Calculates inverse of x and stores in y
 *  \param[in] x        Assume x is setup like x = [a,b,c,d] -> x = |a, b|
 *                                                                  |c, d|
 *  \param[out] y        Inverted 1-d matrix
 *  \returns            void, inverse is stored in y
 */
void TemperatureComponent::invert_1d_2x2_matrix(double * x, double * y) {
    double temp_d = (x[0]*x[3] - x[1]*x[2]);

    if(temp_d == 0) {
        H_THROW("Temperature: Matrix inversion divide by zero.");
    }
    double temp = 1/temp_d;
    y[0] = temp * x[3];
    y[1] = temp * -1 * x[1];
    y[2] = temp * -1 * x[2];
    y[3] = temp * x[0];

    return;
}

//------------------------------------------------------------------------------
// documentation is inherited
void TemperatureComponent::init( Core* coreptr ) {
    logger.open( getComponentName(), false, coreptr->getGlobalLogger().getEchoToFile(), coreptr->getGlobalLogger().getMinLogLevel() );
    H_LOG( logger, Logger::DEBUG ) << "hello " << getComponentName() << std::endl;

    tgaveq.set( 0.0, U_DEGC, 0.0 );
    tgav.set( 0.0, U_DEGC, 0.0 );
    flux_mixed.set( 0.0, U_W_M2, 0.0 );
    flux_interior.set( 0.0, U_W_M2, 0.0 );
    heatflux.set( 0.0, U_W_M2, 0.0 );

    core = coreptr;

    tgav_constrain.allowInterp( true );
    tgav_constrain.name = D_TGAV_CONSTRAIN;

    // Define the doeclim parameters
    diff.set( 0.55, U_CM2_S );    // default ocean heat diffusivity, cm2/s. value is CDICE default (varname is kappa there).
    S.set( 3.0, U_DEGC );         // default climate sensitivity, K (varname is t2co in CDICE).
    alpha.set( 1.0, U_UNITLESS);  // default aerosol scaling, unitless (similar to alpha in CDICE).
    volscl.set(1.0, U_UNITLESS);  // Default volcanic scaling, unitless (works the same way as alpha)

    // Register the data we can provide
    core->registerCapability( D_GLOBAL_TEMP, getComponentName() );
    core->registerCapability( D_LAND_TEMP, getComponentName() );
    core->registerCapability( D_OCEAN_AIR_TEMP, getComponentName() );
    core->registerCapability( D_OCEAN_SURFACE_TEMP, getComponentName() );
    core->registerCapability( D_GLOBAL_TEMPEQ, getComponentName() );
    core->registerCapability( D_FLUX_MIXED, getComponentName() );
    core->registerCapability( D_FLUX_INTERIOR, getComponentName() );
    core->registerCapability( D_HEAT_FLUX, getComponentName() );

    // Register our dependencies
    core->registerDependency( D_RF_TOTAL, getComponentName() );
    core->registerDependency( D_RF_BC, getComponentName() );
    core->registerDependency( D_RF_OC, getComponentName() );
    core->registerDependency( D_RF_SO2d, getComponentName() );
    core->registerDependency( D_RF_SO2i, getComponentName() );
    core->registerDependency( D_RF_VOL, getComponentName() );

    // Register the inputs we can receive from outside
    core->registerInput(D_ECS, getComponentName());
    core->registerInput(D_DIFFUSIVITY, getComponentName());
    core->registerInput(D_AERO_SCALE, getComponentName());
    core->registerInput(D_VOLCANIC_SCALE, getComponentName());
}

//------------------------------------------------------------------------------
// documentation is inherited
unitval TemperatureComponent::sendMessage( const std::string& message,
                                    const std::string& datum,
                                    const message_data info ) throw ( h_exception )
{
    unitval returnval;

    if( message==M_GETDATA ) {          //! Caller is requesting data
        return getData( datum, info.date );
    } else if( message==M_SETDATA ) {   //! Caller is requesting to set data
        setData(datum, info);
    } else {                        //! We don't handle any other messages
        H_THROW( "Caller sent unknown message: "+message );
    }

    return returnval;
}

//------------------------------------------------------------------------------
// documentation is inherited
void TemperatureComponent::setData( const string& varName,
                              const message_data& data ) throw ( h_exception )
{
    using namespace boost;

    H_LOG( logger, Logger::DEBUG ) << "Setting " << varName << "[" << data.date << "]=" << data.value_str << std::endl;

    try {
        if( varName == D_ECS ) {
            H_ASSERT( data.date == Core::undefinedIndex() , "date not allowed" );
            S = data.getUnitval(U_DEGC);
        } else if( varName == D_DIFFUSIVITY ) {
            H_ASSERT( data.date == Core::undefinedIndex(), "date not allowed" );
            diff = data.getUnitval(U_CM2_S);
	} else if( varName == D_AERO_SCALE ) {
            H_ASSERT( data.date == Core::undefinedIndex(), "date not allowed" );
            alpha = data.getUnitval(U_UNITLESS);
        } else if(varName == D_VOLCANIC_SCALE) {
            H_ASSERT( data.date == Core::undefinedIndex(), "date not allowed" );
            volscl = data.getUnitval(U_UNITLESS);
        } else if( varName == D_TGAV_CONSTRAIN ) {
            H_ASSERT( data.date != Core::undefinedIndex(), "date required" );
            tgav_constrain.set(data.date, data.getUnitval(U_DEGC));
        } else {
            H_THROW( "Unknown variable name while parsing " + getComponentName() + ": "
                    + varName );
        }
    } catch( bad_lexical_cast& castException ) {
        H_THROW( "Could not convert var: "+varName+", value: " + data.value_str + ", exception: "
                +castException.what() );
    } catch( h_exception& parseException ) {
        H_RETHROW( parseException, "Could not parse var: "+varName );
    }
}

//------------------------------------------------------------------------------
// documentation is inherited
// TO DO: should we put these in the ini file instead?
void TemperatureComponent::prepareToRun() throw ( h_exception ) {

    H_LOG( logger, Logger::DEBUG ) << "prepareToRun " << std::endl;

    if( tgav_constrain.size() ) {
        Logger& glog = core->getGlobalLogger();
        H_LOG( glog, Logger::WARNING ) << "Temperature will be overwritten by user-supplied values!" << std::endl;
    }

    // Initializing all model components that depend on the number of timesteps (ns)
    ns = core->getEndDate() - core->getStartDate() + 1;

    KT0 = std::vector<double>(ns, 0.0);
    KTA1 = std::vector<double>(ns, 0.0);
    KTB1 = std::vector<double>(ns, 0.0);
    KTA2 = std::vector<double>(ns, 0.0);
    KTB2 = std::vector<double>(ns, 0.0);
    KTA3 = std::vector<double>(ns, 0.0);
    KTB3 = std::vector<double>(ns, 0.0);

    Ker.resize(ns);

    temp.resize(ns);
    temp_landair.resize(ns);
    temp_sst.resize(ns);
    heatflux_mixed.resize(ns);
    heatflux_interior.resize(ns);
    heat_mixed.resize(ns);
    heat_interior.resize(ns);
    forcing.resize(ns);

    for(int i=0; i<3; i++) {
        B[i] = 0.0;
        C[i] = 0.0;
    }

    // DOECLIM parameters calculated from constants set in header
    ocean_area = (1.0 - flnd) * earth_area;    // m2
    cnum = rlam * flnd + bsi * (1.0 - flnd);   // factor from sea-surface climate sensitivity to global mean
    cden = rlam * flnd - ak * (rlam - bsi);    // intermediate parameter
    cfl = flnd * cnum / cden * q2co / S - bk * (rlam - bsi) / cden;      // land climate feedback parameter, W/m2/K
    cfs = (rlam * flnd - ak / (1.0 - flnd) * (rlam - bsi)) * cnum / cden * q2co / S + rlam * flnd / (1.0 - flnd) * bk * (rlam - bsi) / cden;                                // sea climate feedback parameter, W/m2/K
    kls = bk * rlam * flnd / cden - ak * flnd * cnum / cden * q2co / S;  // land-sea heat exchange coefficient, W/m2/K
    keff = kcon * diff;                                                  // ocean heat diffusivity, m2/yr
    taubot = pow(zbot,2) / keff;                                         // ocean bottom diffusion time scale, yr
    powtoheat = ocean_area * secs_per_Year / pow(10.0,22);               // convert flux to total ocean heat
    taucfs = cas / cfs;                                                  // sea climate feedback time scale, yr
    taucfl = cal / cfl;                                                  // land climate feedback time scale, yr
    taudif = pow(cas,2) / pow(csw,2) * M_PI / keff;                      // interior ocean heat uptake time scale, yr
    tauksl  = (1.0 - flnd) * cas / kls;                                  // sea-land heat exchange time scale, yr
    taukls  = flnd * cal / kls;                                          // land-sea heat exchange time scale, yr

    // Components of the analytical solution to the integral found in the temperature difference equation
    // Third order bottom correction terms will be "more than sufficient" for simulations out to 2500
    // (Equation A.25, EK05, or 2.3.23, TK07)

    // First order
    KT0[ns-1] = 4.0 - 2.0 * pow(2.0, 0.5);
    KTA1[ns-1] = -8.0 * exp(-taubot / double(dt)) + 4.0 * pow(2.0, 0.5) * exp(-0.5 * taubot / double(dt));
    KTB1[ns-1] = 4.0 * pow((M_PI * taubot / double(dt)), 0.5) * (1.0 + erf(pow(0.5 * taubot / double(dt), 0.5)) - 2.0 * erf(pow(taubot / double(dt), 0.5)));

    // Second order
    KTA2[ns-1] =  8.0 * exp(-4.0 * taubot / double(dt)) - 4.0 * pow(2.0, 0.5) * exp(-2.0 * taubot / double(dt));
    KTB2[ns-1] = -8.0 * pow((M_PI * taubot / double(dt)), 0.5) * (1.0 + erf(pow((2.0 * taubot / double(dt)), 0.5)) - 2.0 * erf(2.0 * pow((taubot / double(dt)), 0.5)) );

    // Third order
    KTA3[ns-1] = -8.0 * exp(-9.0 * taubot / double(dt)) + 4.0 * pow(2.0, 0.5) * exp(-4.5 * taubot / double(dt));
    KTB3[ns-1] = 12.0 * pow((M_PI * taubot / double(dt)), 0.5) * (1.0 + erf(pow((4.5 * taubot / double(dt)), 0.5)) - 2.0 * erf(3.0 * pow((taubot / double(dt)), 0.5)) );

    // Calculate the kernel component vectors
    for(int i=0; i<(ns-1); i++) {

        // First order
        KT0[i] = 4.0 * pow((double(ns-i)), 0.5) - 2.0 * pow((double(ns+1-i)), 0.5) - 2.0 * pow(double(ns-1-i), 0.5);
        KTA1[i] = -8.0 * pow(double(ns-i), 0.5) * exp(-taubot / double(dt) / double(ns-i)) + 4.0 * pow(double(ns+1-i), 0.5) * exp(-taubot / double(dt) / double(ns+1-i)) + 4.0 * pow(double(ns-1-i), 0.5) * exp(-taubot/double(dt) / double(ns-1-i));
        KTB1[i] =  4.0 * pow((M_PI * taubot / double(dt)), 0.5) * ( erf(pow((taubot / double(dt) / double(ns-1-i)), 0.5)) + erf(pow((taubot / double(dt) / double(ns+1-i)), 0.5)) - 2.0 * erf(pow((taubot / double(dt) / double(ns-i)), 0.5)) );

        // Second order
        KTA2[i] =  8.0 * pow(double(ns-i), 0.5) * exp(-4.0 * taubot / double(dt) / double(ns-i)) - 4.0 * pow(double(ns+1-i), 0.5) * exp(-4.0 * taubot / double(dt) / double(ns+1-i)) - 4.0 * pow(double(ns-1-i), 0.5) * exp(-4.0 * taubot / double(dt) / double(ns-1-i));
        KTB2[i] = -8.0 * pow((M_PI * taubot / double(dt)), 0.5) * ( erf(2.0 * pow((taubot / double(dt) / double(ns-1-i)), 0.5)) + erf(2.0 * pow((taubot / double(dt) / double(ns+1-i)), 0.5)) - 2.0 * erf(2.0 * pow((taubot / double(dt) / double(ns-i)), 0.5)) );

        // Third order
        KTA3[i] = -8.0 * pow(double(ns-i), 0.5) * exp(-9.0 * taubot / double(dt) / double(ns-i)) + 4.0 * pow(double(ns+1-i), 0.5) * exp(-9.0 * taubot / double(dt) / double(ns+1-i)) + 4.0 * pow(double(ns-1-i), 0.5) * exp(-9.0 * taubot / double(dt) / double(ns-1-i));
        KTB3[i] = 12.0 * pow((M_PI * taubot / double(dt)), 0.5) * ( erf(3.0 * pow((taubot / double(dt) / double(ns-1-i)), 0.5)) + erf(3.0 * pow((taubot / double(dt) / double(ns+1-i)), 0.5)) - 2.0 * erf(3.0 * pow((taubot / double(dt) / double(ns-i)), 0.5)) );
    }

    // Sum up the kernel components
    for(int i=0; i<ns; i++) {

        Ker[i] = KT0[i] + KTA1[i] + KTB1[i] + KTA2[i] + KTB2[i] + KTA3[i] + KTB3[i];

    }

    // Correction terms, remove oscillation artefacts due to short-term forcings
    // (Equation 2.3.27, TK07)
    C[0] = 1.0 / pow(taucfl, 2.0) + 1.0 / pow(taukls, 2.0) + 2.0 / taucfl / taukls + bsi / taukls / tauksl;
    C[1] = -1 * bsi / pow(taukls, 2.0) - bsi / taucfl / taukls - bsi / taucfs / taukls - pow(bsi, 2.0) / taukls / tauksl;
    C[2] = -1 * bsi / pow(tauksl, 2.0) - 1.0 / taucfs / tauksl - 1.0 / taucfl / tauksl -1.0 / taukls / tauksl;
    C[3] = 1.0 / pow(taucfs, 2.0) + pow(bsi, 2.0) / pow(tauksl, 2.0) + 2.0 * bsi / taucfs / tauksl + bsi / taukls / tauksl;
    for(int i=0; i<4; i++) {
        C[i] = C[i] * (pow(double(dt), 2.0) / 12.0);
    }

    //------------------------------------------------------------------
    // Matrices of difference equation system B*T(i+1) = Q(i) + A*T(i)
    // T = (TL,TS)
    // (Equations 2.3.24 and 2.3.27, TK07)
    B[0] = 1.0 + double(dt) / (2.0 * taucfl) + double(dt) / (2.0 * taukls);
    B[1] = double(-dt) / (2.0 * taukls) * bsi;
    B[2] = double(-dt) / (2.0 * tauksl);
    B[3] = 1.0 + double(dt) / (2.0 * taucfs) + double(dt) / (2.0 * tauksl) * bsi + 2.0 * fso * pow((double(dt) / taudif), 0.5);

    A[0] = 1.0 - double(dt) / (2.0 * taucfl) - double(dt) / (2.0 * taukls);
    A[1] = double(dt) / (2.0 * taukls) * bsi;
    A[2] = double(dt) / (2.0 * tauksl);
    A[3] = 1.0 - double(dt) / (2.0 * taucfs) - double(dt) / (2.0 * tauksl) * bsi + Ker[ns-1] * fso * pow((double(dt) / taudif), 0.5);
    for (int i=0; i<4; i++) {
        B[i] = B[i] + C[i];
        A[i] = A[i] + C[i];
    }

    // Calculate the inverse of B
    invert_1d_2x2_matrix(B, IB);
}

//------------------------------------------------------------------------------
// documentation is inherited
void TemperatureComponent::run( const double runToDate ) throw ( h_exception ) {
    H_LOG( logger, Logger::DEBUG ) << "temperature run " << runToDate << std::endl;

    // Commented section below is for case of user-specified temperature record.
    // Temperature component can't handle that at the moment!

    //// We track total radiative forcing using internal variable `internal_Ftot`
    //// Need to do this because if we're subject to a user constraint (being forced
    //// to match a temperature record), need to track the Ftot that *would* have
    //// produced the observed temperature record. This way there's a smooth
    //// transition when we exit the constraint period, after which internal_Ftot
    //// will rise in parallel with the value reported by ForcingComponent.

    // If we never had any temperature constraint, `internal_Ftot` will match `Ftot`.
    //
    // If the user has supplied temperature data, use that (except if past its end)
    if( tgav_constrain.size() && runToDate <= tgav_constrain.lastdate() ) {
    //    H_LOG( logger, Logger::NOTICE ) << "** Using user-supplied temperature" << std::endl;
        H_LOG( logger, Logger::SEVERE ) << "** ERROR - Temperature can't currently handle user-supplied temperature" << std::endl;
        H_THROW("User-supplied temperature not yet implemented.")
    }

    // Some needed inputs
    int tstep = runToDate - core->getStartDate();
    double aero_forcing =
        double(core->sendMessage( M_GETDATA, D_RF_BC ).value( U_W_M2 )) + double(core->sendMessage( M_GETDATA, D_RF_OC).value( U_W_M2 )) +
        double(core->sendMessage( M_GETDATA, D_RF_SO2d ).value( U_W_M2 )) + double(core->sendMessage( M_GETDATA, D_RF_SO2i ).value( U_W_M2 ));
    double volcanic_forcing = double(core->sendMessage(M_GETDATA, D_RF_VOL));

    forcing[tstep] = double(core->sendMessage(M_GETDATA, D_RF_TOTAL).value(U_W_M2))
                      - (1.0 - alpha) * aero_forcing
                      - (1.0 - volscl) * volcanic_forcing;

    // Initialize variables for time-stepping through the model
    double DQ1 = 0.0;
    double DQ2 = 0.0;
    double QC1 = 0.0;
    double QC2 = 0.0;
    double DelQL = 0.0;
    double DelQO = 0.0;
    double DPAST1 = 0.0;
    double DPAST2 = 0.0;
    double DTEAUX1 = 0.0;
    double DTEAUX2 = 0.0;

    // Reset the endogenous varibales for this time step
    temp[tstep] = 0.0;
    temp_landair[tstep] = 0.0;
    temp_sst[tstep] = 0.0;
    heat_mixed[tstep] = 0.0;
    heat_interior[tstep] = 0.0;
    heatflux_mixed[tstep] = 0.0;
    heatflux_interior[tstep] = 0.0;

    // Assume land and ocean forcings are equal to global forcing
    std::vector<double> QL = forcing;
    std::vector<double> QO = forcing;

    if (tstep > 0) {

        DelQL = QL[tstep] - QL[tstep - 1];
        DelQO = QO[tstep] - QO[tstep - 1];

        // Assume linear forcing change between tstep and tstep+1
        QC1 = (DelQL/cal*(1.0/taucfl+1.0/taukls)-bsi*DelQO/cas/taukls);
        QC2 = (DelQO/cas*(1.0/taucfs+bsi/tauksl)-DelQL/cal/tauksl);
        QC1 = QC1 * pow(double(dt), 2.0)/12.0;
        QC2 = QC2 * pow(double(dt), 2.0)/12.0;

        // ----------------- Initial Conditions --------------------
        // Initialization of temperature and forcing vector:
        // Factor 1/2 in front of Q in Equation A.27, EK05, and Equation 2.3.27, TK07 is a typo!
        // Assumption: linear forcing change between n and n+1
        DQ1 = 0.5*double(dt)/cal*(QL[tstep]+QL[tstep-1]);
        DQ2 = 0.5*double(dt)/cas*(QO[tstep]+QO[tstep-1]);
        DQ1 = DQ1 + QC1;
        DQ2 = DQ2 + QC2;

        // ---------- SOLVE MODEL ------------------
        // Calculate temperatures
        for(int i = 0; i <= tstep; i++) {
            DPAST2 = DPAST2 + temp_sst[i] * Ker[ns-tstep+i-1];
        }
        DPAST2 = DPAST2 * fso * pow((double(dt)/taudif), 0.5);

        DTEAUX1 = A[0] * temp_landair[tstep-1] + A[1] * temp_sst[tstep-1];
        DTEAUX2 = A[2] * temp_landair[tstep-1] + A[3] * temp_sst[tstep-1];

        temp_landair[tstep] = IB[0] * (DQ1 + DPAST1 + DTEAUX1) + IB[1] * (DQ2 + DPAST2 + DTEAUX2);
        temp_sst[tstep] = IB[2] * (DQ1 + DPAST1 + DTEAUX1) + IB[3] * (DQ2 + DPAST2 + DTEAUX2);
    }
    else {  // Handle the initial conditions
        temp_landair[0] = 0.0;
        temp_sst[0] = 0.0;
    }
    temp[tstep] = flnd * temp_landair[tstep] + (1.0 - flnd) * bsi * temp_sst[tstep];

    // Calculate ocean heat uptake [W/m^2]
    // heatflux[tstep] captures in the heat flux in the period between tstep-1 and tstep.
    // Numerical implementation of Equation 2.7, EK05, or Equation 2.3.13, TK07)
    // ------------------------------------------------------------------------
    if (tstep > 0) {
        heatflux_mixed[tstep] = cas*(temp_sst[tstep] - temp_sst[tstep-1]);
        for (int i=0; i < tstep; i++) {
            heatflux_interior[tstep] = heatflux_interior[tstep] + temp_sst[i]*Ker[ns-tstep+i];
        }
        heatflux_interior[tstep] = cas*fso/pow((taudif*dt), 0.5)*(2.0*temp_sst[tstep] - heatflux_interior[tstep]);
        heat_mixed[tstep] = heat_mixed[tstep-1] + heatflux_mixed[tstep] * (powtoheat*dt);
        heat_interior[tstep] = heat_interior[tstep-1] + heatflux_interior[tstep] * (fso*powtoheat*dt);
    }

    else {   // Handle the initial conditions
        heatflux_mixed[0] = 0.0;
        heatflux_interior[0] = 0.0;
        heat_mixed[0] = 0.0;
        heat_interior[0] = 0.0;
    }

    setoutputs(tstep);
    H_LOG( logger, Logger::DEBUG ) << " tgav=" << tgav << " in " << runToDate << std::endl;
}

//------------------------------------------------------------------------------
// documentation is inherited
unitval TemperatureComponent::getData( const std::string& varName,
                                const double date ) throw ( h_exception ) {

    unitval returnval;

    if(date == Core::undefinedIndex()) {
        // If no date is supplied, return the current value
        if( varName == D_GLOBAL_TEMP ) {
            returnval = tgav;
        } else if( varName == D_GLOBAL_TEMPEQ ) {
            returnval = tgaveq;
        } else if( varName == D_LAND_TEMP ) {
            returnval = tgav_land;
        } else if( varName == D_OCEAN_SURFACE_TEMP ) {
            returnval = tgav_sst;
        } else if( varName == D_OCEAN_AIR_TEMP ) {
            returnval = tgav_oceanair;
        } else if( varName == D_DIFFUSIVITY ) {
            returnval = diff;
        } else if( varName == D_AERO_SCALE ) {
	    returnval = alpha;
        } else if( varName == D_FLUX_MIXED ) {
	    returnval = flux_mixed;
        } else if( varName == D_FLUX_INTERIOR ) {
	    returnval = flux_interior;
        } else if( varName == D_HEAT_FLUX) {
            returnval = heatflux;
        } else if( varName == D_ECS ) {
            returnval = S;
        } else if(varName == D_VOLCANIC_SCALE) {
            returnval = volscl;
        } else {
            H_THROW( "Caller is requesting unknown variable: " + varName );
        }
    }
    else {
        // If a date is supplied, get the needed value from the
        // vectors.  Some values, such as model parameters, don't have
        // time-indexed values, so asking for one of those with a date
        // is an error.
        H_ASSERT(date <= core->getCurrentDate(), "Date must be <= current date.");
        int tstep = date - core->getStartDate();

        if( varName == D_GLOBAL_TEMP ) {
            returnval = unitval(temp[tstep], U_DEGC);
        } else if( varName == D_LAND_TEMP ) {
            returnval = unitval(temp_landair[tstep], U_DEGC);
        } else if( varName == D_OCEAN_SURFACE_TEMP ) {
            returnval = unitval(temp_sst[tstep], U_DEGC);
        } else if( varName == D_OCEAN_AIR_TEMP ) {
            returnval = bsi * unitval(temp_sst[tstep], U_DEGC);
        } else if( varName == D_GLOBAL_TEMPEQ ) {
            returnval = unitval(temp[tstep], U_DEGC);
        } else if( varName == D_FLUX_MIXED ) {
	    returnval = unitval(heatflux_mixed[tstep], U_W_M2);
        } else if( varName == D_FLUX_INTERIOR ) {
	    returnval = unitval(heatflux_interior[tstep], U_W_M2);
        } else if( varName == D_HEAT_FLUX) {
            double value = heatflux_mixed[tstep] + fso*heatflux_interior[tstep];
            returnval = unitval(value, U_W_M2);
        }
    }

    return returnval;
}

void TemperatureComponent::reset(double time) throw(h_exception)
{
    // We take a slightly different approach in this component's reset method than we have in other components.  The
    // temperature component doesn't have its own time counter, and it stores its history in a collection of double
    // vectors.  Therefore, all we do here is set the unitval versions of that stored data to their values from the
    // vectors.
    if(time < core->getStartDate()) // in this case, reset to the starting value.
        time = core->getStartDate();

    int tstep = time - core->getStartDate();
    setoutputs(tstep);
    H_LOG(logger, Logger::NOTICE)
        << getComponentName() << " reset to time= " << time << "\n";
}

//------------------------------------------------------------------------------
// documentation is inherited
void TemperatureComponent::shutDown() {
	H_LOG( logger, Logger::DEBUG ) << "goodbye " << getComponentName() << std::endl;
    logger.close();
}

//------------------------------------------------------------------------------
// documentation is inherited
void TemperatureComponent::accept( AVisitor* visitor ) {
    visitor->visit( this );
}

void TemperatureComponent::setoutputs(int tstep)
{
    double temp_oceanair;

    flux_mixed.set( heatflux_mixed[tstep], U_W_M2, 0.0 );
    flux_interior.set( heatflux_interior[tstep], U_W_M2, 0.0 );
    heatflux.set( heatflux_mixed[tstep] + fso * heatflux_interior[tstep], U_W_M2, 0.0 );
    tgav.set(temp[tstep], U_DEGC, 0.0);
    tgaveq.set(temp[tstep], U_DEGC, 0.0); // per comment line 140 of temperature_component.hpp
    tgav_land.set(temp_landair[tstep], U_DEGC, 0.0);
    tgav_sst.set(temp_sst[tstep], U_DEGC, 0.0);
    temp_oceanair = bsi * temp_sst[tstep];
    tgav_oceanair.set(temp_oceanair, U_DEGC, 0.0);
}

}


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

    @copyright 2016 NumScale SAS

  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_FUNCTION_EPS_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_EPS_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{
  /*!
    @ingroup group-ieee
    This object function computes the distance between its argument and the next
    representable value of its argument type.

    @see ulp, ulpdist, Eps

    @par Header <boost/simd/function/eps.hpp>

    @par Example:

      @snippet eps.cpp eps

    @par Possible output:

      @snippet eps.txt eps

  **/
  Value eps(Value const& x);
} }
#endif

#include <boost/simd/function/scalar/eps.hpp>
#include <boost/simd/function/simd/eps.hpp>

#endif


#include <iostream>
#include <fstream>
#include <vector>
#include <sstream>

#include <boost/timer/timer.hpp>
#include <mathgp/mathgp.h>
#include <mathgp/stdext.h>

#include <CL/opencl.h>

using namespace std;
using namespace mathgp;

const unsigned A_LOT = 200000;

// simple timer
class timer
{
public:
    void start(const char* name)
    {
        _name = name;
        t.start();
    }

    void report()
    {
        t.stop();
        boost::timer::nanosecond_type time = t.elapsed().wall;
        std::cout << _name << ": " << double(time)/1000000 << " ms" << std::endl;
    }

private:
    boost::timer::cpu_timer t;
    const char* _name;
};

// random from 0 to 1
float Rnd01()
{
    return float(rand()) / RAND_MAX;
}

// random from -1 to 1
float Rnd11()
{
    return Rnd01() * 2 - 1;
}

// load an entire file into a string
string LoadFile(const char* filename)
{
    ifstream fin(filename, ios::in | ios::binary);

    if (!fin.is_open())
    {
        ostringstream sout;
        sout << "cannot open " << filename;
        throw runtime_error(sout.str());
    }

    streamoff begin = fin.tellg();
    fin.seekg(0, ios::end);
    size_t fileSize = size_t(fin.tellg() - begin);
    fin.seekg(0, ios::beg);

    char* fileContents = new char[fileSize + 1];

    fin.read(fileContents, fileSize);

    fileContents[fileSize] = 0;

    return fileContents;
}

int main(int argc, char* argv[])
{
    timer t;

    ///////////////////////////////////////////////////////////////////////////////////////
    // initialization
    cl_uint numPlatforms = 0;
    clGetPlatformIDs(0, nullptr, &numPlatforms);
    vector<cl_platform_id> platforms(numPlatforms);
    clGetPlatformIDs(numPlatforms, &platforms.front(), nullptr);
    auto platform = platforms.front();

    auto deviceType = CL_DEVICE_TYPE_GPU;

    if (argc == 1 || strcmp(argv[1], "gpu") == 0)
    {
        cout << "Running on gpu..." << endl;

    }
    else if (strcmp(argv[1], "cpu") == 0)
    {
        cout << "Running on cpu..." << endl;
        deviceType = CL_DEVICE_TYPE_CPU;
    }
    else
    {
        cout << "Error: Unsupported device." << endl;
        return 1;
    }

    cl_uint numDevices = 0;
    clGetDeviceIDs(platform, deviceType, 0, nullptr, &numDevices);
    if (numDevices == 0)
    {
        cout << "Error: No supported devices" << endl;
        return 2;
    }

    vector<cl_device_id> devices(numDevices);
    clGetDeviceIDs(platform, deviceType, numDevices, &devices.front(), nullptr);
    auto device = devices.front();

    auto context = clCreateContext(nullptr, 1, &device, nullptr, nullptr, nullptr);
    auto commandQueue = clCreateCommandQueue(context, device, 0, nullptr);

    ///////////////////////////////////////////////////////////////////////////////////////
    // load cl program
    string clCode = LoadFile("matrix_mul.cl");
    const char* code = clCode.c_str();
    size_t size = clCode.length();
    auto program = clCreateProgramWithSource(context, 1, &code, &size, nullptr);
    cl_int status = clBuildProgram(program, 1, &device, nullptr, nullptr, nullptr);

    if (status != CL_SUCCESS)
    {
        // report build errors if any
        if (status == CL_BUILD_PROGRAM_FAILURE)
        {
            size_t buildLogSize = 0;
            clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, buildLogSize, nullptr, &buildLogSize);

            char* buildLog = new char[buildLogSize + 1];

            clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, buildLogSize, buildLog, NULL);

            cout << "Build errors: " << endl;
            cout << buildLog << endl;
            delete[] buildLog;
            return 5;
        }
    }

    ///////////////////////////////////////////////////////////////////////////////////////
    // generate input
    t.start("Generating input");
    vector<matrix> inputA(A_LOT);
    vector<matrix> inputB(A_LOT);
    for (unsigned i = 0; i < A_LOT; ++i)
    {
        auto& a = inputA[i];
        auto& b = inputB[i];
        for (int j = 0; j < 16; ++j)
        {
            a.at(j) = Rnd11();
            b.at(j) = Rnd11();
        }
    }
    t.report();

    ///////////////////////////////////////////////////////////////////////////////////////
    // software simulation
    t.start("Software simulation");
    vector<matrix> outputCPU(A_LOT);
    for (unsigned i = 0; i < A_LOT; ++i)
    {
        outputCPU[i] = inputA[i] * inputB[i];
    }
    t.report();

    ///////////////////////////////////////////////////////////////////////////////////////
    // prepare input for the cl program
    vector<matrix> outputCL(A_LOT);
    t.start("Working");
    auto inputBufferA = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, A_LOT * sizeof(matrix), inputA.front().as_array(), nullptr);
    auto inputBufferB = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, A_LOT * sizeof(matrix), inputB.front().as_array(), nullptr);
    auto outputBuffer = clCreateBuffer(context, CL_MEM_WRITE_ONLY, A_LOT * sizeof(matrix), nullptr, nullptr);

    auto kernel = clCreateKernel(program, "matrix_mul", nullptr);

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), &inputBufferA);
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), &inputBufferB);
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), &outputBuffer);

    ///////////////////////////////////////////////////////////////////////////////////////
    // execution
    size_t globalWorkItems = (A_LOT * 16)/10000;
    cl_event ndrEvent;
    clEnqueueNDRangeKernel(commandQueue, kernel, 1, nullptr, &globalWorkItems, nullptr, 0, nullptr, &ndrEvent);

    //size_t globalWorkItems[] = { 4, A_LOT * 4 };
    //size_t localWorkItems[] = { 4, 4 };
    //cl_event ndrEvent;
    //clEnqueueNDRangeKernel(commandQueue, kernel, 2, nullptr, globalWorkItems, localWorkItems, 0, nullptr, &ndrEvent);

    clFlush(commandQueue);
    clWaitForEvents(1, &ndrEvent);
    clReleaseEvent(ndrEvent);

    clEnqueueReadBuffer(commandQueue, outputBuffer, CL_TRUE, 0, A_LOT * sizeof(matrix), outputCL.front().as_array(), 0, nullptr, nullptr);

    t.report();
    ///////////////////////////////////////////////////////////////////////////////////////
    // check output
    t.start("Checking output");

    matrix cpu = matrix::zero();
    matrix opencl = matrix::zero();

    for (unsigned i = 0; i < A_LOT; ++i)
    {
        cpu += outputCPU[i];
        opencl += outputCL[i];
    }

    //cout << cpu << endl;
    //cout << opencl << endl;

    cout << "Matching results: " << boolalpha << (cpu == opencl) << endl;

    t.report();

    ///////////////////////////////////////////////////////////////////////////////////////
    // release resources
    clReleaseKernel(kernel);
    clReleaseProgram(program);
    clReleaseMemObject(inputBufferA);
    clReleaseMemObject(inputBufferB);
    clReleaseMemObject(outputBuffer);
    clReleaseCommandQueue(commandQueue);
    clReleaseContext(context);

    return 0;
}