#ifndef SIMULATIONHISTORYPLOTTER_HPP
#define SIMULATIONHISTORYPLOTTER_HPP

#include "elementalmesh.hpp"
#include "simulationresult.hpp"
#include "utilities.hpp"
#include <algorithm>
#include <matplot/matplot.h>

struct SimulationResultsReporter {
  using VertexType = VCGEdgeMesh::VertexType;
  using CoordType = VCGEdgeMesh::CoordType;
  using Vector6d = Vector6d;

  SimulationResultsReporter() {}

  void writeStatistics(const SimulationResults &results,
                       const std::string &reportFolderPath) {

    ofstream file;
    file.open(
        std::filesystem::path(reportFolderPath).append("results.txt").string());

    const size_t numberOfSteps = results.history.numberOfSteps;
    file << "Number of steps " << numberOfSteps << "\n";

    //    file << "Force threshold used " << 1000 << "\n";

    //    assert(numberOfSteps == results.history.potentialEnergy.size() &&
    //           numberOfSteps == results.history.residualForces.size());
    // Write kinetic energies
    const SimulationHistory &history = results.history;
    if (!history.kineticEnergy.empty()) {
      file << "Kinetic energies"
           << "\n";
      for (size_t step = 0; step < numberOfSteps; step++) {
        file << history.kineticEnergy[step] << "\n";
      }
      file << "\n";
    }

    if (!history.residualForces.empty()) {
      file << "Residual forces"
           << "\n";
      for (size_t step = 0; step < numberOfSteps; step++) {
        file << history.residualForces[step] << "\n";
      }
      file << "\n";
    }

    if (!history.potentialEnergies.empty()) {
      file << "Potential energies"
           << "\n";
      for (size_t step = 0; step < numberOfSteps; step++) {
        file << history.potentialEnergies[step] << "\n";
      }
      file << "\n";
    }
    file.close();
  }

  void reportResults(
      const std::vector<SimulationResults> &results,
      const std::string &reportFolderPath,
      const std::string &graphSuffix = std::string(),
      const SimulationResults &groundOfTruthResults = SimulationResults()) {
    if (results.empty()) {
      return;
    }

    //    std::filesystem::remove_all(debuggerFolder);
    std::filesystem::create_directory(reportFolderPath);
    for (const SimulationResults &simulationResult : results) {
      const auto simulationResultPath =
          std::filesystem::path(reportFolderPath)
              .append(simulationResult.label);
      std::filesystem::create_directory(simulationResultPath.string());

      createPlots(simulationResult.history, simulationResultPath.string(),
                  graphSuffix);
      writeStatistics(simulationResult, simulationResultPath);
    }
  }

  static void createPlot(const std::string &xLabel, const std::string &yLabel,
                         const std::vector<double> &YvaluesToPlot,
                         const std::string &saveTo = {}) {
    matplot::xlabel(xLabel);
    matplot::ylabel(yLabel);
    matplot::figure(true);
    //    matplot::hold(matplot::on);
    matplot::grid(matplot::on);
    auto x = matplot::linspace(0, YvaluesToPlot.size(), YvaluesToPlot.size());
    matplot::scatter(x, YvaluesToPlot)
        //        ->marker_indices(history.redMarks)
        //          ->marker_indices(truncatedRedMarks)
        //        .marker_color("r")
        ->marker_size(1);
    //    auto greenMarksY = matplot::transform(
    //        history.greenMarks, [&](auto x) { return history.kineticEnergy[x];
    //        });
    //    matplot::scatter(history.greenMarks, greenMarksY)
    //        ->color("green")
    //        .marker_size(10);
    //    matplot::hold(matplot::off);
    if (!saveTo.empty()) {
      matplot::save(saveTo);
    }
  }

  void createPlots(const SimulationHistory &history,
                   const std::string &reportFolderPath,
                   const std::string &graphSuffix) {
    const auto graphsFolder =
        std::filesystem::path(reportFolderPath).append("Graphs");
    std::filesystem::remove_all(graphsFolder);
    std::filesystem::create_directory(graphsFolder.string());

    if (!history.kineticEnergy.empty()) {
      createPlot("Number of Iterations", "Log of Kinetic Energy",
                 history.kineticEnergy,
                 std::filesystem::path(graphsFolder)
                     .append("KineticEnergy" + graphSuffix + ".png")
                     .string());
    }

    if (!history.residualForces.empty()) {
      createPlot("Number of Iterations", "Residual Forces norm",
                 history.residualForces,
                 std::filesystem::path(graphsFolder)
                     .append("ResidualForces" + graphSuffix + ".png")
                     .string());
    }

    if (!history.potentialEnergies.empty()) {
      createPlot("Number of Iterations", "Potential energy",
                 history.potentialEnergies,
                 std::filesystem::path(graphsFolder)
                     .append("PotentialEnergy" + graphSuffix + ".png")
                     .string());
    }
  }
};

#endif // SIMULATIONHISTORYPLOTTER_HPP