#include "reducedmodeloptimizer.hpp"
#include "bobyqa.h"
#include "flatpattern.hpp"
#include "gradientDescent.h"
#include "simulationhistoryplotter.hpp"
#include "trianglepattterntopology.hpp"
#include <chrono>
#include <dlib/global_optimization.h>
#include <dlib/optimization.h>

const bool gShouldDraw = true;

FormFinder simulator;
std::vector<Eigen::MatrixX3d> g_optimalReducedModelDisplacements;
std::vector<SimulationJob> g_fullPatternSimulationJob;
std::vector<std::shared_ptr<SimulationJob>> g_reducedPatternSimulationJob;
std::unordered_map<ReducedPatternVertexIndex, FullPatternVertexIndex>
    reducedToFullInterfaceViMap;
std::unordered_map<ReducedPatternVertexIndex, FullPatternVertexIndex>
    g_reducedToFullViMap;
matplot::line_handle gPlotHandle;
std::vector<double> gObjectiveValueHistory;
Eigen::Vector4d g_initialX;
std::unordered_set<size_t> g_reducedPatternExludedEdges;
// double g_initialParameters;
Eigen::VectorXd g_initialParameters;
std::vector<ReducedModelOptimizer::SimulationScenario>
    g_simulationScenarioIndices;
std::vector<VectorType> g_innerHexagonVectors{6, VectorType(0, 0, 0)};
double g_innerHexagonInitialPos = 0;
bool g_optimizeInnerHexagonSize{false};
std::vector<SimulationResults> firstOptimizationRoundResults;
int g_firstRoundIterationIndex{0};
double minY{std::numeric_limits<double>::max()};
std::vector<double> minX;

// struct OptimizationCallback {
//  double operator()(const size_t &iterations, const Eigen::VectorXd &x,
//                    const double &fval, Eigen::VectorXd &gradient) const {
//    // run simulation
//    //    SimulationResults reducedModelResults =
//    //        simulator.executeSimulation(reducedModelSimulationJob);
//    //    reducedModelResults.draw(reducedModelSimulationJob);
//    gObjectiveValueHistory.push_back(fval);
//    auto xPlot = matplot::linspace(0, gObjectiveValueHistory.size(),
//                                   gObjectiveValueHistory.size());
//    gPlotHandle = matplot::scatter(xPlot, gObjectiveValueHistory);
//    //    const std::string plotImageFilename = "objectivePlot.png";
//    //    matplot::save(plotImageFilename);
//    //    if (numberOfOptimizationRounds % 30 == 0) {
//    //      std::filesystem::copy_file(
//    //          std::filesystem::path(plotImageFilename),
//    //          std::filesystem::path("objectivePlot_copy.png"));
//    //    }
//    //    std::stringstream ss;
//    //    ss << x;
//    //    reducedModelResults.simulationLabel = ss.str();
//    //    SimulationResultsReporter resultsReporter;
//    //    resultsReporter.reportResults(
//    //        {reducedModelResults},
//    //        std::filesystem::current_path().append("Results"));

//    return true;
//  }
//};

// struct Objective {
//  double operator()(const Eigen::VectorXd &x, Eigen::VectorXd &) const {
//    assert(x.rows() == 4);

//    //      drawSimulationJob(simulationJob);
//    // Set mesh from x
//    std::shared_ptr<SimulationMesh> reducedModel =
//        g_reducedPatternSimulationJob.mesh;
//    for (EdgeIndex ei = 0; ei < reducedModel->EN(); ei++) {
//      if (g_reducedPatternExludedEdges.contains(ei)) {
//        continue;
//      }
//      Element &e = reducedModel->elements[ei];
//      e.axialConstFactor = g_initialStiffnessFactors(ei, 0) * x(0);
//      e.torsionConstFactor = g_initialStiffnessFactors(ei, 1) * x(1);
//      e.firstBendingConstFactor = g_initialStiffnessFactors(ei, 2) * x(2);
//      e.secondBendingConstFactor = g_initialStiffnessFactors(ei, 3) * x(3);
//    }
//    // run simulation
//    SimulationResults reducedModelResults =
//        simulator.executeSimulation(g_reducedPatternSimulationJob);
//    //    std::stringstream ss;
//    //    ss << x;
//    //    reducedModelResults.simulationLabel = ss.str();
//    //    SimulationResultsReporter resultsReporter;
//    //    resultsReporter.reportResults(
//    //        {reducedModelResults},
//    //        std::filesystem::current_path().append("Results"));
//    // compute error and return it
//    double error = 0;
//    for (const auto reducedFullViPair : g_reducedToFullInterfaceViMap) {
//      VertexIndex reducedModelVi = reducedFullViPair.first;
//      Eigen::Vector3d vertexDisplacement(
//          reducedModelResults.displacements[reducedModelVi][0],
//          reducedModelResults.displacements[reducedModelVi][1],
//          reducedModelResults.displacements[reducedModelVi][2]);
//      Eigen::Vector3d errorVector =
//          Eigen::Vector3d(
//              g_optimalReducedModelDisplacements.row(reducedModelVi)) -
//          vertexDisplacement;
//      error += errorVector.norm();
//    }

//    return error;
//  }
//};

double ReducedModelOptimizer::computeError(
    const SimulationResults &reducedPatternResults,
    const Eigen::MatrixX3d &optimalReducedPatternDisplacements) {
  double error = 0;
  for (const auto reducedFullViPair : g_reducedToFullViMap) {
    VertexIndex reducedModelVi = reducedFullViPair.first;
    //    const auto pos =
    //        g_reducedPatternSimulationJob.mesh->vert[reducedModelVi].cP();
    //    std::cout << "Interface vi " << reducedModelVi << " is at position "
    //              << pos[0] << " " << pos[1] << " " << pos[2] << std::endl;
    Eigen::Vector3d vertexDisplacement(
        reducedPatternResults.displacements[reducedModelVi][0],
        reducedPatternResults.displacements[reducedModelVi][1],
        reducedPatternResults.displacements[reducedModelVi][2]);
    if (!std::isfinite(vertexDisplacement[0]) ||
        !std::isfinite(vertexDisplacement[1]) ||
        !std::isfinite(vertexDisplacement[2])) {
      return std::numeric_limits<double>::max();
    }
    Eigen::Vector3d errorVector =
        Eigen::Vector3d(
            optimalReducedPatternDisplacements.row(reducedModelVi)) -
        vertexDisplacement;
    //    error += errorVector.squaredNorm();
    error += errorVector.norm();
  }
  return error;
}

