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

#include "Integratoren2d_force.h"

constexpr double kBT = 1.0;
using Vector = Eigen::Vector2d;


Vector Integratoren2d_force::ortho(Vector e) {
    return Vector{-e[1], e[0]};
}

Eigen::Matrix2d Integratoren2d_force::MatTraf(Vector e) {
    Eigen::Matrix2d mat;
    mat << e[0], -e[1],
            e[1], e[0];
    return mat;
}

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::Set5_MBD(Rod2d &/*rod2D*/, Simulation &/*sim*/) {
}

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)
    //translation
    Vector rand = {sim.getNorm(std), sim.getNorm(std)}; //gaussian noise
    Vector trans_part = rod2D.getDiff_Sqrt() * rand; //translation vector in Particle System
    Vector trans_lab = rod2D.getE_Base_matrix() * trans_part;

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


    //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() / kBT * 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 = rod2D.getE_Base_matrix() * trans_pred_part;

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

    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.
    Vector e = rod2D.getE();
    double M_predict_rot = torque(rod2D.getPos(), rod2D.getE()) * rod2D.getDRot() / kBT * 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_predict = rod2D.getE() + delta_e_predictor;
    e_predict.normalize();


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

    Vector delta_pos_future = Heun_predictor_pos(pred_rod, sim, force);
    Vector delta_e_future = Heun_predictor_rot(pred_rod, sim, torque);

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

void Integratoren2d_force::Set3_Exact(Rod2d &rod2D, Simulation &sim,
                                      std::function<Eigen::Vector2d(Eigen::Vector2d, Eigen::Vector2d)> force,
                                      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 = rod2D.getE_Base_matrix().inverse() * F_lab;

    Vector F_trans = rod2D.getDiff_Sqrt() * F_part;
    F_trans *= sim.getMDeltaT() / kBT;
    Vector trans_lab = rod2D.getE_Base_matrix() * 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() / kBT * sim.getMDeltaT();
    auto correction = -0.5 * pow(rod2D.getDRot(), 2) / pow(kBT, 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,
                                     std::function<Eigen::Vector2d(Eigen::Vector2d, Eigen::Vector2d)> force,
                                     std::function<double(Eigen::Vector2d, Eigen::Vector2d)> /*torque*/) {
    auto std = sim.getSTD();
    //translation
    auto rand = Vector(sim.getNorm(std), sim.getNorm(std));
    auto trans_part = rod2D.getDiff_Sqrt() * rand; //translation vector in Particle System
    Eigen::Rotation2D rotMat(acos(unitVec.dot(rod2D.getE())));
    auto trans_lab = rotMat * trans_part;

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

    Vector F_trans = rod2D.getDiff_Sqrt() * F_part;
    F_trans *= sim.getMDeltaT() / kBT;


    auto rot = Eigen::Rotation2D<double>(sim.getNorm(std) * rod2D.getDRot_Sqrt());
    Eigen::Rotation2Dd rotation(rot);
    auto e_new = (rotation.toRotationMatrix() * rod2D.getE()).normalized();

    // Normalisation should not be necessary if a proper angular representation is used.
    // But with vector e it is done in case of numerical errors

    rod2D.setPos(rod2D.getPos() + trans_lab);
    rod2D.setE(e_new);
}