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

#include "Integratoren2d.h"
#include "tools/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);

    // Heun integration
    Vector e_integrated = rod2D.getE() +
                          0.5 * (ortho(rod2D.getE()) + ortho(e_predictor)) *
                              sim.getNorm(sim.getSTD()) * rod2D.getDRot_Sqrt() +
                          0.5 * (torque_euler(rod2D, sim, torque) +
                                 torque_euler(pred_rod, sim, torque));
    ;
    // 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 = rod2D.getDRot() * 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, {}, {});
}