#include "chronoseulersimulationmodel.hpp"
#include <chrono/fea/ChBeamSectionEuler.h>
#include <chrono/fea/ChBuilderBeam.h>
#include <chrono/fea/ChLoadsBeam.h>
#include <chrono/fea/ChMesh.h>
#include <chrono/fea/ChNodeFEAxyzrot.h>
#include <chrono/physics/ChBody.h>
#include <chrono/physics/ChSystemSMC.h>
#include <chrono/solver/ChIterativeSolverLS.h>
#include "chrono/physics/ChLoadContainer.h"

#include <chrono/assets/ChVisualization.h>
#include <chrono/fea/ChElementBeamEuler.h>
#include <chrono/fea/ChVisualizationFEAmesh.h>

using namespace chrono;
using namespace chrono::fea;
std::shared_ptr<ChMesh> ChronosEulerSimulationModel::convertToChronosMesh_Euler(
    const std::shared_ptr<SimulationEdgeMesh>& pMesh,
    std::vector<std::shared_ptr<ChNodeFEAxyzrot>>&
        edgeMeshVertsToChronosNodes) {
  auto mesh_chronos = chrono_types::make_shared<ChMesh>();
  edgeMeshVertsToChronosNodes.clear();
  edgeMeshVertsToChronosNodes.resize(pMesh->VN(), nullptr);

  // add nodes
  for (int vi = 0; vi < pMesh->VN(); vi++) {
    const auto& vertex = pMesh->vert[vi];
    ChVector<> vertexPos(vertex.cP()[0], vertex.cP()[1], vertex.cP()[2]);
    edgeMeshVertsToChronosNodes[vi] =
        chrono_types::make_shared<ChNodeFEAxyzrot>(ChFrame<>(vertexPos));
    mesh_chronos->AddNode(edgeMeshVertsToChronosNodes[vi]);
  }
  // add elements
  ChBuilderBeamEuler builder;
  for (int ei = 0; ei < pMesh->EN(); ei++) {
    const SimulationEdgeMesh::EdgeType& edge = pMesh->edge[ei];
    // define end-points
    const auto vi0 = pMesh->getIndex(edge.cV(0));
    const auto vi1 = pMesh->getIndex(edge.cV(1));
    // define cross section
    const Element& element = pMesh->elements[ei];
    const double beam_wz = element.dimensions.getDim1();
    const double beam_wy = element.dimensions.getDim2();
    const double E = element.material.youngsModulus;
    //        const double poisson = element.material.poissonsRatio;
    const double density = 1e0;

    //        auto msection =
    //        chrono_types::make_shared<ChBeamSectionEulerAdvanced>();
    auto msection =
        chrono_types::make_shared<ChBeamSectionEulerEasyRectangular>(
            beam_wy, beam_wz, E, element.material.G, density);
    msection->SetArea(element.dimensions.A);
    msection->SetIyy(element.dimensions.inertia.I2);
    msection->SetIzz(element.dimensions.inertia.I3);
    msection->SetJ(element.dimensions.inertia.J);
    //        msection->SetDensity(density);
    //        msection->SetYoungModulus(E);
    //        msection->SetGwithPoissonRatio(poisson);
    //        //        msection->SetBeamRaleyghDamping(0.0);
    //        msection->SetAsRectangularSection(beam_wy, beam_wz);
    builder.BuildBeam(
        mesh_chronos,  // the mesh where to put the created nodes and elements
        msection,  // the ChBeamSectionEuler to use for the ChElementBeamEuler
                   // elements
        1,         // the number of ChElementBeamEuler to create
        edgeMeshVertsToChronosNodes[vi0],  // the 'A' point in space (beginning
                                           // of beam)
        edgeMeshVertsToChronosNodes[vi1],  // the 'B' point in space (end of
                                           // beam)
        ChVector<>(0, 0, 1)
        // ChVector<>(0, cos(rot_rad), sin(rot_rad))
    );  // the 'Y' up direction of the section for the beam
  }

  return mesh_chronos;
}

