#ifndef __LEAST_SQUARE_H__
#define __LEAST_SQUARE_H__

#include <new>

#include "param/matrix.h"
#include "param/complex.h"

class LeastSquare {
  public:
    /**
     * 一次方程式の解を最小二乗法を利用して解く
     * 
     * @f[
     *    y = K x + e
     * @f]
     * で与えられる一次方程式
     * (ただし@f$ x, y, e, K @f$は
     * それぞれ@f$ m \times 1 @f$の求めるべき解、
     * @f$ n \times 1 @f$の観測量、
     * @f$ n \times 1 @f$の誤差、
     * @f$ n \times m @f$の係数行列)
     * の解@f$ x @f$を求める。
     * 
     * 理論としては、
     * @f[
     *   \frac{\partial}{\partial x} \left( \epsilon^{T} \epsilon \right)
     *     = - 2 K^T \left( y - K x \right) \equiv 0
     * @f]
     * の結果を利用している。
     * 
     * @param y 観測量
     * @param K 係数行列
     */
    template <class FloatT>
    Matrix<FloatT> operator()(
        const Matrix<FloatT> &y, const Matrix<FloatT> &K){
      return (K.transpose() * K).inverse() * K.transpose() * y;
    }
};

class LeastSquareCoef {
  public:
    /**
     * 一次方程式の係数を最小二乗法を利用して解く
     * 
     * @f[
     *    y = K x + e
     * @f]
     * で与えられる一次方程式
     * (ただし@f$ x, y, e @f$は
     * それぞれ@f$ m \times 1 @f$の独立変数、
     * @f$ n \times 1 @f$の従属変数、
     * @f$ n \times 1 @f$の誤差、
     * @f$ n \times m @f$の係数行列)の
     * 複数a回の観測
     * @f[
     *    Y = K X + E
     * @f]
     * を利用して、係数Kを決定する。
     * このとき@f$ X, Y @f$はそれぞれ@f$ m \times a @f$、@f$ n \times a @f$ 行列。
     * 
     * @param X @f$ m \times a @f$ 行列、独立変数を集めたもの
     * @param Y @f$ n \times a @f$ 行列、従属変数を集めたもの
     * @return (Matrix<FloatT>) 係数行列
     */
    template <class FloatT>
    Matrix<FloatT> operator()(
        const Matrix<FloatT> &X, const Matrix<FloatT> &Y){
      
      /* 順次係数行列を決定する方法で解く
       * 
       * @f[ y = K x + e @f]を変換して
       * @f[ y_{i} = x^{T} K^{T}_{i} + e_{i} @f] 
       * ただし @f$ K^{T}_{i} @f$は K行列のi列を転置し、
       * @f$ m \times 1 @f$の縦ベクトルにしたもの
       * 
       * これはXを係数行列とし、
       * Kの一部を求める一次方程式とみなすことができるので、
       * 最小二乗法を適用し求める。
       */
      Matrix<FloatT> K(Y.rows(), X.rows());
      LeastSquare ls;
      for(int i(0); i < K.rows(); i++){
        K.rowVector(i).transpose() 
            = ls(Y.rowVector(i).transpose(), X.transpose()); 
      }
      return K;
    }
};

class WeightedLeastSquare {
  public:
    /**
     * 一次方程式の解を最小二乗法を利用して解く(重み付)
     * 
     * @f[
     *    y = K x + e
     * @f]
     * で与えられる一次方程式
     * (ただし@f$ x, y, e, K @f$は
     * それぞれ@f$ m \times 1 @f$の求めるべき解、
     * @f$ n \times 1 @f$の観測量、
     * @f$ n \times 1 @f$の誤差<
     * @f$ n \times m @f$の係数行列)
     * の解@f$ x @f$を求める。
     * 
     * 理論としては、
     * @f[
     *   \frac{\partial}{\partial x} \left( \epsilon^{T} W \epsilon \right)
     *     = - 2 K^T W \left( y - K x \right) \equiv 0
     * @f]
     * の結果を利用している。
     * 
     * @param y 観測量
     * @param K 係数行列
     * @param W 重み
     */
    template <class FloatT>
    Matrix<FloatT> operator()(
        const Matrix<FloatT> &y, const Matrix<FloatT> &K,
        const Matrix<FloatT> &W){
      return (K.transpose() * W * K).inverse() * K.transpose() * W * y;
    }
};

