#include <vector>
#define GL_SILENCE_DEPRECATION
#include <GLFW/glfw3.h>

#include "gltf_util.h"
#include <Eigen/Geometry>

#include "delfem2/glfw/viewer3.h"
#include "delfem2/glfw/util.h"
#include "delfem2/opengl/old/funcs.h"
#include "delfem2/opengl/old/mshuni.h"

class ChannelData {
 public:
  unsigned int ibone;
  std::vector<double> v_input, v_output;
  std::string target_path;
};

/**
 *@brief articulated rigid body for character rigging
 */
class RigBone {
 public:
  RigBone()
      : inverseBindMatrix{}, transformGlobal{} {
    parent_bone_idx = -1;
    transformRelative.setIdentity();
    inverseBindMatrix.setIdentity();
    transformGlobal.setIdentity();
  }

 public:
  std::string name; // name of this bone

  /**
   * @brief Inverse of the binding matrix.
   * @details The binding matrix transform the bone to its reference config.
   * With binding matrix, the origin is mapped to the joint position of this bone in the reference config
   */
  Eigen::Matrix4d inverseBindMatrix;

  /**
   * @brief index of parent bone
   */
  unsigned int parent_bone_idx;

  /**
   * @brief transformation of current bone relative to parent bone
   */
  Eigen::Matrix4d transformRelative;

  /**
   * @brief Affine matrix to transform the bone to the current config
   */
  Eigen::Matrix4d transformGlobal;
};

void LoadBoneFromGltfFile(
    std::vector<RigBone> &bones,
    const tinygltf::Model &model,
    unsigned int inode0,
    int ibone_p,
    const std::vector<unsigned int> &mapNode2Bone) {
  assert(inode0 < model.nodes.size());
  const tinygltf::Node &node = model.nodes[inode0];
  const unsigned int ibone0 = mapNode2Bone[inode0];
  if (ibone0 == UINT_MAX) { return; }
  assert(ibone0 < bones.size());
  bones[ibone0].parent_bone_idx = ibone_p;
  bones[ibone0].name = node.name;
  for (int inode_ch: node.children) {
    LoadBoneFromGltfFile(
        bones,
        model, inode_ch, ibone0, mapNode2Bone);
  }
}

bool LoadGltf(
    std::vector<double> &vtx_xyz_ini,
    std::vector<unsigned int> &tri_vtz,
    std::vector<double> &skinning_sparse_weight,
    std::vector<unsigned int> &skinning_sparse_index,
    std::vector<RigBone> &bones,
    std::vector<ChannelData> &channels,
    unsigned int iskin,
    int imesh,
    int iprimitive,
    int ianimation,
    const std::string &fpath) {
  std::string err;
  std::string warn;
  tinygltf::TinyGLTF loader;
  std::unique_ptr<tinygltf::Model> pModel = std::make_unique<tinygltf::Model>();
  bool ret = loader.LoadASCIIFromFile(
      pModel.get(), &err, &warn,
      fpath);
  if (!warn.empty()) { printf("Warn: %s\n", warn.c_str()); }
  if (!err.empty()) { printf("Err: %s\n", err.c_str()); }
  if (!ret) {
    printf("Failed to parse glTF\n");
    return false;
  }
  tinygltf::Model &model = *(pModel);
  //
  LoadMeshInfo(
      vtx_xyz_ini, tri_vtz,
      skinning_sparse_weight, skinning_sparse_index,
      model, imesh, iprimitive);
  //
  assert(iskin < model.skins.size());
  bones.resize(model.skins[iskin].joints.size());
  unsigned int inode_root = model.skins[iskin].skeleton;
  if (inode_root == UINT_MAX && !model.skins[iskin].joints.empty()) {
    inode_root = model.skins[iskin].joints[0];
  }
  assert(inode_root < model.nodes.size());
  //
  std::vector<unsigned int> mapNode2Bone(model.nodes.size(), UINT_MAX);
  for (unsigned int ij = 0; ij < model.skins[iskin].joints.size(); ++ij) {
    const unsigned int inode = model.skins[iskin].joints[ij];
    assert(inode < model.nodes.size());
    mapNode2Bone[inode] = ij;
  }
  LoadBoneFromGltfFile(
      bones,
      model, inode_root, -1, mapNode2Bone);
  { // set bone invBindMat
    const tinygltf::Skin &skin = model.skins[iskin];
    std::vector<double> M;
    GetArray_Double(M, model, skin.inverseBindMatrices);
    assert(M.size() == bones.size() * 16);
    for (size_t ij = 0; ij < M.size() / 16; ++ij) {
      for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
          bones[ij].inverseBindMatrix(i, j) = M[ij * 16 + j * 4 + i];
        }
      }
      bones[ij].transformGlobal = bones[ij].inverseBindMatrix.inverse().eval();
    }
  }
  { // load bone animation
    std::vector<tinygltf::AnimationSampler> &sampler = pModel->animations[ianimation].samplers;
    std::vector<tinygltf::AnimationChannel> &channel = pModel->animations[ianimation].channels;
    const unsigned int nch = channel.size();
    channels.resize(nch);
    for (unsigned int ich = 0; ich < channel.size(); ++ich) {
      channels[ich].target_path = channel[ich].target_path;
      const int isampler = channel[ich].sampler;
      const int inode = channel[ich].target_node;
      channels[ich].ibone = mapNode2Bone[inode];
      assert(channels[ich].ibone < bones.size());
      GetArray_Double(channels[ich].v_input, model, sampler[isampler].input);
      GetArray_Double(channels[ich].v_output, model, sampler[isampler].output);
    }
  }
  return true;
}

/**
 * set relative transformation to each bone
 */
void SetRelativeTransformationToBones(
    std::vector<RigBone> &bones,
    std::vector<ChannelData> &channels,
    double time) {
  for (auto &bone: bones) {
    bone.transformRelative.setIdentity();
  }
  for (auto &ch: channels) {
    unsigned int idx1 = 0;
    for (; idx1 < ch.v_input.size(); ++idx1) {
      if (ch.v_input[idx1] > time) { break; }
    }
    idx1 = (idx1 == ch.v_input.size()) ? idx1 - 1 : idx1;
    unsigned int idx0 = (idx1 == 0) ? 0 : idx1 - 1;
    double ratio = (idx1 == idx0) ? 0 : (time - ch.v_input[idx0]) / (ch.v_input[idx1] - ch.v_input[idx0]);
    if (ch.target_path == "translation") {  // translation
      const Eigen::Translation3d transl(
          (1 - ratio) * ch.v_output[idx0 * 3 + 0] + ratio * ch.v_output[idx1 * 3 + 0],
          (1 - ratio) * ch.v_output[idx0 * 3 + 1] + ratio * ch.v_output[idx1 * 3 + 1],
          (1 - ratio) * ch.v_output[idx0 * 3 + 2] + ratio * ch.v_output[idx1 * 3 + 2]);
      bones[ch.ibone].transformRelative *= Eigen::Affine3d(transl).matrix();
    }
    if (ch.target_path == "rotation") {  // rotation in quaternion
      const Eigen::Quaterniond quat(
          (1 - ratio) * ch.v_output[idx0 * 4 + 3] + ratio * ch.v_output[idx1 * 4 + 3],
          (1 - ratio) * ch.v_output[idx0 * 4 + 0] + ratio * ch.v_output[idx1 * 4 + 0],
          (1 - ratio) * ch.v_output[idx0 * 4 + 1] + ratio * ch.v_output[idx1 * 4 + 1],
          (1 - ratio) * ch.v_output[idx0 * 4 + 2] + ratio * ch.v_output[idx1 * 4 + 2]);
      bones[ch.ibone].transformRelative *= Eigen::Affine3d(quat).matrix();
    }
  }
}

void LinearBlendSkinning(
    std::vector<RigBone> &bones,
    std::vector<double> &vtx_xyz,
    const std::vector<double> &vtx_xyz_ini,
    const std::vector<double> &skinning_sparse_weight,
    const std::vector<unsigned int> &skinning_sparse_index) {
  // set global transformation using the relative transformation for each bone
  for (std::size_t ibone = 0; ibone < bones.size(); ++ibone) {
    const unsigned int ibone_p = bones[ibone].parent_bone_idx;
    Eigen::Matrix4d transformParent = Eigen::Matrix4d::Identity();
    if (ibone_p != UINT_MAX) { // this is not root bone
      transformParent = bones[ibone_p].transformGlobal;  // set parent bone transformation
    }
    // bones[ibone].transformGlobal = /* write single line code here */
  }
  // deform vertex positions of the triangle mesh using linear blend skinning
  const size_t nvtx = vtx_xyz.size() / 3;
  for (unsigned int ip = 0; ip < nvtx; ++ip) { // loop over vertex
    Eigen::Vector4d pos0( // initial position (red)
        vtx_xyz_ini[ip * 3 + 0],
        vtx_xyz_ini[ip * 3 + 1],
        vtx_xyz_ini[ip * 3 + 2], 1.0);
    Eigen::Vector4d pos1(0, 0, 0, 1);  // deformed position (blue)
    double sum_w = 0.0;
    for (int iibone = 0; iibone < 4; ++iibone) {  // each vertex has weights of 4 bones
      double w = skinning_sparse_weight[ip * 4 + iibone];  // rigging weight
      if (w < 1.0e-30) { continue; }
      unsigned int ibone = skinning_sparse_index[ip * 4 + iibone];  // bone index
      assert (ibone < bones.size());
      Eigen::Matrix4d inverseBindMatrix = bones[ibone].inverseBindMatrix;
      Eigen::Matrix4d transformGlobal = bones[ibone].transformGlobal;
      sum_w += w;
      // pos1 = /* write single line code here */
    }
    assert(fabs(sum_w) > 1.0e-10);
    pos1 /= sum_w;
    vtx_xyz[ip * 3 + 0] = pos1[0];
    vtx_xyz[ip * 3 + 1] = pos1[1];
    vtx_xyz[ip * 3 + 2] = pos1[2];
  }
}

