#ifndef SIMULATIONHISTORY_HPP
#define SIMULATIONHISTORY_HPP

#include "elementalmesh.hpp"
#include "nlohmann/json.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();
  }
};

namespace nlohmann {

template <> struct adl_serializer<std::unordered_map<VertexIndex, Vector6d>> {
  static void to_json(json &j,
                      const std::unordered_map<VertexIndex, Vector6d> &value) {
    // calls the "to_json" method in T's namespace
  }

  static void from_json(const nlohmann::json &j,
                        std::unordered_map<VertexIndex, Vector6d> &m) {
    std::cout << "Entered." << std::endl;
    for (const auto &p : j) {
      m.emplace(p.at(0).template get<VertexIndex>(),
                p.at(1).template get<std::array<double, 6>>());
    }
  }
};
} // namespace nlohmann

class SimulationJob {
  //  const std::unordered_map<VertexIndex, VectorType> nodalForcedNormals;
  // json labels
  struct JSONLabels {
    inline static std::string meshFilename{"mesh filename"};
    inline static std::string constrainedVertices{"fixed vertices"};
    inline static std::string nodalForces{"forces"};
    inline static std::string label{"label"};
    inline static std::string meshLabel{
        "label"}; // TODO: should be in the savePly function of the simulation
                  // mesh class
  } jsonLabels;

public:
  std::shared_ptr<SimulationMesh> pMesh;
  std::string label{"empty_job"};
  std::unordered_map<VertexIndex, std::unordered_set<int>> constrainedVertices;
  std::unordered_map<VertexIndex, Vector6d> nodalExternalForces;
  std::unordered_map<VertexIndex, Eigen::Vector3d> nodalForcedDisplacements;
  SimulationJob(
      const std::shared_ptr<SimulationMesh> &m, const std::string &label,
      const std::unordered_map<VertexIndex, std::unordered_set<int>> &cv,
      const std::unordered_map<VertexIndex, Vector6d> &ef = {},
      const std::unordered_map<VertexIndex, Eigen::Vector3d> &fd = {})
      : pMesh(m), label(label), constrainedVertices(cv),
        nodalExternalForces(ef), nodalForcedDisplacements(fd) {}
  SimulationJob() {}
  SimulationJob(const std::string &jsonFilename) { load(jsonFilename); }
  std::string getLabel() const { return label; }

  std::string toString() const {
    nlohmann::json json;
    if (!constrainedVertices.empty()) {
      json[jsonLabels.constrainedVertices] = constrainedVertices;
    }
    if (!nodalExternalForces.empty()) {
      std::unordered_map<VertexIndex, std::array<double, 6>> arrForces;
      for (const auto &f : nodalExternalForces) {
        arrForces[f.first] = f.second;
      }
      json[jsonLabels.nodalForces] = arrForces;
    }

    return json.dump();
  }

