// Copyright (C) 2007-2011 Garth N. Wells
//
// This file is part of DOLFIN.
//
// DOLFIN is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// DOLFIN is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
//
// Modified by Anders Logg 2008-2011
// Modified by Martin Alnes 2008
//
// First added:  2007-12-10
// Last changed: 2011-11-14

#include <string>
#include <boost/scoped_ptr.hpp>
#include <dolfin/common/NoDeleter.h>
#include <dolfin/fem/FiniteElement.h>
#include <dolfin/function/Function.h>
#include <dolfin/function/FunctionSpace.h>
#include <dolfin/function/GenericFunction.h>
#include <dolfin/log/log.h>
#include <dolfin/log/LogStream.h>
#include <dolfin/mesh/Mesh.h>
#include <dolfin/mesh/MeshData.h>
#include <dolfin/mesh/MeshEntity.h>
#include <dolfin/mesh/MeshFunction.h>
#include "Form.h"

using namespace dolfin;

//-----------------------------------------------------------------------------
Form::Form(uint rank, uint num_coefficients)
  : Hierarchical<Form>(*this),
    dx(*this), ds(*this), dS(*this),
    _function_spaces(rank), _coefficients(num_coefficients), _rank(rank)
{
  // Do nothing
}
//-----------------------------------------------------------------------------
Form::Form(boost::shared_ptr<const ufc::form> ufc_form,
           std::vector<boost::shared_ptr<const FunctionSpace> > function_spaces,
           std::vector<boost::shared_ptr<const GenericFunction> > coefficients)
  : Hierarchical<Form>(*this),
    dx(*this), ds(*this), dS(*this), _ufc_form(ufc_form),
    _function_spaces(function_spaces), _coefficients(coefficients),
    _rank(ufc_form->rank())
{
  // Do nothing
}
//-----------------------------------------------------------------------------
Form::~Form()
{
  // Do nothing
}
//-----------------------------------------------------------------------------
dolfin::uint Form::rank() const
{
  if (!_ufc_form)
    return _rank;
  else
  {
    dolfin_assert(_ufc_form->rank() == _rank);
    return _rank;
  }
}
//-----------------------------------------------------------------------------
dolfin::uint Form::num_coefficients() const
{
  if (!_ufc_form)
    return _coefficients.size();
  else
  {
    dolfin_assert(_ufc_form->num_coefficients() == _coefficients.size());
    return _coefficients.size();
  }
}
//-----------------------------------------------------------------------------
std::vector<dolfin::uint> Form::coloring(uint entity_dim) const
{
  warning("Form::coloring does not properly consider form type.");

  // Get mesh
  const Mesh& mesh = this->mesh();
  const uint cell_dim = mesh.topology().dim();

  std::vector<uint> _coloring;
  if (entity_dim == cell_dim)
  {
    _coloring.push_back(cell_dim);
    _coloring.push_back(0);
    _coloring.push_back(cell_dim);
  }
  else if (entity_dim == cell_dim - 1)
  {
    _coloring.push_back(cell_dim - 1);
    _coloring.push_back(cell_dim);
    _coloring.push_back(0);
    _coloring.push_back(cell_dim);
    _coloring.push_back(cell_dim - 1);
  }
  else
  {
    dolfin_error("Form.cpp",
                 "color form for multicore computing",
                 "Only cell and facet coloring are currently supported");
  }

  return _coloring;
}
//-----------------------------------------------------------------------------
void Form::set_mesh(boost::shared_ptr<const Mesh> mesh)
{
  _mesh = mesh;
}
//-----------------------------------------------------------------------------
const Mesh& Form::mesh() const
{
  // In the case when there are no function spaces (in the case of a
  // a functional) the (generated) subclass must set the mesh directly
  // by calling set_mesh().

  // Extract meshes from function spaces
  std::vector<boost::shared_ptr<const Mesh> > meshes;
  for (uint i = 0; i < _function_spaces.size(); i++)
  {
    if (_function_spaces[i])
    {
      dolfin_assert(_function_spaces[i]->mesh());
      meshes.push_back(_function_spaces[i]->mesh());
    }
  }

  // Add common mesh if any
  if (_mesh)
    meshes.push_back(_mesh);

  // Extract meshes from coefficients. Note that this is only done
  // when we don't already have a mesh sine it may otherwise conflict
  // with existing meshes (if coefficient is defined on another mesh).
  if (meshes.size() == 0)
  {
    for (uint i = 0; i < _coefficients.size(); i++)
    {
      const Function* function = dynamic_cast<const Function*>(&*_coefficients[i]);
      if (function && function->function_space()->mesh())
        meshes.push_back(function->function_space()->mesh());
    }
  }

  // Check that we have at least one mesh
  if (meshes.size() == 0)
  {
    dolfin_error("Form.cpp",
                 "extract mesh from form",
                 "No mesh was found. Try passing mesh to the assemble function");
  }

  // Check that all meshes are the same
  for (uint i = 1; i < meshes.size(); i++)
  {
    if (meshes[i] != meshes[i - 1])
    {
      dolfin_error("Form.cpp",
                   "extract mesh from form",
                   "Non-matching meshes for function spaces");
    }
  }

  // Return first mesh
  dolfin_assert(meshes[0]);
  return *meshes[0];
}
//-----------------------------------------------------------------------------
boost::shared_ptr<const dolfin::Mesh> Form::mesh_shared_ptr() const
{
  return _mesh;
}
//-----------------------------------------------------------------------------
boost::shared_ptr<const FunctionSpace> Form::function_space(uint i) const
{
  dolfin_assert(i < _function_spaces.size());
  return _function_spaces[i];
}
//-----------------------------------------------------------------------------
std::vector<boost::shared_ptr<const FunctionSpace> > Form::function_spaces() const
{
  return _function_spaces;
}
//-----------------------------------------------------------------------------
void Form::set_coefficient(uint i,
                           boost::shared_ptr<const GenericFunction> coefficient)
{
  dolfin_assert(i < _coefficients.size());
  _coefficients[i] = coefficient;
}
//-----------------------------------------------------------------------------
void Form::set_coefficient(std::string name,
                           boost::shared_ptr<const GenericFunction> coefficient)
{
  set_coefficient(coefficient_number(name), coefficient);
}
//-----------------------------------------------------------------------------
void Form::set_coefficients(std::map<std::string, boost::shared_ptr<const GenericFunction> > coefficients)
{
  std::map<std::string, boost::shared_ptr<const GenericFunction> >::iterator it;
  for (it = coefficients.begin(); it != coefficients.end(); ++it)
    set_coefficient(it->first, it->second);
}
//-----------------------------------------------------------------------------
boost::shared_ptr<const GenericFunction> Form::coefficient(uint i) const
{
  dolfin_assert(i < _coefficients.size());
  return _coefficients[i];
}
//-----------------------------------------------------------------------------
boost::shared_ptr<const GenericFunction> Form::coefficient(std::string name) const
{
  return coefficient(coefficient_number(name));
}
//-----------------------------------------------------------------------------
std::vector<boost::shared_ptr<const GenericFunction> > Form::coefficients() const
{
  return _coefficients;
}
//-----------------------------------------------------------------------------
dolfin::uint Form::coefficient_number(const std::string & name) const
{
  // TODO: Dissect name, assuming "wi", and return i.
  dolfin_not_implemented();
  return 0;
}
//-----------------------------------------------------------------------------
std::string Form::coefficient_name(uint i) const
{
  // Create name like "w0", overloaded by Form subclasses generated by form compilers
  std::ostringstream name;
  name << "w" << i;
  return name.str();
}
//-----------------------------------------------------------------------------
boost::shared_ptr<const MeshFunction<dolfin::uint> >
Form::cell_domains_shared_ptr() const
{
  return _cell_domains;
}
//-----------------------------------------------------------------------------
boost::shared_ptr<const MeshFunction<dolfin::uint> >
Form::exterior_facet_domains_shared_ptr() const
{
  return _exterior_facet_domains;
}
//-----------------------------------------------------------------------------
boost::shared_ptr<const MeshFunction<dolfin::uint> >
Form::interior_facet_domains_shared_ptr() const
{
  return _interior_facet_domains;
}
//-----------------------------------------------------------------------------
void Form::set_cell_domains
(boost::shared_ptr<const MeshFunction<uint> > cell_domains)
{
  _cell_domains = cell_domains;
}
//-----------------------------------------------------------------------------
void Form::set_exterior_facet_domains
(boost::shared_ptr<const MeshFunction<uint> > exterior_facet_domains)
{
  _exterior_facet_domains = exterior_facet_domains;
}
//-----------------------------------------------------------------------------
void Form::set_interior_facet_domains
(boost::shared_ptr<const MeshFunction<uint> > interior_facet_domains)
{
  _interior_facet_domains = interior_facet_domains;
}
//-----------------------------------------------------------------------------
boost::shared_ptr<const ufc::form> Form::ufc_form() const
{
  return _ufc_form;
}
//-----------------------------------------------------------------------------
void Form::check() const
{
  dolfin_assert(_ufc_form);

  // Check that the number of argument function spaces is correct
  if (_ufc_form->rank() != _function_spaces.size())
  {
    dolfin_error("Form.cpp",
                 "assemble form",
                 "Expecting %d function spaces (not %d)",
                 _ufc_form->rank(), _function_spaces.size());
  }

  // Check that the number of coefficient function spaces is correct
  if (_ufc_form->num_coefficients() != _coefficients.size())
  {
   dolfin_error("Form.cpp",
                "assemble form",
                "Expecting %d coefficient (not %d)",
                _ufc_form->num_coefficients(), _coefficients.size());
  }

  // Check argument function spaces
  for (uint i = 0; i < _function_spaces.size(); ++i)
  {
    boost::scoped_ptr<ufc::finite_element> element(_ufc_form->create_finite_element(i));
    dolfin_assert(element);
    dolfin_assert(_function_spaces[i]->element());
    if (element->signature() != _function_spaces[i]->element()->signature())
    {
      log(ERROR, "Expected element: %s", element->signature());
      log(ERROR, "Input element:    %s", _function_spaces[i]->element()->signature().c_str());
      dolfin_error("Form.cpp",
                   "assemble form",
                   "Wrong type of function space for argument %d", i);
    }
  }
}
//-----------------------------------------------------------------------------
Equation Form::operator==(const Form& rhs) const
{
  Equation equation(reference_to_no_delete_pointer(*this),
                    reference_to_no_delete_pointer(rhs));
  return equation;
}
//-----------------------------------------------------------------------------
Equation Form::operator==(int rhs) const
{
  Equation equation(reference_to_no_delete_pointer(*this), 0);
  return equation;
}
//-----------------------------------------------------------------------------


// Copyright (c) 2014, Sailing Lab
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice,
// this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the <ORGANIZATION> nor the names of its contributors
// may be used to endorse or promote products derived from this software
// without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
/*
 * Least-squares Matrix Factorization application built on Petuum Parameter Server.
 *
 * Author: qho
 */
#include <vector>
#include <fstream>
#include <cstdint>
#include <random>

#include <petuum_ps_common/include/petuum_ps.hpp>
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <pthread.h>

#include <boost/date_time/posix_time/posix_time.hpp>

#include "matrixloader.hpp"

