#include "beamformfinder.hpp"
#include "polyscope/curve_network.h"
#include "polyscope/histogram.h"
#include "polyscope/polyscope.h"
#include "simulationhistoryplotter.hpp"
#include <algorithm>
#include <chrono>
#include <execution>
#include <omp.h>

void FormFinder::runUnitTests() {
  const std::filesystem::path groundOfTruthFolder{
      "/home/iason/Coding/Libraries/MySources/formFinder_unitTestFiles"};

  FormFinder formFinder;

  // First example of the paper
  VCGEdgeMesh beam;
  //  const size_t spanGridSize = 11;
  //  mesh.createSpanGrid(spanGridSize);
  beam.loadPly(
      std::filesystem::path(groundOfTruthFolder).append("simpleBeam.ply").string());
  std::unordered_map<VertexIndex, std::unordered_set<DoFType>> fixedVertices;
  fixedVertices[0] = std::unordered_set<DoFType>{0, 1, 2, 3};
  fixedVertices[beam.VN() - 1] = std::unordered_set<DoFType>{1, 2};
  std::unordered_map<VertexIndex, Vector6d> nodalForces{
      {beam.VN() / 2, Vector6d({0, 1000, 1000, 0, 0, 0})}};
  // Forced displacements
  std::unordered_map<size_t, Eigen::Vector3d> nodalForcedDisplacements;
  nodalForcedDisplacements.insert({beam.VN() - 1, Eigen::Vector3d(-0.2, 0, 0)});

  SimulationJob beamSimulationJob{
      std::make_shared<SimulationMesh>(beam), "First paper example",
      //      SimulationJob::constructFixedVerticesSpanGrid(spanGridSize,
      //      mesh.VN()),
      fixedVertices, nodalForces, nodalForcedDisplacements};
  beamSimulationJob.pMesh->setBeamMaterial(0.3, 200 * 1e9);
  assert(CrossSectionType::name == CylindricalBeamDimensions::name);

  beamSimulationJob.pMesh->setBeamCrossSection(CrossSectionType{0.03, 0.026});
  Settings settings;
  settings.Dtini = 0.1;
  settings.xi = 0.9969;
  settings.maxDRMIterations = 0;
  settings.totalResidualForcesNormThreshold = 1;
  //  settings.shouldDraw = true;
  settings.beVerbose = true;
  SimulationResults simpleBeam_simulationResults = formFinder.executeSimulation(
      std::make_shared<SimulationJob>(beamSimulationJob), settings);
  simpleBeam_simulationResults.save();
  const std::string simpleBeamGroundOfTruthBinaryFilename =
      std::filesystem::path(groundOfTruthFolder)
          .append("SimpleBeam_displacements.eigenBin").string();
  assert(std::filesystem::exists(
      std::filesystem::path(simpleBeamGroundOfTruthBinaryFilename)));
  Eigen::MatrixXd simpleBeam_groundOfTruthDisplacements;
  Eigen::read_binary(simpleBeamGroundOfTruthBinaryFilename,
                     simpleBeam_groundOfTruthDisplacements);
  if (!simpleBeam_simulationResults.isEqual(
          simpleBeam_groundOfTruthDisplacements)) {
    std::cerr << "Error:Beam simulation produces wrong results." << std::endl;
    //    return;
  }

  // Second example of the paper
  VCGEdgeMesh shortSpanGrid;
  //  const size_t spanGridSize = 11;
  //  mesh.createSpanGrid(spanGridSize);
  shortSpanGrid.loadPly(std::filesystem::path(groundOfTruthFolder)
                            .append("shortSpanGridshell.ply")
                            .string());

  fixedVertices.clear();
  //// Corner nodes
  fixedVertices[0] = std::unordered_set<DoFType>{2};
  fixedVertices[4] = std::unordered_set<DoFType>{2};
  fixedVertices[16] = std::unordered_set<DoFType>{2};
  fixedVertices[20] = std::unordered_set<DoFType>{2};
  //// Center node
  fixedVertices[10] = std::unordered_set<DoFType>{0, 1, 3, 4, 5};

  nodalForcedDisplacements.clear();
  nodalForcedDisplacements.insert({{0, Eigen::Vector3d(0.1, 0.1, 0)},
                                   {4, Eigen::Vector3d(-0.1, 0.1, 0)},
                                   {16, Eigen::Vector3d(0.1, -0.1, 0)},
                                   {20, Eigen::Vector3d(-0.1, -0.1, 0)}});

  SimulationJob shortSpanGridshellSimulationJob{
      std::make_shared<SimulationMesh>(shortSpanGrid),
      "Short span gridshell",
      fixedVertices,
      {},
      nodalForcedDisplacements};
  shortSpanGridshellSimulationJob.pMesh->setBeamMaterial(0.3, 200 * 1e9);
  assert(typeid(CrossSectionType) == typeid(CylindricalBeamDimensions));
  shortSpanGridshellSimulationJob.pMesh->setBeamCrossSection(
      CrossSectionType{0.03, 0.026});
  SimulationResults shortSpanGridshellSimulationResults =
      formFinder.executeSimulation(
          std::make_shared<SimulationJob>(shortSpanGridshellSimulationJob),
          settings);
  shortSpanGridshellSimulationResults.save();

  const std::string groundOfTruthBinaryFilename =
      std::filesystem::path(groundOfTruthFolder)
          .append("ShortSpanGridshell_displacements.eigenBin")
          .string();
  assert(std::filesystem::exists(
      std::filesystem::path(groundOfTruthBinaryFilename)));
  Eigen::MatrixXd shortSpanGridshell_groundOfTruthDisplacements;
  Eigen::read_binary(groundOfTruthBinaryFilename,
                     shortSpanGridshell_groundOfTruthDisplacements);
  //  shortSpanGridshellSimulationResults.registerForDrawing(
  //      shortSpanGridshellSimulationJob);
  //  polyscope::show();
  assert(shortSpanGridshellSimulationResults.isEqual(
      shortSpanGridshell_groundOfTruthDisplacements));
  if (!shortSpanGridshellSimulationResults.isEqual(
          shortSpanGridshell_groundOfTruthDisplacements)) {
    std::cerr << "Error:Short span simulation produces wrong results."
              << std::endl;
    return;
  }

  // Third example of the paper
  VCGEdgeMesh longSpanGrid;
  longSpanGrid.loadPly(std::filesystem::path(groundOfTruthFolder)
                           .append("longSpanGridshell.ply")
                           .string());
  const size_t spanGridSize = 11;

  fixedVertices.clear();
  for (size_t vi = 0; vi < spanGridSize - 1; vi++) {
    fixedVertices[vi] = {0, 2};
  }
  for (size_t vi = longSpanGrid.VN() - 1 - (spanGridSize - 2);
       vi < longSpanGrid.VN(); vi++) {
    fixedVertices[vi] = {0, 2};
  }
  for (size_t vi = spanGridSize - 1;
       vi < longSpanGrid.VN() - 1 - (spanGridSize - 2) - spanGridSize + 1;
       vi += spanGridSize + 1) {
    fixedVertices[vi] = {1, 2};
    fixedVertices[vi + spanGridSize] = {1, 2};
  }

  nodalForcedDisplacements.clear();
  const size_t horizontalOffset = std::floor((spanGridSize - 2) / 2);
  nodalForcedDisplacements.insert(
      {{horizontalOffset, Eigen::Vector3d(0, 0.3, 0)},
       {horizontalOffset + 1, Eigen::Vector3d(0, 0.3, 0)},
       {spanGridSize * (spanGridSize + 1) - 2 + horizontalOffset,
        Eigen::Vector3d(0, -0.3, 0)},
       {spanGridSize * (spanGridSize + 1) - 2 + horizontalOffset + 1,
        Eigen::Vector3d(0, -0.3, 0)},
       {std::floor(spanGridSize / 2) * (spanGridSize + 1) - 2,
        Eigen::Vector3d(0.3, 0, 0)},
       {(std::floor(spanGridSize / 2) + 1) * (spanGridSize + 1) - 2,
        Eigen::Vector3d(0.3, 0, 0)},
       {std::floor(spanGridSize / 2) * (spanGridSize + 1) - 2 + spanGridSize,
        Eigen::Vector3d(-0.3, 0, 0)},
       {(std::floor(spanGridSize / 2) + 1) * (spanGridSize + 1) - 2 +
            spanGridSize,
        Eigen::Vector3d(-0.3, 0, 0)}});

  SimulationJob longSpanGridshell_simulationJob{
      std::make_shared<SimulationMesh>(longSpanGrid),
      "long span gridshell",
      fixedVertices,
      {},
      nodalForcedDisplacements};
  longSpanGridshell_simulationJob.pMesh->setBeamMaterial(0.3, 200 * 1e9);
  if (typeid(CrossSectionType) != typeid(CylindricalBeamDimensions)) {
    std::cerr << "A cylindrical cross section is expected for running the "
                 "paper examples."
              << std::endl;
  }
  longSpanGridshell_simulationJob.pMesh->setBeamCrossSection(
      CrossSectionType{0.03, 0.026});
  SimulationResults longSpanGridshell_simulationResults =
      formFinder.executeSimulation(
          std::make_shared<SimulationJob>(longSpanGridshell_simulationJob),
          settings);
  longSpanGridshell_simulationResults.save();

  const std::string longSpan_groundOfTruthBinaryFilename =
      std::filesystem::path(groundOfTruthFolder)
          .append("LongSpanGridshell_displacements.eigenBin")
          .string();
  assert(std::filesystem::exists(
      std::filesystem::path(longSpan_groundOfTruthBinaryFilename)));
  Eigen::MatrixXd longSpanGridshell_groundOfTruthDisplacements;
  Eigen::read_binary(longSpan_groundOfTruthBinaryFilename,
                     longSpanGridshell_groundOfTruthDisplacements);
  assert(longSpanGridshell_simulationResults.isEqual(
      longSpanGridshell_groundOfTruthDisplacements));
  if (!longSpanGridshell_simulationResults.isEqual(
          longSpanGridshell_groundOfTruthDisplacements)) {
    std::cerr << "Error:Long span simulation produces wrong results."
              << std::endl;
    return;
  }

  std::cout << "Form finder unit tests succesufully passed." << std::endl;

  //  polyscope::show();
}

void FormFinder::reset() {
  mCurrentSimulationStep = 0;
  history.clear();
  constrainedVertices.clear();
  rigidSupports.clear();
  pMesh.reset();
  plotYValues.clear();
  plotHandle.reset();
  checkedForMaximumMoment = false;
  mSettings.shouldUseTranslationalKineticEnergyThreshold = false;
  externalMomentsNorm = 0;
  mSettings.drawingStep = 1;
  Dt = mSettings.Dtini;
  numOfDampings=0;
}

VectorType FormFinder::computeDisplacementDifferenceDerivative(
    const EdgeType &e, const DifferentiateWithRespectTo &dui) const {
  VectorType displacementDiffDeriv(0, 0, 0);
  const DoFType &dofi = dui.dofi;
  const bool differentiateWithRespectToANonEdgeNode =
      e.cV(0) != &dui.v && e.cV(1) != &dui.v;
  if (differentiateWithRespectToANonEdgeNode || dofi > 2) {
    return displacementDiffDeriv;
  }

  if (e.cV(0) == &dui.v) {
    displacementDiffDeriv[dofi] = -1;
  } else if (e.cV(1) == &dui.v) {
    displacementDiffDeriv[dofi] = 1;
  }

  return displacementDiffDeriv;
}

