#include <iostream>
#include <fstream>
#include <sstream>
#include <stdexcept>
#include <utility>
#include <string>
#include <Eigen/Dense>
#include <yaml-cpp/yaml.h>

#include "core/problem/problem.h"
#include <core/ceres_wrapper/solver_ceres.h>
#include "core/sensor/sensor_pose.h" // not necessary
#include "core/frame/frame_base.h"

// Object SLAM
#include "objectslam/processor/processor_tracker_landmark_object.h"
#include "objectslam/capture/capture_object.h"
#include "objectslam/internal/config.h"
#include "objectslam/landmark/landmark_object.h"

// IMU
#include "imu/sensor/sensor_imu.h"
#include "imu/capture/capture_imu.h"
#include "imu/processor/processor_imu.h"
#include "imu/factor/factor_imu.h"

// ROS
#include <rosbag/bag.h>
#include <rosbag/view.h>
#include <wolf_ros_objectslam/CosyObjectArray.h>
#include <wolf_ros_objectslam/CosyObject.h>
#include <sensor_msgs/Imu.h>
#include <std_msgs/Header.h>
#include <boost/foreach.hpp>
#define foreach BOOST_FOREACH

using namespace wolf;
using namespace Eigen;

Isometry3d posevec_to_isometry(Vector7d pose)
{
    return Translation<double, 3>(pose.head<3>()) * Quaterniond(pose.tail<4>());
}

Vector7d isometry_to_posevec(Isometry3d M)
{
    Vector7d pose_vec;
    pose_vec << M.translation(), Quaterniond(M.rotation()).coeffs();
    return pose_vec;
}

