/*
    MIT License

    Copyright (c) 2021 Zhepei Wang (wangzhepei@live.com)

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.
*/

#ifndef FLATNESS_HPP
#define FLATNESS_HPP

#include <Eigen/Eigen>
#include <cmath>

namespace flatness
{
    class FlatnessMap
    {
    public:
        inline void reset(const double &vehicle_mass,
                          const double &gravitational_acceleration,
                          const double &horitonral_drag_coeff,
                          const double &vertical_drag_coeff,
                          const double &parasitic_drag_coeff,
                          const double &speed_smooth_factor)
        {
            mass = vehicle_mass;
            grav = gravitational_acceleration;
            dh = horitonral_drag_coeff;
            dh_over_m = dh / mass;
            dv = vertical_drag_coeff;
            cp = parasitic_drag_coeff;
            veps = speed_smooth_factor;
            return;
        }

        inline void optimizated_forward(const Eigen::Vector3d &vel,
                                        const Eigen::Vector3d &acc,
                                        const Eigen::Vector3d &jer,
                                        Eigen::Vector4d &quat)
        {
            double v0 = vel(0);
            double v1 = vel(1);
            double v2 = vel(2);
            double a0 = acc(0);
            double a1 = acc(1);
            double a2 = acc(2);
            double cp_term = sqrt(v0 * v0 + v1 * v1 + v2 * v2 + veps);
            double w_term = 1.0 + cp * cp_term;
            double w0 = w_term * v0;
            double w1 = w_term * v1;
            double w2 = w_term * v2;
            double zu0 = a0 + dh_over_m * w0;
            double zu1 = a1 + dh_over_m * w1;
            double zu2 = a2 + dh_over_m * w2 + grav;
            double zu_sqr0 = zu0 * zu0;
            double zu_sqr1 = zu1 * zu1;
            double zu_sqr2 = zu2 * zu2;
            double zu_sqr_norm = zu_sqr0 + zu_sqr1 + zu_sqr2;
            double zu_norm = sqrt(zu_sqr_norm);
            double z0 = zu0 / zu_norm;
            double z1 = zu1 / zu_norm;
            double z2 = zu2 / zu_norm;
            double tilt_den = sqrt(2.0 * (1.0 + z2)); 
            quat(0) = 0.5 * tilt_den;
            quat(1) = -z1 / tilt_den;
            quat(2) = z0 / tilt_den;
            quat(3) = 0;
            return;
        }

        inline void optimizated_forward(const Eigen::Vector3d &vel,
                                        const Eigen::Vector3d &acc,
                                        const Eigen::Vector3d &jer,
                                        Eigen::Vector4d &quat,
                                        Eigen::Vector3d &omg)
        {
            double v0 = vel(0);
            double v1 = vel(1);
            double v2 = vel(2);
            double a0 = acc(0);
            double a1 = acc(1);
            double a2 = acc(2);
            double cp_term = sqrt(v0 * v0 + v1 * v1 + v2 * v2 + veps);
            double w_term = 1.0 + cp * cp_term;
            double w0 = w_term * v0;
            double w1 = w_term * v1;
            double w2 = w_term * v2;
            double zu0 = a0 + dh_over_m * w0;
            double zu1 = a1 + dh_over_m * w1;
            double zu2 = a2 + dh_over_m * w2 + grav;
            double zu_sqr0 = zu0 * zu0;
            double zu_sqr1 = zu1 * zu1;
            double zu_sqr2 = zu2 * zu2;
            double zu_sqr_norm = zu_sqr0 + zu_sqr1 + zu_sqr2;
            double zu_norm = sqrt(zu_sqr_norm);
            double z0 = zu0 / zu_norm;
            double z1 = zu1 / zu_norm;
            double z2 = zu2 / zu_norm;
            double tilt_den = sqrt(2.0 * (1.0 + z2)); 
            quat(0) = 0.5 * tilt_den;
            quat(1) = -z1 / tilt_den;
            quat(2) = z0 / tilt_den;
            quat(3) = 0;
            double ng_den = zu_sqr_norm * zu_norm;
            double ng00 = (zu_sqr1 + zu_sqr2) / ng_den;
            double zu01 = zu0 * zu1;
            double zu12 = zu1 * zu2;
            double zu02 = zu0 * zu2;
            double ng01 = -zu01 / ng_den;
            double ng02 = -zu02 / ng_den;
            double ng11 = (zu_sqr0 + zu_sqr2) / ng_den;
            double ng12 = -zu12 / ng_den;
            double ng22 = (zu_sqr0 + zu_sqr1) / ng_den;
            double v_dot_a = v0 * a0 + v1 * a1 + v2 * a2;
            double dw_term = cp * v_dot_a / cp_term;
            double dw0 = w_term * a0 + dw_term * v0;
            double dw1 = w_term * a1 + dw_term * v1;
            double dw2 = w_term * a2 + dw_term * v2;
            double dz_term0 = jer(0) + dh_over_m * dw0;
            double dz_term1 = jer(1) + dh_over_m * dw1;
            double dz_term2 = jer(2) + dh_over_m * dw2;
            double dz0 = ng00 * dz_term0 + ng01 * dz_term1 + ng02 * dz_term2;
            double dz1 = ng01 * dz_term0 + ng11 * dz_term1 + ng12 * dz_term2;
            double dz2 = ng02 * dz_term0 + ng12 * dz_term1 + ng22 * dz_term2;
            double omg_den = z2 + 1.0;     
            double omg_term = dz2 / omg_den; 
            omg(0) = -dz1 + z1 * omg_term;
            omg(1) = dz0 - z0 * omg_term;
            omg(2) = (z1 * dz0 - z0 * dz1) / omg_den; 
            return;
        }

