/* LICENSE:
  =========================================================================
    CMPack'04 Source Code Release for OPEN-R SDK 1.1.5-r2 for ERS7
    Copyright (C) 2004 Multirobot Lab [Project Head: Manuela Veloso]
    School of Computer Science, Carnegie Mellon University
    All rights reserved.
  ========================================================================= */
#include "BallEKFPosVel.h"
#include "../headers/Utility.h"
#include "../headers/Reporting.h"

#ifdef PLATFORM_LINUX
#include <cstdio>
#endif

//========================================================================

static const double GRAVITY = 9800; // mm/s^2

// Ball coefficient of friction
static const double BALL_FRICTION = 0.025;

// This will be dependent on the distance that the ball is observed
//static const double BALL_POSITION_VARIANCE = 10000.0; // Stddev = 100mm
static const double BALL_POSITION_VARIANCE = 5625.0; // Stddev = 75mm
//static const double BALL_POSITION_VARIANCE = 2500.0; // Stddev = 50mm
//static const double BALL_POSITION_VARIANCE = 100.0; // Stddev = 10mm

// This should also be dependent on the distance that the ball is observed

//static const double BALL_VELOCITY_VARIANCE = 100.0; // Stddev = 10mm
//static const double BALL_VELOCITY_VARIANCE = 400.0; // Stddev = 20mm
static const double BALL_VELOCITY_VARIANCE = 2500.0; // Stddev = 50mm
//static const double BALL_VELOCITY_VARIANCE = 10000.0; // 100mm/s
//static const double BALL_VELOCITY_VARIANCE = 40000.0; // 200mm/s
static const double BALL_VELOCITY_VARIANCE_INIT = 250000.0; // 500mm/s

// noiseless dynamics
Matrix& BallEKFPosVel::f(const Matrix &x, Matrix &I) {

  static Matrix _f;

  _f = x; // Copy Matrix

  double &_x = _f.e(0,0);
  double &_y = _f.e(1,0);
  double &_vx = _f.e(2,0);
  double &_vy = _f.e(3,0);

  double _v = sqrt(_vx * _vx + _vy * _vy);
  double _a = min(BALL_FRICTION*GRAVITY,_v/stepsize);
  // Compute the components of the acceleration
  double _ax = (_v == 0.0) ? 0.0 : -_a * _vx / _v;
  double _ay = (_v == 0.0) ? 0.0 : -_a * _vy / _v;
  
  // Update Position
  _x += _vx * stepsize + 0.5 * _ax * stepsize * stepsize;
  _y += _vy * stepsize + 0.5 * _ay * stepsize * stepsize;

  // Update velocity by adding any velocity impulses
  _vx += (_ax * stepsize) + added_vel.x;
  _vy += (_ay * stepsize) + added_vel.y;

  added_vel=vector2d(0,0);

  return _f;
}

// noiseless observation
Matrix& BallEKFPosVel::h(const Matrix &x) {

  static Matrix _h(2,1);
  _h.e(0,0) = x.e(0,0);
  _h.e(1,0) = x.e(1,0);
  return _h;
}

// Covariance of propagation noise
Matrix& BallEKFPosVel::Q(const Matrix &x) {

  // This is going to depend heavily on whether the ball is near 
  // other robots or the walls...

  // This value will depend on some external observation of 
  // some kind...
  
  // This is initialized to the identity matrix
  static Matrix _Q(2);

  _Q.e(0,0)=_Q.e(1,1)=BALL_VELOCITY_VARIANCE;

  return _Q;
}

// Covariance of observation noise
Matrix& BallEKFPosVel::R(const Matrix &x) {

  // This will depend heavily on the distance from the robot to the
  // ball...

  // The observation method needs to change this value when it gets
  // called since we will have a way of obtaining the uncertainty.

  // This is initialized to the identity matrix
  static Matrix _R(2);

  _R.e(0,0)=_R.e(1,1)=BALL_POSITION_VARIANCE; // millimeters

  return _R;
}