VectorType FormFinder::computeDerivativeOfNormal(
    const VertexType &v, const DifferentiateWithRespectTo &dui) const {
  const size_t vi = pMesh->getIndex(v);
  VectorType normalDerivative(0, 0, 0);
  if (&dui.v != &v ||
      (dui.dofi == 0 || dui.dofi == 1 || dui.dofi == 2 || dui.dofi == 5)) {
    return normalDerivative;
  }
  const VectorType &n = v.cN();
  const double &nx = n[0], ny = n[1];
  const double nxnyMagnitude = std::pow(nx, 2) + std::pow(ny, 2);

  if (dui.dofi == 3) {
    if (nxnyMagnitude + 1e-5 >= 1) {
      const double normalDerivativeX =
          1 / sqrt(nxnyMagnitude) -
          std::pow(nx, 2) / std::pow(nxnyMagnitude, 1.5);
      const double normalDerivativeY = -nx * ny / std::pow(nxnyMagnitude, 1.5);
      const double normalDerivativeZ = 0;
      normalDerivative =
          VectorType(normalDerivativeX, normalDerivativeY, normalDerivativeZ);
    } else {
      const double normalDerivativeX = 1;
      const double normalDerivativeY = 0;
      const double normalDerivativeZ = -nx / std::sqrt(1 - nxnyMagnitude);
      normalDerivative =
          VectorType(normalDerivativeX, normalDerivativeY, normalDerivativeZ);
    }
  } else if (dui.dofi == 4) {
    if (nxnyMagnitude + 1e-5 >= 1) {
      const double normalDerivativeX = -nx * ny / std::pow(nxnyMagnitude, 1.5);
      const double normalDerivativeY =
          1 / sqrt(nxnyMagnitude) -
          std::pow(ny, 2) / std::pow(nxnyMagnitude, 1.5);
      const double normalDerivativeZ = 0;
      normalDerivative =
          VectorType(normalDerivativeX, normalDerivativeY, normalDerivativeZ);
    } else {
      const double normalDerivativeX = 0;
      const double normalDerivativeY = 1;
      const double normalDerivativeZ = -ny / std::sqrt(1 - nxnyMagnitude);
      normalDerivative =
          VectorType(normalDerivativeX, normalDerivativeY, normalDerivativeZ);
    }
  }

  const bool normalDerivativeIsFinite = std::isfinite(normalDerivative[0]) &&
                                        std::isfinite(normalDerivative[1]) &&
                                        std::isfinite(normalDerivative[2]);
  assert(normalDerivativeIsFinite);
  const bool shouldBreak =
      mCurrentSimulationStep == 118 && vi == 1 && dui.dofi == 3;

  return normalDerivative;
}

double FormFinder::computeDerivativeElementLength(
    const EdgeType &e, const DifferentiateWithRespectTo &dui) const {
  if (e.cV(0) != &dui.v && e.cV(1) != &dui.v) {
    return 0;
  }

  const VectorType &X_j = e.cP(0);
  const VectorType &X_jplus1 = e.cP(1);
  const VectorType positionVectorDiff = X_jplus1 - X_j;
  const VectorType displacementDiffDeriv =
      computeDisplacementDifferenceDerivative(e, dui);
  const double edgeLength = pMesh->elements[e].length;
  const double L_kDeriv =
      positionVectorDiff * displacementDiffDeriv / edgeLength;
  return L_kDeriv;
}

double FormFinder::computeDerivativeOfNorm(const VectorType &x,
                                           const VectorType &derivativeOfX) {
  return x.dot(derivativeOfX) / x.Norm();
}

VectorType FormFinder::computeDerivativeOfCrossProduct(
    const VectorType &a, const VectorType &derivativeOfA, const VectorType &b,
    const VectorType &derivativeOfB) {
  const auto firstTerm = Cross(derivativeOfA, b);
  const auto secondTerm = Cross(a, derivativeOfB);
  return firstTerm + secondTerm;
}

VectorType
FormFinder::computeDerivativeOfR(const EdgeType &e,
                                 const DifferentiateWithRespectTo &dui) const {

  const VertexType &v_j = *e.cV(0);
  const VertexType &v_jplus1 = *e.cV(1);
  const VectorType normal_j = v_j.cN();
  const VectorType normal_jplus1 = v_jplus1.cN();
  const VectorType derivativeOfNormal_j =
      &v_j == &dui.v && dui.dofi > 2
          ? pMesh->nodes[v_j].derivativeOfNormal[dui.dofi]
          : VectorType(0, 0, 0);
  const VectorType derivativeOfNormal_jplus1 =
      &v_jplus1 == &dui.v && dui.dofi > 2
          ? pMesh->nodes[v_jplus1].derivativeOfNormal[dui.dofi]
          : VectorType(0, 0, 0);

  const VectorType derivativeOfSumOfNormals =
      derivativeOfNormal_j + derivativeOfNormal_jplus1;
  const VectorType sumOfNormals = normal_j + normal_jplus1;
  const double normOfSumOfNormals = sumOfNormals.Norm();
  const double derivativeOfNormOfSumOfNormals =
      computeDerivativeOfNorm(sumOfNormals, derivativeOfSumOfNormals);

  const VectorType derivativeOfR_firstTerm = -sumOfNormals *
                                             derivativeOfNormOfSumOfNormals /
                                             std::pow(normOfSumOfNormals, 2);
  const VectorType derivativeOfR_secondTerm =
      derivativeOfSumOfNormals / normOfSumOfNormals;
  const VectorType derivativeOfR =
      derivativeOfR_firstTerm + derivativeOfR_secondTerm;

  assert(std::isfinite(derivativeOfR[0]) && std::isfinite(derivativeOfR[1]) &&
         std::isfinite(derivativeOfR[2]));

  return derivativeOfR;
}

VectorType
FormFinder::computeDerivativeT1(const EdgeType &e,
                                const DifferentiateWithRespectTo &dui) const {

  const VectorType &X_j = e.cP(0);
  const VectorType &X_jplus1 = e.cP(1);
  const VectorType edgeVector = X_jplus1 - X_j;
  const double L_kDerivative = computeDerivativeElementLength(e, dui);
  const double edgeLength = pMesh->elements[e].length;
  const VectorType firstTerm =
      -edgeVector * L_kDerivative / std::pow(edgeLength, 2);
  const VectorType secondTerm =
      computeDisplacementDifferenceDerivative(e, dui) / edgeLength;
  const VectorType t1Derivative = firstTerm + secondTerm;

  return t1Derivative;
}

VectorType
FormFinder::computeDerivativeT2(const EdgeType &e,
                                const DifferentiateWithRespectTo &dui) const {

  const DoFType dofi = dui.dofi;

  const VertexType &v_j = *e.cV(0);
  const size_t vi_j = pMesh->getIndex(v_j);
  const VertexType &v_jplus1 = *e.cV(1);
  const size_t vi_jplus1 = pMesh->getIndex(v_jplus1);

  const VectorType r = (v_j.cN() + v_jplus1.cN()).Normalize();
  const VectorType derivativeR_j =
      dofi > 2 && &v_j == &dui.v ? pMesh->elements[e].derivativeR_j[dui.dofi]
                                 : VectorType(0, 0, 0);
  const VectorType derivativeR_jplus1 =
      dofi > 2 && &v_jplus1 == &dui.v
          ? pMesh->elements[e].derivativeR_jplus1[dui.dofi]
          : VectorType(0, 0, 0);
  const VectorType derivativeR = derivativeR_j + derivativeR_jplus1;

  const VectorType &t1 = pMesh->elements[e].frame.t1;
  const VectorType derivativeT1_j =
      dofi < 3 && &v_j == &dui.v ? pMesh->elements[e].derivativeT1_j[dui.dofi]
                                 : VectorType(0, 0, 0);
  const VectorType derivativeT1_jplus1 =
      dofi < 3 && &v_jplus1 == &dui.v
          ? pMesh->elements[e].derivativeT1_jplus1[dui.dofi]
          : VectorType(0, 0, 0);
  const VectorType derivativeT1 = derivativeT1_j + derivativeT1_jplus1;

  const VectorType derivativeOfRCrossT1 =
      computeDerivativeOfCrossProduct(r, derivativeR, t1, derivativeT1);
  const VectorType rCrossT1 = Cross(r, t1);
  const double normOfRCrossT1 = rCrossT1.Norm();
  const double derivativeNormRCrossT1 =
      computeDerivativeOfNorm(rCrossT1, derivativeOfRCrossT1);

  const VectorType t2Deriv_firstTerm =
      -(rCrossT1 * derivativeNormRCrossT1) / std::pow(normOfRCrossT1, 2);
  const VectorType t2Deriv_secondTerm = derivativeOfRCrossT1 / normOfRCrossT1;
  const VectorType t2Deriv = t2Deriv_firstTerm + t2Deriv_secondTerm;

  const double t2DerivNorm = t2Deriv.Norm();
  assert(std::isfinite(t2DerivNorm));
  const bool shouldBreak = mCurrentSimulationStep == 118 &&
                           (vi_jplus1 == 1 || vi_j == 1) && dofi == 3;
  return t2Deriv;
}

VectorType
FormFinder::computeDerivativeT3(const EdgeType &e,
                                const DifferentiateWithRespectTo &dui) const {
  const Element &element = pMesh->elements[e];
  const VectorType &t1 = element.frame.t1;
  const VectorType &t2 = element.frame.t2;
  const VectorType t1CrossT2 = Cross(t1, t2);
  const VertexType &v_j = *e.cV(0);
  const size_t vi_j = pMesh->getIndex(v_j);
  const VertexType &v_jplus1 = *e.cV(1);
  const size_t vi_jplus1 = pMesh->getIndex(v_jplus1);
  const VectorType derivativeT1_j =
      dui.dofi < 3 && &v_j == &dui.v
          ? pMesh->elements[e].derivativeT1_j[dui.dofi]
          : VectorType(0, 0, 0);
  const VectorType derivativeT1_jplus1 =
      dui.dofi < 3 && &v_jplus1 == &dui.v
          ? pMesh->elements[e].derivativeT1_jplus1[dui.dofi]
          : VectorType(0, 0, 0);
  const VectorType derivativeT1 = derivativeT1_j + derivativeT1_jplus1;

  //  const VectorType derivativeOfT2 = computeDerivativeT2(e, dui);
  //  const VectorType derivativeT2_j =
  //      &v_j == &dui.v
  //          ? mesh->elements[e].derivativeT2_j[dui.dofi]
  //          : VectorType(0, 0, 0);
  //  const VectorType derivativeT2_jplus1 =
  //       &v_jplus1 == &dui.v
  //          ? mesh->elements[e].derivativeT2_jplus1[dui.dofi]
  //          : VectorType(0, 0, 0);
  VectorType derivativeT2(0, 0, 0);
  if (&v_j == &dui.v) {
    derivativeT2 = pMesh->elements[e].derivativeT2_j[dui.dofi];
  } else if (&v_jplus1 == &dui.v) {
    derivativeT2 = pMesh->elements[e].derivativeT2_jplus1[dui.dofi];
  }

  const VectorType derivativeT1CrossT2 =
      computeDerivativeOfCrossProduct(t1, derivativeT1, t2, derivativeT2);
  const double derivativeOfNormT1CrossT2 =
      computeDerivativeOfNorm(t1CrossT2, derivativeT1CrossT2);
  const double normT1CrossT2 = t1CrossT2.Norm();

  const VectorType t3Deriv_firstTerm =
      -(t1CrossT2 * derivativeOfNormT1CrossT2) / std::pow(normT1CrossT2, 2);
  const VectorType t3Deriv_secondTerm = derivativeT1CrossT2 / normT1CrossT2;
  const VectorType t3Deriv = t3Deriv_firstTerm + t3Deriv_secondTerm;

  assert(std::isfinite(t3Deriv[0]) && std::isfinite(t3Deriv[1]) &&
         std::isfinite(t3Deriv[2]));
  return t3Deriv;
}