// Command-line flags
DEFINE_string(hostfile, "", "Path to Petuum PS server configuration file");
DEFINE_int32(ps_row_in_memory_limit, 1000, "Single-machine version only: max # rows of L, R that can be held in memory");
DEFINE_int32(ps_row_cache_size, 0, "Max # rows of L, R that can be cached per client; set to 0 (default) to cache all rows. If you are running out of memory due to big L, R, consider limiting the cache size. No effect for single-machine version.");
DEFINE_string(datafile, "", "Input sparse matrix");
DEFINE_string(offsetsfile, "", "Offsets file for datafile - if provided, the program will run in disk-based mode to save memory");
DEFINE_string(output_prefix, "", "Output results with this prefix. If the prefix is X, the output will be called X.L and X.R");
DEFINE_double(lambda, 0.0, "L2 regularization strength; 0 disables");
DEFINE_double(init_step_size, 0.5, "SGD step size at iteration t is init_step_size * (step_size_offset + t)^(-step_size_pow)");
DEFINE_double(step_size_offset, 100.0, "SGD step size at iteration t is init_step_size * (step_size_offset + t)^(-step_size_pow)");
DEFINE_double(step_size_pow, 0.5, "SGD step size at iteration t is init_step_size * (step_size_offset + t)^(-step_size_pow)");
DEFINE_int32(rng_seed, 967234, "RNG seed for initializing factor matrices L, R");
DEFINE_int32(num_clients, 1, "Total number of clients");
DEFINE_int32(num_worker_threads, 1, "Number of worker threads per client");
DEFINE_int32(client_id, 0, "This client's ID");
DEFINE_int32(K, 100, "Factorization rank");
DEFINE_int32(num_iterations, 100, "Number of iterations");
DEFINE_int32(staleness, 0, "Staleness");
DEFINE_string(ps_stats_path, "", "Petuum PS statistics output file");
DEFINE_int32(ps_snapshot_clock, 0, "How often to take PS snapshots; 0 disables");
DEFINE_int32(ps_resume_clock, 0, "If nonzero, resume from the PS snapshot with this ID");
DEFINE_string(ps_snapshot_dir, "", "Where to save PS snapshots");
DEFINE_string(ps_resume_dir, "", "Where to retrieve PS snapshots for resuming");

// Typedefs
typedef std::pair<int,float> int_pair_float_t;

// Data variables
matrixloader::AbstractMatrixLoader* data_matrix;

// Returns the number of workers threads across all clients
int get_total_num_workers() {
  return FLAGS_num_clients * FLAGS_num_worker_threads;
}
// Returns the global thread ID of this worker
int get_global_worker_id(int local_thread_id) {
  return (FLAGS_client_id * FLAGS_num_worker_threads) + local_thread_id;
}

// Performs stochastic gradient descent on the (i,j)-th element of the matrix,
// which takes value Xij.
void sgd_element(const int i, const int j, const float Xij, float step_size, int global_worker_id,
                 petuum::Table<float>& L_table,
                 petuum::Table<float>& R_table,
                 petuum::Table<float>& loss_table) {
  // Read L(i,:) and R(:,j) from Petuum PS
  petuum::RowAccessor Li_acc;
  petuum::RowAccessor Rj_acc;
  L_table.Get(i, &Li_acc);
  R_table.Get(j, &Rj_acc);
  const petuum::DenseRow<float>& Li = Li_acc.Get<petuum::DenseRow<float> >();
  const petuum::DenseRow<float>& Rj = Rj_acc.Get<petuum::DenseRow<float> >();

  // Compute L(i,:) * R(:,j)
  float LiRj = 0.0;
  for (int k = 0; k < FLAGS_K; ++k) {
    LiRj += Li[k] * Rj[k];
  }
  // Update the loss function (does not include L2 regularizer term)
  loss_table.Inc(0,global_worker_id,pow(Xij - LiRj, 2));

  // Now update L(i,:) and R(:,j) based on the loss function at X(i,j).
  // The non-regularized loss function at X(i,j) is ( X(i,j) - L(i,:)*R(:,j) )^2.
  //
  // The non-regularized gradient w.r.t. L(i,k) is -2*X(i,j)R(k,j) + 2*L(i,:)*R(:,j)*R(k,j).
  // The non-regularized gradient w.r.t. R(k,j) is -2*X(i,j)L(i,k) + 2*L(i,:)*R(:,j)*L(i,k).
  petuum::UpdateBatch<float> Li_update;
  petuum::UpdateBatch<float> Rj_update;
  for (int k = 0; k < FLAGS_K; ++k) {
    float gradient = 0.0;
    // Compute update for L(i,k)
    gradient = -2 * (Xij - LiRj) * Rj[k] + FLAGS_lambda * 2 * Li[k];
    Li_update.Update(k, -gradient * step_size);
    // Compute update for R(k,j)
    gradient = -2 * (Xij - LiRj) * Li[k] + FLAGS_lambda * 2 * Rj[k];
    Rj_update.Update(k, -gradient * step_size);
  }
  // Commit updates to Petuum PS
  L_table.BatchInc(i,Li_update);
  R_table.BatchInc(j,Rj_update);
}

// Initialize the Matrix Factorization solver
void init_mf(petuum::Table<float>& L_table, petuum::Table<float>& R_table) {
  // Uniform RNG that produces numbers in [-1,1)
  std::mt19937 generator(FLAGS_rng_seed);
  std::uniform_real_distribution<> dist(-1.,1.);
  // Add a random initialization in [-1,1)/num_workers to each element of L and R
  const int num_workers = get_total_num_workers();
  for (int i = 0; i < data_matrix->getN(); ++i) {
    petuum::UpdateBatch<float> L_updates;
    for (int k = 0; k < FLAGS_K; ++k) {
      double init_val = dist(generator) / num_workers;
      L_updates.Update(k,init_val);
    }
    L_table.BatchInc(i,L_updates);
  }
  for (int j = 0; j < data_matrix->getM(); ++j) {
    petuum::UpdateBatch<float> R_updates;
    for (int k = 0; k < FLAGS_K; ++k) {
      double init_val = dist(generator) / num_workers;
      R_updates.Update(k,init_val);
    }
    R_table.BatchInc(j,R_updates);
  }
}

// Outputs L_table and R_table to disk
void output_to_disk(petuum::Table<float>& L_table, petuum::Table<float>& R_table) {
  // L_table
  std::string L_file = FLAGS_output_prefix + ".L";
  std::ofstream L_stream(L_file.c_str());
  for (uint32_t i = 0; i < data_matrix->getN(); ++i) {
    petuum::RowAccessor Li_acc;
    L_table.Get(i, &Li_acc);
    const petuum::DenseRow<float>& Li = Li_acc.Get<petuum::DenseRow<float> >();
    for (uint32_t k = 0; k < FLAGS_K; ++k) {
      L_stream << Li[k] << " ";
    }
    L_stream << "\n";
  }
  L_stream.close();
  // R_table
  std::string R_file = FLAGS_output_prefix + ".R";
  std::ofstream R_stream(R_file.c_str());
  for (uint32_t j = 0; j < data_matrix->getM(); ++j) {
    petuum::RowAccessor Rj_acc;
    R_table.Get(j, &Rj_acc);
    const petuum::DenseRow<float>& Rj = Rj_acc.Get<petuum::DenseRow<float> >();
    for (uint32_t k = 0; k < FLAGS_K; ++k) {
      R_stream << Rj[k] << " ";
    }
    R_stream << "\n";
  }
  R_stream.close();
}

// Main Matrix Factorization routine, called by pthread_create
void* solve_mf(void* args) {
  int local_thread_id = *reinterpret_cast<int*>(args);
  // Register this thread with Petuum PS
  petuum::PSTableGroup::RegisterThread();
  // Get tables
  petuum::Table<float> L_table = petuum::PSTableGroup::GetTableOrDie<float>(0);
  petuum::Table<float> R_table = petuum::PSTableGroup::GetTableOrDie<float>(1);
  petuum::Table<float> loss_table = petuum::PSTableGroup::GetTableOrDie<float>(2);
  // Initialize MF solver
  const int total_num_workers = get_total_num_workers();
  const int global_worker_id = get_global_worker_id(local_thread_id); // global_worker_id lies in the range [0,total_num_workers)
  STATS_APP_INIT_BEGIN();
  if (FLAGS_ps_resume_clock == 0) { // Only initialize L, R tables if not resuming
    init_mf(L_table,R_table);
  } else if (global_worker_id == 0) {
    std::cout << "...Resuming at Iteration " << FLAGS_ps_resume_clock-FLAGS_staleness
              << " (Snapshot " << FLAGS_ps_resume_clock << ")" << std::endl;
  }
  petuum::PSTableGroup::GlobalBarrier();  // Let stale values finish propagating (performs staleness+1 clock()s)
  STATS_APP_INIT_END();
  // Run MF solver
  for (int iter = 0; iter < FLAGS_num_iterations; ++iter) {
    // If resuming, skip iterations until we reach the correct iteration
    if (iter < FLAGS_ps_resume_clock-FLAGS_staleness-1) {
      petuum::PSTableGroup::Clock();
      continue;
    }

    // Begin iteration
    boost::posix_time::ptime beginT = boost::posix_time::microsec_clock::local_time();
    if (global_worker_id == 0) {
      std::cout << "Iteration " << iter+1 << "/" << FLAGS_num_iterations << "... " << std::flush;
    }
    // Clear loss function table
    petuum::RowAccessor loss_acc;
    loss_table.Get(0, &loss_acc);
    const petuum::DenseRow<float>& loss_row = loss_acc.Get<petuum::DenseRow<float> >();
    loss_table.Inc(0,global_worker_id,-loss_row[global_worker_id]);

    // Divide matrix elements across workers, and perform SGD
    bool last_el = false;
    int64_t row,col;
    float val;
    STATS_APP_ACCUM_COMP_BEGIN();
    float step_size = FLAGS_init_step_size * pow(FLAGS_step_size_offset + iter, -FLAGS_step_size_pow);
    while (last_el == false) {
      data_matrix->getNextEl(global_worker_id,row,col,val,last_el);
      sgd_element(row,col,val,step_size,global_worker_id,L_table,R_table,loss_table);
    }
    STATS_APP_ACCUM_COMP_END();

    // Output loss function (does not include L2 regularizer term)
    if (global_worker_id == 0) {
      petuum::RowAccessor loss_acc;
      loss_table.Get(0, &loss_acc);
      const petuum::DenseRow<float>& loss_row = loss_acc.Get<petuum::DenseRow<float> >();
      float loss = 0.0;
      for (int t = 0; t < total_num_workers; ++t) {
        loss += loss_row[t];
      }
      std::cout << "loss function = " << loss << "... " << std::flush;
    }
    // Advance Parameter Server iteration
    petuum::PSTableGroup::Clock();
    boost::posix_time::time_duration elapTime = boost::posix_time::microsec_clock::local_time() - beginT;
    if (global_worker_id == 0) {
      std::cout << "elapsed time = " << ((float) elapTime.total_milliseconds()) / 1000 << std::endl;
    }
  }
  // Let stale values finish propagating (performs staleness+1 clock()s)
  petuum::PSTableGroup::GlobalBarrier();
  // Output results to disk
  if (global_worker_id == 0) {
    std::cout << "Outputting results to prefix " << FLAGS_output_prefix << " ... " << std::flush;
    output_to_disk(L_table,R_table);
    std::cout << "done" << std::endl;
  }
  // Deregister this thread with Petuum PS
  petuum::PSTableGroup::DeregisterThread();
  return 0;
}