// ---------------------------------

int main() {
  std::vector<double> vtx_xyz;  // deformed vertex positions (blue)
  std::vector<unsigned int> tri_vtx;  // triangle index
  std::vector<double> skinning_sparse_weight;  // skinning weights for vertex
  std::vector<unsigned int> skinning_sparse_index;  // bones for vertex
  std::vector<RigBone> bones;
  std::vector<ChannelData> channels;
  LoadGltf(vtx_xyz, tri_vtx,
           skinning_sparse_weight, skinning_sparse_index,
           bones, channels,
           0, 0, 0, 0,
           std::string(SOURCE_DIR) + "/../assets/CesiumMan.gltf");
  const std::vector<double> vtx_xyz_ini = vtx_xyz;  // initial vertex position (red)

  // --------------
  // opengl starts here
  delfem2::glfw::CViewer3 viewer(2);
  viewer.window_title = "task06";
  viewer.view_rotation = std::make_unique<delfem2::ModelView_Ztop>();
  if (!glfwInit()) { exit(EXIT_FAILURE); }
  glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 2);
  glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 1);
  viewer.OpenWindow();
  while (!glfwWindowShouldClose(viewer.window)) {
    double time = glfwGetTime();
    time = time - 2. * int(time / 2.);  // loop with 2 sec interval
    // for each bone set relative rotation from parent bone
    SetRelativeTransformationToBones(
        bones, channels,
        time);
    // linear blend skinning
    LinearBlendSkinning(
        bones, vtx_xyz,
        vtx_xyz_ini,
        skinning_sparse_weight, skinning_sparse_index);
    // --------------------
    viewer.DrawBegin_oldGL();
    ::glDisable(GL_LIGHTING);
    delfem2::opengl::DrawAxis(1);
    // draw reference
    glColor3d(1, 0, 0);
    delfem2::opengl::DrawMeshTri3D_Edge(
        vtx_xyz_ini.data(), vtx_xyz_ini.size() / 3,
        tri_vtx.data(), tri_vtx.size() / 3);
    glColor3d(0, 0, 1);
    delfem2::opengl::DrawMeshTri3D_Edge(
        vtx_xyz.data(), vtx_xyz.size() / 3,
        tri_vtx.data(), tri_vtx.size() / 3);
    viewer.SwapBuffers();
    glfwPollEvents();
  }
  glfwDestroyWindow(viewer.window);
  glfwTerminate();
  exit(EXIT_SUCCESS);
}


#include "boost_test_tools_extensions.h"

#include <boost/test/unit_test.hpp>
#include <k52/common/constants.h>

void CheckComplexEqual(std::complex< double > a, std::complex< double > b)
{
    CheckDoubleEqual(a.real(), b.real());
    CheckDoubleEqual(a.imag(), b.imag());
}

void CheckDoubleEqual(double a, double b)
{
    BOOST_CHECK_SMALL(a - b, ::k52::common::Constants::Eps);
}


#include "mesh/mechanical/beam/fem_submesh.hpp"

#include "mesh/dof_allocator.hpp"
#include "interpolations/interpolation_factory.hpp"
#include "math/transform_expand.hpp"

#include <tbb/parallel_for.h>
#include <Eigen/Geometry>

namespace neon::mechanical::beam
{
fem_submesh::fem_submesh(json const& material_data,
                         json const& simulation_data,
                         std::shared_ptr<material_coordinates>& coordinates,
                         basic_submesh const& submesh)
    : basic_submesh(submesh),
      sf(make_line_interpolation(topology(), simulation_data)),
      coordinates(coordinates),
      profiles(elements() * sf->quadrature().points()),
      view(sf->quadrature().points()),
      variables(std::make_shared<internal_variable_type>(elements() * sf->quadrature().points())),
      cm(std::make_unique<isotropic_linear>(variables, material_data))
{
    dof_allocator(node_indices, dof_indices, traits::dof_order);

    orientations.resize(elements(), vector3::UnitX());
    tangents.resize(elements(), vector3::UnitZ());

    variables->add(variable::second::cauchy_stress,
                   variable::scalar::shear_area_1,
                   variable::scalar::shear_area_2,
                   variable::scalar::cross_sectional_area,
                   variable::scalar::second_moment_area_1,
                   variable::scalar::second_moment_area_2);
}

void fem_submesh::update_internal_variables(double const time_step_size)
{
    for (auto& profile : profiles)
    {
        profile = std::make_unique<geometry::rectangular_bar>(1.0, 1.0);
    }

    auto& A = variables->get(variable::scalar::cross_sectional_area);
    auto& As1 = variables->get(variable::scalar::shear_area_1);
    auto& As2 = variables->get(variable::scalar::shear_area_2);

    auto& area_moment_1 = variables->get(variable::scalar::second_moment_area_1);
    auto& area_moment_2 = variables->get(variable::scalar::second_moment_area_2);

    // Loop over all elements and quadrature points to compute the profile
    // properties for each quadrature point in the beam
    tbb::parallel_for(std::int64_t{0}, elements(), [&, this](auto const element) {
        sf->quadrature().for_each([&, this](auto const& femval, auto const l) {
            [[maybe_unused]] auto const& [N, dN] = femval;

            A.at(view(element, l)) = profiles.at(element)->area();

            auto const [shear_area_1, shear_area_2] = profiles.at(element)->shear_area();
            As1.at(view(element, l)) = shear_area_1;
            As2.at(view(element, l)) = shear_area_2;

            auto const [I1, I2] = profiles.at(element)->second_moment_area();
            area_moment_1.at(view(element, l)) = I1;
            area_moment_2.at(view(element, l)) = I2;
        });
    });

    cm->update_internal_variables();
}

std::pair<index_view, matrix const&> fem_submesh::tangent_stiffness(std::int32_t const element) const
{
    static thread_local matrix ke(12, 12);

    auto const& configuration = coordinates->initial_configuration(local_node_view(element));

    matrix12 const T = rotation_matrix(element);

    ke = T
         * (bending_stiffness(configuration, element) + shear_stiffness(configuration, element)
            + axial_stiffness(configuration, element) + torsional_stiffness(configuration, element))
         * T.transpose();

    return {local_dof_view(element), ke};
}

matrix const& fem_submesh::bending_stiffness(matrix const& configuration,
                                             std::int32_t const element) const
{
    static thread_local matrix2x B_b(2, 6 * sf->nodes());
    static thread_local matrix k_b(6 * sf->nodes(), 6 * sf->nodes());

    B_b.setZero();
    k_b.setZero();

    auto const& D_b = variables->get(variable::second::bending_stiffness);

    sf->quadrature().integrate_inplace(k_b, [&, this](auto const& femval, auto const l) {
        auto const& [N, dN] = femval;

        double const jacobian = configuration.row(0) * dN;

        for (int i = 0; i < sf->nodes(); ++i)
        {
            auto const offset = i * 6;

            B_b(0, 3 + offset) = B_b(1, 4 + offset) = dN(i, l) / jacobian;
        }

        return B_b.transpose() * D_b.at(view(element, l)) * B_b * jacobian;
    });
    return k_b;
}

matrix const& fem_submesh::shear_stiffness(matrix const& configuration, std::int32_t const element) const
{
    static thread_local matrix2x B_s(2, 6 * sf->nodes());
    static thread_local matrix k_s(6 * sf->nodes(), 6 * sf->nodes());

    B_s.setZero();
    k_s.setZero();

    auto const& D_s = variables->get(variable::second::shear_stiffness);

    sf->quadrature().integrate_inplace(k_s, [&, this](auto const& femval, auto const l) {
        auto const& [N, dN] = femval;

        double const jacobian = configuration.row(0) * dN;

        for (int i = 0; i < sf->nodes(); ++i)
        {
            auto const offset = i * 6;

            B_s(0, 0 + offset) = B_s(1, 1 + offset) = dN(i, l) / jacobian;

            B_s(0, 4 + offset) = -N(i, l);
            B_s(1, 3 + offset) = N(i, l);
        }

        return B_s.transpose() * D_s.at(view(element, l)) * B_s * jacobian;
    });
    return k_s;
}

matrix const& fem_submesh::axial_stiffness(matrix const& configuration, std::int32_t const element) const
{
    static thread_local vector B_a(6 * sf->nodes());
    static thread_local matrix k_a(6 * sf->nodes(), 6 * sf->nodes());

    B_a.setZero();
    k_a.setZero();

    auto const D_a = variables->get(variable::scalar::axial_stiffness);

    sf->quadrature().integrate_inplace(k_a, [&, this](auto const& femval, auto const l) {
        auto const& [N, dN] = femval;

        double const jacobian = configuration.row(0) * dN;

        for (int i = 0; i < sf->nodes(); ++i)
        {
            auto const offset = i * 6;

            B_a(2 + offset) = dN(i, l) / jacobian;
        }

        return B_a * D_a.at(view(element, l)) * B_a.transpose() * jacobian;
    });
    return k_a;
}

matrix const& fem_submesh::torsional_stiffness(matrix const& configuration,
                                               std::int32_t const element) const
{
    static thread_local vector B_t(6 * sf->nodes());
    static thread_local matrix k_t(6 * sf->nodes(), 6 * sf->nodes());

    B_t.setZero();
    k_t.setZero();

    auto const D_t = variables->get(variable::scalar::torsional_stiffness);

    sf->quadrature().integrate_inplace(k_t, [&, this](auto const& femval, auto const l) {
        auto const& [N, dN] = femval;

        double const jacobian = configuration.row(0) * dN;

        for (int i = 0; i < sf->nodes(); ++i)
        {
            auto const offset = i * 6;

            B_t(5 + offset) = dN(i, l) / jacobian;
        }

        return B_t * D_t.at(view(element, l)) * B_t.transpose() * jacobian;
    });
    return k_t;
}

matrix12 fem_submesh::rotation_matrix(std::int32_t const element) const
{
    matrix3 rotation;

    // Local coordinate (x1)
    rotation.col(0) = orientations.at(element);

    // Local coordinate (x3)
    rotation.col(2) = tangents.at(element);

    // Local coordinate (x2)
    rotation.col(1) = rotation.col(0).cross(rotation.col(2)).normalized().cwiseAbs();

    matrix12 element_rotation = matrix12::Zero();

    element_rotation.block<3, 3>(0, 0) = rotation;
    element_rotation.block<3, 3>(3, 3) = rotation;
    element_rotation.block<3, 3>(6, 6) = rotation;
    element_rotation.block<3, 3>(9, 9) = rotation;

    return element_rotation;
}
}


