//
// Created by jholder on 21.10.21.
//

#include "Integratoren2d_force.h"
#include "vec_trafos.h"

using Vector = Eigen::Vector2d;
using Matrix = Eigen::Matrix2d;

void Integratoren2d_force::Set0_deterministic(Rod2d &rod2D,
                                              Simulation & /*sim*/) {
    auto trans_lab = Vector({0, 0.01});
    rod2D.setPos(rod2D.getPos() + trans_lab);
    Eigen::Rotation2D<double> rot(0.1);
    rod2D.setE(rot * rod2D.getE());
}

const Vector Integratoren2d_force::unitVec = {1, 0};

void Integratoren2d_force::Set1_Euler(
        Rod2d &rod2D, Simulation &sim,
        const std::function<Vector(Vector, Vector)> &force,
        const std::function<double(Vector, Vector)> &torque) {
    auto std = sim.getSTD();  // sqrt(2*delta_T)
    auto e = rod2D.getE();
    // translation
    Vector rand = {sim.getNorm(std), sim.getNorm(std)};  // gaussian noise
    Vector trans_part = rod2D.getDiff_Sqrt() * rand;
    Vector trans_lab = rotation_Matrix(e) * trans_part;

    // Force
    Vector F_lab = force(rod2D.getPos(), e);
    Vector F_part = rotation_Matrix(e).inverse() * F_lab;
    Vector F_trans = rod2D.getDiff_Sqrt() * F_part;
    F_trans *= sim.getMDeltaT();

    // rotation
    double rot =
        sim.getNorm(std) * rod2D.getDRot_Sqrt();  // rotationsmatrix verwenden.
    double M_rot = torque(rod2D.getPos(), e) * rod2D.getDRot()  * sim.getMDeltaT();
    Vector new_e = e + (rot + M_rot) * ortho(e);
    new_e.normalize();
    // apply
    rod2D.setPos(rod2D.getPos() + trans_lab + F_trans);
    rod2D.setE(new_e);
}

Vector Integratoren2d_force::Heun_predictor_pos(
    const Rod2d &rod2D, Simulation &sim,
    const std::function<Vector(Vector, Vector)> &force) {
    auto standard_dev = sim.getSTD();
    Vector rand_pred = {sim.getNorm(standard_dev),
                        sim.getNorm(standard_dev)};  // gaussian noise
    Vector trans_pred_part =
        rod2D.getDiff_Sqrt() *
        rand_pred;  // translation vector in Particle System
    Vector trans_pred_lab = rotation_Matrix(rod2D.getE()) * trans_pred_part;

    Vector F_pred_lab = force(rod2D.getPos(), rod2D.getE());
    Vector F_pred_part = rotation_Matrix(rod2D.getE()).inverse() * F_pred_lab;
    Vector F_pred_trans = rod2D.getDiff_Sqrt() * F_pred_part;
    F_pred_trans *= sim.getMDeltaT() ;

    return F_pred_trans + trans_pred_lab;
}

Vector Integratoren2d_force::Heun_predictor_rot(
    const Rod2d &rod2D, Simulation &sim,
    const std::function<double(Vector, Vector)> &torque) {
    auto std = sim.getSTD();
    double rot_predict =
        sim.getNorm(std) * rod2D.getDRot_Sqrt();  // rotationsmatrix verwenden.
    const Vector& e = rod2D.getE();
    double M_predict_rot = torque(rod2D.getPos(), rod2D.getE()) *
                           rod2D.getDRot()  * sim.getMDeltaT();
    Vector e_change_predict = (rot_predict + M_predict_rot) * ortho(e);

    return e_change_predict;
}

void Integratoren2d_force::Set2_Heun(
    Rod2d &rod2D, Simulation &sim,
    const std::function<Vector(Vector, Vector)> &force,
    const std::function<double(Vector, Vector)> &torque) {
    Vector delta_pos_predictor = Heun_predictor_pos(rod2D, sim, force);
    Vector pos_predictor = rod2D.getPos() + delta_pos_predictor;

    Vector delta_e_predictor = Heun_predictor_rot(rod2D, sim, torque);
    Vector e_predictor = rod2D.getE() + delta_e_predictor;
    e_predictor.normalize();

    Rod2d pred_rod = rod2D;
    pred_rod.setPos(pos_predictor);
    pred_rod.setE(e_predictor);

    Vector delta_e_future = Heun_predictor_rot(pred_rod, sim, torque);


    Vector e_integrated =
        rod2D.getE() + 0.5 * (delta_e_predictor + delta_e_future);
    // apply
    rod2D.setPos(pos_predictor);
    rod2D.setE(e_integrated.normalized());
}

