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

#include <numeric>
#include <vector>

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

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


namespace pg {

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

  Block() noexcept {
  }

  void Update() noexcept {
    const auto id = "Block | "+
        std::to_string(reinterpret_cast<uintptr_t>(this));

    const auto em = ImGui::GetFontSize();
    ImGui::SetNextWindowSize({40*em, 24*em}, ImGuiCond_Once);
    if (ImGui::Begin(id.c_str())) {
      UpdateParams();
      ImGui::SameLine();
      UpdatePlots();
      if (msg_.size() > 0) {
        ImGui::TextUnformatted(msg_.c_str());
      }
    }
    ImGui::End();
  }
  void UpdateParams() noexcept {
    const auto em = ImGui::GetFontSize();

    bool mod = false;
    ImGui::BeginGroup();
    ImGui::PushItemWidth(6*em);
    mod |= ImGui::DragInt("input_slot", &src_, 1, 0, 1024);
    mod |= ImGui::DragInt("time", &time_, 1, 0, 1024);
    mod |= ImGui::DragInt("dur", &dur_, 1, 1, BLKY_SENSOR_MAX_DUR);
    mod |= ImGui::DragInt2("blocks", block_, 1, 1, 16);
    mod |= ImGui::DragInt2("sensors", sensor_);
    mod |= ImGui::DragFloat("var thresh", &var_thresh_);
    ImGui::PopItemWidth();
    ImGui::EndGroup();

    if (mod) Calc();
  }
  void UpdatePlots() noexcept {
    const auto em = ImGui::GetFontSize();

    ImGui::BeginGroup();

    auto avail = ImGui::GetContentRegionAvail();
    avail.x -= em;
    avail.y -= em;
    if (ImPlot::BeginPlot("correls", avail)) {
      ImPlot::SetupAxisLimits(ImAxis_X1, 0, 1);
      ImPlot::SetupAxisLimits(ImAxis_Y1, 0, 1);
      ImPlot::PlotHeatmap(
          "correls_avg", correl_avg_.data(), block_[1], block_[0],
          0, 1, "%.2f");
      ImPlot::PlotHeatmap(
          "correls_var", correl_var_.data(), block_[1], block_[0],
          0, var_thresh_, "%.4f");
      ImPlot::PlotHeatmap(
          "correls_avg_filtered", correl_avg_filtered_.data(), block_[1], block_[0],
          0, 1, "%.2f");
      ImPlot::EndPlot();
    }
    ImGui::EndGroup();
  }

 private:
  int src_ = 0;
  int time_ = 0;
  int dur_ = 30;
  int block_[2] = {1, 1};
  int sensor_[2] = {1, 1};

  float var_thresh_ = .1f;

  std::string msg_;

  std::vector<double> correls_    = {0};
  std::vector<double> correl_avg_ = {0};
  std::vector<double> correl_var_ = {0};
  std::vector<double> correl_avg_filtered_ = {0};

  void Calc() noexcept
  try {
    auto data = Input::instance().slots(static_cast<size_t>(src_));
    if (!data) throw "missing input";

    const auto block_x = static_cast<size_t>(block_[0]);
    const auto block_y = static_cast<size_t>(block_[1]);
    const auto block_n = block_x*block_y;
    const auto block_w = 1.f / static_cast<float>(block_x);
    const auto block_h = 1.f / static_cast<float>(block_y);

    const auto sensor_x = static_cast<size_t>(sensor_[0]);
    const auto sensor_y = static_cast<size_t>(sensor_[1]);
    const auto sensor_n = sensor_x*sensor_y;
    const auto sensor_interval_x = block_w / static_cast<float>(sensor_x);
    const auto sensor_interval_y = block_h / static_cast<float>(sensor_y);

    correls_.clear();
    correls_.reserve(block_n*sensor_n*3);
    for (size_t by = 0; by < block_y; ++by) {
      for (size_t bx = 0; bx < block_x; ++bx) {
        for (size_t sy = 0; sy < sensor_y; ++sy) {
          for (size_t sx = 0; sx < sensor_x; ++sx) {
            const auto xf =
                block_w*static_cast<float>(bx) +
                sensor_interval_x*(static_cast<float>(sx)+.5f);
            const auto yf =
                block_h*static_cast<float>(by) +
                sensor_interval_y*(static_cast<float>(sy)+.5f);

            for (size_t i = 0; i < 3; ++i) {
              const auto samp = data->FetchSamples(
                  static_cast<size_t>(time_),
                  static_cast<size_t>(dur_),
                  xf, yf, i);

              blky_sensor_t sensor = {};
              blky_sensor_feed(&sensor, samp.data(), samp.size());
              correls_.push_back(sensor.correl);
            }
          }
        }
      }
    }

    correl_avg_.clear();
    correl_avg_.reserve(block_n);
    correl_var_.clear();
    correl_var_.reserve(block_n);
    correl_avg_filtered_.clear();
    correl_avg_filtered_.reserve(block_n);

    const auto correls_per_block = sensor_n*3;
    for (size_t bi = 0; bi < block_n; ++bi) {
      const size_t st = bi*correls_per_block;
      const size_t ed = st + correls_per_block;

      double sum = 0;
      for (size_t i = st; i < ed; ++i) {
        sum += std::abs(correls_[i]);
      }

      const auto avg = sum / static_cast<double>(correls_per_block);
      correl_avg_.push_back(avg);

      double var = 0;
      for (size_t i = st; i < ed; ++i) {
        const auto diff = std::abs(correls_[i]) - avg;
        var += diff*diff;
      }
      var /= static_cast<double>(correls_per_block);
      correl_var_.push_back(var);

      if (var < var_thresh_) {
        correl_avg_filtered_.push_back(avg);
      } else {
        correl_avg_filtered_.push_back(-1);
      }
    }

  } catch (const char* msg) {
    msg_ = msg;
  }
};

}  // namespace pg