void updateMesh(long n, const double *x) {
  std::shared_ptr<SimulationMesh> &pReducedPatternSimulationMesh =
      g_reducedPatternSimulationJob[0]->pMesh;
  //  const Element &elem = g_reducedPatternSimulationJob[0]->mesh->elements[0];
  //  std::cout << elem.axialConstFactor << " " << elem.torsionConstFactor << "
  //  "
  //            << elem.firstBendingConstFactor << " "
  //            << elem.secondBendingConstFactor << std::endl;
  for (EdgeIndex ei = 0; ei < pReducedPatternSimulationMesh->EN(); ei++) {
    Element &e = pReducedPatternSimulationMesh->elements[ei];
    //    if (g_reducedPatternExludedEdges.contains(ei)) {
    //      continue;
    //    }
    //    e.properties.E = g_initialParameters * x[ei];
    //    e.properties.E = g_initialParameters(0) * x[0];
    //    e.properties.G = g_initialParameters(1) * x[1];
    e.properties.A = g_initialParameters(0) * x[0];
    e.properties.J = g_initialParameters(1) * x[1];
    e.properties.I2 = g_initialParameters(2) * x[2];
    e.properties.I3 = g_initialParameters(3) * x[3];
    //    e.properties.G = e.properties.E / (2 * (1 + 0.3));
    e.axialConstFactor = e.properties.E * e.properties.A / e.initialLength;
    e.torsionConstFactor = e.properties.G * e.properties.J / e.initialLength;
    e.firstBendingConstFactor =
        2 * e.properties.E * e.properties.I2 / e.initialLength;
    e.secondBendingConstFactor =
        2 * e.properties.E * e.properties.I3 / e.initialLength;
  }

  //  std::cout << elem.axialConstFactor << " " << elem.torsionConstFactor << "
  //  "
  //            << elem.firstBendingConstFactor << " "
  //            << elem.secondBendingConstFactor << std::endl;
  //  const Element &e = pReducedPatternSimulationMesh->elements[0];
  //  std::cout << e.axialConstFactor << " " << e.torsionConstFactor << " "
  //            << e.firstBendingConstFactor << " " <<
  //            e.secondBendingConstFactor
  //            << std::endl;

  if (g_optimizeInnerHexagonSize) {
    assert(pReducedPatternSimulationMesh->EN() == 12);
    for (VertexIndex vi = 0; vi < pReducedPatternSimulationMesh->VN();
         vi += 2) {
      pReducedPatternSimulationMesh->vert[vi].P() =
          g_innerHexagonVectors[vi / 2] * x[n - 1];
    }

    pReducedPatternSimulationMesh->reset();
    pReducedPatternSimulationMesh->updateEigenEdgeAndVertices();
    //    pReducedPatternSimulationMesh->registerForDrawing("Optimized
    //    hexagon"); polyscope::show();
  }
}

double ReducedModelOptimizer::objective(double x0, double x1, double x2,
                                        double x3) {
  std::vector<double> x{x0, x1, x2, x3};
  return ReducedModelOptimizer::objective(4, x.data());
}

double ReducedModelOptimizer::objective(long n, const double *x) {
  std::cout.precision(17);

  for (size_t parameterIndex = 0; parameterIndex < n; parameterIndex++) {
    std::cout << "x[" + std::to_string(parameterIndex) + "]="
              << x[parameterIndex] << std::endl;
  }
  const Element &e = g_reducedPatternSimulationJob[0]->pMesh->elements[0];
  //  std::cout << e.axialConstFactor << " " << e.torsionConstFactor << " "
  //            << e.firstBendingConstFactor << " " <<
  //            e.secondBendingConstFactor
  //            << std::endl;
  updateMesh(n, x);
  std::cout << e.axialConstFactor << " " << e.torsionConstFactor << " "
            << e.firstBendingConstFactor << " " << e.secondBendingConstFactor
            << std::endl;

  // run simulations
  double error = 0;
  for (const int simulationScenarioIndex : g_simulationScenarioIndices) {
    SimulationResults reducedModelResults = simulator.executeSimulation(
        g_reducedPatternSimulationJob[simulationScenarioIndex]);
    error += computeError(
        reducedModelResults,
        g_optimalReducedModelDisplacements[simulationScenarioIndex]);
  }
  if (error < minY) {
    minY = error;
    minX = std::vector<double>({x[0], x[1], x[2], x[3]});
  }

  // compute error and return it
  gObjectiveValueHistory.push_back(error);
  //  auto xPlot = matplot::linspace(0, gObjectiveValueHistory.size(),
  //                                 gObjectiveValueHistory.size());
  //  std::vector<double> colors(gObjectiveValueHistory.size(), 2);
  //  if (g_firstRoundIterationIndex != 0) {
  //    for_each(colors.begin() + g_firstRoundIterationIndex, colors.end(),
  //             [](double &c) { c = 0.7; });
  //  }
  //  gPlotHandle = matplot::scatter(xPlot, gObjectiveValueHistory, 6, colors);
  std::cout << std::endl;
  SimulationResultsReporter::createPlot("Number of Steps", "Objective value",
                                        gObjectiveValueHistory);

  return error;
}