double FormFinder::computeDerivativeTheta1(const EdgeType &e,
                                           const VertexIndex &evi,
                                           const VertexIndex &dwrt_evi,
                                           const DoFType &dwrt_dofi) const {
  const VertexType &v = *e.cV(evi);
  const size_t vi = pMesh->getIndex(v);
  const Element &element = pMesh->elements[e];
  const VectorType derivativeT1 = element.derivativeT1[dwrt_evi][dwrt_dofi];
  const VectorType derivativeT3 = element.derivativeT3[dwrt_evi][dwrt_dofi];
  const VectorType nDerivative =
      evi != dwrt_evi ? VectorType(0, 0, 0)
                      : pMesh->nodes[v].derivativeOfNormal[dwrt_dofi];
  const VectorType n = v.cN();
  const VectorType &t1 = element.frame.t1;
  const VectorType &t3 = element.frame.t3;
  const double theta1Derivative =
      derivativeT1 * Cross(t3, n) +
      t1 * (Cross(derivativeT3, n) + Cross(t3, nDerivative));
  const bool shouldBreak =
      mCurrentSimulationStep == 118 && vi == 1 && dwrt_dofi == 3;

  return theta1Derivative;
}

double FormFinder::computeDerivativeTheta2(const EdgeType &e,
                                           const VertexIndex &evi,
                                           const VertexIndex &dwrt_evi,
                                           const DoFType &dwrt_dofi) const {
  const VertexType &v = *e.cV(evi);

  const Element &element = pMesh->elements[e];
  const VectorType derivativeT2 = element.derivativeT2[dwrt_evi][dwrt_dofi];
  const VectorType derivativeT3 = element.derivativeT3[dwrt_evi][dwrt_dofi];

  const VectorType n = v.cN();
  const VectorType &t2 = element.frame.t2;
  const VectorType &t3 = element.frame.t3;
  const VectorType nDerivative =
      dwrt_evi == evi ? pMesh->nodes[v].derivativeOfNormal[dwrt_dofi]
                      : VectorType(0, 0, 0);
  const double theta2Derivative =
      derivativeT2 * Cross(t3, n) +
      t2 * (Cross(derivativeT3, n) + Cross(t3, nDerivative));

  return theta2Derivative;
}

double FormFinder::computeTheta3(const EdgeType &e, const VertexType &v) {
  const VertexIndex &vi = pMesh->nodes[v].vi;
  //  assert(e.cV(0) == &v || e.cV(1) == &v);

  const Element &elem = pMesh->elements[e];
  const EdgeIndex &ei = elem.ei;
  const Element::LocalFrame &ef = elem.frame;
  const VectorType &t1 = ef.t1;
  const VectorType &n = v.cN();
  const Node &node = pMesh->nodes[v];
  const double &nR = node.nR; // TODO: This makes the function non-const.
                              // Should be refactored in the future

  double theta3;
  if (&e == node.referenceElement) {
    // Use nR as theta3 only for the first star edge
    return nR;
  }
  //    assert(node.alphaAngles.contains(mesh->getIndex(e)));
  double alphaAngle = node.alphaAngles.find(elem.ei)->second;
  const EdgeType &refElem = *node.referenceElement;
  const double rotationAngle = nR + alphaAngle;

  //    const VectorType &t1_k = computeT1Vector(refElem);
  const VectorType &t1_k = pMesh->elements[refElem].frame.t1;
  const VectorType f1 = (t1_k - (n * (t1_k * n))).Normalize();
  vcg::Matrix44<double> rotationMatrix;
  rotationMatrix.SetRotateRad(rotationAngle, n);
  const double cosRotationAngle = cos(rotationAngle);
  const double sinRotationAngle = sin(rotationAngle);
  const VectorType f2 =
      (f1 * cosRotationAngle + Cross(n, f1) * sinRotationAngle).Normalize();
  const VectorType &t1Current = t1;
  const VectorType f3 = Cross(t1Current, f2);

  Element &element = pMesh->elements[e];
  // Save for computing theta3Derivative
  if (&v == e.cV(0)) {
    element.f1_j = f1;
    element.f2_j = f2;
    element.f3_j = f3;
    element.cosRotationAngle_j = cosRotationAngle;
    element.sinRotationAngle_j = sinRotationAngle;
  } else {
    element.f1_jplus1 = f1;
    element.f2_jplus1 = f2;
    element.f3_jplus1 = f3;
    element.cosRotationAngle_jplus1 = cosRotationAngle;
    element.sinRotationAngle_jplus1 = sinRotationAngle;
  }
  theta3 = f3.dot(n);

  return theta3;
}

double FormFinder::computeDerivativeTheta3(
    const EdgeType &e, const VertexType &v,
    const DifferentiateWithRespectTo &dui) const {
  const Node &node = pMesh->nodes[v];
  const VertexIndex &vi = pMesh->nodes[v].vi;
  const bool isRigidSupport = rigidSupports.contains(vi);
  if (&e == node.referenceElement && !isRigidSupport) {
    if (dui.dofi == DoF::Nr && &dui.v == &v) {
      return 1;
    } else {
      return 0;
    }
  }
  //  assert(e.cV(0) == &v || e.cV(1) == &v);

  const Element &element = pMesh->elements[e];
  const Element::LocalFrame &ef = element.frame;
  const VectorType &t1 = ef.t1;
  const VectorType &n = v.cN();
  const DoFType &dofi = dui.dofi;
  const VertexPointer &vp_j = e.cV(0);
  const VertexPointer &vp_jplus1 = e.cV(1);

  double derivativeTheta3_dofi = 0;
  if (isRigidSupport) {
    const VectorType &t1Initial =
        computeT1Vector(pMesh->nodes[vp_j].initialLocation,
                        pMesh->nodes[vp_jplus1].initialLocation);
    VectorType g1 = Cross(t1, t1Initial);

    const VectorType derivativeInitialT1_dofi(0, 0, 0);
    const VectorType derivativeT1_j = dui.dofi < 3 && vp_j == &dui.v
                                          ? element.derivativeT1_j[dui.dofi]
                                          : VectorType(0, 0, 0);
    const VectorType derivativeT1_jplus1 =
        dui.dofi < 3 && vp_jplus1 == &dui.v
            ? element.derivativeT1_jplus1[dui.dofi]
            : VectorType(0, 0, 0);
    const VectorType derivativeT1_dofi = derivativeT1_j + derivativeT1_jplus1;
    //    VectorType derivativeT1_dofi(0, 0, 0);
    //    if (dui.dofi < 3) {
    //      if (&v_j == &dui.v) {
    //        derivativeT1_dofi = mesh->elements[e].derivativeT1_j[dui.dofi];
    //      } else if (&v_jplus1 == &dui.v) {
    //        derivativeT1_dofi =
    //        mesh->elements[e].derivativeT1_jplus1[dui.dofi];
    //      }
    //    }

    const VectorType derivativeG1_firstTerm =
        Cross(derivativeT1_dofi, t1Initial);
    const VectorType derivativeG1_secondTerm =
        Cross(t1, derivativeInitialT1_dofi);
    const VectorType derivativeG1 =
        derivativeG1_firstTerm + derivativeG1_secondTerm;
    const VectorType derivativeNormal = &v == &dui.v && dui.dofi > 2
                                            ? node.derivativeOfNormal[dui.dofi]
                                            : VectorType(0, 0, 0);
    derivativeTheta3_dofi = derivativeG1 * n + g1 * derivativeNormal;
    return derivativeTheta3_dofi;
  }
  EdgeType &refElem = *node.referenceElement;
  const VectorType &t1_k = pMesh->elements[refElem].frame.t1;
  VectorType f1, f2, f3;
  double cosRotationAngle, sinRotationAngle;
  if (e.cV(0) == &v) {
    f1 = element.f1_j;
    cosRotationAngle = element.cosRotationAngle_j;
    sinRotationAngle = element.sinRotationAngle_j;
    f2 = element.f2_j;
    f3 = element.f3_j;
  } else {
    f1 = element.f1_jplus1;
    cosRotationAngle = element.cosRotationAngle_jplus1;
    sinRotationAngle = element.sinRotationAngle_jplus1;
    f2 = element.f2_jplus1;
    f3 = element.f3_jplus1;
  }
  const VectorType &t1_kplus1 = t1;
  const VectorType f1Normalized = f1 / f1.Norm();

  VectorType derivativeF1(0, 0, 0);
  VectorType derivativeF2(0, 0, 0);
  VectorType derivativeF3(0, 0, 0);
  if (dui.dofi < 3) {

    const VectorType derivativeT1_kplus1_j =
        vp_j == &dui.v ? element.derivativeT1_j[dui.dofi] : VectorType(0, 0, 0);
    const VectorType derivativeT1_kplus1_jplus1 =
        vp_jplus1 == &dui.v ? element.derivativeT1_jplus1[dui.dofi]
                            : VectorType(0, 0, 0);
    const VectorType derivativeT1_kplus1 =
        derivativeT1_kplus1_j + derivativeT1_kplus1_jplus1;

    const VectorType derivativeT1_k_j =
        refElem.cV(0) == &dui.v
            ? pMesh->elements[refElem].derivativeT1_j[dui.dofi]
            : VectorType(0, 0, 0);
    const VectorType derivativeT1_k_jplus1 =
        refElem.cV(1) == &dui.v
            ? pMesh->elements[refElem].derivativeT1_jplus1[dui.dofi]
            : VectorType(0, 0, 0);
    const VectorType derivativeT1_k = derivativeT1_k_j + derivativeT1_k_jplus1;

    derivativeF1 = derivativeT1_k - (n * (derivativeT1_k * n));

    const double f1Norm = f1.Norm();
    const VectorType derivativeF1Normalized =
        -f1 * (f1 * derivativeF1) / (f1Norm * f1Norm * f1Norm) +
        derivativeF1 / f1Norm;

    derivativeF2 = derivativeF1Normalized * cosRotationAngle +
                   Cross(n, derivativeF1Normalized) * sinRotationAngle;
    const VectorType derivativeF3_firstTerm = Cross(derivativeT1_kplus1, f2);
    const VectorType derivativeF3_secondTerm = Cross(t1_kplus1, derivativeF2);
    derivativeF3 = derivativeF3_firstTerm + derivativeF3_secondTerm;
    derivativeTheta3_dofi = derivativeF3 * n;

  } else if (dui.dofi == DoF::Nr && &dui.v == &v) {
    derivativeF2 = f1Normalized * (-sinRotationAngle) +
                   Cross(n, f1Normalized) * cosRotationAngle;
    derivativeF3 = Cross(t1_kplus1, derivativeF2);
    derivativeTheta3_dofi = derivativeF3 * n;
  } else { // 2<dofi<5
    if (&v == &dui.v) {
      const VectorType &derivativeOfNormal = node.derivativeOfNormal[dofi];
      derivativeF1 =
          -(n * (t1_k * derivativeOfNormal) + derivativeOfNormal * (t1_k * n));
      const double f1Norm = f1.Norm();
      const VectorType derivativeF1Normalized =
          -f1 * (f1 * derivativeF1) / (f1Norm * f1Norm * f1Norm) +
          derivativeF1 / f1Norm;

      derivativeF2 = derivativeF1Normalized * cosRotationAngle +

                     (Cross(derivativeOfNormal, f1Normalized) +
                      Cross(n, derivativeF1Normalized)) *
                         sinRotationAngle;
      derivativeF3 = Cross(t1_kplus1, derivativeF2);
      derivativeTheta3_dofi = derivativeF3 * n + f3 * derivativeOfNormal;
    }
  }
  return derivativeTheta3_dofi;
}

