/* Copyright (c) 2016 - 2021, the adamantine authors.
 *
 * This file is subject to the Modified BSD License and may not be distributed
 * without copyright and license information. Please refer to the file LICENSE
 * for the text and further information on this license.
 */

#include <MaterialProperty.hh>

#include <deal.II/base/point.h>
#include <deal.II/base/quadrature_lib.h>
#include <deal.II/dofs/dof_handler.h>
#include <deal.II/fe/fe_values.h>
#include <deal.II/hp/fe_values.h>
#include <deal.II/hp/q_collection.h>

#include <boost/algorithm/string/split.hpp>
#include <boost/optional.hpp>

#include <algorithm>

namespace adamantine
{

template <int dim>
MaterialProperty<dim>::MaterialProperty(
    MPI_Comm const &communicator,
    dealii::parallel::distributed::Triangulation<dim> const &tria,
    boost::property_tree::ptree const &database)
    : _communicator(communicator), _fe(0), _mp_dof_handler(tria)
{
  // Because deal.II cannot easily attach data to a cell. We store the state
  // of the material in distributed::Vector. This allows to use deal.II to
  // compute the new state after refinement of the mesh. However, this
  // requires to use another DoFHandler.
  reinit_dofs();

  // Set the material state to the state defined in the geometry.
  set_initial_state();

  // Fill the _properties map
  fill_properties(database);
}

template <int dim>
double MaterialProperty<dim>::get_cell_value(
    typename dealii::Triangulation<dim>::active_cell_iterator const &cell,
    StateProperty prop) const
{
  unsigned int property = static_cast<unsigned int>(prop);
  auto const mp_dof_index = get_dof_index(cell);

  return _property_values[property][mp_dof_index];
}

template <int dim>
double MaterialProperty<dim>::get_cell_value(
    typename dealii::Triangulation<dim>::active_cell_iterator const &cell,
    Property prop) const
{
  dealii::types::material_id material_id = cell->material_id();
  unsigned int property = static_cast<unsigned int>(prop);

  return _properties(material_id, property);
}

template <int dim>
void MaterialProperty<dim>::reinit_dofs()
{
  _mp_dof_handler.distribute_dofs(_fe);
  // Initialize the state vectors
  for (auto &vec : _state)
    vec.reinit(_mp_dof_handler.locally_owned_dofs(), _communicator);
}

template <int dim>
void MaterialProperty<dim>::update(
    dealii::DoFHandler<dim> const &temperature_dof_handler,
    dealii::LA::distributed::Vector<double> const &temperature)
{
  auto temperature_average =
      compute_average_temperature(temperature_dof_handler, temperature);
  for (auto &val : _property_values)
    val.reinit(temperature_average.get_partitioner());

  std::vector<dealii::types::global_dof_index> mp_dof(1);
  for (auto cell :
       dealii::filter_iterators(_mp_dof_handler.active_cell_iterators(),
                                dealii::IteratorFilters::LocallyOwnedCell()))
  {
    dealii::types::material_id material_id = cell->material_id();

    unsigned int constexpr liquid =
        static_cast<unsigned int>(MaterialState::liquid);
    unsigned int constexpr powder =
        static_cast<unsigned int>(MaterialState::powder);
    unsigned int constexpr solid =
        static_cast<unsigned int>(MaterialState::solid);
    unsigned int constexpr prop_solidus =
        static_cast<unsigned int>(Property::solidus);
    unsigned int constexpr prop_liquidus =
        static_cast<unsigned int>(Property::liquidus);

    double const solidus = _properties(material_id, prop_solidus);
    double const liquidus = _properties(material_id, prop_liquidus);
    cell->get_dof_indices(mp_dof);
    unsigned int const dof = mp_dof[0];

    // First determine the ratio of liquid.
    double liquid_ratio = -1.;
    double powder_ratio = -1.;
    double solid_ratio = -1.;
    if (temperature_average[dof] < solidus)
      liquid_ratio = 0.;
    else if (temperature_average[dof] > liquidus)
      liquid_ratio = 1.;
    else
      liquid_ratio =
          (temperature_average[dof] - solidus) / (liquidus - solidus);
    // Because the powder can only become liquid, the solid can only become
    // liquid, and the liquid can only become solid, the ratio of powder can
    // only decrease.
    powder_ratio = std::min(1. - liquid_ratio, _state[powder][dof]);
    // Use max to make sure that we don't create matter because of round-off.
    solid_ratio = std::max(1 - liquid_ratio - powder_ratio, 0.);

    // Update the value
    _state[liquid][dof] = liquid_ratio;
    _state[powder][dof] = powder_ratio;
    _state[solid][dof] = solid_ratio;

    if (_use_table)
    {
      for (unsigned int property = 0; property < _n_state_properties;
           ++property)
      {
        for (unsigned int material_state = 0;
             material_state < _n_material_states; ++material_state)
        {
          _property_values[property][dof] +=
              _state[material_state][dof] *
              compute_property_from_table(material_id, material_state, property,
                                          temperature_average[dof]);
        }
      }
    }
    else
    {
      for (unsigned int property = 0; property < _n_state_properties;
           ++property)
      {
        for (unsigned int material_state = 0;
             material_state < _n_material_states; ++material_state)
        {
          for (unsigned int i = 0; i <= _polynomial_order; ++i)
          {
            _property_values[property][dof] +=
                _state[material_state][dof] *
                _state_property_polynomials(material_id, material_state,
                                            property, i) *
                std::pow(temperature_average[dof], i);
          }
        }
      }
    }

    // If we are in the mushy state, i.e., part liquid part solid, we need to
    // modify the rho C_p to take into account the latent heat.
    if ((liquid_ratio > 0.) && (liquid_ratio < 1.))
    {
      unsigned int const specific_heat_prop =
          static_cast<unsigned int>(StateProperty::specific_heat);
      unsigned int const latent_heat_prop =
          static_cast<unsigned int>(Property::latent_heat);
      for (unsigned int material_state = 0; material_state < _n_material_states;
           ++material_state)
      {
        _property_values[specific_heat_prop][dof] +=
            _state[material_state][dof] *
            _properties(material_id, latent_heat_prop) / (liquidus - solidus);
      }
    }

    // The radiation heat transfer coefficient is not a real material property
    // but it is derived from other material properties:
    // h_rad = emissitivity * stefan-boltzmann constant * (T + T_infty) (T^2 +
    // T^2_infty).
    unsigned int const emissivity_prop =
        static_cast<unsigned int>(StateProperty::emissivity);
    unsigned int const radiation_heat_transfer_coef_prop =
        static_cast<unsigned int>(StateProperty::radiation_heat_transfer_coef);
    unsigned int const radiation_temperature_infty_prop =
        static_cast<unsigned int>(Property::radiation_temperature_infty);
    double const T = temperature_average[dof];
    double const T_infty =
        _properties(material_id, radiation_temperature_infty_prop);
    double const emissivity = _property_values[emissivity_prop][dof];
    _property_values[radiation_heat_transfer_coef_prop][dof] =
        emissivity * Constant::stefan_boltzmann * (T + T_infty) *
        (T * T + T_infty * T_infty);
  }
}

template <int dim>
void MaterialProperty<dim>::update_boundary_material_properties(
    dealii::DoFHandler<dim> const &temperature_dof_handler,
    dealii::LA::distributed::Vector<double> const &temperature)
{
  auto temperature_average =
      compute_average_temperature(temperature_dof_handler, temperature);
  for (auto &val : _property_values)
    val.reinit(temperature_average.get_partitioner());

  std::vector<dealii::types::global_dof_index> mp_dof(1);
  // We don't need to loop over all the active cells. We only need to loop over
  // the cells at the boundary and at the interface with FE_Nothing. However, to
  // do this we need to use the temperature_dof_handler instead of the
  // _mp_dof_handler.
  for (auto cell :
       dealii::filter_iterators(_mp_dof_handler.active_cell_iterators(),
                                dealii::IteratorFilters::LocallyOwnedCell()))
  {
    dealii::types::material_id material_id = cell->material_id();

    cell->get_dof_indices(mp_dof);
    unsigned int const dof = mp_dof[0];
    if (_use_table)
    {
      // We only care about properties that are used to compute the boundary
      // condition. So we start at 3.
      for (unsigned int property = 3; property < _n_state_properties;
           ++property)
      {
        for (unsigned int material_state = 0;
             material_state < _n_material_states; ++material_state)
        {
          _property_values[property][dof] +=
              _state[material_state][dof] *
              compute_property_from_table(material_id, material_state, property,
                                          temperature_average[dof]);
        }
      }
    }
    else
    {
      // We only care about properties that are used to compute the boundary
      // condition. So we start at  3.
      for (unsigned int property = 3; property < _n_state_properties;
           ++property)
      {
        for (unsigned int material_state = 0;
             material_state < _n_material_states; ++material_state)
        {
          for (unsigned int i = 0; i <= _polynomial_order; ++i)
          {
            _property_values[property][dof] +=
                _state[material_state][dof] *
                _state_property_polynomials(material_id, material_state,
                                            property, i) *
                std::pow(temperature_average[dof], i);
          }
        }
      }
    }

    // The radiation heat transfer coefficient is not a real material property
    // but it is derived from other material properties:
    // h_rad = emissitivity * stefan-boltzmann constant * (T + T_infty) (T^2 +
    // T^2_infty).
    unsigned int const emissivity_prop =
        static_cast<unsigned int>(StateProperty::emissivity);
    unsigned int const radiation_heat_transfer_coef_prop =
        static_cast<unsigned int>(StateProperty::radiation_heat_transfer_coef);
    unsigned int const radiation_temperature_infty_prop =
        static_cast<unsigned int>(Property::radiation_temperature_infty);
    double const T = temperature_average[dof];
    double const T_infty =
        _properties(material_id, radiation_temperature_infty_prop);
    double const emissivity = _property_values[emissivity_prop][dof];
    _property_values[radiation_heat_transfer_coef_prop][dof] =
        emissivity * Constant::stefan_boltzmann * (T + T_infty) *
        (T * T + T_infty * T_infty);
  }
}

template <int dim>
ADAMANTINE_HOST_DEV dealii::VectorizedArray<double>
MaterialProperty<dim>::compute_material_property(
    StateProperty state_property, dealii::types::material_id const *material_id,
    dealii::VectorizedArray<double> const *state_ratios,
    dealii::VectorizedArray<double> temperature) const
{
  dealii::VectorizedArray<double> value = 0.0;
  unsigned int const property_index = static_cast<unsigned int>(state_property);

  if (_use_table)
  {
    for (unsigned int material_state = 0; material_state < _n_material_states;
         ++material_state)
    {
      for (unsigned int n = 0; n < dealii::VectorizedArray<double>::size(); ++n)
      {

        const dealii::types::material_id m_id = material_id[n];

        value[n] += state_ratios[material_state][n] *
                    compute_property_from_table(m_id, material_state,
                                                property_index, temperature[n]);
      }
    }
  }
  else
  {
    for (unsigned int material_state = 0; material_state < _n_material_states;
         ++material_state)
    {
      for (unsigned int n = 0; n < dealii::VectorizedArray<double>::size(); ++n)
      {

        dealii::types::material_id m_id = material_id[n];

        for (unsigned int i = 0; i <= _polynomial_order; ++i)
        {
          value[n] += state_ratios[material_state][n] *
                      _state_property_polynomials(m_id, material_state,
                                                  property_index, i) *
                      std::pow(temperature[n], i);
        }
      }
    }
  }

  return value;
}

template <int dim>
void MaterialProperty<dim>::set_state(
    dealii::Table<2, dealii::VectorizedArray<double>> const &liquid_ratio,
    dealii::Table<2, dealii::VectorizedArray<double>> const &powder_ratio,
    std::map<typename dealii::DoFHandler<dim>::cell_iterator,
             std::pair<unsigned int, unsigned int>> &cell_it_to_mf_cell_map,
    dealii::DoFHandler<dim> const &dof_handler)
{
  auto const powder_state = static_cast<unsigned int>(MaterialState::powder);
  auto const liquid_state = static_cast<unsigned int>(MaterialState::liquid);
  auto const solid_state = static_cast<unsigned int>(MaterialState::solid);
  std::vector<dealii::types::global_dof_index> mp_dof(1.);

  for (auto cell :
       dealii::filter_iterators(dof_handler.active_cell_iterators(),
                                dealii::IteratorFilters::LocallyOwnedCell()))
  {
    typename dealii::Triangulation<dim>::active_cell_iterator cell_tria(cell);
    auto mp_dof_index = get_dof_index(cell_tria);
    auto const &mf_cell_vector = cell_it_to_mf_cell_map[cell];
    unsigned int const n_q_points = dof_handler.get_fe().tensor_degree() + 1;
    double liquid_ratio_sum = 0.;
    double powder_ratio_sum = 0.;
    for (unsigned int q = 0; q < n_q_points; ++q)
    {
      liquid_ratio_sum +=
          liquid_ratio(mf_cell_vector.first, q)[mf_cell_vector.second];
      powder_ratio_sum +=
          powder_ratio(mf_cell_vector.first, q)[mf_cell_vector.second];
    }
    _state[liquid_state][mp_dof_index] = liquid_ratio_sum / n_q_points;
    _state[powder_state][mp_dof_index] = powder_ratio_sum / n_q_points;
    _state[solid_state][mp_dof_index] =
        std::max(1. - _state[liquid_state][mp_dof_index] -
                     _state[powder_state][mp_dof_index],
                 0.);
  }
}

template <int dim>
void MaterialProperty<dim>::set_initial_state()
{
  // Set the material state to the one defined by the user_index
  std::vector<dealii::types::global_dof_index> mp_dof(1);
  for (auto cell :
       dealii::filter_iterators(_mp_dof_handler.active_cell_iterators(),
                                dealii::IteratorFilters::LocallyOwnedCell()))
  {
    cell->get_dof_indices(mp_dof);
    _state[cell->user_index()][mp_dof[0]] = 1.;
  }
}

template <int dim>
void MaterialProperty<dim>::fill_properties(
    boost::property_tree::ptree const &database)
{
  std::array<std::string, _n_material_states> material_state = {
      {"powder", "solid", "liquid"}};
  std::array<std::string, _n_properties> properties = {
      {"liquidus", "solidus", "latent_heat", "radiation_temperature_infty",
       "convection_temperature_infty"}};
  std::array<std::string, _n_state_properties> state_properties = {
      {"density", "specific_heat", "thermal_conductivity", "emissivity",
       "radiation_heat_transfer_coef", "convection_heat_transfer_coef"}};

  // PropertyTreeInput materials.property_format
  std::string property_format = database.get<std::string>("property_format");
  ASSERT_THROW((property_format == "table") ||
                   (property_format == "polynomial"),
               "property_format should be table or polynomial.");
  _use_table = (property_format == "table");
  // PropertyTreeInput materials.n_materials
  unsigned int const n_materials = database.get<unsigned int>("n_materials");
  // Find all the material_ids being used.
  std::vector<dealii::types::material_id> material_ids;
  for (dealii::types::material_id id = 0;
       id < dealii::numbers::invalid_material_id; ++id)
  {
    if (database.count("material_" + std::to_string(id)) != 0)
      material_ids.push_back(id);
    if (material_ids.size() == n_materials)
      break;
  }
  ASSERT_THROW(material_ids.size() == n_materials,
               "Could not find all the material_ids.");

  // When using the polynomial format we allocate one contiguous block of
  // memory. Thus, the largest material_id should be as small as possible
  unsigned int const n_material_ids =
      *std::max_element(material_ids.begin(), material_ids.end()) + 1;
  _properties.reinit(n_material_ids);
  if (_use_table)
  {
    // TODO read table size from input file
    _table_size = 4;
    _state_property_tables.reinit(n_material_ids, _table_size);
    _state_property_tables.set_zero();
  }
  else
  {
    // TODO read polynomial_order from input file
    _polynomial_order = 4;
    _state_property_polynomials.reinit(n_material_ids + 1,
                                       _polynomial_order + 1);
    _state_property_polynomials.set_zero();
  }

  for (auto const material_id : material_ids)
  {
    // Get the material property tree.
    boost::property_tree::ptree const &material_database =
        database.get_child("material_" + std::to_string(material_id));
    // For each material, loop over the possible states.
    bool valid_state = false;
    for (unsigned int state = 0; state < _n_material_states; ++state)
    {
      // The state may or may not exist for the material.
      boost::optional<boost::property_tree::ptree const &> state_database =
          material_database.get_child_optional(material_state[state]);
      if (state_database)
      {
        valid_state = true;
        // For each state, loop over the possible properties.
        for (unsigned int p = 0; p < _n_state_properties; ++p)
        {
          // The property may or may not exist for that state
          boost::optional<std::string> const property =
              state_database.get().get_optional<std::string>(
                  state_properties[p]);
          // If the property exists, put it in the map. If the property does not
          // exist, we have a nullptr.
          if (property)
          {
            // Remove blank spaces
            std::string property_string = property.get();
            property_string.erase(
                std::remove_if(property_string.begin(), property_string.end(),
                               [](unsigned char x) { return std::isspace(x); }),
                property_string.end());
            if (_use_table)
            {
              std::vector<std::string> parsed_property;
              boost::split(parsed_property, property_string,
                           [](char c) { return c == ';'; });
              unsigned int const parsed_property_size = parsed_property.size();
              ASSERT_THROW(parsed_property_size <= _table_size,
                           "Too many coefficients, increase the table size");
              for (unsigned int i = 0; i < parsed_property_size; ++i)
              {
                std::vector<std::string> t_v;
                boost::split(t_v, parsed_property[i],
                             [](char c) { return c == ','; });
                ASSERT(t_v.size() == 2, "Error reading material property.");
                _state_property_tables(material_id, state, p, i, 0) =
                    std::stod(t_v[0]);
                _state_property_tables(material_id, state, p, i, 1) =
                    std::stod(t_v[1]);
              }
              // fill the rest  with the last value
              for (unsigned int i = parsed_property_size; i < _table_size; ++i)
              {
                _state_property_tables(material_id, state, p, i, 0) =
                    _state_property_tables(material_id, state, p, i - 1, 0);
                _state_property_tables(material_id, state, p, i, 1) =
                    _state_property_tables(material_id, state, p, i - 1, 1);
              }
            }
            else
            {
              std::vector<std::string> parsed_property;
              boost::split(parsed_property, property_string,
                           [](char c) { return c == ','; });
              unsigned int const parsed_property_size = parsed_property.size();
              ASSERT_THROW(
                  parsed_property_size <= _polynomial_order,
                  "Too many coefficients, increase the polynomial order");
              for (unsigned int i = 0; i < parsed_property_size; ++i)
              {
                _state_property_polynomials(material_id, state, p, i) =
                    std::stod(parsed_property[i]);
              }
            }
          }
        }
      }
    }
    // Check that there is at least one valid MaterialState
    ASSERT_THROW(
        valid_state == true,
        "Material without any valid state (solid, powder, or liquid).");

    // Check for the properties that are associated to a material but that
    // are independent of an individual state. These properties are duplicated
    // for every state.
    for (unsigned int p = 0; p < _n_properties; ++p)
    {
      // The property may or may not exist for that state
      boost::optional<double> const property =
          material_database.get_optional<double>(properties[p]);
      // If the property exists, put it in the map. If the property does not
      // exist, we use the largest possible value. This is useful if the
      // liquidus and the solidus are not set.
      _properties(material_id, p) =
          property ? property.get() : std::numeric_limits<double>::max();
    }
  }
}

// We need to compute the average temperature on the cell because we need the
// material properties to be uniform over the cell. If there aren't then we have
// problems with the weak form discretization.
template <int dim>
dealii::LA::distributed::Vector<double>
MaterialProperty<dim>::compute_average_temperature(
    dealii::DoFHandler<dim> const &temperature_dof_handler,
    dealii::LA::distributed::Vector<double> const &temperature) const
{
  // TODO: this should probably done in a matrix-free fashion.
  // The triangulation is the same for both DoFHandler
  dealii::LA::distributed::Vector<double> temperature_average(
      _mp_dof_handler.locally_owned_dofs(), temperature.get_mpi_communicator());
  temperature.update_ghost_values();
  temperature_average = 0.;
  auto mp_cell = _mp_dof_handler.begin_active();
  auto mp_end_cell = _mp_dof_handler.end();
  auto enth_cell = temperature_dof_handler.begin_active();
  dealii::hp::FECollection<dim> const &fe_collection =
      temperature_dof_handler.get_fe_collection();
  dealii::hp::QCollection<dim> q_collection;
  q_collection.push_back(dealii::QGauss<dim>(fe_collection.max_degree() + 1));
  q_collection.push_back(dealii::QGauss<dim>(1));
  dealii::hp::FEValues<dim> hp_fe_values(
      fe_collection, q_collection,
      dealii::UpdateFlags::update_values |
          dealii::UpdateFlags::update_quadrature_points |
          dealii::UpdateFlags::update_JxW_values);
  std::vector<dealii::types::global_dof_index> mp_dof_indices(1);
  unsigned int const n_q_points = q_collection.max_n_quadrature_points();
  unsigned int const dofs_per_cell = fe_collection.max_dofs_per_cell();
  std::vector<dealii::types::global_dof_index> enth_dof_indices(dofs_per_cell);
  for (; mp_cell != mp_end_cell; ++enth_cell, ++mp_cell)
  {
    ASSERT(mp_cell->is_locally_owned() == enth_cell->is_locally_owned(),
           "Internal Error");
    if ((mp_cell->is_locally_owned()) && (enth_cell->active_fe_index() == 0))
    {
      hp_fe_values.reinit(enth_cell);
      dealii::FEValues<dim> const &fe_values =
          hp_fe_values.get_present_fe_values();
      mp_cell->get_dof_indices(mp_dof_indices);
      dealii::types::global_dof_index const mp_dof_index = mp_dof_indices[0];
      enth_cell->get_dof_indices(enth_dof_indices);
      double volume = 0.;
      for (unsigned int q = 0; q < n_q_points; ++q)
        for (unsigned int i = 0; i < dofs_per_cell; ++i)
        {
          volume += fe_values.shape_value(i, q) * fe_values.JxW(q);
          temperature_average[mp_dof_index] +=
              fe_values.shape_value(i, q) * temperature[enth_dof_indices[i]] *
              fe_values.JxW(q);
        }
      temperature_average[mp_dof_index] /= volume;
    }
  }

  return temperature_average;
}

template <int dim>
double MaterialProperty<dim>::compute_property_from_table(
    unsigned int const material_id, unsigned int const material_state,
    unsigned int const property, double const temperature) const
{
  if (temperature <=
      _state_property_tables(material_id, material_state, property, 0, 0))
  {
    return _state_property_tables(material_id, material_state, property, 0, 1);
  }
  else
  {
    unsigned int i = 0;
    unsigned int const size = _state_property_tables.extent(3);
    for (; i < size; ++i)
    {
      if (temperature <
          _state_property_tables(material_id, material_state, property, i, 0))
      {
        break;
      }
    }

    if (i >= size - 1)
    {
      return _state_property_tables(material_id, material_state, property,
                                    size - 1, 1);
    }
    else
    {
      auto tempertature_i =
          _state_property_tables(material_id, material_state, property, i, 0);
      auto tempertature_im1 = _state_property_tables(
          material_id, material_state, property, i - 1, 0);
      auto property_i =
          _state_property_tables(material_id, material_state, property, i, 1);
      auto property_im1 = _state_property_tables(material_id, material_state,
                                                 property, i - 1, 1);
      return property_im1 + (temperature - tempertature_im1) *
                                (property_i - property_im1) /
                                (tempertature_i - tempertature_im1);
    }
  }
}

} // namespace adamantine