void ReducedModelOptimizer::computeMaps(
    FlatPattern &fullPattern, FlatPattern &reducedPattern,
    const std::unordered_set<size_t> &reducedModelExcludedEges) {
  // Compute the offset between the interface nodes
  const size_t interfaceSlotIndex = 4; // bottom edge
  assert(slotToNode.find(interfaceSlotIndex) != slotToNode.end() &&
         slotToNode.find(interfaceSlotIndex)->second.size() == 1);
  // Assuming that in the bottom edge there is only one vertex which is also the
  // interface
  const size_t baseTriangleInterfaceVi =
      *(slotToNode.find(interfaceSlotIndex)->second.begin());

  vcg::tri::Allocator<FlatPattern>::PointerUpdater<FlatPattern::VertexPointer>
      pu_fullModel;
  fullPattern.deleteDanglingVertices(pu_fullModel);
  const size_t fullModelBaseTriangleInterfaceVi =
      pu_fullModel.remap.empty() ? baseTriangleInterfaceVi
                                 : pu_fullModel.remap[baseTriangleInterfaceVi];
  const size_t fullModelBaseTriangleVN = fullPattern.VN();
  fullPattern.createFan();
  const size_t duplicateVerticesPerFanPair =
      fullModelBaseTriangleVN - fullPattern.VN() / 6;
  const size_t fullPatternInterfaceVertexOffset =
      fullModelBaseTriangleVN - duplicateVerticesPerFanPair;
  //  std::cout << "Dups in fan pair:" << duplicateVerticesPerFanPair <<
  //  std::endl;

  // Save excluded edges
  g_reducedPatternExludedEdges.clear();
  const size_t fanSize = 6;
  const size_t reducedBaseTriangleNumberOfEdges = reducedPattern.EN();
  for (size_t fanIndex = 0; fanIndex < fanSize; fanIndex++) {
    for (const size_t ei : reducedModelExcludedEges) {
      g_reducedPatternExludedEdges.insert(
          fanIndex * reducedBaseTriangleNumberOfEdges + ei);
    }
  }

  // Construct reduced->full and full->reduced interface vi map
  reducedToFullInterfaceViMap.clear();
  vcg::tri::Allocator<FlatPattern>::PointerUpdater<FlatPattern::VertexPointer>
      pu_reducedModel;
  reducedPattern.deleteDanglingVertices(pu_reducedModel);
  const size_t reducedModelBaseTriangleInterfaceVi =
      pu_reducedModel.remap[baseTriangleInterfaceVi];
  const size_t reducedModelInterfaceVertexOffset =
      reducedPattern.VN() - 1 /*- reducedModelBaseTriangleInterfaceVi*/;
  reducedPattern.createFan();
  for (size_t fanIndex = 0; fanIndex < fanSize; fanIndex++) {
    reducedToFullInterfaceViMap[reducedModelInterfaceVertexOffset * fanIndex +
                                reducedModelBaseTriangleInterfaceVi] =
        fullModelBaseTriangleInterfaceVi +
        fanIndex * fullPatternInterfaceVertexOffset;
  }
  m_fullToReducedInterfaceViMap.clear();
  constructInverseMap(reducedToFullInterfaceViMap,
                      m_fullToReducedInterfaceViMap);

  //  fullPattern.setLabel("FullPattern");
  //  reducedPattern.setLabel("ReducedPattern");
  // Create opposite vertex map
  m_fullPatternOppositeInterfaceViMap.clear();
  for (int fanIndex = fanSize / 2 - 1; fanIndex >= 0; fanIndex--) {
    const size_t vi0 = fullModelBaseTriangleInterfaceVi +
                       fanIndex * fullPatternInterfaceVertexOffset;
    const size_t vi1 = vi0 + (fanSize / 2) * fullPatternInterfaceVertexOffset;
    assert(vi0 < fullPattern.VN() && vi1 < fullPattern.VN());
    m_fullPatternOppositeInterfaceViMap[vi0] = vi1;
  }

  g_reducedToFullViMap = reducedToFullInterfaceViMap;

  const bool debugMapping = false;
  if (debugMapping) {
    reducedPattern.registerForDrawing();
    std::vector<glm::vec3> colors_reducedPatternExcludedEdges(
        reducedPattern.EN(), glm::vec3(0, 0, 0));
    for (const size_t ei : g_reducedPatternExludedEdges) {
      colors_reducedPatternExcludedEdges[ei] = glm::vec3(1, 0, 0);
    }
    const std::string label = reducedPattern.getLabel();
    polyscope::getCurveNetwork(label)
        ->addEdgeColorQuantity("Excluded edges",
                               colors_reducedPatternExcludedEdges)
        ->setEnabled(true);
    polyscope::show();

    std::vector<glm::vec3> nodeColorsOpposite(fullPattern.VN(),
                                              glm::vec3(0, 0, 0));
    for (const std::pair<size_t, size_t> oppositeVerts :
         m_fullPatternOppositeInterfaceViMap) {
      auto color = polyscope::getNextUniqueColor();
      nodeColorsOpposite[oppositeVerts.first] = color;
      nodeColorsOpposite[oppositeVerts.second] = color;
    }
    fullPattern.registerForDrawing();
    polyscope::getCurveNetwork(fullPattern.getLabel())
        ->addNodeColorQuantity("oppositeMap", nodeColorsOpposite)
        ->setEnabled(true);
    polyscope::show();

    std::vector<glm::vec3> nodeColorsReducedToFull_reduced(reducedPattern.VN(),
                                                           glm::vec3(0, 0, 0));
    std::vector<glm::vec3> nodeColorsReducedToFull_full(fullPattern.VN(),
                                                        glm::vec3(0, 0, 0));
    for (size_t vi = 0; vi < reducedPattern.VN(); vi++) {
      if (reducedToFullInterfaceViMap.contains(vi)) {

        auto color = polyscope::getNextUniqueColor();
        nodeColorsReducedToFull_reduced[vi] = color;
        nodeColorsReducedToFull_full[reducedToFullInterfaceViMap[vi]] = color;
      }
    }
    polyscope::getCurveNetwork(reducedPattern.getLabel())
        ->addNodeColorQuantity("reducedToFull_reduced",
                               nodeColorsReducedToFull_reduced)
        ->setEnabled(true);
    polyscope::getCurveNetwork(fullPattern.getLabel())
        ->addNodeColorQuantity("reducedToFull_full",
                               nodeColorsReducedToFull_full)
        ->setEnabled(true);
    polyscope::show();
  }
}