double FormFinder::computeTotalPotentialEnergy() const {

  double totalExternalPotentialEnergy = 0;
  for (const SimulationMesh::VertexType &v : pMesh->vert) {
    const Node &node = pMesh->nodes[v];
    if (!node.force.hasExternalForce()) {
      continue;
    }
    const auto nodeDisplacement = v.cP() - node.initialLocation;
    const SimulationMesh::CoordType externalForce(
        node.force.external[0], node.force.external[1], node.force.external[2]);
    totalExternalPotentialEnergy += externalForce.dot(nodeDisplacement);
  }

  double totalInternalPotentialEnergy = 0;
  for (const SimulationMesh::EdgeType &e : pMesh->edge) {

    const Element &element = pMesh->elements[e];
    const EdgeIndex ei = pMesh->getIndex(e);
    const double e_k = element.length - element.initialLength;
    const double axialPE = pow(e_k, 2) * element.material.youngsModulus *
                           element.A / (2 * element.initialLength);
    const double theta1Diff = element.rotationalDisplacements_jplus1.theta1 -
                              element.rotationalDisplacements_j.theta1;
    const double torsionalPE = element.G * element.J * pow(theta1Diff, 2) /
                               (2 * element.initialLength);
    const double &theta2_j = element.rotationalDisplacements_j.theta2;
    const double &theta2_jplus1 = element.rotationalDisplacements_jplus1.theta2;
    const double firstBendingPE_inBracketsTerm = 4 * pow(theta2_j, 2) +
                                                 4 * theta2_j * theta2_jplus1 +
                                                 4 * pow(theta2_jplus1, 2);
    const double firstBendingPE = firstBendingPE_inBracketsTerm *
                                  element.material.youngsModulus * element.I2 /
                                  (2 * element.initialLength);
    const double &theta3_j = element.rotationalDisplacements_j.theta3;
    const double &theta3_jplus1 = element.rotationalDisplacements_jplus1.theta3;
    const double secondBendingPE_inBracketsTerm = 4 * pow(theta3_j, 2) +
                                                  4 * theta3_j * theta3_jplus1 +
                                                  4 * pow(theta3_jplus1, 2);
    const double secondBendingPE = secondBendingPE_inBracketsTerm * 2 *
                                   element.material.youngsModulus * element.I3 /
                                   element.initialLength;

    totalInternalPotentialEnergy +=
        axialPE + torsionalPE + firstBendingPE + secondBendingPE;
  }
  return totalInternalPotentialEnergy - totalExternalPotentialEnergy;
}

void FormFinder::updateResidualForcesOnTheFly(
    const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
        &fixedVertices) {

  //  std::vector<Vector6d> internalForcesParallel(mesh->vert.size());

  std::vector<std::vector<std::pair<int, Vector6d>>>
      internalForcesContributionsFromEachEdge(
          pMesh->EN(),
          std::vector<std::pair<int, Vector6d>>(4, {-1, Vector6d()}));
  //  omp_lock_t writelock;
  //  omp_init_lock(&writelock);
//#pragma omp parallel for //schedule(static) num_threads(8)
  for (int ei = 0; ei < pMesh->EN(); ei++) {
    const EdgeType &e = pMesh->edge[ei];
    const SimulationMesh::VertexType &ev_j = *e.cV(0);
    const SimulationMesh::VertexType &ev_jplus1 = *e.cV(1);
    const Element &element = pMesh->elements[e];
    const Element::LocalFrame &ef = element.frame;
    const VectorType t3CrossN_j = Cross(ef.t3, ev_j.cN());
    const VectorType t3CrossN_jplus1 = Cross(ef.t3, ev_jplus1.cN());
    const double theta1_j = ef.t1.dot(t3CrossN_j);
    const double theta1_jplus1 = ef.t1.dot(t3CrossN_jplus1);
    const double theta2_j = ef.t2.dot(t3CrossN_j);
    const double theta2_jplus1 = ef.t2.dot(t3CrossN_jplus1);
    const double theta3_j = computeTheta3(e, ev_j);
    const double theta3_jplus1 = computeTheta3(e, ev_jplus1);
    //    element.rotationalDisplacements_j.theta1 = theta1_j;
    //    element.rotationalDisplacements_jplus1.theta1 = theta1_jplus1;
    //    element.rotationalDisplacements_j.theta2 = theta2_j;
    //    element.rotationalDisplacements_jplus1.theta2 = theta2_jplus1;
    //    element.rotationalDisplacements_j.theta3 = theta3_j;
    //    element.rotationalDisplacements_jplus1.theta3 = theta3_jplus1;
    std::vector<std::pair<int, Vector6d>>
        internalForcesContributionFromThisEdge(4, {-1, Vector6d()});
    for (VertexIndex evi = 0; evi < 2; evi++) {
      const SimulationMesh::VertexType &ev = *e.cV(evi);
      const Node &edgeNode = pMesh->nodes[ev];
      internalForcesContributionFromThisEdge[evi].first = edgeNode.vi;

      const VertexPointer &rev_j = edgeNode.referenceElement->cV(0);
      const VertexPointer &rev_jplus1 = edgeNode.referenceElement->cV(1);
      // refElemOtherVertex can be precomputed
      const VertexPointer &refElemOtherVertex =
          rev_j == &ev ? rev_jplus1 : rev_j;
      const Node &refElemOtherVertexNode = pMesh->nodes[refElemOtherVertex];
      if (edgeNode.referenceElement != &e) {
        internalForcesContributionFromThisEdge[evi + 2].first =
            refElemOtherVertexNode.vi;
      }
      const size_t vi = edgeNode.vi;
      //      #pragma omp parallel for schedule(static) num_threads(6)
      for (DoFType dofi = DoF::Ux; dofi < DoF::NumDoF; dofi++) {
        const bool isDofConstrainedFor_ev =
            fixedVertices.contains(edgeNode.vi) &&
            fixedVertices.at(edgeNode.vi).contains(dofi);
        if (!isDofConstrainedFor_ev) {
          DifferentiateWithRespectTo dui{ev, dofi};
          // Axial force computation
          const double e_k = element.length - element.initialLength;
          const double e_kDeriv = computeDerivativeElementLength(e, dui);
          const double axialForce_dofi =
              e_kDeriv * e_k * element.rigidity.axial;

          // Torsional force computation
          const double theta1_j_deriv =
              computeDerivativeTheta1(e, 0, evi, dofi);
          const double theta1_jplus1_deriv =
              computeDerivativeTheta1(e, 1, evi, dofi);
          const double theta1Diff = theta1_jplus1 - theta1_j;
          const double theta1DiffDerivative =
              theta1_jplus1_deriv - theta1_j_deriv;
          const double torsionalForce_dofi =
              theta1DiffDerivative * theta1Diff * element.rigidity.torsional;

          // First bending force computation
          ////theta2_j derivative
          const double theta2_j_deriv =
              computeDerivativeTheta2(e, 0, evi, dofi);
          ////theta2_jplus1 derivative
          const double theta2_jplus1_deriv =
              computeDerivativeTheta2(e, 1, evi, dofi);
          ////1st in bracket term
          const double firstBendingForce_inBracketsTerm_0 =
              theta2_j_deriv * 2 * theta2_j;
          ////2nd in bracket term
          const double firstBendingForce_inBracketsTerm_1 =
              theta2_jplus1_deriv * theta2_j;
          ////3rd in bracket term
          const double firstBendingForce_inBracketsTerm_2 =
              theta2_j_deriv * theta2_jplus1;
          ////4th in bracket term
          const double firstBendingForce_inBracketsTerm_3 =
              2 * theta2_jplus1_deriv * theta2_jplus1;
          // 3rd term computation
          const double firstBendingForce_inBracketsTerm =
              firstBendingForce_inBracketsTerm_0 +
              firstBendingForce_inBracketsTerm_1 +
              firstBendingForce_inBracketsTerm_2 +
              firstBendingForce_inBracketsTerm_3;
          const double firstBendingForce_dofi =
              firstBendingForce_inBracketsTerm * element.rigidity.firstBending;

          // Second bending force computation
          ////theta2_j derivative
          const double theta3_j_deriv = computeDerivativeTheta3(e, ev_j, dui);
          ////theta2_jplus1 derivative
          const double theta3_jplus1_deriv =
              computeDerivativeTheta3(e, ev_jplus1, dui);
          ////1st in bracket term
          const double secondBendingForce_inBracketsTerm_0 =
              theta3_j_deriv * 2 * theta3_j;
          ////2nd in bracket term
          const double secondBendingForce_inBracketsTerm_1 =
              theta3_jplus1_deriv * theta3_j;
          ////3rd in bracket term
          const double secondBendingForce_inBracketsTerm_2 =
              theta3_j_deriv * theta3_jplus1;
          ////4th in bracket term
          const double secondBendingForce_inBracketsTerm_3 =
              2 * theta3_jplus1_deriv * theta3_jplus1;
          // 3rd term computation
          const double secondBendingForce_inBracketsTerm =
              secondBendingForce_inBracketsTerm_0 +
              secondBendingForce_inBracketsTerm_1 +
              secondBendingForce_inBracketsTerm_2 +
              secondBendingForce_inBracketsTerm_3;
          double secondBendingForce_dofi = secondBendingForce_inBracketsTerm *
                                           element.rigidity.secondBending;
          const bool shouldBreak =
              mCurrentSimulationStep == 118 && edgeNode.vi == 1 && dofi == 3;
          internalForcesContributionFromThisEdge[evi].second[dofi] =
              axialForce_dofi + firstBendingForce_dofi +
              secondBendingForce_dofi + torsionalForce_dofi;
        }
        if (edgeNode.referenceElement != &e) {
          const bool isDofConstrainedFor_refElemOtherVertex =
              fixedVertices.contains(refElemOtherVertexNode.vi) &&
              fixedVertices.at(refElemOtherVertexNode.vi).contains(dofi);
          if (!isDofConstrainedFor_refElemOtherVertex) {
            DifferentiateWithRespectTo dui{*refElemOtherVertex, dofi};
            ////theta3_j derivative
            const double theta3_j_deriv = computeDerivativeTheta3(e, ev_j, dui);
            ////theta3_jplus1 derivative
            const double theta3_jplus1_deriv =
                computeDerivativeTheta3(e, ev_jplus1, dui);
            ////1st in bracket term
            const double secondBendingForce_inBracketsTerm_0 =
                theta3_j_deriv * 2 * theta3_j;
            ////2nd in bracket term
            const double secondBendingForce_inBracketsTerm_1 =
                theta3_jplus1_deriv * theta3_j;
            ////3rd in bracket term
            const double secondBendingForce_inBracketsTerm_2 =
                theta3_j_deriv * theta3_jplus1;
            ////4th in bracket term
            const double secondBendingForce_inBracketsTerm_3 =
                theta3_jplus1_deriv * 2 * theta3_jplus1;

            // 4th term computation
            const double secondBendingForce_inBracketsTerm =
                secondBendingForce_inBracketsTerm_0 +
                secondBendingForce_inBracketsTerm_1 +
                secondBendingForce_inBracketsTerm_2 +
                secondBendingForce_inBracketsTerm_3;
            const double secondBendingForce_dofi =
                secondBendingForce_inBracketsTerm *
                element.rigidity.secondBending;
            internalForcesContributionFromThisEdge[evi + 2].second[dofi] =
                secondBendingForce_dofi;
          }
        }
      }
    }
    internalForcesContributionsFromEachEdge[ei] =
        internalForcesContributionFromThisEdge;
  }

  //#pragma omp parallel for schedule(static) num_threads(8)

  for (size_t vi = 0; vi < pMesh->VN(); vi++) {
    Node::Forces &force = pMesh->nodes[vi].force;
    force.residual = force.external;
    force.internal = 0;
  }
  double totalResidualForcesNorm = 0;
  for (size_t ei = 0; ei < pMesh->EN(); ei++) {
    for (int i = 0; i < 4; i++) {
      std::pair<int, Vector6d> internalForcePair =
          internalForcesContributionsFromEachEdge[ei][i];
      int vi = internalForcePair.first;
      if (i > 1 && vi == -1) {
        continue;
      }
      Node::Forces &force = pMesh->nodes[vi].force;
      force.internal = force.internal + internalForcePair.second;
      force.residual = force.residual + (internalForcePair.second * -1);
    }
  }
  Vector6d sumOfResidualForces(0);
  for (size_t vi = 0; vi < pMesh->VN(); vi++) {
    Node::Forces &force = pMesh->nodes[vi].force;
    const Vector6d &nodeResidualForce = force.residual;
    //    sumOfResidualForces = sumOfResidualForces + nodeResidualForce;
    const double residualForceNorm = nodeResidualForce.norm();
    const bool shouldTrigger = std::isnan(residualForceNorm);
    totalResidualForcesNorm += residualForceNorm;
  }
  pMesh->previousTotalResidualForcesNorm = pMesh->totalResidualForcesNorm;
  pMesh->totalResidualForcesNorm = totalResidualForcesNorm;
  //  mesh->totalResidualForcesNorm = sumOfResidualForces.norm();

  //  plotYValues.push_back(totalResidualForcesNorm);
  //  auto xPlot = matplot::linspace(0, plotYValues.size(), plotYValues.size());
  //  plotHandle = matplot::scatter(xPlot, plotYValues);
}

