#ifndef BEAMFORMFINDER_HPP
#define BEAMFORMFINDER_HPP

//#ifdef USE_MATPLOT
//#include "matplot.h"
#include <matplot/matplot.h>
//#endif
#include "simulationmesh.hpp"
#include "simulationmodel.hpp"
#include <Eigen/Dense>
#include <filesystem>
#include <unordered_set>

#ifdef USE_ENSMALLEN
#include <armadillo>
#include <ensmallen.hpp>
#endif

class SimulationJob;

enum DoF { Ux = 0, Uy, Uz, Nx, Ny, Nr, NumDoF };
using DoFType = int;
using EdgeType = VCGEdgeMesh::EdgeType;
using VertexType = VCGEdgeMesh::VertexType;

struct DifferentiateWithRespectTo {
  const VertexType &v;
  const DoFType &dofi;
};

class DRMSimulationModel : public SimulationModel
{
public:
    struct Settings
    {
        bool useTranslationalKineticEnergyForKineticDamping{true};
        bool useTotalRotationalKineticEnergyForKineticDamping{false};
        bool shouldDraw{false};
        bool beVerbose{false};
        bool shouldCreatePlots{false};
        //        int drawingStep{0};
        //        double residualForcesMovingAverageDerivativeNormThreshold{1e-8};
        //        double residualForcesMovingAverageNormThreshold{1e-8};
        double Dtini{0.1};
        double xi{0.9969};
        std::optional<double> translationalKineticEnergyThreshold;
        int gradualForcedDisplacementSteps{50};
        //        int desiredGradualExternalLoadsSteps{1};
        double gamma{0.8};
        double totalResidualForcesNormThreshold{1e-20};
        double totalExternalForcesNormPercentageTermination{1e-5};
        std::optional<int> maxDRMIterations;
        std::optional<int> debugModeStep;
        std::optional<double> totalTranslationalKineticEnergyThreshold;
        std::optional<double> averageResidualForcesCriterionThreshold;
        std::optional<double> linearGuessForceScaleFactor;
        //        std::optional<int> intermediateResultsSaveStep;
        std::optional<bool> saveIntermediateBestStates;
        std::optional<double> viscousDampingFactor;
        Settings() {}
        void save(const std::filesystem::path &jsonFilePath) const;
        bool load(const std::filesystem::path &filePath);
        struct JsonLabels
        {
            const std::string shouldDraw{"shouldDraw"};
            const std::string beVerbose{"beVerbose"};
            const std::string shouldCreatePlots{"shouldCreatePlots"};
            const std::string Dtini{"DtIni"};
            const std::string xi{"xi"};
            const std::string gamma{"gamma"};
            const std::string totalResidualForcesNormThreshold;
            const std::string maxDRMIterations{"maxIterations"};
            const std::string debugModeStep{"debugModeStep"};
            const std::string totalTranslationalKineticEnergyThreshold{
                "totalTranslationaKineticEnergyThreshold"};
            const std::string averageResidualForcesCriterionThreshold{
                "averageResidualForcesThreshold"};
            const std::string linearGuessForceScaleFactor{"linearGuessForceScaleFactor"};
            const std::string viscousDampingFactor{"viscousDampingFactor"};
        };
        static JsonLabels jsonLabels;
    };

private:
    Settings mSettings;
    double Dt{mSettings.Dtini};
    bool checkedForMaximumMoment{false};
    bool shouldTemporarilyDampForces{false};
    bool shouldTemporarilyAmplifyForces{true};
    double externalMomentsNorm{0};
    size_t mCurrentSimulationStep{0};
    matplot::line_handle plotHandle;
    std::vector<double> plotYValues;
    size_t numOfDampings{0};
    int externalLoadStep{1};
    std::vector<bool> isVertexConstrained;
    std::vector<bool> isRigidSupport;
    double minTotalResidualForcesNorm{std::numeric_limits<double>::max()};