void ReducedModelOptimizer::createSimulationMeshes(FlatPattern &fullModel,
                                                   FlatPattern &reducedModel) {
  if (typeid(CrossSectionType) != typeid(RectangularBeamDimensions)) {
    std::cerr << "Error: A rectangular cross section is expected." << std::endl;
    terminate();
  }
  m_pReducedPatternSimulationMesh =
      std::make_shared<SimulationMesh>(reducedModel);
  m_pReducedPatternSimulationMesh->setBeamCrossSection(
      CrossSectionType{0.002, 0.002});
  m_pReducedPatternSimulationMesh->setBeamMaterial(0.3, 1);
  m_pFullPatternSimulationMesh = std::make_shared<SimulationMesh>(fullModel);
  m_pFullPatternSimulationMesh->setBeamCrossSection(
      CrossSectionType{0.002, 0.002});
  m_pFullPatternSimulationMesh->setBeamMaterial(0.3, 1);
}

ReducedModelOptimizer::ReducedModelOptimizer(
    const std::vector<size_t> &numberOfNodesPerSlot) {
  FlatPatternTopology::constructNodeToSlotMap(numberOfNodesPerSlot, nodeToSlot);
  FlatPatternTopology::constructSlotToNodeMap(nodeToSlot, slotToNode);
}

void ReducedModelOptimizer::initializePatterns(
    FlatPattern &fullPattern, FlatPattern &reducedPattern,
    const std::unordered_set<size_t> &reducedModelExcludedEdges) {
  fullPattern.setLabel("full_pattern_" + fullPattern.getLabel());
  reducedPattern.setLabel("reduced_pattern_" + reducedPattern.getLabel());
  assert(fullPattern.VN() == reducedPattern.VN() &&
         fullPattern.EN() >= reducedPattern.EN());
  polyscope::removeAllStructures();
  // Create copies of the input models
  FlatPattern copyFullPattern;
  FlatPattern copyReducedPattern;
  copyFullPattern.copy(fullPattern);
  copyReducedPattern.copy(reducedPattern);
  g_optimizeInnerHexagonSize = copyReducedPattern.EN() == 2;
  if (g_optimizeInnerHexagonSize) {
    const double h = copyReducedPattern.getBaseTriangleHeight();
    double baseTriangle_bottomEdgeSize = 2 * h / tan(vcg::math::ToRad(60.0));
    VectorType baseTriangle_leftBottomNode(-baseTriangle_bottomEdgeSize / 2, -h,
                                           0);

    const int fanSize = 6;
    const CoordType rotationAxis(0, 0, 1);
    for (int rotationCounter = 0; rotationCounter < fanSize;
         rotationCounter++) {
      VectorType rotatedVector =
          vcg::RotationMatrix(rotationAxis,
                              vcg::math::ToRad(rotationCounter * 60.0)) *
          baseTriangle_leftBottomNode;
      g_innerHexagonVectors[rotationCounter] = rotatedVector;
    }
    const double innerHexagonInitialPos_x =
        copyReducedPattern.vert[0].cP()[0] / g_innerHexagonVectors[0][0];
    const double innerHexagonInitialPos_y =
        copyReducedPattern.vert[0].cP()[1] / g_innerHexagonVectors[0][1];
    g_innerHexagonInitialPos = innerHexagonInitialPos_x;
  }
  computeMaps(copyFullPattern, copyReducedPattern, reducedModelExcludedEdges);
  createSimulationMeshes(copyFullPattern, copyReducedPattern);
  initializeStiffnesses(m_pReducedPatternSimulationMesh);
}

void ReducedModelOptimizer::initializeStiffnesses(
    const std::shared_ptr<SimulationMesh> &mesh) {
  g_initialParameters.resize(g_optimizeInnerHexagonSize ? 5 : 4);
  // Save save the beam stiffnesses
  //  for (size_t ei = 0; ei < pReducedModelElementalMesh->EN(); ei++) {
  //    Element &e = pReducedModelElementalMesh->elements[ei];
  //    if (g_reducedPatternExludedEdges.contains(ei)) {
  //      const double stiffnessFactor = 5;
  //      e.axialConstFactor *= stiffnessFactor;
  //      e.torsionConstFactor *= stiffnessFactor;
  //      e.firstBendingConstFactor *= stiffnessFactor;
  //      e.secondBendingConstFactor *= stiffnessFactor;
  //    }
  //  g_initialParameters =
  //      m_pReducedPatternSimulationMesh->elements[0].properties.E;
  //  for (size_t ei = 0; ei < m_pReducedPatternSimulationMesh->EN(); ei++) {
  //  }
  //  g_initialParameters(0) = mesh->elements[0].properties.E;
  //  g_initialParameters(1) = mesh->elements[0].properties.G;
  g_initialParameters(0) = mesh->elements[0].properties.A;
  g_initialParameters(1) = mesh->elements[0].properties.J;
  g_initialParameters(2) = mesh->elements[0].properties.I2;
  g_initialParameters(3) = mesh->elements[0].properties.I3;
  //  }
}

void ReducedModelOptimizer::computeReducedModelSimulationJob(
    const SimulationJob &simulationJobOfFullModel,
    const std::unordered_map<size_t, size_t> &simulationJobFullToReducedMap,
    SimulationJob &simulationJobOfReducedModel) {
  assert(simulationJobOfReducedModel.pMesh->VN() != 0);
  std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
      reducedModelFixedVertices;
  for (auto fullModelFixedVertex :
       simulationJobOfFullModel.constrainedVertices) {
    reducedModelFixedVertices[simulationJobFullToReducedMap.at(
        fullModelFixedVertex.first)] = fullModelFixedVertex.second;
  }

  std::unordered_map<VertexIndex, Vector6d> reducedModelNodalForces;
  for (auto fullModelNodalForce :
       simulationJobOfFullModel.nodalExternalForces) {
    reducedModelNodalForces[simulationJobFullToReducedMap.at(
        fullModelNodalForce.first)] = fullModelNodalForce.second;
  }

  //  std::unordered_map<VertexIndex, VectorType>
  //  reducedModelNodalForcedNormals; for (auto fullModelNodalForcedRotation :
  //       simulationJobOfFullModel.nodalForcedNormals) {
  //    reducedModelNodalForcedNormals[simulationJobFullToReducedMap.at(
  //        fullModelNodalForcedRotation.first)] =
  //        fullModelNodalForcedRotation.second;
  //  }
  simulationJobOfReducedModel.constrainedVertices = reducedModelFixedVertices;
  simulationJobOfReducedModel.nodalExternalForces = reducedModelNodalForces;
  simulationJobOfReducedModel.label = simulationJobOfFullModel.getLabel();
  //  simulationJobOfReducedModel.nodalForcedNormals =
  //      reducedModelNodalForcedNormals;
}