ChronosEulerSimulationModel::ChronosEulerSimulationModel() {}

// SimulationResults ChronosEulerSimulationModel::executeSimulation(
//    const std::shared_ptr<SimulationJob> &pJob)
//{}

// chrono::ChSystemSMC convertToChronosSystem(const
// std::shared_ptr<SimulationJob> &pJob)
//{
//    chrono::ChSystemSMC my_system;
//}

void ChronosEulerSimulationModel::parseForces(
    const std::shared_ptr<chrono::fea::ChMesh>& mesh_chronos,
    const std::vector<std::shared_ptr<chrono::fea::ChNodeFEAxyzrot>>&
        edgeMeshVertsToChronoNodes,
    const std::unordered_map<VertexIndex, Vector6d>& nodalExternalForces) {
  mesh_chronos->SetAutomaticGravity(false);
  for (const std::pair<VertexIndex, Vector6d>& externalForce :
       nodalExternalForces) {
    const int& forceVi = externalForce.first;
    const Vector6d& force = externalForce.second;
    edgeMeshVertsToChronoNodes[forceVi]->SetForce(
        ChVector<>(force[0], force[1], force[2]));
    edgeMeshVertsToChronoNodes[forceVi]->SetTorque(
        ChVector<>(force[3], force[4], force[5]));
  }
}

void ChronosEulerSimulationModel::parseForcedDisplacements(
    const std::shared_ptr<const SimulationJob>& pJob,
    const std::vector<std::shared_ptr<chrono::fea::ChNodeFEAxyzrot>>&
        edgeMeshVertsToChronoNodes,
    chrono::ChSystemSMC& my_system) {
  assert(!edgeMeshVertsToChronoNodes.empty());

  ChState x;
  ChStateDelta v;
  double t;
  for (const std::pair<VertexIndex, Eigen::Vector3d>& forcedDisplacement :
       pJob->nodalForcedDisplacements) {
    assert(false);
    std::cerr << "Forced displacements dont work" << std::endl;
    //    std::terminate();
    const int& constrainedVi = forcedDisplacement.first;
    ChVector<double> displacementVector(
        pJob->nodalForcedDisplacements.at(constrainedVi));
    const std::shared_ptr<chrono::fea::ChNodeFEAxyzrot>& constrainedChronoNode =
        edgeMeshVertsToChronoNodes[constrainedVi];
    constrainedChronoNode->NodeIntStateGather(0, x, 0, v, t);
    std::cout << "state rows:" << x.rows() << std::endl;
    std::cout << "state cols:" << x.cols() << std::endl;
    //    x = x + displacementVector;
    constrainedChronoNode->NodeIntStateScatter(0, x, 0, v, t);
  }
}
void ChronosEulerSimulationModel::parseConstrainedVertices(
    const std::shared_ptr<const SimulationJob>& pJob,
    const std::vector<std::shared_ptr<chrono::fea::ChNodeFEAxyzrot>>&
        edgeMeshVertsToChronoNodes,
    chrono::ChSystemSMC& my_system) {
  assert(!edgeMeshVertsToChronoNodes.empty());
  for (const std::pair<VertexIndex, std::unordered_set<int>>&
           constrainedVertex : pJob->constrainedVertices) {
    const int& constrainedVi = constrainedVertex.first;
    const std::unordered_set<int>& constrainedDoF = constrainedVertex.second;
    // Create a truss
    auto truss = chrono_types::make_shared<ChBody>();
    truss->SetBodyFixed(true);
    my_system.Add(truss);
    const std::shared_ptr<chrono::fea::ChNodeFEAxyzrot>& constrainedChronoNode =
        edgeMeshVertsToChronoNodes[constrainedVi];
    auto constraint = chrono_types::make_shared<ChLinkMateGeneric>();
    constraint->SetConstrainedCoords(
        constrainedDoF.contains(0), constrainedDoF.contains(1),
        constrainedDoF.contains(2), constrainedDoF.contains(3),
        constrainedDoF.contains(4), constrainedDoF.contains(5));
    constraint->Initialize(constrainedChronoNode, truss, false,
                           constrainedChronoNode->Frame(),
                           constrainedChronoNode->Frame());
    my_system.Add(constraint);
  }
}

