#include "depth_guide_v2.h"
#include "core/imgui_utility.hpp"
#include "image_process/process_funcs.h"

#include <glm/gtx/transform.hpp>

app_depth_guide_v2::app_depth_guide_v2(const create_config &_conf) {
    main_conf = _conf;
    auto conf = main_conf.ext_conf;

    // initialize object manager
    OBJ_SAVE(img_color, image_u8c3());
    OBJ_SAVE(img_depth, image_f32c1());
    OBJ_SAVE(img_remap, image_ptr());
    OBJ_SAVE(pc_raw, pc_ptr());
    OBJ_SAVE(scene_out, scene_ptr());
    OBJ_SAVE(img_out_left, image_ptr());
    OBJ_SAVE(img_out_right, image_ptr());
    OBJ_SAVE(img_out, image_ptr());

    // initialize meta
    OBJ_PIN_META(pc_raw, META_SERIES_NAME, meta_hash("pc_raw"));
    OBJ_PIN_META(img_remap, META_SERIES_NAME, meta_hash("img_remap"));
    OBJ_PIN_META(pc_raw, META_DEPTH_MIN, (float) 0.0f); // TODO: make adjustable
    OBJ_PIN_META(pc_raw, META_DEPTH_MAX, (float) 3000.0f);

    OBJ_SIG(img_color)->connect([](auto _) {
        gen_pc_rgbd::call(
                {img_color, img_depth, img_remap,
                 pc_raw, default_cuda_stream});
    });
    OBJ_SIG(pc_raw)->connect([this](auto _) {
        aug_manager->update(aug_helper->get_camera());

        auto out_conf = camera_augment_helper_v2::output_config{
                .size = aug_stereo->get_suggest_size(), .stream = default_cuda_stream,
        };
        out_conf.img_name = img_out_left;
        aug_left->render_image(out_conf);
        out_conf.img_name = img_out_right;
        aug_right->render_image(out_conf);
    });

    // initialize sophiar
    sophiar_thread = std::make_unique<std::thread>([conf_path = LOAD_STR("sophiar_config")] {
        sophiar::run_sophiar(conf_path);
    });
    sophiar_conn = std::make_unique<sophiar_conn_type>();
    sophiar_start_var = LOAD_STR("sophiar_start_var");

    // initialize modules
    auto orb_cam_conf = orb_camera_ui::create_config{
            .cf_name = img_color, .df_name = img_depth,
            .remap_name = img_remap,
            .ctx = main_conf.asio_ctx, .stream = default_cuda_stream,
    };
    orb_cam = std::make_unique<orb_camera_ui>(orb_cam_conf);

    auto aug_list_v1 = augment_manager::item_list_from_yaml(LOAD_LIST("augment_list"));
    auto aug_list = augment_manager_v2::item_list_from_v1(aug_list_v1);
    auto aug_conf = augment_manager_v2::create_config{
            .scene_name = scene_out,
            .sophiar_conn = sophiar_conn.get(),
            .stream = default_cuda_stream,
    };
    aug_conf.item_list.push_back(
            {.type = augment_manager_v2::ITEM_PC,
                    .disp_name = "PC Raw", .name = pc_raw,
                    .transform_var = LOAD_STR("camera_transform_var")});
    std::copy(aug_list.begin(), aug_list.end(), std::back_inserter(aug_conf.item_list));
    aug_manager = std::make_unique<augment_manager_v2>(aug_conf);

    auto guide_conf = cursor_guide::create_config{
            .item_list = cursor_guide::item_list_from_aug(aug_list),
            .sophiar_conn = sophiar_conn.get(),
    };
    aug_guide = std::make_unique<cursor_guide>(guide_conf);

    auto view_conf = camera_augment_helper_v2::create_config{
            .camera = camera_augment_helper_v2::create_config::freedom_camera_config{
                    .transform_var = LOAD_STR("camera_transform_var")},
            .sophiar_conn = sophiar_conn.get(),
            .manager = aug_manager.get(),
            .act_guide = aug_guide.get(),
            .ctx = main_conf.asio_ctx,
    };
    aug_helper = std::make_unique<camera_augment_helper_v2>(view_conf);
    aug_helper->set_active_guidance(true);

    view_conf.camera = camera_augment_helper_v2::create_config::relative_camera_config{
            .parent = aug_helper.get(),
            .transform  = glm::translate(glm::vec3(-baseline_dis / 2, 0, 0))};
    aug_left = std::make_unique<camera_augment_helper_v2>(view_conf);
    aug_left->sync_with(aug_helper.get());

    view_conf.camera = camera_augment_helper_v2::create_config::relative_camera_config{
            .parent = aug_helper.get(),
            .transform  = glm::translate(glm::vec3(baseline_dis / 2, 0, 0))};
    aug_right = std::make_unique<camera_augment_helper_v2>(view_conf);
    aug_right->sync_with(aug_helper.get());

    auto output_size = cv::Size(
            LOAD_NUMBER(int, "output_width"),
            LOAD_NUMBER(int, "output_height"));
    auto stereo_conf = stereo_augment_helper::create_config{
            .left_name = img_out_left, .right_name = img_out_right,
            .out_name = img_out, .fbo_size = output_size,
            .flip_image = true, .stream = default_cuda_stream,
    };
    aug_stereo = std::make_unique<stereo_augment_helper>(stereo_conf);

    auto out_conf = image_streamer::create_config{
            .img_name = img_out, .encode_scene = false,
            .asio_ctx = main_conf.asio_ctx,
            .cuda_ctx = main_conf.cuda_ctx, .stream = default_cuda_stream
    };
    out_streamer = std::make_unique<image_streamer>(out_conf);

    auto reg_conf = registration_config{
            .conn = sophiar_conn.get(),
            .probe_model_path = LOAD_STR("probe_model"),
    };
    reg.reset(registration::create(reg_conf));
    for (auto reg_item: LOAD_LIST("registration_list")) {
        auto item_conf = registration_target::from_yaml(reg_item);
        reg->add_target(item_conf);
    }
}

