/**
 * 激光定位接口，发布定位位姿
 * 
 * ＠2020年5月22日　成都河狸智能科技有限责任公司.保留所有权利.
 * 
 * 作者: lumin
 */

#include <ros/ros.h>
#include <ros/package.h>
#include <sensor_msgs/PointCloud.h>
#include <geometry_msgs/Point32.h>
#include <geometry_msgs/Pose2D.h>
#include <sensor_msgs/LaserScan.h>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/MapMetaData.h>
#include <nav_msgs/GetMap.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <tf/transform_broadcaster.h>
#include <tf/transform_listener.h>
#include <space_lib/space_lib.h>

#include <std_msgs/Float32.h>

#include <yaml-cpp/yaml.h>
#include <key/key.h>
#include <atomic>
#include <thread>
#include <mutex>
#include <fstream>
#include <sys/io.h>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>

#include "base_msgs/global_position.h"

#include <location_base/location_base.h>

namespace {

using namespace space_location::sensor;
using namespace space_location;

ros::NodeHandle* nh_;

ros::Subscriber location_msgs_sub_;
ros::Subscriber odom_data_sub_;
ros::Publisher map_pose_pub_;
ros::Publisher map_pub_;

std::vector<ros::Subscriber> subs_;

std::thread matchThread_;
std::atomic<bool> is_calibration_;
std::atomic<bool> is_located_;

std::string odom_frame_id_ = "odom";
std::string base_frame_id_ = "base_link";
std::string global_frame_id_ = "map";
ros::Duration transform_tolerance_;
tf::Transform tf_transform_;
std::unique_ptr<tf::TransformBroadcaster> trans_tfb_;
std::unique_ptr<tf::TransformListener> tf_;

float linear_search_; // 线性方向搜索
float angular_search_; // 角度方向搜索
float angular_perturbation_step_size_; // 角度搜索精度
float init_linear_search_;
float init_angular_search_;
float init_angular_step_;

float occupied_space_weight_; // 点云匹配权重
float translation_weight_;  // 里程计位移权重
float rotation_weight_; // 旋转权重
float self_stab_weight_; // 自稳定旋转权重

std::map<std::string, Eigen::Matrix<float, 4, 4>> transforms_;
std::mutex start_location_mutex_;
Eigen::Vector3f init_location_pose_;

#define START_LOCATION_LOCK() std::lock_guard<std::mutex> lockGuard(start_location_mutex_);
#define MAP_IDX(sx, i, j) ((sx) * (j) + (i))

template <typename T>
static void operator>>(const YAML::Node &node, T &i)
{
    i = node.as<T>();
}

struct PointCloudWeight {
    PointCloud point_cloud;
    std::vector<float> weights;
    ros::Time stamp;
};

static void __laserDataCallback(const sensor_msgs::LaserScanConstPtr& msg);
static void __odomDataCallback(const nav_msgs::OdometryConstPtr& msg);
static void __pubMapPose(const float& x, const float& y, const float& yaw);
static void __tfMapToOdom(const float& x, const float& y, const float& yaw, ros::Time stamp);
static void __mapPub(std::string map_file, std::string config_path);

static bool __readLocationConfig() 
{
    std::string path = space::home() + "configure/location_config.json";

    Json::Value c = space::readConfig(path);
    Json::Value compair = c;

    space::jsonValueTo(c["linear_search"], linear_search_) ?
        false : (c["linear_search"] = linear_search_ = 0.5);
    space::jsonValueTo(c["angular_search"], angular_search_) ?
        false : (c["angular_search"] = angular_search_ = 0.67);
    space::jsonValueTo(c["angular_perturbation_step"], angular_perturbation_step_size_) ?
        false : (c["angular_perturbation_step"] = angular_perturbation_step_size_ = 0.017);
    space::jsonValueTo(c["init_linear_search"], init_linear_search_) ?
        false : (c["init_linear_search"] = init_linear_search_ = 5.0);
    space::jsonValueTo(c["init_angular_search"], init_angular_search_) ?
        false : (c["init_angular_search"] = init_angular_search_ = M_PI);
    space::jsonValueTo(c["init_angular_step"], init_angular_step_) ?
        false : (c["init_angular_step"] = init_angular_step_ = 0.108);

    space::jsonValueTo(c["occupied_space_weight"], occupied_space_weight_) ?
        false : (c["occupied_space_weight"] = occupied_space_weight_ = 1.0);
    space::jsonValueTo(c["translation_weight"], translation_weight_) ?
        false : (c["translation_weight"] = translation_weight_ = 2.0);
    space::jsonValueTo(c["rotation_weight"], rotation_weight_) ?
        false : (c["rotation_weight"] = rotation_weight_ = 3.0);
    space::jsonValueTo(c["self_stab_weight"], self_stab_weight_) ?
        false : (c["self_stab_weight"] = self_stab_weight_ = 1.0);
    
    if (compair != c) {
        space::writeConfig(path, c);
    }
    return true;
}

static Json::Value __readLaserConfig() 
{
    std::string path = space::home() + "configure/laser_config.json";

    Json::Value c = space::readConfig(path);
    Json::Value laser;
    
    space::jsonValueTo(c, laser);
    
    return laser;
}

static bool __CacluLaserToBaseLinkTransform() 
{
    transforms_.clear();

    Json::Value lasers = __readLaserConfig();

    for (int i = 0;i < lasers.size();i++) {
        float tf_x, tf_y, tf_z, tf_roll, tf_pitch, tf_yaw;
        std::string name;

        bool ok =   space::jsonValueTo(lasers[i]["tf_x"], tf_x) &&
                    space::jsonValueTo(lasers[i]["tf_y"], tf_y) &&
                    space::jsonValueTo(lasers[i]["tf_z"], tf_z) &&
                    space::jsonValueTo(lasers[i]["tf_roll"], tf_roll) &&
                    space::jsonValueTo(lasers[i]["tf_pitch"], tf_pitch) &&
                    space::jsonValueTo(lasers[i]["tf_yaw"], tf_yaw) &&
                    space::jsonValueTo(lasers[i]["name"], name);
        if (ok) {
            Eigen::AngleAxisd rollAngle(Eigen::AngleAxisd(tf_roll, Eigen::Vector3d::UnitX()));
            Eigen::AngleAxisd pitchAngle(Eigen::AngleAxisd(tf_pitch, Eigen::Vector3d::UnitY()));
            Eigen::AngleAxisd yawAngle(Eigen::AngleAxisd(tf_yaw, Eigen::Vector3d::UnitZ()));
            
            Eigen::Matrix3f rotation_matrix;
            Eigen::Matrix<float, 4, 4> transform;

            rotation_matrix = (yawAngle * pitchAngle * rollAngle).cast<float>();
            transform.rightCols<1>() << tf_x, tf_y, tf_z, 1.0;
            transform.block<3, 3>(0, 0, 3, 3) = rotation_matrix;
            transform.bottomRows<1>() << 0, 0, 0, 1.0;

            transforms_.emplace(std::make_pair(name, transform));
        }
    }
    if (!transforms_.empty()) {
        return true;
    }
    return false;
}

static void __match(PointCloud point_cloud, std::vector<float> weights, ros::Time stamp) 
{
    if (!is_located_) {
        if (location_base::initLocation(
                point_cloud, 
                init_location_pose_(0),
                init_location_pose_(1), 
                init_location_pose_(2), 
                init_linear_search_, 
                init_angular_search_, 
                init_angular_step_) == location_base::State::OK) {
            is_located_ = true; 
        }
    }else {
        if (!is_calibration_) {
            location_base::Pose2D pose_2d;
            
            location_base::State state = location_base::scanMatchPrediction(point_cloud, weights, pose_2d);
            __pubMapPose(pose_2d.x, pose_2d.y, pose_2d.yaw);
            __tfMapToOdom(pose_2d.x, pose_2d.y, pose_2d.yaw, stamp);
        }else {
            double estimate[7];
            location_base::Pose2D pose_2d;
            location_base::State state = location_base::clibreation(point_cloud, weights, estimate, pose_2d);

            if (state == location_base::State::CLIBREATION_OK) {
                LOG(INFO)<<"estimate 0: "<<estimate[0];
                LOG(INFO)<<"estimate 1: "<<estimate[1];
                LOG(INFO)<<"estimate 2: "<<estimate[2];

                Eigen::Quaterniond quaternion(estimate[6], estimate[3], estimate[4], estimate[5]);

                LOG(INFO)<<"z: "<<quaternion.matrix().eulerAngles(2, 1, 0)(0);
            }else {
                // LOG(INFO)<<"state: "<<state;
            }
        }
    }
}

static void __initLocation(const base_msgs::global_position position) 
{
    START_LOCATION_LOCK()
    enum LOCMODE {
        QUIT = 0,
        FAST,
        EXACT,
        CALIBRATION
    };
    if (__readLocationConfig()) {
        location_base::setCeresWeight(
            occupied_space_weight_, translation_weight_, rotation_weight_, self_stab_weight_
            );
        location_base::setSearchParam(
            linear_search_, angular_search_, angular_perturbation_step_size_
        );
    }
    {
        init_location_pose_ = Eigen::Vector3f(
            position.init_position.pose.position.x,
            position.init_position.pose.position.y,
            tf::getYaw(position.init_position.pose.orientation)
        );
        is_located_ = false;
        if (position.mode == QUIT) {   //退出定位
            location_base::quitLocation();
            
            return;
        }else if (position.mode == FAST) { //指定位置启动
            location_base::setInitPose(
                init_location_pose_(0),
                init_location_pose_(1),
                init_location_pose_(2)
            );
            __CacluLaserToBaseLinkTransform();
            is_located_ = true;
        }else if (position.mode == EXACT) {
            __CacluLaserToBaseLinkTransform();
        }else if (position.mode == CALIBRATION || is_calibration_) {
            transforms_.clear();
            Eigen::Matrix<float, 4, 4> transform;

            transform.setIdentity();
            transforms_.emplace(std::make_pair("laser", transform));
        }else {
            return;
        }
        if (!location_base::initMap(position.map_path + ".png", position.yaml_path + ".json")) {   // 读取地图失败，退出
            is_located_ = false;
            LOG(INFO)<<"LOC[ERROR]: 读取地图失败!!!: "<<position.map_path;
            return;
        }else {
            __mapPub(position.map_path + ".png", position.yaml_path + ".json");
        }
        Json::Value lasers = __readLaserConfig();

        subs_.clear();
        for (int i = 0;i < lasers.size();i++) {
            std::string topic = "";

            if (space::jsonValueTo(lasers[i]["topic"], topic)) {
                subs_.push_back(nh_->subscribe(topic, 1, &__laserDataCallback));
            }
        }
        odom_data_sub_ = nh_->subscribe<nav_msgs::Odometry>("odom", 1, &__odomDataCallback);
    }
}

PointCloud __transformRangeData2PointCloudAndWeights(const sensor_msgs::LaserScan& laser_data, 
                                                    std::vector<float>& weights, Eigen::Matrix<float, 4, 4> transform) 
{
    PointCloud point_cloud;
    float theta = laser_data.angle_max;
    float sum = 0;
    float offseta = laser_data.angle_max + laser_data.angle_min;

    point_cloud.clear();
    for (const float& range : laser_data.ranges) {
        theta -= laser_data.angle_increment;
        
        if (range > laser_data.range_max) continue;
        if (range < laser_data.range_min) continue;
        if (theta < laser_data.angle_min) continue;
        if (theta > laser_data.angle_max) continue;
        if (std::isnan(range)) continue;

        const Eigen::Matrix<float, 4, 1> point(range * cos(theta - offseta),
                                                -range * sin(theta - offseta),
                                                0.0,
                                                1.0);
        // TimePointCloud time_point;
        
        // time_point.position = Eigen::Vector3f(base_link_coord(0),
        //                                       base_link_coord(1), 
        //                                       0);
        // time_point.time = std::chrono::system_clock::now();
        // point_cloud.push_back(time_point);
    }
    float revers = 1.0 / sum;

    for (float& value : weights) {
        value *= revers;
    }
    Eigen::MatrixXf map;

    map.setZero(3000, 3000);
    map(0, 0);

    return point_cloud;
}

static void __laserDataCallback(const sensor_msgs::LaserScanConstPtr& msg)
{
    std::map<std::string, Eigen::Matrix<float, 4, 4>>::iterator it = 
        transforms_.find(msg->header.frame_id);

    if (it == transforms_.end()) {
        return;
    }
    static std::map<std::string, PointCloudWeight> pointcloud_weights;
    std::vector<float> weights;
    PointCloud match_cloud = __transformRangeData2PointCloudAndWeights(*msg, weights, it->second);

    pointcloud_weights[msg->header.frame_id] = PointCloudWeight({match_cloud, weights, msg->header.stamp});
    if (pointcloud_weights.size() >= transforms_.size()) {
        weights.clear();
        match_cloud.clear();

        for (std::map<std::string, PointCloudWeight>::iterator pt = pointcloud_weights.begin();pt != pointcloud_weights.end();pt++) {
            if ((ros::Time::now() - pt->second.stamp).toSec() < 0.5) {
                weights.insert(weights.end(), pt->second.weights.begin(), pt->second.weights.end());
                match_cloud.insert(match_cloud.end(), pt->second.point_cloud.begin(), pt->second.point_cloud.end());
            }
        }
        pointcloud_weights.clear();
    }else {
        return;
    }
    ros::Time stamp = msg->header.stamp;
    if (!matchThread_.joinable()) {
        matchThread_ = std::thread([match_cloud, weights, stamp] {
            START_LOCATION_LOCK()
            __match(match_cloud, weights, stamp);
            usleep(1000 * 5);
            matchThread_.detach();
        });
    }
}

static void __mapPub(std::string map_file, std::string config_path)
{
    static nav_msgs::GetMap::Response map_res;
    cv::Mat image;
    float x_origin, y_origin, resolution, occupied_thresh, free_thresh;

    image = cv::imread(map_file, cv::IMREAD_GRAYSCALE);

    tf::Quaternion q;
    q.setRPY(0, 0, 0);

    Json::Value c = space::readConfig(config_path);
    
    bool ok =   space::jsonValueTo(c["x_origin"], x_origin) &&
                space::jsonValueTo(c["y_origin"], y_origin) &&
                space::jsonValueTo(c["resolution"], resolution) && 
                space::jsonValueTo(c["occupied_thresh"], occupied_thresh) &&
                space::jsonValueTo(c["free_thresh"], free_thresh);

    map_res.map.info.map_load_time = ros::Time::now();
    map_res.map.header.frame_id = "map";
    map_res.map.header.stamp = ros::Time::now();
    map_res.map.info.width = image.cols;
    map_res.map.info.height = image.rows;
    map_res.map.info.resolution = resolution;
    map_res.map.info.origin.position.x = -x_origin;
    map_res.map.info.origin.position.y = -y_origin;
    map_res.map.info.origin.position.z = 0.0;
    map_res.map.info.origin.orientation.x = q.x();
    map_res.map.info.origin.orientation.y = q.y();
    map_res.map.info.origin.orientation.z = q.z();
    map_res.map.info.origin.orientation.w = q.w();
    map_res.map.data.resize(map_res.map.info.width * map_res.map.info.height);

    unsigned char *p;
    unsigned char value;
    unsigned char pix = 0;

    for (size_t x = 0; x < image.rows; ++x) {
        p = image.ptr<uchar>(x);
        for(size_t y = 0; y < image.cols; ++y) {
            pix = *(p + y);
            float cm = float(pix) / 255.f;

            if (cm < 1 - occupied_thresh) {
                value = (1 - cm) * 255;
            }else if (cm > 1 - free_thresh) {
                value = 0;
            }else {
                value = 0;
            }

            map_res.map.data[MAP_IDX(map_res.map.info.width, y, map_res.map.info.height - x - 1)] = value;
        }
    }

    LOG(INFO)<<"Read a "<<map_res.map.info.width<<" * "<<map_res.map.info.height
                <<" @ "<<map_res.map.info.resolution<<" Map!";
    map_pub_.publish(map_res.map);
}

static void __tfMapToOdom(const float& x, const float& y, const float& yaw, ros::Time stamp) 
{
    tf::Stamped<tf::Pose> odom_to_map;
    ros::Time copy_stamp = stamp;
    
    transform_tolerance_.fromSec(0.3);
    try {
        /*tmp_tf,从map原点看base_link的位置*/
        tf::Transform tmp_tf(tf::createQuaternionFromYaw(yaw),
                                tf::Vector3(x, y, 0.0));
        /*tmp_tf.inverse(),以为Mp为原点，map原点的位置，
            就是在base_link坐标系下map 原点的位置*/
        tf::Stamped<tf::Pose> tmp_tf_stamped(tmp_tf.inverse(),
                                                copy_stamp, //tf_time
                                                base_frame_id_);
        /*进行 base_link坐标系下的点转换到odom坐标系，
            也就是把map原点转换到odom坐标系下，
            等于从odom原点看map原点的位置。放到tf_transform_变量里面*/
        tf_->transformPose(odom_frame_id_, tmp_tf_stamped, odom_to_map);
    }catch (tf::TransformException) {
        // LOG(INFO)<<"Location Failed to subtract base to odom transform";
        return;
    }

    tf_transform_ = tf::Transform(tf::Quaternion(odom_to_map.getRotation()),
                                    tf::Point(odom_to_map.getOrigin()));

    ros::Time transform_expiration = (copy_stamp + transform_tolerance_);

    /*发布tf转换关系，latest_tf_.inverse() 得到的是从map原点看odom原点的位置，也就是 odom->map的转换关系*/
    tf::StampedTransform tmp_tf_stamped(tf_transform_.inverse(),
                                        transform_expiration,
                                        global_frame_id_, odom_frame_id_);

    trans_tfb_->sendTransform(tmp_tf_stamped);
}

inline static void __quitLocation() 
{
    for (auto sub : subs_) {
        sub.shutdown();
    }
    odom_data_sub_.shutdown();
    location_base::quitLocation();
}

static void __odomDataCallback(const nav_msgs::OdometryConstPtr& msg) 
{
    // return;
    location_base::addOdomData(msg->pose.pose.position.x, msg->pose.pose.position.y,
                tf::getYaw(msg->pose.pose.orientation));
}

static void __locationMsgsCallback(const base_msgs::global_positionConstPtr& msg) 
{
    START_LOCATION_LOCK()
    
    __quitLocation();

    std::thread monitorThread(&__initLocation, *msg);
    monitorThread.detach();
}

static void __pubMapPose(const float& x, const float& y, const float& yaw) 
{
    geometry_msgs::Pose2D map_pose;

    map_pose.x = x;
    map_pose.y = y;
    map_pose.theta = yaw;

    map_pose_pub_.publish(map_pose);
}

}

int main(int argc, char** argv) 
{
    ros::init(argc, argv, "space_location");

    is_calibration_.store(false);
    is_located_.store(false);
    
    nh_ = new ros::NodeHandle;
    trans_tfb_ = std::unique_ptr<tf::TransformBroadcaster>(new tf::TransformBroadcaster());
    tf_ = std::unique_ptr<tf::TransformListener>(new tf::TransformListener());
    map_pub_ = nh_->advertise<nav_msgs::OccupancyGrid>("map", 1, true);
    map_pose_pub_ = nh_->advertise<geometry_msgs::Pose2D>("map_pose", 1);
    
    location_msgs_sub_ = 
        nh_->subscribe<base_msgs::global_position>("location_msgs", 1, &__locationMsgsCallback);

    ros::spin();

    delete nh_;
}