// Extended Kalman Filter class for SCAAT navigation

import Jama.Matrix; // http://math.nist.gov/javanumerics/jama/

class EKF
{
  Matrix x;  // n x 1 state vector
  Matrix A;  // n x n state transition matrix
  Matrix P;  // n x n error covariance matrix
  Matrix Q;  // n x n process noise covariance matrix
  float eta; // process noise variance
  float R;   // measurement noise variance
  double e;  // innovation residual
  float t;   // time
  
  public EKF(float var, float eta, float R)
  {
    println("new EKF(" + var + ", " + eta + ", " + R + ")");
    // init state: x = [ 0, 0, 5, 0, 0, 0]'
    x = new Matrix(6, 1);
    x.set(2, 0, 5);
    // init error covariance matrix P = [ 10, 10, 1, 0.1, 0.1, 0.1 ]
    P = new Matrix(6, 6);
    P.set(0, 0, 10);
    P.set(1, 1, 10);
    P.set(2, 2, 1);
    P.set(3, 3, 0.1);
    P.set(4, 4, 0.1);
    P.set(5, 5, 0.1);
    P = P.times(var);
    // init state transition matrix 
    A = Matrix.identity(6, 6);
    // init process noise covariance matrix
    Q = new Matrix(6, 6);
    // init process noise variance
    this.eta = eta;
    // init measurement noise variance
    this.R = R;
    // init time
    t = 0;
  }

  public void predict(float dt)
  {
    // predict the [n x 1] state
    A.set(0, 3, dt);
    A.set(1, 4, dt);
    A.set(2, 5, dt);
    x = A.times(x);
    // prevent negative z
    double z = x.get(2, 0);
    if(z < 0)
    {
      x.set(2, 0, -z);
      x.set(5, 0, -x.get(5, 0));
    }
    
    // predict the [n x n] covariance matrix
    double f1 = eta * dt;
    double f2 = f1 * dt;
    double f3 = f2 * dt;
    f2 /= 2;
    f3 /= 3;
    Q.set(0, 0, f3);
    Q.set(1, 1, f3);
    Q.set(2, 2, f3);
    Q.set(3, 3, f1);
    Q.set(4, 4, f1);
    Q.set(5, 5, f1);
    Q.set(0, 3, f2);
    Q.set(1, 4, f2);
    Q.set(2, 5, f2);
    Q.set(3, 0, f2);
    Q.set(4, 1, f2);
    Q.set(5, 2, f2);
    P = A.times(P).times(A.transpose()).plus(Q);
    
    // update time
    t += dt;
  }
  
  public void innovate(double d, double bx, double by, double bz)
  {
    // predict the measurement
    double x1 = x.get(0, 0) - bx;
    double x2 = x.get(1, 0) - by;
    double x3 = x.get(2, 0) - bz;
    double z = Math.sqrt(x1*x1 + x2*x2 + x3*x3);
    
    // compute the [1 x n] jacobian
    Matrix H = new Matrix(1, 6);
    H.set(0, 0, x1 / z);
    H.set(0, 1, x2 / z);
    H.set(0, 2, x3 / z);
    
    // compute the (scalar) innovation residual
    e = d - z;
    
    // compute the (scalar) innovation variance
    double Re = H.times(P).times(H.transpose()).get(0, 0) + R;
    
    // compute the [n x 1] kalman gain
    Matrix K = P.times(H.transpose()).times(1.0 / Re);
    
    // correct the predicted state and covariance matrix
    x = x.plus(K.times(e));
    P = Matrix.identity(6, 6).minus(K.times(H)).times(P);
  }

  public float getResidual()
  {
    return (float)e;
  }
  
  public float getX()
  {
    return (float)x.get(0, 0);
  }
  
  public float getY()
  {
    return (float)x.get(1, 0);
  }
  
  public float getZ()
  {
    return (float)x.get(2, 0);
  }
  
  public float getX(float t)
  {
    return getX() + getdX() * (t - this.t);
  }

  public float getY(float t)
  {
    return getY() + getdY() * (t - this.t);
  }

  public float getZ(float t)
  {
    return getZ() + getdZ() * (t - this.t);
  }

  public float getdX()
  {
    return (float)x.get(3, 0);
  }
  
  public float getdY()
  {
    return (float)x.get(4, 0);
  }
  
  public float getdZ()
  {
    return (float)x.get(5, 0);
  }
  
  public PVector getPosition()
  {
    return new PVector(getX(), getY(), getZ());
  }
  
  public float getVelocity()
  {
    return dist(getdX(), getdY(), getdZ(), 0, 0, 0);
  }
  
  public float getDistance(PVector p)
  {
    return dist(getX(), getY(), getZ(), p.x, p.y, p.z);
  }

  public void trace()
  {
    print("x:");
    x.print(6, 2);

    print("P:");
    P.print(6, 2);
  }
};