#include <boost/python/module.hpp>
#include <boost/python/def.hpp>
#include <dials/array_family/reflection_table.h>
#include <cctbx/miller.h>

namespace xfel {
namespace merging {
namespace error_model {
namespace sdfac_refine {

typedef
 scitbx::af::versa<cctbx::miller::index<>, scitbx::af::flex_grid<> > shared_miller;

using namespace dials::af;

static scitbx::af::shared<double>
compute_normalized_deviations(reflection_table ISIGI, shared_miller hkl_list) {
  /*
   * This formulation of the normalized deviations of a set of intensities and sigmas is similar to that
   * described in Evans 2011, but includes the nn term as currently implmented by aimless
   *
   */
  SCITBX_ASSERT(ISIGI.contains("scaled_intensity"));
  SCITBX_ASSERT(ISIGI.contains("meanprime_scaled_intensity"));
  SCITBX_ASSERT(ISIGI.contains("isigi"));
  SCITBX_ASSERT(ISIGI.contains("nn"));

  scitbx::af::shared<double>         result(ISIGI.size(), 0);
  scitbx::af::const_ref<double>      scaled_intensity = ISIGI["scaled_intensity"];
  scitbx::af::const_ref<double>      meanprime_scaled_intensity = ISIGI["meanprime_scaled_intensity"];
  scitbx::af::const_ref<double>      isigi = ISIGI["isigi"];
  scitbx::af::const_ref<double>      nn = ISIGI["nn"];

  double sigma;
  #pragma omp parallel for private(sigma)
  for (int i = 0; i < ISIGI.size(); i++) {
    if (isigi[i] == 0 || scaled_intensity[i] == 0) continue;
    sigma = scaled_intensity[i] / isigi[i];
    result[i] = std::sqrt(nn[i]) * (scaled_intensity[i] - meanprime_scaled_intensity[i]) / sigma;
  }
  return result;
}

void
apply_sd_error_params(reflection_table ISIGI, const double sdfac, const double sdb, const double sdadd, const bool squared_params) {
  /*
   * Apply a set of sd params (sdfac, sdb and sdd) to an ISIGI reflection table

   Squared not only uses the squared formulation of sigma', but also fixes 2 bugs:
   1) Use meanI not meanIprime
   2) When returning isigi, don't multiply by slope

   If using squared, it is assumed that sdfac, sdb and sdadd have already been squared
   */
  SCITBX_ASSERT(ISIGI.contains("scaled_intensity"));
  SCITBX_ASSERT(ISIGI.contains("mean_scaled_intensity"));
  SCITBX_ASSERT(ISIGI.contains("meanprime_scaled_intensity"));
  SCITBX_ASSERT(ISIGI.contains("isigi"));
  SCITBX_ASSERT(ISIGI.contains("slope"));
  SCITBX_ASSERT(ISIGI.contains("miller_id"));

  scitbx::af::const_ref<double>      scaled_intensity = ISIGI["scaled_intensity"];
  scitbx::af::const_ref<double>      mean_scaled_intensity = ISIGI["mean_scaled_intensity"];
  scitbx::af::const_ref<double>      meanprime_scaled_intensity = ISIGI["meanprime_scaled_intensity"];
  scitbx::af::ref<double>            isigi = ISIGI["isigi"];
  scitbx::af::const_ref<double>      slope = ISIGI["slope"];
  scitbx::af::shared<double>         sigmas(ISIGI.size(), 0);
  scitbx::af::const_ref<std::size_t> miller_id = ISIGI["miller_id"];

  std::size_t max_miller_id = scitbx::af::max(miller_id);
  scitbx::af::shared<std::size_t>    n_refl(max_miller_id+1, 0);
  scitbx::af::shared<double>         isum(max_miller_id+1, 0);

  double tmp = 0;
  double sigma_corrected = 0;
  double meanI, meanIprime, minimum;
  #pragma omp parallel for private(meanI, meanIprime, minimum)
  for (int i = 0; i < ISIGI.size(); i++) {
    // scaled intensity (iobs/slope)
    // corrected sigma (original sigma/slope)
    sigmas[i] = scaled_intensity[i] / isigi[i];

    // apply correction parameters
    if (squared_params) {
      // use meanI, which is the mean of all observations of this hkl
      meanI = mean_scaled_intensity[i];
      tmp = std::pow(sigmas[i],2) + sdb * meanI + sdadd * std::pow(meanI,2);
    }
    else {
      // use meanIprime, which for each observation, is the mean of all other observations of this hkl
      meanIprime = meanprime_scaled_intensity[i];
      tmp = std::pow(sigmas[i],2) + sdb * meanIprime + std::pow(sdadd*meanIprime,2);
    }

    // avoid rare negatives
    minimum = 0.1 * std::pow(sigmas[i],2);
    if (tmp < minimum)
      tmp = minimum;

    if (squared_params) {
      sigma_corrected = std::sqrt(sdfac * tmp);
      SCITBX_ASSERT(sigma_corrected != 0.0);
      isigi[i] = scaled_intensity[i] / sigma_corrected;
    }
    else {
      sigma_corrected = sdfac * std::sqrt(tmp);
      SCITBX_ASSERT(sigma_corrected != 0.0);
      isigi[i] = scaled_intensity[i] * slope[i]/ sigma_corrected;
    }
  }
}

scitbx::af::shared<double> df_dpsq(scitbx::af::shared<double>all_sigmas_normalized,
                                   scitbx::af::shared<double>sigma_prime,
                                   scitbx::af::shared<double>dsigmasq_dpsq,
                                   reflection_table ISIGI,
                                   scitbx::af::shared<int>bin_indices,
                                   size_t n_bins) {
  scitbx::af::shared<double> g(n_bins, 0);
  scitbx::af::shared<double> dsigmanormsq_dpsq(ISIGI.size());

  scitbx::af::const_ref<double> scaled_intensity = ISIGI["scaled_intensity"];
  scitbx::af::const_ref<double> meanprime_scaled_intensity = ISIGI["meanprime_scaled_intensity"];
  scitbx::af::const_ref<double> nn = ISIGI["nn"];

  scitbx::af::shared<int> counts(n_bins);
  scitbx::af::shared<double> bnssq(n_bins);
  scitbx::af::shared<double> t3(n_bins);

  double tmp, c; int b;
  // OpenMP doesn't work here. Need a reduction. Easy in OpenMP 4.5 which isn't available yet for my gcc
  for (int i = 0; i < ISIGI.size(); i++) {
    tmp = scaled_intensity[i]-meanprime_scaled_intensity[i];
    c = nn[i] * tmp * tmp;
    dsigmanormsq_dpsq[i] = -c / std::pow(sigma_prime[i], 4) * dsigmasq_dpsq[i];

    b = bin_indices[i];
    counts[b]++;
    bnssq[b] += all_sigmas_normalized[i] * all_sigmas_normalized[i];
    t3[b] += dsigmanormsq_dpsq[i];
  }

  int n; double t1, t2;
  #pragma omp parallel for private(n, t1, t2)
  for (int b = 0; b < n_bins; b++) {
    n = counts[b];

    if (!(n == 0 || bnssq[b] == 0.)) {
      bnssq[b] /= n;
      t1 = 2 * (1 - std::sqrt(bnssq[b]));
      t2 = -0.5 / std::sqrt(bnssq[b]);
      t3[b] /= n;
      g[b] = t1 * t2 * t3[b];
    }
  }
  return g;
}

/* Jiffy function to compute statistics needed downstream
  For every observation, computes:
  mean_scaled_intensity: mean of all observations of this miller index
  meanprime_scaled_intensity: mean of all observations of this miller index except this observation
  n_refl: count of observed reflections for this miller index
  nn: n_refl-1/n_refl
*/
void setup_isigi_stats(reflection_table ISIGI, scitbx::af::const_ref<cctbx::miller::index<> > indices) {
  scitbx::af::shared<double> sumI                = scitbx::af::shared<double>(indices.size(), 0);
  scitbx::af::shared<double> n_refl              = scitbx::af::shared<double>(indices.size(), 0);
  scitbx::af::const_ref<size_t> miller_id        = ISIGI["miller_id"];
  scitbx::af::const_ref<double> scaled_intensity = ISIGI["scaled_intensity"];

  for (size_t i = 0; i < ISIGI.size(); i++) {
    size_t hkl_id = miller_id[i];
    sumI[hkl_id] += scaled_intensity[i];
    n_refl[hkl_id] += 1;
  }

  scitbx::af::shared<double> all_meanI      = scitbx::af::shared<double>(ISIGI.size(), 0);
  scitbx::af::shared<double> all_n_refl     = scitbx::af::shared<double>(ISIGI.size(), 0);
  scitbx::af::shared<double> nn             = scitbx::af::shared<double>(ISIGI.size(), 0);
  scitbx::af::shared<double> all_imeanprime = scitbx::af::shared<double>(ISIGI.size(), 0);
  for (size_t i = 0; i < ISIGI.size(); i++) {
    size_t hkl_id = miller_id[i];
    all_meanI[i] = sumI[hkl_id]/n_refl[hkl_id];
    all_n_refl[i] = n_refl[hkl_id];
    nn[i] = (n_refl[hkl_id]-1)/n_refl[hkl_id];
    SCITBX_ASSERT(n_refl[hkl_id] > 0);
    if (n_refl[hkl_id] > 1)
      all_imeanprime[i] = (sumI[hkl_id]-scaled_intensity[i])/(n_refl[hkl_id]-1);
  }
  ISIGI["mean_scaled_intensity"] = all_meanI;
  ISIGI["n_refl"] = all_n_refl;
  ISIGI["nn"] = nn;
  ISIGI["meanprime_scaled_intensity"] = all_imeanprime ;
}

namespace boost_python { namespace {
  void
  init_module() {
    using namespace boost::python;
    def("compute_normalized_deviations", &compute_normalized_deviations);
    def("apply_sd_error_params", &apply_sd_error_params);
    def("df_dpsq", &df_dpsq);
    def("setup_isigi_stats", &setup_isigi_stats);
}
}}
}}}} // namespace

