// file: $isip_ifc/class/algo/KalmanFilter/kf_05.cc
// version: $Id: kf_05.cc 10602 2006-08-17 17:43:12Z ss754 $
//

// isip include files
//
#include "KalmanFilter.h"

#ifndef ISIP_PREDICTION
#include <Prediction.h>
#endif

// method: apply
//
// arguments:
// Vector<AlgorithmData>& output: (output) output data
// const Vector<CircularBuffer<AlgorithmData< >& input: (input) input data
//
// return: a bolean value indicating status
//
boolean KalmanFilter::apply(Vector<AlgorithmData>& output_a,
		      const Vector<CircularBuffer<AlgorithmData> >& input_a) {

  // calculate the number of channels and work on each channel at a
  // time
  //
  long len = input_a.length();
  
  output_a.setLength(len);
  boolean res = true;

  // start the debugging output
  //
  displayStart(this);

  for ( long c = 0; c < len; c++) {

    // display the channel number
    //
    displayChannel(c);

    if (input_a(c)(0).getDataType() == AlgorithmData::VECTOR_DOUBLE) {
      
      // call AlgorithmData::makeVectorDouble to force the output
      // vector for this channel to be a VectorDouble, call
      // AlgorithmData::getVectorDouble on the input for this channel
      // to check that the input is already a VectorDouble and return
      // that vector.
      //
      res &= compute(output_a(c).makeVectorDouble(),
		     input_a(c)(0).getVectorDouble(),
		     input_a(c)(0).getCoefType(),
		     c);
    }

    // set the coefficient type for the output
    //            
    output_a(c).setCoefType(AlgorithmData::SIGNAL);
    
  }
  
  // finish the debugging output
  //
  displayFinish(this);

  // exit gracefully
  //
  return res;
}

// method: compute
//
// arguments:
// VectorDouble& output: (output) output data
// const VectorDouble& input: (input) input data
// AlgorithmData::COEF_TYPE input_coef_type: (input) type of input
// long channel_index: (inout) channel index
//
// return: a boolean value indicating status
//
// this method computes the filtered signal for the specified
// algorithm and implementation
//
boolean KalmanFilter::compute(VectorDouble& output_a,
			      const VectorDouble& input_a,
			      AlgorithmData::COEF_TYPE input_coef_type_a,
			      long channel_index_a) {

  // declare local variables
  //
  boolean status = false;

  if (cmode_d != FRAME_INTERNAL) {

    return Error::handle(name(), L"compute", ERR_UNSUPM, __FILE__, __LINE__);
  }
  else {
    frame_index_d++;
    if ((!is_valid_d) || (frame_index_d == 0)) {
      init();
    }
    
    // algorithm: CONVENTIONAL and implementation: AR_SCALAR
    //
    if (implementation_d == AR_SCALAR) {
      if (algorithm_d == CONVENTIONAL) { 
	status = computeStates(input_a, output_a);
      }
      else if(algorithm_d == UNSCENTED) {
	status = computeStates(input_a, output_a);
	Console::put(L"unscented Kalman filter -- ");
	Console::put(L"yet to complete the implementation kf_05.cc line 128");
      }
      else {
	// actually unknown algorithm
	//
	return Error::handle(name(), L"compute",
			     ERR_UNKIMP, __FILE__, __LINE__);
      }
    }
    // error: unknown implementation
    //
    else {
      return Error::handle(name(), L"compute",
			   ERR_UNKIMP, __FILE__, __LINE__);
    }
    
    // possibly display the data
    //
    display(output_a, input_a, name());
  }
  
  // exit gracefully
  //
  return status;
}

