#include "kalmanFilter.h"

KalmanFilter::KalmanFilter()
{
    dt_ = 1.0/20.0; //20hz
    int n = 3; // Number of states
    int m = 2; // Number of measurements
    Eigen::MatrixXd A(n, n); //3x3
    Eigen::MatrixXd C(m, n); //2x3
    Eigen::MatrixXd Q(n, n); //3x3
    Eigen::MatrixXd R(m, m); //2x2
    Eigen::MatrixXd P(n, n); //3x3
    A << 1, dt_, 0, 0, 1, dt_, 0, 0, 1;
    C << 1, 0, 0,
        0, 1, 0;
    Q << 1e-4, 1e-4, 1e-4,
        1e-4, 1e-2, 1e-3,
        1e-4, 1e-3, 1e-1;
    R << 1e-4, 1e-3,
        1e-3, 1e-2;
    P << 1e-3, 1e-3, 1e-3,
        1e-3, 1e1, 1e-1,
        1e-3, 1e-1, 1e3;
    init(A, C, Q, R, P);
}

void KalmanFilter::init(
    const Eigen::MatrixXd& A,
    const Eigen::MatrixXd& C,
    const Eigen::MatrixXd& Q,
    const Eigen::MatrixXd& R,
    const Eigen::MatrixXd& P)
{
    m_ = C.rows();
    n_ = A.rows();
    tf::Quaternion quat(0,0,0,1);
    tf::Vector3 T(0,0,0);
    m_last_pose=tf::Transform(quat,T);
    m_last_time=clock();
    velocity.setZero();
    for (int i = 0; i < 3; i++)
    {
        A_[i] = A;
        C_[i] = C;
        Q_[i] = Q;
        R_[i] = R;
        P_[i] = P;
        Eigen::MatrixXd I(n_, n_);
        I_[i] = I;
        I_[i].setIdentity();
        Eigen::MatrixXd temp1(n_, 1), temp2(n_, 1);
        x_hat[i] = temp1;
        x_hat[i].setZero();
        x_hat_new[i] = temp2;
    }
}

KalmanFilter::~KalmanFilter()
{
}

void KalmanFilter::posesCombine()
{
    tf::Vector3 weightedPosition(0,0,0);
    tf::Quaternion weightedQuat(0,0,0,1);
    double weightedTheta = 0;
    double sum=0, angSum=0;
    for (int i = 0; i < poses_vec.size(); i++)
    {
        // weightedPosition += poses_vec[i].pose.getOrigin() * (1.0 / poses_vec[i].error);
        // weightedTheta += poses_vec[i].theta * (1.0 / poses_vec[i].error);
        // sum += 1.0 / poses_vec[i].error;

        double dist = (poses_vec[i].pose.getOrigin() - m_last_pose.getOrigin()).length2();
        double angDist = poses_vec[i].error;

        // double dist = (poses_vec[i].pose.getOrigin() - m_last_pose.getOrigin()).length2();
        // double angDist = (poses_vec[i].pose.getRotation() - m_last_pose.getRotation()).length2();

        if(dist==0)
            dist = 1.0 / 1e-7;
        else
            dist = 1.0 / dist;

        if(angDist==0)
            angDist = 1.0 / 1e-7;
        else
            angDist = 1.0 / angDist;
        
        weightedPosition += poses_vec[i].pose.getOrigin() * dist;
        weightedTheta += poses_vec[i].theta * angDist;
        sum += dist;
        angSum += angDist;
        // std::cout<<"--trans: "<<poses_vec[i].pose.getOrigin()[0]<<","<<poses_vec[i].pose.getOrigin()[1]<<", "<<poses_vec[i].pose.getOrigin()[2]<<", "<<dist<<std::endl;
        std::cout<<"--theta: "<<poses_vec[i].theta<<", "<<angDist<<std::endl;
        // printf("pose: %.3f, %.3f, %.3f\n",poses_vec[i].pose.getOrigin()[0], poses_vec[i].pose.getOrigin()[1], poses_vec[i].pose.getOrigin()[2]);
        // std::cout<<"pose: \n"<<poses_vec[i].pose.getOrigin()<<std::endl;
    }
    weightedPosition /= (sum==0?1e-7:sum);
    weightedTheta /= (angSum==0?1e-7:angSum);
    printf("pose: %.3f, %.3f, %.3f %.3f\n",weightedPosition[0],weightedPosition[1],weightedPosition[2],weightedTheta);
    weightedQuat.setW(cos(weightedTheta/2.0));
    weightedQuat.setZ(sin(weightedTheta/2.0));
    m_last_pose.setRotation(weightedQuat);
    m_last_pose.setOrigin(weightedPosition);
    poses_vec.clear();
    m_last_time=clock();
}

void KalmanFilter::addData(tf::Pose pose, double error)
{
    Eigen::MatrixXd y(m_,1);
    // std::cout<<"add 0 "<<std::endl;
    
    for (int i = 0; i < 3; i++)
    {
        P_[i] = A_[i] * P_[i] * A_[i].transpose() + Q_[i];
        K_[i] = P_[i] * C_[i].transpose() * (C_[i] * P_[i] * C_[i].transpose() + R_[i]).inverse();
        //std::cout << "K_: " << K_[i] << std::endl;
        y(0,0) = pose.getOrigin()[i];
        y(1,0) = (y(0,0)-m_last_pose.getOrigin()[i])/dt_;
        //std::cout<<"hat "<<i<<" : "<<x_hat[i]<<std::endl;
        x_hat_new[i] = A_[i] * x_hat[i] + K_[i] * (y - C_[i] * A_[i] * x_hat[i]);
        // std::cout<<"part1: "<<A_ * x_hat[i]<<std::endl;
        // std::cout<<"part2: "<<K_ * (y - C_ * A_ * x_hat[i])<<std::endl;
        // std::cout<<"inner: "<<y - C_ * A_ * x_hat[i]<<std::endl;
        // std::cout<<"y: "<<y<<std::endl;
        // std::cout<<"inner: "<<C_ * A_ * x_hat[i]<<std::endl;
        // std::cout<<x_hat[i].rows()<<x_hat[i].cols()<<std::endl;
        P_[i] = (I_[i] - K_[i] * C_[i]) * P_[i];
        x_hat[i] = x_hat_new[i];
    }
    // std::cout<<"add 2"<<std::endl;
    tf::Vector3 Position(x_hat[0](0,0), x_hat[1](0,0), x_hat[2](0,0));
    // tf::Quaternion Quat(0, 0, 0, 1);
    // Quat.setRPY(0, 0, atan2(x_hat[1](0,0), x_hat[0](0,0)));
    // std::cout<<"add 3"<<std::endl;
    poseInfo pInfo;
    // printf("quat: %.3f, %.3f, %.3f, %.3f\n",Quat.getX(),Quat.getY(),Quat.getZ(),Quat.getW());
    pInfo.pose.setOrigin(Position);
    pInfo.pose.setRotation(pose.getRotation());
    pInfo.theta = atan2(pose.getRotation().getZ(), pose.getRotation().getW());
    pInfo.error = error==0?1e-10:error;
    poses_vec.push_back(pInfo);
    printf("input: %.3f, %.3f, %.3f, %.3f\n",pose.getOrigin()[0], pose.getOrigin()[1], pose.getOrigin()[2], pInfo.theta);
    // printf("origin: %.3f, %.3f, %.3f\n",x_hat[0](0,0), x_hat[1](0,0), x_hat[2](0,0));
    // std::cout<<"add 4"<<std::endl;
}

tf::Pose KalmanFilter::filter()
{
    posesCombine();
    return m_last_pose;

}