#include <iostream>
#include <boost/astronomy/coordinate/coordinate.hpp>
#include <boost/astronomy/coordinate/sky_point.hpp>
#include <boost/astronomy/coordinate/frame.hpp>
#include <boost/units/systems/si/length.hpp>
#include <boost/units/systems/si/time.hpp>
#include <boost/units/systems/si/velocity.hpp>
#include <boost/units/systems/si/plane_angle.hpp>
#include <boost/units/systems/angle/degrees.hpp>
#include <boost/units/io.hpp>

using namespace boost::astronomy::coordinate;
using namespace boost::units;
using namespace boost::units::si;
namespace bud = boost::units::degree;

typedef spherical_representation<double, quantity<degree::plane_angle>, quantity<degree::plane_angle>,
                                 quantity<si::length>>
    spherical_rep;
typedef spherical_coslat_differential<double, quantity<degree::plane_angle>,
                                      quantity<degree::plane_angle>, quantity<si::velocity>>
    spherical_coslat_diff;

int main()
{
    //first we create a point we want to represent in the sky
    auto star = make_spherical_representation(45.0 * bud::degrees, 45.0 * bud::degrees, 3.0 * meters);

    //we define motion it has (if it has)
    auto motion = make_spherical_coslat_differential(10.0 * bud::degrees, 15.0 * bud::degrees, 1.0 * meters / seconds);

    //we create a frame
    icrs<spherical_rep, spherical_coslat_diff> frame(star, motion);

    //creating a sky point object
    sky_point<icrs<spherical_rep, spherical_coslat_diff>> object(frame);

    return 0;
}


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

  @copyright 2015 NumScale SAS
  @copyright 2015 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_FUNCTION_SCALAR_EXPM1_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_FUNCTION_SCALAR_EXPM1_HPP_INCLUDED
#include <boost/simd/function/std.hpp>

#ifndef BOOST_SIMD_NO_INVALIDS
#include <boost/simd/constant/nan.hpp>
#include <boost/simd/function/scalar/is_nan.hpp>
#endif
#include <boost/simd/arch/common/detail/generic/expm1_kernel.hpp>
#include <boost/simd/constant/inf.hpp>
#include <boost/simd/constant/logeps.hpp>
#include <boost/simd/constant/maxlog.hpp>
#include <boost/simd/constant/mone.hpp>
#include <boost/simd/function/scalar/is_greater.hpp>
#include <boost/simd/function/scalar/is_less.hpp>
#include <boost/dispatch/function/overload.hpp>
#include <boost/config.hpp>
#include <cmath>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;
  BOOST_DISPATCH_OVERLOAD ( expm1_
                          , (typename A0)
                          , bd::cpu_
                          , bd::generic_< bd::floating_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() (A0 a0) const BOOST_NOEXCEPT
    {
    #ifndef BOOST_SIMD_NO_INVALIDS
      if(is_nan(a0)) return Nan<A0>();
    #endif
      if((a0 < Logeps<A0>())) return Mone<A0>();
      if((a0 > Maxlog<A0>())) return  Inf<A0>();
      return detail::expm1_kernel<A0>::expm1(a0);
    }
  };

  BOOST_DISPATCH_OVERLOAD ( expm1_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::floating_<A0> >
                          , bs::std_tag
                          )
  {
    BOOST_FORCEINLINE A0 operator() (A0 a0, std_tag const&) const BOOST_NOEXCEPT
    {
      return std::expm1(a0);
    }
  };
} } }

#endif


#include "cnn/nodes.h"
#include "cnn/cnn.h"
#include "cnn/training.h"
#include "cnn/timing.h"
#include "cnn/rnn.h"
#include "cnn/gru.h"
#include "cnn/lstm.h"
#include "cnn/dict.h"
#include "cnn/expr.h"

#include <algorithm>
#include <iostream>
#include <fstream>
#include <sstream>
#include <cstdlib>
#include <set>

#include <boost/archive/text_iarchive.hpp>
#include <boost/archive/text_oarchive.hpp>

#include <boost/algorithm/string.hpp>

using namespace std;
using namespace cnn;

unsigned LAYERS = 1;
unsigned INPUT_DIM = 64;
unsigned XCRIBE_DIM = 64;
unsigned SEG_DIM = 48;
unsigned H1DIM = 64;
unsigned H2DIM = 48;
unsigned TAG_DIM = 16;
unsigned TAG_SIZE = 0;
unsigned VOCAB_SIZE = 0;
unsigned DURATION_DIM = 4;

bool eval = false;
cnn::Dict d;
cnn::Dict td;
int kNONE;
int kSOS;
int kEOS;

bool use_pretrained_embeding = true;

// returns embeddings of labels
struct SymbolEmbedding {
  SymbolEmbedding(Model& m, unsigned n, unsigned dim) {
    p_labels = m.add_lookup_parameters(n, {dim});
  }
  void new_graph(ComputationGraph& g) { cg = &g; }
  Expression embed(unsigned label_id) {
    return lookup(*cg, p_labels, label_id);
  }
  ComputationGraph* cg;
  LookupParameters* p_labels;
};

struct DurationEmbedding {
  virtual ~DurationEmbedding();
  virtual void new_graph(ComputationGraph& g) = 0;
  virtual Expression embed(unsigned dur) = 0;
};

DurationEmbedding::~DurationEmbedding() {}

struct MLPDurationEmbedding : public DurationEmbedding {
  MLPDurationEmbedding(Model& m, unsigned hidden, unsigned dim) {
    p_zero = m.add_parameters({dim});
    p_d2h = m.add_parameters({hidden, 2});
    p_hb = m.add_parameters({hidden});
    p_h2o = m.add_parameters({dim, hidden});
    p_ob = m.add_parameters({dim});
    dur_xs.resize(10000, vector<float>(2));
    for (unsigned i = 1; i < dur_xs.size(); ++i) {
      dur_xs[i][0] = i;
      dur_xs[i][1] = float(i) / logf(2);
    }
  }
  void new_graph(ComputationGraph& g) override {
    zero = parameter(g, p_zero);
    d2h = parameter(g, p_d2h);
    hb = parameter(g, p_hb);
    h2o = parameter(g, p_h2o);
    ob = parameter(g, p_ob);
    cg = &g;
  }
  Expression embed(unsigned dur) override {
    if (dur) {
      Expression x = input(*cg, {2}, &dur_xs[dur]);
      Expression h = rectify(affine_transform({hb, d2h, x}));
      return affine_transform({ob, h2o, h});
    } else {
      return zero;
    }
  }

  ComputationGraph* cg;
  vector<vector<float>> dur_xs;
  Expression zero, d2h, hb, h2o, ob;
  Parameters* p_zero;
  Parameters* p_d2h;
  Parameters* p_hb;
  Parameters* p_h2o;
  Parameters* p_ob;
};

struct BinnedDurationEmbedding : public DurationEmbedding {
  BinnedDurationEmbedding(Model& m, unsigned dim, unsigned num_bins = 8) : max_bin(num_bins - 1) {
    assert(num_bins > 0);
    p_e = m.add_lookup_parameters(num_bins, {dim});
  }
  void new_graph(ComputationGraph& g) override {
    cg = &g;
  }
  Expression embed(unsigned dur) override {
    if (dur) dur = static_cast<unsigned>(log(dur) / log(1.6)) + 1;
    if (dur > max_bin) { dur = max_bin; }
    return lookup(*cg, p_e, dur);
  }
  ComputationGraph* cg;
  unsigned max_bin;
  LookupParameters* p_e;
};

template <class Builder>
struct BiTrans {
  Builder l2rbuilder;
  Builder r2lbuilder;
  Parameters* p_f2c;
  Parameters* p_r2c;
  Parameters* p_cb;

  explicit BiTrans(Model& model) :
      l2rbuilder(LAYERS, INPUT_DIM, XCRIBE_DIM, &model),
      r2lbuilder(LAYERS, INPUT_DIM, XCRIBE_DIM, &model) {
    p_f2c = model.add_parameters({XCRIBE_DIM, INPUT_DIM});
    p_r2c = model.add_parameters({XCRIBE_DIM, INPUT_DIM});
    p_cb = model.add_parameters({XCRIBE_DIM});
  }

  vector<Expression> transcribe(ComputationGraph& cg, const vector<Expression>& x) {
    l2rbuilder.new_graph(cg);
    l2rbuilder.start_new_sequence();
    r2lbuilder.new_graph(cg);
    r2lbuilder.start_new_sequence();
    Expression f2c = parameter(cg, p_f2c);
    Expression r2c = parameter(cg, p_r2c);
    Expression cb = parameter(cg, p_cb);

    const int len = x.size();
    vector<Expression> fwd(len), rev(len), res(len);
    for (int i = 0; i < len; ++i)
      fwd[i] = l2rbuilder.add_input(x[i]);
    for (int i = len - 1; i >= 0; --i)
      rev[i] = r2lbuilder.add_input(x[i]);
    for (int i = 0; i < len; ++i)
      res[i] = affine_transform({cb, f2c, fwd[i], r2c, rev[i]});
    return res;
  }
};

