#ifndef __SI_COMMON_H
#define __SI_COMMON_H

/** @file
 * @brief システム同定をする際の共通部分を切り出したヘッダ
 *
 */

#include <strstream>
#include <string>
#include <cstring>
#include <vector>
#include <map>
#include <set>

#include "param/matrix.h"
#include "algorithm/kalman.h"

template <class FloatT>
void extract_values(
    std::vector<char *> &orig, std::map<std::string, FloatT> &extracted){
  for(std::vector<char *>::iterator it(orig.begin());
      it != orig.end();
      it++){
    // strtokを使いたくない!!
    FloatT v(0);
    unsigned i(0);
    unsigned len(strlen(*it));
    for(; i < len; i++){
      if((*it)[i] == ','){(*it)[i++] = '\0'; break;}
    }
    if(i < len){std::stringstream(std::string(&((*it)[i]))) >> v;}
    extracted[std::string(*it)] = v;
  }
}

template <class FloatT, class MaskT>
struct DiagonalElementOfMatrix {
  const Matrix<FloatT> &_mat;
  const MaskT &_mask;
  const char *_separator;
  const unsigned _precision;
  DiagonalElementOfMatrix(
      const Matrix<FloatT> &mat, const MaskT &mask,
      const char *separator, const unsigned precision)
        : _mat(mat), _mask(mask), _separator(separator), _precision(precision) {}
    ~DiagonalElementOfMatrix(){}
  friend std::ostream& operator<<(std::ostream &out, const DiagonalElementOfMatrix &target){
    for(unsigned i(0), j(0); j < (target._mat.rows() + target._mask.size()); j++){
      out << std::setprecision(target._precision);
      if(j != 0){out << target._separator;}
      if(target._mask.find(j) == target._mask.end()){
        out << const_cast<Matrix<FloatT> &>(target._mat)(i, i);
        i++;
      }else{
        out << 0;
      }
    }
    return out;
  }
};

template <class FloatT, class MaskT>
DiagonalElementOfMatrix<FloatT, MaskT> diagonal(
    const Matrix<FloatT> &mat, const MaskT &mask,
    const char *separator = ",", const unsigned precision = 8){
  return DiagonalElementOfMatrix<FloatT, MaskT>(mat, mask, separator, precision);
}

// フィルタの基底クラス
template <
    class EqMotion, class EqObs,
    class States, class Input, class Observed,
    class FloatT = double>
class AbstractKalmanFilter {
  protected:
    typedef EqMotion dynamics_t;
    typedef EqObs obs_t;
    typedef States states_t;
    typedef Input input_t;
    typedef Observed observed_t;
    dynamics_t &dynamics;
    obs_t &obs;
  public:
    AbstractKalmanFilter(dynamics_t &d, obs_t &o) : dynamics(d), obs(o) {}
    ~AbstractKalmanFilter() {}
    virtual KalmanFilter<FloatT> &filter() = 0;
    virtual void predict(states_t &, input_t &) = 0;
    virtual void correct(states_t &, observed_t &, Matrix<FloatT> &) = 0;
    virtual void dump(
        std::ostream &out, const FloatT &t, const states_t &states){
      out << t << ","
          << states.join(",") << ","
          << diagonal(filter().getP(), states_t::masked()) << endl;
    }
};

// UKFの設定、アップデート手順
template <
    class EqMotion, class EqObs,
    class State, class Input, class Observed,
    class FloatT = double>
class UKF : public AbstractKalmanFilter<
    EqMotion, EqObs,
    State, Input, Observed, FloatT> {
  protected:
    typedef AbstractKalmanFilter<
        EqMotion, EqObs,
        State, Input, Observed, FloatT> super_t;
    UnscentedKalmanFilter<FloatT> ukf;
  public:
    UKF(dynamics_t &d, obs_t &o) :
        super_t(d, o),
        ukf(
            Matrix<FloatT>(states_t::variables(), states_t::variables()),
            Matrix<FloatT>(input_t::variables(), input_t::variables())) {};
    ~UKF(){};
    KalmanFilter<FloatT> &filter() {return ukf;}
    void predict(states_t &x, input_t &u){
      ukf.predict(super_t::dynamics, x, u);
    }
    void correct(states_t &x, observed_t &z, Matrix<FloatT> &R){
      ukf.correct(super_t::obs, x, z, R);
    }
};

