#include "chronoseulersimulationmodel.hpp"
#include <chrono/fea/ChBeamSectionEuler.h>
#include <chrono/fea/ChBuilderBeam.h>
#include <chrono/fea/ChMesh.h>
#include <chrono/fea/ChNodeFEAxyzrot.h>
#include <chrono/physics/ChSystemSMC.h>
#include <chrono/solver/ChIterativeSolverLS.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];
        edgeMeshVertsToChronosNodes[vi] = std::make_shared<ChNodeFEAxyzrot>(
            ChFrame<>(ChVector<>(vertex.cP()[0], vertex.cP()[1], vertex.cP()[2]),
                      ChVector<>(0, 1, 0)));
        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 = 1e4;

        auto msection = chrono_types::make_shared<ChBeamSectionEulerAdvanced>();
        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
                       5,        // 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 my_system;
    my_system.SetTimestepperType(chrono::ChTimestepper::Type::EULER_IMPLICIT_LINEARIZED);
    std::vector<std::shared_ptr<ChNodeFEAxyzrot>> edgeMeshVertsToChronosNodes;
    auto mesh_chronos = convertToChronosMesh_Euler(pJob->pMesh, edgeMeshVertsToChronosNodes);
    //    mesh_chronos->SetAutomaticGravity(false);
    //parse constrained vertices
    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;
        const bool vertexIsFullyConstrained = constrainedDoF.size() == 6
                                              && constrainedDoF.contains(0)
                                              && constrainedDoF.contains(1)
                                              && constrainedDoF.contains(2)
                                              && constrainedDoF.contains(3)
                                              && constrainedDoF.contains(4)
                                              && constrainedDoF.contains(5);
        if (vertexIsFullyConstrained) {
            edgeMeshVertsToChronosNodes[constrainedVi]->SetFixed(true);
        } else {
            std::cerr << "Currently only rigid vertices are handled." << std::endl;
            SimulationResults simulationResults;
            simulationResults.converged = false;
            assert(false);
            return simulationResults;
        }
    }
    //    std::dynamic_pointer_cast<std::shared_ptr<ChNodeFEAxyzrot>>(mesh_chronos->GetNode(1))
    //        ->SetFixed(true);
    //parse external forces
    for (const std::pair<VertexIndex, Vector6d> &externalForce : pJob->nodalExternalForces) {
        const int &forceVi = externalForce.first;
        const Vector6d &force = externalForce.second;
        if (Eigen::Vector3d(force[3], force[4], force[5]).norm() > 1e-10) {
            std::cerr << "Moments are currently not supported." << std::endl;
            SimulationResults simulationResults;
            simulationResults.converged = false;
            assert(false);
            return simulationResults;
        } else {
            edgeMeshVertsToChronosNodes[forceVi]->SetForce(ChVector<>(force[0], force[1], force[2]));
        }
    }
    //    edgeMeshVertsToChronosNodes[0]->SetFixed(true);
    //    edgeMeshVertsToChronosNodes[1]->SetForce(ChVector<>(1e-10, 0, 0));
    //    edgeMeshVertsToChronosNodes[0]->SetFixed(true);
    //    builder.GetLastBeamNodes().front()->SetFixed(true);
    //    builder.GetLastBeamNodes().back()->SetForce(ChVector<>(0, 0, 1e1));

    //    edgeMeshVertsToChronosNodes[3]->SetFixed(true);
    //    edgeMeshVertsToChronosNodes[7]->SetFixed(true);
    //    edgeMeshVertsToChronosNodes[11]->SetFixed(true);
    //    edgeMeshVertsToChronosNodes[15]->SetForce(ChVector<>(0.0, 0, 1e-1));
    //    edgeMeshVertsToChronosNodes[19]->SetForce(ChVector<>(0.0, 0, 1e-1));
    //    edgeMeshVertsToChronosNodes[23]->SetForce(ChVector<>(0.0, 0, 1e-1));
    //    builder.GetLastBeamNodes().back()->SetForce(ChVector<>(0.0, 1e2, 0.0));

    my_system.Add(mesh_chronos);

    auto solver = chrono_types::make_shared<ChSolverMINRES>();
    solver->SetMaxIterations(5000);
    solver->SetTolerance(1e-10);
    solver->EnableDiagonalPreconditioner(true);
    solver->SetVerbose(true);
    my_system.SetSolver(solver);
    SimulationResults simulationResults;
    simulationResults.converged = my_system.DoStaticNonlinear();
    if (!simulationResults.converged) {
        std::cerr << "Simulation failed" << std::endl;
        assert(false);
        return simulationResults;
    }
    //    my_system.DoStaticLinear();
    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++) {
        std::shared_ptr<ChNodeFEAxyzrot> node_chronos = edgeMeshVertsToChronosNodes[vi];
        const auto posDisplacement = node_chronos->GetPos() - node_chronos->GetX0().GetPos();
        //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();
        simulationResults.displacements[vi][3] = eulerAngles[0];
        simulationResults.displacements[vi][4] = eulerAngles[1];
        simulationResults.displacements[vi][5] = eulerAngles[2];
    }
    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;
}