#include "reducedmodelevaluator.hpp"

#include <execution>
#include <filesystem>

#include "hexagonremesher.hpp"
#include "reducedmodel.hpp"
#include "reducedmodeloptimizer.hpp"
#include "simulationmodelfactory.hpp"
#include "trianglepatterngeometry.hpp"

using PatternVertexIndex = VertexIndex;
using ReducedModelVertexIndex = VertexIndex;

ReducedModelEvaluator::ReducedModelEvaluator() {
  pTileIntoSurface = [&]() {
    std::istringstream inputStream_tileIntoTriSurface(
        tileIntoSurfaceFileContent);
    VCGTriMesh tileInto_triMesh;
    const bool surfaceLoadSuccessfull =
        tileInto_triMesh.load(inputStream_tileIntoTriSurface);
    tileInto_triMesh.setLabel("instantMeshes_plane_34");
    assert(surfaceLoadSuccessfull);
    return PolygonalRemeshing::remeshWithPolygons(tileInto_triMesh);
  }();
}

ReducedModelEvaluator::Results ReducedModelEvaluator::evaluateReducedModel(
    ReducedModelOptimization::Results& optimizationResult) {
  const std::filesystem::path scenariosDirectoryPath =
      "/home/iason/Coding/Projects/Approximating shapes with flat "
      "patterns/ReducedModelEvaluator/Scenarios";
  const std::filesystem::path fullPatternTessellatedResultsDirectoryPath =
      "/home/iason/Coding/Projects/Approximating shapes with flat "
      "patterns/ReducedModelEvaluator/TessellatedResults";

  return evaluateReducedModel(optimizationResult, scenariosDirectoryPath,
                              fullPatternTessellatedResultsDirectoryPath);
}

void ReducedModelEvaluator::printResults(const Results& evaluationResults,
                                         const std::string& resultsLabel) {
  csvFile csvOutputToCout({}, true);
  Settings exportSettings;
  exportSettings.exportingDirection = Vertical;
  exportSettings.shouldWriteHeader = false;
  exportSettings.resultsLabel = resultsLabel;
  printResults(evaluationResults, exportSettings, csvOutputToCout);
}

void ReducedModelEvaluator::printResults(const Results& evaluationResults,
                                         const Settings& settings,
                                         csvFile& csvOutput) {
  if (settings.shouldWriteHeader) {
    csvOutput << csvExportingDataStrings[settings.exportingData];
    printHeader(settings, csvOutput);
    csvOutput << endrow;
  }
  if (!settings.resultsLabel.empty()) {
    csvOutput << settings.resultsLabel;
  }
  if (settings.exportingDirection == Vertical) {
    printResultsVertically(evaluationResults, csvOutput);
  } else {
    printResultsHorizontally(evaluationResults, csvOutput);
  }
}

void ReducedModelEvaluator::printHeader(const Settings& settings,
                                        csvFile& csvOutput) {
  if (settings.exportingDirection == Horizontal) {
    //        csvOutput << "Job label";
    for (int jobIndex = 0;
         jobIndex < ReducedModelEvaluator::scenariosTestSetLabels.size();
         jobIndex++) {
      const std::string& jobLabel =
          ReducedModelEvaluator::scenariosTestSetLabels[jobIndex];
      csvOutput << jobLabel;
    }
  } else {
    std::cout << "Vertical header not defined" << std::endl;
    assert(false);
    std::terminate();
  }
}

