#ifndef SIMULATIONHISTORY_HPP
#define SIMULATIONHISTORY_HPP

#include "elementalmesh.hpp"

struct SimulationHistory {
  SimulationHistory() {}

  size_t numberOfSteps{0};
  std::string label;
  std::vector<double> residualForces;
  std::vector<double> kineticEnergy;
  std::vector<double> potentialEnergies;
  std::vector<size_t> redMarks;
  std::vector<double> greenMarks;

  void markRed(const size_t &stepNumber) { redMarks.push_back(stepNumber); }

  void markGreen(const size_t &stepNumber) { greenMarks.push_back(stepNumber); }

  void stepPulse(const SimulationMesh &mesh) {
    kineticEnergy.push_back(log(mesh.currentTotalKineticEnergy));
    //    potentialEnergy.push_back(mesh.totalPotentialEnergykN);
    residualForces.push_back(mesh.totalResidualForcesNorm);
  }

  void clear() {
    residualForces.clear();
    kineticEnergy.clear();
    potentialEnergies.clear();
  }
};

struct SimulationJob {
  std::shared_ptr<SimulationMesh> mesh;
  std::unordered_map<VertexIndex, std::unordered_set<int>> fixedVertices;
  std::unordered_map<VertexIndex, Vector6d> nodalExternalForces;
  std::unordered_map<VertexIndex, Eigen::Vector3d> nodalForcedDisplacements;
  std::unordered_map<VertexIndex, VectorType> nodalForcedNormals;

  void registerForDrawing() const {
    initPolyscope();
    const std::string polyscopeName = mesh->getLabel() + "_Simulation Job";
    polyscope::registerCurveNetwork(polyscopeName, mesh->getEigenVertices(),
                                    mesh->getEigenEdges());
    //  polyscope::registerPointCloud("Undeformed edge mesh PC",
    //                                job.edgeMesh.getEigenVertices());

    std::vector<std::array<double, 3>> nodeColors(mesh->VN());
    for (auto fixedVertex : fixedVertices) {
      nodeColors[fixedVertex.first] = {0, 0, 1};
    }
    if (!nodalForcedDisplacements.empty()) {
      for (std::pair<VertexIndex, Eigen::Vector3d> viDisplPair :
           nodalForcedDisplacements) {
        const VertexIndex vi = viDisplPair.first;
        nodeColors[vi][0] += 1;
        nodeColors[vi][0] /= 2;
        nodeColors[vi][1] += 0;
        nodeColors[vi][1] /= 2;
        nodeColors[vi][2] += 0;
        nodeColors[vi][2] /= 2;
      }
    }
    std::for_each(nodeColors.begin(), nodeColors.end(),
                  [](std::array<double, 3> &color) {
                    const double norm =
                        sqrt(std::pow(color[0], 2) + std::pow(color[1], 2) +
                             std::pow(color[2], 2));
                    if (norm > std::pow(10, -7)) {
                      color[0] /= norm;
                      color[1] /= norm;
                      color[2] /= norm;
                    }
                  });

    if (!nodeColors.empty()) {
      polyscope::getCurveNetwork(polyscopeName)
          ->addNodeColorQuantity("Boundary conditions", nodeColors);
      polyscope::getCurveNetwork(polyscopeName)
          ->getQuantity("Boundary conditions")
          ->setEnabled(true);
    }

    // per node external forces
    std::vector<std::array<double, 3>> externalForces(mesh->VN());
    for (const auto &forcePair : nodalExternalForces) {
      auto index = forcePair.first;
      auto force = forcePair.second;
      externalForces[index] = {force[0], force[1], force[2]};
    }

    if (!externalForces.empty()) {
      polyscope::getCurveNetwork(polyscopeName)
          ->addNodeVectorQuantity("External force", externalForces);
      polyscope::getCurveNetwork(polyscopeName)
          ->getQuantity("External force")
          ->setEnabled(true);
    }
  }
  static std::unordered_map<size_t, std::unordered_set<int>>
  constructFixedVerticesSpanGrid(const size_t &spanGridSize,
                                 const size_t &gridVertices) {
    std::unordered_map<size_t, std::unordered_set<int>> fixedVertices;
    for (size_t vi = 0; vi < spanGridSize - 1; vi++) {
      fixedVertices[vi] = {0, 1, 2};
    }
    for (size_t vi = gridVertices - 1 - (spanGridSize - 2); vi < gridVertices;
         vi++) {
      fixedVertices[vi] = {0, 1, 2};
    }
    for (size_t vi = spanGridSize - 1;
         vi < gridVertices - 1 - (spanGridSize - 2) - spanGridSize + 1;
         vi += spanGridSize + 1) {
      fixedVertices[vi] = {0, 1, 2};
      fixedVertices[vi + spanGridSize] = {0, 1, 2};
    }

    return fixedVertices;
  }
};