// uni-directional segment embeddings
// first call construct_chart(sequence), then access the embeddings of the spans with operator()(i,j)
template <class Builder>
struct SegEmbedUni {
  Parameters* p_h0;
  int len;
  Builder builder;
  vector<vector<Expression>> h;  // h[i][length of segment - 1]
  explicit SegEmbedUni(Model& m) :
      builder(LAYERS, XCRIBE_DIM, SEG_DIM, &m) {
    p_h0 = m.add_parameters({XCRIBE_DIM});
  }
  void construct_chart(ComputationGraph& cg, const vector<Expression>& c, int max_seg_len = 0) {
    len = c.size();
    h.clear();
    h.resize(len);
    Expression h0 = parameter(cg, p_h0);
    builder.new_graph(cg);
    for (int i = 0; i < len; ++i) {
      int max_j = i + len;
      if (max_seg_len) max_j = i + max_seg_len;
      if (max_j > len) max_j = len;
      int seg_len = max_j - i;
      auto& hi = h[i];
      hi.resize(seg_len);

      builder.start_new_sequence();
      builder.add_input(h0);
      for (int k = 0; k < seg_len; ++k)
        hi[k] = builder.add_input(c[i+k]);
    }
  }
  // returns the embedding of segment [i,j]
  const Expression& operator()(int i, int j) const {
    assert(j < len);
    assert(j >= i);
    return h[i][j - i];
  }
};

// uni-directional segment embeddings
// first call construct_chart(sequence), then access the embeddings of the spans with operator()(i,j)
template <class Builder>
struct SegEmbedBi {
  int len;
  vector<vector<pair<Expression, Expression>>> h;
  SegEmbedUni<Builder> fwd, rev;
  explicit SegEmbedBi(Model& m) : fwd(m), rev(m) {}
  void construct_chart(ComputationGraph& cg, const vector<Expression>& c, int max_seg_len = 0) {
    len = c.size();
    fwd.construct_chart(cg, c, max_seg_len);
    vector<Expression> rc(len);
    for (int i = 0; i < len; ++i) rc[i] = c[len - 1 - i];
    rev.construct_chart(cg, rc, max_seg_len);
    h.clear();
    h.resize(len);
    for (int i = 0; i < len; ++i) {
      int max_j = i + len;
      if (max_seg_len) max_j = i + max_seg_len;
      if (max_j > len) max_j = len;
      int seg_len = max_j - i;
      auto& hi = h[i];
      hi.resize(seg_len);
      for (int k = 0; k < seg_len; ++k) {
        int j = i + k;
        const Expression& fe = fwd(i, j);
        const Expression& re = rev(len-1-j, len-1-i);
        hi[k] = make_pair(fe,re);
      }
    }
  }
  const pair<Expression, Expression>& operator()(int i, int j) const {
    assert(j >= i);
    assert(j < len);
    return h[i][j-i];
  }
};

template <class Builder>
struct ZerothOrderSemiCRF {
  SymbolEmbedding* xe;
  SymbolEmbedding* ye;
  DurationEmbedding* de;
  BiTrans<Builder> bt;
  SegEmbedBi<Builder> seb;
  explicit ZerothOrderSemiCRF(Model& model) :
      bt(model), seb(model) {
    xe = new SymbolEmbedding(model, VOCAB_SIZE, INPUT_DIM);

    if (use_pretrained_embeding) {
       ifstream fin("/usr1/corpora/LDC2011T13/vec-cwin.txt");
       string s;
       while( getline(fin,s) )
       {
        vector <string> fields;
        boost::algorithm::trim(s);
        boost::algorithm::split( fields, s, boost::algorithm::is_any_of( " " ) );
        string word = fields[0];
        vector<float> p_embeding;
        for (int ind = 1; ind < fields.size(); ++ind){
          p_embeding.push_back(std::stod(fields[ind]));
        }
        if (d.Contains(word)){
          // cout << "init" << endl;
          xe->p_labels->Initialize(d.Convert(word), p_embeding);
        }

      }
    }

    ye = new SymbolEmbedding(model, TAG_SIZE, TAG_DIM);
    de = new BinnedDurationEmbedding(model, DURATION_DIM);

    // potential lower layer
    p_d2h1 = model.add_parameters({H1DIM, DURATION_DIM});
    p_y2h1 = model.add_parameters({H1DIM, TAG_DIM});
    p_fwd2h1 = model.add_parameters({H1DIM, SEG_DIM});
    p_rev2h1 = model.add_parameters({H1DIM, SEG_DIM});
    p_h1b = model.add_parameters({H1DIM});

    // potential upper layer
    p_h12h2 = model.add_parameters({H2DIM, H1DIM});
    p_h2b = model.add_parameters({H2DIM});

    // potential output layer
    p_h22o = model.add_parameters({1, H2DIM});
    p_ob = model.add_parameters({1});
  }

  // return Expression of total loss
  Expression SupervisedLoss(const vector<int>& x,
                            const vector<pair<int,int>>& yz,  // .first = y, .second = duration (z)
                            ComputationGraph& cg,
                            int max_seg_len = 0) {
    int len = x.size();

    // is_ref[i][j][y] returns true iff, in the reference, span(i,j) is labeled y
    vector<vector<vector<bool>>> is_ref(len, vector<vector<bool>>(len+1, vector<bool>(TAG_SIZE, false)));
    unsigned cur = 0;
    for (unsigned ri = 0; ri < yz.size(); ++ri) {
      assert(cur < len);
      int y = yz[ri].first;
      int dur = yz[ri].second;
      if (max_seg_len && dur > max_seg_len) {
        cerr << "max_seg_len=" << max_seg_len << " but reference duration is " << dur << endl;
        abort();
      }
      unsigned j = cur + dur;
      assert(j <= len);
      is_ref[cur][j][y] = true;
      // cerr << "Span[" << cur << "," << j << ")=" << td.Convert(y) << endl;
      cur = j;
    }
    assert(cur == len);

    xe->new_graph(cg);
    ye->new_graph(cg);
    de->new_graph(cg);
    Expression d2h1 = parameter(cg, p_d2h1);
    Expression y2h1 = parameter(cg, p_y2h1);
    Expression fwd2h1 = parameter(cg, p_fwd2h1);
    Expression rev2h1 = parameter(cg, p_rev2h1);
    Expression h1b = parameter(cg, p_h1b);
    Expression h12h2 = parameter(cg, p_h12h2);
    Expression h2b = parameter(cg, p_h2b);
    Expression h22o = parameter(cg, p_h22o);
    Expression ob = parameter(cg, p_ob);
    vector<Expression> xins(x.size());
    for (int i = 0; i < len; ++i)
      xins[i] = xe->embed(x[i]);
    vector<Expression> c = bt.transcribe(cg, xins);
    seb.construct_chart(cg, c, max_seg_len);
    vector<Expression> fwd(len+1);  // fwd trellis for model
    vector<Expression> ref_fwd(len+1); // fwd trellis for reference
    vector<Expression> f, fr;

    // careful: in the other algorithms spans are [i,j], here they are [i,j)
    for (int j = 1; j <= len; ++j) {
      // fwd[j] is the total unnoramlized probability for all segmentations / labels
      // ending (and including) the symbol at position j-1
      f.clear(); // f stores all additive contributions to item [j]
      fr.clear();
      for (unsigned tag = 0; tag < TAG_SIZE; ++tag) {
        Expression y = ye->embed(tag);
        const int i_start = max_seg_len ? max(0, j - max_seg_len) : 0;
        for (int i = i_start; i < j; ++i) {  // i is the starting position
          bool matches_ref = is_ref[i][j][tag];
          auto seg_embedding_ij = seb(i, j-1); // pair<expr, expr>
          Expression d = de->embed(j - i);
          // factor includes: fwd embedding, rev embedding, duration embedding, label embedding
          Expression h1 = rectify(affine_transform({h1b, d2h1, d, y2h1, y, fwd2h1, seg_embedding_ij.first, rev2h1, seg_embedding_ij.second}));
          Expression h2 = tanh(affine_transform({h2b, h12h2, h1}));
          //Expression p = exp(affine_transform({ob, h22o, h2}));
          Expression p = affine_transform({ob, h22o, h2});
          if (i == 0) { // fwd[0] is the path up and including -1, so it's the empty set, i.e., its probability is 1
            f.push_back(p);
            if (matches_ref) fr.push_back(p);
          } else {
            // f.push_back(p * fwd[i]);
            // if (matches_ref) fr.push_back(p * ref_fwd[i]);
            f.push_back(p + fwd[i]);
            if (matches_ref) fr.push_back(p + ref_fwd[i]);
          }
        }
      }
      // cerr << "info\t" << j << "\tsize: " << f.size() << endl;
      // for(auto& e : f){
      //   cerr << "node\t" << as_scalar(cg.get_value(e.i)) << endl;
      // }
      // cerr << "size of vector f at j = " << j << " equals " << f.size() << endl;
      // cerr << "info end" << endl;
      fwd[j] = logsumexp(f);

      // cerr << "shock me\t" << as_scalar(cg.get_value((fwd[j]))) << endl;

      if (fr.size()) ref_fwd[j] = logsumexp(fr);
    }
    //return log(fwd.back()) - log(ref_fwd.back());
    return fwd.back() - ref_fwd.back();
  }