void ReducedModelEvaluator::printResultsHorizontally(
    const Results& evaluationResults,
    csvFile& csvOutput) {
  // print header
  // print raw error
  constexpr bool shouldPrintRawError = false;
  if (shouldPrintRawError) {
    double sumOfFull2Reduced = 0;
    int numOfNonEvaluatedScenarios = 0;
    for (int jobIndex = 0;
         jobIndex < ReducedModelEvaluator::scenariosTestSetLabels.size();
         jobIndex++) {
      const double& distance_fullDrmToReduced =
          evaluationResults.distances_drm2reduced[jobIndex];
      if (distance_fullDrmToReduced == -1) {
        csvOutput << "notEvaluated";
        numOfNonEvaluatedScenarios++;
      } else {
        csvOutput << distance_fullDrmToReduced;
        sumOfFull2Reduced += distance_fullDrmToReduced;
      }
    }
  }
  // print normalized error
  double sumOfNormalizedFull2Reduced = 0;
  for (int jobIndex = 0;
       jobIndex < ReducedModelEvaluator::scenariosTestSetLabels.size();
       jobIndex++) {
    const double& distance_normalizedFullDrmToReduced =
        evaluationResults.distances_normalizedDrm2reduced[jobIndex];
    if (distance_normalizedFullDrmToReduced == -1) {
      csvOutput << "notEvaluated";
    } else {
      csvOutput << distance_normalizedFullDrmToReduced;
      sumOfNormalizedFull2Reduced += distance_normalizedFullDrmToReduced;
    }
  }
}

void ReducedModelEvaluator::printResultsVertically(
    const Results& evaluationResults,
    csvFile& csvOutput) {
#ifdef POLYSCOPE_DEFINED
  csvOutput << "pattern2Reduced"
            << "norm_pattern2Reduced";
#else
  csvOutput << "Job Label"
            << "drm2Reduced"
            << "norm_drm2Reduced";

#endif
  csvOutput << endrow;
  double sumOfFull2Reduced = 0;
  double sumOfNormalizedFull2Reduced = 0;
  int numOfNonEvaluatedScenarios = 0;
  for (int jobIndex = 0;
       jobIndex < ReducedModelEvaluator::scenariosTestSetLabels.size();
       jobIndex++) {
    const double& distance_fullDrmToReduced =
        evaluationResults.distances_drm2reduced[jobIndex];
    const double& distance_normalizedFullDrmToReduced =
        evaluationResults.distances_normalizedDrm2reduced[jobIndex];
#ifndef POLYSCOPE_DEFINED
    const std::string& jobLabel =
        ReducedModelEvaluator::scenariosTestSetLabels[jobIndex];
    csvOutput << jobLabel;
#endif
    if (distance_fullDrmToReduced == -1) {
      csvOutput << "notEvaluated"
                << "notEvaluated";
      numOfNonEvaluatedScenarios++;
    } else {
      csvOutput << distance_fullDrmToReduced
                << distance_normalizedFullDrmToReduced;
      sumOfFull2Reduced += distance_fullDrmToReduced;
      sumOfNormalizedFull2Reduced += distance_normalizedFullDrmToReduced;
    }
    csvOutput << endrow;
  }
  const int numOfEvaluatedScenarios =
      ReducedModelEvaluator::scenariosTestSetLabels.size() -
      numOfNonEvaluatedScenarios;
  const double averageDistance_full2Reduced =
      sumOfFull2Reduced / numOfEvaluatedScenarios;
  const double averageDistance_normalizedFull2Reduced =
      sumOfNormalizedFull2Reduced / numOfEvaluatedScenarios;
#ifndef POLYSCOPE_DEFINED
  csvOutput << "Average error";
#endif
  csvOutput << averageDistance_full2Reduced
            << averageDistance_normalizedFull2Reduced;
  csvOutput << endrow;
#ifndef POLYSCOPE_DEFINED
  csvOutput << "Cumulative error";
#endif
  csvOutput << sumOfFull2Reduced << sumOfNormalizedFull2Reduced;
  csvOutput << endrow;
  csvOutput << endrow;
}