#include "UGCPopularity.hpp"
#include "ContentElement.hpp"
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_01.hpp>
#include <boost/random/gamma_distribution.hpp>
#include "boost/random/uniform_real_distribution.hpp"
#include <boost/math/distributions/lognormal.hpp>
#include "boost/random/uniform_int.hpp"
#include <algorithm>

extern boost::mt19937 gen;

UGCPopularity::UGCPopularity(unsigned int totalRounds,
        bool perturbations) {
  this->totalRounds = totalRounds;
  this->ttpDist = new boost::math::ttpDistribution(totalRounds);
  this->perturbations = perturbations;
}

UGCPopularity::~UGCPopularity() {
  delete ttpDist;
}

// generate views for all the content elements at a given phase for this round only
unsigned int UGCPopularity::generateViews(std::list<ContentElement*> contentList,
      peakingPhase phase) {
  unsigned int bodyCount, tailCount;
  tailCount = std::floor(contentList.size() / 10);
  bodyCount = contentList.size() - tailCount;
  std::list<unsigned int> views;
  if (phase == BEFORE_PEAK) {
    for (unsigned int i = 0; i < tailCount; i++) {
      views.push_back(this->getBeforePeakTailViews());
    }
    for (unsigned int i = 0; i < bodyCount; i++) {
      views.push_back(this->getBeforePeakViews());
    }
  } else if (phase == AT_PEAK) {
    for (unsigned int i = 0; i < tailCount; i++) {
      views.push_back(this->getAtPeakTailViews());
    }
    for (unsigned int i = 0; i < bodyCount; i++) {
      views.push_back(this->getAtPeakViews());
    }
  } else { // phase == AFTER_PEAK
    for (unsigned int i = 0; i < tailCount; i++) {
      views.push_back(this->getAfterPeakTailViews());
    }
    for (unsigned int i = 0; i < bodyCount; i++) {
      views.push_back(this->getAfterPeakViews());
    }
  }
  // sort the views in descending order and assign them to the contentList
  // (which is ordered by ranking) so that relative ranking remains unaltered
  views.sort(std::greater<unsigned int>());
  std::list<unsigned int>::iterator viewsIt, mIt;

  // introduce perturbations in the number of views, by switching views between
  // videos in the same phase and at a maximum distance g
  if (this->perturbations && phase != AT_PEAK) {
    uint g = 12; // suggested value from Borghol et al.
    std::vector<std::pair<uint, uint> > winLimits;
    uint minWin, maxWin;
    uint maxViews = views.front();
    // compute valid switching windows for each content
    BOOST_FOREACH(uint v, views) {
      minWin = std::floor(v / g);
      maxWin = std::min(maxViews, v * g);
      winLimits.push_back(std::make_pair(minWin,maxWin));
    }
    // switch views a "sufficiently large number of times" (!) between contents
    // with compatible switching windows

    // NOTE: the following code is so ugly it makes me want to kill kittens.
    // try to rewrite it in a half-decent way, ok?
    boost::random::uniform_int_distribution<> uniDist(0, contentList.size()-1);
    for (uint i = 0; i < contentList.size() / 3; i++) {
      uint randC = uniDist(gen);
      viewsIt = views.begin();
      for (uint j = 0; j < randC; j++)
        viewsIt++;
      mIt = viewsIt;
      int j = 0;
      while (mIt != views.begin() && *mIt >= winLimits[randC].first) {
        j--;
        mIt--;
      }
      if (j != 0) j++;
      mIt = viewsIt;
      int k = 0;
      while (mIt != views.end() && *mIt <= winLimits[randC].second) {
        mIt++;
        k++;
      }
      if (k != 0) k--;
      if (j != 0 || k != 0) {
        boost::random::uniform_int_distribution<> switchDist(0,k-j);
        j = switchDist(gen) + j;
        if (j < 0) {
          std::swap(winLimits[randC], winLimits[randC+j]);
          mIt = viewsIt;
          while (j < 0) {
            mIt--;
            j++;
          }
          std::iter_swap(viewsIt, mIt);
        }
        else if (j > 0) {
          std::swap(winLimits[randC], winLimits[randC+j]);
          mIt = viewsIt;
          while (j > 0) {
            mIt++;
            j--;
          }
          std::iter_swap(viewsIt, mIt);
        }
      }
    }
  }
  // assign views in the intended order (sorted or permuted)
  uint totalViews = 0;
  viewsIt = views.begin();
  BOOST_FOREACH (ContentElement* content, contentList) {
    content->setViewsThisRound(*viewsIt);
    totalViews += *viewsIt;
    viewsIt++;
  }
  return totalViews;
}