  Expression ViterbiDecode(const vector<int>& x,
                            const vector<pair<int,int>>& yz_gold,  // .first = y, .second = duration (z)
                            ComputationGraph& cg,
                            vector<pair<int, int>> &yz_pred,
                            int max_seg_len = 0) {
    yz_pred.clear();
    int len = x.size();

    xe->new_graph(cg);
    ye->new_graph(cg);
    de->new_graph(cg);
    Expression d2h1 = parameter(cg, p_d2h1);
    Expression y2h1 = parameter(cg, p_y2h1);
    Expression fwd2h1 = parameter(cg, p_fwd2h1);
    Expression rev2h1 = parameter(cg, p_rev2h1);
    Expression h1b = parameter(cg, p_h1b);
    Expression h12h2 = parameter(cg, p_h12h2);
    Expression h2b = parameter(cg, p_h2b);
    Expression h22o = parameter(cg, p_h22o);
    Expression ob = parameter(cg, p_ob);
    vector<Expression> xins(x.size());
    for (int i = 0; i < len; ++i)
      xins[i] = xe->embed(x[i]);
    vector<Expression> c = bt.transcribe(cg, xins);
    seb.construct_chart(cg, c, max_seg_len);
    vector<Expression> fwd(len+1);  // fwd trellis for model
    vector<Expression> f, fr;
    vector<std::tuple<int, int, unsigned>> ijt; // ijt stores positions where we get the max (i.e. i, j and tag)
    vector<std::tuple<int, int, unsigned>> it; // it stores positions where we get the max ending at j
    it.push_back(make_tuple(0,0,0)); // push one to make the index consistent, now it[j] means j rather than j+1
    // careful: in the other algorithms spans are [i,j], here they are [i,j)
    for (int j = 1; j <= len; ++j) {
      // fwd[j] is the total unnoramlized probability for all segmentations / labels
      // ending (and including) the symbol at position j-1
      f.clear(); // f stores all additive contributions to item [j]
      fr.clear();
      ijt.clear();
      for (unsigned tag = 0; tag < TAG_SIZE; ++tag) {
        Expression y = ye->embed(tag);
        const int i_start = max_seg_len ? max(0, j - max_seg_len) : 0;
        for (int i = i_start; i < j; ++i) {  // i is the starting position
          auto seg_embedding_ij = seb(i, j-1); // pair<expr, expr>
          Expression d = de->embed(j - i);
          // factor includes: fwd embedding, rev embedding, duration embedding, label embedding
          Expression h1 = rectify(affine_transform({h1b, d2h1, d, y2h1, y, fwd2h1, seg_embedding_ij.first, rev2h1, seg_embedding_ij.second}));
          Expression h2 = tanh(affine_transform({h2b, h12h2, h1}));
          //Expression p = exp(affine_transform({ob, h22o, h2}));
          Expression p = affine_transform({ob, h22o, h2});

          if (i == 0) { // fwd[0] is the path up and including -1, so it's the empty set, i.e., its probability is 1
            f.push_back(p);
          } else {
            f.push_back(p + fwd[i]);
          }
          ijt.push_back(std::make_tuple(i, j, tag));
        }
      }
      // cerr << "size of vector f at j = " << j << " equals " << f.size() << endl;
      unsigned max_ind = 0;
      auto max_val = as_scalar(cg.get_value(f[0].i));
      for(auto ind = 1; ind < f.size(); ++ind){
        auto val = as_scalar(cg.get_value(f[ind].i));
        // cerr << val << endl;
        if (max_val < val){
          max_val = val;
          max_ind = ind;
        }
      }
      fwd[j] = f[max_ind];
      it.push_back(ijt[max_ind]);
      // cerr << "max value = \t" << as_scalar(cg.get_value(fwd[j])) << endl;
    }
    // for(auto j = 1; j <= len; ++j) {
    //   cerr << j << "\t" << std::get<0>(it[j]) << "\t" << std::get<1>(it[j]) << "\t" << std::get<2>(it[j]) << endl;
    // }
    auto cur_j = len;
    vector<std::tuple<int, int, unsigned>> pred;
    while(cur_j > 0){
      auto cur_i = std::get<0>(it[cur_j]);
      pred.push_back(make_tuple(cur_i, cur_j, std::get<2>(it[cur_j])));
      yz_pred.push_back(make_pair(std::get<2>(it[cur_j]), cur_j - cur_i));
      cur_j = cur_i;
    }
    std::reverse(pred.begin(),pred.end());
    std::reverse(yz_pred.begin(),yz_pred.end());
    cout << "==pred==" << endl;
    // for(auto e : pred){
    //   cout << std::get<0>(e) << "\t" << std::get<1>(e) << "\t" << std::get<2>(e) << endl;
    // }

    for(auto e: yz_pred){
      cout << e.first << "," << e.second << "\t";
    }
    cout << endl;

    cout << "==gold==" << endl;
    for(auto e: yz_gold){
      cout << e.first << "," << e.second << "\t";
    }
    cout << endl;

    // abort();
    return fwd.back();
  }

  Parameters* p_d2h1, *p_y2h1, *p_fwd2h1, *p_rev2h1, *p_h1b;
  Parameters* p_h12h2, *p_h2b;
  Parameters* p_h22o, *p_ob;
};

// a a 0 1 a ||| O:1 O:1 N:2 O:1
pair<vector<int>,vector<pair<int,int>>> ParseTrainingInstance(const std::string& line) {
  std::istringstream in(line);
  std::string word;
  std::string sep = "|||";
  vector<int> x;
  vector<pair<int,int>> yz;
  while(1) {
    in >> word;
    if (word == sep) break;
    x.push_back(d.Convert(word));
  }
  while(1) {
    in >> word;
    if (!in) break;
    size_t p = word.rfind(':');
    if (p == string::npos || p == 0 || p == (word.size() - 1)) {
      cerr << "mal-formed label: " << word << endl;
      abort();
    }
    int y = td.Convert(word.substr(0, p));
    int z = atoi(word.substr(p+1).c_str());
    yz.push_back(make_pair(y,z));
  }
  return make_pair(x, yz);
}

bool inline check_max_seg(const vector<pair<int,int>>& yz, int max_seg_len = 0){
  if (max_seg_len == 0) return true;
  for (unsigned ri = 0; ri < yz.size(); ++ri) {
    int dur = yz[ri].second;
    if (max_seg_len && dur > max_seg_len) {
      cerr << "SKIP: max_seg_len=" << max_seg_len << " but reference duration is " << dur << endl;
      return false;
    }
  }
  return true;
}

double compute_f_score(vector<vector<pair<int,int>>>& yz_preds, vector<vector<pair<int,int>>>& yz_golds){
  assert(yz_preds.size() == yz_golds.size());
  int p_correct = 0;
  int r_correct = 0;
  int p_total = 0;
  int r_total = 0;
  for (unsigned i = 0; i < yz_preds.size(); i++){
    // for sentence i
    std::set<pair<int,int>> gold;
    std::set<pair<int,int>> pred;
    vector<pair<int,int>>& yz_pred = yz_preds[i];
    vector<pair<int,int>>& yz_gold = yz_golds[i];
    int cur = 0;
    for(auto e : yz_pred){
      pred.insert(make_pair(cur, cur + e.second));
      cur += e.second;
    }
    cur = 0;
    for(auto e : yz_gold){
      gold.insert(make_pair(cur, cur + e.second));
      cur += e.second;
    }

    for (auto e : pred){
      if(gold.find(e) != gold.end()){
        p_correct++;
      }
      p_total++;
    }
    for (auto e : gold){
      if(pred.find(e) != pred.end()){
        r_correct++;
      }
      r_total++;
    }
  }
  double p = (double)(p_correct) / (double)(p_total);
  double r = (double)(r_correct) / (double)(r_total);
  double f = 2.0 * ((p * r) / (p + r));
  cerr << "p: " << p << "\tr: " << r << "\tf: " << f << endl;
  return f;

}

int main(int argc, char** argv) {
  cnn::Initialize(argc, argv);
  if (argc != 3 && argc != 4) {
    cerr << "Usage: " << argv[0] << " corpus.txt dev.txt [model.params]\n";
    return 1;
  }
  vector<pair<vector<int>,vector<pair<int,int>>>> training, dev;
  string line;
  cerr << "Reading training data from " << argv[1] << "...\n";
  {
    ifstream in(argv[1]);
    assert(in);
    while(getline(in, line)) {
      training.push_back(ParseTrainingInstance(line));
    }
  }
  d.Freeze();  // no new word types allowed
  td.Freeze(); // no new tag types allowed

  d.SetUnk("<UNK>"); // set UNK to allow the unseen character in the dev and test set

  cerr << "Reading dev data from " << argv[2] << "...\n";
  {
    ifstream in(argv[2]);
    assert(in);
    while(getline(in, line)) {
      dev.push_back(ParseTrainingInstance(line));
    }
  }
  VOCAB_SIZE = d.size();
  TAG_SIZE = td.size();

  cerr << "TAG_SIZE = " << td.size() << endl;
  ostringstream os;
  os << "0scrf"
     << '_' << LAYERS
     << '_' << INPUT_DIM
     << '_' << XCRIBE_DIM
     << "-pid" << getpid();
  const string fname = os.str() + ".params";
  cerr << "Parameters will be written to: " << fname << endl;
  double best = 9e+99;

  const string f_d_name = os.str() + ".dict";
  ofstream out_d(f_d_name);
  boost::archive::text_oarchive oa_d(out_d);
  oa_d << d;

  const string f_td_name = os.str() + ".tdict";
  ofstream out_td(f_td_name);
  boost::archive::text_oarchive oa_td(out_td);
  oa_td << td;

  Model model;
  auto sgd = new SimpleSGDTrainer(&model);
  int max_seg_len = 10;

  ZerothOrderSemiCRF<LSTMBuilder> crf(model);
  if (argc == 4) {
    string fname = argv[3];
    ifstream in(fname);
    boost::archive::text_iarchive ia(in);
    ia >> model;
  }

  unsigned report_every_i = 10;
  unsigned dev_every_i_reports = 100;
  unsigned si = training.size();
  vector<unsigned> order(training.size());
  for (unsigned i = 0; i < order.size(); ++i) order[i] = i;
  bool first = true;
  int report = 0;
  unsigned lines = 0;
  while(1) {
    Timer iteration("completed in");
    double loss = 0;
    unsigned ttags = 0;
    double correct = 0;
    for (unsigned i = 0; i < report_every_i; ++i) {
      if (si == training.size()) {
        si = 0;
        if (first) { first = false; } else { sgd->update_epoch(); }
        cerr << "**SHUFFLE\n";
        // shuffle(order.begin(), order.end(), *rndeng);
      }

      // build graph for this instance
      ComputationGraph cg;
      auto& sent = training[order[si]];
      ++si;
      if(check_max_seg(sent.second, max_seg_len)){
        crf.SupervisedLoss(sent.first, sent.second, cg, max_seg_len);
        ttags += sent.second.size();
        loss += as_scalar(cg.forward());
        cg.backward();
        sgd->update(1.0);
      }
      ++lines;
    }
    sgd->status();
    cerr << " E = " << (loss / ttags) << " ppl=" << exp(loss / ttags) << " (acc=" << (correct / ttags) << ") ";
    report++;
    if (report % dev_every_i_reports == 0) {
      double dloss = 0;
      unsigned dtags = 0;
      double f_best = 0;
      vector<vector<pair<int,int>>> yz_preds;
      vector<vector<pair<int,int>>> yz_golds;
      for (auto& sent : dev) {
        ComputationGraph cg;
        if(check_max_seg(sent.second, max_seg_len)){
          crf.SupervisedLoss(sent.first, sent.second, cg, max_seg_len);
          dtags += sent.second.size();
          dloss += as_scalar(cg.forward());
        }
        vector<pair<int, int>> yz_pred;
        crf.ViterbiDecode(sent.first, sent.second, cg, yz_pred, max_seg_len);
        yz_golds.push_back(sent.second);
        yz_preds.push_back(yz_pred);
      }
      cout << "iter_end" << endl;

      double f = compute_f_score(yz_preds, yz_golds);

      if (f > f_best) {
        f_best = f;
        ofstream out(fname);
        boost::archive::text_oarchive oa(out);
        oa << model;
      }
      cerr << "\n***DEV [epoch=" << (lines / (double)training.size()) << "] E = " << (dloss / dtags) << " ppl=" << exp(dloss / dtags) << ' ';
    }
  }
  delete sgd;
}