void FormFinder::updateNodalExternalForces(
    const std::unordered_map<VertexIndex, Vector6d> &nodalForces,
    const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
        &fixedVertices) {

  externalMomentsNorm = 0;
  for (const std::pair<VertexIndex, Vector6d> &nodalForce : nodalForces) {
    const VertexIndex nodeIndex = nodalForce.first;
    const bool isNodeConstrained = fixedVertices.contains(nodeIndex);
    Node &node = pMesh->nodes[nodeIndex];
    Vector6d nodalExternalForce(0);
    for (DoFType dofi = 0; dofi < 6; dofi++) {
      const bool isDofConstrained =
          isNodeConstrained && fixedVertices.at(nodeIndex).contains(dofi);
      if (isDofConstrained) {
        continue;
      }
      nodalExternalForce[dofi] = nodalForce.second[dofi];
    }
    externalMomentsNorm += std::sqrt(pow(nodalExternalForce[3], 2) +
                                     pow(nodalExternalForce[4], 2));
    //    CoordType v = (mesh->vert[nodeIndex].cP() -
    //    mesh->vert[0].cP()).Normalize(); const double forceMagnitude = 0.1;
    //    nodalExternalForce[3] = v[0] * forceMagnitude;
    //    nodalExternalForce[4] = v[1] * forceMagnitude;
    node.force.external = nodalExternalForce;
  }
}

void FormFinder::updateResidualForces() {

  pMesh->totalResidualForcesNorm = 0;
  for (const VertexType &v : pMesh->vert) {
    Node &node = pMesh->nodes[v];
    node.force.residual = node.force.external - node.force.internal;
    const double residualForceNorm = (node.force.residual).norm();
    pMesh->totalResidualForcesNorm += residualForceNorm;
  }
}

void FormFinder::computeRigidSupports() {

  for (const VertexType &v : pMesh->vert) {
    const VertexIndex vi = pMesh->nodes[v].vi;
    const bool isVertexConstrained = constrainedVertices.contains(vi);
    if (isVertexConstrained) {
      auto constrainedDoFType = constrainedVertices.at(vi);
      const bool hasAllDoFTypeConstrained =
          constrainedDoFType.contains(DoF::Ux) &&
          constrainedDoFType.contains(DoF::Uy) &&
          constrainedDoFType.contains(DoF::Uz) &&
          constrainedDoFType.contains(DoF::Nx) &&
          constrainedDoFType.contains(DoF::Ny) &&
          constrainedDoFType.contains(DoF::Nr);
      if (hasAllDoFTypeConstrained) {
        rigidSupports.insert(vi);
      }
    }
  }
}

void FormFinder::updateNormalDerivatives() {
  for (const VertexType &v : pMesh->vert) {
    for (DoFType dofi = DoF::Nx; dofi < DoF::NumDoF; dofi++) {
      DifferentiateWithRespectTo dui{v, dofi};
      pMesh->nodes[v].derivativeOfNormal[dofi] =
          computeDerivativeOfNormal(v, dui);
    }
  }
}

void FormFinder::updateT1Derivatives() {
  for (const EdgeType &e : pMesh->edge) {
    for (DoFType dofi = DoF::Ux; dofi < DoF::Nx; dofi++) {
      DifferentiateWithRespectTo dui_v0{*e.cV(0), dofi};
      Element &element = pMesh->elements[e];
      element.derivativeT1_j[dofi] = computeDerivativeT1(e, dui_v0);
      DifferentiateWithRespectTo dui_v1{*e.cV(1), dofi};
      element.derivativeT1_jplus1[dofi] = computeDerivativeT1(e, dui_v1);

      element.derivativeT1[0][dofi] = element.derivativeT1_j[dofi];
      element.derivativeT1[1][dofi] = element.derivativeT1_jplus1[dofi];
    }
  }
}

void FormFinder::updateT2Derivatives() {
  for (const EdgeType &e : pMesh->edge) {
    for (DoFType dofi = DoF::Ux; dofi < DoF::NumDoF; dofi++) {
      DifferentiateWithRespectTo dui_v0{*e.cV(0), dofi};
      pMesh->elements[e].derivativeT2_j[dofi] = computeDerivativeT2(e, dui_v0);
      DifferentiateWithRespectTo dui_v1{*e.cV(1), dofi};
      pMesh->elements[e].derivativeT2_jplus1[dofi] =
          computeDerivativeT2(e, dui_v1);

      Element &element = pMesh->elements[e];
      element.derivativeT2[0][dofi] = element.derivativeT2_j[dofi];
      element.derivativeT2[1][dofi] = element.derivativeT2_jplus1[dofi];
    }
  }
}

void FormFinder::updateT3Derivatives() {
  for (const EdgeType &e : pMesh->edge) {
    for (DoFType dofi = DoF::Ux; dofi < DoF::NumDoF; dofi++) {
      DifferentiateWithRespectTo dui_v0{*e.cV(0), dofi};
      Element &element = pMesh->elements[e];
      element.derivativeT3_j[dofi] = computeDerivativeT3(e, dui_v0);
      DifferentiateWithRespectTo dui_v1{*e.cV(1), dofi};
      element.derivativeT3_jplus1[dofi] = computeDerivativeT3(e, dui_v1);

      element.derivativeT3[0][dofi] = element.derivativeT3_j[dofi];
      element.derivativeT3[1][dofi] = element.derivativeT3_jplus1[dofi];
    }
  }
}

void FormFinder::updateRDerivatives() {
  for (const EdgeType &e : pMesh->edge) {
    for (DoFType dofi = DoF::Nx; dofi < DoF::NumDoF; dofi++) {
      DifferentiateWithRespectTo dui_v0{*e.cV(0), dofi};
      pMesh->elements[e].derivativeR_j[dofi] = computeDerivativeOfR(e, dui_v0);
      DifferentiateWithRespectTo dui_v1{*e.cV(1), dofi};
      pMesh->elements[e].derivativeR_jplus1[dofi] =
          computeDerivativeOfR(e, dui_v1);
    }
  }
}

void FormFinder::updateElementalLengths() { pMesh->updateElementalLengths(); }

FormFinder::FormFinder() {}

void FormFinder::updateNodalMasses() {
  const double gamma = 0.8;
  for (VertexType &v : pMesh->vert) {
    const size_t vi = pMesh->getIndex(v);
    double translationalSumSk = 0;
    double rotationalSumSk_I2 = 0;
    double rotationalSumSk_I3 = 0;
    double rotationalSumSk_J = 0;
    for (const EdgePointer &ep : pMesh->nodes[v].incidentElements) {
      const size_t ei = pMesh->getIndex(ep);
      const Element &elem = pMesh->elements[ep];
      const double SkTranslational =
          elem.material.youngsModulus * elem.A / elem.length;
      translationalSumSk += SkTranslational;
      const double lengthToThe3 = std::pow(elem.length, 3);
      const long double SkRotational_I2 =
          elem.material.youngsModulus * elem.I2 /
          lengthToThe3; // TODO: I2->t2,I3->t3,t1->polar inertia
      const long double SkRotational_I3 =
          elem.material.youngsModulus * elem.I3 /
          lengthToThe3; // TODO: I2->t2,I3->t3,t1->polar inertia
      const long double SkRotational_J =
          elem.material.youngsModulus * elem.J /
          lengthToThe3; // TODO: I2->t2,I3->t3,t1->polar inertia
      rotationalSumSk_I2 += SkRotational_I2;
      rotationalSumSk_I3 += SkRotational_I3;
      rotationalSumSk_J += SkRotational_J;
      assert(rotationalSumSk_I2 != 0);
      assert(rotationalSumSk_I3 != 0);
      assert(rotationalSumSk_J != 0);
    }
    pMesh->nodes[v].translationalMass =
        gamma * pow(mSettings.Dtini, 2) * 2 * translationalSumSk;
    pMesh->nodes[v].rotationalMass_I2 =
        gamma * pow(mSettings.Dtini, 2) * 8 * rotationalSumSk_I2;
    pMesh->nodes[v].rotationalMass_I3 =
        gamma * pow(mSettings.Dtini, 2) * 8 * rotationalSumSk_I3;
    pMesh->nodes[v].rotationalMass_J =
        gamma * pow(mSettings.Dtini, 2) * 8 * rotationalSumSk_J;

    assert(std::pow(mSettings.Dtini, 2.0) * translationalSumSk /
               pMesh->nodes[v].translationalMass <
           2);
    assert(std::pow(mSettings.Dtini, 2.0) * rotationalSumSk_I2 /
               pMesh->nodes[v].rotationalMass_I2 <
           2);
    assert(std::pow(mSettings.Dtini, 2.0) * rotationalSumSk_I3 /
               pMesh->nodes[v].rotationalMass_I3 <
           2);
    assert(std::pow(mSettings.Dtini, 2.0) * rotationalSumSk_J /
               pMesh->nodes[v].rotationalMass_J <
           2);
  }
}