// Returns the round at which a content will peak in terms of views
unsigned int UGCPopularity::generatePeakRound() {
  boost::uniform_01<> uni;
  double randVal = uni(gen);
//  std::cout << "Uniform Value: " << randVal << std::endl;
  return boost::math::quantile(*ttpDist, randVal);
}

unsigned int UGCPopularity::getBeforePeakTailViews() {
  boost::math::lognormal logDist(2.000, 2.135);
  boost::random::uniform_real_distribution<> uniDist(0.903, 1);
  double randValue = uniDist(gen);
  double retValue = boost::math::quantile(logDist, randValue);
  return static_cast<unsigned int>(0.5+retValue);
}

unsigned int UGCPopularity::getAtPeakTailViews() {
  boost::math::lognormal logDist(-3.826, 3.477);
  boost::random::uniform_real_distribution<> uniDist(0.997, 1);
  double randValue = uniDist(gen);
  double retValue = boost::math::quantile(logDist, randValue);
  return static_cast<unsigned int>(0.5+retValue);
}

unsigned int UGCPopularity::getAfterPeakTailViews() {
  boost::math::lognormal logDist(-0.356, 2.533);
  boost::random::uniform_real_distribution<> uniDist(0.931, 1);
  double randValue = uniDist(gen);
  double retValue = boost::math::quantile(logDist, randValue);
  return static_cast<unsigned int>(0.5+retValue);
}

unsigned int UGCPopularity::getBeforePeakViews() {
  const unsigned int xMin = 0, xThresh = 119;
  unsigned int retValue = 0;
  boost::uniform_01<> uniDist;
  double randVal = uniDist(gen);
  // std::cout << "Uniform Value: " << randVal << std::endl;
  boost::math::beta_distribution<> betaDist(0.191, 1.330);
  retValue = std::floor(0.5+xMin+(xThresh-xMin)*
          boost::math::quantile(betaDist, randVal));
  return retValue;
}

unsigned int UGCPopularity::getAtPeakViews() {
  const unsigned int xMin = 4, xThresh = 297;
  unsigned int retValue = 0;
  boost::uniform_01<> uniDist;
  double randVal = uniDist(gen);
  // std::cout << "Uniform Value: " << randVal << std::endl;
  boost::math::beta_distribution<> betaDist(0.543, 2.259);
  retValue = std::floor(0.5+xMin+(xThresh-xMin)*
            boost::math::quantile(betaDist, randVal));
  return retValue;
}

unsigned int UGCPopularity::getAfterPeakViews() {
  const unsigned int xMin = 0, xThresh = 30;
  unsigned int retValue = 0;
  boost::uniform_01<> uniDist;
  double randVal = uniDist(gen);
  // std::cout << "Uniform Value: " << randVal << std::endl;
  boost::math::beta_distribution<> betaDist(0.077, 0.968);
  retValue = std::floor(0.5+xMin+(xThresh-xMin)*
            boost::math::quantile(betaDist, randVal));
  return retValue;
}


/*
 * =====================================================================================
 *
 *       Filename:  EnvelopFinder.cpp
 *
 *    Description:
 *
 *        Version:  1.0
 *        Created:  9/6/2012 4:48:20 PM
 *       Revision:  none
 *       Compiler:  msvc
 *
 *         Author:  Han Hu (hh1985@bu.edu),
 *   Organization:  Boston University
 *
 * =====================================================================================
 */

#include <algorithm>
#include <GAGPL/MISC/Param.h>
#include <GAGPL/SPECTRUM/EnvelopFinder.h>
#include <GAGPL/SPECTRUM/EnvelopReference.h>
#include <GAGPL/SPECTRUM/IsotopicDistribution.h>
#include <boost/utility.hpp>
//#include <boost/math/distributions/normal.hpp>

namespace gag
{
	EnvelopFinder::EnvelopFinder(RichList& spec)
		: spectrum(spec), param(Param::instance())
	{
		/*
		Guess the range of average formula. Usually we assume the difference between composition estimated from the base peak and the one estimated from monoisotopic peak can be safely ignored.
			No Sulfur (100D)		High Sulfur (100D)
			C: 3.7238523        1.9961864
			H: 5.4425534        2.9942796
			O: 2.8645018        3.1938983
			N: 0.2864502        0.1996186
			S: 0                0.5988559
		*/
		ave1.insert(std::make_pair("C", 3.7238523));
		ave1.insert(std::make_pair("H", 5.4425534));
		ave1.insert(std::make_pair("O", 2.8645018));
		ave1.insert(std::make_pair("N", 0.2864502));
		ave1.insert(std::make_pair("S", 0.0));

		ave2.insert(std::make_pair("C", 1.9961864));
		ave2.insert(std::make_pair("H", 2.9942796));
		ave2.insert(std::make_pair("O", 3.1938983));
		ave2.insert(std::make_pair("N", 0.1996186));
		ave2.insert(std::make_pair("S", 0.5988559));

		run();
	}