void ReducedModelOptimizer::computeDesiredReducedModelDisplacements(
    const SimulationResults &fullModelResults,
    const std::unordered_map<size_t, size_t> &displacementsReducedToFullMap,
    Eigen::MatrixX3d &optimalDisplacementsOfReducedModel) {

  assert(optimalDisplacementsOfReducedModel.rows() != 0 &&
         optimalDisplacementsOfReducedModel.cols() == 3);
  optimalDisplacementsOfReducedModel.setZero(
      optimalDisplacementsOfReducedModel.rows(),
      optimalDisplacementsOfReducedModel.cols());

  for (auto reducedFullViPair : displacementsReducedToFullMap) {
    const VertexIndex fullModelVi = reducedFullViPair.second;
    const Vector6d fullModelViDisplacements =
        fullModelResults.displacements[fullModelVi];
    optimalDisplacementsOfReducedModel.row(reducedFullViPair.first) =
        Eigen::Vector3d(fullModelViDisplacements[0],
                        fullModelViDisplacements[1],
                        fullModelViDisplacements[2]);
  }
}

Eigen::VectorXd ReducedModelOptimizer::runOptimization(
    double (*pObjectiveFunction)(long, const double *)) {

  gObjectiveValueHistory.clear();

  const size_t n = g_optimizeInnerHexagonSize ? 5 : 4;
  //  const size_t npt = (n + 1) * (n + 2) / 2;
  //  //      ((n + 2) + ((n + 1) * (n + 2) / 2)) / 2;
  //  assert(npt <= (n + 1) * (n + 2) / 2 && npt >= n + 2);
  //  assert(npt <= 2 * n + 1 && "The choice of the number of interpolation "
  //                             "conditions is not recommended.");
  // Set initial guess of solution

  const size_t initialGuess = 1;
  std::vector<double> x(n, initialGuess);
  if (g_optimizeInnerHexagonSize) {
    x[n - 1] = g_innerHexagonInitialPos;
  }
  /*if (!initialGuess.empty()) {
    x = g_optimizationInitialGuess;
  }*/ //      {0.10000000000000	001,
  //  2, 1.9999999971613847, 6.9560343643347764};
  //      {1, 5.9277};
  //      {0.0001, 2, 2.000000005047502, 1.3055270196964464};
  //      {initialGuess(0), initialGuess(1), initialGuess(2),
  //      initialGuess(3)};
  const double xMin = 1e-2;
  const double xMax = 1e2;
  //  assert(x.end() == find_if(x.begin(), x.end(), [&](const double &d) {
  //           return d >= xMax || d <= xMin;
  //         }));
  //  std::vector<double> xLow(x.size(), xMin);
  //  std::vector<double> xUpper(x.size(), xMax);
  //  if (g_optimizeInnerHexagonSize) {
  //    xLow[n - 1] = 0.1;
  //    xUpper[n - 1] = 0.9;
  //  }
  //  const double maxX = *std::max_element(
  //      x.begin(), x.end(),
  //      [](const double &a, const double &b) { return abs(a) < abs(b); });
  //  const double rhobeg = std::min(0.95, 0.2 * maxX);
  //  double rhobeg = 1;
  //  double rhoend = rhobeg * 1e-8;
  //  const size_t wSize = (npt + 5) * (npt + n) + 3 * n * (n + 5) / 2;
  //  std::vector<double> w(wSize);
  //  const size_t maxFun = std::min(100.0 * (x.size() + 1), 1000.0);
  const size_t maxFun = 150;

  double (*objF)(double, double, double, double) = &objective;
  auto start = std::chrono::system_clock::now();
  dlib::function_evaluation result = dlib::find_min_global(
      objF, {xMin, xMin, xMin, xMin}, {xMax, xMax, xMax, xMax},
      dlib::max_function_calls(maxFun));
  auto end = std::chrono::system_clock::now();
  auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(end - start);
  std::cout << "Finished optimization with dlib" << endl;
  std::cout << "Solution x:" << endl;
  std::cout << result.x << endl;
  std::cout << "Solution y:" << endl;
  std::cout << result.y << endl;
  std::cout << minY << endl;
  std::cout << "Time(sec):" << elapsed.count() << std::endl;
  std::cout << "Max function evaluations:" << maxFun << std::endl;
  std::cout << "[" + std::to_string(xMin) + "," + std::to_string(xMax) + "]"
            << std::endl;
  std::cout << "Initial guess:" << initialGuess << std::endl;
  //  const size_t maxFun = 200;
  //  bobyqa(pObjectiveFunction, n, npt, x.data(), xLow.data(), xUpper.data(),
  //         rhobeg, rhoend, maxFun, w.data());
  //  std::cout << "Finished first optimization round" << std::endl;
  //  firstOptimizationRoundResults.resize(6);
  //  for (int simulationScenarioIndex = SimulationScenario::Axial;
  //       simulationScenarioIndex !=
  //       SimulationScenario::NumberOfSimulationScenarios;
  //       simulationScenarioIndex++) {
  //    SimulationResults reducedModelResults = simulator.executeSimulation(
  //        g_reducedPatternSimulationJob[simulationScenarioIndex], false,
  //        false);
  //    reducedModelResults.setLabelPrefix("FirstRound");
  //    firstOptimizationRoundResults[simulationScenarioIndex] =
  //        std::move(reducedModelResults);
  //  }
  //  g_firstRoundIterationIndex = gObjectiveValueHistory.size();
  //  rhobeg *= 1e1;
  //  //  rhoend *= 1e2;
  //  bobyqa(pObjectiveFunction, n, npt, x.data(), xLow.data(), xUpper.data(),
  //         rhobeg, rhoend, maxFun, w.data());
  //  std::cout << "Finished second optimization round" << std::endl;

  Eigen::VectorXd eigenX(x.size(), 1);
  for (size_t xi = 0; xi < x.size(); xi++) {
    eigenX(xi) = minX[xi];
    //    eigenX(xi) = x[xi];
    //    eigenX(xi) = result.x(xi);
  }
  //  for (size_t xi = 0; xi < x.size(); xi++) {
  //  }
  return eigenX;
}