// Jacobian of f w.r.t. x
Matrix& BallEKFPosVel::A(const Matrix &x) {

  static Matrix _A;

#if 0
  // This is taken from the smallsize and needs to be redone since it
  // doesn't take into account the ball dynamics
  if (!_A.nrows()) {
    _A.identity(4);
    _A.e(0,2) = stepsize;
    _A.e(1,3) = stepsize;
  }
#else
  // This is a first-order attempt to compute the Jacobian of the ball
  // given a constant coefficient of friction.
  //
  // This friction model is probably not correct, however, and should
  // instead be some sort of an exponential model

  Matrix _f = x;

  _A.identity(4);

  double &_vx = _f.e(2,0);
  double &_vy = _f.e(3,0);

  double vel = (_vx*_vx + _vy*_vy);

  // Watch out for singularities
  if (0.0==vel) {
    _A.e(0,2) = stepsize;
    _A.e(1,3) = stepsize;
    return _A;
  }

  // Compute some quantities and cache them
  // The constant 'c' is not computed correctly since it is now dependent
  // on the velocities. This needs to be rederived.
  double c = -1 * min(BALL_FRICTION*GRAVITY,vel/stepsize);
  double pos_c = (-0.5*stepsize*stepsize*c);
  double vel_c = (-stepsize*c);
  double inv_vel_1 = 1/sqrt(vel);
  double inv_vel_2 = 1/(vel);
  double inv_vel = inv_vel_1*inv_vel_2;

  _A.e(0,2) = stepsize + pos_c*inv_vel_1 + pos_c*_vx*(-_vx*inv_vel);
  _A.e(0,3) = pos_c*_vx*(-_vy*inv_vel);

  _A.e(1,2) = pos_c*_vy*(-_vx*inv_vel);
  _A.e(1,3) = stepsize + pos_c*inv_vel_1 + pos_c*_vy*(-_vy*inv_vel);

  _A.e(2,2) = 1 + vel_c*inv_vel_1 + vel_c*_vx*(-_vx*inv_vel);
  _A.e(2,3) = vel_c*_vx*(-_vy*inv_vel);

  _A.e(3,2) = vel_c*_vy*(-_vx*inv_vel);
  _A.e(3,3) = 1 + vel_c*inv_vel_1 + vel_c*_vy*(-_vy*inv_vel);
  
  // Whew...

#endif

  return _A;
}

// Jacobian of f w.r.t. noise
Matrix& BallEKFPosVel::W(const Matrix &x) {
  // This needs to be redone as well

  static Matrix _W("[0, 0; 0, 0; 1, 0; 0, 1]");
  return _W;
}

// Jacobian of h w.r.t. x
Matrix& BallEKFPosVel::H(const Matrix &x) {
  // This needs to be redone to take into account the nonlinearities in
  // how the AIBO calculates the distances

  static Matrix _H("[ 1, 0, 0, 0; "
                  "  0, 1, 0, 0 ] ");
  return _H;
}

// Jacobian of h w.r.t. noise
Matrix& BallEKFPosVel::V(const Matrix &x) {
  static Matrix _V("[ 1, 0; "
		   "  0, 1 ]");
  return _V;
}

//========================================================================
// Constructor
// The Kalman filter has a state vector of size 4, an observation of size
// 2 and the timestep comes in at a 30th of a second
BallEKFPosVel::BallEKFPosVel() : Kalman(4,2,1.0/30.0){
  reset(vector2d(0,0),0);
}

void BallEKFPosVel::reset(vector2d pos,ulong new_time) {
  _reset=true;
  added_vel=vector2d(0,0);

  Matrix x(4,1), P(4);

  x.e(0,0) = 0.0;
  x.e(1,0) = 0.0;
  x.e(2,0) = 0.0;
  x.e(3,0) = 0.0;
  
  P.e(0,0) = BALL_POSITION_VARIANCE;
  P.e(1,1) = BALL_POSITION_VARIANCE;
  P.e(2,2) = BALL_VELOCITY_VARIANCE_INIT;
  P.e(3,3) = BALL_VELOCITY_VARIANCE_INIT;
  
  initial(TimeToSec(new_time), x, P);
}

// Return the number of hypotheses
int BallEKFPosVel::numHypotheses() {
  return 1;
}

void BallEKFPosVel::propagate(ulong timestamp) {
  if (!_reset) {
    double _timestamp = TimeToSec(timestamp);
    if (_timestamp > time) {
      tick(_timestamp - time);
    }
  }
}

void BallEKFPosVel::perturb(vector2d pos,int idx,ulong timestamp) {
#ifdef PLATFORM_LINUX
  printf("BallEKFPosVel::perturb() NOT IMPLEMENTED YET!!!\n");
#endif
#ifdef PLATFORM_APERIOS
  pprintf(TextOutputStream,"BallEKFPosVel::perturb() NOT IMPLEMENTED YET!!!\n");
#endif
}