// Main function
int main(int argc, char *argv[]) {
  google::ParseCommandLineFlags(&argc, &argv, true);
  google::InitGoogleLogging(argv[0]);
  boost::posix_time::ptime beginT = boost::posix_time::microsec_clock::local_time();

  // Configure Petuum PS
  petuum::TableGroupConfig table_group_config;
  // Global parameters for the whole PS system
  table_group_config.num_total_server_threads = FLAGS_num_clients;  // 1 server thread per client
  table_group_config.num_total_bg_threads = FLAGS_num_clients;  // 1 background thread per client
  table_group_config.num_total_clients = FLAGS_num_clients;
  table_group_config.num_tables = 3;  // L_table, R_table, loss_table
  // Single-node-PS versus multi-machine PS
#ifdef PETUUM_SINGLE_NODE
  table_group_config.ooc_path_prefix = FLAGS_output_prefix + ".matrixfact_sn.localooc";
  table_group_config.consistency_model = petuum::LocalOOC;
#else
  petuum::GetHostInfos(FLAGS_hostfile, &table_group_config.host_map);
  petuum::GetServerIDsFromHostMap(&table_group_config.server_ids, table_group_config.host_map);
  table_group_config.consistency_model = petuum::SSPPush;
#endif
  // Local parameters for this process
  table_group_config.num_local_server_threads = 1;
  table_group_config.num_local_bg_threads = 1;
  table_group_config.num_local_app_threads = FLAGS_num_worker_threads + 1;  // +1 for main() thread
  table_group_config.client_id = FLAGS_client_id;
  // Stats and snapshots
  table_group_config.stats_path = FLAGS_ps_stats_path;
  table_group_config.snapshot_clock = FLAGS_ps_snapshot_clock;  // Only for multi-machine PS
  table_group_config.resume_clock = FLAGS_ps_resume_clock;  // Only for multi-machine PS
  table_group_config.snapshot_dir = FLAGS_ps_snapshot_dir;  // Only for multi-machine PS
  table_group_config.resume_dir = FLAGS_ps_resume_dir;  // Only for multi-machine PS
  // Configure PS row types
  petuum::PSTableGroup::RegisterRow<petuum::DenseRow<float> >(0);  // Register dense rows as ID 0
  // Start PS
  // IMPORTANT: This command starts up the name node service on client 0.
  //            We therefore do it ASAP, before other lengthy actions like
  //            loading data.
  petuum::PSTableGroup::Init(table_group_config, false);  // Initializing thread does not need table access

  // Load Data
  STATS_APP_LOAD_DATA_BEGIN();
  if (FLAGS_offsetsfile.size() == 0) { // Load the full matrix
    if (FLAGS_client_id == 0) {
      std::cout << "Data mode: Loading matrix " << FLAGS_datafile << " into memory..." << std::endl;
    }
    data_matrix = new matrixloader::StandardMatrixLoader(FLAGS_datafile,get_total_num_workers());
  } else { // Disk-based loader
    if (FLAGS_client_id == 0) {
      std::cout << "Data mode: Streaming matrix " << FLAGS_datafile << " from disk using offsets " << FLAGS_offsetsfile << std::endl;
    }
    auto dml = new matrixloader::DiskMatrixLoader(FLAGS_datafile,FLAGS_offsetsfile,get_total_num_workers());
    for (int t = 0; t < FLAGS_num_worker_threads; ++t) {  // Initialize workers
      dml->initWorker(get_global_worker_id(t));
    }
    data_matrix = dml;
  }
  STATS_APP_LOAD_DATA_END();

  // Output statistics
  if (FLAGS_client_id == 0) {
    std::cout << "Matrix dimensions: " << data_matrix->getN() << " by " << data_matrix->getM() << std::endl;
    std::cout << "# non-missing entries: " << data_matrix->getNNZ() << std::endl;
    std::cout << "Factorization rank: " << FLAGS_K << std::endl;
    std::cout << "# client machines: " << FLAGS_num_clients << std::endl;
    std::cout << "# worker threads per client: " << FLAGS_num_worker_threads << std::endl;
    std::cout << "SSP staleness: " << FLAGS_staleness << std::endl;
    std::cout << "Step size formula: " << FLAGS_init_step_size << " * (" << FLAGS_step_size_offset << " + t)^(-" << FLAGS_step_size_pow << ")" << std::endl;
    std::cout << "Regularization strength lambda: " << FLAGS_lambda << std::endl;
    std::cout << "  (Note: displayed loss function does not include regularization term)" << std::endl;
  }

  // Configure PS tables
  petuum::ClientTableConfig table_config;
  table_config.table_info.row_type = 0; // Dense rows
  table_config.oplog_capacity = 100;
  // L_table (N by K)
  table_config.table_info.table_staleness = FLAGS_staleness;
  table_config.table_info.row_capacity = FLAGS_K;
#ifdef PETUUM_SINGLE_NODE
  table_config.process_cache_capacity = FLAGS_ps_row_in_memory_limit;
#else
  if (FLAGS_ps_row_cache_size == 0) {
    table_config.process_cache_capacity = data_matrix->getN();
  } else {
    table_config.process_cache_capacity = FLAGS_ps_row_cache_size;
  }
#endif
  petuum::PSTableGroup::CreateTable(0,table_config);
  // R_table (M by K)
  table_config.table_info.table_staleness = FLAGS_staleness;
  table_config.table_info.row_capacity = FLAGS_K;
#ifdef PETUUM_SINGLE_NODE
  table_config.process_cache_capacity = FLAGS_ps_row_in_memory_limit;
#else
  if (FLAGS_ps_row_cache_size == 0) {
    table_config.process_cache_capacity = data_matrix->getM();
  } else {
    table_config.process_cache_capacity = FLAGS_ps_row_cache_size;
  }
#endif
  petuum::PSTableGroup::CreateTable(1,table_config);
  // loss_table (1 by total # workers)
  table_config.table_info.table_staleness = 0;  // No staleness for loss table
  table_config.table_info.row_capacity = get_total_num_workers();
  table_config.process_cache_capacity = 1;
  petuum::PSTableGroup::CreateTable(2,table_config);
  // Finished creating tables
  petuum::PSTableGroup::CreateTableDone();

  // Run Petuum PS-based MF solver
  pthread_t* threads = new pthread_t[FLAGS_num_worker_threads];
  int* local_thread_ids = new int[FLAGS_num_worker_threads];
  for (int t = 0; t < FLAGS_num_worker_threads; ++t) {
    local_thread_ids[t] = t;
    pthread_create(threads+t, NULL, solve_mf, local_thread_ids+t);
  }
  petuum::PSTableGroup::WaitThreadRegister();
  for (int t = 0; t < FLAGS_num_worker_threads; ++t) {
    pthread_join(threads[t], NULL);
  }
  delete threads;
  delete local_thread_ids;
  delete data_matrix;

  // Cleanup and output runtime
  petuum::PSTableGroup::ShutDown();
  boost::posix_time::time_duration elapTime = boost::posix_time::microsec_clock::local_time() - beginT;
  if (FLAGS_client_id == 0) {
    std::cout << "total runtime = " << ((float) elapTime.total_milliseconds()) / 1000 << "s" << std::endl;
  }
  return 0;
}


/*
  Author(s):      Robert Patterson and Markus Sander
  Project:        sweepc (population balance solver)
  Sourceforge:    http://sourceforge.net/projects/mopssuite

  Copyright (C) 2008 Matthew S Celnik.

  File purpose:
    Implementation of the DimerInception class declared in the
    swp_DimerInception.h header file.

  Licence:
    This file is part of "sweepc".

    sweepc is free software; you can redistribute it and/or
    modify it under the terms of the GNU Lesser General Public License
    as published by the Free Software Foundation; either version 2
    of the License, or (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

  Contact:
    Prof Markus Kraft
    Dept of Chemical Engineering
    University of Cambridge
    New Museums Site
    Pembroke Street
    Cambridge
    CB2 3RA
    UK

    Email:       mk306@cam.ac.uk
    Website:     http://como.cheng.cam.ac.uk
*/

#include "swp_dimer_inception.h"

#include "swp_mechanism.h"
#include "swp_primary.h"

#include <boost/random/uniform_01.hpp>

using namespace Sweep;
using namespace Sweep::Processes;
using namespace std;

// CONSTRUCTORS AND DESTRUCTORS.

// Default constructor (protected).
DimerInception::DimerInception(void)
: Inception(), m_kfm(0.0), m_ksf1(0.0), m_ksf2(0.0), m_efm(2.2)
{
    m_name = "DimerInception";
}

// Initialising constructor.
DimerInception::DimerInception(const Sweep::Mechanism &mech)
: Inception(mech), m_kfm(0.0), m_ksf1(0.0), m_ksf2(0.0),
  m_efm(mech.GetEnhancementFM())
{
    m_name = "DimerInception";
}

// Copy constructor.
DimerInception::DimerInception(const DimerInception &copy)
: m_efm(copy.m_efm)
{
    *this = copy;
}

// Stream-reading constructor.
DimerInception::DimerInception(std::istream &in, const Sweep::Mechanism &mech)
: m_efm(mech.GetEnhancementFM())
{
    Deserialize(in, mech);
}

// Default destructor.
DimerInception::~DimerInception(void)
{
}

// OPERATOR OVERLOADS.

// Assignment operator.
DimerInception &DimerInception::operator =(const DimerInception &rhs)
{
    if (this != &rhs) {
        Inception::operator =(rhs);
        m_kfm  = rhs.m_kfm;
        m_ksf1 = rhs.m_ksf1;
        m_ksf2 = rhs.m_ksf2;
    }
    return *this;
}

/*!
 * Create a new particle and add it to the ensemble with position uniformly
 * distributed over the grid cell.
 *
 * The iterm parameter is included because it will be needed for many process
 * types and this function is meant to have a general signature.
 *
 * \param[in]       t               Time
 * \param[in]       local_geom      Details of geometry around current location
 * \param[in,out]   sys             System to update
 * \param[in]       iterm           Process term responsible for this event
 * \param[in,out]   rng             Random number generator
 *
 * \return      0 on success, otherwise negative.
 */
int DimerInception::Perform(const double t, Cell &sys,
                            const Geometry::LocalGeometry1d &local_geom,
                            const unsigned int iterm,
                            rng_type &rng) const {

    // This routine performs the inception on the given chemical system.

    // Check if hybrid particle-number/particle model
    // If not, add particle to ensemble.
    if (!m_mech->IsHybrid())
    {
    // Create a new particle of the type specified
    // by the system ensemble.
    Particle *sp = m_mech->CreateParticle(t);
	//Incepting first PAH from inception list.

    // Get the cell vertices
    fvector vertices = local_geom.cellVertices();

    // Sample a uniformly distributed position, note that this method
    // works whether the vertices come in increasing or decreasing order,
    // but 1d is assumed for now.
    double posn = vertices.front();

    const double width = vertices.back() - posn;
    boost::uniform_01<rng_type&, double> uniformGenerator(rng);
    posn += width * uniformGenerator();

    sp->setPositionAndTime(posn, t);

    // Initialise the new particle.
    sp->Primary()->SetComposition(ParticleComp());
    sp->Primary()->SetValues(ParticleTrackers());
    sp->UpdateCache();

	double spwt = sp->getStatisticalWeight();

    // Add particle to system's ensemble.
    sys.Particles().Add(*sp, rng);

    // Update gas-phase chemistry of system.
    if (!sys.GetIsAdiabaticFlag())
        adjustGas(sys, spwt);
    // Update gas-phase chemistry and temperature of system.
    else
        adjustParticleTemperature(sys, spwt, 1, ParticleComp()[0], 1);
    }
    else
    {
        // Adjust particle number properties
        sys.Particles().UpdateNumberAtIndex(ParticleComp()[0], 1);
        sys.Particles().UpdateTotalParticleNumber(1);
        sys.Particles().UpdateTotalsWithIndex(ParticleComp()[0], 1.0);

        // Update gas-phase chemistry of system.
        if (!sys.GetIsAdiabaticFlag())
            adjustGas(sys, 1);
        // Update gas-phase chemistry and temperature of system.
        else
            adjustParticleTemperature(sys, 1, 1, ParticleComp()[0], 1);
    }

    return 0;
}