	void EnvelopFinder::run()
	{
		// Get/Set parameters.
		unsigned int precursor_charge = param.getParameter<unsigned int>("precursor_charge").first;

		// Divide the spectrum into several independent islands.
		IslandList is_list = this->getPeakIslands(spectrum);

		// Iterate over all independent islands.
		for(IslandList::iterator iter = is_list.begin(); iter != is_list.end(); iter++)
		{
			std::set<EnvelopPtr> temp_env_set;

			RichList island_pks = *iter;
			std::cout << "Found island: " << std::endl;
			island_pks.printPeakList<peak_mz>();
			// Starts from the peak with highest intensity. Assuming it is the base peak.
			RichPeakListByResolution& pks_area = island_pks.getPeakListByType<peak_resolution>();

			RichPeakListByResolution::iterator iter_area = pks_area.begin();
			RichPeakListByResolution::iterator end_iter = pks_area.end();
			// The last peak cannot be considered as a base peak.
			end_iter--;

			/* Step 1. Envelop identification */
			for(; iter_area != end_iter; iter_area++)
			{
				// Iterate over all possible charge states.
				// Notice that the sign of the charge is not used here.
				for(unsigned int z=precursor_charge; z>=1; z--)
				{
					// All candidate envelops should be stored for records.
					this->findNextEnvelops(island_pks, iter_area, z, temp_env_set);
				}

			}

			std::cout << "Before optimizing, output all the envelop information..." << std::endl;
			BOOST_FOREACH(EnvelopPtr env, temp_env_set)
				env_ref.printEnvelopInformation(env);

			std::cout << std::endl;

			// Adjust the fitting scores and splitting scores for envelop set.
			this->optimizeEnvelopSet(temp_env_set);

		}
	}

	IslandList EnvelopFinder::getPeakIslands(RichList& pks)
	{
		double island_space = param.getParameter<double>("island_space").first;

		IslandList is_list;

		RichList* temp_spec = new RichList();

		RichPeakListByMZ& pks_mz = pks.getPeakListByType<peak_mz>();

		for(RichPeakListByMZ::const_iterator iter = pks_mz.begin(); iter != pks_mz.end(); iter++)
		{
			if(iter == pks_mz.begin()) {
				temp_spec->addPeak(*iter);
				continue;
			}

			if((*iter)->mz - (*boost::prior(iter))->mz < island_space) {
				temp_spec->addPeak(*iter);
			} else { //
				is_list.push_back(*temp_spec);

				// Delete the container. Be careful of this section.
				delete temp_spec;

				temp_spec = new RichList();
				temp_spec->addPeak(*iter);
			}
		}
		is_list.push_back(*temp_spec);
		delete temp_spec;

		return is_list;
	}

	void EnvelopFinder::findNextEnvelops(RichList& island, RichPeakListByResolution::iterator iter_area, unsigned int charge, std::set<EnvelopPtr>& env_set)
	{
		// 1. Find candidate envelops which extend from current peak.

		// Project the area iterator into corresponding mz iterator.
		RichPeakListByMZ::iterator iter_mz = island.getPeakContainer().project<peak_mz>(iter_area);

		std::cout << "Identify base peak: " << (*iter_mz)->mz << " Charge: " << charge << std::endl;

		int mode = param.getParameter<int>("mode").first;
		double pre_mz = param.getParameter<double>("precursor_mz").first;
		int pre_z = param.getParameter<int>("precursor_charge").first;

		// Convert m/z value to corresponding mass.
		double pre_mass = calculateMass(pre_mz, pre_z * mode);

		// If the calculated mass is larger than precursor mass, there is no need to continue identification.
		double mass = calculateMass((*iter_mz)->mz, charge * mode);

		if(mass > pre_mass)
			return;

		std::pair<PeakSet,size_t> pk_set = this->extendPeakSet(island, iter_mz, charge);

		if(pk_set.first.getSize()==1)
			return;

    // Container for temporary storage of envelops.
		std::vector<EnvelopPtr> env_store;
		for(size_t i=0; i < pk_set.second; i++) {
			EnvelopPtr env = createEnvelop(charge, pk_set.first.getBoundary());
			env_store.push_back(env);
		}

		// 2. TBD: Estimate the validity of the pk_set.
		// e.g. the overall shifts and the continuity of the shifts.

		// 2.a If the base peak has been reported as a distal peak of other envelops.
		// 2.b If envelops which share the base peak have the same charge
		// 2.c No event report for the base peak.
		// 2.d no contribution of new peaks.
		// 2.e fitting score is worse than expected.

		// Get the maximum number of envelops.

		// Dynamically convert pk_set into vector of envelops.
		ProtoEnv& peak_information = pk_set.first.getPeakInformation();

		ProtoEnv::iterator pk_iter = peak_information.begin();
		// Iterate over all shift, decide if all envelop will share.
		for(; pk_iter != peak_information.end(); )
		{
			int count = peak_information.count(pk_iter->first);
			std::pair<ProtoEnv::iterator, ProtoEnv::iterator> shift_pair = peak_information.equal_range(pk_iter->first);
			ProtoEnv::iterator shift_iter = shift_pair.first;
			ProtoEnv::iterator last_iter = shift_pair.second;
			last_iter--;

			for(size_t i=0; i<env_store.size(); i++)
			{
				// Clean the definition of the internal structure.
				int shift = shift_iter->first;
				RichPeakPtr pk = shift_iter->second.first;
				InfoPeakPtr info = shift_iter->second.second;
				// Register the (peak, env) pair, shift and peak info into the envelop reference table.
				InfoPeakPtr info_copy(new InfoPeak(*info));
				std::cout << "Register peak " << pk->mz << "for envelop " << env_store[i]->id << std::endl;
				env_ref.addDictionaryReference(pk, env_store.at(i), shift, info_copy);

				// Non_unique_key.
				if(shift_iter != last_iter)
					shift_iter++;

			}

			pk_iter = shift_pair.second;
		}

		// Update envelop lambda information for each envelop.
		this->updateEnvelopParameter(env_store);

		// Store the envelops.
		env_set.insert(env_store.begin(), env_store.end());
	}

std::pair<PeakSet, size_t> EnvelopFinder::extendPeakSet(RichList& island, RichPeakListByMZ::iterator iter_mz, unsigned int charge)
	{

		// Candidate base peak. Notice that this might not be the monoisotopic peak for molecules with high mass values.
		RichPeakPtr expr_base_pk = *iter_mz;
		int mode = param.getParameter<int>("mode").first;
		double mass = calculateMass(expr_base_pk->mz, charge * mode);
		EnvelopBoundary bound = this->createBoundary(mass, charge);

		// The reason of creating an object called envelop is because during the extension of peakset towards different directions, there might be multiple matching for given shift.
		PeakSet pk_set(bound);

		InfoPeakPtr pk_infor = boost::make_shared<InfoPeak>(expr_base_pk->resolution, 0.0, 0.5, 0.0);

		pk_set.addPeak(0, expr_base_pk, pk_infor);

		RichPeakListByMZ& pks_mz = island.getPeakListByType<peak_mz>();

		// Shift towards right along the mz axis.
		size_t max_env = this->extendPeak(pk_set, pks_mz, iter_mz, charge, 1);
		// TBD: Shift towards left along the mz axis.
		//this->extendPeak(pk_set, pks_mz, iter_mz, charge, -1);

		return std::make_pair(pk_set,max_env);
	}

	EnvelopBoundary EnvelopFinder::createBoundary(const double mass, unsigned int charge)
	{
		// 1. Estimate the rounded version of compositions.
		double coef = mass / 100.0;
		Composition compo1, compo2;

		AveragineFormulae::iterator iter1 = ave1.begin();
		for(; iter1 != ave1.end(); iter1++)
			compo1.addElement(iter1->first, (int)floor(iter1->second * coef + 0.5));

		AveragineFormulae::iterator iter2 = ave2.begin();
		for(; iter2 != ave2.end(); iter2++)
			compo2.addElement(iter2->first, (int)floor(iter2->second * coef + 0.5));

		int mode = param.getParameter<int>("mode").first;

		int z = mode > 0 ? charge : (-1)*charge;
		// 2. Get the range of the isotopic distribution
		IsotopicDistribution iso_up(compo1);
		AggregatedIsotopicVariants peakset1 = iso_up.getAggregatedIsotopicVariants(z);
		IsotopicDistribution iso_down(compo2);
		AggregatedIsotopicVariants peakset2 = iso_down.getAggregatedIsotopicVariants(z);

		return EnvelopBoundary(peakset1, peakset2);
	}