BOOST_PYTHON_MODULE(xfel_sdfac_refine_ext)
{
  xfel::merging::error_model::sdfac_refine::boost_python::init_module();

}


#include <Eigen/Geometry>
#include <algorithm>
#include <elasty/cloth-sim-object.hpp>
#include <elasty/constraint.hpp>
#include <elasty/particle.hpp>
#include <elasty/utils.hpp>
#include <iostream>
#include <tiny_obj_loader.h>

elasty::ClothSimObject::ClothSimObject(const unsigned           resolution,
                                       const double             in_plane_stiffness,
                                       const double             in_plane_compliance,
                                       const double             out_of_plane_stiffness,
                                       const double             out_of_plane_compliance,
                                       const double             dt,
                                       const Eigen::Affine3d&   transform,
                                       const InPlaneStrategy    in_plane_strategy,
                                       const OutOfPlaneStrategy out_of_plane_strategy)
{
    std::istringstream obj_data_stream(generateClothMeshObjData(2.0, 2.0, resolution, resolution));

    tinyobj::attrib_t                attrib;
    std::vector<tinyobj::shape_t>    shapes;
    std::vector<tinyobj::material_t> materials;

    std::string warn;
    std::string err;
    const bool  return_value = tinyobj::LoadObj(&attrib, &shapes, &materials, &warn, &err, &obj_data_stream);

    if (!warn.empty())
    {
        std::cerr << warn << std::endl;
    }
    if (!err.empty())
    {
        std::cerr << err << std::endl;
    }
    if (!return_value)
    {
        throw std::runtime_error("");
    }
    if (attrib.vertices.empty())
    {
        throw std::runtime_error("");
    }

    const auto& shape = shapes[0];

    assert(shapes.size() == 1);
    assert(attrib.vertices.size() % 3 == 0);
    assert(shape.mesh.indices.size() % 3 == 0);

    m_triangle_list.resize(shape.mesh.indices.size() / 3, 3);
    for (unsigned int i = 0; i < shape.mesh.indices.size() / 3; ++i)
    {
        m_triangle_list(i, 0) = shape.mesh.indices[i * 3 + 0].vertex_index;
        m_triangle_list(i, 1) = shape.mesh.indices[i * 3 + 1].vertex_index;
        m_triangle_list(i, 2) = shape.mesh.indices[i * 3 + 2].vertex_index;
    }

    if (!attrib.texcoords.empty())
    {
        m_uv_list.resize(shape.mesh.indices.size() / 3, 2 * 3);
        for (unsigned int i = 0; i < shape.mesh.indices.size() / 3; ++i)
        {
            m_uv_list(i, 2 * 0 + 0) = attrib.texcoords[2 * shape.mesh.indices[i * 3 + 0].texcoord_index + 0];
            m_uv_list(i, 2 * 0 + 1) = attrib.texcoords[2 * shape.mesh.indices[i * 3 + 0].texcoord_index + 1];
            m_uv_list(i, 2 * 1 + 0) = attrib.texcoords[2 * shape.mesh.indices[i * 3 + 1].texcoord_index + 0];
            m_uv_list(i, 2 * 1 + 1) = attrib.texcoords[2 * shape.mesh.indices[i * 3 + 1].texcoord_index + 1];
            m_uv_list(i, 2 * 2 + 0) = attrib.texcoords[2 * shape.mesh.indices[i * 3 + 2].texcoord_index + 0];
            m_uv_list(i, 2 * 2 + 1) = attrib.texcoords[2 * shape.mesh.indices[i * 3 + 2].texcoord_index + 1];
        }
    }

    std::map<unsigned int, std::shared_ptr<Particle>> map_from_obj_vertex_index_to_particle;
    for (unsigned int i = 0; i < attrib.vertices.size() / 3; ++i)
    {
        const Eigen::Vector3d position{
            attrib.vertices[3 * i + 0], attrib.vertices[3 * i + 1], attrib.vertices[3 * i + 2]};

        const Eigen::Vector3d x = transform * position;
        const Eigen::Vector3d v = Eigen::Vector3d::Zero();
        const double          m = 1.0 / double(attrib.vertices.size());

        auto particle = std::make_shared<elasty::Particle>(x, v, m);

        map_from_obj_vertex_index_to_particle[i] = particle;

        m_particles.push_back(particle);
    }

    for (unsigned int i = 0; i < shape.mesh.indices.size() / 3; ++i)
    {
        const auto p_0 = map_from_obj_vertex_index_to_particle[shape.mesh.indices[i * 3 + 0].vertex_index];
        const auto p_1 = map_from_obj_vertex_index_to_particle[shape.mesh.indices[i * 3 + 1].vertex_index];
        const auto p_2 = map_from_obj_vertex_index_to_particle[shape.mesh.indices[i * 3 + 2].vertex_index];

        if (in_plane_strategy == InPlaneStrategy::EdgeDistance || in_plane_strategy == InPlaneStrategy::Both)
        {
            m_constraints.push_back(std::make_shared<elasty::DistanceConstraint>(
                p_0, p_1, in_plane_stiffness, in_plane_compliance, dt, (p_0->x - p_1->x).norm()));
            m_constraints.push_back(std::make_shared<elasty::DistanceConstraint>(
                p_0, p_2, in_plane_stiffness, in_plane_compliance, dt, (p_0->x - p_2->x).norm()));
            m_constraints.push_back(std::make_shared<elasty::DistanceConstraint>(
                p_1, p_2, in_plane_stiffness, in_plane_compliance, dt, (p_1->x - p_2->x).norm()));
        }

        if (in_plane_strategy == InPlaneStrategy::ContinuumTriangle || in_plane_strategy == InPlaneStrategy::Both)
        {
            m_constraints.push_back(std::make_shared<elasty::ContinuumTriangleConstraint>(
                p_0, p_1, p_2, in_plane_stiffness, in_plane_compliance, dt, 1000.0, 0.10));
        }
    }

    using vertex_t   = unsigned int;
    using triangle_t = unsigned int;
    using edge_t     = std::pair<vertex_t, vertex_t>;
    std::map<edge_t, std::vector<triangle_t>> edges_and_triangles;
    for (unsigned int i = 0; i < shape.mesh.indices.size() / 3; ++i)
    {
        const vertex_t index_0 = shape.mesh.indices[i * 3 + 0].vertex_index;
        const vertex_t index_1 = shape.mesh.indices[i * 3 + 1].vertex_index;
        const vertex_t index_2 = shape.mesh.indices[i * 3 + 2].vertex_index;

        const edge_t e_01 = std::make_pair(std::min(index_0, index_1), std::max(index_0, index_1));
        const edge_t e_02 = std::make_pair(std::min(index_0, index_2), std::max(index_0, index_2));
        const edge_t e_12 = std::make_pair(std::min(index_1, index_2), std::max(index_1, index_2));

        auto register_edge = [&](const edge_t& edge)
        {
            if (edges_and_triangles.find(edge) == edges_and_triangles.end())
            {
                edges_and_triangles[edge] = {i};
            }
            else
            {
                edges_and_triangles[edge].push_back(i);
            }
        };

        register_edge(e_01);
        register_edge(e_02);
        register_edge(e_12);
    }

    for (const auto& key_value : edges_and_triangles)
    {
        const edge_t&                  edge      = key_value.first;
        const std::vector<triangle_t>& triangles = key_value.second;

        assert(triangles.size() == 1 || triangles.size() == 2);

        // Boundary
        if (triangles.size() == 1)
        {
            continue;
        }

        auto obtain_another_vertex = [&](const triangle_t& triangle, const edge_t& edge)
        {
            const vertex_t vertex_0 = shape.mesh.indices[3 * triangle + 0].vertex_index;
            const vertex_t vertex_1 = shape.mesh.indices[3 * triangle + 1].vertex_index;
            const vertex_t vertex_2 = shape.mesh.indices[3 * triangle + 2].vertex_index;

            if (vertex_0 != edge.first && vertex_0 != edge.second)
            {
                return vertex_0;
            }
            else if (vertex_1 != edge.first && vertex_1 != edge.second)
            {
                return vertex_1;
            }
            else
            {
                assert(vertex_2 != edge.first && vertex_2 != edge.second);
                return vertex_2;
            }
        };

        const triangle_t another_vertex_0 = obtain_another_vertex(triangles[0], edge);
        const triangle_t another_vertex_1 = obtain_another_vertex(triangles[1], edge);

        switch (out_of_plane_strategy)
        {
            case OutOfPlaneStrategy::Bending:
            {
                const auto p_0 = map_from_obj_vertex_index_to_particle[edge.first];
                const auto p_1 = map_from_obj_vertex_index_to_particle[edge.second];
                const auto p_2 = map_from_obj_vertex_index_to_particle[another_vertex_0];
                const auto p_3 = map_from_obj_vertex_index_to_particle[another_vertex_1];

                const Eigen::Vector3d& x_0 = p_0->x;
                const Eigen::Vector3d& x_1 = p_1->x;
                const Eigen::Vector3d& x_2 = p_2->x;
                const Eigen::Vector3d& x_3 = p_3->x;

                const Eigen::Vector3d p_10 = x_1 - x_0;
                const Eigen::Vector3d p_20 = x_2 - x_0;
                const Eigen::Vector3d p_30 = x_3 - x_0;

                const Eigen::Vector3d n_0 = p_10.cross(p_20).normalized();
                const Eigen::Vector3d n_1 = p_10.cross(p_30).normalized();

                assert(!n_0.hasNaN());
                assert(!n_1.hasNaN());

                // Typical value is 0.0 or pi
                const double dihedral_angle = std::acos(std::clamp(n_0.dot(n_1), -1.0, 1.0));

                assert(!std::isnan(dihedral_angle));

                m_constraints.push_back(std::make_shared<elasty::BendingConstraint>(
                    p_0, p_1, p_2, p_3, out_of_plane_stiffness, out_of_plane_compliance, dt, dihedral_angle));

                break;
            }
            case OutOfPlaneStrategy::IsometricBending:
            {
                const auto p_0 = map_from_obj_vertex_index_to_particle[edge.first];
                const auto p_1 = map_from_obj_vertex_index_to_particle[edge.second];
                const auto p_2 = map_from_obj_vertex_index_to_particle[another_vertex_0];
                const auto p_3 = map_from_obj_vertex_index_to_particle[another_vertex_1];

                m_constraints.push_back(std::make_shared<elasty::IsometricBendingConstraint>(
                    p_0, p_1, p_2, p_3, out_of_plane_stiffness, out_of_plane_compliance, dt));

                break;
            }
            case OutOfPlaneStrategy::Cross:
            {
                const auto p_2 = map_from_obj_vertex_index_to_particle[another_vertex_0];
                const auto p_3 = map_from_obj_vertex_index_to_particle[another_vertex_1];

                const Eigen::Vector3d& x_2 = p_2->x;
                const Eigen::Vector3d& x_3 = p_3->x;

                m_constraints.push_back(std::make_shared<elasty::DistanceConstraint>(
                    p_2, p_3, out_of_plane_stiffness, out_of_plane_compliance, dt, (x_2 - x_3).norm()));

                break;
            }
        }
    }

    calculateAreas();
}