// PERFORMING THE PROCESS.

// INCEPTION KERNEL.

/*!
 * The inception rate is based on a transition coagulation kernel that
 * is calculated using a molecule diameter and mass.  These values do not
 * have to correspond the the physical properties of the molecule, but they
 * are fed into the transition regime kernel.
 *
 * @param[in]    m1    mass of first molecule
 * @param[in]    m2    mass of second molecule
 * @param[in]    d1    diameter of first molecule
 * @param[in]    d2    diameter of second molecule
 */
void DimerInception::SetInceptingSpecies(double m1, double m2, double d1, double d2)
{
    // The free mol part can be handled by the free mol specific method
    SetInceptingSpeciesFreeMol(m1, m2, d1, d2);

    // Now the slip flow part
    double invd1=1.0/d1, invd2=1.0/d2;
    m_ksf1 = CSF * (d1+d2);
    m_ksf2 = 2.0 * 1.257 * m_ksf1 * ((invd1*invd1) + (invd2*invd2));
    m_ksf1 = m_ksf1 * (invd1+invd2);
}

/*!
 * The inception rate is based on a free molecular coagulation kernel only,
 * contrast \ref SetInceptingSpecies that
 * is calculated using a molecule diameter and mass.  These values do not
 * have to correspond the the physical properties of the molecule, but they
 * are fed into the transition regime kernel.
 *
 * @param[in]    m1    mass of first molecule
 * @param[in]    m2    mass of second molecule
 * @param[in]    d1    diameter of first molecule
 * @param[in]    d2    diameter of second molecule
 */
void DimerInception::SetInceptingSpeciesFreeMol(double m1, double m2, double d1, double d2)
{
    // This routine sets the free-mol and slip flow kernel parameters given
    // the mass and diameter of the incepting species.
    m_kfm  = m_efm * CFM * sqrt((1.0/m1) + (1.0/m2)) * (d1+d2) * (d1+d2);
    m_ksf1 = 0.0;
    m_ksf2 = 0.0;
}

// TOTAL RATE CALCULATIONS.

// Returns rate of the process for the given system.
double DimerInception::Rate(double t, const Cell &sys, const Geometry::LocalGeometry1d &local_geom) const
{
    // Get the current chemical conditions.
    double T = sys.GasPhase().Temperature();
    double P = sys.GasPhase().Pressure();

	// Calculate the rate.
	return Rate(sys.GasPhase(), sqrt(T),
                MeanFreePathAir(T,P),
                sys.SampleVolume());
}

/*!
 * Calculate inception rate using a the transition coagulation kernel
 * with the values provided by the user.  The result is this value
 * multiplied by the square of the number concentration of the gas
 * phase species.
 *
 * Requires all the parameters that would otherwise be calculated by
 * the routine to be passed as arguments.
 *
 * @param[in]    gas      Gas phase mixture
 * @param[in]    sqrtT    square root of temperature
 * @param[in]    MFP      mean free path in gas
 * @param[in]    vol      sample volume
 *
 * @return    Inception rate for a cell of size vol. (\f$ \mathrm{s}^{-1}\f$)
 */
double DimerInception::Rate(const EnvironmentInterface &gas, double sqrtT,
                     double MFP, double vol) const
{
    double rate = A() * vol * chemRatePart(gas);

    const double fm   = sqrtT * m_kfm;
    if((m_ksf1 > 0) || (m_ksf2 > 0))  {
        const double Temperature = gas.Temperature();

        // Temperature divided by viscosity
        double T_viscosity = Temperature / gas.Viscosity();

        // Transition regime
        double sf   = T_viscosity  * (m_ksf1 + (MFP*m_ksf2));
        rate *= ((fm*sf) / (fm+sf));
    }
    else {
        // Free mol regime only
        rate *= fm;
    }

    return rate;
}

/*!
 * Calculates the gas-phase chemistry contribution to the rate
 * expression.  This is overloaded as Avogadro's number must be
 * included in the terms for inception processes.
 *
 * @param[in]    gas      Gas phase mixture
 */
double DimerInception::chemRatePart(const EnvironmentInterface &gas) const
{
    // Factor of 0.5 adjusts for doubling counting of pairs of molecules in the number
    // of possible collisions.
    double rate = 0.5;

    Sprog::StoichMap::const_iterator i;
    for (i=m_reac.begin(); i!=m_reac.end(); ++i) {
        //std::cerr << "Mole frac to use " << fracs[i->first] << std::endl;
        double conc = gas.SpeciesConcentration(i->first);
        for (int j=0; j!=i->second; ++j) {
            rate *= (NA * conc);
        }
    }

    return rate;
}

// RATE TERM CALCULATIONS.

// Returns the number of rate terms for this process (one).
unsigned int DimerInception::TermCount(void) const {return 1;}

// Calculates the rate terms given an iterator to a double vector. The
// iterator is advanced to the position after the last term for this
// process.  Returns the sum of all terms.
double DimerInception::RateTerms(const double t, const Cell &sys,
                               const Geometry::LocalGeometry1d &local_geom,
                               fvector::iterator &iterm) const
{
    // Get the current chemical conditions.
    double T = sys.GasPhase().Temperature();
    double P = sys.GasPhase().Pressure();

    // Calculate the single rate term and advance iterator.

    if (sys.ParticleModel()->Postprocessing() == ParticleModel::wdotA4) {
        double Rate = NA * sys.GasPhase().PropertyValue(1007) * sys.SampleVolume();

        if (Rate < 0.0)
            Rate = 0.0;

        *iterm = Rate;
	}else{
        *iterm = Rate(sys.GasPhase(), sqrt(T),
                      MeanFreePathAir(T,P),
                      sys.SampleVolume());
    }

    return *(iterm++);
}

// READ/WRITE/COPY.

// Creates a copy of the inception.
DimerInception *const DimerInception::Clone(void) const {return new DimerInception(*this);}

// Returns the process type.  Used to identify different
// processes and for serialisation.
ProcessType DimerInception::ID(void) const {return Dimer_Inception_ID;}

// Writes the object to a binary stream.
void DimerInception::Serialize(std::ostream &out) const
{
    if (out.good()) {
        // Output the version ID (=0 at the moment).
        const unsigned int version = 0;
        out.write((char*)&version, sizeof(version));

        // Serialize base class.
        Inception::Serialize(out);

        // Write free-mol parameter.
        double v = (double)m_kfm;
        out.write((char*)&v, sizeof(v));

        // Write slip-flow parameters.
        v = (double)m_ksf1;
        out.write((char*)&v, sizeof(v));
        v = (double)m_ksf2;
        out.write((char*)&v, sizeof(v));

    } else {
        throw invalid_argument("Output stream not ready "
                               "(Sweep, DimerInception::Serialize).");
    }
}

// Reads the object from a binary stream.
void DimerInception::Deserialize(std::istream &in, const Sweep::Mechanism &mech)
{
    if (in.good()) {
        // Read the output version.  Currently there is only one
        // output version, so we don't do anything with this variable.
        // Still needs to be read though.
        unsigned int version = 0;
        in.read(reinterpret_cast<char*>(&version), sizeof(version));

        double val = 0.0;

        switch (version) {
            case 0:
                // Deserialize base class.
                Inception::Deserialize(in, mech);

                // Read free-mol parameter.
                in.read(reinterpret_cast<char*>(&val), sizeof(val));
                m_kfm = (double)val;

                // Read slip-flow parameters.
                in.read(reinterpret_cast<char*>(&val), sizeof(val));
                m_ksf1 = (double)val;
                in.read(reinterpret_cast<char*>(&val), sizeof(val));
                m_ksf2 = (double)val;

                break;
            default:
                throw runtime_error("Serialized version number is invalid "
                                    "(Sweep, DimerInception::Deserialize).");
        }
    } else {
        throw invalid_argument("Input stream not ready "
                               "(Sweep, DimerInception::Deserialize).");
    }
}


// This file is part of DM-HEOM (https://github.com/noma/dm-heom)
//
// Copyright (c) 2015-2019 Matthias Noack, Zuse Institute Berlin
//
// Licensed under the 3-clause BSD License, see accompanying LICENSE,
// CONTRIBUTORS.md, and README.md for further information.

#ifndef heom_hierarchy_graph_hpp
#define heom_hierarchy_graph_hpp

#include <unordered_map>

#include <boost/functional/hash.hpp>
#include <boost/math/special_functions/binomial.hpp>

#include "heom/common.hpp"

namespace heom {

using uid_t = std::int32_t; // unique node id

// represents a tuple in the HEOM sense (kind of an excitation vector) which uniquely describes one ado_tuple/matrix/hierarchy node
using tuple = std::vector<int_t>; // excitation vector, value type of each node
// hash function for tuple
//auto tuple_hash = [](const tuple& v) { return boost::hash_range(v.begin(), v.end()); }; // TODO: report compiler bug, causes multiple definition linker error with Intel compiler, tested with 16.0.3 and 17.0.1
std::size_t tuple_hash(const tuple& v);

class hierarchy_graph
{
public:
	//using tuple_to_uid_map = std::unordered_map<tuple, uid_t, decltype(tuple_hash)>;
	using tuple_to_uid_map = std::unordered_map<tuple, uid_t, std::function<decltype(tuple_hash)>>;
	using node_list = std::vector<tuple>;
	using adj_list = std::vector<uid_t>;
	using edge_list = std::vector<adj_list>;

	static const uid_t invalid_uid;

	hierarchy_graph(std::int64_t baths_number, std::int64_t baths_matsubaras, std::int64_t depth);
	hierarchy_graph(std::int64_t width, std::int64_t depth);

	int64_t width() const { return width_; }
	int64_t depth() const { return depth_; }
	int64_t nodes() const { return nodes_; }
	int64_t compute_nodes() const { return compute_nodes_; }
	int64_t halo_nodes() const { return nodes_ - compute_nodes_; }
	int64_t edges() const { return plus_edge_count_ + minus_edge_count_; }
	int64_t plus_edge_count() const { return plus_edge_count_; }
	int64_t minus_edge_count() const { return minus_edge_count_; }

	uid_t first_non_plus_uid() const { return first_non_plus_uid_; }

	const node_list& tuples() const { return tuples_; }
	const edge_list& plus_edges() const { return plus_edges_; }
	const edge_list& minus_edges() const { return minus_edges_; }

	const std::vector<int_t>& nodes_per_depth() const { return nodes_per_depth_; }

	// compute depth of a tuple
	int_t depth(const tuple& t) const;
	// compute depth of an uid
	int_t depth(const uid_t& uid) const;

	/**
	 * Write graph to file, METIS and Boost.Graph compatible adjacency list
	 * format, i.e. first row contains two space separated numbers for the
	 * node and edge count. Every other row i cantains the adjacency list for
	 * node i (count starting at 1).
	 *
	 * @param filename name/path of the output file.
	 */
	void write_graph_file(const std::string& filename);

protected:
	hierarchy_graph() = default;
	hierarchy_graph(int64_t width, int64_t depth, int64_t nodes, int64_t compute_nodes, uid_t first_non_plus_uid);

//private:
	int64_t width_ = 0; // max. number of outgoing plus or minus edges per node
	int64_t depth_ = 0; // number of hierarchy layers within the graph
	int64_t nodes_ = 0; // number of nodes
	int64_t compute_nodes_ = 0; // nodes to be computed
	int64_t plus_edge_count_ = 0; // number of plus edges
	int64_t minus_edge_count_ = 0; // number of minu edges