// Add an observation to the tracker
void BallEKFPosVel::observe(vector2d obs, 
			    double confidence,
			    ulong timestamp, 
			    bool imp_filter) {
  // Currently, only new observations are supported, as opposed to
  // adding old observations
  
  double _timestamp = TimeToSec(timestamp);

  if (_reset) {
    // Initialize the Kalman filter
    Matrix x(4,1), P(4);

    x.e(0,0) = obs.x;
    x.e(1,0) = obs.y;
    x.e(2,0) = 0.0;
    x.e(3,0) = 0.0;

    P.e(0,0) = BALL_POSITION_VARIANCE;
    P.e(1,1) = BALL_POSITION_VARIANCE;
    P.e(2,2) = BALL_VELOCITY_VARIANCE_INIT;
    P.e(3,3) = BALL_VELOCITY_VARIANCE_INIT;

    initial(_timestamp, x, P);

    _reset=false;

  } else {
    // We need to make a number of checks here to make sure that the
    // various confidence values are valid
    
    // If the timestamp is in the past, return... 
    if (_timestamp < time) { 
#ifdef PLATFORM_LINUX
      printf("Not a new observation!\n");
#endif
      return;
    }
    
    if (_timestamp > time) {
      // Propagate the model as necessary
      tick(_timestamp - time);
    }

    // Pass the update to the filter
    Matrix z(2,1);
    z.e(0,0) = obs.x;
    z.e(1,0) = obs.y;

    if (imp_filter) {
      // Filter out improbable values here
      // 
      // One notion is to analyze the mahalanobis distance and reject
      // values higher than 3 sigma
      
      Matrix x = xs.front();
      Matrix P = Ps.front();
      Matrix res=(z - h(x));
      Matrix &_H = H(x);
      Matrix &_V = V(x);
      Matrix &_R = R(x);
      Matrix S = _H*P*transpose(_H) + _V*_R*transpose(_V);
      S.inverse();
      Matrix mah = transpose(res)*S*res;
      
#ifdef PLATFORM_LINUX
      printf("Mahalanobis distance for new observation = %f\n",mah.e(0,0));
#endif
      //#ifdef PLATFORM_APERIOS_SDK
      //      pprintf(TextOutputStream,"Mahalanobis: %f\n",mah.e(0,0));
      //#endif
      if (3 < mah.e(0,0)) {
	// We reject the observation
#ifdef PLATFORM_APERIOS_SDK
	pprintf(TextOutputStream,"Mahalanobis reject %f\n",mah.e(0,0));
#endif
	return;
      }
    }
    update(z,3);
  }
}

// Return the estimated ball position
Gaussian2 BallEKFPosVel::position(ulong dt) {
  // Include forward prediction
  Matrix x = predict(TimeToSec(dt));
  Matrix P = predict_cov(TimeToSec(dt));
  Gaussian2 pos;
  pos.mean=vector2d(x.e(0),x.e(1));
  pos.setsxsy(P.e(0,0),P.e(1,1),P.e(1,0));
  return pos;
}

// Return the estimated ball velocity
Gaussian2 BallEKFPosVel::velocity(ulong dt) {
  Matrix x = predict(TimeToSec(dt));
  Matrix P = predict_cov(TimeToSec(dt));
  Gaussian2 vel;
  vel.mean=vector2d(x.e(2),x.e(3));
  vel.setsxsy(P.e(2,2),P.e(3,3),P.e(3,2));
  return vel;
}

// Add velocity to the current state (should also add a covariance est)
void BallEKFPosVel::addVelocity(vector2d vel) {
  added_vel = vel;
}

// A generic accessor function based on the input string argument
bool BallEKFPosVel::getState(Matrix & val, const string & param, int idx,ulong dt) {
  Matrix x;
  if (string("position_mean") == param) {
    x = predict(TimeToSec(dt));
    val.resize(2,1);
    val.e(0,0)=x.e(0,0);
    val.e(0,1)=x.e(1,0);
  } else if (string("velocity_mean") == param) {
    x = predict(TimeToSec(dt));
    val.resize(2,1);
    val.e(0,0)=x.e(2,0);
    val.e(0,1)=x.e(3,0);
  } else if (string("position_cov") == param) {
    x = predict_cov(TimeToSec(dt));
    val.resize(2,2);
    val.e(0,0)=x.e(0,0);
    val.e(0,1)=x.e(0,1);
    val.e(1,0)=x.e(1,0);
    val.e(1,1)=x.e(1,1);
  } else if (string("velocity_cov") == param) {
    x = predict_cov(TimeToSec(dt));
    val.resize(2,2);
    val.e(0,0)=x.e(2,2);
    val.e(0,1)=x.e(2,3);
    val.e(1,0)=x.e(3,2);
    val.e(1,1)=x.e(3,3);
  } else {
    return false;
  }
  return true;
}

bool BallEKFPosVel::getState(Gaussian2 & val, const string & param, int idx, ulong dt) {
  if (string("position") == param) {
    val = position(dt);
  } else if (string("velocity") == param) {
    val = velocity(dt);
  } else {
    return false;
  }
  return true;
}

bool BallEKFPosVel::getState(vector2d & val, const string & param, int idx, ulong dt) {
  if (string("position") == param) {
    val = position(dt).mean;
  } else if (string("velocity") == param) {
    val = velocity(dt).mean;
  } else {
    return false;
  }
  return true;
}