// method: computeStates
// algorithm: CONVENTIONAL
// implementation: AR_SCALAR
// 
boolean KalmanFilter::computeStates(const VectorDouble& obsn_a,
				    VectorDouble& output_a) {

  boolean status = true;
  
  // assuming that the data is copied in observation data member
  // time will be equal to the number of observation actually, it
  // will be number of iterations that the calculation will go
  // through
  //
  long time = obsn_a.length();
  initTransitionModel(obsn_a);
   int i = 0;

  if ((time > 0) && (algorithm_d == CONVENTIONAL)) {
    
    // for first frame
    //
    if (frame_index_d == 0) {
      
      // first 'order' elements will be assumed as clean signal
      //
      if (state_mem_d(0) == (double)-1.0) {
	state_mem_d.move(obsn_a, (int)order_d, 0, 0);
      }
      
      data_filt_d.move(state_mem_d, (int)order_d, 0, 0);
      i = order_d;
    }
    else {
      i = 0;
    }

    // same for loop for all the frames
    // i = order_d (for first frame) or
    // i = ZERO (otherwise)
    //
    for ( ; i < time; i++) {

      // predicting state value
      //
      // x^'(k) = F * x^(k-1)
      trn_d.multv(state_mem_d, state_mem_d);

      // predicting the estimate error covariance
      //
      // P`(k) = F * P(k-1) * F_t + Q
      // where,
      // F_t -> transpose of F
      //
      MatrixDouble trn_t(trn_d.getNumColumns(), trn_d.getNumRows());
      MatrixDouble mat1(covar_err_est_d.getNumRows(),
		       covar_err_est_d.getNumColumns());

      trn_t.transpose(trn_d);
      mat1.mult(trn_d, covar_nse_pro_d);
      mat1.mult(trn_t);
      covar_err_est_d.add(mat1, covar_err_est_d);

      // finding kalman gain
      //
      // k(k) = P'(k) * H_t * [H * P`(k) * H_t + R]inverse
      //  where,
      //  H_t -> transpose of H
      //
      VectorDouble kalman_gain((int)order_d);
      VectorDouble vec1((int)order_d);

      covar_err_est_d.multv(kalman_gain, obs_d);

      // inside the inverse bracket
      //
      covar_err_est_d.vmult(vec1, obs_d);

      double value = vec1.dotProduct(obs_d);
      value += (double) covar_nse_obs_d;
      kalman_gain.div(value);

      // correcting the state value by innovation
      // innovation is associated with kalman gain and
      // error in the predicted observation and actual one
      //
      // x^(k) = x^`(k) + k(k) * [ y(k) - H * x^`(k)]
      //
      double obsn_est = obs_d.dotProduct(state_mem_d);
      obsn_est = (double) obsn_a(i) - obsn_est;
      kalman_gain.mult(obsn_est);
      state_mem_d.add(kalman_gain);

      // correcting the estimate error covariance
      //
      // P(k) = [ I - k(k) * H] * P`(k)
      //
      MatrixDouble mat2(kalman_gain.length(), 1);
      MatrixDouble mat3(1, kalman_gain.length());
      MatrixDouble mat4(covar_err_est_d.getNumRows(),
		       covar_err_est_d.getNumColumns());

      for (int j = 0; j < kalman_gain.length(); j++) {
	mat2.setValue(j, 0, (double)kalman_gain(i));
	mat3.setValue(0, j, (double)obs_d(i));
	mat4.setValue(j, j, 1);
      }

      mat1.mult(mat2, mat3);
      mat1.sub(mat4, mat1);
      covar_err_est_d.mult(mat1, covar_err_est_d);
      data_filt_d(i) = obs_d.dotProduct(state_mem_d);
    }
  }
  else if ((time > 0) && (algorithm_d == UNSCENTED)) {

    int aug_len = state_mem_d.length();
    int sigma_len = 2 * aug_len + 1;
    MatrixDouble sigma_all(aug_len, sigma_len);
    MatrixDouble sigma_state((int)order_d, sigma_len);

    // computing weights neccessary for 
    VectorDouble wts_c(sigma_len);
    VectorDouble wts_m(sigma_len);

    double lambda = (double)alpha_d * (double)alpha_d
      * (aug_len + (double)kappa_d) - aug_len;
    wts_m(0) = lambda / (aug_len + lambda);
    wts_c(0) = wts_m(0) + ( 1 - (double)alpha_d * (double)alpha_d
			    + (double)beta_d);

    for (int j = 1; j < sigma_len; j++) {
      wts_m(j) = 1.0 / (2.0 * (aug_len + lambda));
      wts_c(j) = wts_m(j);
    }      

    // another constant to the sigma points
    //
    Double sqrtL_n_lambda = aug_len + lambda;
    sqrtL_n_lambda.sqrt();
    
    // for first frame
    //
    if (frame_index_d == 0) {
      // first 'order' elements will be assumed as clean signal
      //
      if (state_mem_d(0) == (double)-1.0) {
	state_mem_d.move(obsn_a, (int)order_d, 0, 0);
      }
      data_filt_d.move(state_mem_d, (int)order_d, 0, 0);
      i = order_d;
    }
    else { i = 0; }

    covar_err_est_d.setValue((int)order_d + 1,
			     (int)order_d + 1, (double)covar_nse_obs_d);
    covar_err_est_d.setValue((int)order_d, (int)order_d,
			     (double)covar_nse_pro_d.getValue(0, 0));
    
    for ( ; i < time; i++) {

      // square root of the estimate error covariance matrix
      //
      covar_err_est_d.mult((double)sqrtL_n_lambda);

      MatrixDouble sqrt_covar;
      sqrt_covar.setDimensions(aug_len, aug_len);
      matrixSqrt(covar_err_est_d, sqrt_covar);
      covar_err_est_d.assign(sqrt_covar);

      for ( int j = 0; j < sigma_len; j++){
	for( int k = 0; k < (int)order_d; k++) {
	  sigma_all.setValue(k, j, state_mem_d(k));
	}
      }
      
      // getting sigma points
      // based on data_filt_d and covar_err_est_d square root
      //
      for (int j = 0; j < aug_len; j++) {
	for ( int k = 0; k < aug_len; k++) {
	  double val = sigma_all.getValue(k, 0) -
	               covar_err_est_d.getValue(k, j);
	  sigma_all.setValue(k, j+1, val);
	  val = sigma_all.getValue(k, 0) +
	               covar_err_est_d.getValue(k, j);
	  sigma_all.setValue(k, aug_len + j + 1, val);
	}
      }

      for (int j = 0; j < sigma_len; j++)
	for (int k = 0; k < (int)order_d; k++) 
	  sigma_state.setValue(k, j, (double)sigma_all.getValue(k, j));
      
      // getting time update for sigma state
      //
      VectorDouble state_current(aug_len);
      VectorDouble state_update(order_d);

      // get a column of state sigma,
      // multiply trn_s with this column,
      // put back the result (which is a vector) back into 
      // the same column.
      // repeat the above process for all columns
      //
      for ( int j = 0; j < sigma_len; j++) {
	sigma_state.getColumn(state_update, j);
	trn_d.multv(state_update, state_update);
	sigma_state.setColumn(j, state_update);
      }

      VectorDouble pro_nse_sigma(aug_len);
      sigma_all.getRow(pro_nse_sigma, (int)order_d);
      
      sigma_state.getRow(state_current, (int)order_d - 1);
      state_current.add(pro_nse_sigma);
      sigma_state.setRow( (int)order_d - 1, state_current);
      
      // averaging the state sigmas
      // get a column of state sigma (now it is time updated)
      // multiply it by respective weights
      //
      for (int j = 0; j < state_update.length(); j++)
	state_update(j) = 0.0;
      
      for (int j = 0; j < state_current.length(); j++)
	state_current(j) = 0.0;
      
      for (int j = 0; j < sigma_len; j++) {
	sigma_state.getColumn(state_current, j);
	state_current.mult((double)wts_m(j));
	state_update.add(state_current);
      }
      
      for ( int j = 0; j < state_current.length(); j++)
	state_current(j) = 0.0;

      // state covariance matrix
      //
      MatrixDouble state_covar;
      state_covar.setDimensions((int)order_d, (int)order_d);

      MatrixDouble state_diff((int)order_d, sigma_len);
      MatrixDouble state_diff_T(sigma_len, (int)order_d);

      // to find the difference between the state sigms and mean of
      // time updated state sigmas
      //
      for (int j = 0; j < sigma_len; j++) {
	sigma_state.getColumn(state_current, j);
	state_current.sub(state_update);
	state_diff_T.setRow(j, state_current);
	state_current.mult((double)wts_c(j));
	state_diff.setColumn(j, state_current);
      }

      state_diff.mult(state_covar, state_diff_T);

      // observation sigmas
      //
      VectorDouble sigma_obsNoise(sigma_len);
      sigma_state.getRow(state_current, (int)order_d - 1);
      sigma_all.getRow(sigma_obsNoise, (int)order_d + 1);
      state_current.add(sigma_obsNoise);

      double obsn_est = 0.0;
      for (int j =0; j < sigma_len; j++)
	obsn_est = obsn_est + (double)wts_m(j) * (double)state_current(j);

      // measurement update equations
      // obsn covariance matrix
      // find difference between the sigma point and estimated obsn
      // multiply the difference by the wts_c
      //
      Double covar_obs = 0.0;
      state_current.sub((double)obsn_est);

      for (int j = 0; j < sigma_len; j++)
	covar_obs = (double)covar_obs + (double)wts_c(j)
	              * (double)state_current(j)  * (double)state_current(j);

      // computing cross covariance
      //
      VectorDouble cross_covar((int)order_d);
      for (int k = 0; k < (int)order_d; k++) {
	cross_covar(k) = 0.0;
	for(int j = 0; j < sigma_len; j++) {
	  cross_covar(k) = cross_covar(k) + wts_c(j) *
	                   (sigma_state.getValue(k, j) - state_update(k)) *
	                   state_current(j);
	}
      }

      // computing kalman gain
      //
      VectorDouble kg((int)order_d);
      kg.div(cross_covar, covar_obs);

      // setting up the addition term for measurement update of state
      // filt_data_d is our measurement update of state term
      //
      kg.mult((double)obsn_a(i) - obsn_est);
      state_update.add(state_update, kg);

      state_mem_d.assign(state_update);
      data_filt_d(i) = state_update((int)order_d - 1);
      
      // computing measurement update state_covar
      //
      kg.div((double)obsn_a(i) - obsn_est);

      MatrixDouble kg_mat;
      kg_mat.setDimensions((int)order_d, 1);
      MatrixDouble kg_mat_T;
      kg_mat.setColumn(0, kg);
      kg_mat_T.transpose(kg_mat);

      MatrixDouble kg_x_kg_T;
      kg_x_kg_T.setDimensions((int)order_d, (int)order_d);
      kg_x_kg_T.mult(kg_mat, kg_mat_T);
      kg_x_kg_T.mult(covar_obs);
      state_covar.sub(state_covar, kg_x_kg_T);

      for (int j = 0; j < (int)order_d; j++)
	for (int k = 0; k < (int)order_d; k++)
	  covar_err_est_d.setValue(j, k, (double)state_covar.getValue(j, k));
      covar_err_est_d.setValue((int)order_d, (int)order_d, covar_nse_obs_d);
      covar_err_est_d.setValue((int)order_d + 1, (int)order_d + 1,
			       covar_nse_pro_d.getValue(0, 0));
      
    } // looping through all the sample in a frame
    
  }
  
  return status;
}