	// first node that has no outgoing downward (plus) edges
	uid_t first_non_plus_uid_ = invalid_uid;

	node_list tuples_;
	edge_list plus_edges_;
	edge_list minus_edges_;

	std::vector<int_t> nodes_per_depth_;
};

} // namespace heom

#endif // heom_hierarchy_graph_hpp


/*
  Author(s):      Matthew Celnik (msc37)
  Project:        sweepc (population balance solver)
  Sourceforge:    http://sourceforge.net/projects/mopssuite

  Copyright (C) 2008 Matthew S Celnik.

  File purpose:
    Implementation of the ParticleImage class declared in the
    swp_particle_image.h header file.

  Licence:
    This file is part of "sweepc".

    sweepc is free software; you can redistribute it and/or
    modify it under the terms of the GNU Lesser General Public License
    as published by the Free Software Foundation; either version 2
    of the License, or (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

  Contact:
    Dr Markus Kraft
    Dept of Chemical Engineering
    University of Cambridge
    New Museums Site
    Pembroke Street
    Cambridge
    CB2 3RA
    UK

    Email:       mk306@cam.ac.uk
    Website:     http://como.cheng.cam.ac.uk
*/

#include "swp_particle_image.h"
#include "swp_bintree_primary.h"
#include "swp_surfvol_primary.h"
#include "swp_PAH_primary.h"
#include "string_functions.h"
#include <boost/random/uniform_01.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <fstream>
#include <vector>
#include <stdexcept>
using namespace Sweep;
using namespace Sweep::Imaging;
using namespace std;
using namespace Strings;

const double ParticleImage::m_necking = 1.000;

// CONSTRUCTORS AND DESTRUCTORS.

// Default contructor.
ParticleImage::ParticleImage(void)
{
}

// Initialising constructor.
/*ParticleImage::ParticleImage(const Particle &sp, const ParticleModel &model)
{
    Construct(sp, model);
}*/

// Destructor.
ParticleImage::~ParticleImage(void)
{
}

// PARTICLE IMAGE DATA CONSTRUCTION.

/*!
 * @brief           Construct a particle image
 *
 * @param sp        Particle object
 * @param model     Particle model describing object
 */
void ParticleImage::Construct(const Particle &sp, const ParticleModel &model)
{
    // Initialise new RNG here as the calling functions do not have access to
    // the original one
    rng_type rng(size_t(100));

    // Clear the current image data structure.
    m_root.Clear();

    if (model.AggModel() == AggModels::Spherical_ID) {
        // Spherical particle model, just draw
        // a single sphere.

        m_root.Insert(sp.SphDiameter() * 0.5e9); // Convert to nm.

    } else if (model.AggModel() == AggModels::SurfVol_ID) {
        // Surface-volume model, construct a tree
        // of identical primaries, estimating the primary
        // count and diameter.

        const AggModels::SurfVolPrimary *svp = NULL;
        svp = dynamic_cast<const AggModels::SurfVolPrimary*>(sp.Primary());

        m_root.Clear();
        for (unsigned int i=0; i!=svp->PP_Count(); ++i) {
            // Convert to radius in nm
            m_root.Insert(svp->PP_Diameter() * 0.5e9);
        }

        //! At the moment tracking of the distance between the centres of
        //! primary particles does not apply to the surface-volume model.
        calc_FM(m_root, *(&rng), false);

    } else if (model.AggModel() == AggModels::BinTree_ID ||
            model.AggModel() == AggModels::BinTreeSilica_ID) {
        // Binary Tree generic model

        const AggModels::BinTreePrimary *p;
        p = dynamic_cast<const AggModels::BinTreePrimary*>(sp.Primary());
        ConstructTree(p, *(&rng), model.getTrackPrimaryCoordinates());

    } else if (model.AggModel() == AggModels::PAH_KMC_ID) {
        // PAHPP (binary tree like) model

        const AggModels::PAHPrimary *p;
        p = dynamic_cast<const AggModels::PAHPrimary*>(sp.Primary());
        ConstructTree(p, *(&rng), model.getTrackPrimaryCoordinates());
    } else {
        throw std::runtime_error("Unknown particle model. (ParticleImage::Construct)");
    }

    // Set to TEM-style projection
    //m_root.Project();
}

/*!
 *  This function is like a limited assignment operator, except that the
 *  children are not copied and the pointers to the particles may need
 *  adjusting after this method has finished.
 *
 *  @param[in,out] node   Pointer to node of ImgNode tree.
 *  @param[in]     source Pointer to the primary to be copied.
 */
template <class ParticleClass>
void ParticleImage::CopyParts(ImgNode &node, const ParticleClass *source)
{
	//! Since m_cen_bshp and m_cen_mass are vectors do unit conversion in
    //! function.
    node.setBoundSph(source->m_cen_bsph);
    node.setCOM(source->m_cen_mass);

	//! Units of nm.
    node.setRadius(source->m_primarydiam*0.5e9);
    node.setDistance(source->m_distance_centreToCentre*1.0e9);
}

/*!
 *  Recursively copy the tree for non-leaf nodes.
 *
 *  @param[in,out] node   Pointer to node of ImgNode tree.
 *  @param[in]     source Pointer to the primary to be copied.
 */
template <class ParticleClass>
void ParticleImage::CopyTree(ImgNode &node, const ParticleClass *source)
{
    //! Create the new left and right children with nothing in them.
    node.m_leftchild = new ImgNode();
    node.m_rightchild = new ImgNode();

    //! Copy the properties such as the volume, surface area and list of
    //! constituent PAH molecules.
    CopyParts(*node.m_leftchild, source->m_leftchild);
    CopyParts(*node.m_rightchild, source->m_rightchild);

    // Set the pointers to the parent.
    node.m_leftchild->m_parent=&node;
    node.m_rightchild->m_parent=&node;

    //! The left and right particle are set further down in UpdateAllPointers.
    //! These are the pointers that specify which primary particles touch each
    //! other in the aggregate structure.
    node.m_leftparticle=NULL;
    node.m_rightparticle=NULL;

    //! Recurse to copy the subtrees.
    if (source->m_leftchild->m_leftchild!=NULL)
        CopyTree(*node.m_leftchild, source->m_leftchild);

    if (source->m_rightchild->m_leftchild!=NULL)
        CopyTree(*node.m_rightchild, source->m_rightchild);

    //! Set the leftparticle and rightparticle.
    UpdateAllPointers(node, source);
}

//! Generates a projection on the zx plane (set all y to 0)
void ParticleImage::Project()
{
    // Call to ImgNode
    m_root.Project();
}

/*!
 * @brief       Writes a 3dout format file
 *
 * This made use of some custom-made OpenGL renderer by Markus Sander
 * which has since been lost. (wjm34 27/07/2012)
 *
 * @param file  File output stream
 * @param x     ?
 * @param y     ?
 * @param z     ?
 */
void ParticleImage::Write3dout(std::ofstream &file, double x, double y, double z)
{
    if (file.good()) {
        string line;
        double val = 0.0;

        // vector of arrays to store primary coordinates.  First
        // 3 values are the cartesian coordinates, final value
        // is the primary radius.
        vector<fvector> coords;

        // Get the primary coordinates from the aggregate tree.
        m_root.GetPriCoords(coords);
		file.write(line.c_str(), line.length());
        //write the radius of gyration
        double Rg=RadiusofGyration();
        line = cstr(0) + " " + cstr(0) +
               " " + cstr(0)+"\n ";
        file.write(line.c_str(), line.length());
        line = cstr(Rg)+"\n";							//write radius of gyration
        file.write(line.c_str(), line.length());
        // Write the primaries to the 3dout file.
        for (unsigned int i=0; i!=coords.size(); ++i) {
            val  = coords[i][3] * m_necking;
            line = cstr(coords[i][0]+x) + " " + cstr(coords[i][1]+y) +
                   " " + cstr(coords[i][2]+z)+"\n ";
            file.write(line.c_str(), line.length());
            line = cstr(val)+"\n";							//write radius
            file.write(line.c_str(), line.length());
        }

    } else {
        throw invalid_argument("Output stream not ready "
                               "(Sweep, ParticleImage::Write3dout).");
    }
}

/*
void ParticleImage::LengthWidth(double &L, double &W)
{
    //align the particle along the z axis
	double xmax=0;
	double ymax=0;
	double xmin=0;
	double ymin=0;
    double zmax=0;
	double zmin=0;
    double dx=0,dy=0,dz=0;
    //the radius of outher spheres
	double yminrad=0;
	double xminrad=0;
	double ymaxrad=0;
	double xmaxrad=0;
    double zmaxrad=0;
    double zminrad=0;
	vector<fvector> coords;
    m_root.GetPriCoords(coords);
    double dist=0, distmax=0;
	for (unsigned int i=0; i!=coords.size(); ++i) {
        dist=sqrt(coords[i][0]*coords[i][0]+coords[i][1]*coords[i][1]+coords[i][2]*coords[i][2])+coords[i][3];
		if (dist>distmax) {
            xmax=coords[i][0];ymax=coords[i][1];zmax=coords[i][2];
            distmax=dist;
        }
	}
    //added to debug ms785
    if (m_root.m_left!=NULL)
    {
        double phi=atan2(ymax,xmax);
        double theta=(PI/2)-atan(zmax/sqrt((xmax*xmax)+(ymax*ymax)));

      //  m_root.Translate(-dx,-dy,-dz);
        m_root.RotateOrigin(0,-phi-PI/2);
        ofstream out;
       // out.open("particlebeforeshift.3d");
       // this->Write3dout(out,0,0,0);
       // out.close();
        m_root.RotateOrigin(-theta,0);

      //  out.open("particleaftershift.3d");
      //  this->Write3dout(out,0,0,0);
      //  out.close();
    }
    coords.clear();
    m_root.GetPriCoords(coords);
	xmax=coords[0][0];
	xmin=xmax;
	ymax=coords[0][1];
	ymin=ymax;
    zmax=coords[0][2];
	zmin=zmax;
	for (unsigned int i=0; i!=coords.size(); ++i) {
		if (xmax<=coords[i][0]+coords[i][3]) {xmax=coords[i][0]+coords[i][3];xmaxrad=coords[i][3];}
		if (ymax<=coords[i][1]+coords[i][3]) {ymax=coords[i][1]+coords[i][3];ymaxrad=coords[i][3];}
        if (zmax<=coords[i][2]+coords[i][3]) {zmax=coords[i][2]+coords[i][3];zmaxrad=coords[i][3];}
		if (xmin>=coords[i][0]-coords[i][3]) {xmin=coords[i][0]-coords[i][3];xminrad=coords[i][3];}
		if (ymin>=coords[i][1]-coords[i][3]) {ymin=coords[i][1]-coords[i][3];yminrad=coords[i][3];}
        if (zmin>=coords[i][2]-coords[i][3]) {zmin=coords[i][2]-coords[i][3];zminrad=coords[i][3];}
	}
    dx=abs(xmax-xmin);
    dy=abs(ymax-ymin);
    dz=abs(zmax-zmin);
	L=dz;

   // W=max(dy,dx);
    W=dx;

}*/

//! Calculates the radius of gyration of a particle
double ParticleImage::RadiusofGyration()
{
    double sum=0;
    double mass;
    double totalmass=0;
    double r2;
    double Rg;
    vector<fvector> coords;
    m_root.GetPriCoords(coords);
	for (unsigned int i=0; i!=coords.size(); ++i) {
        //mass is proportional to the cube of the radius
        mass=coords[i][3]*coords[i][3]*coords[i][3];
        r2=coords[i][0]*coords[i][0]+coords[i][1]*coords[i][1]+coords[i][2]*coords[i][2];
        sum+=mass*r2;
        totalmass+=mass;
	}
    Rg=sqrt(sum/totalmass);
    return Rg;
}

