#include "registration.h"
#include "core/utility.hpp"
#include "render_v3/vtk_viewer.h"

#include <vtkActor.h>
#include <vtkCellLocator.h>
#include <vtkIterativeClosestPointTransform.h>
#include <vtkLandmarkTransform.h>
#include <vtkMatrix4x4.h>
#include <vtkNamedColors.h>
#include <vtkPolyData.h>
#include <vtkProperty.h>
#include <vtkSmartPointer.h>

#include <queue>
#include <vector>

using namespace vtk_viewer_helper;

struct registration::impl {

    static constexpr auto MIN_REG_POINTS = 3;
    static constexpr auto DIS_LIMIT_1 = 0.5;
    static constexpr auto DIS_LIMIT_2 = 1.0;
    static constexpr auto NAN_VALUE = std::numeric_limits<float>::quiet_NaN();

    enum show_mode {
        CONFIG,
        COLLECTING,
        FINISHED
    };

    struct target_store_type {
        std::string name;
        vtkSmartPointer<vtkPolyData> model;
        vtkSmartPointer<vtkActor> model_actor;
        vtkSmartPointer<vtkCellLocator> model_locator;
        std::unique_ptr<smart_point_sets> all_points;
        std::unique_ptr<smart_point_sets> pending_points;
        std::unique_ptr<smart_point_sets> current_point;
        std::unique_ptr<smart_point_sets> finished_points;
        std::unique_ptr<smart_point_sets> level_points[3];

        std::string target_var_name;
        std::string point_var_name;
        std::string collect_obj_name;
        std::string probe_var_name;

        bool is_finished = false;
        float max_error = 0;
    };

    int cur_target_id = -1;
    std::vector<target_store_type> targets;
    target_store_type *cur_target = nullptr;

    std::unique_ptr<vtk_viewer> viewer;
    bool is_picking = false;
    bool is_collecting = false;
    std::vector<Eigen::Vector3d> collected_points;
    std::queue<std::function<void()>> eq;

    sophiar::local_connection *conn = nullptr;
    vtkSmartPointer<vtkActor> probe_actor;
    float probe_test_error = NAN_VALUE;

    // colors
    vtkColor4d target_color;
    vtkColor4d idle_color;
    vtkColor4d pending_color;
    vtkColor4d current_color;
    vtkColor4d finished_color;
    vtkColor4d level_color[3]; // - 0.5mm - 1.0mm -

    impl() {
        viewer = std::make_unique<vtk_viewer>();

        // preload colors
        vtkNew<vtkNamedColors> colors;
        target_color = colors->GetColor4d("silver");
        idle_color = colors->GetColor4d("white");
        pending_color = colors->GetColor4d("red");
        current_color = colors->GetColor4d("yellow");
        finished_color = colors->GetColor4d("lime");
        level_color[0] = colors->GetColor4d("lime");
        level_color[1] = colors->GetColor4d("yellow");
        level_color[2] = colors->GetColor4d("red");
    }

    void switch_viewer_mode(show_mode mode) {
        assert(cur_target != nullptr);
        viewer->clear_actor();
        viewer->add_actor(cur_target->model_actor);
        viewer->add_actor(probe_actor);
        switch (mode) {
            case CONFIG: {
                viewer->add_actor(cur_target->all_points->get_actor());
                break;
            }
            case COLLECTING: {
                viewer->add_actor(cur_target->pending_points->get_actor());
                viewer->add_actor(cur_target->current_point->get_actor());
                viewer->add_actor(cur_target->finished_points->get_actor());
                break;
            }
            case FINISHED: {
                for (auto &points: cur_target->level_points) {
                    viewer->add_actor(points->get_actor());
                }
                break;
            }
        }
    }

    void add_target(const registration_target &conf) {
        auto &target = targets.emplace_back();
        target.name = conf.name;
        target.model = load_any(conf.model_path);
        target.model_actor = create_actor(target.model);
        target.model_locator = vtkSmartPointer<vtkCellLocator>::New();
        target.model_locator->SetDataSet(target.model);
        target.model_locator->BuildLocator();

        // copy information
        target.target_var_name = conf.target_var_name;
        target.collect_obj_name = conf.collect_obj_name;
        target.point_var_name = conf.collect_var_name;
        target.probe_var_name = conf.probe_var_name;

        // create point sets
        target.all_points = std::make_unique<smart_point_sets>();
        target.pending_points = std::make_unique<smart_point_sets>();
        target.current_point = std::make_unique<smart_point_sets>();
        target.finished_points = std::make_unique<smart_point_sets>();
        for (auto &points: target.level_points) {
            points = std::make_unique<smart_point_sets>();
        }

        // set colors
        target.model_actor->GetProperty()->SetColor(target_color.GetData());
        target.all_points->get_actor()->GetProperty()->SetColor(idle_color.GetData());
        target.pending_points->get_actor()->GetProperty()->SetColor(pending_color.GetData());
        target.current_point->get_actor()->GetProperty()->SetColor(current_color.GetData());
        target.finished_points->get_actor()->GetProperty()->SetColor(finished_color.GetData());
        for (auto k = 0; k < 3; ++k) {
            target.level_points[k]->get_actor()->GetProperty()->SetColor(level_color[k].GetData());
        }
    }