void elasty::ClothSimObject::applyAerodynamicForces(const Eigen::Vector3d& global_velocity,
                                                    const double           drag_coeff,
                                                    const double           lift_coeff)
{
    constexpr double rho = 1.225; // Taken from Wikipedia: https://en.wikipedia.org/wiki/Density_of_air

    assert(drag_coeff >= lift_coeff);

    const int num_triangles = m_triangle_list.rows();

    for (int i = 0; i < num_triangles; ++i)
    {
        const auto& x_0 = m_particles[m_triangle_list.row(i)(0)]->x;
        const auto& x_1 = m_particles[m_triangle_list.row(i)(1)]->x;
        const auto& x_2 = m_particles[m_triangle_list.row(i)(2)]->x;

        const auto& v_0 = m_particles[m_triangle_list.row(i)(0)]->v;
        const auto& v_1 = m_particles[m_triangle_list.row(i)(1)]->v;
        const auto& v_2 = m_particles[m_triangle_list.row(i)(2)]->v;

        const auto& m_0 = m_particles[m_triangle_list.row(i)(0)]->m;
        const auto& m_1 = m_particles[m_triangle_list.row(i)(1)]->m;
        const auto& m_2 = m_particles[m_triangle_list.row(i)(2)]->m;

        const double m_sum = m_0 + m_1 + m_2;

        // Calculate the weighted average of the particle velocities
        const Eigen::Vector3d v_triangle = (m_0 * v_0 + m_1 * v_1 + m_2 * v_2) / m_sum;

        // Calculate the relative velocity of the triangle
        const Eigen::Vector3d v_rel         = v_triangle - global_velocity;
        const double          v_rel_squared = v_rel.squaredNorm();

        const auto            cross         = (x_1 - x_0).cross(x_2 - x_0);
        const double          area          = 0.5 * cross.norm();
        const auto            n_either_side = cross.normalized();
        const Eigen::Vector3d n             = (n_either_side.dot(v_rel) > 0.0) ? n_either_side : -n_either_side;

        const double coeff = 0.5 * rho * area;

        // Note: This wind force model was proposed by [Wilson+14]
        const Eigen::Vector3d f =
            -coeff * ((drag_coeff - lift_coeff) * v_rel.dot(n) * v_rel + lift_coeff * v_rel_squared * n);

        m_particles[m_triangle_list.row(i)(0)]->f += (m_0 / m_sum) * f;
        m_particles[m_triangle_list.row(i)(1)]->f += (m_1 / m_sum) * f;
        m_particles[m_triangle_list.row(i)(2)]->f += (m_2 / m_sum) * f;
    }
}