    const std::string meshPolyscopeLabel{"Simulation mesh"};
    std::unique_ptr<SimulationMesh> pMesh;
    std::unordered_map<VertexIndex, std::unordered_set<DoFType>> constrainedVertices;
    SimulationHistory history;
    //  Eigen::Tensor<double, 4> theta3Derivatives;
    //  std::unordered_map<MyKeyType, double, key_hash> theta3Derivatives;
    bool shouldApplyInitialDistortion = false;
    //#ifdef USE_ENSMALLEN
    //    std::shared_ptr<SimulationJob> pJob;
    //#endif
    void reset(const std::shared_ptr<SimulationJob> &pJob, const Settings &settings);
    void updateNodalInternalForces(
        const std::unordered_map<VertexIndex, std::unordered_set<DoFType>> &fixedVertices);
    void updateNodalExternalForces(
        const std::unordered_map<VertexIndex, Vector6d> &nodalForces,
        const std::unordered_map<VertexIndex, std::unordered_set<DoFType>> &fixedVertices);
    void updateResidualForces();
    void updateRotationalDisplacements();
    void updateElementalLengths();

    void updateNodalMasses();

    void updateNodalAccelerations();

    void updateNodalVelocities();

    void updateNodalDisplacements();

    void applyForcedDisplacements(
        const std::unordered_map<VertexIndex, Eigen::Vector3d> &nodalForcedDisplacements);

    Vector6d computeElementTorsionalForce(const Element &element,
                                          const Vector6d &displacementDifference,
                                          const std::unordered_set<DoFType> &constrainedDof);

    //  BeamFormFinder::Vector6d computeElementFirstBendingForce(
    //      const Element &element, const Vector6d &displacementDifference,
    //      const std::unordered_set<gsl::index> &constrainedDof);

    //  BeamFormFinder::Vector6d computeElementSecondBendingForce(
    //      const Element &element, const Vector6d &displacementDifference,
    //      const std::unordered_set<gsl::index> &constrainedDof);

    void updateKineticEnergy();

    void resetVelocities();

    SimulationResults computeResults(const std::shared_ptr<SimulationJob> &pJob);

    void updateNodePosition(
        VertexType &v,
        const std::unordered_map<VertexIndex, std::unordered_set<DoFType>> &fixedVertices);

    void applyDisplacements(
        const std::unordered_map<VertexIndex, std::unordered_set<DoFType>> &fixedVertices);

#ifdef POLYSCOPE_DEFINED
    void draw(const std::shared_ptr<SimulationJob> &pJob, const std::string &screenshotsFolder = {});
#endif
  void
  updateNodalInternalForce(Vector6d &nodalInternalForce,
                           const Vector6d &elementInternalForce,
                           const std::unordered_set<DoFType> &nodalFixedDof);

  Vector6d computeElementInternalForce(
      const Element &elem, const Node &n0, const Node &n1,
      const std::unordered_set<DoFType> &n0ConstrainedDof,
      const std::unordered_set<DoFType> &n1ConstrainedDof);

  Vector6d computeElementAxialForce(const ::EdgeType &e) const;
  VectorType computeDisplacementDifferenceDerivative(
      const EdgeType &e, const DifferentiateWithRespectTo &dui) const;
  double
  computeDerivativeElementLength(const EdgeType &e,
                                 const DifferentiateWithRespectTo &dui) const;

  VectorType computeDerivativeT1(const EdgeType &e,
                                 const DifferentiateWithRespectTo &dui) const;

  VectorType
  computeDerivativeOfNormal(const VertexType &v,
                            const DifferentiateWithRespectTo &dui) const;

  VectorType computeDerivativeT3(const EdgeType &e,
                                 const DifferentiateWithRespectTo &dui) const;

  VectorType computeDerivativeT2(const EdgeType &e,
                                 const DifferentiateWithRespectTo &dui) const;

  double computeDerivativeTheta2(const EdgeType &e, const VertexIndex &evi,
                                 const VertexIndex &dwrt_evi,
                                 const DoFType &dwrt_dofi) const;

  void updateElementalFrames();

  VectorType computeDerivativeOfR(const EdgeType &e,
                                  const DifferentiateWithRespectTo &dui) const;


