extern "C" {
#include <liblocky.h>
}

#include <sstream>
#include <string>

#include <imgui.h>
#include <imgui_stdlib.h>
#include <implot.h>

#include "app.hh"
#include "input.hh"


namespace pg {

class Extractor final : public App {
 public:
  static inline TypeInfo kType = TypeInfo::Create<Extractor>("Extractor");

  Extractor() noexcept {
  }

  void Update() noexcept override {
    const auto em = ImGui::GetFontSize();
    const auto id = " Extractor | "+
        std::to_string(reinterpret_cast<uintptr_t>(this));

    ImGui::SetNextWindowSize({32*em, 32*em}, ImGuiCond_Once);
    if (ImGui::Begin(id.c_str())) {
      ImGui::DragInt("src", &src_, 1, 0, 1024);
      ImGui::DragInt2("block num", block_num_, 1, 1, 1024);
      ImGui::DragInt2("sensor num", sensor_num_, 1, 1, 1024);
      ImGui::DragInt("utime", &utime_, 1, 5, 64);
      ImGui::DragInt("len", &len_, 1, 1, 1024);
      ImGui::DragFloat("correl max var", &correl_max_var_, 1e-3f, 0, 1);
      ImGui::DragFloat("correl min avg", &correl_min_avg_, 1e-3f, 0, 1);
      ImGui::DragFloat("feature min probability", &feat_min_prob_, 1e-3f, 0, 1);
      if (ImGui::Button("extract")) {
        Extract();
      }

      time_ = std::clamp(time_, 0, len_-1);
      ImGui::SliderInt("time", &time_, 0, len_-1);

      const auto offset = static_cast<size_t>(block_num_[0]*block_num_[1]*time_);
      if (probs_.size() > offset) {
        auto avail = ImGui::GetContentRegionAvail();
        avail.x -= 1*em;
        avail.y -= 1*em;
        if (ImPlot::BeginPlot("probs", avail)) {
          ImPlot::SetupAxisLimits(ImAxis_X1, 0, 1);
          ImPlot::SetupAxisLimits(ImAxis_Y1, 0, 1);

          ImPlot::PlotHeatmap(
              "probs", probs_.data()+offset, block_num_[1], block_num_[0],
              0, 1, "%.3f");
          ImPlot::EndPlot();
        }
      } else {
        ImGui::TextUnformatted("no data");
      }
    }
    ImGui::End();
  }

 private:
  int src_ = 0;
  int block_num_[2] = {2, 2};
  int sensor_num_[2] = {2, 2};
  int utime_ = 10;
  int len_ = 1;

  float correl_max_var_ = 1e-3f;
  float correl_min_avg_ = 0.8f;
  float feat_min_prob_  = .5f;

  int time_;

  std::vector<double> probs_;


  void Extract() noexcept {
    const auto src = Input::instance().slots(static_cast<size_t>(src_));
    if (!src) return;

    const uint64_t dur = static_cast<uint64_t>(utime_ * len_);

    blky_extractor_t ex = {};

    ex.block_num_x = static_cast<uint32_t>(block_num_[0]);
    ex.block_num_y = static_cast<uint32_t>(block_num_[1]);
    ex.sensor_num_block_x = static_cast<uint32_t>(sensor_num_[0]);
    ex.sensor_num_block_y = static_cast<uint32_t>(sensor_num_[1]);
    ex.samples_per_pix = 3;
    ex.pix_stride = 4;
    ex.utime = static_cast<uint32_t>(utime_);
    ex.correl_max_var = static_cast<double>(correl_max_var_);
    ex.correl_min_avg = static_cast<double>(correl_min_avg_);
    blky_extractor_init(&ex);

    static const double verts[] = {
      0, 0,
      0, 1,
      1, 1,
      1, 0,
    };

    const uint64_t block_num = ex.block_num_x * ex.block_num_y;
    probs_.clear();
    probs_.reserve(block_num*static_cast<uint64_t>(len_));
    for (uint64_t t = 0; t < dur; ++t) {
      const auto f = src->Fetch(t);
      if (!f.rgba) break;

      const bool pop = blky_extractor_feed(
          &ex, f.rgba,
          static_cast<uint32_t>(f.w),
          static_cast<uint32_t>(f.h),
          verts);
      if (pop) {
        probs_.insert(probs_.end(), ex.probs, ex.probs+block_num);
      }
    }
    blky_extractor_deinit(&ex);
  }
};

}  // namespace pg