	size_t EnvelopFinder::extendPeak(PeakSet& pk_set, RichPeakListByMZ& pks_mz,RichPeakListByMZ::iterator iter_mz, unsigned int charge, int direction)
	{
		//int dist = direction > 0 ? std::distance(iter_mz, pks_mz.end())-1 : std::distance(pks_mz.begin(), iter_mz);

		unsigned int shift = 1;
		//int prev_dist = 0;

		// Experimental candidate base peak.
		RichPeakPtr expr_base_pk = *iter_mz;

		int mode = param.getParameter<int>("mode").first;
		// For confidence estimation, using internal_accuracy.
		double internal_accuracy = param.getParameter<double>("internal_accuracy").first;

		// For peak finding and lambda estimation, using external_accuracy.
		double external_accuracy = param.getParameter<double>("external_accuracy").first;

		// Calculate the mass of the candidate base peak.
		// Notice that there will be mechanism where there is electron capture.
		double mass = calculateMass(expr_base_pk->mz, charge * mode);

		// Create theoretical boundary. The boundary has been adjusted to fit the charge.
		//EnvelopBoundary env_bound = this->createBoundary(mass, charge);
		EnvelopBoundary& env_bound = pk_set.getBoundary();

		//std::cout << "Up Boundary: " << std::endl;
		//env_bound.first.printPeakList<peak_mz>();
		//std::cout << "Down Boundary: " << std::endl;
		//env_bound.second.printPeakList<peak_mz>();

		// Get a copy of the base peak iterator for further move operation.
		//RichPeakListByMZ::iterator iter = iter_mz;
		int stop_flag = 1;

		// Dynamically convert pk_set into vector of envelops.
		std::pair<unsigned int, size_t> shift_count(0,1);
		size_t max_env = 0;

		//std::advance(iter_mz, direction);
		if(direction < 0 && iter_mz == pks_mz.begin())
			return 0;
		std::advance(iter_mz, direction);
		if(direction > 0 && iter_mz == pks_mz.end())
			return 0;

		// TBD: the program should efficiently location the position of most likely peak. instead of sequentially search for it.
		while(1)
		{

			std::cout << "Shift: " << shift << std::endl;

			// Get the boundary peaks at current shift.
			PeakPtr theo_pk1 = env_bound.getUpperBound().getPeakByShift<peak_intensity>(direction * shift);
			PeakPtr theo_pk2 = env_bound.getLowerBound().getPeakByShift<peak_intensity>(direction * shift);

			// No extension any more.
			if(theo_pk1->mz == 0.0 && theo_pk2->mz == 0.0)
				break;

			RichPeakPtr current_pk = *iter_mz;
			std::cout << "Examine peak: " << current_pk->mz << std::endl;

			// Experimental distance from current peak to base peak.
			// Notice that the value might be negative.
			double expr_dist = abs(current_pk->mz - expr_base_pk->mz);

			// Theoretical distance from current peak to base peak. Adjusted for charged peak.
			double theo_dist1 = abs(env_bound.getUpperBound().getMassDifferenceByShift<peak_intensity>(0, shift));
			double theo_dist2 = abs(env_bound.getLowerBound().getMassDifferenceByShift<peak_intensity>(0, shift));

			double max_dist = theo_dist1;
			double min_dist = theo_dist2;

			if(max_dist < min_dist) {
				swap(min_dist, max_dist);
			}

			double error1 = 1e6 * (expr_dist - min_dist)/expr_base_pk->mz;
			double error2 = 1e6 * (expr_dist - max_dist)/expr_base_pk->mz;

			if(error1 < -1 * external_accuracy) {
				// Before the shift region: keep moving the iterator.
				// std::advance(iter_mz, direction); continue;
			} else if(error2 > external_accuracy) {
				// Beyond the shift region: update shift.
				if(stop_flag == 0) {
					shift++; stop_flag = 1; continue;
				} else {
					// break means no missing peak is allowed in the middle. This might be controlled by some parameter.
					stop_flag = 1;
					break;
				}
			} else {
				// A matching peak. Notice the effect of charge state.
				std::cout << "A matching peak!" << std::endl;

				if(shift == shift_count.first) {
					shift_count.second++;
				} else {
					// Reset shift_count.
					shift_count = std::make_pair(shift, 1);
				}

				// Update maximum number of envelops.
				if(shift_count.second > max_env)
					max_env = shift_count.second;

				stop_flag = 0;
				// Estimate the confidence of the mz in terms of estimating lambda.
				double diff = max_dist - min_dist;

				// Error window is used for confidence estimation.
				double err_win = internal_accuracy * expr_base_pk->mz * 1e-6;
				double prob = diff/(diff + 2*err_win);

				// Indicate the sign of the error window.
				int err_sign = (theo_dist1 < theo_dist2 ? -1 : 1);

				//: lambda_mz should always be in (0, 1). If lambda is closed to 1, the pattern is closed to no-sulfate boundary, and 0 for high-sulfate boundary.
				double lambda_mz = (expr_dist - theo_dist2+err_sign*err_win)/(theo_dist1 - theo_dist2 + 2*err_sign*err_win);

				// Round the abnormal lambda_mz to the boundary.
				if(lambda_mz > 1) {
					lambda_mz = 1.0;
					prob = 1.0;
				} else if(lambda_mz < 0) {
					lambda_mz = 0.0;
					prob = 1.0;
				}

				// TBD: estimate lambda_abd (within 5%) and delta_lambda.
				// Notice: lambda_abd can be < 0 or > 1.
				double lambda_abd = (current_pk->resolution/expr_base_pk->resolution - theo_pk2->intensity)/(theo_pk1->intensity - theo_pk2->intensity);

				InfoPeakPtr pk_infor = boost::make_shared<InfoPeak>(current_pk->resolution, lambda_mz, lambda_abd, prob);
				// Notice: The abundance information should be initialized during the establish of global lambda value.
				//pk_infor.adjusted_abundance = env_bound.getTheoreticalPeak(lambda_abd, shift).intensity;

				pk_set.addPeak(shift, current_pk, pk_infor);

			}
			if(direction < 0 && iter_mz == pks_mz.begin())
				break;
			std::advance(iter_mz, direction);
			if(direction > 0 && iter_mz == pks_mz.end())
				break;
		}
		return max_env;
	}

	void EnvelopFinder::updateEnvelopParameter(EnvelopPtr env)
	{
		EnvEntry base_entry = env_ref.getEntryByShift(env, 0);

		// 1.a Find lambda value for each of the peak.
		std::vector<EnvEntry> peak_entries = env_ref.getEntryByEnvelop(env);

		double current_prob = 0.0;
		BOOST_FOREACH(EnvEntry& entry, peak_entries)
		{
			if(entry.info->prob > current_prob) {
				current_prob = entry.info->prob;
				env->lambda_mz = entry.info->lambda_mz;
			}

		}
		//// 1.b Normalize the confidence value.
		//double sum_confidence = 0.0;
		//for(std::vector<EnvEntry>::iterator iter = peak_entries.begin();
		//	iter != peak_entries.end(); iter++)
		//	sum_confidence += pow(iter->info->prob,2);

		//double lambda_mz = 0.0;
		//for(std::vector<EnvEntry>::iterator iter = peak_entries.begin();
		//	iter != peak_entries.end(); iter++) {
		//		// Pass base peak.
		//		if(iter->getShift() == 0)
		//			continue;

		//		lambda_mz += pow(iter->info->prob,2)/sum_confidence *  iter->info->lambda_mz;
		//}
		//env->lambda_mz = lambda_mz;

		// 2. Update theoretical abundance from the lambda.
		for(std::vector<EnvEntry>::iterator iter = peak_entries.begin();
			iter != peak_entries.end(); iter++) {
			// The adjusted abundance is initialized as the exp data.
			iter->info->adjusted_abundance = iter->pk->area;
		}
	}

	void EnvelopFinder::updateEnvelopParameter( std::vector<EnvelopPtr> env_list )
	{
		BOOST_FOREACH(EnvelopPtr env, env_list) {
			this->updateEnvelopParameter(env);

			//this->updateSplittingPotential(env);
		}
	}

	void EnvelopFinder::optimizeEnvelopSet(std::set<EnvelopPtr>& env_set, double last_score)
	{
		/* New version */
		this->updateSplittingPotential(env_set);
		this->updateFittingScore(env_set);

		// Get all the peaks.
		RichList pk_list = env_ref.getRichPeakList(env_set);

		RichPeakListByMZ& pks_mz = pk_list.getPeakListByType<peak_mz>();

		for(RichPeakListByMZ::iterator iter = pks_mz.begin(); iter != pks_mz.end(); iter++)
		{
			// Split the peak based on fitting score & estimated abundance.
			// Notice that the identification of FP will also happen here.
			this->splitPeak(*iter);
		}
		// Calculate the overall fitting scores.
		double next_score = this->updateFittingScore(env_set);
		std::cout << "Old score: " << last_score << "\t" << "New score: " << next_score << std::endl;
		if(abs(next_score - last_score)/last_score < 1e-5)
			return;
		else
			this->optimizeEnvelopSet(env_set, next_score);
		/* Old version */

		//std::set<RichPeak> processed_peak_set;
		//std::set<RichPeak>::iterator peak_iter;

		//// Generate fitting score for each envelop.
		//// Estimate scaling factor for each envelop.
		//this->updateSplittingPotential(env_pool);
		//this->updateFittingScore(env_pool);

		//sort(env_pool.begin(), env_pool.end(), Envelop::scoreLarger);

		//// Iterate over all envelops with intensity from high to low.
		//BOOST_FOREACH(EnvelopPtr& env, env_pool)
		//{
		//	// If the base peak has been processed.
		//	EnvEntry base_entry = env_ref.getBasePeakForEnvelop(env);
		//
		//	peak_iter = processed_peak_set.find(base_entry->pk);
		//	if(peak_iter != processed_peak_set.end())
		//		continue;

		//	// Split the base peak.
		//	bool status = this->splitPeak(base_entry->pk);
		//	if(status) {
		//		this->updateSplittingPotential(env);
		//		this->updateFittingScore(env);
		//	}

		//	// Get all belonging peaks.
		//	std::vector<EnvEntry> entry_vec = env_ref.getEntryByEnvelop(env);

		//	// Iterate over all peaks with intensity from high to low.
		//	BOOST_FOREACH(EnvEntry& entry, entry_vec)
		//	{
		//		if(entry->getShift()==0)
		//			continue;
		//		// Split the peak based on 1. splitting potential and 2. scaling factor.
		//		this->splitPeak(entry->pk);
		//		// Update the records.
		//		processed_peak_set.insert(entry->pk);
		//	}

		//}
	}

	void EnvelopFinder::updateSplittingPotential(EnvelopPtr env)
	{

		//double scaling_factor = env->scaling_factor;

		// Entry for base peak.
		EnvEntry& base_entry = env_ref.getEntryByShift(env, 0);
		std::cout << "The base peak is: " << base_entry.pk->mz << std::endl;

		// Get all peaks included by the envelop.
		std::vector<EnvEntry> entry_vec = env_ref.getEntryByEnvelop(env);

		// Update the scaling factor. The value for base peak should be 1.
		BOOST_FOREACH(EnvEntry& entry, entry_vec)
		{
			double expected_abundance = env_ref.getTheoreticalAbundance(entry);
			// Adjusted abunance / Theoretical abundance.
			double scaling_factor = entry.info->adjusted_abundance/expected_abundance;
			if(scaling_factor < env->scaling_factor)
				env->scaling_factor = scaling_factor;
		}

		// Update the adjusted abundance.
		// 1. Update base entry.
		base_entry.info->adjusted_abundance *= env->scaling_factor;
		// 2. Update all the rest of the peaks.
		BOOST_FOREACH(EnvEntry& entry, entry_vec)
		{
			if(entry.getShift()==0)
				continue;

			entry.info->adjusted_abundance = env_ref.getTheoreticalAbundance(entry);
		}

		//// Iterate over all theoretical peaks (except base peak), calculate their relative shift.
		//double total_share = 0.0;
		//std::map<int, PeakPtr> pk_map = env->getTheoreticalPeaks();
		//for(std::map<int, PeakPtr>::iterator iter = pk_map.begin();
		//	iter != pk_map.end(); iter++)
		//{
		//	total_share += iter->second->intensity;
		//}
		//
		//double base_adjusted_abd = 0.0;
		//for(std::map<int, PeakPtr>::iterator iter = pk_map.begin();
		//	iter != pk_map.end(); iter++)
		//{
		//	int shift = iter->first; PeakPtr pk = iter->second;

		//	// Get entry corresponding to the shift. Notice the problem if shift does not exist for the envelop..
		//	EnvEntry entry = env_ref.getEntryByShift(env, shift);

		//	// If EnvEntry is empty. Which means the experimental peak corresponds to the shift doesn't show up.
		//	if(entry.isEmpty()) continue;

		//	entry.info->relative_share = pk->intensity / total_share;

		//	double expected_abundance = env_ref.getTheoreticalAbundance(entry);

		//	if(expected_abundance > entry.info->adjusted_abundance) {
		//		// Potential overlapping event for base peak. Scale the base peak.
		//		base_adjusted_abd += base_entry.info->adjusted_abundance * (entry.info->adjusted_abundance / expected_abundance) * entry.info->relative_share;
		//	}
		//	// Update adjusted abundance.
		//	entry.info->adjusted_abundance = expected_abundance;
		//}

		//// Update base peak information.
		//base_entry.info->adjusted_abundance = base_adjusted_abd;
	}