void ReducedModelOptimizer::setInitialGuess(std::vector<double> v) {
  initialGuess = v;
}

std::vector<std::shared_ptr<SimulationJob>>
ReducedModelOptimizer::createScenarios(
    const std::shared_ptr<SimulationMesh> &pMesh) {
  std::vector<std::shared_ptr<SimulationJob>> scenarios;
  scenarios.resize(SimulationScenario::NumberOfSimulationScenarios);
  std::unordered_map<VertexIndex, std::unordered_set<DoFType>> fixedVertices;
  std::unordered_map<VertexIndex, Vector6d> nodalForces;
  const double forceMagnitude = 1;
  // Assuming the patterns lays on the x-y plane
  const CoordType patternPlaneNormal(0, 0, 1);
  // Make the first interface node lay on the x axis
  //  const size_t fullPatternFirstInterfaceNodeIndex =
  //      m_fullPatternOppositeInterfaceViMap.begin()->second;
  //  CoordType fullPatternFirstInterfaceNodePosition =
  //      m_pFullModelSimulationMesh->vert[fullPatternFirstInterfaceNodeIndex].cP();
  //  CoordType centerOfMass(0, 0, 0);
  //  for (size_t vi = 0; vi < pMesh->VN(); vi++) {
  //    centerOfMass = centerOfMass + pMesh->vert[vi].P();
  //  }
  //  centerOfMass /= pMesh->VN();
  //  vcg::tri::UpdatePosition<SimulationMesh>::Translate(
  //      *m_pFullModelSimulationMesh, -centerOfMass);
  //  vcg::tri::UpdatePosition<SimulationMesh>::Translate(
  //      *m_pReducedPatternSimulationMesh, centerOfMass);

  //  const vcg::Matrix33d R = vcg::RotationMatrix(
  //      fullPatternFirstInterfaceNodePosition,
  //      CoordType(fullPatternFirstInterfaceNodePosition.Norm(), 0, 0), false);
  //  std::for_each(m_pFullModelSimulationMesh->vert.begin(),
  //                m_pFullModelSimulationMesh->vert.end(), [&](auto &v) {
  //                  v.P() = R * v.P();
  //                  v.N() = R * v.N();
  //                });
  //  std::for_each(m_pReducedPatternSimulationMesh->vert.begin(),
  //                m_pReducedPatternSimulationMesh->vert.end(), [&](auto &v) {
  //                  v.P() = R * v.P();
  //                  v.N() = R * v.N();
  //                });

  //  m_pFullModelSimulationMesh->updateEigenEdgeAndVertices();
  //  m_pReducedPatternSimulationMesh->updateEigenEdgeAndVertices();

  //// Axial
  SimulationScenario scenarioName = SimulationScenario::Axial;
  for (const auto &viPair : m_fullPatternOppositeInterfaceViMap) {
    CoordType forceDirection =
        (pMesh->vert[viPair.first].cP() - pMesh->vert[viPair.second].cP())
            .Normalize();
    nodalForces[viPair.first] = Vector6d({forceDirection[0], forceDirection[1],
                                          forceDirection[2], 0, 0, 0}) *
                                forceMagnitude * 10;
    fixedVertices[viPair.second] =
        std::unordered_set<DoFType>{0, 1, 2, 3, 4, 5};
  }
  scenarios[scenarioName] = std::make_shared<SimulationJob>(
      SimulationJob(pMesh, simulationScenarioStrings[scenarioName],
                    fixedVertices, nodalForces, {}));

  //// Shear
  scenarioName = SimulationScenario::Shear;
  fixedVertices.clear();
  nodalForces.clear();
  for (const auto &viPair : m_fullPatternOppositeInterfaceViMap) {
    CoordType v =
        (pMesh->vert[viPair.first].cP() - pMesh->vert[viPair.second].cP())
            .Normalize();
    CoordType forceDirection = (v ^ patternPlaneNormal).Normalize();
    nodalForces[viPair.first] = Vector6d({forceDirection[0], forceDirection[1],
                                          forceDirection[2], 0, 0, 0}) *
                                0.40 * forceMagnitude;
    fixedVertices[viPair.second] =
        std::unordered_set<DoFType>{0, 1, 2, 3, 4, 5};
  }
  scenarios[scenarioName] = std::make_shared<SimulationJob>(
      SimulationJob(pMesh, simulationScenarioStrings[scenarioName],
                    fixedVertices, nodalForces, {}));

  //  //// Torsion
  //  fixedVertices.clear();
  //  nodalForces.clear();
  //  for (auto viPairIt = m_fullPatternOppositeInterfaceViMap.begin();
  //       viPairIt != m_fullPatternOppositeInterfaceViMap.end(); viPairIt++) {
  //    const auto &viPair = *viPairIt;
  //    if (viPairIt == m_fullPatternOppositeInterfaceViMap.begin()) {
  //      CoordType v =
  //          (pMesh->vert[viPair.first].cP() - pMesh->vert[viPair.second].cP())
  //              .Normalize();
  //      CoordType normalVDerivativeDir = (v ^ patternPlaneNormal).Normalize();
  //      nodalForces[viPair.first] = Vector6d{
  //          0, 0, 0, normalVDerivativeDir[0], normalVDerivativeDir[1], 0};
  //      fixedVertices[viPair.second] =
  //          std::unordered_set<DoFType>{0, 1, 2, 3, 4, 5};
  //      fixedVertices[viPair.first] = std::unordered_set<DoFType>{0, 1, 2};
  //    } else {
  //      fixedVertices[viPair.first] = std::unordered_set<DoFType>{0, 1, 2};
  //      fixedVertices[viPair.second] = std::unordered_set<DoFType>{0, 1, 2};
  //    }
  //  }
  //  scenarios.push_back({pMesh, fixedVertices, nodalForces});

  ////  Bending
  scenarioName = SimulationScenario::Bending;
  fixedVertices.clear();
  nodalForces.clear();
  for (const auto &viPair : m_fullPatternOppositeInterfaceViMap) {
    nodalForces[viPair.first] = Vector6d({0, 0, forceMagnitude, 0, 0, 0}) * 1;
    fixedVertices[viPair.second] =
        std::unordered_set<DoFType>{0, 1, 2, 3, 4, 5};
  }
  scenarios[scenarioName] = std::make_shared<SimulationJob>(
      SimulationJob(pMesh, simulationScenarioStrings[scenarioName],
                    fixedVertices, nodalForces, {}));

  //// Double using moments
  scenarioName = SimulationScenario::Dome;
  fixedVertices.clear();
  nodalForces.clear();
  for (auto viPairIt = m_fullPatternOppositeInterfaceViMap.begin();
       viPairIt != m_fullPatternOppositeInterfaceViMap.end(); viPairIt++) {
    const auto viPair = *viPairIt;
    if (viPairIt == m_fullPatternOppositeInterfaceViMap.begin()) {
      fixedVertices[viPair.first] = std::unordered_set<DoFType>{0, 1, 2};
      fixedVertices[viPair.second] = std::unordered_set<DoFType>{0, 2};
    } else {
      fixedVertices[viPair.first] = std::unordered_set<DoFType>{2};
      fixedVertices[viPair.second] = std::unordered_set<DoFType>{2};
    }
    CoordType v =
        (pMesh->vert[viPair.first].cP() - pMesh->vert[viPair.second].cP())
            .Normalize();
    nodalForces[viPair.first] =
        Vector6d({0, 0, 0, v[0], v[1], 0}) * forceMagnitude * 0.1;
    nodalForces[viPair.second] =
        Vector6d({0, 0, 0, -v[0], -v[1], 0}) * forceMagnitude * 0.1;
  }
  scenarios[scenarioName] = std::make_shared<SimulationJob>(
      SimulationJob(pMesh, simulationScenarioStrings[scenarioName],
                    fixedVertices, nodalForces, {}));

  //// Saddle
  scenarioName = SimulationScenario::Saddle;
  fixedVertices.clear();
  nodalForces.clear();
  for (auto viPairIt = m_fullPatternOppositeInterfaceViMap.begin();
       viPairIt != m_fullPatternOppositeInterfaceViMap.end(); viPairIt++) {
    const auto &viPair = *viPairIt;
    CoordType v =
        (pMesh->vert[viPair.first].cP() - pMesh->vert[viPair.second].cP())
            .Normalize();
    if (viPairIt == m_fullPatternOppositeInterfaceViMap.begin()) {
      nodalForces[viPair.first] =
          Vector6d({0, 0, 0, v[0], v[1], 0}) * 0.02 * forceMagnitude;
      nodalForces[viPair.second] =
          Vector6d({0, 0, 0, -v[0], -v[1], 0}) * 0.02 * forceMagnitude;
    } else {
      fixedVertices[viPair.first] = std::unordered_set<DoFType>{2};
      fixedVertices[viPair.second] = std::unordered_set<DoFType>{0, 1, 2};

      nodalForces[viPair.first] =
          Vector6d({0, 0, 0, -v[0], -v[1], 0}) * 0.01 * forceMagnitude;
      nodalForces[viPair.second] =
          Vector6d({0, 0, 0, v[0], v[1], 0}) * 0.01 * forceMagnitude;
    }
  }
  scenarios[scenarioName] = std::make_shared<SimulationJob>(
      SimulationJob(pMesh, simulationScenarioStrings[scenarioName],
                    fixedVertices, nodalForces, {}));

  return scenarios;
}