void FormFinder::updateNodalAccelerations() {
  for (VertexType &v : pMesh->vert) {
    Node &node = pMesh->nodes[v];
    const VertexIndex vi = pMesh->getIndex(v);
    for (DoFType dofi = DoF::Ux; dofi < DoF::NumDoF; dofi++) {
      if (dofi == DoF::Ux || dofi == DoF::Uy || dofi == DoF::Uz) {
        node.acceleration[dofi] =
            node.force.residual[dofi] / node.translationalMass;
      } else if (dofi == DoF::Nx) {
        node.acceleration[dofi] =
            node.force.residual[dofi] / node.rotationalMass_J;
      } else if (dofi == DoF::Ny) {
        node.acceleration[dofi] =
            node.force.residual[dofi] / node.rotationalMass_I3;
      } else if (dofi == DoF::Nr) {
        node.acceleration[dofi] =
            node.force.residual[dofi] / node.rotationalMass_I2;
      }
    }
  }
}

void FormFinder::updateNodalVelocities() {
  for (VertexType &v : pMesh->vert) {
    const VertexIndex vi = pMesh->getIndex(v);
    Node &node = pMesh->nodes[v];
    node.velocity = node.velocity + node.acceleration * Dt;
  }
  updateKineticEnergy();
}

void FormFinder::updateNodalDisplacements() {
  for (VertexType &v : pMesh->vert) {
    Node &node = pMesh->nodes[v];
    node.displacements = node.displacements + node.velocity * Dt;
    //    if (mSettings.beVerbose) {
    //      std::cout << "Node " << node.vi << ":" << endl;
    //      std::cout << node.displacements[0] << " " << node.displacements[0]
    //      << " "
    //                << node.displacements[1] << " " << node.displacements[2]
    //                << " "
    //                << node.displacements[3] << " " << node.displacements[4]
    //                << " "
    //                << node.displacements[5] << std::endl;
    //    }
  }
}

void FormFinder::updateNodePosition(
    VertexType &v,
    const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
        &fixedVertices) {
  Node &node = pMesh->nodes[v];
  CoordType previousLocation = v.cP();
  const VertexIndex &vi = pMesh->nodes[v].vi;

  VectorType displacementVector(0, 0, 0);
  if (!fixedVertices.contains(vi) || !fixedVertices.at(vi).contains(0)) {
    displacementVector += VectorType(node.displacements[0], 0, 0);
  }
  if (!fixedVertices.contains(vi) || !fixedVertices.at(vi).contains(1)) {
    displacementVector += VectorType(0, node.displacements[1], 0);
  }
  if (!fixedVertices.contains(vi) || !fixedVertices.at(vi).contains(2)) {
    displacementVector += VectorType(0, 0, node.displacements[2]);
  }
  node.previousLocation = previousLocation;
  v.P() = node.initialLocation + displacementVector;
  if (shouldApplyInitialDistortion && mCurrentSimulationStep < 40) {
    VectorType desiredInitialDisplacement(0, 0, 0.01);
    v.P() = v.P() + desiredInitialDisplacement;
  }
}

void FormFinder::updateNodeNormal(
    VertexType &v,
    const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
        &fixedVertices) {

  Node &node = pMesh->nodes[v];
  const VertexIndex &vi = node.vi;
  //  if (vi == 1) {
  //    std::cout << "PRE:" << mesh->vert[1].N()[0] << " " <<
  //    mesh->vert[1].N()[1]
  //              << " " << mesh->vert[1].N()[2] << std::endl;
  //  }
  VectorType normalDisplacementVector(0, 0, 0);
  if (!fixedVertices.contains(vi) || !fixedVertices.at(vi).contains(3)) {
    normalDisplacementVector += VectorType(node.displacements[3], 0, 0);
  }
  if (!fixedVertices.contains(vi) || !fixedVertices.at(vi).contains(4)) {
    normalDisplacementVector += VectorType(0, node.displacements[4], 0);
  }
  const double nxnyMagnitudePre = std::pow(v.N()[0], 2) + std::pow(v.N()[1], 2);
  v.N() = node.initialNormal + normalDisplacementVector;
  const double &nx = v.N()[0], ny = v.N()[1], nz = v.N()[2];
  const double nxnyMagnitude = std::pow(nx, 2) + std::pow(ny, 2);
  VectorType n = v.N();
  const bool shouldBreak = mCurrentSimulationStep == 118 && vi == 3;
  if (nxnyMagnitude > 1) {
    VectorType newNormal(nx / std::sqrt(nxnyMagnitude),
                         ny / std::sqrt(nxnyMagnitude), 0);
    v.N() = newNormal;

    /*If an external moment caused the normal to lay on the xy plane this means
     * that in order to disable its effect a greater internal force is needed
     * than what is possible (the constraint on the z of the normals imposes a
     * constraint on the maximum internal force). Because of that the
     * totalResidualForcesNorm can't drop below the magnitude of external moment
     * applied on vertex vi. In order to allow termination of the simulation
     * when the described phenomenon happens we allow the termination of the
     * algorithm if the kinetic energy of the system drops below the set
     * threshold.
     * */
    const bool viHasMoments =
        node.force.external[3] != 0 || node.force.external[4] != 0;
    if (!checkedForMaximumMoment && viHasMoments) {
      mSettings.shouldUseTranslationalKineticEnergyThreshold = true;
      if (mSettings.beVerbose) {
        std::cout
            << "Maximum moment reached.The Kinetic energy of the system will "
               "be used as a convergence criterion"
            << std::endl;
      }
      checkedForMaximumMoment = true;
    }

  } else {
    const double nzSquared = 1.0 - nxnyMagnitude;
    const double nz = std::sqrt(nzSquared);
    VectorType newNormal(nx, ny, nz);
    v.N() = newNormal;
  }
  if (!rigidSupports.contains(vi)) {
    node.nR = node.displacements[5];
  } else {
    const EdgePointer &refElem = node.referenceElement;
    const VectorType &refT1 = pMesh->elements[refElem].frame.t1;

    const VectorType &t1Initial =
        computeT1Vector(pMesh->nodes[refElem->cV(0)].initialLocation,
                        pMesh->nodes[refElem->cV(1)].initialLocation);
    VectorType g1 = Cross(refT1, t1Initial);
    node.nR = g1.dot(v.cN());
  }
  //  if (vi == 1) {
  //    std::cout << "AFTER:" << mesh->vert[1].N()[0] << " " <<
  //    mesh->vert[1].N()[1]
  //              << " " << mesh->vert[1].N()[2] << std::endl;
  //  }
}

void FormFinder::applyDisplacements(
    const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
        &fixedVertices) {
  for (VertexType &v : pMesh->vert) {
    updateNodePosition(v, fixedVertices);
    updateNodeNormal(v, fixedVertices);
  }
  updateElementalFrames();
  if (mSettings.shouldDraw) {
    pMesh->updateEigenEdgeAndVertices();
  }
}

void FormFinder::updateElementalFrames() {
  for (const EdgeType &e : pMesh->edge) {
    const VectorType elementNormal =
        (e.cV(0)->cN() + e.cV(1)->cN()).Normalize();
    pMesh->elements[e].frame =
        computeElementFrame(e.cP(0), e.cP(1), elementNormal);
  }
}

void FormFinder::applyForcedDisplacements(
    const std::unordered_map<VertexIndex, Eigen::Vector3d>
        nodalForcedDisplacements) {
  for (const std::pair<VertexIndex, Eigen::Vector3d>
           vertexIndexDisplacementPair : nodalForcedDisplacements) {
    const VertexIndex vi = vertexIndexDisplacementPair.first;
    const Eigen::Vector3d vertexDisplacement =
        vertexIndexDisplacementPair.second;
    Node &node = pMesh->nodes[vi];
    VectorType displacementVector(vertexDisplacement(0), vertexDisplacement(1),
                                  vertexDisplacement(2));
    pMesh->vert[vi].P() = node.initialLocation + displacementVector;
    node.displacements = Vector6d(
        {vertexDisplacement(0), vertexDisplacement(1), vertexDisplacement(2),
         node.displacements[3], node.displacements[4], node.displacements[5]});
  }

  if (mSettings.shouldDraw) {
    pMesh->updateEigenEdgeAndVertices();
  }
}

void FormFinder::applyForcedNormals(
    const std::unordered_map<VertexIndex, VectorType> nodalForcedRotations) {
  for (const std::pair<VertexIndex, VectorType> vertexIndexDesiredNormalPair :
       nodalForcedRotations) {
    const VertexIndex vi = vertexIndexDesiredNormalPair.first;

    Node &node = pMesh->nodes[vi];
    pMesh->vert[vi].N() = vertexIndexDesiredNormalPair.second;
    node.displacements = Vector6d(
        {node.displacements[0], node.displacements[1], node.displacements[2],
         vertexIndexDesiredNormalPair.second[0] - node.initialNormal[0],
         vertexIndexDesiredNormalPair.second[1] - node.initialNormal[1],
         node.displacements[5]});
  }
}

// NOTE: Is this correct? Should the kinetic energy be computed like that?
void FormFinder::updateKineticEnergy() {
  pMesh->previousTotalKineticEnergy = pMesh->currentTotalKineticEnergy;
  pMesh->currentTotalKineticEnergy = 0;
  pMesh->currentTotalTranslationalKineticEnergy = 0;
  for (const VertexType &v : pMesh->vert) {
    Node &node = pMesh->nodes[v];
    node.kineticEnergy = 0;

    const double translationalVelocityNorm = std::sqrt(
        std::pow(node.velocity[0], 2) + std::pow(node.velocity[1], 2) +
        std::pow(node.velocity[2], 2));
    const double nodeTranslationalKineticEnergy =
        0.5 * node.translationalMass * pow(translationalVelocityNorm, 2);

    const double nodeRotationalKineticEnergy =
        0.5 * (node.rotationalMass_J * pow(node.velocity[3], 2) +
               +node.rotationalMass_I3 * pow(node.velocity[4], 2) +
               +node.rotationalMass_I2 * pow(node.velocity[5], 2));

    node.kineticEnergy +=
        nodeTranslationalKineticEnergy /*+ nodeRotationalKineticEnergy*/;
    assert(node.kineticEnergy < 10000000000000000000);

    pMesh->currentTotalKineticEnergy += node.kineticEnergy;
    pMesh->currentTotalTranslationalKineticEnergy +=
        nodeTranslationalKineticEnergy;
  }
  //  assert(mesh->currentTotalKineticEnergy < 100000000000000);
}

void FormFinder::resetVelocities() {
  for (const VertexType &v : pMesh->vert) {
    pMesh->nodes[v].velocity = 0;
    //            mesh->nodes[v].force.residual * 0.5 * Dt /
    //            mesh->nodes[v].mass; // NOTE: Do I reset the previous
    //            velocity?
    //        reset
    // current to 0 or this?
  }
  updateKineticEnergy();
}