  void registerForDrawing(const std::string &meshLabel) const {
    initPolyscope();
    if (meshLabel.empty()) {
      assert(false);
      std::cerr << "Expects a mesh label on which to draw the simulation job."
                << std::endl;
      return;
    }

    auto structs = polyscope::state::structures;
    if (!polyscope::hasCurveNetwork(meshLabel)) {
      assert(false);
      std::cerr << "Expects mesh already being registered to draw the "
                   "simulation job. No struct named " +
                       meshLabel
                << std::endl;
      return;
    }
    std::vector<std::array<double, 3>> nodeColors(pMesh->VN());
    for (auto fixedVertex : constrainedVertices) {
      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;
                    }
                  });

    auto cn = polyscope::getCurveNetwork(meshLabel);
    if (!nodeColors.empty()) {
      polyscope::getCurveNetwork(meshLabel)
          ->addNodeColorQuantity("Boundary conditions_" + label, nodeColors)
          ->setEnabled(false);
    }

    // per node external forces
    std::vector<std::array<double, 3>> externalForces(pMesh->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(meshLabel)
          ->addNodeVectorQuantity("External force_" + label, externalForces)
          ->setEnabled(false);
    }
  }

  bool load(const std::string &jsonFilename) {
    if (std::filesystem::path(jsonFilename).extension() != ".json") {
      std::cerr << "A json file is expected as input. The given file has the "
                   "following extension:"
                << std::filesystem::path(jsonFilename).extension() << std::endl;
      assert(false);
      return false;
    }

    if (!std::filesystem::exists(std::filesystem::path(jsonFilename))) {
      std::cerr << "The json file does not exist. Json file provided:"
                << jsonFilename << std::endl;
      assert(false);
      return false;
    }

    std::cout << "Loading json file:" << jsonFilename << std::endl;
    nlohmann::json json;
    std::ifstream ifs(jsonFilename);
    json << ifs;

    pMesh = std::make_shared<SimulationMesh>();
    if (json.contains(jsonLabels.meshFilename)) {
      pMesh->loadPly(json[jsonLabels.meshFilename]);
    }

    if (json.contains(jsonLabels.constrainedVertices)) {
      constrainedVertices =
          //      auto conV =
<<<<<<< HEAD
          std::unordered_map<VertexIndex, std::unordered_set<int>>(
              json[jsonLabels.constrainedVertices]);
=======
              json[jsonLabel_constrainedVertices].get<std::unordered_map<VertexIndex, std::unordered_set<int>>>();
>>>>>>> 2599d472614f403f8249060272ce288033a919d4
      std::cout << "Loaded constrained vertices. Number of constrained "
                   "vertices found:"
                << constrainedVertices.size() << std::endl;
    }

<<<<<<< HEAD
    if (json.contains(jsonLabels.nodalForces)) {
      auto f = std::unordered_map<VertexIndex, std::array<double, 6>>(
          json[jsonLabels.nodalForces]);
=======
    if (json.contains(jsonLabel_nodalForces)) {
      auto f (
          json[jsonLabel_nodalForces].get<std::unordered_map<VertexIndex, std::array<double, 6>>>());
>>>>>>> 2599d472614f403f8249060272ce288033a919d4
      for (const auto &forces : f) {
        nodalExternalForces[forces.first] = Vector6d(forces.second);
      }
      std::cout << "Loaded forces. Number of forces found:"
                << nodalExternalForces.size() << std::endl;
    }

    if (json.contains(jsonLabels.label)) {
      label = json[jsonLabels.label];
    }

    if (json.contains(jsonLabels.meshLabel)) {
      pMesh->setLabel(json[jsonLabels.meshLabel]);
    }

    return true;
  }

  bool save(const std::string &folderDirectory) const {
    const std::filesystem::path pathFolderDirectory(folderDirectory);
    if (!std::filesystem::is_directory(pathFolderDirectory)) {
      std::cerr << "A folder directory is expected for saving the simulation "
                   "job. Exiting.."
                << std::endl;
      return false;
    }

    bool returnValue = true;
    const std::string meshFilename =
        std::filesystem::absolute(
            std::filesystem::canonical(
                std::filesystem::path(pathFolderDirectory)))
            .append(pMesh->getLabel() + ".ply").string();
    returnValue = pMesh->savePly(meshFilename);
    nlohmann::json json;
    json[jsonLabels.meshFilename] = meshFilename;
    if (!constrainedVertices.empty()) {
      json[jsonLabels.constrainedVertices] = constrainedVertices;
    }
    if (!nodalExternalForces.empty()) {
      std::unordered_map<VertexIndex, std::array<double, 6>> arrForces;
      for (const auto &f : nodalExternalForces) {
        arrForces[f.first] = f.second;
      }
      json[jsonLabels.nodalForces] = arrForces;
    }

    if (!label.empty()) {
      json[jsonLabels.label] = label;
    }
    if (!pMesh->getLabel().empty()) {
      json[jsonLabels.meshLabel] = pMesh->getLabel();
    }

    std::string jsonFilename(
        std::filesystem::path(pathFolderDirectory)
<<<<<<< HEAD
            .append(label + "_" + pMesh->getLabel() + "_simulationJob.json"));
=======
            .append(pMesh->getLabel() + "_simScenario.json").string());
>>>>>>> 2599d472614f403f8249060272ce288033a919d4
    std::ofstream jsonFile(jsonFilename);
    jsonFile << json;
    std::cout << "Saved simulation job as:" << jsonFilename << std::endl;

    return returnValue;
  }
};

