package awa.imu;

/*
 * The source code of this file or parts of it were copied from 
 * the autopilot onboard code package.
 *  
 * Autopilot is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * Autopilot is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with the source code; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

public class Kalman
{
    
    Euler angles = new Euler();

    /**
     *  Kalman weighting matrices
     */
    double[][] Q = { { 0.00015, 0.00000, 0.00000, 0.00000 }, 
                    { 0.00000, 0.00015, 0.00000, 0.00000 }, 
                    { 0.00000, 0.00000, 0.00015, 0.00000 },
                    { 0.00000, 0.00000, 0.00000, 0.00015 }, };

    double[][] R = { { 0.086, 0.000, 0.000 }, 
                     { 0.000, 0.086, 0.000 }, 
                     { 0.000, 0.000, 0.015 }, };

    /**
     * Covariance matrix
     */
    double[][] P = new double[4][4];

    /**
     * Position vector
     * Initial position:  Level, facing north
     *
     * Estimated position is derived from this.  Unless ahrs_init() is called
     * the filter is likely to spike.
     */
    double[] X = { 1.0, 0.0, 0.0, 0.0 };


    /**
     *  Vector helpers related to the AHRS computation
     */
    static double sq(double d)
    {
        return d * d;
    }


    /**
     * This will normalize a quaternion vector q
     * q/norm(q)
     * q(4,1)
     */
    void v_normq(double[] q_out, double[] q)
    {
        KVec.v_scale(q_out, q, 1.0 / Math.sqrt(sq(q[0]) + sq(q[1]) + sq(q[2]) + sq(q[3])), 4);
    }


    /**
     * This will convert from quaternions to euler angles
     * q(4,1) -> euler[phi;theta;psi] (rad)
     */
    private Euler quat2euler(double[] q)
    {
        Euler euler = new Euler();
        
        double q0 = q[0];
        double q1 = q[1];
        double q2 = q[2];
        double q3 = q[3];

        //double q02 = q0 * q0;
        double q12 = q1 * q1;
        double q22 = q2 * q2;
        double q32 = q3 * q3;

        /* phi */
        euler.phi = Math.atan2(2.0 * (q2 * q3 + q0 * q1), 1.0 - 2.0 * (q12 + q22));

        /* theta */
        euler.theta = -Math.asin(2.0 * (q1 * q3 - q0 * q2));

        /* psi */
        euler.psi = Math.atan2(2.0 * (q1 * q2 + q0 * q3), 1.0 - 2.0 * (q22 + q32));
        
        return euler;
    }


    /**
     * This will convert from euler angles to quaternion vector
     * phi, theta, psi -> q(4,1)
     * euler angles in radians
     */
    void euler2quat(Euler euler)
    {
        double sphi = Math.sin(0.5 * euler.phi);
        double stheta = Math.sin(0.5 * euler.theta);
        double spsi = Math.sin(0.5 * euler.psi);

        double cphi = Math.cos(0.5 * euler.phi);
        double ctheta = Math.cos(0.5 * euler.theta);
        double cpsi = Math.cos(0.5 * euler.psi);

        X[0] = cphi * ctheta * cpsi + sphi * stheta * spsi;
        X[1] = -cphi * stheta * spsi + sphi * ctheta * cpsi;
        X[2] = cphi * stheta * cpsi + sphi * ctheta * spsi;
        X[3] = cphi * ctheta * spsi - sphi * stheta * cpsi;
    }


    /**
     * This will convert from accelerometer and compass readings into
     * [phi, theta, psi] (in radians).  The accelerometer readings are
     * assumed to be in m/s^2, but the actual units don't matter.
     */
    public Euler accel2euler(double ax, double ay, double az, double headingDeg)
    {
        Euler euler = new Euler();
        
        double gAbs = Math.sqrt(ax*ax + ay*ay + az*az);
        
        // prevent division by zero
        if (gAbs == 0.0) gAbs = 1.0;
        
        euler.phi = Math.atan2(ay, az);    // bank angle
        euler.theta = Math.asin(ax / gAbs);
        euler.psi = Euler.toNormRad(headingDeg);
        
        return euler;
    }


    /**
     * This will construct a direction cosine matrix from 
     * quaternions in the standard rotation sequence
     * [phi][theta][psi]
     *
     * body = tBL(3,3)*NED
     * q(4,1)
     */
    private double[][] quat2dcm(double[] q)
    {
        double[][] dcm = new double[3][3];
        
        double q0 = q[0];
        double q1 = q[1];
        double q2 = q[2];
        double q3 = q[3];

        //double q02 = q0 * q0;
        double q12 = q1 * q1;
        double q22 = q2 * q2;
        double q32 = q3 * q3;

        dcm[0][0] = 1.0 - 2.0 * (q22 + q32);
        dcm[0][1] = 2.0 * (q1 * q2 + q0 * q3);
        dcm[0][2] = 2.0 * (q1 * q3 - q0 * q2);

        dcm[1][0] = 2.0 * (q1 * q2 - q0 * q3);
        dcm[1][1] = 1.0 - 2.0 * (q12 + q32);
        dcm[1][2] = 2.0 * (q2 * q3 + q0 * q1);

        dcm[2][0] = 2.0 * (q1 * q3 + q0 * q2);
        dcm[2][1] = 2.0 * (q2 * q3 - q0 * q1);
        dcm[2][2] = 1.0 - 2.0 * (q12 + q22);
        
        return dcm;
    }


    /**
     * This will construct the quaternion omega matrix
     * @return W(4,4)
     * @param q, q, r (rad/sec)
     */
    private double[][] rotate2omega(double p, double q, double r)
    {
        double[][] W = new double[4][4];
        
        W[0][0] = 0.0;
        W[0][1] = -p / 2.0;
        W[0][2] = -q / 2.0;
        W[0][3] = -r / 2.0;

        W[1][0] = p / 2.0;
        W[1][1] = 0.0;
        W[1][2] = r / 2.0;
        W[1][3] = -q / 2.0;

        W[2][0] = q / 2.0;
        W[2][1] = -r / 2.0;
        W[2][2] = 0.0;
        W[2][3] = p / 2.0;

        W[3][0] = r / 2.0;
        W[3][1] = q / 2.0;
        W[3][2] = -p / 2.0;
        W[3][3] = 0.0;
        
        return W;
    }

    /**
     *  F[3] is a temp vector of d( arcsin(X) ) to reduce the complexity
     * of the C matrix calculations.
     */
    private double[] compute_f(double[][] DCM)
    {
        double[] F = new double[3];
        
        F[0] = 2.0 / (sq(DCM[2][2]) + sq(DCM[1][2]));
        F[1] = -1.0 / Math.sqrt(1.0 - sq(DCM[0][2]));
        F[2] = 2.0 / (sq(DCM[0][0]) + sq(DCM[0][1]));
        
        return F;
    }


    /**
     *  C is the measurement matrix that has the measured state
     * and the estimated state.  
     *
     * C = [
     *  [ dphi/dq0  dhpi/dq1    ... ]
     *  [ dtheta/qd0    dtheta/dq1  ... ]
     *  [ dpsi/dq0  dpsi/dq1    ... ]
     *     ]
     */
    private double[][] compute_c(double[] X)
    {
        double[][] C = new double[4][4];
        
        /* Compute the new DCM matrix from the quaternion position */
        double[][] DCM = quat2dcm(X);

        /* Build our temporary matrix */
        double[] F = compute_f(DCM);

        /* Compute the estimated state */
        C[0][0] = F[0] * (X[1] * DCM[2][2]);
        C[0][1] = F[0] * (X[0] * DCM[2][2] + 2.0 * X[1] * DCM[1][2]);
        C[0][2] = F[0] * (X[3] * DCM[2][2] + 2.0 * X[2] * DCM[1][2]);
        C[0][3] = F[0] * (X[2] * DCM[2][2]);

        C[1][0] = F[1] * -2.0 * X[2];
        C[1][1] = F[1] * 2.0 * X[3];
        C[1][2] = F[1] * -2.0 * X[0];
        C[1][3] = F[1] * 2.0 * X[1];

        C[2][0] = F[2] * (X[3] * DCM[0][0]);
        C[2][1] = F[2] * (X[2] * DCM[0][0]);
        C[2][2] = F[2] * (X[1] * DCM[0][0] + 2.0 * X[2] * DCM[0][1]);
        C[2][3] = F[2] * (X[0] * DCM[0][0] + 2.0 * X[3] * DCM[0][1]);
        
        return C;
    }


    /**
     * Kalman filter update
     *
     * P is [4,4] and holds the covariance matrix
     * R is [3,3] and holds the measurement weights
     * C is [4,3] and holds the euler to quat tranform
     * X is [4,1] and holds the estimated state as a quaternion
     * Xe is [3,1] and holds the euler angles of the estimated state
     * Xm is [3,1] and holds the measured euler angles
     */
    void kalman_update(double[][] P, double[][] R, double[][] C, double[] X, Euler Xe, Euler Xm)
    {
        double[][] T1 = new double[4][4];
        double[][] T2 = new double[4][4];
        double[][] E = new double[4][4];
        double[][] K = new double[4][4];
        double[] err;
        double[] temp = new double[4];

        /* E[3,3] = C[3,4] P[4,4] C'[4,3] + R[3,3] */
        KVec.m_mult(T1, C, P, 3, 4, 4);
        KVec.m_transpose(T2, C, 3, 4);
        KVec.m_mult(E, T1, T2, 3, 4, 3);
        KVec.m_add(E, R, E, 3, 3);

        /* K[4,3] = P[4,4] C'[4,3] inv(E)[3,3] */
        KVec.m_mult(T1, P, T2, 4, 4, 3);
        //KVec.m33_inv(E, T2);  // Bug in original code: E is overwritten here
        KVec.m33_inv(T2, E);
        KVec.m_mult(K, T1, T2, 4, 3, 3);

        /* X[4] = X[4] + K[4,3] ( Xm[3] - Xe[3] ) */
        err = KVec.v_sub(Xm, Xe);
        KVec.mv_mult(temp, K, err, 4, 3);
        KVec.v_add(X, temp, X, 4);

        /* P[4,4] = P[4,4] - K[4,3] C[3,4] P[4,4] */
        KVec.m_mult(T1, K, C, 4, 3, 4);
        KVec.m_mult(T2, T1, P, 4, 4, 4);
        KVec.m_sub(P, T2, P, 4, 4);
    }


    /**
     *  ahrs_init() takes the initial values from the IMU and converts
     * them into the quaternion state vector.
     */
    void ahrs_init(double ax, double ay, double az, double heading)
    {
        Euler angles = accel2euler(ax, ay, az, heading);

        // sets X
        euler2quat(angles);
    }

    /**
     * Method for testing rotational update without kalman filtering
     * @param p
     * @param q
     * @param r
     */
    public Euler addRotation(double dt, double p, double q, double r) 
    {
        double[] X_dot = new double[4];
       
        /* Construct the quaternion omega matrix */
        double[][] A = rotate2omega(p, q, r);

        /* X_dot = AX */
        KVec.mv_mult(X_dot, A, X, 4, 4);

        /* X += X_dot dt */
        KVec.v_scale(X_dot, X_dot, dt, 4);

        KVec.v_add(X, X_dot, X, 4);
        v_normq(X, X);

        /*
         * Compute the euler angles of the estimated attitude
         * Xe = quat2euler( X )
         */
        return quat2euler(X);
    }

    /**
     *  ahrs_step() takes the values from the IMU and produces the
     * new estimated attitude.
     */
    void ahrs_step(double dt, double p, double q, double r, double ax, double ay, double az, double heading)
    {
        double[][] AT = new double[4][4];
        double[] X_dot = new double[4];
        double[][] temp = new double[4][4];

        /* Construct the quaternion omega matrix */
        double[][] A = rotate2omega(p, q, r);

        /* X_dot = AX */
        KVec.mv_mult(X_dot, A, X, 4, 4);

        /* X += X_dot dt */
        KVec.v_scale(X_dot, X_dot, dt, 4);

        KVec.v_add(X, X_dot, X, 4);
        v_normq(X, X);

        /* AT = transpose(A) */
        KVec.m_transpose(AT, A, 4, 4);

        /* P = A * P * AT + Q */
        KVec.m_mult(temp, A, P, 4, 4, 4);
        KVec.m_mult(P, temp, AT, 4, 4, 4);
        KVec.m_add(P, Q, P, 4, 4);

        double[][] C = compute_c(X);

        /* Compute the euler angles of the measured attitude */
        Euler Xm = accel2euler(ax, ay, az, heading);

        /*
         * Compute the euler angles of the estimated attitude
         * Xe = quat2euler( X )
         */
        Euler Xe = quat2euler(X);

        /* Kalman filter update */
        kalman_update(P, R, C, X, Xe, Xm);

        /* Estimated attitude extraction */
        angles = quat2euler(X);
        
        //System.out.println("Angles: " + angles.phi + " " + angles.theta + " " + angles.psi);
    }
}