#include <cstdint>
#include "./imgui.h"
#include "./imgui_internal.h"

class Histogram {

public:
  unsigned int count[4][256] = {0};
  int width;
  int height;
  int channels;

  Histogram(int width, int height, int channels) {
    this->width = width;
    this->height = height;
    this->channels = channels;
  }

  void Load(uint8_t * image) {
    for (int l = 0; l < height * width; l++)
    {
      count[0][*image++]++;
      count[1][*image++]++;
      count[2][*image++]++;
      count[3][*image++]++;
    }
  };

  void Draw() {
    ImGui::InvisibleButton("histogram", ImVec2(512, 256));

    unsigned int maxv = count[0][0];
    unsigned int* pCount = &count[0][0];
    for (int i = 0; i < 3 * 256; i++, pCount++)
    {
      maxv = (maxv > *pCount) ? maxv : *pCount;
    }

    ImDrawList* drawList = ImGui::GetWindowDrawList();
    const ImVec2 rmin = ImGui::GetItemRectMin();
    const ImVec2 rmax = ImGui::GetItemRectMax();
    const ImVec2 size = ImGui::GetItemRectSize();
    const float hFactor = size.y / float(maxv);

    for (int i = 0; i <= 10; i++)
    {
      float ax = rmin.x + (size.x / 10.f) * float(i);
      float ay = rmin.y + (size.y / 10.f) * float(i);
      drawList->AddLine(ImVec2(rmin.x, ay), ImVec2(rmax.x, ay), 0x80808080);
      drawList->AddLine(ImVec2(ax, rmin.y), ImVec2(ax, rmax.y), 0x80808080);
    }

    const float barWidth = (size.x / 256.f);
    for (int j = 0; j < 256; j++)
    {
      // pixel count << 2 + color index(on 2 bits)
      uint32_t cols[3] = {(count[0][j] << 2), (count[1][j] << 2) + 1, (count[2][j] << 2) + 2};
      if (cols[0] > cols[1])
        ImSwap(cols[0], cols[1]);
      if (cols[1] > cols[2])
        ImSwap(cols[1], cols[2]);
      if (cols[0] > cols[1])
        ImSwap(cols[0], cols[1]);
      float heights[3];
      uint32_t colors[3];
      uint32_t currentColor = 0xFFFFFFFF;
      for (int i = 0; i < 3; i++)
      {
        heights[i] = rmax.y - (cols[i] >> 2) * hFactor;
        colors[i] = currentColor;
        currentColor -= 0xFF << ((cols[i] & 3) * 8);
      }

      float currentHeight = rmax.y;
      const float left = rmin.x + barWidth * float(j);
      const float right = left + barWidth;
      for (int i = 0; i < 3; i++)
      {
        if (heights[i] >= currentHeight)
        {
          continue;
        }
        drawList->AddRectFilled(ImVec2(left, currentHeight), ImVec2(right, heights[i]), colors[i]);
        currentHeight = heights[i];
      }
    }
  };
};