class WeightedLeastSquareCoef {
  public:
    /**
     * 一次方程式の係数を最小二乗法を利用して解く
     * 
     * @param X @f$ m \times a @f$ 行列、独立変数を集めたもの
     * @param Y @f$ n \times a @f$ 行列、従属変数を集めたもの
     * @param W 従属変数に対する重みを返すファンクタ
     * @return (Matrix<FloatT>) 係数行列
     * @see LeastSquareCoef#operator(X, Y)
     */
    template <class FloatT, class WeightCollection>
    Matrix<FloatT> operator()(
        const Matrix<FloatT> &X, const Matrix<FloatT> &Y,
        const WeightCollection &W){
      Matrix<FloatT> K(Y.rows(), X.rows());
      WeightedLeastSquare wls;
      for(int i(0); i < K.rows(); i++){
        K.rowVector(i).transpose() 
            = wls(Y.rowVector(i).transpose(), X.transpose(), W[i]); 
      }
      return K;
    }
};

template <class FloatT>
class RecursiveWeightedLeastSquare {
  protected:
    Matrix<FloatT> m_P;
  public:
    // TODO: operator new [] をなんとか定義してPなしでの初期化を回避したい
    RecursiveWeightedLeastSquare() : m_P() {}
    
    RecursiveWeightedLeastSquare(
        Matrix<FloatT> &P) : m_P(P) {
      
    }
    ~RecursiveWeightedLeastSquare(){}
  
  protected:
    /**
     * 新たな観測量を用いて更新する場合に用いるゲイン行列を求める
     * 
     * @param H 観測行列
     * @param r 観測の正しさ(誤差共分散行列)
     * @return (Matrix<FloatT>) ゲイン行列
     */
    Matrix<FloatT> operator()(
        const Matrix<FloatT> &h, const Matrix<FloatT> &r) const {
      Matrix<FloatT> h_t(h.transpose());
      return m_P * h_t * ((h * m_P * h_t) + r).inverse();
    }
    
  public:
    /**
     * 再帰的に推定値を更新
     * 
     * @param z 新たな観測量
     * @param H 観測行列
     * @param r 観測の正しさ(誤差共分散行列)
     * @param orig_x 更新前の推定値
     * @return (Matrix<FloatT>) 更新された結果の推定値
     */
    Matrix<FloatT> &operator()(
        const Matrix<FloatT> &z, const Matrix<FloatT> &h,
        const Matrix<FloatT> &r, 
        Matrix<FloatT> &orig_x){
      Matrix<FloatT> K(operator()(h, r));
      m_P = (Matrix<FloatT>::getI(K.rows()) - K * h) * m_P;
      orig_x += K * (z - h * orig_x);
      return orig_x;
    }
    
    Matrix<FloatT> &getP() {return m_P;}
};

template <class FloatT>
class RecursiveWeightedLeastSquareCoef {
  protected:
    RecursiveWeightedLeastSquare<FloatT> *rwls;
    
  public:
    template <class PCollection>
    RecursiveWeightedLeastSquareCoef(PCollection &Ps)
        : rwls(new RecursiveWeightedLeastSquare<FloatT> [Ps.size()]) {
      for(int i(0); i < Ps.size(); i++){
        new(&rwls[i]) RecursiveWeightedLeastSquare<FloatT>(Ps[i]);
      }
    }
    ~RecursiveWeightedLeastSquareCoef(){
      delete [] rwls;
    }
    
    /**
     * 再帰的に推定値を更新
     * 
     * @param y 新たに得られた従属変数
     * @param x 新たに得られた独立変数
     * @param r 新たに得られた従属変数の正しさ
     *     (例えば1次元配列、各従属変数間の正しさは独立である)
     * @param orig_K 更新前の係数行列の推定値
     * @return (Matrix<FloatT>) 更新された係数行列の推定値
     */
    template <class WeightCollection>
    Matrix<FloatT> &operator()(
        const Matrix<FloatT> &y, const Matrix<FloatT> &x,
        const WeightCollection &w, 
        Matrix<FloatT> &orig_K){
      Matrix<FloatT> x_t(x.transpose());
      for(unsigned int i(0); i < y.rows(); i++){
        /*
         * 通常の再帰的重み付最小二乗法に帰着させる、その際、
         * 
         * z <= y(i, 0)
         * h <= x.transpose()
         * r <= y(i, 0) の推定される分散 w[i]
         * orig_x <= orig_K.rowVector(i).transpose()
         */
        (rwls[i])(y.partial(1, y.columns(), i, 0), x_t, 
            w[i], orig_K.rowVector(i).transpose());
      }
      return orig_K;
    }
};

#endif /* __LEAST_SQUARE_H__ */