void Integratoren2d_force::Set3_Exact(
    Rod2d &rod2D, Simulation &sim,
    const std::function<Eigen::Vector2d(Eigen::Vector2d, Eigen::Vector2d)>& force,
    const std::function<double(Eigen::Vector2d, Eigen::Vector2d)>& torque) {
    auto std = sim.getSTD();  // sqrt(2*delta_T)
    // translation
    Vector rand = {sim.getNorm(std), sim.getNorm(std)};  // gaussian noise
    Vector trans_part =
        rod2D.getDiff_Sqrt() * rand;  // translation vector in Particle System

    Vector F_lab = force(rod2D.getPos(), rod2D.getE());
    Vector F_part = rotation_Matrix(rod2D.getE()).inverse() * F_lab;

    Vector F_trans = rod2D.getDiff_Sqrt() * F_part;
    F_trans *= sim.getMDeltaT();
    Vector trans_lab = rotation_Matrix(rod2D.getE()) * trans_part;

    // rotation
    double rot =
        sim.getNorm(std) * rod2D.getDRot_Sqrt();  // rotationsmatrix verwenden.
    Vector e = rod2D.getE();
    double M_rot = torque(rod2D.getPos(), rod2D.getE()) * rod2D.getDRot() * sim.getMDeltaT();
    auto correction =
        -0.5 * pow(rod2D.getDRot(), 2)  * sim.getMDeltaT();
    Vector new_e = e + (rot + M_rot) * ortho(e) + correction * e;
    new_e.normalize();
    // apply
    rod2D.setPos(rod2D.getPos() + trans_lab + F_trans);
    rod2D.setE(new_e);
}

void Integratoren2d_force::Set4_BDAS(
    Rod2d &rod2D, Simulation &sim,
    const std::function<Eigen::Vector2d(Eigen::Vector2d, Eigen::Vector2d)>& force,
    const std::function<double(Eigen::Vector2d, Eigen::Vector2d)>& torque) {
    auto std = sim.getSTD();  // sqrt(2*delta_T)
    Vector e = rod2D.getE();
    // translation
    Vector rand = {sim.getNorm(std), sim.getNorm(std)};  // gaussian noise
    Vector trans_part = rod2D.getDiff_Sqrt() * rand;
    Vector trans_lab = rotation_Matrix(rod2D.getE()) * trans_part;

    // Force
    Vector F_lab = force(rod2D.getPos(), rod2D.getE());
    Vector F_part = rotation_Matrix(rod2D.getE()).inverse() * F_lab;
    Vector F_trans = rod2D.getDiff_Sqrt() * F_part;
    F_trans *= sim.getMDeltaT();

    // rotation
    double rot =
            sim.getNorm(std) * rod2D.getDRot_Sqrt();  // rotationsmatrix verwenden.

    double M_rot = torque(rod2D.getPos(), rod2D.getE()) * rod2D.getDRot()  * sim.getMDeltaT();
    Vector new_e = Eigen::Rotation2D<double>(rot + M_rot) * e;
    new_e.normalize();
    // apply
    rod2D.setPos(rod2D.getPos() + trans_lab + F_trans);
    rod2D.setE(new_e);
}

void Integratoren2d_force::Set5_MBD(Rod2d &rod2D, Simulation &sim,
                                    const std::function<Eigen::Vector2d(Eigen::Vector2d, Eigen::Vector2d)>& force,
                                    const std::function<double(Eigen::Vector2d, Eigen::Vector2d)>&  torque) {

    Vector delta_pos_predictor = Heun_predictor_pos(rod2D, sim, force);
   [[maybe_unused]] Vector pos_predictor = rod2D.getPos() + delta_pos_predictor;

    Vector delta_e_predictor = Heun_predictor_rot(rod2D, sim, torque);
    Vector e_predictor = rod2D.getE() + delta_e_predictor;
    e_predictor.normalize();


    auto std = sim.getSTD();
    Eigen::Vector2d e = rod2D.getE();


    Eigen::Matrix2d Diff_combined =  0.5*(rod2D.getDiff() + rotation_Matrix(e).inverse() * rotation_Matrix(e_predictor)*rod2D.getDiff());
    Eigen::Matrix2d D_combined_sqrt =  Diff_combined.llt().matrixL();

    auto rand = Eigen::Vector2d(sim.getNorm(std), sim.getNorm(std));
    Eigen::Vector2d pos_integrated =rod2D.getPos() + rotation_Matrix(e) * D_combined_sqrt * rand;

    double rot = sim.getNorm(std) * rod2D.getDRot_Sqrt();
    Eigen::Vector2d e_integrated = e + 0.5 * (rot * ortho(e) + rot * ortho(e_predictor));
    //apply
    rod2D.setPos(pos_integrated);
    rod2D.setE(e_integrated.normalized());
}