void ReducedModelOptimizer::runBeamOptimization() {
  // load beams
  VCGEdgeMesh fullBeam;
  fullBeam.loadFromPly("/home/iason/Models/simple_beam_model_10elem_1m.ply");
  VCGEdgeMesh reducedBeam;
  reducedBeam.loadFromPly("/home/iason/Models/simple_beam_model_4elem_1m.ply");
  fullBeam.registerForDrawing();
  reducedBeam.registerForDrawing();
  //  polyscope::show();
  // maps
  std::unordered_map<size_t, size_t> displacementReducedToFullMap;
  displacementReducedToFullMap[reducedBeam.VN() / 2] = fullBeam.VN() / 2;
  g_reducedToFullViMap = displacementReducedToFullMap;
  std::unordered_map<size_t, size_t> jobFullToReducedMap;
  jobFullToReducedMap[0] = 0;
  jobFullToReducedMap[fullBeam.VN() - 1] = reducedBeam.VN() - 1;

  // full model simuilation job
  auto pFullPatternSimulationMesh = std::make_shared<SimulationMesh>(fullBeam);
  pFullPatternSimulationMesh->setBeamCrossSection(CrossSectionType{0.02, 0.02});
  pFullPatternSimulationMesh->setBeamMaterial(0.3, 1);
  std::unordered_map<VertexIndex, std::unordered_set<int>> fixedVertices;
  fixedVertices[0] = ::unordered_set<int>({0, 1, 2, 3, 4, 5});
  std::unordered_map<VertexIndex, Vector6d> forces;
  forces[fullBeam.VN() - 1] = Vector6d({0, 0, 10, 0, 0, 0});
  const std::string fullBeamSimulationJobLabel = "Pull_Z";
  std::shared_ptr<SimulationJob> pFullModelSimulationJob =
      make_shared<SimulationJob>(SimulationJob(pFullPatternSimulationMesh,
                                               fullBeamSimulationJobLabel,
                                               fixedVertices, forces));
  FormFinder formFinder;
  auto fullModelResults = formFinder.executeSimulation(pFullModelSimulationJob);

  // Optimal reduced model displacements
  const size_t numberOfSimulationScenarios = 1;
  g_optimalReducedModelDisplacements.resize(numberOfSimulationScenarios);
  g_optimalReducedModelDisplacements[numberOfSimulationScenarios - 1].resize(
      reducedBeam.VN(), 3);
  computeDesiredReducedModelDisplacements(
      fullModelResults, displacementReducedToFullMap,
      g_optimalReducedModelDisplacements[numberOfSimulationScenarios - 1]);

  // reduced model simuilation job
  auto reducedSimulationMesh = std::make_shared<SimulationMesh>(reducedBeam);
  reducedSimulationMesh->setBeamCrossSection(CrossSectionType{0.02, 0.02});
  reducedSimulationMesh->setBeamMaterial(0.3, 1);
  g_reducedPatternSimulationJob.resize(numberOfSimulationScenarios);
  SimulationJob reducedSimJob;
  computeReducedModelSimulationJob(*pFullModelSimulationJob,
                                   jobFullToReducedMap, reducedSimJob);
  reducedSimJob.nodalExternalForces[reducedBeam.VN() - 1] =
      reducedSimJob.nodalExternalForces[reducedBeam.VN() - 1] * 0.1;
  g_reducedPatternSimulationJob[numberOfSimulationScenarios - 1] =
      make_shared<SimulationJob>(
          reducedSimulationMesh, fullBeamSimulationJobLabel,
          reducedSimJob.constrainedVertices, reducedSimJob.nodalExternalForces);
  initializeStiffnesses(reducedSimulationMesh);

  //  const std::string simulationJobsPath = "SimulationJobs";
  //  std::filesystem::create_directory(simulationJobsPath);
  //  g_reducedPatternSimulationJob[0].save(simulationJobsPath);
  //  g_reducedPatternSimulationJob[0].load(
  //      std::filesystem::path(simulationJobsPath)
  //          .append(g_reducedPatternSimulationJob[0].mesh->getLabel() +
  //                  "_simScenario.json"));

  runOptimization(&objective);

  fullModelResults.registerForDrawing();
  SimulationResults reducedModelResults = simulator.executeSimulation(
      g_reducedPatternSimulationJob[numberOfSimulationScenarios - 1]);
  double error = computeError(
      reducedModelResults,
      g_optimalReducedModelDisplacements[numberOfSimulationScenarios - 1]);
  reducedModelResults.registerForDrawing();
  std::cout << "Error between beams:" << error << endl;
  //    registerWorldAxes();
  polyscope::show();
  fullModelResults.unregister();
  reducedModelResults.unregister();
}