struct SimulationResults {
  SimulationHistory history;
  std::vector<Vector6d> displacements;
  double executionTime{0};
  std::string simulationLabel{"NoLabel"};

  void registerForDrawing(const SimulationJob &job) const {
    polyscope::options::groundPlaneEnabled = false;
    polyscope::view::upDir = polyscope::view::UpDir::ZUp;
    const std::string branchName = "Branch:Polyscope";
    polyscope::options::programName = branchName;
    if (!polyscope::state::initialized) {
      polyscope::init();
    } /* else {
       polyscope::removeAllStructures();
     }*/
    const std::string undeformedMeshName = "Undeformed_" + simulationLabel;
    polyscope::registerCurveNetwork(undeformedMeshName,
                                    job.mesh->getEigenVertices(),
                                    job.mesh->getEigenEdges());

    const std::string deformedMeshName = "Deformed_" + simulationLabel;
    polyscope::registerCurveNetwork(deformedMeshName,
                                    job.mesh->getEigenVertices(),
                                    job.mesh->getEigenEdges());
    Eigen::MatrixX3d nodalDisplacements(job.mesh->VN(), 3);
    for (VertexIndex vi = 0; vi < job.mesh->VN(); vi++) {
      const Vector6d &nodalDisplacement = displacements[vi];
      nodalDisplacements.row(vi) = Eigen::Vector3d(
          nodalDisplacement[0], nodalDisplacement[1], nodalDisplacement[2]);
    }
    polyscope::getCurveNetwork(deformedMeshName)
        ->updateNodePositions(job.mesh->getEigenVertices() +
                              nodalDisplacements);

    std::vector<std::array<double, 3>> nodeColors(job.mesh->VN());
    for (auto fixedVertex : job.fixedVertices) {
      nodeColors[fixedVertex.first] = {0, 0, 1};
    }
    if (!job.nodalForcedDisplacements.empty()) {
      for (std::pair<VertexIndex, Eigen::Vector3d> viDisplPair :
           job.nodalForcedDisplacements) {
        const VertexIndex vi = viDisplPair.first;
        nodeColors[vi][0] += 1;
        nodeColors[vi][0] /= 2;
        nodeColors[vi][1] += 0;
        nodeColors[vi][1] /= 2;
        nodeColors[vi][2] += 0;
        nodeColors[vi][2] /= 2;
      }
    }
    std::for_each(nodeColors.begin(), nodeColors.end(),
                  [](std::array<double, 3> &color) {
                    const double norm =
                        sqrt(std::pow(color[0], 2) + std::pow(color[1], 2) +
                             std::pow(color[2], 2));
                    if (norm > std::pow(10, -7)) {
                      color[0] /= norm;
                      color[1] /= norm;
                      color[2] /= norm;
                    }
                  });
    // per node external forces
    std::vector<std::array<double, 3>> externalForces(job.mesh->VN());
    std::vector<std::array<double, 3>> externalMoments(job.mesh->VN());
    for (const auto &forcePair : job.nodalExternalForces) {
      auto index = forcePair.first;
      auto force = forcePair.second;
      externalForces[index] = {force[0], force[1], force[2]};
      externalMoments[index] = {force[3], force[4], 0};
    }

    polyscope::getCurveNetwork(undeformedMeshName)
        ->addNodeColorQuantity("Boundary conditions", nodeColors)
        ->setEnabled(true);

    polyscope::getCurveNetwork(undeformedMeshName)
        ->addNodeVectorQuantity("External force", externalForces)
        ->setEnabled(true);

    polyscope::getCurveNetwork(undeformedMeshName)
        ->addNodeVectorQuantity("External moments", externalMoments)
        ->setEnabled(true);

    polyscope::getCurveNetwork(deformedMeshName)
        ->addNodeColorQuantity("Boundary conditions", nodeColors)
        ->setEnabled(false);

    polyscope::getCurveNetwork(deformedMeshName)
        ->addNodeVectorQuantity("External force", externalForces)
        ->setEnabled(true);

    //    polyscope::show();
  }
};

#endif // SIMULATIONHISTORY_HPP