void app_depth_guide_v2::show_ui() {
    if (ImGui::Begin("Depth Guide Control")) {
        ImGui::PushItemWidth(200);

        if (ImGui::CollapsingHeader("Camera")) {
            auto id_guard = imgui_id_guard("camera");
            orb_cam->show();
        }

        if (ImGui::CollapsingHeader("Augment")) {
            ImGui::SeparatorText("Tracker");
            if (ImGui::Button("Start")) {
                start_tracking();
            }
            ImGui::SameLine();
            ImGui::Checkbox("Registration", &enable_reg);
            {
                ImGui::SeparatorText("Scene");
                auto id_guard = imgui_id_guard("augment_scene");
                aug_manager->show();
            }
            {
                ImGui::SeparatorText("Camera");
                auto id_guard = imgui_id_guard("augment_camera");
                aug_helper->show();
            }
            {
                ImGui::SeparatorText("Stereo");
                auto id_guard = imgui_id_guard("stereo");
                aug_stereo->show();
            }
        }

        if (ImGui::CollapsingHeader("Streamer")) {
            auto id_guard = imgui_id_guard("streamer");
            out_streamer->show();
        }

        ImGui::PopItemWidth();
    }
    ImGui::End();

    if (enable_reg) {
        reg->process();
        reg->show();
    }
}

void app_depth_guide_v2::render_background() {
    aug_helper->render();
}

void app_depth_guide_v2::start_tracking() {
    // sophiar_conn->start_object(sophiar_start_var);
    // work in another thread to prevent blocking
    auto t = std::thread([this] {
        auto conn = sophiar::local_connection();
        conn.start_object(sophiar_start_var);
    });
    t.detach();
}

app_depth_guide_v2::~app_depth_guide_v2() {
    // sophiar
    sophiar::stop_sophiar();
    sophiar_thread->join();
}