    void change_target() {
        assert(cur_target_id != -1);
        cur_target = &targets[cur_target_id];
        if (cur_target->is_finished) {
            switch_viewer_mode(FINISHED);
        } else {
            switch_viewer_mode(CONFIG);
        }
        viewer->reset_camera();
    }

    bool progress_reg_point() {
        if (!cur_target->current_point->empty()) {
            cur_target->finished_points->add_point(
                    cur_target->current_point->pop_front());
        }
        if (cur_target->pending_points->empty()) return false;
        cur_target->current_point->add_point(
                cur_target->pending_points->pop_front());
        return true;
    }

    void start() {
        assert(cur_target != nullptr);
        assert(!is_collecting);
        is_collecting = true;
        cur_target->is_finished = false;
        conn->mark_variable_disposal(cur_target->point_var_name);
        CALL_CHECK(conn->start_object(cur_target->collect_obj_name));
        collected_points.clear();

        // copy points
        cur_target->all_points->for_each([this](void *, const Eigen::Vector3d &point) {
            cur_target->pending_points->add_point(point);
        });
        progress_reg_point();

        // switch actor
        switch_viewer_mode(COLLECTING);
    }

    void stop() {
        assert(is_collecting);
        is_collecting = false;
        CALL_CHECK(conn->stop_object(cur_target->collect_obj_name));

        // clear point set
        cur_target->pending_points->clear();
        cur_target->current_point->clear();
        cur_target->finished_points->clear();

        // switch actor
        if (cur_target->is_finished) {
            switch_viewer_mode(FINISHED);
        } else {
            switch_viewer_mode(CONFIG);
        }
    }

    auto calc_closest_point(const Eigen::Vector3d &point) {
        Eigen::Vector3d close_point;
        vtkIdType cell_id;
        int sub_id;
        double dis2;
        cur_target->model_locator->FindClosestPoint(
                point.data(), close_point.data(), cell_id, sub_id, dis2);
        return std::make_tuple(close_point, std::sqrt(dis2));
    }

    // return if it needs to continue
    bool calc_result() {
        // prepare landmark
        auto num_points = collected_points.size();
        if (num_points < MIN_REG_POINTS) return true;
        auto source_points = Eigen::Matrix3Xd{3, num_points};
        auto target_points = Eigen::Matrix3Xd{3, num_points};
        for (auto k = 0; k < num_points; ++k) {
            source_points.col(k) = (*cur_target->all_points)[k];
        }
        for (auto k = 0; k < num_points; ++k) {
            target_points.col(k) = collected_points[k];
        }

        // calculate landmark
        auto result = (Eigen::Isometry3d) Eigen::umeyama(source_points, target_points, false);

        // prepare icp
        vtkNew<vtkPoints> icp_points;
        auto landmark_inv = result.inverse();
        for (auto k = 0; k < num_points; ++k) {
            Eigen::Vector3d point = landmark_inv * collected_points[k];
            icp_points->InsertNextPoint(point.data());
        }

        // calculate icp
        vtkNew<vtkIterativeClosestPointTransform> icp;
        vtkNew<vtkPolyData> tmp_poly;
        tmp_poly->SetPoints(icp_points);
        icp->GetLandmarkTransform()->SetModeToRigidBody();
        icp->SetSource(tmp_poly);
        icp->SetTarget(cur_target->model);
        icp->Modified();
        icp->Update();

        // refine result
        Eigen::Isometry3d trans_delta;
        for (auto i = 0; i < 4; ++i)
            for (auto j = 0; j < 4; ++j) {
                trans_delta(i, j) = icp->GetMatrix()->GetElement(i, j);
            }
        result = result * trans_delta.inverse();

        // commit result
        conn->update_transform_variable(cur_target->target_var_name, result);

        // calculate error only when all points are collected
        if (num_points != cur_target->all_points->size()) return true;

        // calculate error
        cur_target->max_error = 0;
        for (auto &points: cur_target->level_points) {
            points->clear();
        }
        auto result_inv = result.inverse();
        for (auto k = 0; k < num_points; ++k) {
            Eigen::Vector3d point = result_inv * collected_points[k];

            // find the closest point
            double dis;
            std::tie(std::ignore, dis) = calc_closest_point(point);

            // update results
            if (dis > cur_target->max_error) {
                cur_target->max_error = dis;
            }
            if (dis < DIS_LIMIT_1) {
                cur_target->level_points[0]->add_point(point);
            } else if (dis < DIS_LIMIT_2) {
                cur_target->level_points[1]->add_point(point);
            } else {
                cur_target->level_points[2]->add_point(point);
            }
        }

        cur_target->is_finished = true;
        return false;
    }

    void try_picking() {
        assert(is_picking);
        auto val = viewer->get_picked_point();
        if (!val.has_value()) return;
        cur_target->all_points->add_point(val.value());
    }