int main()
{
    // Config init
    std::string wolf_root = _WOLF_OBJECTSLAM_ROOT_DIR;
    std::string yaml_path = wolf_root + "/demos/yaml/cosyslam_imu.yaml";
    YAML::Node config = YAML::LoadFile(yaml_path);

    // Retrieve parameters from config file
    std::string bagfile = config["bag"].as<std::string>();
    bool unfix_extrinsic = config["unfix_extrinsic"].as<bool>();
    std::vector<double> i_p_c_vec = config["i_p_c"].as<std::vector<double>>();
    std::vector<double> i_q_c_vec = config["i_q_c"].as<std::vector<double>>();
    Eigen::Map<Vector3d> i_p_c(i_p_c_vec.data());
    Eigen::Map<Vector4d> i_q_c(i_q_c_vec.data());

    // Rosbag opening
    rosbag::Bag bag;
    std::string bagpath = wolf_root + bagfile;
    bag.open(bagpath, rosbag::bagmode::Read);

    ////////////////////////////
    // setup the wolf problem //
    ////////////////////////////

    ProblemPtr problem = Problem::create("POV", 3);
    SolverCeresPtr solver = std::make_shared<SolverCeres>(problem);
    solver->getSolverOptions().max_num_iterations = 200;

    // Cosypose processor
    ParamsSensorPosePtr params_sp = std::make_shared<ParamsSensorPose>();
    Vector7d extr_cam; extr_cam << i_p_c, i_q_c;
    auto sen = problem->installSensor("SensorPose", "sensor_pose", extr_cam, params_sp);
    if (unfix_extrinsic){
        sen->unfixExtrinsics();
    }
    else{
        sen->fixExtrinsics();
    }
    auto prc = problem->installProcessor("ProcessorTrackerLandmarkObject", "objects_wrapper", "sensor_pose", wolf_root + "/demos/yaml/processor_tracker_landmark_object.yaml");
    auto prc_obj = std::static_pointer_cast<ProcessorTrackerLandmarkObject>(prc);

    // IMU processor
    SensorBasePtr sen_imu_base = problem->installSensor("SensorImu", "SenImu",(Vector7d() << 0, 0, 0, 0, 0, 0, 1).finished(), wolf_root + "/demos/yaml/realsense_imu.yaml");
    SensorImuPtr sen_imu = std::static_pointer_cast<SensorImu>(sen_imu_base);
    ProcessorBasePtr proc_imu_base = problem->installProcessor("ProcessorImu", "imu pre-integrator", "SenImu", wolf_root + "/demos/yaml/processor_imu.yaml");
    ProcessorImuPtr proc_imu = std::static_pointer_cast<ProcessorImu>(proc_imu_base);
    
    // Rosbag topics target
    std::vector<std::string> topics;
    topics.push_back(std::string("/imu"));
    topics.push_back(std::string("/cosypose"));
    rosbag::View view(bag, rosbag::TopicQuery(topics));
    ros::Time bag_begin_time = view.getBeginTime();
    TimeStamp init_ts(bag_begin_time.toSec());
    int frame_id = 0;

    /////////////////////////////////
    // Initial orientation with IMU//
    /////////////////////////////////

    Vector3d acc_mean;
    int N_iter = 10;
    foreach(rosbag::MessageInstance const m, view)
    {
        if (frame_id < N_iter){
            sensor_msgs::Imu::ConstPtr cosyIMU = m.instantiate<sensor_msgs::Imu>();
            if (cosyIMU != nullptr){
                frame_id++;
                Vector3d acc;
                acc << cosyIMU->linear_acceleration.x,
                    cosyIMU->linear_acceleration.y,
                    cosyIMU->linear_acceleration.z;
                acc_mean = acc_mean + acc;
            }
        }
    }
    acc_mean = acc_mean/N_iter;

    // we use the rodriguez formula to align the acceleration to g
    Vector3d w_g; w_g << 0,0,-9.81;
    Vector3d w_g_norm = w_g/w_g.norm();
    Vector3d acc_norm = acc_mean/acc_mean.norm();
    Vector3d v = (-w_g_norm).cross(acc_norm);
    double c = (-w_g_norm).dot(acc_norm);

    Matrix3d v_skew;
    v_skew << 0, -v(2), v(1),
        v(2), 0, -v(0),
        -v(1), v(0), 0;
    Matrix3d b_R_w = Matrix3d::Identity() + v_skew + (v_skew*v_skew)/(1+c);
    Vector4d q_init = Quaterniond(b_R_w.transpose()).coeffs();


    ///////////////
    // SLAM loop //
    ///////////////

    // Variables init
    Matrix6d cov;
    Vector6d sig;
    Vector7d poseVec;
    ObjectDetections dets;
    int number_of_kf = 0;
    int counter = 0;
    frame_id = 0;

    VectorComposite x0("POV", { Vector3d::Zero(), q_init, Vector3d::Zero()});
    VectorComposite sig0("POV", {Vector3d::Ones()*0.001, Vector3d::Ones()*0.1, Vector3d::Ones()*0.05});
    FrameBasePtr F1 = problem->setPriorFactor(x0, sig0, 0.0, 0.1);

    foreach(rosbag::MessageInstance const m, view)
    {
        std::cout << "--------FRAME " << frame_id << "-------------" << std::endl;

        wolf_ros_objectslam::CosyObjectArray::ConstPtr cosyArray = m.instantiate<wolf_ros_objectslam::CosyObjectArray>();
        if (cosyArray != nullptr){
            for (auto object : cosyArray->objects){
                poseVec << object.pose.position.x,
                    object.pose.position.y,
                    object.pose.position.z,
                    object.pose.orientation.x,
                    object.pose.orientation.y,
                    object.pose.orientation.z,
                    object.pose.orientation.w;

                // default covariance value : the real value is computed in the processor
                sig << 0.02, 0.02, 0.02, 0.05, 0.05, 0.05;
                cov = sig.array().matrix().asDiagonal();
                ObjectDetection det = {.measurement = poseVec, .detection_score = object.score, .meas_cov = cov, .object_type = object.name};
                dets.push_back(det);
            }
            TimeStamp ts(m.getTime().toSec());
            CaptureObjectPtr cap_obj = std::make_shared<CaptureObject>(ts - init_ts, sen, dets);
            cap_obj->process();
            dets.clear();
        }

        sensor_msgs::Imu::ConstPtr cosyIMU = m.instantiate<sensor_msgs::Imu>();
        if (cosyIMU != nullptr){
            Vector6d acc_gyr_meas;
            acc_gyr_meas << cosyIMU->linear_acceleration.x,
                cosyIMU->linear_acceleration.y,
                cosyIMU->linear_acceleration.z,
                cosyIMU->angular_velocity.x,
                cosyIMU->angular_velocity.y,
                cosyIMU->angular_velocity.z;
            Matrix6d acc_gyr_cov = sen_imu->getNoiseCov();
            TimeStamp ts(m.getTime().toSec());
            CaptureImuPtr cap_imu = std::make_shared<CaptureImu>(ts - init_ts, sen_imu, acc_gyr_meas, acc_gyr_cov);
            cap_imu->process();
        }

        // Solve if a KF was voted
        number_of_kf = problem->getTrajectory()->getFrameMap().size();
        if (number_of_kf > counter){
            std::string report = solver->solve(SolverCeres::ReportVerbosity::FULL);
            counter++;
        }
        frame_id ++;
    }

    bag.close();

    //////////////////////
    // results analysis //
    //////////////////////

    std::string report = solver->solve(SolverCeres::ReportVerbosity::FULL);
    std::cout << report << std::endl;

    problem->print(4,true,true,true);

    // Save the results
    std::fstream file_res;
    std::fstream file_map;
    std::fstream file_init;
    std::string res_filename = wolf_root + "/demos/result.csv";
    std::string map_filename = wolf_root + "/demos/map.csv";
    std::string init_filename = wolf_root + "/demos/init.csv";
    file_init.open(init_filename, std::fstream::out);
    file_res.open(res_filename, std::fstream::out);
    file_map.open(map_filename, std::fstream::out);

    // Write headers for csv files
    std::string header_res = "id,t,px,py,pz,qx,qy,qz,qw,\n";
    file_res << header_res;
    std::string header_map = "id,px,py,pz,qx,qy,qz,qw,\n";
    file_map << header_map;
    std::string header_init = "id,ts,object_id,pose,image_coord,\n";
    file_init << header_init;

    VectorComposite state_est;
    VectorComposite state_lmk;
    CaptureBasePtr cap_cosy;
    counter = 0;

    for (auto& elt: problem->getTrajectory()->getFrameMap()){ 
        TimeStamp t = elt.first; 
        FrameBasePtr kf = elt.second; 
        state_est = kf->getState();
        file_res << counter << "," 
                 << t << ","
                 << state_est['P'](0) << ","
                 << state_est['P'](1) << ","
                 << state_est['P'](2) << "," 
                 << state_est['O'](0) << "," 
                 << state_est['O'](1) << "," 
                 << state_est['O'](2) << "," 
                 << state_est['O'](3) << "\n";

        // World to robot
        Vector7d pose_rob;
        pose_rob << state_est['P'], state_est['O']; 
        Isometry3d w_M_c = posevec_to_isometry(pose_rob);


        int good_lmk = 0;
        for (auto& lmk: problem->getMap()->getLandmarkList()){
            auto lmk_obj = std::dynamic_pointer_cast<LandmarkObject>(lmk);

            // World to lmk
            state_lmk = lmk_obj->getState();
            Vector7d pose_lmk;
            pose_lmk << state_lmk['P'], state_lmk['O'];
            Isometry3d w_M_lmk = posevec_to_isometry(pose_lmk);

            // Cam to lmk
            Isometry3d c_M_lmk =  w_M_c.inverse() * w_M_lmk;

            // coordinates of the object in the image
            Matrix3d intrinsic = lmk_obj->getIntrinsic();
            Vector3d c_t_lmk = c_M_lmk.translation();
            Vector3d image_coord = intrinsic * c_t_lmk;
            image_coord = image_coord/image_coord(2);
            
            file_init << 3*counter << ","
                    << t << ","
                    << lmk_obj->getObjectType() << ","
                    << c_M_lmk.matrix().row(0) << " "
                    << c_M_lmk.matrix().row(1) << " "
                    << c_M_lmk.matrix().row(2) << " "
                    << c_M_lmk.matrix().row(3) << "  ,"
                    << image_coord(0) << " "
                    << image_coord(1) << " "
                    << "\n";

            good_lmk++;
        }
        counter++; 
    }

    counter = 0;
    for (auto& elt: problem->getMap()->getLandmarkList()){
        state_est = elt->getState();
        if (elt->getConstrainedByList().size() < 15){
            continue;
        }
        file_map << counter << ","
                 << state_est['P'](0) << ","
                 << state_est['P'](1) << ","
                 << state_est['P'](2) << "," 
                 << state_est['O'](0) << "," 
                 << state_est['O'](1) << "," 
                 << state_est['O'](2) << "," 
                 << state_est['O'](3) << "\n";
        counter++; 
    } 
    file_res.close();
    file_map.close();
    file_init.close();

    return 0;
}