        inline void forward(const Eigen::Vector3d &vel,
                            const Eigen::Vector3d &acc,
                            const Eigen::Vector3d &jer,
                            const double &psi, 
                            const double &dpsi, 
                            double &thr,
                            Eigen::Vector4d &quat,
                            Eigen::Vector3d &omg)
        {
            double v0 = vel(0);
            double v1 = vel(1);
            double v2 = vel(2);
            double a0 = acc(0);
            double a1 = acc(1);
            double a2 = acc(2);
            double cp_term = sqrt(v0 * v0 + v1 * v1 + v2 * v2 + veps);
            double w_term = 1.0 + cp * cp_term;
            double w0 = w_term * v0;
            double w1 = w_term * v1;
            double w2 = w_term * v2;
            double zu0 = a0 + dh_over_m * w0;
            double zu1 = a1 + dh_over_m * w1;
            double zu2 = a2 + dh_over_m * w2 + grav;
            double zu_sqr0 = zu0 * zu0;
            double zu_sqr1 = zu1 * zu1;
            double zu_sqr2 = zu2 * zu2;
            double zu01 = zu0 * zu1;
            double zu12 = zu1 * zu2;
            double zu02 = zu0 * zu2;
            double zu_sqr_norm = zu_sqr0 + zu_sqr1 + zu_sqr2;
            double zu_norm = sqrt(zu_sqr_norm);
            double z0 = zu0 / zu_norm;
            double z1 = zu1 / zu_norm;
            double z2 = zu2 / zu_norm;
            double ng_den = zu_sqr_norm * zu_norm;
            double ng00 = (zu_sqr1 + zu_sqr2) / ng_den;
            double ng01 = -zu01 / ng_den;
            double ng02 = -zu02 / ng_den;
            double ng11 = (zu_sqr0 + zu_sqr2) / ng_den;
            double ng12 = -zu12 / ng_den;
            double ng22 = (zu_sqr0 + zu_sqr1) / ng_den;
            double v_dot_a = v0 * a0 + v1 * a1 + v2 * a2;
            double dw_term = cp * v_dot_a / cp_term;
            double dw0 = w_term * a0 + dw_term * v0;
            double dw1 = w_term * a1 + dw_term * v1;
            double dw2 = w_term * a2 + dw_term * v2;
            double dz_term0 = jer(0) + dh_over_m * dw0;
            double dz_term1 = jer(1) + dh_over_m * dw1;
            double dz_term2 = jer(2) + dh_over_m * dw2;
            double dz0 = ng00 * dz_term0 + ng01 * dz_term1 + ng02 * dz_term2;
            double dz1 = ng01 * dz_term0 + ng11 * dz_term1 + ng12 * dz_term2;
            double dz2 = ng02 * dz_term0 + ng12 * dz_term1 + ng22 * dz_term2;
           
            double f_term0 = mass * a0 + dv * w0; 
            double f_term1 = mass * a1 + dv * w1;
            double f_term2 = mass * (a2 + grav) + dv * w2;
            thr = z0 * f_term0 + z1 * f_term1 + z2 * f_term2;
            double tilt_den = sqrt(2.0 * (1.0 + z2));
            double tilt0 = 0.5 * tilt_den;
            double tilt1 = -z1 / tilt_den;
            double tilt2 = z0 / tilt_den;
            double c_half_psi = cos(0.5 * psi); 
            double s_half_psi = sin(0.5 * psi); 
            quat(0) = tilt0 * c_half_psi;
            quat(1) = tilt1 * c_half_psi + tilt2 * s_half_psi;
            quat(2) = tilt2 * c_half_psi - tilt1 * s_half_psi;
            quat(3) = tilt0 * s_half_psi;
            double c_psi = cos(psi);             
            double s_psi = sin(psi);             
            double omg_den = z2 + 1.0;           
            double omg_term = dz2 / omg_den;     
            omg(0) = dz0 * s_psi - dz1 * c_psi - 
                     (z0 * s_psi - z1 * c_psi) * omg_term;
            omg(1) = dz0 * c_psi + dz1 * s_psi - 
                     (z0 * c_psi + z1 * s_psi) * omg_term;
            omg(2) = (z1 * dz0 - z0 * dz1) / omg_den + dpsi; 

            return;
        }
      
