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

#include "Integratoren2d.h"
#include "vec_trafos.h"

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

Vector
force_euler(const Rod2d &rod2D, const Simulation &sim, const std::function<Vector(Vector, Vector)> &force) {
    if (not force) {
        return {0.0, 0.0};
    }
    const auto &e = rod2D.getE();
    Vector F_lab = force(rod2D.getPos(), e);
    Vector F_part = rotation_Matrix(e).inverse() * F_lab;
    Vector F_trans = rotation_Matrix(e) * rod2D.getDiff_Sqrt() * F_part;
    F_trans *= sim.getMDeltaT();
    return F_trans;
}

Vector torque_euler(const Rod2d &rod2D, const Simulation &sim, const std::function<double(Vector, Vector)> &torque) {
    if (not torque) {
        return {0.0, 0.0};
    }
    const auto &e = rod2D.getE();
    return torque(rod2D.getPos(), e) * rod2D.getDRot() * sim.getMDeltaT() * ortho(rod2D.getE());
}

Vector
rand_trans_euler(const Rod2d &rod2D, Simulation &sim) {
    const auto std = sim.getSTD();
    const auto &e = rod2D.getE();
    Vector rand = {sim.getNorm(std), sim.getNorm(std)};  // gaussian noise
    Vector trans_part = rod2D.getDiff_Sqrt() * rand;
    return rotation_Matrix(e) * trans_part;
}

Vector rand_rot_euler(const Rod2d &rod2D, Simulation &sim) {
    const auto std = sim.getSTD();
    return sim.getNorm(std) * rod2D.getDRot_Sqrt() * ortho(rod2D.getE());
}

void Integratoren2d::Set1_Euler_f(
        Rod2d &rod2D, Simulation &sim,
        const std::function<Vector(Vector, Vector)> &force,
        const std::function<double(Vector, Vector)> &torque) {

    // Gaussian - translation
    auto trans = rand_trans_euler(rod2D, sim);

    // Force
    auto F_trans = force_euler(rod2D, sim, force);

    // Gaussian rotation
    auto rot = rand_rot_euler(rod2D, sim);

    // torque
    auto M_rot = torque_euler(rod2D, sim, torque);

    Vector integrated_e = rod2D.getE() + rot + M_rot;
    integrated_e.normalize();
    Vector integrated_pos = rod2D.getPos() + trans + F_trans;

    // apply
    rod2D.setPos(integrated_pos);
    rod2D.setE(integrated_e);
}

void Integratoren2d::Set1_Euler(
        Rod2d &rod2D, Simulation &sim) {
    Set1_Euler_f(rod2D, sim, {}, {});
}