ReducedModelEvaluator::Results ReducedModelEvaluator::evaluateReducedModel(
    ReducedModelOptimization::Results& optimizationResult,
    const std::filesystem::path& scenariosDirectoryPath,
    const std::filesystem::path& patternTessellatedResultsDirectoryPath) {
  // Apply optimization results to the reduced model
  ReducedModel reducedModel;
  reducedModel.deleteDanglingVertices();
  std::unordered_map<std::string, double> optimalXVariables_set(
      optimizationResult.optimalXNameValuePairs.begin(),
      optimizationResult.optimalXNameValuePairs.end());
  reducedModel.updateBaseTriangleGeometry_R(optimalXVariables_set.at("R"));
  reducedModel.updateBaseTriangleGeometry_theta(
      optimalXVariables_set.at("Theta"));

  //  reducedModel.registerForDrawing();

  // Scale tile-into surface
  pTileIntoSurface->moveToCenter();
  const double scaleFactor =
      optimizationResult.settings.targetBaseTriangleSize /
      pTileIntoSurface->getAverageFaceRadius();
  vcg::tri::UpdatePosition<VCGPolyMesh>::Scale(*pTileIntoSurface, scaleFactor);
#ifdef POLYSCOPE_DEFINED
  pTileIntoSurface->registerForDrawing(color_tileIntoSurface);
#endif

  // Tile full pattern into surface
  std::vector<PatternGeometry> patterns(1);
  patterns[0].copy(optimizationResult.baseTrianglePattern);
  patterns[0].interfaceNodeIndex = 3;
  patterns[0].deleteDanglingVertices();
  std::vector<int> perSurfaceFacePatternIndices(pTileIntoSurface->FN(), 0);
  std::vector<std::vector<size_t>> perPatternIndexTiledFullPatternEdgeIndices;
  std::vector<size_t> tileIntoEdgeToTiledFullVi;
  std::shared_ptr<PatternGeometry> pTilled_pattern =
      PatternGeometry::tilePattern(patterns, {}, *pTileIntoSurface,
                                   perSurfaceFacePatternIndices,
                                   tileIntoEdgeToTiledFullVi,
                                   perPatternIndexTiledFullPatternEdgeIndices);
  pTilled_pattern->setLabel("Tilled_pattern");

  // Tile reduced pattern into surface
  std::vector<PatternGeometry> reducedModels(1);
  reducedModels[0].copy(reducedModel);
  std::vector<std::vector<size_t>>
      perPatternIndexTiledReducedPatternEdgeIndices;
  std::vector<size_t> tileIntoEdgeToTiledReducedVi;
  std::shared_ptr<PatternGeometry> pTilled_reducedModel =
      PatternGeometry::tilePattern(
          reducedModels, {0}, *pTileIntoSurface, perSurfaceFacePatternIndices,
          tileIntoEdgeToTiledReducedVi,
          perPatternIndexTiledReducedPatternEdgeIndices);
  pTilled_reducedModel->setLabel("Tilled_reduced_model");

  std::unordered_map<PatternVertexIndex, ReducedModelVertexIndex>
      fullToReducedViMap;  // of only the common vertices
  std::unordered_map<ReducedModelVertexIndex, PatternVertexIndex>
      reducedToFullViMap;  // of only the common vertices
  for (int ei = 0; ei < pTileIntoSurface->EN(); ei++) {
    fullToReducedViMap[tileIntoEdgeToTiledFullVi[ei]] =
        tileIntoEdgeToTiledReducedVi[ei];
  }
  constructInverseMap(fullToReducedViMap, reducedToFullViMap);

  std::vector<size_t> tilledFullPatternInterfaceVi;
  tilledFullPatternInterfaceVi.clear();
  tilledFullPatternInterfaceVi.resize(fullToReducedViMap.size());
  size_t viIndex = 0;
  for (std::pair<size_t, size_t> fullToReducedPair : fullToReducedViMap) {
    tilledFullPatternInterfaceVi[viIndex++] = fullToReducedPair.first;
  }

  // Create simulation meshes
  ////Tessellated full pattern simulation mesh
  std::shared_ptr<SimulationEdgeMesh> pSimulationEdgeMesh_tilledPattern =
      std::make_shared<SimulationEdgeMesh>(*pTilled_pattern);
  pSimulationEdgeMesh_tilledPattern->setBeamCrossSection(
      optimizationResult.settings.beamDimensions_pattern);
  if (optimizationResult.settings.youngsModulus_pattern == 0) {
    std::cerr << "Full pattern's young modulus not found." << std::endl;
    std::terminate();
  }
  pSimulationEdgeMesh_tilledPattern->setBeamMaterial(
      0.3, optimizationResult.settings.youngsModulus_pattern);
  pSimulationEdgeMesh_tilledPattern->reset();
  //    optimizationResult.draw();
#ifdef POLYSCOPE_DEFINED
  pSimulationEdgeMesh_tilledPattern->registerForDrawing(
      color_tesselatedPatterns);
#endif

  ////Tessellated reduced pattern simulation mesh
  std::shared_ptr<SimulationEdgeMesh> pSimulationEdgeMesh_tilledReducedModel;
  pSimulationEdgeMesh_tilledReducedModel =
      std::make_shared<SimulationEdgeMesh>(*pTilled_reducedModel);
  ReducedModelOptimization::Results::applyOptimizationResults_elements(
      optimizationResult, pSimulationEdgeMesh_tilledReducedModel);
  pSimulationEdgeMesh_tilledReducedModel->reset();
#ifdef POLYSCOPE_DEFINED
  pSimulationEdgeMesh_tilledReducedModel->registerForDrawing(
      color_tesselatedReducedModels);
  polyscope::show();
#endif

  Results evaluationResults;
  evaluationResults.distances_drm2reduced.fill(-1);
  evaluationResults.distances_normalizedDrm2reduced.fill(-1);
  DRMSimulationModel::Settings drmSimulationSettings;
  //  drmSimulationSettings.threshold_residualToExternalForcesNorm = 1e-3;
  //    drmSimulationSettings.load(drmSettingsFilePath);
  drmSimulationSettings.beVerbose = true;
  drmSimulationSettings.maxDRMIterations = 5e6;
  drmSimulationSettings.debugModeStep = 100000;
  drmSimulationSettings.threshold_totalTranslationalKineticEnergy = 1e-14;
  //  drmSimulationSettings.threshold_currentToFirstPeakTranslationalKineticEnergy
  //  =
  //      1e-10;
  drmSimulationSettings.threshold_averageResidualToExternalForcesNorm = 1e-5;
//  drmSimulationSettings.linearGuessForceScaleFactor = 0.8;
//  drmSimulationSettings.viscousDampingFactor = 7e-3;
//  drmSimulationSettings.xi = 0.9999;
//  drmSimulationSettings.gamma = 0.25;
#ifdef POLYSCOPE_DEFINED
  //    drmSimulationSettings.shouldDraw = true;
  drmSimulationSettings.shouldCreatePlots = false;
  constexpr bool shouldDrawScenarioResults = true;
  if (shouldDrawScenarioResults) {
    pSimulationEdgeMesh_tilledPattern->registerForDrawing(
        ReducedModelOptimization::Colors::patternInitial);
  }
#endif
  const std::string& simulationModelLabel_pattern =
      optimizationResult.settings.simulationModelLabel_groundTruth;
  const std::string& simulationModelLabel_reducedModel =
      optimizationResult.settings.simulationModelLabel_reducedModel;
  const bool shouldRerunFullPatternSimulation = [&]() {
    //    if (simulationModelLabel_pattern == DRMSimulationModel::label) {
    return false;
    //    }
    //    return true;
  }();
  std::for_each(
#ifndef POLYSCOPE_DEFINED
      std::execution::par_unseq,
#endif
      scenariosTestSetLabels.begin(), scenariosTestSetLabels.end(),
      [&](const std::string& jobLabel) {
        // check if reduced model scenario exists
        //        const std::string &jobLabel =
        //        scenariosTestSetLabels[jobIndex];
        const std::filesystem::path tiledReducedPatternJobFilePath =
            std::filesystem::path(scenariosDirectoryPath)
                .append(pTileIntoSurface->getLabel())
                .append(jobLabel)
                .append("ReducedJob")
                .append(SimulationJob::jsonDefaultFileName);
        if (!std::filesystem::exists(tiledReducedPatternJobFilePath)) {
          std::cerr << "Scenario " << jobLabel
                    << " not found in:" << tiledReducedPatternJobFilePath
                    << std::endl;
          //            continue; //if not move on to the next scenario
          return;
        }
        // Map the reduced job to the job on the pattern tessellation
        // set jobs
        std::shared_ptr<SimulationJob> pJob_tiledReducedModel;
        pJob_tiledReducedModel =
            std::make_shared<SimulationJob>(SimulationJob());
        pJob_tiledReducedModel->load(tiledReducedPatternJobFilePath, false);
        pJob_tiledReducedModel->pMesh = pSimulationEdgeMesh_tilledReducedModel;
        std::shared_ptr<SimulationJob> pJob_tilledPattern;
        pJob_tilledPattern = std::make_shared<SimulationJob>(SimulationJob());
        pJob_tilledPattern->pMesh = pSimulationEdgeMesh_tilledPattern;
        pJob_tiledReducedModel->remap(reducedToFullViMap, *pJob_tilledPattern);
        //            pJob_tiledReducedPattern->registerForDrawing(pTiledReducedPattern->getLabel());
        //            pJob_tiledFullPattern->registerForDrawing(pTiledFullPattern->getLabel());
        //            polyscope::show();
        const std::filesystem::path surfaceFolderPath =
            std::filesystem::path(patternTessellatedResultsDirectoryPath)
                .append(simulationModelLabel_pattern + "_" +
                        pTileIntoSurface->getLabel());
        const std::string scenarioLabel = pJob_tilledPattern->getLabel();
        const std::filesystem::path scenarioDirectoryPath =
            std::filesystem::path(surfaceFolderPath).append(scenarioLabel);
        // Save reduced job
        constexpr bool exportReducedJob = false;
        if (exportReducedJob) {
          const std::filesystem::path reducedJobDirectoryPath =
              std::filesystem::path(scenarioDirectoryPath).append("ReducedJob");
          std::filesystem::create_directories(reducedJobDirectoryPath);
          pJob_tiledReducedModel->save(reducedJobDirectoryPath);
        }
        // Check if the drm simulation of the full pattern has already been
        // computed
        ////Full
        const std::string& patternLabel = [&]() {
          const std::string patternLabel =
              optimizationResult.baseTrianglePattern.getLabel();
          if (patternLabel.find("_") == std::string::npos) {
            return std::to_string(optimizationResult.baseTrianglePattern.EN()) +
                   "_" + patternLabel;
          } else {
            return patternLabel;
          }
        }();
        const auto tilledPatternResultsFolderPath =
            std::filesystem::path(scenarioDirectoryPath)
                .append(patternLabel)
                .append("Results");
        if (shouldRerunFullPatternSimulation &&
            std::filesystem::exists(tilledPatternResultsFolderPath)) {
          std::filesystem::remove_all(tilledPatternResultsFolderPath);
        }

        const std::filesystem::path fullPatternJobFolderPath =
            std::filesystem::path(scenarioDirectoryPath)
                .append(patternLabel)
                .append("SimulationJob");
        SimulationResults simulationResults_tilledPattern;
        if (std::filesystem::exists(tilledPatternResultsFolderPath)) {
          // Load full pattern results
          assert(std::filesystem::exists(fullPatternJobFolderPath));
          simulationResults_tilledPattern.load(tilledPatternResultsFolderPath,
                                               fullPatternJobFolderPath);
#ifdef POLYSCOPE_DEFINED
          simulationResults_tilledPattern.registerForDrawing(
              color_tesselatedPatterns);
          //          std::ifstream ifs("CameraSettings.json");
          //          nlohmann::json json;
          //          ifs >> json;
          //          polyscope::view::setCameraFromJson(json.dump(), false);
          polyscope::show();
          const std::string cameraJson = polyscope::view::getCameraJson();
          std::filesystem::path jsonFilePath("CameraSettings.json");
          std::ofstream jsonFile_cameraSettings(jsonFilePath.string());
          jsonFile_cameraSettings << cameraJson;
          jsonFile_cameraSettings.close();
          std::filesystem::create_directories("screenshots");
          const std::string screenshotOutputFilePath =
              (std::filesystem::current_path()
                   .append("screenshots")
                   .append(optimizationResult.label + "_" +
                           pJob_tilledPattern->getLabel()))
                  .string() +
              ".png";
          std::cout << "Saving image to:" << screenshotOutputFilePath
                    << std::endl;
          polyscope::screenshot(screenshotOutputFilePath, false);
          simulationResults_tilledPattern.unregister();
#endif

          simulationResults_tilledPattern.converged = true;
        } else {
          std::cout << "Tilled pattern simulation results not found in:"
                    << tilledPatternResultsFolderPath << std::endl;
          // Full
          std::cout << "Executing:" << jobLabel << std::endl;
          SimulationModelFactory factory;
          std::unique_ptr<SimulationModel> pTilledPatternSimulationModel =
              factory.create(simulationModelLabel_pattern);
          // TODO: since the drm simulation model does not have a common
          // interface with the rest of simulation models I need to cast it in
          // order to pass simulation settings. Fix it by removing the
          // SimulationSettings argument
          if (simulationModelLabel_pattern == DRMSimulationModel::label) {
            simulationResults_tilledPattern =
                static_cast<DRMSimulationModel*>(
                    pTilledPatternSimulationModel.get())
                    ->executeSimulation(pJob_tilledPattern,
                                        drmSimulationSettings);
          } else {
            simulationResults_tilledPattern =
                pTilledPatternSimulationModel->executeSimulation(
                    pJob_tilledPattern);
          }
        }
        if (!simulationResults_tilledPattern.converged) {
          std::cerr << "Full pattern simulation failed." << std::endl;
          std::cerr << "Not saving results" << std::endl;
          //            continue;
          return;
        }
        std::filesystem::create_directories(tilledPatternResultsFolderPath);
        const std::filesystem::path drmResultsOutputPath =
            std::filesystem::path(scenarioDirectoryPath).append(patternLabel);
        simulationResults_tilledPattern.save(drmResultsOutputPath);

        //        LinearSimulationModel linearSimulationModel;
        SimulationModelFactory factory;
        std::unique_ptr<SimulationModel> pSimulationModel_tilledReducedModel =
            factory.create(simulationModelLabel_reducedModel);
        SimulationResults simulationResults_tiledReducedModel =
            pSimulationModel_tilledReducedModel->executeSimulation(
                pJob_tiledReducedModel);

    //        ChronosEulerNonLinearSimulationModel
    //            debug_chronosNonLinearSimulationModel;
    //        const auto debug_chronosResults =
    //            debug_chronosNonLinearSimulationModel.executeSimulation(
    //                pJob_tilledPattern);
    //        LinearSimulationModel debug_linearSimulationModel;
    //        const auto debug_linearSimResults =
    //            debug_linearSimulationModel.executeSimulation(pJob_tilledPattern);

#ifdef POLYSCOPE_DEFINED
        if (shouldDrawScenarioResults) {
          simulationResults_tiledReducedModel.registerForDrawing(
              ReducedModelOptimization::Colors::reducedDeformed, true,
              simulationResults_tilledPattern.pJob->pMesh
                  ->getBeamDimensions()[0]
                  .getDrawingRadius());
          simulationResults_tilledPattern.registerForDrawing(
              ReducedModelOptimization::Colors::patternDeformed);
          //                debug_chronosResults.registerForDrawing();
          //                debug_linearSimResults.registerForDrawing();
          polyscope::show();
          //                debug_linearSimResults.unregister();
          simulationResults_tiledReducedModel.unregister();
          simulationResults_tilledPattern.unregister();
          //                debug_chronosResults.unregister();
        }
#endif

        // compute the full2reduced distance
        const double distance_patternToReduced =
            simulationResults_tilledPattern.computeDistance(
                simulationResults_tiledReducedModel, fullToReducedViMap);
        double distance_patternSumOfAllVerts = 0;
        for (std::pair<size_t, size_t> fullToReducedPair : fullToReducedViMap) {
          distance_patternSumOfAllVerts +=
              simulationResults_tilledPattern
                  .displacements[fullToReducedPair.first]
                  .getTranslation()
                  .norm();
        }
        const double distance_normalizedPatternToReduced =
            distance_patternToReduced / distance_patternSumOfAllVerts;
        const int jobIndex = &jobLabel - &scenariosTestSetLabels[0];
        evaluationResults.distances_drm2reduced[jobIndex] =
            distance_patternToReduced;
        evaluationResults.distances_normalizedDrm2reduced[jobIndex] =
            distance_normalizedPatternToReduced;
      });

  return evaluationResults;
}