  static double computeDerivativeOfNorm(const VectorType &x,
                                        const VectorType &derivativeOfX);
  static VectorType computeDerivativeOfCrossProduct(
      const VectorType &a, const VectorType &derivativeOfA, const VectorType &b,
      const VectorType &derivativeOfB);

  double computeTheta3(const EdgeType &e, const VertexType &v);
  double computeDerivativeTheta3(const EdgeType &e, const VertexType &v,
                                 const DifferentiateWithRespectTo &dui) const;
  double computeTotalPotentialEnergy();
  void computeRigidSupports();
  void updateNormalDerivatives();
  void updateT1Derivatives();
  void updateT2Derivatives();
  void updateT3Derivatives();
  void updateRDerivatives();

  double computeDerivativeTheta1(const EdgeType &e, const VertexIndex &evi,
                                 const VertexIndex &dwrt_evi,
                                 const DoFType &dwrt_dofi) const;

  //  void updatePositionsOnTheFly(
  //      const std::unordered_map<VertexIndex,
  //      std::unordered_set<gsl::index>>
  //          &fixedVertices);

  void updateResidualForcesOnTheFly(
      const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
          &fixedVertices);

  void updatePositionsOnTheFly(
      const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
          &fixedVertices);

  void updateNodeNormal(
      VertexType &v,
      const std::unordered_map<VertexIndex, std::unordered_set<DoFType>>
          &fixedVertices);

  void applyForcedNormals(
      const std::unordered_map<VertexIndex, VectorType> nodalForcedRotations);

  void printCurrentState() const;

  void printDebugInfo() const;

  double computeTotalInternalPotentialEnergy();

  void applySolutionGuess(const SimulationResults &solutionGuess,
                          const std::shared_ptr<SimulationJob> &pJob);

  void updateNodeNr(VertexType &v);

  void applyKineticDamping(const std::shared_ptr<SimulationJob> &pJob);

  virtual SimulationResults executeSimulation(const std::shared_ptr<SimulationJob> &pJob) override;
  void setStructure(const std::shared_ptr<SimulationMesh> &pMesh) override;

  void reset(const std::shared_ptr<SimulationJob> &pJob);

  std::vector<std::array<Vector6d, 4>> computeInternalForces(
      const std::unordered_map<VertexIndex, std::unordered_set<DoFType>> &fixedVertices);

  public:
  DRMSimulationModel();
  SimulationResults executeSimulation(const std::shared_ptr<SimulationJob> &pJob,
                                      const Settings &settings,
                                      const SimulationResults &solutionGuess = SimulationResults());
  //#ifdef USE_ENSMALLEN
  //  std::shared_ptr<SimulationJob> pJob;
  //  double EvaluateWithGradient(const arma::mat &x, arma::mat &g);
  //  void setJob(const std::shared_ptr<SimulationJob> &pJob);
  //  SimulationMesh *getDeformedMesh(const arma::mat &x, const std::shared_ptr<SimulationJob> &pJob);
  //  double Evaluate(const arma::mat &x);
  //#endif

  static void runUnitTests();
};

inline DRMSimulationModel::Settings::JsonLabels DRMSimulationModel::Settings::jsonLabels;

template <typename PointType> PointType Cross(PointType p1, PointType p2) {
  return p1 ^ p2;
}

inline size_t currentStep{0};
inline bool TriggerBreakpoint(const VertexIndex &vi, const EdgeIndex &ei,
                              const DoFType &dofi) {
  const size_t numberOfVertices = 10;
  const VertexIndex middleNodeIndex = numberOfVertices / 2;
  //  return vi == middleNodeIndex && dofi == 1;
  return dofi == 1 && ((vi == 1 && ei == 0) || (vi == 9 && ei == 9));
}

inline bool TriggerBreakpoint(const VertexIndex &vi, const EdgeIndex &ei) {
  const size_t numberOfVertices = 10;
  const VertexIndex middleNodeIndex = numberOfVertices / 2;
  return (vi == middleNodeIndex);
  //  return (vi == 0 || vi == numberOfVertices - 1) && currentStep == 1;
  return (vi == 1 && ei == 0) || (vi == 9 && ei == 9);
}

#endif // BEAMFORMFINDER_HPP