	void EnvelopFinder::updateSplittingPotential( std::vector<EnvEntry> entry_vec )
	{
		BOOST_FOREACH(EnvEntry entry, entry_vec)
			this->updateSplittingPotential(entry.env);
	}
	void EnvelopFinder::updateSplittingPotential(std::set<EnvelopPtr>& env_set)
	{
		BOOST_FOREACH(EnvelopPtr env, env_set)
			this->updateSplittingPotential(env);
	}

	double EnvelopFinder::updateFittingScore(EnvelopPtr env)
	{
		std::vector<EnvEntry> entry_vec = env_ref.getEntryByEnvelop(env);
		//RichPeakPtr base_pk = env_ref.getBasePeakForEnvelop(env);
		EnvEntry base_entry = env_ref.getEntryByShift(env, 0);

		double fitting_score = 0.0;
		BOOST_FOREACH(EnvEntry& entry, entry_vec)
		{
			// For base peak, intensity_shift and mz_shift is always 1.
			if(entry.getShift()==0) {
				fitting_score += sqrt(base_entry.info->adjusted_abundance);
				continue;
			}

			double expected_abundance = env_ref.getTheoreticalAbundance(entry);
			// The fitting score comes from 3 parts:
			// 1. the shift of intensity: abs(E-T)/T.
			double intensity_shift = 1.0 - abs(expected_abundance - entry.info->adjusted_abundance)/expected_abundance;
			if(intensity_shift < 0)
				intensity_shift = 0.0;
			//std::cout << "Intensity shift: " << intensity_shift << std::endl;

			// 2. the relative shift of mz.
			double internal_accuracy = param.getParameter<double>("internal_accuracy").first;

			double mz_shift = 1.0 - abs(entry.pk->mz - base_entry.pk->mz - entry.env->getTheoreticalMZDistance(0, entry.getShift()))/(entry.env->getWindowSizeByShift(entry.getShift()));
			//std::cout << "MZ shift: " << mz_shift << std::endl;

			// 3. the square root of expected_abundance.
			double weight = sqrt(expected_abundance);
			//std::cout << "Weight: " << weight << std::endl;
			fitting_score += weight * intensity_shift * mz_shift;
		}
		//std::cout << "Fitting score: " << fitting_score << std::endl;
		env->fitting_score = fitting_score;
		return fitting_score;
	}

	double EnvelopFinder::updateFittingScore( std::vector<EnvEntry> entry_vec)
	{
		double sum_score = 0.0;
		BOOST_FOREACH(EnvEntry& entry, entry_vec)
			sum_score += this->updateFittingScore(entry.env);
		return sum_score;
	}

	double EnvelopFinder::updateFittingScore(std::set<EnvelopPtr>& env_set)
	{
		double sum_score = 0.0;
		BOOST_FOREACH(EnvelopPtr env, env_set)
			sum_score += this->updateFittingScore(env);
		return sum_score;
	}

	void EnvelopFinder::adjustLambdaFromProfile( EnvelopPtr env )
	{
		double last_score = this->updateFittingScore(env);
		// The value of lambda_bound has to be in between 0.0 and 1.0
		double lambda_bound = 0.0;

		while(1)
		{
			// Adjust lambda. Notice the direction.
			env->lambda_mz = (lambda_bound + env->lambda_mz)/2;

			// Update splitting score.
			this->updateSplittingPotential(env);

			double new_score = this->updateFittingScore(env);

			double score_diff = new_score - last_score;

			// Converge.
			if(abs(score_diff) < 1e-6)
				break;

			if(score_diff > 0.0) {
					env->lambda_mz = (lambda_bound + env->lambda_mz)/2;
			} else {
					// Change the bound.
					lambda_bound = 1.0-lambda_bound;
					env->lambda_mz = (lambda_bound + env->lambda_mz)/2;
			}

		}

	}

	void EnvelopFinder::splitPeak(RichPeakPtr pk)
	{
		// Get all connecting envelops.
		std::vector<EnvEntry> entry_vec = env_ref.getEntryByPeak(pk);

		if(entry_vec.size() == 1) // No need to split for single envelop.
			return;
		else if(entry_vec.size() == 0)
			throw std::runtime_error("Undefined peak in the dictionary");
		else
			;

		double splitting_score = 0.0;

		BOOST_FOREACH(EnvEntry& entry, entry_vec)
		{
			// splitting_score += entry.info->adjusted_abundance;
			splitting_score += sqrt(entry.env->fitting_score * entry.info->adjusted_abundance);
		}

		// Split the peak.
		BOOST_FOREACH(EnvEntry& entry, entry_vec)
		{
			double scale_coef = sqrt(entry.env->fitting_score * entry.info->adjusted_abundance)/ splitting_score;

			// Update the estimated abundance.
			entry.info->adjusted_abundance = pk->resolution * scale_coef;

			env_ref.printEnvelopInformation(entry.env);
		}

	}

	void EnvelopFinder::printEnvelop( EnvelopPtr env )
	{
		std::cout << "Envelop ID: " << env->id << std::endl;
		std::cout << "Charge: " << env->charge_state << std::endl;
		std::cout << "Lambda: " << env->lambda_mz << std::endl;
		std::cout << "Score: " << env->fitting_score << std::endl;
		std::cout << std::endl;
		// Peak information.
		std::cout << "Covered peaks:" << std::endl;
		std::cout << "Shift\tMZ\tEXP_ABD\tFIT_ABD" << std::endl;

		// Get all peaks.
		std::vector<EnvEntry> pk_entries = env_ref.getEntryByEnvelop(env);
		BOOST_FOREACH(EnvEntry& pk_entry, pk_entries)
		{
			std::cout << pk_entry.getShift() << "\t" << pk_entry.pk->mz << "\t"
				<< pk_entry.pk->resolution << "\t" << pk_entry.info->adjusted_abundance
				<< std::endl;
		}
		std::cout << std::endl;
	}

	//void EnvelopFinder::optimizeEnvelopSet(RichList& pk_list)
	//{
	//	// 1. Iterate over peaks with m/z from low to high
	//	RichPeakListByMZ& pks_by_mz = pk_list.getPeakListSortedByType<peak_mz>();

	//	RichPeakListByMZ::iterator iter_mz = pks_by_mz.begin();

	//	// Envelop set and corresponding fitting score.
	//	EnvelopSetCollection last_frontier;

	//	// TBD: how to decide the cluster group will end.
	//	for(; iter_mz != pks_by_mz.end(); iter_mz++)
	//	{
	//		// 2. For each peak, find all connecting envelops.
	//		std::vector<EnvEntry> entry_vec = env_ref.getEntryByPeak(*iter_mz);

	//		// Empty information. Consider storing the information with a separate container.
	//		if(entry_vec.size() == 0)
	//			continue;

	//		std::vector<std::set<EnvelopPtr> > env_set_collection;
	//		for(size_t i = 0; i < entry_vec.size()-1; i++)
	//			for(size_t j=i; j<entry_vec.size(); j++)
	//			{
	//				std::set<EnvelopPtr> temp_env;
	//				temp_env.insert(entry_vec[i]->env);

	//				if(i != j) // Not just pairs.
	//					temp_env.insert(entry_vec[j]->env);

	//				env_set_collection.insert(temp_env);
	//			}

	//		// New frontier.
	//		if(last_frontier.size() == 0)
	//		{
	//			std::vector<std::set<EnvelopPtr> >::iterator iter = env_set_collection.begin();
	//			for(; iter != env_set_collection.end(); iter++)
	//			{
	//				double temp_score = this->calculateFittingScore(*iter);
	//				last_frontier.insert(std::make_pair(temp_score, *iter));
	//			}
	//			continue;
	//		}

	//		EnvelopSetCollection next_frontier;
	//		// Update current frontier.
	//		for(EnvelopSetCollection::iterator fron_iter = last_frontier.begin(); fron_iter != last_frontier.end(); fron_iter++)
	//		{
	//			BOOST_FOREACH(std::set<EnvelopPtr>& env_set, env_set_collection)
	//			{
	//				std::set<EnvelopPtr> temp_set(fron_iter->second);
	//				// Append the set.
	//				set_union(temp_set.begin(), temp_set.end(),
	//					env_set.begin(), env_set.end(),
	//					inserter(temp_set, temp_set.begin()));

	//				// Compare the fitting score between std::vector<EnvelopPtr> with and without new current envelops.
	//				double score1 = fron_iter->first;
	//				double score2 = this->calculateFittingScore(new_set);

	//				if(score2 > score1)
	//					next_frontier.insert(std::make_pair(score2, temp_set));
	//				else
	//					next_frontier.insert(std::make_pair(score1, fron_iter->second));
	//			}
	//		}

	//		last_frontier = next_frontier;
	//
	//	}

	//	// Select the set with maximum score.
	//	std::set<EnvelopPtr>& winner_set = last_frontier.rbegin()->second;
	//
	//	env_pool.insert(env_pool.end(), winner_set.begin(), winner_set.end());
	//
	//}

	//double EnvelopFinder::calculateFittingScore( std::set<EnvelopPtr> env_set )
	//{
	//	/* Estimate the abundance based on current set selection. */
	//

	//	double total_score = 0.0;
	//	BOOST_FOREACH(EnvelopPtr& env, env_set)
	//	{
	//		total_score += this->calculateFittingScore(env);
	//	}
	//	return total_score;
	//}

	//double EnvelopFinder::calculateFittingScore(EnvelopPtr env)
	//{
	//	std::vector<EnvEntry> entry_vec = env_ref.getEntryByEnvelop(env);
	//
	//	if(entry_vec.size()==0)
	//		return 0.0;

	//	// Get base entry.
	//	EnvEntry base_entry = env_ref.getEntryByShift(env, 0);

	//	double fitting_score = 0.0;
	//	BOOST_FOREACH(EnvEntry& entry, entry_vec)
	//	{
	//		// For base peak, intensity_shift and mz_shift is always 1.
	//		if(entry.getShift()==0) {
	//			fitting_score += sqrt(base_entry.info->adjusted_abundance);
	//			continue;
	//		}

	//		double expected_abundance = env_ref.getTheoreticalAbundance(entry);
	//
	//		// The fitting score comes from 3 parts:
	//		// 1. the shift of intensity: abs(E-T)/T.
	//		double intensity_shift = 1.0 - abs(expected_abundance - entry.info->adjusted_abundance)/expected_abundance;
	//		if(intensity_shift < 0)
	//			intensity_shift = 0.0;