/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*
 Copyright (C) 2006, 2007 Ferdinando Ametrano
 Copyright (C) 2006 Katiuscia Manzoni
 Copyright (C) 2005 Eric Ehlers
 Copyright (C) 2005 Plamen Neykov

 This file is part of QuantLib, a free-software/open-source library
 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it
 under the terms of the QuantLib license.  You should have received a
 copy of the license along with this program; if not, please email
 <quantlib-dev@lists.sf.net>. The license is also available online at
 <http://quantlib.org/license.shtml>.

 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 license for more details.
*/

#if defined(HAVE_CONFIG_H)
    #include <qlo/config.hpp>
#endif

#include <qlo/indexes/ibor/libor.hpp>

#include <ql/indexes/ibor/eurlibor.hpp>
#include <ql/indexes/ibor/usdlibor.hpp>
#include <ql/indexes/ibor/gbplibor.hpp>
#include <ql/indexes/ibor/sonia.hpp>
#include <ql/indexes/ibor/chflibor.hpp>
#include <ql/indexes/ibor/jpylibor.hpp>

#include <ql/utilities/dataparsers.hpp>

#include <boost/algorithm/string/case_conv.hpp>
using boost::algorithm::to_upper_copy;
using ObjectHandler::ValueObject;
using boost::shared_ptr;

namespace QuantLibAddin {

    Libor::Libor(const shared_ptr<ValueObject>& properties,
                 const QuantLib::Currency& currency,
                 const std::string& p,
                 const QuantLib::Handle<QuantLib::YieldTermStructure>& h,
                 bool permanent) : IborIndex(properties, permanent)
    {
        switch (currency.numericCode()) {
          case 978: // EUR
            if (to_upper_copy(p)=="ON")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorEURLibor(0, h));
            else if (to_upper_copy(p)=="1D")
                QL_FAIL("1D is ambigous: please specify ON, TN, or SN");
            else if (to_upper_copy(p)=="TN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorEURLibor(1, h));
            else if (to_upper_copy(p)=="SN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorEURLibor(2, h));
            else if (to_upper_copy(p)=="SW")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::EURLibor(1*QuantLib::Weeks, h));
            else {
                QuantLib::Period pp = QuantLib::PeriodParser::parse(p);
                pp.normalize();
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::EURLibor(pp, h));
            }
            break;
          case 840: // USD
            if (to_upper_copy(p)=="ON")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorUSDLibor(0, h));
            else if (to_upper_copy(p)=="1D")
                QL_FAIL("1D is ambigous: please specify ON, TN, or SN");
            else if (to_upper_copy(p)=="TN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorUSDLibor(1, h));
            else if (to_upper_copy(p)=="SN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorUSDLibor(2, h));
            else if (to_upper_copy(p)=="SW")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::USDLibor(1*QuantLib::Weeks, h));
            else {
                QuantLib::Period pp = QuantLib::PeriodParser::parse(p);
                pp.normalize();
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::USDLibor(pp, h));
            }
            break;
          case 826: // GBP
            if (to_upper_copy(p)=="ON")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorGBPLibor(0, h));
            else if (to_upper_copy(p)=="1D")
                QL_FAIL("1D is ambigous: please specify ON, TN, or SN");
            else if (to_upper_copy(p)=="TN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorGBPLibor(1, h));
            else if (to_upper_copy(p)=="SN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorGBPLibor(2, h));
            else if (to_upper_copy(p)=="SW")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::GBPLibor(1*QuantLib::Weeks, h));
            else {
                QuantLib::Period pp = QuantLib::PeriodParser::parse(p);
                pp.normalize();
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::GBPLibor(pp, h));
            }
            break;
          case 756: // CHF
            if (to_upper_copy(p)=="ON")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorCHFLibor(0, h));
            else if (to_upper_copy(p)=="1D")
                QL_FAIL("1D is ambigous: please specify ON, TN, or SN");
            else if (to_upper_copy(p)=="TN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorCHFLibor(1, h));
            else if (to_upper_copy(p)=="SN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorCHFLibor(2, h));
            else if (to_upper_copy(p)=="SW")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::CHFLibor(1*QuantLib::Weeks, h));
            else {
                QuantLib::Period pp = QuantLib::PeriodParser::parse(p);
                pp.normalize();
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::CHFLibor(pp, h));
            }
            break;
          case 392: // JPY
            if (to_upper_copy(p)=="ON")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorJPYLibor(0, h));
            else if (to_upper_copy(p)=="1D")
                QL_FAIL("1D is ambigous: please specify ON, TN, or SN");
            else if (to_upper_copy(p)=="TN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorJPYLibor(1, h));
            else if (to_upper_copy(p)=="SN")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::DailyTenorJPYLibor(2, h));
            else if (to_upper_copy(p)=="SW")
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::JPYLibor(1*QuantLib::Weeks, h));
            else {
                QuantLib::Period pp = QuantLib::PeriodParser::parse(p);
                pp.normalize();
                libraryObject_ = shared_ptr<QuantLib::IborIndex>(new
                    QuantLib::JPYLibor(pp, h));
            }
            break;
          default:
              QL_FAIL("Unhandled currency " << currency);
        }
    }

    Sonia::Sonia(const shared_ptr<ValueObject>& properties,
                 const QuantLib::Handle<QuantLib::YieldTermStructure>& h,
                 bool permanent)
    : OvernightIndex(properties, permanent)
    {
        libraryObject_ = shared_ptr<QuantLib::Sonia>(new
            QuantLib::Sonia(h));
    }

}


#include "la.hpp"
#include "types.hpp"

#include <iostream>
#include <iterator>
#include <algorithm>
#include <vector>
#include <cmath>
#include <boost/format.hpp>

#include <eigen3/Eigen/Dense>