void FormFinder::updatePositionsOnTheFly(
    const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
        &fixedVertices) {
  const double gamma = 0.8;
  for (VertexType &v : pMesh->vert) {
    double translationalSumSk = 0;
    double rotationalSumSk_I2 = 0;
    double rotationalSumSk_I3 = 0;
    double rotationalSumSk_J = 0;
    for (const EdgePointer &ep : pMesh->nodes[v].incidentElements) {
      const Element &elem = pMesh->elements[ep];
      const double SkTranslational =
          elem.material.youngsModulus * elem.A / elem.length;
      translationalSumSk += SkTranslational;
      const double lengthToThe3 = std::pow(elem.length, 3);
      const double SkRotational_I2 =
          elem.material.youngsModulus * elem.I2 /
          lengthToThe3; // TODO: I2->t2,I3->t3,t1->polar inertia
      const double SkRotational_I3 =
          elem.material.youngsModulus * elem.I3 /
          lengthToThe3; // TODO: I2->t2,I3->t3,t1->polar inertia
      const double SkRotational_J =
          elem.material.youngsModulus * elem.J /
          lengthToThe3; // TODO: I2->t2,I3->t3,t1->polar inertia
      rotationalSumSk_I2 += SkRotational_I2;
      rotationalSumSk_I3 += SkRotational_I3;
      rotationalSumSk_J += SkRotational_J;
      //      assert(rotationalSumSk_I2 != 0);
      //      assert(rotationalSumSk_I3 != 0);
      //      assert(rotationalSumSk_J != 0);
    }
    pMesh->nodes[v].translationalMass =
        gamma * pow(mSettings.Dtini, 2) * 2 * translationalSumSk;
    pMesh->nodes[v].rotationalMass_I2 =
        gamma * pow(mSettings.Dtini, 2) * 8 * rotationalSumSk_I2;
    pMesh->nodes[v].rotationalMass_I3 =
        gamma * pow(mSettings.Dtini, 2) * 8 * rotationalSumSk_I3;
    pMesh->nodes[v].rotationalMass_J =
        gamma * pow(mSettings.Dtini, 2) * 8 * rotationalSumSk_J;

    //    assert(std::pow(mSettings.Dtini, 2.0) * translationalSumSk /
    //               mesh->nodes[v].translationalMass <
    //           2);
    //    assert(std::pow(mSettings.Dtini, 2.0) * rotationalSumSk_I2 /
    //               mesh->nodes[v].rotationalMass_I2 <
    //           2);
    //    assert(std::pow(mSettings.Dtini, 2.0) * rotationalSumSk_I3 /
    //               mesh->nodes[v].rotationalMass_I3 <
    //           2);
    //    assert(std::pow(mSettings.Dtini, 2.0) * rotationalSumSk_J /
    //               mesh->nodes[v].rotationalMass_J <
    //           2);
  }

  for (VertexType &v : pMesh->vert) {
    Node &node = pMesh->nodes[v];
    for (DoFType dofi = DoF::Ux; dofi < DoF::NumDoF; dofi++) {
      if (dofi == DoF::Ux || dofi == DoF::Uy || dofi == DoF::Uz) {
        node.acceleration[dofi] =
            node.force.residual[dofi] / node.translationalMass;
      } else if (dofi == DoF::Nx) {
        node.acceleration[dofi] =
            node.force.residual[dofi] / node.rotationalMass_J;
      } else if (dofi == DoF::Ny) {
        node.acceleration[dofi] =
            node.force.residual[dofi] / node.rotationalMass_I3;
      } else if (dofi == DoF::Nr) {
        node.acceleration[dofi] =
            node.force.residual[dofi] / node.rotationalMass_I2;
      }
    }
  }

  for (VertexType &v : pMesh->vert) {
    Node &node = pMesh->nodes[v];
    node.velocity = node.velocity + node.acceleration * Dt;
  }
  updateKineticEnergy();

  for (VertexType &v : pMesh->vert) {
    Node &node = pMesh->nodes[v];
    node.displacements = node.displacements + node.velocity * Dt;
  }

  for (VertexType &v : pMesh->vert) {
    updateNodePosition(v, fixedVertices);
    Node &node = pMesh->nodes[v];
    const VertexIndex &vi = node.vi;
    VectorType normalDisplacementVector(0, 0, 0);
    if (!fixedVertices.contains(vi) || !fixedVertices.at(vi).contains(3)) {
      normalDisplacementVector += VectorType(node.displacements[3], 0, 0);
    }
    if (!fixedVertices.contains(vi) || !fixedVertices.at(vi).contains(4)) {
      normalDisplacementVector += VectorType(0, node.displacements[4], 0);
    }
    v.N() = node.initialNormal + normalDisplacementVector;
    const double &nx = v.N()[0], ny = v.N()[1];
    const double nxnyMagnitude = std::pow(nx, 2) + std::pow(ny, 2);
    if (nxnyMagnitude > 1) {
      v.N() = VectorType(nx / std::sqrt(nxnyMagnitude),
                         ny / std::sqrt(nxnyMagnitude), 0);
    } else {
      const double nzSquared = 1.0 - nxnyMagnitude;
      const double nz = std::sqrt(nzSquared);
      VectorType newNormal(nx, ny, nz);
      v.N() = newNormal;
    }
    if (!rigidSupports.contains(vi)) {
      node.nR = node.displacements[5];
    } else {
    }
  }
  updateElementalFrames();
}

SimulationResults
FormFinder::computeResults(const std::shared_ptr<SimulationJob> &pJob) {
  std::vector<Vector6d> displacements(pMesh->VN());
  for (size_t vi = 0; vi < pMesh->VN(); vi++) {
    displacements[vi] = pMesh->nodes[vi].displacements;
  }
  history.numberOfSteps = mCurrentSimulationStep;

  return SimulationResults{true, pJob, history, displacements};
}

