#pragma once

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

#include "lodepng.h"

using namespace std;

template<typename T>
struct array2d {
  T *values;
  size_t w;
  size_t h;

  array2d() : array2d(0, 0) {};

  array2d(const size_t w, const size_t h) noexcept : w(w), h(h) {
    values = (T*) calloc(w * h, sizeof(T));
  }

  array2d(const array2d<T> &o) : w(o.w), h(o.h) {
    values = (T*) calloc(w * h, sizeof(T));
    memcpy(values, o.values, w * h * sizeof(T));
  }

  array2d(array2d<T> &&o) : w(o.w), h(o.h), values(o.values) {
    o.values = nullptr;
  }

  ~array2d() noexcept {
    if (values != nullptr) free(values);
  }

  array2d<T>& operator=(const array2d<T> &o) = delete;

  array2d<T>& operator=(array2d<T> &&o) {
    w = o.w;
    h = o.h;
    T *t = values;
    values = o.values;
    o.values = t;
    return *this;
  }

  T& operator()(size_t i) noexcept {
    return values[i];
  }

  const T& operator()(size_t i) const noexcept {
    return values[i];
  }

  T& operator()(size_t i, size_t j) {
    if (__builtin_expect(i > w or j > h, 0)) {
      char *msg;
      asprintf(&msg, "bad access in array2d: i=%llu j=%llu, w=%llu h=%llu",
               i, j, w, h);
      throw out_of_range(msg);
    }
    return values[w * j + i];
  }

  const T& operator()(size_t i, size_t j) const {
    if (__builtin_expect(i > w or j > h, 0)) {
      char *msg;
      asprintf(&msg, "bad access in array2d: i=%llu j=%llu, w=%llu h=%llu",
               i, j, w, h);
      throw out_of_range(msg);
    }
    return values[w * j + i];
  }

  T* data() noexcept {
    return values;
  }

  const T* data() const noexcept {
    return values;
  }

  void fill(T val) {
    const size_t n = w * h;
    for (size_t i = 0; i < n; i++) {
      values[i] = val;
    }
  }

  array2d<T> transposed() const {
    array2d<T> res(h, w);
    for (size_t i = 0; i < w; i++) {
      for (size_t j = 0; j < h; j++) {
        res(j, i) = operator()(i, j);
      }
    }
    return res;
  }
};

bool to_png(string fname, const array2d<double> &grays) {
  vector<uint8_t> normed;
  const size_t n = grays.w * grays.h;
  normed.reserve(n);
  for (size_t i = 0; i < n; i++) {
    double clamped = grays(i);
    clamped = clamped > 1 ? 1 : clamped < 0 ? 0 : clamped;
    normed.push_back((uint8_t) (clamped * 255));
  }
  auto err = lodepng::encode(fname, normed.data(), grays.w, grays.h, LCT_GREY);
  if (err) {
    cout << "failed to write png: " << err << endl;
    return false;
  }
  return true;
}

bool to_png(string fname, const array2d<uint8_t> &vals) {
  auto err = lodepng::encode(fname, vals.data(), vals.w, vals.h, LCT_GREY);
  if (err) {
    cout << "failed to write png: " << err << endl;
    return false;
  }
  return true;
}