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

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


namespace pg {
namespace {

class Encoder final : public App, public Input::Data {
 public:
  static inline TypeInfo kType = TypeInfo::Create<Encoder>("Encoder");

  Encoder() noexcept : Data("encoder") {
  }

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

    if (ImGui::Begin(id.c_str())) {
      ImGui::DragInt("src", &src_, 1, 0, 1024);
      ImGui::DragInt("w", &w_, 1, 1, 1024);
      ImGui::DragInt("h", &h_, 1, 1, 1024);
      ImGui::DragInt2("block num", block_num_, 1, 1, 1024);
      ImGui::DragInt("utime", &utime_, 1, 5, 64);
      ImGui::InputTextMultiline("feats", &feats_);
      if (ImGui::Button("encode")) {
        Encode();
      }
    }
    ImGui::End();
  }

  Frame Fetch(size_t n) noexcept override {
    if (n >= frames_) return {};
    const auto w = static_cast<size_t>(w_);
    const auto h = static_cast<size_t>(h_);
    return { .w = w, .h = h, .rgba = buf_.data()+4*w*h*n };
  }
  size_t frames() noexcept override {
    return frames_;
  }

 private:
  int src_ = 0;
  int w_ = 128, h_ = 128;
  int block_num_[2] = {2, 2};
  int utime_ = 10;
  std::string feats_;

  std::vector<uint8_t> buf_;
  size_t frames_ = 0;


  struct Feat final {
   public:
    size_t block;
    char type;  // C: const, L: linear
  };

  void Encode() {
    const auto src = Input::instance().slots(static_cast<size_t>(src_));
    if (!src) throw "missing input";

    const auto feats = ParseFeats(feats_);

    const auto w     = static_cast<size_t>(w_);
    const auto h     = static_cast<size_t>(h_);
    const auto utime = static_cast<size_t>(utime_);
    buf_.resize(w*h*4*feats.size()*utime);
    frames_ = utime * feats.size();

    const auto b_num_x = static_cast<size_t>(block_num_[0]);
    const auto b_num_y = static_cast<size_t>(block_num_[1]);

    const auto dst_bw   = w/b_num_x;
    const auto dst_bh   = h/b_num_y;
    std::vector<std::pair<uint8_t, uint8_t>> block_rgba(dst_bw*dst_bh*4);

    const auto dst_bw_f = static_cast<float>(dst_bw);
    const auto dst_bh_f = static_cast<float>(dst_bh);

    auto dst_ptr = buf_.data();
    for (size_t t = 0; t < feats.size(); ++t) {
      const auto& f = feats[t];

      const auto bx = f.block%b_num_x;
      const auto by = f.block/b_num_x;
      if (by < b_num_y) {
        {
          // get pixels in start of the feature
          const auto frame = src->Fetch(t*utime);
          if (!frame.rgba) throw "got an empty frame";

          const auto bw    = frame.w / b_num_x;
          const auto bh    = frame.h / b_num_y;
          const auto bw_f  = static_cast<float>(bw);
          const auto bh_f  = static_cast<float>(bh);
          const auto off_x = bw*bx;
          const auto off_y = bh*by;

          auto itr = block_rgba.begin();
          for (size_t y = 0; y < dst_bh; ++y) {
            const auto yf = static_cast<float>(y) / dst_bh_f;
            for (size_t x = 0; x < dst_bw; ++x) {
              const auto xf = static_cast<float>(x) / dst_bw_f;

              const auto srcx = static_cast<size_t>(bw_f*xf) + off_x;
              const auto srcy = static_cast<size_t>(bh_f*yf) + off_y;

              auto ptr = frame.rgba + 4*(frame.w*srcy+srcx);
              for (size_t i = 0; i < 4; ++i) {
                itr->first  = *ptr;
                itr->second = *ptr;
                ++itr, ++ptr;
              }
            }
          }
        }
        if (f.type == 'L') {
          // get pixels in end frame if it's linear
          const auto frame = src->Fetch((t+1)*utime);
          if (!frame.rgba) throw "got an empty frame";

          const auto bw    = frame.w / b_num_x;
          const auto bh    = frame.h / b_num_y;
          const auto bw_f  = static_cast<float>(bw);
          const auto bh_f  = static_cast<float>(bh);
          const auto off_x = bw*bx;
          const auto off_y = bh*by;

          auto itr = block_rgba.begin();
          for (size_t y = 0; y < dst_bh; ++y) {
            const auto yf = static_cast<float>(y) / dst_bh_f;
            for (size_t x = 0; x < dst_bw; ++x) {
              const auto xf = static_cast<float>(x) / dst_bw_f;

              const auto srcx = static_cast<size_t>(bw_f*xf) + off_x;
              const auto srcy = static_cast<size_t>(bh_f*yf) + off_y;

              auto ptr = frame.rgba + 4*(frame.w*srcy+srcx);
              for (size_t i = 0; i < 4; ++i) {
                (itr++)->second = *(ptr++);
              }
            }
          }
        }
      }

      const auto bx_off = dst_bw*bx;
      const auto by_off = dst_bh*by;

      for (size_t bt = 0; bt < utime; ++bt) {
        const auto bt_f = static_cast<float>(bt) / static_cast<float>(utime);

        const auto frame = src->Fetch(t*utime+bt);
        if (!frame.rgba) throw "got an empty frame";
        for (size_t y = 0; y < w; ++y) {
          for (size_t x = 0; x < h; ++x) {
            const bool is_target =
                bx_off <= x && x < bx_off+dst_bw && by_off <= y && y < by_off+dst_bh;
            if (is_target) {
              const auto* ptr = &block_rgba[4*((x-bx_off)+(y-by_off)*dst_bw)];
              for (size_t i = 0; i < 4; ++i) {
                const auto a = static_cast<float>(ptr->second-ptr->first);
                const auto b = static_cast<float>(ptr->first);
                *(dst_ptr++) = static_cast<uint8_t>(bt_f*a+b);
                ++ptr;
              }
            } else {
              const auto xf = static_cast<float>(x) / static_cast<float>(w);
              const auto yf = static_cast<float>(y) / static_cast<float>(h);

              const auto src_x = static_cast<size_t>(static_cast<float>(frame.w)*xf);
              const auto src_y = static_cast<size_t>(static_cast<float>(frame.h)*yf);

              auto ptr = frame.rgba + 4*(src_x+src_y*frame.w);
              for (size_t i = 0; i < 4; ++i) {
                *(dst_ptr++) = (*ptr++);
              }
            }
          }
        }
      }
    }
  }
  static std::vector<Feat> ParseFeats(const std::string& script) {
    std::vector<Feat> ret;

    const char* begin = &*script.c_str();
    const char* end   = begin + script.size();

    size_t time = 0;
    for (auto itr = begin; itr < end; ++itr, ++time) {
      char* end;
      const auto idx = std::strtol(&*itr, &end, 0);
      if (itr == end) throw "invalid format: expected block index";
      if (*end != 'C' && *end != 'L') {
        throw "invalid format: unknown type";
      }
      itr = end+1;
      ret.push_back(Feat { .block = static_cast<size_t>(idx), .type = *end, });
    }
    return ret;
  }
};

}
}  // namespace pg