/*!
 *  Each node contains two pointers (m_leftparticle and m_rightparticle)
 *  to primary particles that are connected by this node.
 *  This function is used when the entire particle tree is duplicated.
 *  It sets the pointers in the copied node, that the connectivity
 *  of the primary particles in this node is the same as in the original node.
 *
 *  @todo give this method a more accurate name.
 *
 *  @param[in,out] node     Pointer to node of ImgNode tree.
 *  @param[in]     original Pointer to the primary to be copied.
 */
template <class ParticleClass>
void ParticleImage::UpdateAllPointers(ImgNode &node, const ParticleClass *original)
{
    //! The primary has no children => there are no left and right particles.
    if (original->m_leftchild == NULL) {
        //! Since this is not a connecting node it does not have left and
        //!right particles.
        node.m_leftparticle=NULL;
        node.m_rightparticle=NULL;
    } else {
        //! Find the route to m_leftparticle in the original tree.
        std::stack<bool> route = recordPath(original->m_leftparticle, original);

        //! Now follow the same route down the new tree to find the new left particle.
        node.m_leftparticle = descendPath(&node, route);

        //! Find the route to m_rightparticle in the original tree.
        route = recordPath(original->m_rightparticle, original);

        //! Now follow the same route down the new tree to find the new right particle.
        node.m_rightparticle = descendPath(&node, route);
    }
}

/*!
 *  This is a helper function for UpdateAllPointers.
 *  It climbs up the tree from bottom to top recording a route
 *  suitable for use in call to @see descendPath.
 *
 *  @param[in] bottom Tree node from which to start climbing.
 *  @param[in] top    Tree node at which to stop climbing.
 *
 *  @pre top must be above bottom in a tree.
 *
 *  @return Stack that can be used to descend the same path by moving to the
 *          left child each time the top of the stack is true.
 */
template <class ParticleClass>
std::stack<bool> ParticleImage::recordPath(const ParticleClass* bottom, const ParticleClass* const top) {
    std::stack<bool> wasLeftChild;

    while(bottom != top) {
        //! check whether bottom was a left child of its parent.
        wasLeftChild.push(bottom == bottom->m_parent->m_leftchild);

        //! Climb one level up the tree.
        bottom = bottom->m_parent;
    }
    return wasLeftChild;
}

/*!
 *  @param[in]     here           Point in tree from which to start descent.
 *  @param[in,out] takeLeftBranch Instructions for which child to move to at each level.
 *
 *  @return The node at the bottom of the path.
 *
 *  @pre  here must be a node of tree in which takeLeftBranch is a valid path.
 *  @post takeLeftBranch.empty() == true.
 */
template <class ParticleClass>
ParticleClass* ParticleImage::descendPath(ParticleClass *here, std::stack<bool> &takeLeftBranch) {
    while(!takeLeftBranch.empty()) {
        //! Move one step down the tree in the instructed direction.
        if(takeLeftBranch.top())
            here = here->m_leftchild;
        else
            here = here->m_rightchild;

        //! This instuction has now been processed.
        takeLeftBranch.pop();
    }
    return here;
}

// RENDERING FUNCTIONS.

// Draws the particle image to a POVRAY file.
void ParticleImage::WritePOVRAY(std::ofstream &file)
{
    if (file.good()) {
        string line;
        double val = 0.0;

        // vector of arrays to store primary coordinates.  First
        // 3 values are the cartesian coordinates, final value
        // is the primary radius.
        vector<fvector> coords;

        // Write ParticleDiameter argument to POV file.
        val  = m_root.Radius() * 2.0;
        line = "#declare ParticleDiameter = " + cstr(val) + ";\n";
        file.write(line.c_str(), line.length());

        // Write aggregate opening declaration.
        line = "#declare MyParticle = blob {\n";
        file.write(line.c_str(), line.length());

        // Write threshold radius based on necking parameter.
        val  = max(pow(1.0 - (1.0/(m_necking*m_necking)), 2.0), 1.0e-4);
        line = "  threshold " + cstr(val) + "\n";
        file.write(line.c_str(), line.length());

        // Get the primary coordinates from the aggregate tree.
        m_root.GetPriCoords(coords);

        // Write the primaries to the POV-RAY file.
        for (unsigned int i=0; i!=coords.size(); ++i) {
            val  = coords[i][3] * m_necking;
            line = "sphere {<" + cstr(coords[i][0]) + ", " + cstr(coords[i][1]) +
                   ", " + cstr(coords[i][2]) + ">, " + cstr(val) + ", 1.0}\n";
            file.write(line.c_str(), line.length());
        }

        // Write closing brace for MyParticle declaration.
        line = "}\n";
        file.write(line.c_str(), line.length());

    } else {
        throw invalid_argument("Output stream not ready "
                               "(Sweep, ParticleImage::WritePOVRAY).");
    }
}

// AGGREGATE SPHERE-TREE CONSTRUCTORS (FREE-MOLECULAR).

/*!
 *  @brief Generate the free-molecular structure of a particle.
 *
 *  Calculates the aggregate structure down from the given
 *  node. Assumes that the tree leaves have been initialised
 *  with the correct radii, and recalculates their positions.
 *
 *  @param[in,out] node                   Pointer to node of ImgNode tree.
 *  @param[in]     rng                    Random number generator.
 *  @param[in]     trackPrimaryCoordinates Flag used to indicate whether to track primary coordinates.
 */
void ParticleImage::calc_FM(ImgNode &node, Sweep::rng_type &rng, const bool trackPrimaryCoordinates)
{
    ImgNode *target = node.m_leftchild;
    ImgNode *bullet = node.m_rightchild;

    if ((target != NULL) && (bullet != NULL)) {
        //! Pass calculation down binary tree left & right branches.
        calc_FM(*target, rng, trackPrimaryCoordinates);
        calc_FM(*bullet, rng, trackPrimaryCoordinates);

        //! The first part of the collision algorithm is to
        //! randomly orientate both left and right aggregates.
        //! They are both then placed so that their bounding
        //! spheres are at the origin.

        //! Rotate left node randomly about CoM.
        //! Generate a random number on [0,1)-double-interval.
        boost::uniform_01<rng_type&, double> uniformGenerator(rng);
        double phi1   = uniformGenerator() * 2.0 * PI;
        double theta1 = ((2.0*uniformGenerator())-1.0) * PI;
        target->RotateCOM(theta1, phi1);

        //! Rotate right node randomly about CoM.
        double phi2   = uniformGenerator() * 2.0 * PI;
        double theta2 = ((2.0*uniformGenerator())-1.0) * PI;
        bullet->RotateCOM(theta2, phi2);

        //! Move both spheres so that the bounding spheres
        //! sit at the origin.
        target->CentreBoundSph();
        bullet->CentreBoundSph();

        //! Perform the collision of the left and right nodes.
        //! This may require several iterations if the chosen
        //! x-y displacement means that the aggregates cannot
        //! collide in the z-direction.
        Coords::Vector D;
        double sumr=0.0;
        bool hit = false;
        while (!hit) {
            //! Need to reset target and bullet here, in case
            //! they have been changed by the tree traversal
            //! code below.
            target = node.m_leftchild;
            bullet = node.m_rightchild;

            if (!trackPrimaryCoordinates) {
                sumr = target->Radius() + bullet->Radius();
            } else {
                sumr = node.m_distance_centreToCentre;
            }

            //! Create a random displacement of the bullet node
            //! in the x-y plane.  The displacement is never
            //! greater than the sum of the radii, therefore they
            //! should always touch.
            D[0] = ((2.0 * uniformGenerator()) - 1.0) * sumr;
            D[1] = ((2.0 * uniformGenerator()) - 1.0) * sumr;

            //// Calculate the z-position for the collision of the
            //// target and bullet.  We do this in case both the
            //// target and bullet are leaf nodes, and hence the
            //// binary tree traversal won't happen.
            // hit = calcCollZ(target->m_cen_mass, target->m_r,
            //                 bullet->m_cen_mass, bullet->m_r,
            //                 D[0], D[1], dz1);
            // if (!hit) continue; // Should never happen.
            // D[2] = target->m_cen_bsph[2] + dz1;

            if (!trackPrimaryCoordinates) {
                //! The next code determines the displacement along the z-axis
                //! required for the target and bullet aggregates to touch.  This
                //! requires falling down the tree progressively recalculating
                //! the nearest nodes at each level, until the leaf nodes
                //! are reached.
                //! This next code calculates the minimum distance between the
                //! target's children and bullet's children, or the target or
                //! bullet if they have no children.  The two children with
                //! the smallest separation are chosen as the next target
                //! and bullet.
                hit = minCollZ(*target, *bullet, D[0], D[1], D[2]);
            } else {
                //! By including pointers to a node's left and right particles,
                //! we can directly calculate the displacement in the z
                //! direction for these two particles which maintains the
                //! particle's connectivity thus negating the need to call
                //! calcCollZ function through minCollZ.
                hit = calcCollZ(node.m_leftparticle->BoundSphCentre(), node.m_leftparticle->Radius(),
                                node.m_rightparticle->BoundSphCentre(), node.m_rightparticle->Radius(),
                                D[0], D[1], D[2], sumr, trackPrimaryCoordinates);
            }
        }

        //! We have a new location for the bullet (right node), so move it.
        node.m_rightchild->Translate(D[0], D[1], D[2]);

        //! Calculate properties of this node.
        node.CalcBoundSph();
        node.CalcCOM();
        node.CentreBoundSph();
    }
}

// Calculates the minimum collision distance between
// a target and a bullet node by moving down the
// binary tree.  If the nodes collide then returns
// true, otherwise returns false.
/*!
 * @brief           Calculates the minimum collision distance
 *
 * Calculates the minimum collision distance between
 * a target and a bullet node by moving down the
 * binary tree.  If the nodes collide then returns
 * true, otherwise returns false.
 *
 * @param target    Target node
 * @param bullet    Bullet node
 * @param dx        ?
 * @param dy        ?
 * @param dz        ?
 * @return          Have the nodes collided?
 */
