#include "openvslam/data/keypoint.h"
#include "openvslam/data/map_database.h"
#include "openvslam/publish/frame_publisher.h"

#include <spdlog/spdlog.h>
#include <opencv2/imgproc.hpp>

namespace openvslam {
namespace publish {

frame_publisher::frame_publisher(const std::shared_ptr<config>& cfg, data::map_database* map_db,
                                 const unsigned int img_width)
    : cfg_(cfg), map_db_(map_db), img_width_(img_width),
      img_(cv::Mat(480, img_width_, CV_8UC3, cv::Scalar(0, 0, 0))) {
    spdlog::debug("CONSTRUCT: publish::frame_publisher");
}

frame_publisher::~frame_publisher() {
    spdlog::debug("DESTRUCT: publish::frame_publisher");
}

cv::Mat frame_publisher::draw_frame(const bool draw_text) {
    cv::Mat img;
    tracker_state_t tracking_state;
    std::map<int, data::keypoint> init_keypts;
    std::map<int, std::pair<data::keypoint, data::keypoint>> init_matches;
    data::keypoint_container curr_keypts;
    double elapsed_ms;
    bool mapping_is_enabled;
    std::map<int, bool> is_tracked;

    // copy to avoid memory access conflict
    {
        std::lock_guard<std::mutex> lock(mtx_);
        // TODO pali: optimization possible?
        img_.copyTo(img);

        tracking_state = tracking_state_;

        // copy tracking information
        if (tracking_state == tracker_state_t::Initializing) {
            init_keypts = init_keypts_;
            init_matches = init_matches_;
        }

        curr_keypts = curr_keypts_;

        elapsed_ms = elapsed_ms_;

        mapping_is_enabled = mapping_is_enabled_;

        is_tracked = is_tracked_;
    }

    // resize image
    const float mag = (img_width_ < img_.cols) ? static_cast<float>(img_width_) / img.cols : 1.0;
    if (mag != 1.0) {
        cv::resize(img, img, cv::Size(), mag, mag, cv::INTER_NEAREST);
    }

    // to draw COLOR information
    if (img.channels() < 3) {
        cvtColor(img, img, cv::COLOR_GRAY2BGR);
    }

    // draw keypoints
    unsigned int num_tracked = 0;
    switch (tracking_state) {
        case tracker_state_t::Initializing: {
            num_tracked = draw_initial_points(img, init_matches, mag);
            break;
        }
        case tracker_state_t::Tracking: {
            num_tracked = draw_tracked_points(img, curr_keypts, is_tracked, mapping_is_enabled, mag);
            break;
        }
        default: {
            break;
        }
    }

    if (draw_text) {
        // draw tracking info
        draw_info_text(img, tracking_state, num_tracked, elapsed_ms, mapping_is_enabled);
    }

    return img;
}

unsigned int frame_publisher::draw_initial_points(cv::Mat &img,
                                                  const std::map<int, std::pair<data::keypoint, data::keypoint>> &init_matches,
                                                  const float mag) const {
    unsigned int num_tracked = 0;

    for (auto match : init_matches) {
        const cv::Point_<float> &init_point = match.second.first.get_cv_keypoint().pt * mag;
        cv::circle(img, init_point, 2, mapping_color_, -1);

        const cv::Point_<float> &current_frame_point = match.second.second.get_cv_keypoint().pt * mag;
        cv::circle(img, current_frame_point, 2, mapping_color_, -1);

        cv::line(img, init_point, current_frame_point, mapping_color_);

        ++num_tracked;
    }

    return num_tracked;
}

unsigned int frame_publisher::draw_tracked_points(cv::Mat& img, const data::keypoint_container &curr_keypts,
                                                  const std::map<int, bool>& is_tracked, const bool mapping_is_enabled,
                                                  const float mag) const {
    constexpr float radius = 5;

    unsigned int num_tracked = 0;

    for (const auto &keypoint : curr_keypts) {
        if (is_tracked.find(keypoint.first) == is_tracked.end()) {
            continue;
        }

        const cv::Point2f pt_begin{keypoint.second.get_cv_keypoint().pt.x * mag - radius, keypoint.second.get_cv_keypoint().pt.y * mag - radius};
        const cv::Point2f pt_end{keypoint.second.get_cv_keypoint().pt.x * mag + radius, keypoint.second.get_cv_keypoint().pt.y * mag + radius};

        if (mapping_is_enabled) {
            cv::rectangle(img, pt_begin, pt_end, mapping_color_);
            cv::circle(img, keypoint.second.get_cv_keypoint().pt * mag, 2, mapping_color_, -1);
        }
        else {
            cv::rectangle(img, pt_begin, pt_end, localization_color_);
            cv::circle(img, keypoint.second.get_cv_keypoint().pt * mag, 2, localization_color_, -1);
        }

        ++num_tracked;
    }

    return num_tracked;
}

void frame_publisher::draw_info_text(cv::Mat& img, const tracker_state_t tracking_state, const unsigned int num_tracked,
                                     const double elapsed_ms, const bool mapping_is_enabled) const {
    // create text info
    std::stringstream ss;
    switch (tracking_state) {
        case tracker_state_t::NotInitialized: {
            ss << "WAITING FOR IMAGES ";
            break;
        }
        case tracker_state_t::Initializing: {
            ss << "INITIALIZE | ";
            ss << "Frame " << frame_id << ", ";
            ss << "KP: " << curr_keypts_.size() << ", ";
            ss << "TKP: " << num_tracked << ", ";
            ss << "track time: " << std::fixed << std::setprecision(0) << elapsed_ms << "ms";
            break;
        }
        case tracker_state_t::Tracking: {
            ss << (mapping_is_enabled ? "MAPPING | " : "LOCALIZATION | ");
            ss << "Frame " << frame_id << ", ";
            ss << "KF: " << map_db_->get_num_keyframes() << ", ";
            ss << "LM: " << map_db_->get_num_landmarks() << ", ";
            ss << "TLM: " << is_tracked_.size() << ", ";
            ss << "KP: " << curr_keypts_.size() << ", ";
            ss << "track time: " << std::fixed << std::setprecision(0) << elapsed_ms << "ms";
            break;
        }
        case tracker_state_t::Lost: {
            ss << "LOST | ";
            ss << "Frame " << frame_id << ", ";
            ss << "track time: " << std::fixed << std::setprecision(0) << elapsed_ms << "ms";
            break;
        }
    }

    int baseline = 0;
    const cv::Size text_size = cv::getTextSize(ss.str(), cv::FONT_HERSHEY_PLAIN, 1, 1, &baseline);

    // create text area
    constexpr float alpha = 0.6;
    cv::Mat overlay(img.rows, img.cols, img.type());
    img.copyTo(overlay);
    cv::rectangle(overlay, cv::Point2i{0, img.rows - text_size.height - 10}, cv::Point2i{img.cols, img.rows}, cv::Scalar{0, 0, 0}, -1);
    cv::addWeighted(overlay, alpha, img, 1 - alpha, 0, img);
    // put text
    cv::putText(img, ss.str(), cv::Point(5, img.rows - 5), cv::FONT_HERSHEY_PLAIN, 1, cv::Scalar{255, 255, 255}, 1, 8);
}

void frame_publisher::update(tracking_module* tracker) {
    std::lock_guard<std::mutex> lock(mtx_);

    tracker->img_gray_.copyTo(img_);

    curr_keypts_ = tracker->curr_frm_.keypts_;
    elapsed_ms_ = tracker->elapsed_ms_;
    mapping_is_enabled_ = tracker->get_mapping_module_status();
    tracking_state_ = tracker->last_tracking_state_;
    frame_id = tracker->curr_frm_.id_;

    switch (tracking_state_) {
        case tracker_state_t::Initializing: {
            init_keypts_ = tracker->get_initial_keypoints();
            init_matches_ = tracker->get_initial_matches();
            break;
        }
        case tracker_state_t::Tracking: {
            is_tracked_.clear();
            for (auto &lm : tracker->curr_frm_.landmarks_) {
                if (lm.second->is_outlier()) {
                    continue;
                }

                if (0 < lm.second->num_observations()) {
                    is_tracked_[lm.first] = true;
                }
            }
            break;
        }
        default: {
            break;
        }
    }
}

} // namespace publish
} // namespace openvslam