namespace la {

using namespace types;

/**
 * Implementation
 */

Int round(Real r) {
    return (r > 0.0) ? floor(r + 0.5) : ceil(r - 0.5);
}

Cell::Cell(const std::vector<Direction> &directions) {
    for(std::vector<Direction>::const_iterator dIt = directions.begin();
            dIt != directions.end(); ++dIt) {
        Direction temp = *dIt;
        temp.scaleVector(temp.getNElements());
        vectors.push_back(temp.getIncrementVector());
    }
}

template<unsigned int dim>
Cell Grid<dim>::cell() const {
    return Cell(directions);
}

template<unsigned int dim>
bool Grid<dim>::hasSameGrid(const Grid<dim> &g) const {
    if(this->directions == g.directions &&
            this->originVector == g.originVector)
        return true;
    else return false;
}

bool operator==(const Direction& d1, const Direction& d2) {
    return (d1.incrementVector == d2.incrementVector &&
            d1.nElements == d2.nElements);
}

bool Direction::checkRange(Uint index) const {
    if(index < nElements || periodic) return true;
#ifdef LA_STRICT
    else throw std::range_error("Index out of range.");
#else
    else return false;
#endif
}

void Direction::stride(types::Uint s) {
    this->scaleVector(s);
    if(nElements % s == 0) nElements = nElements/s;
    else nElements = nElements/s + 1;
}

void Direction::scaleVector(types::Real factor) {
    for(std::vector<types::Real>::iterator vIt = incrementVector.begin();
            vIt != incrementVector.end(); ++vIt) {
        *vIt *= factor;
    }
}

template<unsigned int dim>
const Grid<dim> & Grid<dim>::operator+=(const Grid<dim> &g) {
    if(!this->hasSameGrid(g))
        throw types::runtimeError() <<
                                    types::errinfo_runtime("Cannot add data on different grids.");

    else if(this->data.size() == g.data.size() && this->data.size() > 0) {
        std::vector<Real>::iterator dataIt = data.begin(), dataEnd = data.end();
        std::vector<Real>::const_iterator gIt = g.data.begin();

        while(dataIt != dataEnd) {
            *dataIt += *gIt;
            ++dataIt;
            ++gIt;
        }
    }
    return *this;
}

template<unsigned int dim>
void Grid<dim>::stride(std::vector<Uint> strides) {
    Uint dimension = strides.size();
    checkDimension(dimension);

    // Get size for preallocation
    Uint size = 1;
    std::vector<Direction>::iterator dirIt = directions.begin();
    std::vector<Uint>::iterator sIt = strides.begin();
    std::vector<Uint> rests;
    std::vector<Uint> newCounts;
    std::vector<Direction> newDirections;
    while(dirIt != directions.end()) {
        Direction d = *dirIt;
        d.stride(*sIt);
        size *= d.getNElements();
        newDirections.push_back(d);
        newCounts.push_back(d.getNElements());

        // Handle "rest"
        Uint r = dirIt->getNElements() % *sIt;
        if(r > 0) rests.push_back(r);
        else rests.push_back(*sIt);

        ++dirIt;
        ++sIt;
    }
    std::vector<Real> newData;
    newData.resize(size);

    // Starting with direction 0, the slowest direction
    std::vector<Real>::const_iterator dataIt = data.begin();
    std::vector<Real>::iterator newIt = newData.begin();
    this->strideRecursive(0, dataIt, newIt, newCounts, strides, rests);

    // Update data and nElements
    this->data = newData;
    this->directions = newDirections;
}

template<unsigned int dim>
void Grid<dim>::strideRecursive(
    Uint directionIndex,
    std::vector<Real>::const_iterator &oldIt,
    std::vector<Real>::iterator &newIt,
    const std::vector<Uint> &newCounts,
    const std::vector<Uint> &strides,
    const std::vector<Uint> &rests) {

    Uint dimension = strides.size();

    // If directionIndex is out of bounds, we are finished
    if( directionIndex >= dimension) return;

    Uint oldCount = directions[directionIndex].getNElements(),
         newCount = newCounts[directionIndex];

    Uint weight = 1;
    for(Uint i = dimension -1; i > directionIndex; --i) {
        weight *= directions[i].getNElements();
    }
    Uint stride = weight * strides[directionIndex];
    Uint rest = weight * rests[directionIndex];

    // If we are at the last dimension, we copy and move the iterators
    if( directionIndex == dimension - 1) {
        for(Uint i = 0; i < newCount; ++i) {
            *newIt = *oldIt;
            ++newIt;
            oldIt += stride;
        }
    }
    else {
        for(Uint i = 0; i < newCount; ++i) {
            this->strideRecursive(directionIndex + 1, oldIt, newIt, newCounts, strides, rests);
            // The above already moved oldIt by 1xweigth
            oldIt += stride - weight;
        }
    }

    // The above is sufficient, if stride divides nElements
    // In this case, rest = stride
    oldIt += rest;
    oldIt -= stride;
}

/** Can cut down as well as enlarge (fill with zeros)
 *
 * Nd implementation sadly needs recursion (depth = number of dimensions).
 * I should test the speed and maybe think about a specialization for 3d.
 * Maybe one can do the recursion during compile time via templates.
 */
template<unsigned int dim>
void Grid<dim>::resize(const std::vector<Uint>& nElements) {
    Uint dimension = nElements.size();
    checkDimension(dimension);

    // Get size for preallocation
    Uint size = 1;
    std::vector<Uint>::const_iterator incIt = nElements.begin(),
                                      incEnd = nElements.end();
    while(incIt != incEnd) {
        size *= *incIt;
        incIt++;
    }
    std::vector<Real> newData;
    newData.resize(size);

    // Starting with direction 0, the slowest direction
    std::vector<Real>::const_iterator dataIt = data.begin();
    std::vector<Real>::iterator newIt = newData.begin();
    this->copyRecursive(0, dataIt, newIt, nElements);

    // Update data and nElements
    this->data = newData;
    std::vector<Direction>::iterator dirIt = directions.begin();
    for(incIt = nElements.begin(); incIt != incEnd; ++incIt, ++dirIt) {
        Direction temp = Direction(dirIt->getIncrementVector(), *incIt);
        *dirIt = temp;
    }
}

/** Used by resize (recursively).
 *
 * The *last* dimension is the fast one (as in 3d cube file format)
 */
template<unsigned int dim>
void Grid<dim>::copyRecursive(
    Uint directionIndex,
    std::vector<Real>::const_iterator &oldIt,
    std::vector<Real>::iterator &newIt,
    const std::vector<Uint> &newCounts) {

    Uint dimension = newCounts.size();

    // If directionIndex is out of bounds, we are finished
    if( directionIndex >= dimension) return;

    Uint oldCount = directions[directionIndex].getNElements(),
         newCount = newCounts[directionIndex];
    Uint copyCount = std::min(oldCount, newCount),
         iterCount = std::max(oldCount, newCount);

    // If we are at the last dimension, we copy and move the iterators
    if( directionIndex == dimension - 1) {
        copy(oldIt, oldIt + copyCount, newIt);
        oldIt += oldCount;
        newIt += newCount;
    }
    // If we are not at the last dimension, we call recursively
    else {

        for(Uint i = 0; i < iterCount; ++i) {
            // If we need to copy...
            if(i < copyCount) {
                this->copyRecursive(directionIndex + 1, oldIt, newIt, newCounts);
            }
            // Else countNew or countOld is > countCopy and we just need to move iterators.
            else {
                if(newCount > i)
                    newIt += newCount;
                else
                    oldIt += oldCount;
            }
        }

    }
}

template<unsigned int dim>
void Grid<dim>::squareValues() {
    std::vector<Real>::iterator it = data.begin(), end = data.end();
    while(it != end) {
        (*it) *= (*it);
        ++it;
    }
}

template<unsigned int dim>
void Grid<dim>::abs() {
    std::vector<Real>::iterator it = data.begin(), end = data.end();
    while(it != end) {
        *it = std::abs(*it);
        ++it;
    }
}
template<unsigned int dim>
void Grid<dim>::sqrt() {
    std::vector<Real>::iterator it = data.begin(), end = data.end();
    while(it != end) {
        *it = std::sqrt(std::abs(*it));
        ++it;
    }
}

template<unsigned int dim>
void Grid<dim>::signsqrt() {
    std::vector<Real>::iterator it = data.begin(), end = data.end();
    while(it != end) {
        if(*it < 0) *it = -std::sqrt(std::abs(*it));
        else        *it = std::sqrt(*it);
        ++it;
    }
}

template<unsigned int dim>
bool Grid<dim>::checkDimension(Uint size) const {
    if (size != dim) {
        throw std::range_error("Number of given indices does not equal number of grid directions.");
    }

    return true;
}

template<unsigned int dim>
bool Grid<dim>::checkRange(const std::vector<Uint>& indices) const {

    checkDimension(indices.size());
    std::vector<Uint>::const_iterator indexIt = indices.end(), indexEnd = indices.end();
    std::vector<Direction>::const_iterator dirIt = directions.begin();
    while(indexIt != indexEnd) {
        if( !dirIt->checkRange(*indexIt) ) return false;
        ++indexIt;
        ++dirIt;
    }
    return true;
}

template<unsigned int dim>
Real Grid<dim>::getNearestDataPoint(std::vector<Real>& coordinates) const {
    std::vector<Uint> indices;
    this->getNearestIndices(coordinates, indices);
    return this->getDataPoint(indices);
}

// This works for n dimensions,
// assuming that they are ordered by increasing fastness
template<unsigned int dim>
Real Grid<dim>::getDataPoint(std::vector<Uint> indices) const {
    wrapAround(indices);
    checkRange(indices);

    // Note: end() points *after* the last item
    std::vector<Uint>::const_iterator indexIt = indices.end(), indexBegin = indices.begin();
    std::vector<Direction>::const_iterator dirIt = directions.end(), dirFirst = directions.begin();

    --indexIt;
    std::vector<Real>::const_iterator dataIt = data.begin();
    dataIt += *indexIt;
    Uint weight = 1;
    while(dirIt != dirFirst) {
        --dirIt;
        --indexIt;
        weight *= dirIt->getNElements();
        dataIt += (*indexIt) * weight;
    }

    return *dataIt;
}

// This works for nd grids
template<unsigned int dim>
bool Grid<dim>::getFractionalIndices(std::vector<types::Real>& cartesianCoordinates, std::vector<types::Real>& fractionalIndices) const {
    if (!checkDimension(dim)) return false;

    std::vector<Real> basisVectors;
    std::back_insert_iterator< std::vector<Real> > it(basisVectors);
    std::vector<Direction>::const_iterator dirIt = directions.begin(), dirEnd = directions.end();
    while(dirIt != dirEnd) {
        std::vector<Real> temp = dirIt->getIncrementVector();
        copy(temp.begin(), temp.end(), it);
        ++dirIt;
    }

    using Eigen::MatrixXd;
    using Eigen::VectorXd;
    using Eigen::Map;
    Map<MatrixXd> basisMatrix(&basisVectors[0], dim, dim);
    Map<VectorXd> cartesianEig(&cartesianCoordinates[0], dim);
    VectorXd gridEig = basisMatrix.inverse() * cartesianEig;

    std::vector<Real> gridStl(gridEig.data(), gridEig.data() + gridEig.size());
    fractionalIndices = gridStl;

    return true;
}

template<unsigned int dim>
Real Grid<dim>::interpolateDataPoint(std::vector<Real> coordinates) const {
    // Get indices of relevant unit cell and the relative location inside the cell
    std::vector<Real> fractionalIndices;
    getFractionalIndices(coordinates, fractionalIndices);

    std::vector<Real>::const_iterator it = fractionalIndices.begin(),
        end = fractionalIndices.end();
    std::vector<Uint> unitCellIndices;
    std::vector<Real> shiftedFractionalIndices;
    while(it != end){
        unitCellIndices.push_back( Uint(*it) );
        shiftedFractionalIndices.push_back( *it - Uint(*it));
        ++it;
    }

    // Loop over corners of unit cell
    std::vector< std::vector<Uint> > corners;

    Uint nCorners = 1 << dim; // 1 * 2^dim
    Uint bit;
    Real sum = 0, summand;
    for(Uint corner = 0; corner < nCorners; ++corner) {
        std::vector<Uint> cornerIndices = unitCellIndices;
        summand = 1;

        // Get binary representation of corner integer
        for(Uint i = 0; i < dim; ++i){
            bit = (corner >> i) & 1; // Either 0 or 1
            cornerIndices[i] += bit;
            if(bit == 0) summand *= shiftedFractionalIndices[i];
            else         summand *= (1 - shiftedFractionalIndices[i]);
        }

        summand *= getDataPoint(cornerIndices);
        sum     += summand;

    }

    return sum;
}

template<unsigned int dim>
void Grid<dim>::wrapAround(std::vector<Uint> &indices) const {
    std::vector<Uint>::iterator it = indices.begin(),
        end = indices.end();
    std::vector<Direction>::const_iterator dirIt = directions.begin();
    while(it != end){
        if(*it > dirIt->getNElements() && dirIt->isPeriodic())
            *it = *it % dirIt->getNElements();
        ++it;
    }
}

template<unsigned int dim>
bool Grid<dim>::getNearestIndices(std::vector<Real>& cartesianCoordinates, std::vector<Uint>& indices) const {
    std::vector<Real> fractionalIndices;
    this->getFractionalIndices(cartesianCoordinates, fractionalIndices);

    indices.clear();
    for(std::vector<Real>::iterator it = fractionalIndices.begin(); it != fractionalIndices.end(); ++it) {
        indices.push_back(round(*it));
    }

    return true;
}

template<unsigned int dim>
const Grid<dim> & Grid<dim>::operator*=(types::Real x) {
    std::vector<Real>::iterator dataIt = data.begin(),
                                dataEnd = data.end();
    while(dataIt != dataEnd) {
        *dataIt *= x;
        ++dataIt;
    }
    return *this;
}

template<unsigned int dim>
Real operator*(const Grid<dim>& g1, const Grid<dim>& g2) {
    if(! g1.hasSameGrid(g2)) throw types::runtimeError() <<
            types::errinfo_runtime("Cannot multiply data on different grids.");

    std::vector<Real>::const_iterator
        it1 = g1.data.begin(), end1= g1.data.end(),
        it2 = g2.data.begin(), end2= g2.data.end();

    Real sum = 0;
    while(it1 != end1) {
        sum += *it1 * *it2;
        ++it1; ++it2;
    }
    return sum;
}

template<unsigned int dim>
Uint Grid<dim>::countPoints() const {
    std::vector<Direction>::const_iterator it;
    Uint points = 1;
    for(it = directions.begin(); it!= directions.end(); ++it) {
        points *= it->getNElements();
    }
    return points;
}

types::Real Cell::getExtent(types::Uint index) const {
    std::vector<Real>::const_iterator it = vectors[index].begin(),
       end = vectors[index].end();
    types::Real sum = 0;
    while(it != end){
        sum += *it * *it;
        ++it;
    }
    return std::sqrt(sum);
}

std::vector<Real> Cell::extents() const {
    std::vector<Real> extents;
    for(Uint i = 0; i < vectors.size(); ++i){
        extents.push_back(getExtent(i));
    }

    return extents;
}

// TODO: replace by Kahan sum
template<unsigned int dim>
types::Real Grid<dim>::sum() const{
    std::vector<Real>::const_iterator it = data.begin(),
        end = data.end();
    Real sum = 0;
    while(it != end){
        sum += *it;
        ++it;
    }
    return sum;
}

/**
 * Required template specializations
 * (need to be given *after* implementation)
 */

template class Grid<3>;
template Real operator*(const Grid<3>& g1, const Grid<3>& g2);

}


