#pragma once

#include "array2d.hpp"
#include "lodepng.h"

#include <vector>

using namespace std;

template<typename T>
struct pixels {
  array2d<T> r;
  array2d<T> g;
  array2d<T> b;
  size_t n;
  size_t w;
  size_t h;

  pixels(size_t w, size_t h) noexcept
    : r(w, h)
    , g(w, h)
    , b(w, h)
    , n(w * h)
    , w(w)
    , h(h)
  {}

  pixels() : pixels(0, 0) {};
  pixels(const pixels&) = delete;
  pixels(pixels&&) = default;
  pixels<T>& operator=(pixels<T> &&o) = default;

  operator bool() const {
    return n > 0;
  }
};

pair<pixels<uint8_t>*, bool> from_png(char *fname) {
  // 4 components per pixel, rgba
  vector<uint8_t> rawdata;
  uint32_t width, height;

  auto err = lodepng::decode(rawdata, width, height, fname);
  if (err) {
    cout << "couldn't load png: " << err << endl;
    return make_pair(nullptr, false);
  }
  if (rawdata.size() % 4 != 0) {
    cout << "bad size for rawdata " << rawdata.size() << endl;
    return make_pair(nullptr, false);
  }

  pixels<uint8_t> * const inpx = new pixels<uint8_t>(width, height);

  #pragma omp parallel for
  for (size_t i = 0; i < inpx->n; i++) {
    inpx->r(i) = rawdata[4 * i + 0];
    inpx->g(i) = rawdata[4 * i + 1];
    inpx->b(i) = rawdata[4 * i + 2];
  }

  return make_pair(inpx, true);
}

bool to_png(string fname, const pixels<uint8_t> &img) {
  uint8_t *data = (uint8_t*) malloc(3 * img.n);
  for (size_t i = 0; i < img.n; i++) {
    data[3 * i + 0] = img.r(i);
    data[3 * i + 1] = img.g(i);
    data[3 * i + 2] = img.b(i);
  }

  auto err = lodepng::encode(fname, data, img.w, img.h, LCT_RGB);
  if (err) {
    cout << "couldn't save pixel image: " << err << endl;
    return false;
  }
  return true;
}