namespace Eigen {
template <class Matrix>
void write_binary(const std::string &filename, const Matrix &matrix) {
  std::ofstream out(filename,
                    std::ios::out | std::ios::binary | std::ios::trunc);
  typename Matrix::Index rows = matrix.rows(), cols = matrix.cols();
  out.write((char *)(&rows), sizeof(typename Matrix::Index));
  out.write((char *)(&cols), sizeof(typename Matrix::Index));
  out.write((char *)matrix.data(),
            rows * cols * sizeof(typename Matrix::Scalar));
  out.close();
}
template <class Matrix>
void read_binary(const std::string &filename, Matrix &matrix) {
  std::ifstream in(filename, std::ios::in | std::ios::binary);
  typename Matrix::Index rows = 0, cols = 0;
  in.read((char *)(&rows), sizeof(typename Matrix::Index));
  in.read((char *)(&cols), sizeof(typename Matrix::Index));
  matrix.resize(rows, cols);
  in.read((char *)matrix.data(), rows * cols * sizeof(typename Matrix::Scalar));
  in.close();
}
} // namespace Eigen

struct SimulationResults {
  bool converged{true};
  std::shared_ptr<SimulationJob> job;
  SimulationHistory history;
  std::vector<Vector6d> displacements;
  double executionTime{0};
  std::string labelPrefix{"deformed"};
  inline static char deliminator{' '};

  void setLabelPrefix(const std::string &lp) {
    labelPrefix += deliminator + lp;
  }
  std::string getLabel() const {
    return labelPrefix + deliminator + job->pMesh->getLabel() + deliminator +
           job->getLabel();
  }
  void unregister() const {
    if (!polyscope::hasCurveNetwork(getLabel())) {
      std::cerr << "No curve network registered with a name: " << getLabel()
                << std::endl;
      std::cerr << "Nothing to remove." << std::endl;
      return;
    }
    polyscope::removeCurveNetwork(getLabel());
  }
  void registerForDrawing() 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_" + label;
    //    polyscope::registerCurveNetwork(
    //        undeformedMeshName, mesh->getEigenVertices(),
    //        mesh->getEigenEdges());

    const std::shared_ptr<SimulationMesh> &mesh = job->pMesh;
    polyscope::registerCurveNetwork(getLabel(), mesh->getEigenVertices(),
                                    mesh->getEigenEdges())
        ->setRadius(0.0007, false);
    Eigen::MatrixX3d nodalDisplacements(mesh->VN(), 3);
    for (VertexIndex vi = 0; vi < mesh->VN(); vi++) {
      const Vector6d &nodalDisplacement = displacements[vi];
      nodalDisplacements.row(vi) = Eigen::Vector3d(
          nodalDisplacement[0], nodalDisplacement[1], nodalDisplacement[2]);
    }
    polyscope::getCurveNetwork(getLabel())
        ->updateNodePositions(mesh->getEigenVertices() + nodalDisplacements);

    job->registerForDrawing(getLabel());
  }

  void saveDeformedModel() {
    VCGEdgeMesh m;
    vcg::tri::Append<VCGEdgeMesh, SimulationMesh>::MeshCopy(m, *job->pMesh);

    for (int vi = 0; vi < m.VN(); vi++) {
      m.vert[vi].P() =
          m.vert[vi].P() + CoordType(displacements[vi][0], displacements[vi][1],
                                     displacements[vi][2]);
    }
    m.savePly(getLabel() + ".ply");
  }

  void save() {
    const std::string filename(getLabel() + "_displacements.eigenBin");

    Eigen::MatrixXd m = Utilities::toEigenMatrix(displacements);
    Eigen::write_binary(filename, m);
  }

  // The comparison of the results happens comparing the 6-dof nodal
  // displacements
  bool isEqual(const Eigen::MatrixXd &nodalDisplacements) {
    assert(nodalDisplacements.cols() == 6);
    Eigen::MatrixXd eigenDisplacements =
        Utilities::toEigenMatrix(this->displacements);
    const double errorNorm = (eigenDisplacements - nodalDisplacements).norm();
    return errorNorm < 1e-10;
    //    return eigenDisplacements.isApprox(nodalDisplacements);
  }
};

#endif // SIMULATIONHISTORY_HPP