void elasty::ClothSimObject::calculateAreas()
{
    const int num_triangles = m_triangle_list.rows();

    m_area_list = Eigen::VectorXd(num_triangles);

    for (int i = 0; i < num_triangles; ++i)
    {
        const auto& x_0 = m_particles[m_triangle_list.row(i)[0]]->x;
        const auto& x_1 = m_particles[m_triangle_list.row(i)[1]]->x;
        const auto& x_2 = m_particles[m_triangle_list.row(i)[2]]->x;

        const double area = 0.5 * (x_1 - x_0).cross(x_2 - x_0).norm();

        m_area_list(i) = area;
    }
}


/*    Copyright (c) 2010-2018, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 */
#include <map>

#include <functional>

#include <Eigen/Core>

#include "Tudat/Astrodynamics/Propagators/variationalEquations.h"
#include "Tudat/Astrodynamics/Propagators/rotationalMotionQuaternionsStateDerivative.h"

#include "Tudat/Astrodynamics/OrbitDetermination/AccelerationPartials/accelerationPartial.h"

namespace tudat
{

namespace propagators
{

template< typename StateScalarType >
void VariationalEquations::getBodyInitialStatePartialMatrix(
        const Eigen::Matrix< StateScalarType, Eigen::Dynamic, Eigen::Dynamic >& stateTransitionAndSensitivityMatrices,
        Eigen::Block< Eigen::Matrix< StateScalarType, Eigen::Dynamic, Eigen::Dynamic > > currentMatrixDerivative )
{
    setBodyStatePartialMatrix( );

    // Add partials of body positions and velocities.
    currentMatrixDerivative.block( 0, 0, totalDynamicalStateSize_, numberOfParameterValues_ ) =
            ( variationalMatrix_.template cast< StateScalarType >( ) * stateTransitionAndSensitivityMatrices );
}

//! Calculates matrix containing partial derivatives of state derivatives w.r.t. body state.
void VariationalEquations::setBodyStatePartialMatrix( )
{
    // Initialize partial matrix
    variationalMatrix_.setZero( );

    if( dynamicalStatesToEstimate_.count( propagators::translational_state ) > 0 )
    {
        int startIndex = stateTypeStartIndices_.at( propagators::translational_state );
        for( unsigned int i = 0; i < dynamicalStatesToEstimate_.at( propagators::translational_state ).size( ); i++ )
        {
            variationalMatrix_.block( startIndex + i * 6, startIndex + i * 6 + 3, 3, 3 ).setIdentity( );
        }
    }

    if( dynamicalStatesToEstimate_.count( propagators::rotational_state ) > 0 )
    {
         Eigen::VectorXd rotationalStates = currentStatesPerTypeInConventionalRepresentation_.at(
                     propagators::rotational_state );

        int startIndex = stateTypeStartIndices_.at( propagators::rotational_state );
        for( unsigned int i = 0; i < dynamicalStatesToEstimate_.at( propagators::rotational_state ).size( ); i++ )
        {
            variationalMatrix_.block( startIndex + i * 7, startIndex + i * 7 , 4, 4 ) =
                    getQuaterionToQuaternionRateMatrix( rotationalStates.segment( 7 * i + 4, 3 ) );
            variationalMatrix_.block( startIndex + i * 7, startIndex + i * 7 + 4, 4, 3 ) =
                    getAngularVelocityToQuaternionRateMatrix( rotationalStates.segment( 7 * i, 4 ) );
        }
    }

    // Iterate over all bodies undergoing accelerations for which initial condition is to be estimated.
    for( std::map< IntegratedStateType, std::vector< std::multimap< std::pair< int, int >,
         std::function< void( Eigen::Block< Eigen::MatrixXd > ) > > > >::iterator
         typeIterator = statePartialList_.begin( ); typeIterator != statePartialList_.end( ); typeIterator++ )
    {
        int startIndex = stateTypeStartIndices_.at( typeIterator->first );
        int currentStateSize = getSingleIntegrationSize( typeIterator->first );
        int entriesToSkipPerEntry = currentStateSize - getGeneralizedAccelerationSize( typeIterator->first );

        for( unsigned int i = 0; i < typeIterator->second.size( ); i++ )
        {
            // Iterate over all bodies exerting an acceleration on this body.
            for( statePartialIterator_ = typeIterator->second.at( i ).begin( );
                 statePartialIterator_ != typeIterator->second.at( i ).end( );
                 statePartialIterator_++ )
            {
                statePartialIterator_->second(
                            variationalMatrix_.block(
                                startIndex + entriesToSkipPerEntry + i* currentStateSize, statePartialIterator_->first.first,
                                currentStateSize - entriesToSkipPerEntry, statePartialIterator_->first.second ) );

            }
        }
    }

   for( unsigned int i = 0; i < statePartialAdditionIndices_.size( ); i++ )
   {
       variationalMatrix_.block( 0, statePartialAdditionIndices_.at( i ).second, totalDynamicalStateSize_, 3 ) +=
               variationalMatrix_.block( 0, statePartialAdditionIndices_.at( i ).first, totalDynamicalStateSize_, 3 );
   }

   for( unsigned int i = 0; i < inertiaTensorsForMultiplication_.size( ); i++ )
   {
       variationalMatrix_.block( inertiaTensorsForMultiplication_.at( i ).first, 0, 3, totalDynamicalStateSize_ ) =
               ( inertiaTensorsForMultiplication_.at( i ).second( ).inverse( ) ) *
               variationalMatrix_.block( inertiaTensorsForMultiplication_.at( i ).first, 0, 3, totalDynamicalStateSize_ ).eval( );
   }

}

//! Function to clear reference/cached values of state derivative partials.
void VariationalEquations::clearPartials( )
{
    for( stateDerivativeTypeIterator_ = stateDerivativePartialList_.begin( );
         stateDerivativeTypeIterator_ != stateDerivativePartialList_.end( );
         stateDerivativeTypeIterator_++ )
    {
        for( unsigned int i = 0; i < stateDerivativeTypeIterator_->second.size( ); i++ )
        {
            for( unsigned int j = 0; j < stateDerivativeTypeIterator_->second.at( i ).size( ); j++ )
            {
                stateDerivativeTypeIterator_->second.at( i ).at( j )->resetTime( TUDAT_NAN );
            }

        }
    }
}

//! Function (called by constructor) to set up the statePartialList_ member from the state derivative partials
void VariationalEquations::setStatePartialFunctionList( )
{
    std::pair< std::function< void( Eigen::Block< Eigen::MatrixXd > ) >, int > currentDerivativeFunction;

    // Iterate over all state types
    for( std::map< propagators::IntegratedStateType,
         orbit_determination::StateDerivativePartialsMap >::iterator
         stateDerivativeTypeIterator_ = stateDerivativePartialList_.begin( );
         stateDerivativeTypeIterator_ != stateDerivativePartialList_.end( );
         stateDerivativeTypeIterator_++ )
    {
        // Iterate over all bodies undergoing 'accelerations' for which initial state is to be estimated.
        for( unsigned int i = 0; i < stateDerivativeTypeIterator_->second.size( ); i++ )
        {
            std::multimap< std::pair< int, int >, std::function< void( Eigen::Block< Eigen::MatrixXd > ) > >
                    currentBodyPartialList;

            // Iterate over all 'accelerations' from single body on other single body
            for( unsigned int j = 0; j < stateDerivativeTypeIterator_->second.at( i ).size( ); j++ )
            {
                for( std::map< propagators::IntegratedStateType,
                     std::vector< std::pair< std::string, std::string > > >::iterator
                     estimatedStateIterator = dynamicalStatesToEstimate_.begin( );
                     estimatedStateIterator != dynamicalStatesToEstimate_.end( );
                     estimatedStateIterator++ )
                {
                    // Iterate over all bodies to see if body exerting acceleration is also to be estimated (cross-terms)
                    for( unsigned int k = 0; k < estimatedStateIterator->second.size( ); k++ )
                    {
                        currentDerivativeFunction = stateDerivativeTypeIterator_->second.at( i ).at( j )->
                                getDerivativeFunctionWrtStateOfIntegratedBody(
                                    estimatedStateIterator->second.at( k ), estimatedStateIterator->first );

                        // If function is not-empty: add to list.
                        if( currentDerivativeFunction.second != 0 )
                        {
                            currentBodyPartialList.insert(
                                        std::make_pair(
                                            std::make_pair( k * getSingleIntegrationSize( estimatedStateIterator->first ) +
                                                            stateTypeStartIndices_.at( estimatedStateIterator->first ),
                                                            getSingleIntegrationSize( estimatedStateIterator->first ) ),
                                            currentDerivativeFunction.first ) );
                        }
                    }
                }
            }
            statePartialList_[ stateDerivativeTypeIterator_->first ].push_back( currentBodyPartialList );
        }
    }
}

template void VariationalEquations::getBodyInitialStatePartialMatrix< double >(
        const Eigen::Matrix< double, Eigen::Dynamic, Eigen::Dynamic >& stateTransitionAndSensitivityMatrices,
        Eigen::Block< Eigen::Matrix< double, Eigen::Dynamic, Eigen::Dynamic > > currentMatrixDerivative );

//#if( BUILD_EXTENDED_PRECISION_PROPAGATION_TOOLS )
template void VariationalEquations::getBodyInitialStatePartialMatrix< long double >(
        const Eigen::Matrix< long double, Eigen::Dynamic, Eigen::Dynamic >& stateTransitionAndSensitivityMatrices,
        Eigen::Block< Eigen::Matrix< long double, Eigen::Dynamic, Eigen::Dynamic > > currentMatrixDerivative );
//#endif

} // namespace propagators

} // namespace tudat