	//		// 2. the relative shift of mz. Here the accuracy depends on the position of the peak in the envelop.
	//		double mz_shift = 1.0 - abs(entry.pk->mz - base_entry.pk->mz - entry.env->getTheoreticalMZDistance(0, entry.getShift()))/(entry.env->getWindowSizeByShift(entry.getShift()));

	//		// 3. the square root of expected_abundance.
	//		double weight = sqrt(expected_abundance);

	//		fitting_score += weight * intensity_shift * mz_shift;
	//	}

	//	env->fitting_score = fitting_score;
	//	return fitting_score;
	//}
}


/** kmeans.cc
    Jeremy Barnes, 16 December 2014
    Copyright (c) 2014 mldb.ai inc.  All rights reserved.

    This file is part of MLDB. Copyright 2015 mldb.ai inc. All rights reserved.
    Implementation of an KMEANS algorithm for embedding of a dataset.
*/

#include "mldb/utils/distribution_simd.h"
#include "kmeans_interface.h"
#include "mldb/builtin/matrix.h"
#include "mldb/core/mldb_engine.h"
#include "mldb/core/dataset.h"
#include "mldb/utils/distribution.h"
#include <boost/multi_array.hpp>
#include "mldb/base/scope.h"
#include "mldb/base/parallel.h"
#include "mldb/utils/pair_utils.h"
#include "mldb/arch/timers.h"
#include "mldb/arch/simd_vector.h"
#include "mldb/utils/vector_utils.h"
#include "mldb/types/basic_value_descriptions.h"
#include "mldb/plugins/jml/kmeans.h"
#include "mldb/sql/sql_expression.h"
#include "mldb/builtin/sql_config_validator.h"
#include "mldb/core/analytics.h"
#include "mldb/types/any_impl.h"
#include "mldb/types/optional_description.h"
#include "mldb/vfs/fs_utils.h"
#include "mldb/vfs/filter_streams.h"

using namespace std;

namespace MLDB {

DEFINE_STRUCTURE_DESCRIPTION(KmeansConfig);

KmeansConfigDescription::
KmeansConfigDescription()
{
    Optional<PolyConfigT<Dataset> > optional;
    optional.emplace(PolyConfigT<Dataset>().
                     withType(KmeansConfig::defaultOutputDatasetType));

    addField("trainingData", &KmeansConfig::trainingData,
             "Specification of the data for input to the k-means procedure.  This should be "
             "organized as an embedding, with each selected row containing the same "
             "set of columns with numeric values to be used as coordinates.  The select statement "
             "does not support groupby and having clauses.");
    addField("outputDataset", &KmeansConfig::output,
             "Dataset for cluster assignment.  This dataset will contain the same "
             "row names as the input dataset, but the coordinates will be replaced "
             "by a single column giving the cluster number that the row was assigned to.",
              optional);
    addField("centroidsDataset", &KmeansConfig::centroids,
             "Dataset in which the centroids will be recorded.  This dataset will "
             "have the same coordinates (columns) as those selected from the input "
             "dataset, but will have one row per cluster, providing the centroid of "
             "the cluster.",
             optional);
    addField("numInputDimensions", &KmeansConfig::numInputDimensions,
             "Number of dimensions from the input to use (-1 = all).  This limits "
             "the number of columns used.  Columns will be ordered alphabetically "
             "and the lowest ones kept.",
             -1);
    addField("numClusters", &KmeansConfig::numClusters,
             "Number of clusters to create.  This will provide the total number of "
             "centroids created.  There must be at least as many rows selected as "
             "clusters.", 10);
    addField("maxIterations", &KmeansConfig::maxIterations,
             "Maximum number of iterations to perform.  If no convergeance is "
             "reached within this number of iterations, the current clustering "
             "will be returned.", 100);
    addField("metric", &KmeansConfig::metric,
             "Metric space in which the k-means distances will be calculated. "
             "Normally this will be Cosine for an orthonormal basis, and "
             "Euclidian for another basis",
             METRIC_COSINE);
    addField("modelFileUrl", &KmeansConfig::modelFileUrl,
             "URL where the model file (with extension '.kms') should be saved. "
             "This file can be loaded by the ![](%%doclink kmeans function). "
             "This parameter is optional unless the `functionName` parameter is used.");
    addField("functionName", &KmeansConfig::functionName,
             "If specified, an instance of the ![](%%doclink kmeans function) of this name will be created using "
             "the trained model. Note that to use this parameter, the `modelFileUrl` must "
             "also be provided.");
    addParent<ProcedureConfig>();

    onPostValidate = chain(validateQuery(&KmeansConfig::trainingData,
                                         MustContainFrom(),
                                         NoGroupByHaving()),
                           validateFunction<KmeansConfig>());
}

// TODO: see http://www.eecs.tufts.edu/~dsculley/papers/fastkmeans.pdf

namespace {

ML::KMeansMetric * makeMetric(MetricSpace metric)
{
    switch (metric) {
    case METRIC_EUCLIDEAN:
        return new ML::KMeansEuclideanMetric();
        break;
    case METRIC_COSINE:
        return new ML::KMeansCosineMetric();
        break;
    default:
        throw MLDB::Exception("Unknown kmeans metric type");
    }
}

} // file scope

/*****************************************************************************/
/* KMEANS PROCEDURE                                                           */
/*****************************************************************************/

KmeansProcedure::
KmeansProcedure(MldbEngine * owner,
               PolyConfig config,
               const std::function<bool (const Json::Value &)> & onProgress)
    : Procedure(owner)
{
    kmeansConfig = config.params.convert<KmeansConfig>();
}

Any
KmeansProcedure::
getStatus() const
{
    return Any();
}

RunOutput
KmeansProcedure::
run(const ProcedureRunConfig & run,
      const std::function<bool (const Json::Value &)> & onProgress) const
{
    auto runProcConf = applyRunConfOverProcConf(kmeansConfig, run);

    // an empty url is allowed but other invalid urls are not
    if(!runProcConf.modelFileUrl.empty() && !runProcConf.modelFileUrl.valid()) {
        throw AnnotatedException(400, "modelFileUrl \"" +
                                  runProcConf.modelFileUrl.toUtf8String()
                                  + "\" is not valid");
    }

    if (!runProcConf.modelFileUrl.empty()) {
        checkWritability(runProcConf.modelFileUrl.toDecodedString(), "modelFileUrl");
    }

    auto onProgress2 = [&] (const Json::Value & progress)
        {
            Json::Value value;
            value["dataset"] = progress;
            return onProgress(value);
        };

    SqlExpressionMldbScope context(engine);

    ConvertProgressToJson convertProgressToJson(onProgress);
    auto embeddingOutput = getEmbedding(*runProcConf.trainingData.stm,
                                        context,
                                        runProcConf.numInputDimensions,
                                        convertProgressToJson);

    std::vector<std::tuple<RowHash, RowPath, std::vector<double>,
                           std::vector<ExpressionValue> > > & rows
        = embeddingOutput.first;
    std::vector<KnownColumn> & vars = embeddingOutput.second;

    std::vector<ColumnPath> columnNames;
    for (auto & v: vars) {
        columnNames.push_back(v.columnName);
    }

    std::vector<distribution<float> > vecs;

    for (unsigned i = 0;  i < rows.size();  ++i) {
        vecs.emplace_back(distribution<float>(std::get<2>(rows[i]).begin(),
                                                  std::get<2>(rows[i]).end()));
    }

    if (vecs.size() == 0)
        throw AnnotatedException(400, "Kmeans training requires at least 1 datapoint. "
                                  "Make sure your dataset is not empty and that your WHERE expression "
                                  "does not filter all the rows");

    ML::KMeans kmeans;
    kmeans.metric.reset(makeMetric(runProcConf.metric));

    vector<int> inCluster;

    int numClusters = runProcConf.numClusters;
    int numIterations = runProcConf.maxIterations;

    kmeans.train(vecs, inCluster, numClusters, numIterations);

    bool saved = false;
    if (!runProcConf.modelFileUrl.empty()) {
        try {
            makeUriDirectory(runProcConf.modelFileUrl.toDecodedString());
            Json::Value md;
            md["algorithm"] = "MLDB k-Means model";
            md["version"] = 1;
            md["columnNames"] = jsonEncode(columnNames);

            filter_ostream stream(runProcConf.modelFileUrl);
            stream << md.toString();
            MLDB::DB::Store_Writer writer(stream);
            kmeans.serialize(writer);
            saved = true;
        }
        catch (const std::exception & exc) {
            throw AnnotatedException(400, "Error saving kmeans centroids at location'" +
                                      runProcConf.modelFileUrl.toString() + "': " +
                                      exc.what());
        }
    }

    if (runProcConf.output.get()) {

        PolyConfigT<Dataset> outputDataset = *runProcConf.output;
        if (outputDataset.type.empty())
            outputDataset.type = KmeansConfig::defaultOutputDatasetType;

        auto output = createDataset(engine, outputDataset, onProgress2, true /*overwrite*/);

        Date applyDate = Date::now();

        for (unsigned i = 0;  i < rows.size();  ++i) {
            std::vector<std::tuple<ColumnPath, CellValue, Date> > cols;
            cols.emplace_back(ColumnPath("cluster"), inCluster[i], applyDate);
            output->recordRow(std::get<1>(rows[i]), cols);
        }

        output->commit();
    }

    if (runProcConf.centroids.get()) {

        PolyConfigT<Dataset> centroidsDataset = *runProcConf.centroids;
        if (centroidsDataset.type.empty())
            centroidsDataset.type = KmeansConfig::defaultOutputDatasetType;

        auto centroids = createDataset(engine, centroidsDataset, onProgress2, true /*overwrite*/);

        Date applyDate = Date::now();

        for (unsigned i = 0;  i < kmeans.clusters.size();  ++i) {
            auto & cluster = kmeans.clusters[i];

            std::vector<std::tuple<ColumnPath, CellValue, Date> > cols;

            for (unsigned j = 0;  j < cluster.centroid.size();  ++j) {
                cols.emplace_back(columnNames[j], cluster.centroid[j], applyDate);
            }

            centroids->recordRow(RowPath(MLDB::format("%i", i)), cols);
        }

        centroids->commit();
    }

    if(!runProcConf.functionName.empty()) {
        if (saved) {
            KmeansFunctionConfig funcConf;
            funcConf.modelFileUrl = runProcConf.modelFileUrl;

            PolyConfig kmeansFuncPC;
            kmeansFuncPC.type = "kmeans";
            kmeansFuncPC.id = runProcConf.functionName;
            kmeansFuncPC.params = funcConf;

            createFunction(engine, kmeansFuncPC, onProgress, true);
        } else {
            throw AnnotatedException(400, "Can't create kmeans function '" +
                                      runProcConf.functionName.rawString() +
                                      "'. Have you provided a valid modelFileUrl?",
                                      "modelFileUrl", runProcConf.modelFileUrl.toString());
        }
    }

    return Any();
}

DEFINE_STRUCTURE_DESCRIPTION(KmeansFunctionConfig);

KmeansFunctionConfigDescription::
KmeansFunctionConfigDescription()
{
    addField("modelFileUrl", &KmeansFunctionConfig::modelFileUrl,
             "URL of the model file (with extension '.kms') to load. "
             "This file is created by the ![](%%doclink kmeans.train procedure).");

    onPostValidate = [] (KmeansFunctionConfig * cfg,
                         JsonParsingContext & context) {
        // this includes empty url
        if(!cfg->modelFileUrl.valid()) {
            throw MLDB::Exception("modelFileUrl \"" + cfg->modelFileUrl.toString()
                                + "\" is not valid");
        }
    };
}

/*****************************************************************************/
/* KMEANS FUNCTION                                                           */
/*****************************************************************************/

DEFINE_STRUCTURE_DESCRIPTION(KmeansFunctionArgs);

KmeansFunctionArgsDescription::
KmeansFunctionArgsDescription()
{
    addField("embedding", &KmeansFunctionArgs::embedding,
             "Embedding values for the k-means function.  The column names in "
             "this embedding must match those used in the original kmeans "
             "dataset.");
}

DEFINE_STRUCTURE_DESCRIPTION(KmeansExpressionValue);

KmeansExpressionValueDescription::KmeansExpressionValueDescription()
{
    addField("cluster", &KmeansExpressionValue::cluster,
             "Index of the row in the `centroids` dataset whose columns describe "
             "the point which is closest to the input according to the `metric` "
             "specified in training.");
}

struct KmeansFunction::Impl {
    ML::KMeans kmeans;
    std::vector<ColumnPath> columnNames;