bool ParticleImage::minCollZ(const ImgNode &target,
                             const ImgNode &bullet,
                             double dx, double dy, double &dz)
{
    bool hit=false, hit1=false;
    double dz2=0.0, dz3=0.0, dz4=0.0;

    if (target.IsLeaf()) {
        // Target is a leaf
        if (bullet.IsLeaf()) {
            // Bullet is a leaf (both leaves).
           return calcCollZ(target.BoundSphCentre(), target.Radius(),
                            bullet.BoundSphCentre(), bullet.Radius(),
                            dx, dy, dz, 0.0, false);
        } else {
            // Bullet is not a leaf, call sub-nodes.
            // Calculate minimum dz for the target and the bullet left subnode.
            hit1 = calcCollZ(target.BoundSphCentre(), target.Radius(),
                             bullet.m_leftchild->BoundSphCentre(), bullet.m_leftchild->Radius(),
                             dx, dy, dz, 0.0, false);
            if (hit1) hit = minCollZ(target, *bullet.m_leftchild, dx, dy, dz);
            // Calculate minimum dz for the target and the bullet right subnode.
            hit1 = calcCollZ(target.BoundSphCentre(), target.Radius(),
                             bullet.m_rightchild->BoundSphCentre(), bullet.m_rightchild->Radius(),
                             dx, dy, dz2, 0.0, false);
            if (hit1) hit = minCollZ(target, *bullet.m_rightchild, dx, dy, dz2) || hit;
            // Return minimum dz.
            dz = min(dz, dz2);
            return hit;
        }
    } else {
        // Target is not a leaf.
        if (bullet.IsLeaf()) {
            // Bullet is a leaf, call target sub-nodes..
            // Calculate minimum dz for the target left subnode and the bullet.
            hit1 = calcCollZ(target.m_leftchild->BoundSphCentre(), target.m_leftchild->Radius(),
                             bullet.BoundSphCentre(), bullet.Radius(),
                             dx, dy, dz, 0.0, false);
            if (hit1) hit = minCollZ(*target.m_leftchild, bullet, dx, dy, dz);
            // Calculate minimum dz for the target right subnode and the bullet.
            hit1 = calcCollZ(target.m_rightchild->BoundSphCentre(), target.m_rightchild->Radius(),
                             bullet.BoundSphCentre(), bullet.Radius(),
                             dx, dy, dz2, 0.0, false);
            if (hit1) hit = minCollZ(*target.m_rightchild, bullet, dx, dy, dz2) || hit;
            // Return minimum dz.
            dz = min(dz, dz2);
            return hit;
        } else {
            // Bullet is not a leaf (neither is a leaf), check all left/right
            // collision combinations.
            // Target left and bullet left.
            hit1 = calcCollZ(target.m_leftchild->BoundSphCentre(), target.m_leftchild->Radius(),
                             bullet.m_leftchild->BoundSphCentre(), bullet.m_leftchild->Radius(),
                             dx, dy, dz, 0.0, false);
            if (hit1) hit = minCollZ(*target.m_leftchild, *bullet.m_leftchild, dx, dy, dz);
            // Target left and bullet right.
            hit1 = calcCollZ(target.m_leftchild->BoundSphCentre(), target.m_leftchild->Radius(),
                             bullet.m_rightchild->BoundSphCentre(), bullet.m_rightchild->Radius(),
                             dx, dy, dz2, 0.0, false);
            if (hit1) hit = minCollZ(*target.m_leftchild, *bullet.m_rightchild, dx, dy, dz2) || hit;
            // Target right and bullet left.
            hit1 = calcCollZ(target.m_rightchild->BoundSphCentre(), target.m_rightchild->Radius(),
                             bullet.m_leftchild->BoundSphCentre(), bullet.m_leftchild->Radius(),
                             dx, dy, dz3, 0.0, false);
            if (hit1) hit = minCollZ(*target.m_rightchild, *bullet.m_leftchild, dx, dy, dz3) || hit;
            // Target right and bullet right.
            hit1 = calcCollZ(target.m_rightchild->BoundSphCentre(), target.m_rightchild->Radius(),
                             bullet.m_rightchild->BoundSphCentre(), bullet.m_rightchild->Radius(),
                             dx, dy, dz4, 0.0, false);
            if (hit1) hit = minCollZ(*target.m_rightchild, *bullet.m_rightchild, dx, dy, dz4) || hit;
            // Returns minimum dz.
            dz = min(min(dz, dz2), min(dz3, dz4));
            return hit;
        }
    }
}

/*!
 *  @brief Calculates the z-displacement of a sphere.
 *
 *  Calculates the z-displacement of a bullet sphere for a +ve
 *  collision with a target sphere. Returns true if the
 *  spheres collide, otherwise false.
 *
 *  @param[in]  p1                     Coordinates of sphere 1.
 *  @param[in]  r1                     Radius of sphere 1.
 *  @param[in]  p2                     Coordinates of sphere 2.
 *  @param[in]  r2                     Radius of sphere 2
 *  @param[in]  dx                     Bullet x displacement.
 *  @param[in]  dy                     Bullet y displacement.
 *  @param[out] dz                     Bullet z displacement.
 *  @param[in]  distanceCentreToCentre Distance between the centres of neighbouring primary particles.
 *  @param[in]  trackPrimaryCoordinates Flag used to indicate whether to track primary coordinates.
 *
 *  @return Have the nodes collided?
 */
bool ParticleImage::calcCollZ(const Coords::Vector &p1, double r1,
                              const Coords::Vector &p2, double r2,
                              double dx, double dy, double &dz,
                              double distanceCentreToCentre, const bool trackPrimaryCoordinates)
{
    double sumrsqr;

    if (!trackPrimaryCoordinates) {
        sumrsqr = r1 + r2;
    } else {
        sumrsqr = distanceCentreToCentre;
    }

    //! Calculate the square of the sum of the radii, or the distance between
    //! the centres of neighbouring primary particles if tracking primary
    //! separation.
    sumrsqr *= sumrsqr;

    //! Calculate dx, dy and dz. Remember to include
    //! argument contributions.
    double xdev = p2[0] - p1[0] + dx;
    double ydev = p2[1] - p1[1] + dy;
    double zdev = p2[2] - p1[2];

    //! Calculate dx, dy and dz squared.
    double dxsqr = xdev * xdev;
    double dysqr = ydev * ydev;
    double dzsqr = zdev * zdev;

    // Calculate quadratic terms.
    double b = 2.0 * zdev;
    double c = dxsqr + dysqr + dzsqr - sumrsqr;

    //! Calculate discriminant.
    double dis = (b*b) - (4.0*c);

    if (dis >= 0.0) {
        //! Spheres intersect.
        dz = - 0.5 * (b + sqrt(dis));
        return true;
    } else {
        //! Spheres do not intersect.
        dz = 1.0e10; //!< A large number.
        return false;
    }
}


// Copyright 2015-2022 The ALMA Project Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.

/// @file
/// Computes conductivities versus temperature.
/// Material information and sweep settings provided via XML input.

#include <iostream>
#include <functional>
#include <boost/filesystem.hpp>
#include <boost/mpi.hpp>
#include <utilities.hpp>
#include <vasp_io.hpp>
#include <qpoint_grid.hpp>
#include <processes.hpp>
#include <isotopic_scattering.hpp>
#include <bulk_hdf5.hpp>
#include <analytic1d.hpp>
#include <io_utils.hpp>
#include <bulk_properties.hpp>
#include <vector>
#include <string>
#include <Eigen/Dense>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/xml_parser.hpp>
#include <beyondRTA.hpp>