    void try_collect() {
        assert(is_collecting);
        auto val = conn->query_scalarxyz_variable(cur_target->point_var_name);
        if (!val.has_value()) return;
        collected_points.emplace_back(val.value());
        SPDLOG_INFO("Collected point ({}, {}, {}).", val->x(), val->y(), val->z());
        calc_result();
        if (!progress_reg_point()) {
            stop();
        }
    }

    void show() {
        if (ImGui::Begin("Registration Control")) {
            if (is_collecting || is_picking) {
                ImGui::BeginDisabled();
            }
            const char *target_name = (cur_target == nullptr) ? nullptr : cur_target->name.c_str();
            if (ImGui::BeginCombo("Target", target_name)) {
                for (auto k = 0; k < targets.size(); ++k) {
                    bool is_selected = (k == cur_target_id);
                    if (ImGui::Selectable(targets[k].name.c_str(), is_selected)) {
                        cur_target_id = k;
                        eq.emplace([this] { change_target(); });
                    }
                    if (is_selected) {
                        ImGui::SetItemDefaultFocus();
                    }
                }
                ImGui::EndCombo();
            }
            if (is_collecting || is_picking) {
                ImGui::EndDisabled();
            }
            if (cur_target != nullptr) {
                auto point_set = cur_target->all_points.get();
                if (ImGui::CollapsingHeader("Actions")) {
                    if (is_collecting) {
                        ImGui::BeginDisabled();
                    }
                    if (ImGui::Checkbox("Config Points", &is_picking)) {
                        if (is_picking) {
                            eq.emplace([this] {
                                switch_viewer_mode(CONFIG);
                                viewer->start_picking();
                            });
                        } else {
                            eq.emplace([this] { viewer->stop_picking(); });
                        }
                    }
                    if (is_collecting) {
                        ImGui::EndDisabled();
                    }
                    if (!is_picking && point_set->size() >= MIN_REG_POINTS) {
                        ImGui::SameLine();
                        if (!is_collecting) {
                            if (ImGui::Button("Start")) {
                                eq.emplace([this] { start(); });
                            }
                        } else {
                            if (ImGui::Button("Stop")) {
                                eq.emplace([this] { stop(); });
                            }
                        }
                    }

                }
                if (ImGui::CollapsingHeader("Infos")) {
                    void *token_delete = nullptr;
                    if (ImGui::TreeNode("Fiducial Points")) {
                        ImGui::PushItemWidth(200);
                        point_set->for_each([&](void *token, const Eigen::Vector3d &point) {
                            Eigen::Vector3f point_f = point.cast<float>();
                            ImGui::PushID(token);
                            ImGui::Bullet();
                            ImGui::BeginDisabled();
                            ImGui::InputFloat3("", point_f.data());
                            ImGui::EndDisabled();
                            if (is_picking) {
                                ImGui::SameLine();
                                if (ImGui::SmallButton("Delete")) {
                                    token_delete = token;
                                }
                            }
                            ImGui::PopID();
                        });
                        ImGui::PopItemWidth();
                        ImGui::TreePop();
                    }
                    if (token_delete != nullptr) {
                        eq.emplace([=] { point_set->remove_point(token_delete); });
                    }
                    ImGui::PushItemWidth(100);
                    ImGui::Bullet();
                    ImGui::InputFloat("Probe Test Error (mm)", &probe_test_error);
                    if (cur_target->is_finished) {
                        ImGui::Bullet();
                        ImGui::InputFloat("Max Fiducial Error (mm)", &cur_target->max_error);
                    }
                    ImGui::PopItemWidth();
                }
            }
        }
        ImGui::End();
        if (cur_target == nullptr) return;

        if (ImGui::Begin("Registration View", nullptr, vtk_viewer::no_scroll_flag)) {
            viewer->show();
        }
        ImGui::End();
    }

    void process() {
        while (!eq.empty()) {
            eq.front()();
            eq.pop();
        }
        if (is_picking) {
            try_picking();
        } else if (is_collecting) {
            try_collect();
        }

        // update probe transform
        if (cur_target != nullptr) {
            auto trans = conn->query_transform_variable(cur_target->probe_var_name);
            update_actor_pose(probe_actor, trans);
            if (trans.has_value()) {
                Eigen::Vector3d point = trans.value().translation();
                std::tie(std::ignore, probe_test_error) = calc_closest_point(point);
            } else {
                probe_test_error = NAN_VALUE;
            }
        }
    }

};

registration::~registration() = default;

registration *registration::create(const registration_config &conf) {
    auto pimpl = new impl{};
    pimpl->conn = conf.conn;
    pimpl->probe_actor = create_actor(conf.probe_model_path);
    auto ret = new registration{};
    ret->pimpl.reset(pimpl);
    return ret;
}

void registration::add_target(const registration_target &item) {
    pimpl->add_target(item);
}

void registration::show() {
    pimpl->show();
}

void registration::process() {
    pimpl->process();
}