#include "chronoseulersimulationmodel.hpp"
#include "chrono/physics/ChLoadContainer.h"
#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/ChSystemSMC.h>
#include <chrono/solver/ChIterativeSolverLS.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<SimulationMesh> &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 SimulationMesh::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.b;
        const double beam_wy = element.dimensions.h;
        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->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::parseConstrainedVertices(
    const std::shared_ptr<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 auto &constrainedChronoNode = edgeMeshVertsToChronoNodes[constrainedVi];
        // Create a constraint at the end of the beam
        auto constr_a = chrono_types::make_shared<ChLinkMateGeneric>();
        constr_a->SetConstrainedCoords(constrainedDoF.contains(0),
                                       constrainedDoF.contains(1),
                                       constrainedDoF.contains(2),
                                       constrainedDoF.contains(3),
                                       constrainedDoF.contains(4),
                                       constrainedDoF.contains(5));
        constr_a->Initialize(constrainedChronoNode,
                             truss,
                             false,
                             constrainedChronoNode->Frame(),
                             constrainedChronoNode->Frame());
        //        const auto frameNode = constrainedChronoNode->Frame();
        my_system.Add(constr_a);

        //        edgeMeshVertsToChronoNodes[constrainedVi]->SetFixed(true);
        //        if (vertexIsFullyConstrained) {
        //        } else {
        //            std::cerr << "Currently only rigid vertices are handled." << std::endl;
        //            //            SimulationResults simulationResults;
        //            //            simulationResults.converged = false;
        //            //            assert(false);
        //            //            return simulationResults;
        //        }
    }
}

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
    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(100);
    //        solver->SetTolerance(1e-8);
    solver->EnableWarmStart(true); // IMPORTANT for convergence when using EULER_IMPLICIT_LINEARIZED
    solver->EnableDiagonalPreconditioner(true);
    my_system.SetSolverForceTolerance(1e-6);
    solver->SetVerbose(false);

    SimulationResults simulationResults;
    simulationResults.converged = my_system.DoStaticNonlinear();
    //    simulationResults.converged = my_system.DoStaticLinear();
    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 ChVector<double> eulerAngles = rotQuat.Q_to_Euler123();
        //        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.setLabelPrefix("chronos");
    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<SimulationMesh> &pMesh)
{
    mesh_chronos = convertToChronosMesh_Euler(pMesh, edgeMeshVertsToChronoNodes);
}