SimulationResults ChronosEulerSimulationModel::executeSimulation(
    const std::shared_ptr<SimulationJob>& pJob) {
  assert(pJob->pMesh->VN() != 0);
  //  const bool structureInitialized = mesh_chronos != nullptr;
  //  const bool wasInitializedWithDifferentStructure =
  //      structureInitialized &&
  //      (pJob->pMesh->EN() != mesh_chronos->GetNelements() ||
  //       pJob->pMesh->VN() != mesh_chronos->GetNnodes());
  //  if (!structureInitialized || wasInitializedWithDifferentStructure) {
  setStructure(pJob->pMesh);
  //  }
  chrono::ChSystemSMC my_system;
  //    chrono::irrlicht::ChIrrApp application(&my_system,
  //                                           L"Irrlicht FEM visualization",
  //                                           irr::core::dimension2d<irr::u32>(800,
  //                                           600),
  //                                           chrono::irrlicht::VerticalDir::Z,
  //                                           false,
  //                                           true);
  //    const std::string chronoDataFolderPath =
  //    "/home/iason/Coding/build/external "
  //                                             "dependencies/CHRONO-src/data/";
  //    application.AddTypicalLogo(chronoDataFolderPath +
  //    "logo_chronoengine_alpha.png");
  //    application.AddTypicalSky(chronoDataFolderPath + "skybox/");
  //    application.AddTypicalLights();
  //    application.AddTypicalCamera(irr::core::vector3df(0, (irr::f32) 0.6,
  //    -1));
  //    my_system.SetTimestepperType(chrono::ChTimestepper::Type::EULER_IMPLICIT_LINEARIZED);
  // parse forces
  parseForcedDisplacements(pJob, edgeMeshVertsToChronoNodes, my_system);
  parseForces(mesh_chronos, edgeMeshVertsToChronoNodes,
              pJob->nodalExternalForces);
  // parse constrained vertices
  parseConstrainedVertices(pJob, edgeMeshVertsToChronoNodes, my_system);
  //    std::dynamic_pointer_cast<std::shared_ptr<ChNodeFEAxyzrot>>(mesh_chronos->GetNode(1))
  //        ->SetFixed(true);
  // and load containers must be added to your system
  //    auto mvisualizemesh =
  //    chrono_types::make_shared<ChVisualizationFEAmesh>(*(mesh_chronos.get()));
  //    mvisualizemesh->SetFEMdataType(ChVisualizationFEAmesh::E_PLOT_NODE_DISP_NORM);
  //    mvisualizemesh->SetColorscaleMinMax(0.0, 5.50);
  //    mvisualizemesh->SetShrinkElements(false, 0.85);
  //    mvisualizemesh->SetSmoothFaces(false);
  //    mesh_chronos->AddAsset(mvisualizemesh);

  //    application.AssetBindAll();
  //    application.AssetUpdateAll();
  my_system.Add(mesh_chronos);

  auto solver = chrono_types::make_shared<ChSolverMINRES>();
  my_system.SetSolver(solver);
  solver->SetMaxIterations(1e5);
  solver->EnableWarmStart(
      true);  // IMPORTANT for convergence when using EULER_IMPLICIT_LINEARIZED
  solver->EnableDiagonalPreconditioner(true);
  //  solver->SetTolerance(1e-15);
  my_system.SetSolverForceTolerance(1e-15);

  SimulationResults simulationResults;
  //#ifdef POLYSCOPE_DEFINED
  //  solver->SetVerbose(true);
  //  //  edgeMeshVertsToChronoNodes[10]->Frame().Move({0, 0, 1e-1});
  //  simulationResults.converged = my_system.DoEntireKinematics(5e2);
  ////  edgeMeshVertsToChronoNodes[10]->SetForce({0, 0, 1e6});
  //#else
  solver->SetVerbose(false);
  if (settings.analysisType == Settings::AnalysisType::Linear) {
    simulationResults.converged = my_system.DoStaticLinear();
    //    simulationResults.converged = my_system.DoStaticRelaxing();
  } else {
    simulationResults.converged = my_system.DoStaticNonlinear();
    //    simulationResults.converged = my_system.DoStaticRelaxing();
  }
  //#endif
  if (!simulationResults.converged) {
    std::cerr << "Simulation failed" << std::endl;
    assert(false);
    return simulationResults;
  }

  //    my_system.SetTimestepperType(ChTimestepper::Type::EULER_IMPLICIT_LINEARIZED);
  //    application.SetTimestep(0.001);

  //    while (application.GetDevice()->run()) {
  //        application.BeginScene();
  //        application.DrawAll();
  //        application.EndScene();
  //    }
  simulationResults.pJob = pJob;
  simulationResults.displacements.resize(pJob->pMesh->VN());
  simulationResults.rotationalDisplacementQuaternion.resize(pJob->pMesh->VN());
  for (size_t vi = 0; vi < pJob->pMesh->VN(); vi++) {
    const auto node_chronos = edgeMeshVertsToChronoNodes[vi];
    const auto posDisplacement =
        node_chronos->Frame().GetPos() - node_chronos->GetX0().GetPos();
    //        std::cout << "Node " << vi << " coordinate x= " <<
    //        node_chronos->Frame().GetPos().x()
    //                  << "    y=" << node_chronos->Frame().GetPos().y()
    //                  << "    z=" << node_chronos->Frame().GetPos().z() <<
    //                  "\n";
    // Translations
    simulationResults.displacements[vi][0] = posDisplacement[0];
    simulationResults.displacements[vi][1] = posDisplacement[1];
    simulationResults.displacements[vi][2] = posDisplacement[2];
    // Rotations
    chrono::ChQuaternion<double> rotQuat = node_chronos->GetRot();
    simulationResults.rotationalDisplacementQuaternion[vi] =
        Eigen::Quaternion<double>(rotQuat.e0(), rotQuat.e1(), rotQuat.e2(),
                                  rotQuat.e3());
    const Eigen::Vector3d eulerAngles =
        simulationResults.rotationalDisplacementQuaternion[vi]
            .toRotationMatrix()
            .eulerAngles(0, 1, 2);
    //        std::cout << "Euler angles:" << eulerAngles << std::endl;
    simulationResults.displacements[vi][3] = eulerAngles[0];
    simulationResults.displacements[vi][4] = eulerAngles[1];
    simulationResults.displacements[vi][5] = eulerAngles[2];
  }

  simulationResults.simulationModelUsed = label;
  return simulationResults;

  //    VCGEdgeMesh deformedMesh;
  //    deformedMesh.copy(*(pJob->pMesh));
  //    for (size_t vi = 0; vi < pJob->pMesh->VN(); vi++) {
  //        const std::shared_ptr<ChNodeFEAxyzrot> node_chronos =
  //        edgeMeshVertsToChronosNodes[vi]; deformedMesh.vert[vi].P() =
  //        CoordType(node_chronos->GetPos()[0],
  //                                              node_chronos->GetPos()[1],
  //                                              node_chronos->GetPos()[2]);
  //    }

  //    deformedMesh.updateEigenEdgeAndVertices();
  //    deformedMesh.setLabel("deformed");
  //    deformedMesh.registerForDrawing();
  //    polyscope::show();
  //    return simulationResults;
}

void ChronosEulerSimulationModel::setStructure(
    const std::shared_ptr<SimulationEdgeMesh>& pMesh) {
  mesh_chronos = convertToChronosMesh_Euler(pMesh, edgeMeshVertsToChronoNodes);
}