void Integratoren2d::Set2_Heun_f(
        Rod2d &rod2D, Simulation &sim,
        const std::function<Vector(Vector, Vector)> &force,
        const std::function<double(Vector, Vector)> &torque) {
    //position
    Vector pos_predictor = rod2D.getPos() + rand_trans_euler(rod2D, sim) + force_euler(rod2D, sim, force);
    //rotation
    Vector delta_e_predictor = rand_rot_euler(rod2D, sim) + torque_euler(rod2D, sim, torque);
    Vector e_predictor = rod2D.getE() + delta_e_predictor;
    e_predictor.normalize();

    //make predicted Particle
    Rod2d pred_rod = rod2D;
    pred_rod.setPos(pos_predictor);
    pred_rod.setE(e_predictor);

    //integrate euler for future
    Vector delta_e_future = rand_rot_euler(rod2D, sim) + torque_euler(rod2D, sim, torque);

    //Heun integration
    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::Set2_Heun(
        Rod2d &rod2D, Simulation &sim) {
    Set2_Heun_f(rod2D, sim, {}, {});
}

void Integratoren2d::Set3_Exact_f(
        Rod2d &rod2D, Simulation &sim,
        const std::function<Eigen::Vector2d(Eigen::Vector2d, Eigen::Vector2d)> &force,
        const std::function<double(Eigen::Vector2d, Eigen::Vector2d)> &torque) {
    // Gaussian - translation
    auto trans = rand_trans_euler(rod2D, sim);

    // Force
    auto F_trans = force_euler(rod2D, sim, force);

    Vector integrated_pos = rod2D.getPos() + trans + F_trans;

    // Gaussian rotation
    auto rot = rand_rot_euler(rod2D, sim);

    // torque
    auto M_rot = torque_euler(rod2D, sim, torque);

    auto correction =
            -0.5 * pow(rod2D.getDRot(), 2) * sim.getMDeltaT() * rod2D.getE();

    Vector integrated_e = rod2D.getE() + rot + M_rot + correction;
    integrated_e.normalize();
    // apply
    rod2D.setPos(integrated_pos);
    rod2D.setE(integrated_e);
}

void Integratoren2d::Set3_Exact(Rod2d &rod2D, Simulation &sim) {
    Set3_Exact_f(rod2D, sim, {}, {});
}

void Integratoren2d::Set4_BDAS_f(
        Rod2d &rod2D, Simulation &sim,
        const std::function<Eigen::Vector2d(Eigen::Vector2d, Eigen::Vector2d)> &force,
        const std::function<double(Eigen::Vector2d, Eigen::Vector2d)> &torque) {
    // Gaussian - translation
    auto trans = rand_trans_euler(rod2D, sim);
    // Force
    auto F_trans = force_euler(rod2D, sim, force);
    Vector integrated_pos = rod2D.getPos() + trans + F_trans;

    // rotation
    double rot = sim.getNorm(sim.getSTD()) * rod2D.getDRot_Sqrt();

    double M_rot = torque ? torque(rod2D.getPos(), rod2D.getE()) * rod2D.getDRot() * sim.getMDeltaT() : 0.0;

    Vector integrated_e = Eigen::Rotation2D<double>(rot + M_rot) * rod2D.getE();
    integrated_e.normalize();
    // apply
    rod2D.setPos(integrated_pos);
    rod2D.setE(integrated_e);
}

void Integratoren2d::Set4_BDAS(Rod2d &rod2D, Simulation &sim) {
    Set4_BDAS_f(rod2D, sim, {}, {});
}


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

    //position
    Vector pos_predictor = rod2D.getPos() + rand_trans_euler(rod2D, sim) + force_euler(rod2D, sim, force);
    //rotation
    Vector delta_e_predictor = rand_rot_euler(rod2D, sim) + torque_euler(rod2D, sim, torque);
    Vector e_predictor = rod2D.getE() + delta_e_predictor;
    e_predictor.normalize();

    //make predicted Particle
    Rod2d pred_rod = rod2D;
    pred_rod.setPos(pos_predictor);
    pred_rod.setE(e_predictor);

    auto std = sim.getSTD();
    Vector e = rod2D.getE();

    auto rot_Mat = [](const Matrix& mat,const Matrix& rotMat) { return rotMat * mat * rotMat.transpose(); };
    Eigen::Matrix2d Diff_combined = 0.5 * (rot_Mat(rod2D.getDiff(), rotation_Matrix(e)) +
                                           rot_Mat(rod2D.getDiff(), rotation_Matrix(e_predictor)));
    Eigen::Matrix2d D_combined_sqrt = Diff_combined.llt().matrixL();
    auto rand = Eigen::Vector2d(sim.getNorm(std), sim.getNorm(std));
    double rot = sim.getNorm(std) * rod2D.getDRot_Sqrt();
    Eigen::Vector2d e_integrated = e + 0.5 * (rot * ortho(e) + rot * ortho(e_predictor));

    Eigen::Vector2d pos_integrated = rod2D.getPos() + D_combined_sqrt * rand + 0.5*(force_euler(rod2D,sim,force)+force_euler(pred_rod,sim,force));

    //apply
    rod2D.setPos(pos_integrated);
    rod2D.setE(e_integrated.normalized());
}

void Integratoren2d::Set5_MBD(Rod2d &rod2D, Simulation &sim) {
    Set5_MBD_f(rod2D, sim, {}, {});
}