#pragma once

#include <boost/hana/range.hpp>
#include <boost/hana/transform.hpp>
#include <boost/hana/tuple.hpp>

namespace lys::str
{

template <size_t N>
constexpr auto repeat = [](auto && str) { //
    return boost::hana::transform(boost::hana::to_tuple(boost::hana::range_c<int, 0, N>), [&](auto /*unused*/) { return str; });
};
}


#include <scitbx/array_family/boost_python/flex_fwd.h>

#include <scitbx/array_family/boost_python/flex_wrapper.h>
#include <scitbx/serialization/single_buffered.h>
#include <scitbx/matrix/transpose_multiply.h>
#include <scitbx/math/utils.h>
#include <scitbx/mat3.h>
#include <boost/python/make_constructor.hpp>
#include <boost/python/args.hpp>
#include <boost/python/return_arg.hpp>
#include <boost/format.hpp>
#include <scitbx/array_family/boost_python/flex_helpers.h>

namespace scitbx { namespace serialization { namespace single_buffered {

  inline
  char* to_string(char* start, scitbx::mat3<double> const& value)
  {
    return
      to_string(to_string(to_string(
        to_string(to_string(to_string(
          to_string(to_string(to_string(
            start,
              value[0]), value[1]), value[2]),
              value[3]), value[4]), value[5]),
              value[6]), value[7]), value[8]);
  }

  template <>
  struct from_string<scitbx::mat3<double> >
  {
    from_string(const char* start)
    {
      end = start;
      for(std::size_t i=0;i<9;i++) {
        from_string<double> proxy(end);
        value[i] = proxy.value;
        end = proxy.end;
      }
    }

    scitbx::mat3<double> value;
    const char* end;
  };

}}} // namespace scitbx::serialization::single_buffered

#include <scitbx/array_family/boost_python/flex_pickle_single_buffered.h>

namespace scitbx { namespace af { namespace boost_python {

  template <>
  struct flex_default_element<scitbx::mat3<double> >
  {
    static scitbx::mat3<double>
    get() { return scitbx::mat3<double>(0,0,0,0,0,0,0,0,0); }
  };

}}}

namespace scitbx { namespace af {
namespace {

  scitbx::af::flex<scitbx::mat3<double> >::type*
  from_double(
    scitbx::af::const_ref<double> const& x)
  {
    SCITBX_ASSERT(x.size() % 9 == 0);
    std::size_t result_size = x.size() / 9;
    scitbx::af::shared<scitbx::mat3<double> > result(result_size,
      scitbx::af::init_functor_null< scitbx::mat3<double> >());
    const double* d = x.begin();
    for(std::size_t i=0;i<result_size;i++) {
      for (std::size_t j=0;j<9;++j) {
        result[i][j] = *d++;
      }
    }
    return new scitbx::af::flex<scitbx::mat3<double> >::type(result, result.size());
  }

  scitbx::af::flex_double
  as_double(scitbx::af::flex<scitbx::mat3<double> >::type const& a)
  {
    SCITBX_ASSERT(a.accessor().is_trivial_1d());
    scitbx::af::flex_double result(a.size()*9, scitbx::af::init_functor_null<double>());
    double* r = result.begin();
    scitbx::af::const_ref<scitbx::mat3<double> > a_ref = a.const_ref().as_1d();
    for(std::size_t i=0;i<a_ref.size();i++) {
      for(std::size_t j=0;j<9;j++) {
        *r++ = a_ref[i][j];
      }
    }
    return result;
  }

  scitbx::af::shared<mat3<double> >
  mul_a_scalar(
    scitbx::af::const_ref<mat3<double> > const& a,
    double f)
  {
    scitbx::af::shared<mat3<double> > result((af::reserve(a.size())));
    for(std::size_t i=0;i<a.size();i++) {
      result.push_back(a[i] * f);
    }
    return result;
  }

  scitbx::af::shared<mat3<double> >
  mul_a_a_scalar(
    scitbx::af::const_ref<mat3<double> > const& lhs,
    scitbx::af::const_ref<double> const& rhs)
  {
    SCITBX_ASSERT(lhs.size() == rhs.size());
    scitbx::af::shared<mat3<double> > result((af::reserve(lhs.size())));
    for(std::size_t i=0;i<lhs.size();i++) {
      result.push_back(lhs[i] * rhs[i]);
    }
    return result;
  }

  scitbx::af::shared<vec3<double> >
  mul_a_vec3(
    scitbx::af::const_ref<mat3<double> > const& a,
    vec3<double> const& m)
  {
    af::shared<vec3<double> > result((af::reserve(a.size())));
    for(std::size_t i=0;i<a.size();i++) {
      result.push_back(a[i] * m);
    }
    return result;
  }

  scitbx::af::shared<vec3<double> >
  mul_a_a_vec3(
    scitbx::af::const_ref<mat3<double> > const& lhs,
    scitbx::af::const_ref<vec3<double> > const& rhs)
  {
    SCITBX_ASSERT(lhs.size() == rhs.size());
    af::shared<vec3<double> > result((af::reserve(lhs.size())));
    for(std::size_t i=0;i<lhs.size();i++) {
      result.push_back(lhs[i] * rhs[i]);
    }
    return result;
  }

  scitbx::af::shared<mat3<double> >
  mul_a_mat3(
    scitbx::af::const_ref<mat3<double> > const& a,
    mat3<double> const& m)
  {
    af::shared<mat3<double> > result((af::reserve(a.size())));
    for(std::size_t i=0;i<a.size();i++) {
      result.push_back(a[i] * m);
    }
    return result;
  }

  scitbx::af::shared<mat3<double> >
  mul_a_a_mat3(
    scitbx::af::const_ref<mat3<double> > const& lhs,
    scitbx::af::const_ref<mat3<double> > const& rhs)
  {
    SCITBX_ASSERT(lhs.size() == rhs.size());
    af::shared<mat3<double> > result((af::reserve(lhs.size())));
    for(std::size_t i=0;i<lhs.size();i++) {
      result.push_back(lhs[i] * rhs[i]);
    }
    return result;
  }

} // namespace <anonymous>

namespace boost_python {

  void wrap_flex_mat3_double()
  {
    using namespace boost::python;
    using boost::python::arg;
    typedef scitbx::af::boost_python::flex_wrapper<scitbx::mat3<double> > f_w;
    f_w::plain("mat3_double")
      .def_pickle(scitbx::af::boost_python::flex_pickle_single_buffered<
        scitbx::mat3<double>,
        9*scitbx::af::boost_python::pickle_size_per_element<double>::value>())
      .def("__init__", make_constructor(from_double))
      .def("__mul__", mul_a_scalar)
      .def("__mul__", mul_a_a_scalar)
      .def("__mul__", mul_a_vec3)
      .def("__mul__", mul_a_a_vec3)
      .def("__mul__", mul_a_mat3)
      .def("__mul__", mul_a_a_mat3)
      .def("as_double", as_double);
    ;
  }

}}} // namespace scitbx::af::boost_python