void FormFinder::draw(const std::string &screenshotsFolder = {}) {
  // update positions
  //  polyscope::getCurveNetwork("Undeformed edge mesh")->setEnabled(false);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->updateNodePositions(pMesh->getEigenVertices());

  // Vertex quantities
  std::vector<double> kineticEnergies(pMesh->VN());
  std::vector<std::array<double, 3>> nodeNormals(pMesh->VN());
  std::vector<std::array<double, 3>> internalForces(pMesh->VN());
  std::vector<std::array<double, 3>> externalForces(pMesh->VN());
  std::vector<std::array<double, 3>> externalMoments(pMesh->VN());
  std::vector<double> internalForcesNorm(pMesh->VN());
  std::vector<std::array<double, 3>> internalAxialForces(pMesh->VN());
  std::vector<std::array<double, 3>> internalFirstBendingTranslationForces(
      pMesh->VN());
  std::vector<std::array<double, 3>> internalFirstBendingRotationForces(
      pMesh->VN());
  std::vector<std::array<double, 3>> internalSecondBendingTranslationForces(
      pMesh->VN());
  std::vector<std::array<double, 3>> internalSecondBendingRotationForces(
      pMesh->VN());
  std::vector<double> nRs(pMesh->VN());
  std::vector<double> theta2(pMesh->VN());
  std::vector<double> theta3(pMesh->VN());
  std::vector<std::array<double, 3>> residualForces(pMesh->VN());
  std::vector<double> residualForcesNorm(pMesh->VN());
  std::vector<double> accelerationX(pMesh->VN());
  for (const VertexType &v : pMesh->vert) {
    kineticEnergies[pMesh->getIndex(v)] = pMesh->nodes[v].kineticEnergy;
    const VectorType n = v.cN();
    nodeNormals[pMesh->getIndex(v)] = {n[0], n[1], n[2]};
    // per node internal forces
    const Vector6d nodeforce = pMesh->nodes[v].force.internal * (-1);
    internalForces[pMesh->getIndex(v)] = {nodeforce[0], nodeforce[1],
                                          nodeforce[2]};
    internalForcesNorm[pMesh->getIndex(v)] = nodeforce.norm();
    // External force
    const Vector6d nodeExternalForce = pMesh->nodes[v].force.external;
    externalForces[pMesh->getIndex(v)] = {
        nodeExternalForce[0], nodeExternalForce[1], nodeExternalForce[2]};
    externalMoments[pMesh->getIndex(v)] = {nodeExternalForce[3],
                                           nodeExternalForce[4], 0};
    internalAxialForces[pMesh->getIndex(v)] = {nodeforce[0], nodeforce[1],
                                               nodeforce[2]};
    const Node &node = pMesh->nodes[v];
    const Vector6d nodeForceFirst = node.force.internalFirstBending * (-1);
    internalFirstBendingTranslationForces[pMesh->getIndex(v)] = {
        nodeForceFirst[0], nodeForceFirst[1], nodeForceFirst[2]};
    internalFirstBendingRotationForces[pMesh->getIndex(v)] = {
        nodeForceFirst[3], nodeForceFirst[4], 0};

    const Vector6d nodeForceSecond = node.force.internalSecondBending * (-1);
    internalSecondBendingTranslationForces[pMesh->getIndex(v)] = {
        nodeForceSecond[0], nodeForceSecond[1], nodeForceSecond[2]};
    internalSecondBendingRotationForces[pMesh->getIndex(v)] = {
        nodeForceSecond[3], nodeForceSecond[4], 0};
    nRs[pMesh->getIndex(v)] = node.nR;
    const Vector6d nodeResidualForce = pMesh->nodes[v].force.residual;
    residualForces[pMesh->getIndex(v)] = {
        nodeResidualForce[0], nodeResidualForce[1], nodeResidualForce[2]};
    residualForcesNorm[pMesh->getIndex(v)] = nodeResidualForce.norm();
    accelerationX[pMesh->getIndex(v)] = pMesh->nodes[v].acceleration[0];
  }
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addNodeScalarQuantity("Kinetic Energy", kineticEnergies)
      ->setEnabled(false);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addNodeVectorQuantity("Node normals", nodeNormals)
      ->setEnabled(true);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addNodeVectorQuantity("Internal force", internalForces)
      ->setEnabled(false);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addNodeVectorQuantity("External force", externalForces)
      ->setEnabled(true);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addNodeVectorQuantity("External Moments", externalMoments)
      ->setEnabled(true);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addNodeScalarQuantity("Internal force norm", internalForcesNorm)
      ->setEnabled(true);
  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeVectorQuantity("Internal Axial force", internalAxialForces)
  //      ->setEnabled(false);
  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeVectorQuantity("First bending force-Translation",
  //                              internalFirstBendingTranslationForces)
  //      ->setEnabled(false);
  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeVectorQuantity("First bending force-Rotation",
  //                              internalFirstBendingRotationForces)
  //      ->setEnabled(false);

  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeVectorQuantity("Second bending force-Translation",
  //                              internalSecondBendingTranslationForces)
  //      ->setEnabled(false);
  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeVectorQuantity("Second bending force-Rotation",
  //                              internalSecondBendingRotationForces)
  //      ->setEnabled(false);
  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeScalarQuantity("nR", nRs)
  //      ->setEnabled(false);
  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeScalarQuantity("theta3", theta3)
  //      ->setEnabled(false);
  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeScalarQuantity("theta2", theta2)
  //      ->setEnabled(false);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addNodeVectorQuantity("Residual force", residualForces)
      ->setEnabled(false);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addNodeScalarQuantity("Residual force norm", residualForcesNorm)
      ->setEnabled(false);
  //  polyscope::getCurveNetwork(meshPolyscopeLabel)
  //      ->addNodeScalarQuantity("Node acceleration x", accelerationX);

  // Edge quantities
  std::vector<double> A(pMesh->EN());
  std::vector<double> J(pMesh->EN());
  std::vector<double> I2(pMesh->EN());
  std::vector<double> I3(pMesh->EN());
  for (const EdgeType &e : pMesh->edge) {
    const size_t ei = pMesh->getIndex(e);
    A[ei] = pMesh->elements[e].A;
    J[ei] = pMesh->elements[e].J;
    I2[ei] = pMesh->elements[e].I2;
    I3[ei] = pMesh->elements[e].I3;
  }

  polyscope::getCurveNetwork(meshPolyscopeLabel)->addEdgeScalarQuantity("A", A);
  polyscope::getCurveNetwork(meshPolyscopeLabel)->addEdgeScalarQuantity("J", J);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addEdgeScalarQuantity("I2", I2);
  polyscope::getCurveNetwork(meshPolyscopeLabel)
      ->addEdgeScalarQuantity("I3", I3);

  // Specify the callback
  polyscope::state::userCallback = [&]() {
    // Since options::openImGuiWindowForUserCallback == true by default,
    // we can immediately start using ImGui commands to build a UI

    ImGui::PushItemWidth(100); // Make ui elements 100 pixels wide,
                               // instead of full width. Must have
                               // matching PopItemWidth() below.

    ImGui::InputInt("Simulation drawing step",
                    &mSettings.drawingStep); // set a int variable
    ImGui::Checkbox("Enable drawing",
                    &mSettings.shouldDraw); // set a int variable
    ImGui::Text("Number of simulation steps: %zu", mCurrentSimulationStep);

    ImGui::PopItemWidth();
  };

  if (!screenshotsFolder.empty()) {
    static bool firstDraw = true;
    if (firstDraw) {
      for (const auto &entry :
           std::filesystem::directory_iterator(screenshotsFolder))
        std::filesystem::remove_all(entry.path());
      //      polyscope::view::processZoom(5);
      firstDraw = false;
    }
    polyscope::screenshot(
        std::filesystem::path(screenshotsFolder)
            .append(std::to_string(mCurrentSimulationStep) + ".png")
            .string(),
        false);
  } else {
    polyscope::show();
  }
}

void FormFinder::printCurrentState() const {
  std::cout << "Simulation steps executed:" << mCurrentSimulationStep
            << std::endl;
  std::cout << "Residual forces norm: " << pMesh->totalResidualForcesNorm
            << std::endl;
  std::cout << "Kinetic energy:" << pMesh->currentTotalKineticEnergy
            << std::endl;
}

void FormFinder::printDebugInfo() const {
  std::cout << pMesh->elements[0].rigidity.toString() << std::endl;
  std::cout << "Number of dampings:" << numOfDampings << endl;
  printCurrentState();
}

SimulationResults
FormFinder::executeSimulation(const std::shared_ptr<SimulationJob> &pJob,
                              const Settings &settings) {
  assert(pJob->pMesh.operator bool());
  auto t1 = std::chrono::high_resolution_clock::now();
  reset();
  mSettings = settings;

  //  if (!pJob->nodalExternalForces.empty()) {
  //    double externalForcesNorm = 0;
  //    for (const auto &externalForce : pJob->nodalExternalForces) {
  //      externalForcesNorm += externalForce.second.norm();
  //    }
  //    mSettings.totalResidualForcesNormThreshold = externalForcesNorm * 1e-2;
  //  }

  constrainedVertices = pJob->constrainedVertices;
  if (!pJob->nodalForcedDisplacements.empty()) {
    for (std::pair<VertexIndex, Eigen::Vector3d> viDisplPair :
         pJob->nodalForcedDisplacements) {
      const VertexIndex vi = viDisplPair.first;
      constrainedVertices[vi].insert({0, 1, 2});
    }
  }
  //  if (!pJob->nodalForcedNormals.empty()) {
  //    for (std::pair<VertexIndex, Eigen::Vector3d> viNormalPair :
  //         pJob->nodalForcedDisplacements) {
  //      assert(viNormalPair3second[2] >= 0);
  //    }
  //  }

  pMesh.reset();
  pMesh = std::make_unique<SimulationMesh>(*pJob->pMesh);
  if (mSettings.beVerbose  ) {
    std::cout << "Executing simulation for mesh with:" << pMesh->VN()
              << " nodes and " << pMesh->EN() << " elements." << std::endl;
  }
  computeRigidSupports();

  if (mSettings.shouldDraw ) {
    initPolyscope();
    polyscope::registerCurveNetwork(
        meshPolyscopeLabel, pMesh->getEigenVertices(), pMesh->getEigenEdges());
    //    registerWorldAxes();
  }
  for (auto fixedVertex : pJob->constrainedVertices) {
    assert(fixedVertex.first < pMesh->VN());
  }

  updateElementalFrames();
  updateNodalMasses();
  if (!pJob->nodalForcedDisplacements.empty() &&
      pJob->nodalExternalForces.empty()) {
    shouldApplyInitialDistortion = true;
  }
  updateNodalExternalForces(pJob->nodalExternalForces, constrainedVertices);
  while (settings.maxDRMIterations == 0 ||
         mCurrentSimulationStep < settings.maxDRMIterations) {
    //  while (true) {
    updateNormalDerivatives();
    updateT1Derivatives();
    updateRDerivatives();
    updateT2Derivatives();
    updateT3Derivatives();
    updateResidualForcesOnTheFly(constrainedVertices);

    // TODO: write parallel function for updating positions
    // TODO: Make a single function out of updateResidualForcesOnTheFly
    // updatePositionsOnTheFly
    //    updatePositionsOnTheFly(constrainedVertices);
    updateNodalMasses();
    updateNodalAccelerations();
    updateNodalVelocities();
    updateNodalDisplacements();
    applyDisplacements(constrainedVertices);
    if (!pJob->nodalForcedDisplacements.empty()) {
      applyForcedDisplacements(

          pJob->nodalForcedDisplacements);
    }
    //    if (!pJob->nodalForcedNormals.empty()) {
    //      applyForcedNormals(pJob->nodalForcedNormals);
    //    }
    updateElementalLengths();
    pMesh->previousTotalPotentialEnergykN =
        pMesh->currentTotalPotentialEnergykN;
    pMesh->currentTotalPotentialEnergykN = computeTotalPotentialEnergy() / 1000;

    if (mCurrentSimulationStep != 0) {
      history.stepPulse(*pMesh);
    }

    if (std::isnan(pMesh->currentTotalKineticEnergy)) {
      std::time_t now = time(0);
      std::tm *ltm = std::localtime(&now);
      std::string dir = "../ProblematicSimulationJobs/" +
                        std::to_string(ltm->tm_mday) + "_" +
                        std::to_string(1 + ltm->tm_mon) + "_" +
                        std::to_string(1900 + ltm->tm_year);
      const std::string subDir = std::filesystem::path(dir)
                                     .append(std::to_string(ltm->tm_hour) +
                                             ":" + std::to_string(ltm->tm_min))
                                     .string();
      std::filesystem::create_directories(subDir);
      pJob->save(subDir);
      std::cout << "Infinite kinetic energy detected.Exiting.." << std::endl;
      FormFinder debug;
      FormFinder::Settings settings;
      settings.shouldDraw = true;
      settings.beVerbose = true;
      debug.executeSimulation(pJob, settings);
      std::terminate();
      break;
    }

    if (mSettings.shouldDraw &&
        mCurrentSimulationStep % mSettings.drawingStep == 0) /* &&
currentSimulationStep > maxDRMIterations*/
    {
      //      std::string saveTo = std::filesystem::current_path()
      //                               .append("Debugging_files")
      //                               .append("Screenshots")
      //                               .string();
      //      draw(saveTo);
      draw();

      //      yValues = history.kineticEnergy;
      //      plotHandle = matplot::scatter(xPlot, yValues);
      //      label = "Log of Kinetic energy";
      //      plotHandle->legend_string(label);

      //      shouldUseKineticEnergyThreshold = true;
    }
    if (mSettings.shouldCreatePlots && mCurrentSimulationStep % 10 == 0) {
      printCurrentState();
      SimulationResultsReporter::createPlot(
          "Number of Steps", "Log of Kinetic energy", history.kineticEnergy);
    }
    //    t = t + Dt;
    mCurrentSimulationStep++;
    //    std::cout << "Kinetic energy:" << mesh.currentTotalKineticEnergy
    //              << std::endl;
    //    std::cout << "Residual forces norm:" << mesh.totalResidualForcesNorm
    //              << std::endl;
    // Kinetic energy convergence
    if ((mSettings.shouldUseTranslationalKineticEnergyThreshold ||
         mCurrentSimulationStep > 100000) &&
        pMesh->currentTotalTranslationalKineticEnergy <
            mSettings.totalTranslationalKineticEnergyThreshold) {
      if (mSettings.beVerbose) {
        std::cout << "Simulation converged." << std::endl;
        printCurrentState();
        std::cout << "Total potential:" << pMesh->currentTotalPotentialEnergykN
                  << " kNm" << std::endl;
        std::cout << "Warning: The kinetic energy of the system was "
                     " used as a convergence criterion"
                  << std::endl;
      }
      break;
    }
    // Residual forces norm convergence
    if (pMesh->previousTotalKineticEnergy > pMesh->currentTotalKineticEnergy
        /*||
mesh->previousTotalPotentialEnergykN >
mesh->currentTotalPotentialEnergykN*/
        /*|| mesh->currentTotalPotentialEnergykN < minPotentialEnergy*/) {
      if (pMesh->totalResidualForcesNorm <
          mSettings.totalResidualForcesNormThreshold) {
        if (mSettings.beVerbose) {
          std::cout << "Simulation converged." << std::endl;
          printCurrentState();
          std::cout << "Total potential:"
                    << pMesh->currentTotalPotentialEnergykN << " kNm"
                    << std::endl;
        }
        break;
        //        }
      }
      //      for (VertexType &v : mesh->vert) {
      //        Node &node = mesh->nodes[v];
      //        node.displacements = node.displacements - node.velocity * Dt;
      //      }
      //      applyDisplacements(constrainedVertices);
      //      if (!pJob->nodalForcedDisplacements.empty()) {
      //        applyForcedDisplacements(

      //            pJob->nodalForcedDisplacements);
      //      }
      //      updateElementalLengths();
      resetVelocities();
      Dt = Dt * mSettings.xi;
      ++numOfDampings;
    }

    if (mSettings.debugMessages) {
      printDebugInfo();
    }
  }

  SimulationResults results = computeResults(pJob);

  if (mCurrentSimulationStep == settings.maxDRMIterations &&
      mCurrentSimulationStep != 0) {
    std::cout << "Did not reach equilibrium before reaching the maximum number "
                 "of DRM steps ("
              << settings.maxDRMIterations << "). Breaking simulation"
              << std::endl;
    results.converged = false;
  } else if (std::isnan(pMesh->currentTotalKineticEnergy)) {
    results.converged = false;

  } else if (mSettings.beVerbose ) {
    auto t2 = std::chrono::high_resolution_clock::now();
    double simulationDuration =
        std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1).count();
    simulationDuration /= 1000;
    std::cout << "Simulation converged after " << simulationDuration << "s"
              << std::endl;
    std::cout << "Time/(node*iteration) "
              << simulationDuration /
                     (static_cast<double>(mCurrentSimulationStep) * pMesh->VN())
              << "s" << std::endl;
    std::cout << "Number of dampings:" << numOfDampings << endl;
  }
  //    mesh.printVertexCoordinates(mesh.VN() / 2);
  if (results.converged) {
    if (mSettings.shouldCreatePlots) {
      SimulationResultsReporter reporter;
      reporter.reportResults({results}, "Results", pJob->pMesh->getLabel());
    }
    if (mSettings.shouldDraw) {
      draw();
    }
  }
  if (polyscope::hasCurveNetwork(meshPolyscopeLabel)) {
    polyscope::removeCurveNetwork(meshPolyscopeLabel);
  }
  return results;
}