// method: initTransitionModel
// 
boolean KalmanFilter::initTransitionModel(const VectorDouble& obsn_a) {

  // 1. first compute the autocorrelation coefficients
  // 2. then compute the LP coefficients
  // 
  int i, j;
  VectorFloat lpc;
  
  int frame_length = (int) obsn_a.length();
  VectorFloat autoc((int)order_d + 1);
  for (i = 0; i < (int)order_d + 1; i++) {
    autoc(i) = 0;
    for (j = 0; j < frame_length - i; j++) {
      autoc(i) = autoc(i) + (float) obsn_a(j) * (float)obsn_a(j + i);
    }
  }

  for ( i = order_d; i >= 0; i--) {
    autoc(i) = autoc(i) / autoc(0);
  }

  Prediction lp(Prediction::AUTOCORRELATION,
		Prediction::DURBIN, order_d, -60);
  lp.compute(lpc, autoc);
  //  lpc.debug(L"LP coefficients:");
  //  return true;

  // set up the parameters for the calculation
  // transition model will have the last row as lp coefficients
  //
  //  next_state =   transition_model * current_state
  //
  //   x(k)      =      trn_d         *  x(k-1)
  //  _    _          _          _         _    _ 
  // |      |        |            |       |      |
  // | x(2) |        | 0  1  0  0 |       | x(1) |
  // | x(3) |        | 0  0  1  0 |       | x(2) |
  // | x(4) |    =   | 0  0  0  1 |   *   | x(3) |
  // | x(5) |        | a4 a3 a2 a1|       | x(4) |
  // |_    _|        |_          _|       |_    _|
  //
  // where,
  // a4, a3, a2, a1 are LP coefficients ~ 4th order AR model
  // x(5) is the element desired from the AR modeled kalman filter
  //
  double lpc_value = 0;
  for ( j = 0; j < order_d; j++) {
    lpc_value = (double) lpc((int)order_d - j);
    trn_d.setValue((int)order_d - 1, j, lpc_value);
  }

  //  trn_d.debug(L"transition model : ");
  return true;
}

// method: matrixSqrt
// to find a square root of a matrix by Denman-Beavers square root iteration
// ref: http://en.wikipedia.org/wiki/Square_root_of_a_matrix
//
boolean KalmanFilter::matrixSqrt(MatrixDouble& mat_in, MatrixDouble& mat_sqrt) {

  int row = mat_in.getNumRows();
  int col = mat_in.getNumColumns();

  MatrixDouble mat1;
  MatrixDouble mat2;
  MatrixDouble mat3;
  MatrixDouble mat4;

  mat1.setDimensions(row, col);
  mat2.setDimensions(row, col);
  mat3.setDimensions(row, col);
  mat4.setDimensions(row, col);

  mat1.assign(mat_in);
  mat2.setDiagonal(1.0);

  for ( int i = 0; i < 15; i++) {
    mat3.inverse(mat2);
    mat3.add(mat1);
    mat3.mult(0.5);

    mat4.inverse(mat1);
    mat4.add(mat2);
    mat4.mult(0.5);

    mat2.assign(mat4);
    mat1.assign(mat3);    
  }
  
  return mat_sqrt.assign(mat1);
}