    Impl(const Url & modelFileUrl)
    {
        filter_istream stream(modelFileUrl);
        std::string firstLine;
        std::getline(stream, firstLine);
        Json::Value md = Json::parse(firstLine);
        if (md["algorithm"] != "MLDB k-Means model") {
            throw AnnotatedException(400, "Model file is not a k-means model");
        }
        if (md["version"].asInt() != 1) {
            throw AnnotatedException(400, "k-Means model version is wrong");
        }
        columnNames = jsonDecode<std::vector<ColumnPath> >(md["columnNames"]);

        MLDB::DB::Store_Reader store(stream);
        kmeans.reconstitute(store);
    }
};

KmeansFunction::
KmeansFunction(MldbEngine * owner,
               PolyConfig config,
               const std::function<bool (const Json::Value &)> & onProgress)
    : BaseT(owner, config)
{
    functionConfig = config.params.convert<KmeansFunctionConfig>();

    impl.reset(new Impl(functionConfig.modelFileUrl));

    dimension = impl->kmeans.clusters[0].centroid.size();
}

KmeansExpressionValue
KmeansFunction::
call(KmeansFunctionArgs input) const
{
    Date ts = input.embedding.getEffectiveTimestamp();

    return {ExpressionValue
            (impl->kmeans.assign
             (input.embedding.getEmbedding(impl->columnNames.data(),
                                           impl->columnNames.size())
              .cast<float>()),
             ts)};
}

namespace {

RegisterProcedureType<KmeansProcedure, KmeansConfig>
regKmeans(builtinPackage(),
          "Simple clustering algorithm based on cluster centroids in embedding space",
          "procedures/KmeansProcedure.md.html");

RegisterFunctionType<KmeansFunction, KmeansFunctionConfig>
regKmeansFunction(builtinPackage(),
                  "kmeans",
                  "Apply a k-means clustering to new data",
                  "functions/Kmeans.md.html");

} // file scope

} // namespace MLDB


// ====================================================================
// This file is part of FlexibleSUSY.
//
// FlexibleSUSY is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published
// by the Free Software Foundation, either version 3 of the License,
// or (at your option) any later version.
//
// FlexibleSUSY 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
// General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with FlexibleSUSY.  If not, see
// <http://www.gnu.org/licenses/>.
// ====================================================================

#ifndef SUM_H
#define SUM_H

#include <type_traits>
#include <cstddef>
#include <Eigen/Core>

namespace flexiblesusy {

#define SUM(...) (get_sum(__VA_ARGS__)(__VA_ARGS__))

#define get_sum(...) get_sum_macro(__VA_ARGS__, sum_user_t, sum_ptrdiff_t,)

#define get_sum_macro(_1, _2, _3, _4, _5, name, ...) name

#define sum_ptrdiff_t(idx, ini, fin, expr)      \
    sum<std::ptrdiff_t>((ini), (fin), [&](std::ptrdiff_t (idx)) { return (expr); })

#define sum_user_t(type, idx, ini, fin, expr)	\
    sum<type>((ini), (fin), [&](type (idx)) { return (expr); })

template<typename T>
struct is_eigen_type
{
    static constexpr auto value =
	std::is_base_of<Eigen::EigenBase<T>, T>::value;
};

template<typename Idx, typename Function, bool isEigenType>
struct EvalEigenXprImpl {
    static auto eval(Idx i, Function f) -> decltype(f(i)) {
	return f(i);
    }
};

template<typename Idx, typename Function>
struct EvalEigenXprImpl<Idx, Function, true> {
    static auto eval(Idx i, Function f) ->
	typename std::remove_reference<decltype(f(i).eval())>::type
    {
	return f(i).eval();
    }
};

template<typename Idx, typename Function>
auto EvalEigenXpr(Idx i, Function f) ->
    decltype(
	EvalEigenXprImpl<Idx, Function, is_eigen_type<decltype(f(i))>::value>::
	eval(i, f))
{
    return
	EvalEigenXprImpl<Idx, Function, is_eigen_type<decltype(f(i))>::value>::
	eval(i, f);
}

template<typename T, bool isEigenType>
struct create_zero {
    static const T zero() {
	return T();
    }
};

template<typename T>
struct create_zero<T, true> {
    static const T zero() {
	T z;
	z.setZero();
	return z;
    }
};

template<class Idx, class Function>
auto sum(Idx ini, Idx fin, Function f) -> decltype(EvalEigenXpr<Idx>(ini, f))
{
    using Evaled = decltype(EvalEigenXpr<Idx>(ini, f));
    using Acc = typename std::remove_cv<Evaled>::type;
    Acc s = create_zero<Acc, is_eigen_type<Evaled>::value>::zero();
    for (Idx i = ini; i <= fin; i++) s += f(i);
    return s;
}

} // namespace flexiblesusy

#endif // sum_hpp


/*
 * This file is part of the M. xanthus polarity simulator
 * Copyright (c) 2020 Filipe Tostevin
 *
 * This file contains the main routine that creates all simulator
 * objects and calls the propagation routines.
 */

#include <iostream>
#include <cstdlib>
#include <boost/property_tree/ini_parser.hpp>
#include <boost/program_options.hpp>
#include "state.h"
#include "propagator.h"

using namespace std;
namespace po = boost::program_options;
namespace pt = boost::property_tree;

int main(
	int argc,
	char* argv[]
){
	// Command line options and help
	po::options_description desc("Allowed options");
	desc.add_options()
		("help,h", "print help message")
		("conf,c", po::value<string>(), "parameter configuration file")
		("mutants,m", "simulate deletion mutants")
	;
	po::variables_map vm;
	po::store(po::parse_command_line(argc, argv, desc), vm);
	po::notify(vm);

	if(vm.count("help")) {
		cout << desc << "\n";
		return(EXIT_SUCCESS);
	}

	// Read the input file given by the "--conf" option, or otherwise from
	// "params.in"
	pt::ptree ptree;
	string defaultconfig="params.in";
	if(vm.count("conf")) {
		cout << "Reading parameters from: " << vm["conf"].as<string>() << endl;
		pt::read_ini(vm["conf"].as<string>(), ptree);
	} else {
		cout << "No parameter file specified. Reading parameters from default file: " << defaultconfig << endl;
		pt::read_ini(defaultconfig, ptree);
	}

	// Get total simulation time from input file
	const double t_total=ptree.get<double>("t_total");
	const double sampling_interval=ptree.get<double>("sampling_interval");

	// Define the system state variable and initialize all elements to zero
	state conc(NSPEC,NPOOL,0);

	// Set up simulation infrastructure
	propagator P;
	// Read in parameter values from input file
	for(auto p: ptree.get_child("model_parameters")) {
		P.params.emplace(p.first, ptree.get<double>("model_parameters."+p.first));
	}
	//Run a test call of the propagator to check that all required parameters
	//exist
	try{ P(conc, conc, 0); }
	catch(...) {
		cout << "Error in model definition, check that all parameters are specified in the input file." << endl;
		exit(EXIT_FAILURE);
	}

	// Deletion mutants are defined by an integer that when represented as a
	// binary string defines the presence/absence of each protein,
	// i.e. if the number of proteins in NSPEC=3, then
	// mutant_number=1==b100 indicates only the first-listed protein is present,
	// mutant_number=7==b111 indicates that all proteins are present
	// Here we define the first simulation to run depending on whether we are
	// simulating all mutants, or just the wild-type
	int mutant_number;
	if(vm.count("mutants")) {
		cout << "Simulating deletion mutants" << endl;
		mutant_number=1;
	} else {
		cout << "Simulating only wild-type" << endl;
		mutant_number=(1<<NSPEC) -1;
	}

	// Loop over all mutants until we reach wild-type
	while(mutant_number < 1<<NSPEC) {
		cout << mutant_number << endl;
		//Each mutant will be output to a separate data file.
		//Decode mutant number bitwise into whether protein is present or not and
		//reconstruct readable output string from it, to use in the name of the
		//output file
		string mutant_string="";
		for(int s=0; s<NSPEC; ++s) {
			if( !( (mutant_number >> s) & 1) ) {
				mutant_string+=".d" + species_names[s];
			}
		}
		ofstream datafile("data.txt"+mutant_string);

		//Headings for output file
		datafile << "t ";
		for(int s=0; s<NSPEC; ++s) {
			datafile << species_names[s] << "1 "
			         << species_names[s] << "2 "
			         << species_names[s] << "t ";
		} datafile << endl;

		for(int s=0; s<NSPEC; ++s) {
			// Initialize total amounts of each protein to 1 or 0 according to the
			// deletion condition
			if( (mutant_number >> s) & 1 ) {
				conc(s,TOTAL)=1;
			} else {
				conc(s,TOTAL)=0;
			}
			// Initialize polar fractions
			conc(s,POLE1)=ptree.get<double>("initial_fractions_"+species_names[s]+".pole1")*conc(s,TOTAL);
			conc(s,POLE2)=ptree.get<double>("initial_fractions_"+species_names[s]+".pole2")*conc(s,TOTAL);
		}

		// Start at t=0
		double t=0;
		write_state(conc, t, datafile);

		// Run until the end of the simulation time t_total
		while(t<t_total) {
			// Call the ODE integrator to run until the next output time point
			boost::numeric::odeint::integrate(boost::ref(P), conc, t, t+sampling_interval, 0.2*sampling_interval);
			t+=sampling_interval;

			write_state(conc, t, datafile);
		}

		datafile.close();
		++mutant_number;
	}

	return EXIT_SUCCESS;
}