// EKFの設定、アップデート手順
template <
    class EqMotion, class EqObs,
    class State, class Input, class Observed,
    class FloatT = double>
class EKF : public AbstractKalmanFilter<
    EqMotion, EqObs,
    State, Input, Observed, FloatT> {
  protected:
    typedef AbstractKalmanFilter<
            EqMotion, EqObs,
            State, Input, Observed, FloatT> super_t;
    KalmanFilter<FloatT> ekf;
  public:
    EKF(dynamics_t &d, obs_t &o) :
        super_t(d, o),
        ekf(
            Matrix<FloatT>(states_t::variables(), states_t::variables()),
            Matrix<FloatT>(input_t::variables(), input_t::variables())) {};
    ~EKF(){};
    KalmanFilter<FloatT> &filter() {return ekf;}
    void predict(states_t &x, input_t &u){
      //std::cerr << "A:" << super_t::dynamics.getA(x, u) << std::endl;
      //std::cerr << "B:" << super_t::dynamics.getB(x, u) << std::endl;
      ekf.predict(
          super_t::dynamics.getA(x, u), super_t::dynamics.getB(x, u),
          super_t::dynamics.delta_t);
      x = super_t::dynamics(x, u);  // 時間更新
    }
    void correct(states_t &x, observed_t &z, Matrix<FloatT> &R){
      observed_t tilde_z(super_t::obs(x));
      //std::cerr << "H:" << super_t::obs.getH(x) << std::endl;
      Matrix<FloatT> K(ekf.correct(super_t::obs.getH(x, z), R));
      Matrix<FloatT> delta_z(observed_t::variables(), 1);
      for(unsigned i(0); i < delta_z.rows(); i++){
        delta_z(i, 0) = tilde_z[i] - z[i];
      }
      Matrix<FloatT> mod_x(K * delta_z);
      for(unsigned i(0); i < states_t::variables(); i++){
        x[i] -= mod_x(i, 0);
      }
    }
};

// 状態量の時間更新のみ行い、修正はしない => 運動方程式のシミュレータとして使える
template <
    class EqMotion, class EqObs,
    class State, class Input, class Observed,
    class FloatT = double>
class Simulator : public AbstractKalmanFilter<
    EqMotion, EqObs,
    State, Input, Observed, FloatT> {
  protected:
    typedef AbstractKalmanFilter<
        EqMotion, EqObs,
        State, Input, Observed, FloatT> super_t;
    KalmanFilter<FloatT> dummy;
  public:
    Simulator(dynamics_t &d, obs_t &o) :
        super_t(d, o),
        dummy(
            Matrix<FloatT>(states_t::variables(), states_t::variables()),
            Matrix<FloatT>(input_t::variables(), input_t::variables())) {};
    ~Simulator(){};
    KalmanFilter<FloatT> &filter() {return dummy;}
    void predict(states_t &x, input_t &u){
      x = super_t::dynamics(x, u);  // 時間更新
    }
    void correct(states_t &x, observed_t &z, Matrix<FloatT> &R){
      // 状態量の修正はしない
      // かわりに観測方程式をといて、観測量を求めなおす
    }
    virtual void dump(
        std::ostream &out, const FloatT &t, const states_t &states){
      // シミュレータの場合、出力を時刻、観測量、状態量にする
      // 本当はシステム同定の入力用ファイルと同じ書式にしたい
      out << t;
      if(super_t::obs.input){
        out << "," << super_t::obs(states).join(",");
      }else{
        // 開始前なので、観測量を求めることは出来ない
        for(unsigned i(0); i < observed_t::variables(); i++){
          out << "," << 0;
        }
      }
      out << "," << states.join(",") << std::endl;
    }
};

#endif /* __SI_COMMON_H */