        inline void backwardthreadsafe(
            const Eigen::Vector3d &vel,
            const Eigen::Vector3d &acc,
            const Eigen::Vector3d &jer,
            const Eigen::Vector3d &pos_grad,
            const Eigen::Vector3d &vel_grad,
            const Eigen::Vector4d &quat_grad,
            const Eigen::Vector3d &omg_grad,
            Eigen::Vector3d &pos_total_grad,
            Eigen::Vector3d &vel_total_grad, 
            Eigen::Vector3d &acc_total_grad, 
            Eigen::Vector3d &jer_total_grad)
        {
            double v0 = vel(0);
            double v1 = vel(1);
            double v2 = vel(2);
            double a0 = acc(0);
            double a1 = acc(1);
            double a2 = acc(2);
            double cp_term = sqrt(v0 * v0 + v1 * v1 + v2 * v2 + veps);
            double w_term = 1.0 + cp * cp_term;
            double w0 = w_term * v0;
            double w1 = w_term * v1;
            double w2 = w_term * v2;
            double zu0 = a0 + dh_over_m * w0;
            double zu1 = a1 + dh_over_m * w1;
            double zu2 = a2 + dh_over_m * w2 + grav;
            double zu_sqr0 = zu0 * zu0;
            double zu_sqr1 = zu1 * zu1;
            double zu_sqr2 = zu2 * zu2;
            double zu01 = zu0 * zu1;
            double zu12 = zu1 * zu2;
            double zu02 = zu0 * zu2;
            double zu_sqr_norm = zu_sqr0 + zu_sqr1 + zu_sqr2;
            double zu_norm = sqrt(zu_sqr_norm);
            double z0 = zu0 / zu_norm;
            double z1 = zu1 / zu_norm;
            double z2 = zu2 / zu_norm;
            double ng_den = zu_sqr_norm * zu_norm;
            double ng00 = (zu_sqr1 + zu_sqr2) / ng_den;
            double ng01 = -zu01 / ng_den;
            double ng02 = -zu02 / ng_den;
            double ng11 = (zu_sqr0 + zu_sqr2) / ng_den;
            double ng12 = -zu12 / ng_den;
            double ng22 = (zu_sqr0 + zu_sqr1) / ng_den;
            double v_dot_a = v0 * a0 + v1 * a1 + v2 * a2;
            double dw_term = cp * v_dot_a / cp_term;
            double dw0 = w_term * a0 + dw_term * v0;
            double dw1 = w_term * a1 + dw_term * v1;
            double dw2 = w_term * a2 + dw_term * v2;
            double dz_term0 = jer(0) + dh_over_m * dw0;
            double dz_term1 = jer(1) + dh_over_m * dw1;
            double dz_term2 = jer(2) + dh_over_m * dw2;
            double dz0 = ng00 * dz_term0 + ng01 * dz_term1 + ng02 * dz_term2;
            double dz1 = ng01 * dz_term0 + ng11 * dz_term1 + ng12 * dz_term2;
            double dz2 = ng02 * dz_term0 + ng12 * dz_term1 + ng22 * dz_term2;
       
            double f_term0 = mass * a0 + dv * w0; 
            double f_term1 = mass * a1 + dv * w1;
            double f_term2 = mass * (a2 + grav) + dv * w2;
            double tilt_den = sqrt(2.0 * (1.0 + z2)); 
            double tilt0 = 0.5 * tilt_den;
            double tilt1 = -z1 / tilt_den;
            double tilt2 = z0 / tilt_den;
            double omg_den = z2 + 1.0;       
            double omg_term = dz2 / omg_den; 

            double w0b, w1b, w2b, dw0b, dw1b, dw2b;
            double z0b, z1b, z2b, dz0b, dz1b, dz2b;
            double v_sqr_normb, cp_termb, w_termb;
            double zu_sqr_normb, zu_normb, zu0b, zu1b, zu2b;
            double zu_sqr0b, zu_sqr1b, zu_sqr2b, zu01b, zu12b, zu02b;
            double ng00b, ng01b, ng02b, ng11b, ng12b, ng22b, ng_denb;
            double dz_term0b, dz_term1b, dz_term2b;
            double tilt_denb, tilt0b, tilt1b, tilt2b, head0b, head3b;
            double cpsib, spsib, omg_denb, omg_termb;
            double tempb, tilt_den_sqr;

            tilt0b = quat_grad(0); 
            head3b = tilt0 * (quat_grad(3)) + tilt2 * (quat_grad(1)) - tilt1 * (quat_grad(2));
            tilt2b = quat_grad(2); 
            head0b = tilt2 * (quat_grad(2)) + tilt1 * (quat_grad(1)) + tilt0 * (quat_grad(0));
            tilt1b = quat_grad(1);
            tilt_den_sqr = tilt_den * tilt_den;
            tilt_denb = (z1 * tilt1b - z0 * tilt2b) / tilt_den_sqr + 0.5 * tilt0b;
            omg_termb = -z0 * omg_grad(1) +
                        z1 * omg_grad(0); 
            tempb = omg_grad(2) / omg_den;
            z1b = dz0 * tempb;
        
            dz0b = z1 * tempb + omg_grad(1);
            z0b = -(dz1 * tempb);
            dz1b = -z0 * tempb - omg_grad(0); 
            omg_denb = -((z1 * dz0 - z0 * dz1) * tempb / omg_den) -
                       dz2 * omg_termb / (omg_den * omg_den);
            tempb = -(omg_term * (omg_grad(1)));
            cpsib = dz0 * (omg_grad(1)) + z0 * tempb;
            spsib = dz1 * (omg_grad(1)) + z1 * tempb;

            z0b += tempb;

            tempb = -(omg_term * (omg_grad(0)));
            spsib += dz0 * (omg_grad(0)) + z0 * tempb;
            cpsib += -dz1 * (omg_grad(0)) - z1 * tempb;
            z0b += tilt2b / tilt_den;
            z1b += -tempb - tilt1b / tilt_den;
            dz2b = omg_termb / omg_den;
            z2b = omg_denb + tilt_denb / tilt_den;
    

            ng02b = dz_term0 * dz2b + dz_term2 * dz0b;
            dz_term0b = ng02 * dz2b + ng01 * dz1b + ng00 * dz0b;
            ng12b = dz_term1 * dz2b + dz_term2 * dz1b;
            dz_term1b = ng12 * dz2b + ng11 * dz1b + ng01 * dz0b;
            ng22b = dz_term2 * dz2b;
            dz_term2b = ng22 * dz2b + ng12 * dz1b + ng02 * dz0b;
            ng01b = dz_term0 * dz1b + dz_term1 * dz0b;
            ng11b = dz_term1 * dz1b;
            ng00b = dz_term0 * dz0b;
            jer_total_grad(2) = dz_term2b;
            dw2b = dh_over_m * dz_term2b;
            jer_total_grad(1) = dz_term1b;
            dw1b = dh_over_m * dz_term1b;
            jer_total_grad(0) = dz_term0b;
            dw0b = dh_over_m * dz_term0b;
            tempb = cp * (v2 * dw2b + v1 * dw1b + v0 * dw0b) / cp_term;
            acc_total_grad(2) = w_term * dw2b + v2 * tempb;
            acc_total_grad(1) = w_term * dw1b + v1 * tempb;
            acc_total_grad(0) = w_term * dw0b + v0 * tempb;
            vel_total_grad(2) = dw_term * dw2b + a2 * tempb;
            vel_total_grad(1) = dw_term * dw1b + a1 * tempb;
            vel_total_grad(0) = dw_term * dw0b + a0 * tempb;
            cp_termb = -(v_dot_a * tempb / cp_term);
            tempb = ng22b / ng_den;
            zu_sqr0b = tempb;
            zu_sqr1b = tempb;
            ng_denb = -((zu_sqr0 + zu_sqr1) * tempb / ng_den);
            zu12b = -(ng12b / ng_den);
            tempb = ng11b / ng_den;
            ng_denb += zu12 * ng12b / (ng_den * ng_den) -
                       (zu_sqr0 + zu_sqr2) * tempb / ng_den;
            zu_sqr0b += tempb;
            zu_sqr2b = tempb;
            zu02b = -(ng02b / ng_den);
            zu01b = -(ng01b / ng_den);
            tempb = ng00b / ng_den;
            ng_denb += zu02 * ng02b / (ng_den * ng_den) +
                       zu01 * ng01b / (ng_den * ng_den) -
                       (zu_sqr1 + zu_sqr2) * tempb / ng_den;
            zu_normb = zu_sqr_norm * ng_denb -
                       (zu2 * z2b + zu1 * z1b + zu0 * z0b) / zu_sqr_norm;
            zu_sqr_normb = zu_norm * ng_denb + zu_normb / (2.0 * zu_norm);
            tempb += zu_sqr_normb;
            zu_sqr1b += tempb;
            zu_sqr2b += tempb;
            zu2b = z2b / zu_norm + zu0 * zu02b + zu1 * zu12b + 2 * zu2 * zu_sqr2b;
            w2b = dh_over_m * zu2b;
            zu1b = z1b / zu_norm + zu2 * zu12b + zu0 * zu01b + 2 * zu1 * zu_sqr1b;
            w1b = dh_over_m * zu1b;
            zu_sqr0b += zu_sqr_normb;
            zu0b = z0b / zu_norm + zu2 * zu02b + zu1 * zu01b + 2 * zu0 * zu_sqr0b;
            w0b = dh_over_m * zu0b;
            w_termb = a2 * dw2b + a1 * dw1b + a0 * dw0b +
                      v2 * w2b + v1 * w1b + v0 * w0b;
            acc_total_grad(2) += zu2b;
            acc_total_grad(1) += zu1b;
            acc_total_grad(0) += zu0b;
            cp_termb += cp * w_termb;
            v_sqr_normb = cp_termb / (2.0 * cp_term);
            vel_total_grad(2) += w_term * w2b + 2 * v2 * v_sqr_normb + vel_grad(2);
            vel_total_grad(1) += w_term * w1b + 2 * v1 * v_sqr_normb + vel_grad(1);
            vel_total_grad(0) += w_term * w0b + 2 * v0 * v_sqr_normb + vel_grad(0);
            pos_total_grad(2) = pos_grad(2);
            pos_total_grad(1) = pos_grad(1);
            pos_total_grad(0) = pos_grad(0);
            return;
        }
  
    private:
        double mass, grav, dh, dv, cp, veps; // thread safe
        double dh_over_m;
    };
}

#endif