//---------------------------------Spheral++----------------------------------//
// ProbabilisticDamageModel
// A damage model based on Weibull statistics that uses volume based
// probabilities per node to decide when damage starts to accrue.  Should
// generate similar results to the classic Benz-Asphaug (Grady-Kipp) model
// without generating explicit flaws.  Also appropriate for use with varying
// resolution materials.
//
// Created by JMO, Tue Apr 13 15:58:08 PDT 2021
//----------------------------------------------------------------------------//
#include "FileIO/FileIO.hh"
#include "ProbabilisticDamageModel.hh"
#include "TensorStrainPolicy.hh"
#include "ProbabilisticDamagePolicy.hh"
#include "YoungsModulusPolicy.hh"
#include "LongitudinalSoundSpeedPolicy.hh"
#include "DamageGradientPolicy.hh"
#include "Strength/SolidFieldNames.hh"
#include "NodeList/SolidNodeList.hh"
#include "DataBase/DataBase.hh"
#include "DataBase/State.hh"
#include "DataBase/StateDerivatives.hh"
#include "DataBase/ReplaceState.hh"
#include "Hydro/HydroFieldNames.hh"
#include "Field/FieldList.hh"
#include "Boundary/Boundary.hh"
#include "Neighbor/Neighbor.hh"
#include "Utilities/mortonOrderIndices.hh"
#include "Utilities/allReduce.hh"
#include "Utilities/uniform_random.hh"