void ReducedModelOptimizer::visualizeResults(
    const std::vector<std::shared_ptr<SimulationJob>>
        &fullPatternSimulationJobs,
    const std::vector<SimulationScenario> &simulationScenarios) {
  m_pFullPatternSimulationMesh->registerForDrawing();
  double error = 0;
  for (const int simulationScenarioIndex : simulationScenarios) {
    const std::shared_ptr<SimulationJob> &pFullPatternSimulationJob =
        fullPatternSimulationJobs[simulationScenarioIndex];
    pFullPatternSimulationJob->registerForDrawing(
        m_pFullPatternSimulationMesh->getLabel());
    SimulationResults fullModelResults =
        simulator.executeSimulation(pFullPatternSimulationJob);
    fullModelResults.registerForDrawing();
    fullModelResults.saveDeformedModel();
    const std::shared_ptr<SimulationJob> &pReducedPatternSimulationJob =
        g_reducedPatternSimulationJob[simulationScenarioIndex];
    SimulationResults reducedModelResults =
        simulator.executeSimulation(pReducedPatternSimulationJob);
    error += computeError(
        reducedModelResults,
        g_optimalReducedModelDisplacements[simulationScenarioIndex]);
    std::cout << "Error of simulation scenario "
              << simulationScenarioStrings[simulationScenarioIndex] << " is "
              << error << std::endl;
    reducedModelResults.registerForDrawing();
    //    firstOptimizationRoundResults[simulationScenarioIndex].registerForDrawing();
    //    reducedModelResults.saveDeformedModel();
    //    registerWorldAxes();
    const std::string screenshotFilename =
        "/home/iason/Coding/Projects/Approximating shapes with flat "
        "patterns/RodModelOptimizationForPatterns/build/OptimizationResults/" +
        m_pFullPatternSimulationMesh->getLabel() + "_" +
        simulationScenarioStrings[simulationScenarioIndex];
    polyscope::show();
    polyscope::screenshot(screenshotFilename, false);
    fullModelResults.unregister();
    reducedModelResults.unregister();
    //    firstOptimizationRoundResults[simulationScenarioIndex].unregister();
  }
}

void ReducedModelOptimizer::optimize(
    const std::vector<SimulationScenario> &simulationScenarios) {

  g_simulationScenarioIndices = simulationScenarios;
  if (g_simulationScenarioIndices.empty()) {
    g_simulationScenarioIndices = {
        SimulationScenario::Axial, SimulationScenario::Shear,
        SimulationScenario::Bending, SimulationScenario::Dome,
        SimulationScenario::Saddle};
  }

  std::vector<std::shared_ptr<SimulationJob>> simulationJobs =
      createScenarios(m_pFullPatternSimulationMesh);
  g_optimalReducedModelDisplacements.resize(6);
  g_reducedPatternSimulationJob.resize(6);
  g_firstRoundIterationIndex = 0;
  minY = std::numeric_limits<double>::max();
  //  polyscope::removeAllStructures();

  for (int simulationScenarioIndex : g_simulationScenarioIndices) {
    const std::shared_ptr<SimulationJob> &pFullPatternSimulationJob =
        simulationJobs[simulationScenarioIndex];
    SimulationResults fullModelResults =
        simulator.executeSimulation(pFullPatternSimulationJob);
    g_optimalReducedModelDisplacements[simulationScenarioIndex].resize(
        m_pReducedPatternSimulationMesh->VN(), 3);
    computeDesiredReducedModelDisplacements(
        fullModelResults, g_reducedToFullViMap,
        g_optimalReducedModelDisplacements[simulationScenarioIndex]);
    SimulationJob reducedPatternSimulationJob;
    reducedPatternSimulationJob.pMesh = m_pReducedPatternSimulationMesh;
    computeReducedModelSimulationJob(*pFullPatternSimulationJob,
                                     m_fullToReducedInterfaceViMap,
                                     reducedPatternSimulationJob);
    g_reducedPatternSimulationJob[simulationScenarioIndex] =
        std::make_shared<SimulationJob>(reducedPatternSimulationJob);
  }

  Eigen::VectorXd optimalParameters = runOptimization(&objective);
  updateMesh(4, optimalParameters.data());
  visualizeResults(simulationJobs, g_simulationScenarioIndices);
}