int main(int argc, char** argv) {
    // set up MPI environment
    boost::mpi::environment env;
    boost::mpi::communicator world;

    if (world.size() > 1) {
        std::cout
            << "*** ERROR: kappa_Tsweep does not run in parallel mode. ***"
            << std::endl;
        world.abort(1);
    }

    if (argc < 2) {
        std::cout << "USAGE: kappa_Tsweep <inputfile.xml>" << std::endl;
        return 1;
    }

    else {
        // define variables
        std::string target_directory = "AUTO";
        std::string target_filename = "AUTO";
        std::string h5_repository = ".";
        std::string mat_directory;
        std::string mat_base;
        int gridDensityA = -1;
        int gridDensityB = -1;
        int gridDensityC = -1;
        bool logsweep = false;
        double Tmin = -1.0;
        double Tmax = -1.0;
        int NT = -1;
        bool projectConductivity = false;
        Eigen::Vector3d uvector(0.0, 0.0, 0.0);
        Eigen::Vector3d u_norm(0.0, 0.0, 0.0);
        bool fullBTE = false;
        bool fullBTE_iterative = true;
        bool outputCapacity = false;
        bool superlattice = false;
        std::string superlattice_UID = "NULL";

        std::cout << "*************************************" << std::endl;
        std::cout << "This is ALMA/kappa_Tsweep version " << ALMA_VERSION_MAJOR
                  << "." << ALMA_VERSION_MINOR << std::endl;
        std::cout << "*************************************" << std::endl;

        // verify that input file exists.
        if (!boost::filesystem::exists(boost::filesystem::path{argv[1]})) {
            std::cout << "ERROR: input file " << argv[1] << " does not exist."
                      << std::endl;
            exit(1);
        }

        /////////////////////////
        /// PARSE INPUT FILE  ///
        /////////////////////////

        std::string xmlfile(argv[1]);
        std::cout << "PARSING " << xmlfile << " ..." << std::endl;

        // Create empty property tree object
        boost::property_tree::ptree tree;

        // Parse XML input file into the tree
        boost::property_tree::read_xml(xmlfile, tree);

        for (const auto& v : tree.get_child("Tsweep")) {
            if (v.first == "H5repository") {
                h5_repository =
                    alma::parseXMLfield<std::string>(v, "root_directory");
            }

            if (v.first == "target") {
                target_directory =
                    alma::parseXMLfield<std::string>(v, "directory");
                target_filename = alma::parseXMLfield<std::string>(v, "file");
            }

            if (v.first == "compound") {
                mat_directory =
                    alma::parseXMLfield<std::string>(v, "directory");
                mat_base = alma::parseXMLfield<std::string>(v, "base");

                gridDensityA = alma::parseXMLfield<int>(v, "gridA");
                gridDensityB = alma::parseXMLfield<int>(v, "gridB");
                gridDensityC = alma::parseXMLfield<int>(v, "gridC");
            }

            if (v.first == "superlattice") {
                superlattice_UID = alma::parseXMLfield<std::string>(v, "UID");
            }

            if (v.first == "sweep") {
                std::string sweepID =
                    alma::parseXMLfield<std::string>(v, "type");

                if (sweepID.compare("log") == 0) {
                    logsweep = true;
                }

                Tmin = alma::parseXMLfield<double>(v, "start");
                Tmax = alma::parseXMLfield<double>(v, "stop");
                NT = alma::parseXMLfield<int>(v, "points");
            }

            if (v.first == "fullBTE") {
                fullBTE = true;
                fullBTE_iterative = alma::parseXMLfield<bool>(v, "iterative");
            }

            if (v.first == "transportAxis") {
                projectConductivity = true;
                double ux = alma::parseXMLfield<double>(v, "x");
                double uy = alma::parseXMLfield<double>(v, "y");
                double uz = alma::parseXMLfield<double>(v, "z");

                uvector << ux, uy, uz;
            }

            if (v.first == "outputHeatCapacity") {
                outputCapacity = true;
            }
        } // end XML parsing

        // Ensure that provided information is within expected bounds

        bool badinput = false;

        if (gridDensityA < 1) {
            std::cout << "ERROR: provided gridA is " << gridDensityA
                      << std::endl;
            std::cout << "Value must be at least 1." << std::endl;
            badinput = true;
        }

        if (gridDensityB < 1) {
            std::cout << "ERROR: provided gridB is " << gridDensityB
                      << std::endl;
            std::cout << "Value must be at least 1." << std::endl;
            badinput = true;
        }

        if (gridDensityC < 1) {
            std::cout << "ERROR: provided gridC is " << gridDensityC
                      << std::endl;
            std::cout << "Value must be at least 1." << std::endl;
            badinput = true;
        }

        if (Tmin <= 0.0) {
            std::cout << "ERROR: provided start temperature is " << Tmin
                      << std::endl;
            std::cout << "Value must be positive." << std::endl;
            badinput = true;
        }

        if ((Tmax < Tmin) || (Tmax <= 0.0)) {
            std::cout << "ERROR: provided end temperature is " << Tmax
                      << std::endl;
            std::cout << "Value must be positive and >= start temperature."
                      << std::endl;
            badinput = true;
        }

        if (NT <= 0) {
            std::cout << "ERROR: provided number of temperature values is "
                      << NT << std::endl;
            std::cout << "Value must be at least 1." << std::endl;
            badinput = true;
        }

        if (projectConductivity && (uvector.norm() < 1e-12)) {
            std::cout << "ERROR: provided transport axis vector has zero norm."
                      << std::endl;
            badinput = true;
        }

        if (badinput) {
            world.abort(1);
        }

        // Initialise file system and verify that directories actually exist
        auto launch_path = boost::filesystem::current_path();
        auto basedir = boost::filesystem::path(h5_repository);

        if (!(boost::filesystem::exists(boost::filesystem::path(basedir)))) {
            std::cout << "ERROR:" << std::endl;
            std::cout << "Repository directory " << basedir
                      << " does not exist." << std::endl;
            world.abort(1);
        }

        if (!(boost::filesystem::exists(
                boost::filesystem::path(basedir / mat_directory)))) {
            std::cout << "ERROR:" << std::endl;
            std::cout << "Material directory " << mat_directory
                      << " does not exist within the HDF5 repository."
                      << std::endl;
            world.abort(1);
        }

        // Resolve name of HDF5 file
        std::stringstream h5namebuilder;
        h5namebuilder << mat_base << "_" << gridDensityA << "_" << gridDensityB
                      << "_" << gridDensityC << ".h5";
        std::string h5filename = h5namebuilder.str();

        // obtain phonon data from HDF5 file
        auto hdf5_path = basedir / boost::filesystem::path(mat_directory) /
                         boost::filesystem::path(h5filename);

        if (!(boost::filesystem::exists(hdf5_path))) {
            std::cout << "ERROR:" << std::endl;
            std::cout << "H5 file " << h5filename
                      << " does not exist within the material directory."
                      << std::endl;
            world.abort(1);
        }

        std::cout << "Opening HDF5 file " << hdf5_path << std::endl;

        auto hdf5_data =
            alma::load_bulk_hdf5(hdf5_path.string().c_str(), world);
        auto description = std::get<0>(hdf5_data);
        auto poscar = std::move(std::get<1>(hdf5_data));
        auto syms = std::move(std::get<2>(hdf5_data));
        auto grid = std::move(std::get<3>(hdf5_data));
        auto processes = std::move(std::get<4>(hdf5_data));

        if (processes->size() == 0) {
            std::cout << "ERROR:" << std::endl;
            std::cout << "List of 3-phonon processes is missing in H5 file."
                      << std::endl;
            world.abort(1);
        }

        // Check if we are dealing with a superlattice.
        // If so, load the applicable scattering data.

        auto subgroups =
            alma::list_scattering_subgroups(hdf5_path.string().c_str(), world);

        int superlattice_count = 0;

        for (std::size_t ngroup = 0; ngroup < subgroups.size(); ngroup++) {
            if (subgroups.at(ngroup).find("superlattice") !=
                std::string::npos) {
                superlattice_count++;
            }
        }

        superlattice_count /= 2;

        if (superlattice_count > 0) {
            superlattice = true;
        }

        Eigen::ArrayXXd w0_SLdisorder;
        Eigen::ArrayXXd w0_SLbarriers;

        if (superlattice) {
            // complain if there are multiple possibilities

            if ((superlattice_count > 1) && (superlattice_UID == "NULL")) {
                std::cout << "ERROR:" << std::endl;
                std::cout << "H5 file contains scattering information for "
                             "multiple superlattices."
                          << std::endl;
                std::cout << "Must provide the superlattice UID via the "
                             "<superlattice> XML tag."
                          << std::endl;
                world.abort(1);
            }

            // if the user provided a UID, verify that corresponding data exists

            if (superlattice_UID != "NULL") {
                int UIDcount = 0;

                for (std::size_t ngroup = 0; ngroup < subgroups.size();
                     ngroup++) {
                    if (subgroups.at(ngroup).find(superlattice_UID) !=
                        std::string::npos) {
                        UIDcount++;
                    }
                }

                if (UIDcount != 2) {
                    std::cout << "ERROR:" << std::endl;
                    std::cout << "H5 file does not contain any superlattice "
                                 "data with provided UID "
                              << superlattice_UID << "." << std::endl;
                    world.abort(1);
                }
            }

            // load the scattering rates from the H5 file

            bool UIDmatch = true;

            for (std::size_t ngroup = 0; ngroup < subgroups.size(); ngroup++) {
                bool contains_SLdisorder =
                    (subgroups.at(ngroup).find("superlattice") !=
                     std::string::npos) &&
                    (subgroups.at(ngroup).find("disorder") !=
                     std::string::npos);

                bool contains_SLbarriers =
                    (subgroups.at(ngroup).find("superlattice") !=
                     std::string::npos) &&
                    (subgroups.at(ngroup).find("barriers") !=
                     std::string::npos);

                if (superlattice_count > 1) {
                    UIDmatch = (subgroups.at(ngroup).find(superlattice_UID) !=
                                std::string::npos);
                }

                if (contains_SLdisorder && UIDmatch) {
                    auto mysubgroup = alma::load_scattering_subgroup(
                        hdf5_path.string().c_str(),
                        subgroups.at(ngroup),
                        world);
                    w0_SLdisorder = mysubgroup.w0;
                }

                if (contains_SLbarriers && UIDmatch) {
                    auto mysubgroup = alma::load_scattering_subgroup(
                        hdf5_path.string().c_str(),
                        subgroups.at(ngroup),
                        world);
                    w0_SLbarriers = mysubgroup.w0;
                }
            }
        }

        // Build list of temperature values
        Eigen::VectorXd Tlist;

        if (logsweep) {
            Tlist = alma::logSpace(Tmin, Tmax, NT);
        }

        else {
            Tlist.setLinSpaced(NT, Tmin, Tmax);
        }

        // Create output writer
        std::stringstream outputbuffer;

        // Write file header

        if (projectConductivity) {
            outputbuffer << "Temp[K],kappaRTA<" << uvector(0) << ","
                         << uvector(1) << "," << uvector(2) << ">[W/m-K]";

            if (fullBTE) {
                outputbuffer << ",kappaBTE<" << uvector(0) << "," << uvector(1)
                             << "," << uvector(2) << ">[W/m-K]";
            }
        }

        else {
            if (!fullBTE) {
                outputbuffer << "Temp[K],kappaRTA_xx[W/m-K],kappaRTA_xy[W/"
                                "m-K],kappaRTA_xz[W/m-K],kappaRTA_yy[W/"
                                "m-K],kappaRTA_yz[W/m-K],kappaRTA_zz[W/m-K]";
            }
            else {
                outputbuffer << "Temp[K],kappaRTA_xx[W/m-K],kappaBTE_xx[W/"
                                "m-K],kappaRTA_xy[W/m-K],kappaBTE_xy[W/"
                                "m-K],kappaRTA_xz[W/m-K],kappaBTE_xz[W/"
                                "m-K],kappaRTA_yy[W/m-K],kappaBTE_yy[W/"
                                "m-K],kappaRTA_yz[W/m-K],kappaBTE_yz[W/"
                                "m-K],kappaRTA_zz[W/m-K]kappaBTE_zz[W/m-K],";
            }
        }

        if (outputCapacity) {
            outputbuffer << ",Cbulk[J/m^3-K]";
        }

        outputbuffer << std::endl;

        // RUN CALCULATIONS

        std::cout << "Running Tsweep for " << mat_base << std::endl;

        auto twoph_processes = alma::find_allowed_twoph(*grid, world);

        // temperature-independent scattering rates
        Eigen::ArrayXXd w_elastic;

        if (superlattice) {
            w_elastic = w0_SLbarriers.array() + w0_SLdisorder.array();
        }
        else {
            w_elastic =
                alma::calc_w0_twoph(*poscar, *grid, twoph_processes, world);
        }

        if (projectConductivity) {
            u_norm = uvector.array() / uvector.norm();
        }

        for (int nT = 0; nT < NT; nT++) {
            double T = Tlist(nT);

            std::cout << " Processing " << T << " K (temperature " << nT + 1
                      << " of " << NT << ")" << std::endl;

            Eigen::ArrayXXd w3(
                alma::calc_w0_threeph(*grid, *processes, T, world));
            Eigen::ArrayXXd w(w3 + w_elastic);

            Eigen::Matrix3d kappa_RTA = alma::calc_kappa(*poscar, *grid, w, T);
            Eigen::Matrix3d kappa_BTE;

            if (fullBTE) {
                kappa_BTE = alma::beyondRTA::calc_kappa(*poscar,
                                                        *grid,
                                                        *syms,
                                                        *processes,
                                                        twoph_processes,
                                                        w,
                                                        T,
                                                        fullBTE_iterative,
                                                        world);
            }

            if (projectConductivity) { // write conductivity along chosen
                                       // direction
                Eigen::MatrixXd kappa_u =
                    u_norm.transpose() * kappa_RTA.matrix() * u_norm;
                outputbuffer << T << "," << kappa_u(0, 0);

                if (fullBTE) {
                    Eigen::MatrixXd kappa_u =
                        u_norm.transpose() * kappa_BTE.matrix() * u_norm;
                    outputbuffer << "," << kappa_u(0, 0);
                }
            }

            else { // write tensor components
                if (!fullBTE) {
                    outputbuffer << T << "," << kappa_RTA(0, 0) << ","
                                 << kappa_RTA(0, 1) << ",";
                    outputbuffer << kappa_RTA(0, 2) << "," << kappa_RTA(1, 1)
                                 << ",";
                    outputbuffer << kappa_RTA(1, 2) << "," << kappa_RTA(2, 2);
                }

                else {
                    outputbuffer << T << "," << kappa_RTA(0, 0) << ","
                                 << kappa_BTE(0, 0) << ",";
                    outputbuffer << kappa_RTA(0, 1) << "," << kappa_BTE(0, 1)
                                 << ",";
                    outputbuffer << kappa_RTA(0, 2) << "," << kappa_BTE(0, 2)
                                 << ",";
                    outputbuffer << kappa_RTA(1, 1) << "," << kappa_BTE(1, 1)
                                 << ",";
                    outputbuffer << kappa_RTA(1, 2) << "," << kappa_BTE(1, 2)
                                 << ",";
                    outputbuffer << kappa_RTA(2, 2) << "," << kappa_BTE(2, 2);
                }
            }

            if (outputCapacity) {
                outputbuffer << "," << 1e27 * alma::calc_cv(*poscar, *grid, T);
            }

            outputbuffer << std::endl;
        } // END CALCULATIONS

        // WRITE FILE

        // go to the launch directory
        boost::filesystem::current_path(launch_path);

        // resolve output directory if AUTO is selected
        if (target_directory.compare("AUTO") == 0) {
            target_directory = "output/kappa_Tsweep";
        }

        // create output directory if it doesn't exist yet
        auto outputfolder = boost::filesystem::path(target_directory);

        if (!(boost::filesystem::exists(outputfolder))) {
            boost::filesystem::create_directories(outputfolder);
        }

        boost::filesystem::current_path(launch_path);

        // resolve file name if AUTO selected

        if (target_filename.compare("AUTO") == 0) {
            std::stringstream filenamebuilder;
            filenamebuilder << mat_base << "_" << gridDensityA << "_"
                            << gridDensityB << "_" << gridDensityC;

            if (projectConductivity) {
                filenamebuilder << "_" << uvector(0) << "," << uvector(1) << ","
                                << uvector(2);
            }

            filenamebuilder << "_" << Tmin << "_" << Tmax;

            filenamebuilder << ".Tsweep";
            target_filename = filenamebuilder.str();
        }

        // save file

        std::cout << std::endl;
        std::cout << "Writing to file " << target_filename << std::endl;
        std::cout << "in subdirectory " << target_directory << std::endl;
        std::cout << "under root directory " << launch_path << std::endl;

        std::ofstream outputwriter;
        outputwriter.open("./" + target_directory + "/" + target_filename);
        outputwriter << outputbuffer.str();
        outputwriter.close();

        std::cout << std::endl << "[DONE.]" << std::endl;

        return 0;
    }
}