#include <boost/functional/hash.hpp>  // hash_combine

#include <string>
#include <vector>
#include <algorithm>
#include <limits>
using std::vector;
using std::string;
using std::pair;
using std::make_pair;
using std::cout;
using std::cerr;
using std::endl;
using std::min;
using std::max;
using std::abs;

namespace Spheral {

//------------------------------------------------------------------------------
// Constructor.
//------------------------------------------------------------------------------
template<typename Dimension>
ProbabilisticDamageModel<Dimension>::
ProbabilisticDamageModel(SolidNodeList<Dimension>& nodeList,
                         const TableKernel<Dimension>& W,
                         const double kWeibull,
                         const double mWeibull,
                         const size_t seed,
                         const size_t minFlawsPerNode,
                         const double crackGrowthMultiplier,
                         const double volumeMultiplier,
                         const DamageCouplingAlgorithm damageCouplingAlgorithm,
                         const TensorStrainAlgorithm strainAlgorithm,
                         const bool damageInCompression,
                         const double criticalDamageThreshold,
                         const Field<Dimension, int>& mask):
  DamageModel<Dimension>(nodeList, W, crackGrowthMultiplier, damageCouplingAlgorithm),
  mStrainAlgorithm(strainAlgorithm),
  mDamageInCompression(damageInCompression),
  mkWeibull(kWeibull),
  mmWeibull(mWeibull),
  mVolumeMultiplier(volumeMultiplier),
  mVmin(std::numeric_limits<double>::max()),
  mVmax(std::numeric_limits<double>::min()),
  mCriticalDamageThreshold(criticalDamageThreshold),
  mSeed(seed),
  mMinFlawsPerNode(minFlawsPerNode),
  mNumFlaws(SolidFieldNames::numFlaws, nodeList),
  mMask(mask),
  mMinFlaw(SolidFieldNames::minFlaw, nodeList),
  mMaxFlaw(SolidFieldNames::maxFlaw, nodeList),
  mInitialVolume(SolidFieldNames::initialVolume, nodeList),
  mYoungsModulus(SolidFieldNames::YoungsModulus, nodeList),
  mLongitudinalSoundSpeed(SolidFieldNames::longitudinalSoundSpeed, nodeList),
  mDdamageDt(ProbabilisticDamagePolicy<Dimension>::prefix() + SolidFieldNames::scalarDamage, nodeList),
  mStrain(SolidFieldNames::strainTensor, nodeList),
  mEffectiveStrain(SolidFieldNames::effectiveStrainTensor, nodeList) {
}

//------------------------------------------------------------------------------
// Destructor.
//------------------------------------------------------------------------------
template<typename Dimension>
ProbabilisticDamageModel<Dimension>::
~ProbabilisticDamageModel() {
}

//------------------------------------------------------------------------------
// initializeProblemStartup
//
// After all initial state has been initialize (node positions, masses, etc),
// but before we try to run any physics cycles.  This is when we initialize a
// lot of our state for the damage model.
//------------------------------------------------------------------------------
template<typename Dimension>
void
ProbabilisticDamageModel<Dimension>::
initializeProblemStartup(DataBase<Dimension>& dataBase) {

  // How many points are actually being damaged?
  // We have to be careful to use an unsigned size_t here due to overflow
  // problems with large numbers of points.
  size_t nused_local = 0u;
  for (auto i = 0u; i < mMask.numInternalElements(); ++i) {
    if (mMask[i] == 1) ++nused_local;
  }
  const size_t nused_global = allReduce(nused_local, MPI_SUM, Communicator::communicator());

  // Compute the Morton-ordering for hashing with the global seed to seed each
  // point-wise random number generator.
  typedef KeyTraits::Key Key;
  const FieldList<Dimension, Key> keyList = mortonOrderIndices(dataBase);
  const auto& nodes = this->nodeList();
  CHECK(keyList.fieldForNodeList(nodes) < keyList.end());
  const Field<Dimension, Key>& keys = **(keyList.fieldForNodeList(nodes));

  // Compute the initial volumes and random seeds for each node.  We hash the
  // morton index of each point with the global seed value to create unique
  // but reproducible seeds for each points random number generator.
  const auto& mass = nodes.mass();
  const auto& rho = nodes.massDensity();
  const auto  nlocal = nodes.numInternalNodes();
  vector<uniform_random> randomGenerators(nlocal);
#pragma omp parallel for
  for (auto i = 0u; i < nlocal; ++i) {
    if (mMask(i) == 1) {
      CHECK(mass(i) > 0.0 and rho(i) > 0.0);
      mInitialVolume(i) = mass(i)/rho(i) * mVolumeMultiplier;
      mVmin = std::min(mVmin, mInitialVolume(i));
      mVmax = std::max(mVmax, mInitialVolume(i));
      Key seedi = mSeed;
      boost::hash_combine(seedi, keys(i));
      randomGenerators[i].seed(seedi);      // starting out generating in [0,1)
      randomGenerators[i]();                // Recommended to discard first value in sequence
    }
  }
  mVmin = allReduce(mVmin, MPI_MIN, Communicator::communicator());
  mVmax = allReduce(mVmax, MPI_MAX, Communicator::communicator());

  // Generate min/max ranges of flaws for each point.
  const auto mInv = 1.0/mmWeibull;
  size_t minNumFlaws = std::numeric_limits<size_t>::max();
  size_t maxNumFlaws = 0u;
  size_t totalNumFlaws = 0u;
  auto epsMin = std::numeric_limits<double>::max();
  auto epsMax = std::numeric_limits<double>::min();
  auto numFlawsRatio = 0.0;
#pragma omp parallel for
  for (auto i = 0u; i < nlocal; ++i) {
    if (mMask(i) == 1) {
      const auto Nflaws = size_t(mMinFlawsPerNode*mInitialVolume(i)/mVmin + 0.5);  // Target number of flaws on this point
      CHECK(Nflaws >= mMinFlawsPerNode);
      const auto Ai = pow(Nflaws/(mkWeibull*mInitialVolume(i)), mInv);
      mMinFlaw(i) = 1.0;
      mMaxFlaw(i) = 0.0;
      auto ntries = 0u;
      while (mMinFlaw(i) > mMaxFlaw(i) and ntries++ < 100u) {
        mMinFlaw(i) = Ai*pow(1.0 - pow(randomGenerators[i](), 1.0/Nflaws), mInv);
        mMaxFlaw(i) = Ai*pow(randomGenerators[i](), 1.0/(mmWeibull*Nflaws));
      }
      CHECK(mMaxFlaw(i) > mMinFlaw(i));
      mNumFlaws(i) = std::max(1, 1 + int(mInitialVolume(i)*mkWeibull*(pow(mMaxFlaw(i), mmWeibull) - pow(mMinFlaw(i), mmWeibull))));

      // Gather statistics
#pragma omp critical
      {
        minNumFlaws = min(minNumFlaws, size_t(mNumFlaws(i)));
        maxNumFlaws = max(maxNumFlaws, size_t(mNumFlaws(i)));
        totalNumFlaws += size_t(mNumFlaws(i));
        epsMin = std::min(epsMin, mMinFlaw(i));
        epsMax = std::max(epsMax, mMaxFlaw(i));
        numFlawsRatio += double(mNumFlaws(i))/Nflaws;
      }
    }
  }

  // Some diagnostic output.
  if (nused_global > 0) {
    minNumFlaws = allReduce(minNumFlaws, MPI_MIN, Communicator::communicator());
    maxNumFlaws = allReduce(maxNumFlaws, MPI_MAX, Communicator::communicator());
    totalNumFlaws = allReduce(totalNumFlaws, MPI_SUM, Communicator::communicator());
    epsMin = allReduce(epsMin, MPI_MIN, Communicator::communicator());
    epsMax = allReduce(epsMax, MPI_MAX, Communicator::communicator());
    numFlawsRatio = allReduce(numFlawsRatio, MPI_SUM, Communicator::communicator())/nused_global;
    if (Process::getRank() == 0) {
      cerr << "ProbabilisticDamageModel for " << nodes.name() << ":" << endl
           << " Min, max, max/min volumes: " << mVmin << " " << mVmax << " " << mVmax*safeInv(mVmin) << endl
           << "    Min num flaws per node: " << minNumFlaws << endl
           << "    Max num flaws per node: " << maxNumFlaws << endl
           << "    Total num flaws       : " << totalNumFlaws << endl
           << "    Avg flaws per node    : " << totalNumFlaws/nused_global << endl
           << "    Min flaw strain       : " << epsMin << endl
           << "    Max flaw strain       : " << epsMax << endl
           << "    Avg Neff/Nflaws       : " << numFlawsRatio << endl;
    }
  }
}

//------------------------------------------------------------------------------
// Evaluate derivatives.
//
// In this model we compute the scalar damage derivative assuming unresolved
// crack growth for every point. However, that is not applied in the tensor
// damage update policy unless the flaws are actually activated.
//------------------------------------------------------------------------------
template<typename Dimension>
void
ProbabilisticDamageModel<Dimension>::
evaluateDerivatives(const Scalar time,
                    const Scalar dt,
                    const DataBase<Dimension>& dataBase,
                    const State<Dimension>& state,
                    StateDerivatives<Dimension>& derivs) const {

  // Set the scalar magnitude of the damage evolution.
  const auto* nodeListPtr = &(this->nodeList());
  auto&       DDDt = derivs.field(state.buildFieldKey(ProbabilisticDamagePolicy<Dimension>::prefix() + SolidFieldNames::scalarDamage, nodeListPtr->name()), 0.0);
  this->computeScalarDDDt(dataBase,
                          state,
                          time,
                          dt,
                          DDDt);
}

//------------------------------------------------------------------------------
// Vote on a time step.
//------------------------------------------------------------------------------
template<typename Dimension>
typename ProbabilisticDamageModel<Dimension>::TimeStepType
ProbabilisticDamageModel<Dimension>::
dt(const DataBase<Dimension>& /*dataBase*/,
   const State<Dimension>& /*state*/,
   const StateDerivatives<Dimension>& /*derivs*/,
   const Scalar /*currentTime*/) const {

  // // Look at how quickly we're trying to change the damage.
  // double dt = DBL_MAX;
  // const Field<Dimension, SymTensor>& damage = this->nodeList().damage();
  // const ConnectivityMap<Dimension>& connectivityMap = dataBase.connectivityMap();
  // const vector<const NodeList<Dimension>*>& nodeLists = connectivityMap.nodeLists();
  // const size_t nodeListi = distance(nodeLists.begin(), find(nodeLists.begin(), nodeLists.end(), &(this->nodeList())));
  // for (typename ConnectivityMap<Dimension>::const_iterator iItr = connectivityMap.begin(nodeListi);
  //      iItr != connectivityMap.end(nodeListi);
  //      ++iItr) {
  //   const int i = *iItr;
  //   const double D0 = damage(i).Trace() / Dimension::nDim;
  //   dt = min(dt, 0.8*max(D0, 1.0 - D0)/
  //            std::sqrt(mDdamageDt(i)*mDdamageDt(i) + 1.0e-20));
  // }
  // return TimeStepType(dt, "Rate of damage change");

  return TimeStepType(1.0e100, "Rate of damage change -- NO VOTE.");
}

//------------------------------------------------------------------------------
// Register our state.
//------------------------------------------------------------------------------
template<typename Dimension>
void
ProbabilisticDamageModel<Dimension>::
registerState(DataBase<Dimension>& dataBase,
              State<Dimension>& state) {

  typedef typename State<Dimension>::PolicyPointer PolicyPointer;

  // Register Youngs modulus and the longitudinal sound speed.
  PolicyPointer EPolicy(new YoungsModulusPolicy<Dimension>());
  PolicyPointer clPolicy(new LongitudinalSoundSpeedPolicy<Dimension>());
  state.enroll(mYoungsModulus, EPolicy);
  state.enroll(mLongitudinalSoundSpeed, clPolicy);

  // Set the initial values for the Youngs modulus, sound speed, and pressure.
  typename StateDerivatives<Dimension>::PackageList dummyPackages;
  StateDerivatives<Dimension> derivs(dataBase, dummyPackages);
  EPolicy->update(state.key(mYoungsModulus), state, derivs, 1.0, 0.0, 0.0);
  clPolicy->update(state.key(mLongitudinalSoundSpeed), state, derivs, 1.0, 0.0, 0.0);

  // Register the strain and effective strain.
  PolicyPointer effectiveStrainPolicy(new TensorStrainPolicy<Dimension>(mStrainAlgorithm));
  state.enroll(mStrain);
  state.enroll(mEffectiveStrain, effectiveStrainPolicy);

  // Register the damage and state it requires.
  // Note we are overriding the default no-op policy for the damage
  // as originally registered by the SolidSPHHydroBase class.
  auto& damage = this->nodeList().damage();
  PolicyPointer damagePolicy(new ProbabilisticDamagePolicy<Dimension>(mDamageInCompression,
                                                                      mkWeibull,
                                                                      mmWeibull));
  state.enroll(damage, damagePolicy);
  state.enroll(mNumFlaws);
  state.enroll(mMinFlaw);
  state.enroll(mMaxFlaw);
  state.enroll(mInitialVolume);

  // Mask out nodes beyond the critical damage threshold from setting the timestep.
  auto maskKey = state.buildFieldKey(HydroFieldNames::timeStepMask, this->nodeList().name());
  auto& mask = state.field(maskKey, 0);
  const auto nlocal = this->nodeList().numInternalNodes();
#pragma omp parallel for
  for (auto i = 0u; i < nlocal; ++i) {
    if (damage(i).Trace() > mCriticalDamageThreshold) mask(i) = 0;
  }
}

//------------------------------------------------------------------------------
// Register the derivatives.
//------------------------------------------------------------------------------
template<typename Dimension>
void
ProbabilisticDamageModel<Dimension>::
registerDerivatives(DataBase<Dimension>& /*dataBase*/,
                    StateDerivatives<Dimension>& derivs) {
  derivs.enroll(mDdamageDt);
}

//------------------------------------------------------------------------------
// Apply the boundary conditions to the ghost nodes.
//------------------------------------------------------------------------------
template<typename Dimension>
void
ProbabilisticDamageModel<Dimension>::
applyGhostBoundaries(State<Dimension>& state,
                     StateDerivatives<Dimension>& /*derivs*/) {

  // Grab this models damage field from the state.
  typedef typename State<Dimension>::KeyType Key;
  const Key nodeListName = this->nodeList().name();
  const Key DKey = state.buildFieldKey(SolidFieldNames::tensorDamage, nodeListName);
  CHECK(state.registered(DKey));
  auto& D = state.field(DKey, SymTensor::zero);

  // Apply ghost boundaries to the damage.
  for (auto boundaryItr = this->boundaryBegin();
       boundaryItr < this->boundaryEnd();
       ++boundaryItr) {
    (*boundaryItr)->applyGhostBoundary(D);
  }
}

//------------------------------------------------------------------------------
// Enforce boundary conditions for the physics specific fields.
//------------------------------------------------------------------------------
template<typename Dimension>
void
ProbabilisticDamageModel<Dimension>::
enforceBoundaries(State<Dimension>& state,
                  StateDerivatives<Dimension>& /*derivs*/) {

  // Grab this models damage field from the state.
  typedef typename State<Dimension>::KeyType Key;
  const Key nodeListName = this->nodeList().name();
  const Key DKey = state.buildFieldKey(SolidFieldNames::tensorDamage, nodeListName);
  CHECK(state.registered(DKey));
  auto& D = state.field(DKey, SymTensor::zero);

  // Enforce!
  for (auto boundaryItr = this->boundaryBegin();
       boundaryItr < this->boundaryEnd();
       ++boundaryItr) {
    (*boundaryItr)->enforceBoundary(D);
  }
}

//------------------------------------------------------------------------------
// Dump the current state to the given file.
//------------------------------------------------------------------------------
template<typename Dimension>
void
ProbabilisticDamageModel<Dimension>::
dumpState(FileIO& file, const string& pathName) const {
  DamageModel<Dimension>::dumpState(file, pathName);
  file.write(mNumFlaws, pathName + "/numFlaws");
  file.write(mMinFlaw, pathName + "/minFlaw");
  file.write(mMaxFlaw, pathName + "/maxFlaw");
  file.write(mStrain, pathName + "/strain");
  file.write(mEffectiveStrain, pathName + "/effectiveStrain");
  file.write(mDdamageDt, pathName + "/DdamageDt");
  file.write(mMask, pathName + "/mask");
}

//------------------------------------------------------------------------------
// Restore the state from the given file.
//------------------------------------------------------------------------------
template<typename Dimension>
void
ProbabilisticDamageModel<Dimension>::
restoreState(const FileIO& file, const string& pathName) {
  DamageModel<Dimension>::restoreState(file, pathName);
  file.read(mNumFlaws, pathName + "/numFlaws");
  file.read(mMinFlaw, pathName + "/minFlaw");
  file.read(mMaxFlaw, pathName + "/maxFlaw");
  file.read(mStrain, pathName + "/strain");
  file.read(mEffectiveStrain, pathName + "/effectiveStrain");
  file.read(mDdamageDt, pathName + "/DdamageDt");
  file.read(mMask, pathName + "/mask");
}

}