#define _GNU_SOURCE

#include "lib/array2d.hpp"
#include "lib/pixels.hpp"
#include "lib/lodepng.cpp"
#include "lib/hsvrgb.cpp"

#include <algorithm>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <random>
#include <sstream>
#include <unordered_map>
#include <vector>

using namespace std;

constexpr double pexp = 1.8;
constexpr double ppmm = 150.0 / 25.4; // pixels per mm
constexpr double pointfill = 0.35; // number of output points per input pixel
constexpr double bucketsize_mm = 3;
constexpr int seed = 0x54481;

/* scale used for sample and sortimage */
constexpr double sscale = 1;
constexpr double ppmm_sortimage = ppmm * sscale;

template<typename T>
constexpr double proportion(T v) {
  return ((double) v) / ((double) std::numeric_limits<T>::max());
}

constexpr double P(uint8_t ri, uint8_t gi, uint8_t bi) {
#if defined(MULTIPLY)
  double v =
    (1.0 - proportion(ri)) *
    (1.0 - proportion(gi)) *
    (1.0 - proportion(bi));
#elif defined(RED)
  double v =
    1.0 - proportion(ri);
#elif defined(WEIGHTED)
  constexpr double rweight = 1.0;
  constexpr double gweight = 0.4;
  constexpr double bweight = 0.1;
  double v = 1.0;
  v *= (1.0 - rweight) + rweight * (1.0 - proportion(ri));
  v *= (1.0 - gweight) + gweight * (1.0 - proportion(gi));
  v *= (1.0 - bweight) + bweight * (1.0 - proportion(bi));
#else
#error no grayscale mode defined
#endif
  return pow(v, pexp);
}

static pixels<uint8_t> image;
static pixels<uint8_t> sortimage;
static array2d<double> probability;
static array2d<uint8_t> samples;
static vector<pair<double, double>> points;

bool load(char *fname) {
  pixels<uint8_t> *loaded;
  bool ok;
  tie(loaded, ok) = from_png(fname);
  if (not ok) return false;
  image = std::move(*loaded);

  probability = array2d<double>(image.w, image.h);
  samples = array2d<uint8_t>(image.w, image.h);
  return true;
}

void genprob() {
  #pragma omp parallel for
  for (size_t i = 0; i < probability.w; i++) {
    for (size_t j = 0; j < probability.h; j++) {
      probability(i, j) = P(image.r(i, j), image.g(i, j), image.b(i, j));
    }
  }
}

inline void emit(ostream &nc,
                 int64_t istart,
                 int64_t iend,
                 int64_t j,
                 bool transpose) {
  double xs, ys, xe, ye;
  if (transpose) {
    xs = ((double) j) / ppmm;
    ys = ((double) istart) / ppmm;
    xe = xs;
    ye = ((double) iend) / ppmm;
  } else {
    xs = ((double) istart) / ppmm;
    ys = ((double) j) / ppmm;
    xe = ((double) iend) / ppmm;
    ye = ys;
  }
  nc << "G0 X" << xs << " Y" << ys << "\n"
     << "G1 Z0\n";
  if (istart != iend)
    nc << "G1 X" << xe << " Y" << ye << "\n";
  nc << "G1 Z1\n";
}

void scatter_sample() {
  const double xmax = image.w / ppmm;
  const double ymax = image.h / ppmm;
  const size_t npoints = pointfill * image.n;

  std::mt19937 gen(seed);
  std::uniform_real_distribution<double> xdist(0.0, xmax);
  std::uniform_real_distribution<double> ydist(0.0, ymax);
  std::uniform_real_distribution<double> pdist(0.0, 1.0);

  points.clear();
  points.reserve(npoints);

  cout << "filling " << image.w << "x" << image.h
       << " (" << (image.w * image.h) << "px) with "
       << npoints << " points ("
       << ((npoints * 100.0) / ((double) image.w * image.h))
       << "% fill)"
       << endl;

  size_t p = 0;
  while (p < npoints) {
    double x = xdist(gen);
    double y = ydist(gen);
    double xpx = x * ppmm;
    double ypx = y * ppmm;

    int xlo = floor(xpx);
    int xhi = ceil(xpx);
    int ylo = floor(ypx);
    int yhi = ceil(ypx);

    if (yhi >= (int64_t) probability.h) yhi = ylo;
    if (xhi >= (int64_t) probability.h) xhi = xlo;

    double s11 = probability(xlo, ylo);
    double s12 = probability(xlo, yhi);
    double s22 = probability(xhi, yhi);
    double s21 = probability(xhi, ylo);

    double y1 = (xhi - xpx) * s11 + (xpx - xlo) * s21;
    double y2 = (xhi - xpx) * s12 + (xpx - xlo) * s22;
    double prob = (yhi - ypx) * y1 + (ypx - ylo) * y2;
    prob = pow(prob, pexp);

    if (pdist(gen) < prob) {
      points.push_back(make_pair(x, y));
      p++;
    }
  }

  samples = array2d<uint8_t>(sscale * image.w, sscale * image.h);
  sortimage = pixels<uint8_t>(sscale * image.w, sscale * image.h);

  for (const auto &pair : points) {
    const double &xmm = pair.first;
    const double &ymm = pair.second;
    int xpx = xmm * ppmm_sortimage;
    int ypx = ymm * ppmm_sortimage;
    samples(xpx, ypx) = 255;
  }
}

void linefeeder_sample() {
  std::mt19937 gen(seed);
  std::uniform_real_distribution<double> pdist(0.0, 1.0);
#pragma omp parallel for
  for (size_t i = 0; i < image.w; i++) {
    for (size_t j = 0; j < image.h; j++) {
      samples(i, j) = pdist(gen) < probability(i, j) ? 255 : 0;
    }
  }
}

void linefeeder() {
  ostream &nc = cout;
  nc << setprecision(3)
     << fixed
     << "G0 Z1\n"
     << "G0 X0 Y0\n";
  int64_t w = (int64_t) samples.w;
  int64_t h = (int64_t) samples.h;
  bool transpose = h > w;
  int64_t j = 0;
  int64_t i, runstart, runend;
  while (j < h) {
    i = 0;
    while (i < w) {
      if (not samples(i, j)) {
        i++;
        continue;
      }
      runstart = i;
      do {
        i++;
      } while (i < w and samples(i, j));
      runend = i - 1;
      emit(nc, runstart, runend, j, transpose);
    }
    j++;
    if (j >= h) break;
    i = w - 1;
    while (i >= 0) {
      if (not samples(i, j)) {
        i--;
        continue;
      }
      runstart = i;
      do {
        i--;
      } while (i >= 0 and samples(i, j));
      runend = i + 1;
      emit(nc, runstart, runend, j, transpose);
    }
    j++;
  }
}

typedef uint64_t bucketidx;

bucketidx make_idx(double x, double y) {
  x /= bucketsize_mm;
  y /= bucketsize_mm;

  if (x < 0 || x > std::numeric_limits<uint32_t>::max()) {
    throw std::out_of_range("bad value for x");
  }
  if (y < 0 || y > std::numeric_limits<uint32_t>::max()) {
    throw std::out_of_range("bad value for y");
  }

  uint64_t xx = static_cast<uint64_t>(x);
  uint64_t yy = static_cast<uint64_t>(y);
  return (yy << 32) | xx;
}

void arcsort() {
  typedef vector<pair<double, double>> bucket;
  unordered_map<bucketidx, bucket> buckets;

  for (const auto &p : points) {
    buckets[make_idx(p.first, p.second)].push_back(p);
  }

  vector<bucketidx> sorted_bucket_indexes;
  sorted_bucket_indexes.reserve(buckets.size());
  for (const auto &bentry : buckets) {
    sorted_bucket_indexes.push_back(bentry.first);
  }
  sort(sorted_bucket_indexes.begin(), sorted_bucket_indexes.end(),
       [](const auto &b1, const auto &b2) {
         uint64_t y1 = b1 & 0xFFFFFFFF00000000;
         uint64_t y2 = b2 & 0xFFFFFFFF00000000;
         if (y1 < y2) return true;
         if (y1 > y2) return false;
         if (y1 % 2) return b1 > b2;
         return b1 < b2;
       });
  for (auto &bentry : buckets) {
    sort(bentry.second.begin(), bentry.second.end(),
         [](const auto &p1, const auto &p2) {
           if (p1.first < p2.first) return true;
           if (p1.first > p2.first) return false;
           return p1.second < p2.second;
         });
  }

  size_t pi = 0;
  for (const auto &bidx : sorted_bucket_indexes) {
    const bucket &bucket = buckets.at(bidx);
    for (const auto &p : bucket) {
      double xpx = p.first * ppmm_sortimage;
      double ypx = p.second * ppmm_sortimage;
      double t = (double) pi++ / (double) points.size();
      uint8_t r, g, b;
      rainbow(t, r, g, b);
      sortimage.r(xpx, ypx) = r;
      sortimage.g(xpx, ypx) = g;
      sortimage.b(xpx, ypx) = b;
    }
  }

  ostream &nc = cout;
  nc << setprecision(3) << fixed
     << "G0 Z1\nG0 X0 Y0\n";
  for (const auto &bidx : sorted_bucket_indexes) {
    const bucket &bucket = buckets.at(bidx);
    // put the conditional outside the loop for maximum loopage
    if (image.h > image.w) {
      for (const auto &p : bucket) {
        nc << "G0 X" << p.second << " Y" << p.first
           << "\nG1 Z0\nG1 Z1\n";
      }
    } else {
      for (const auto &p : bucket) {
        nc << "G0 X" << p.first << " Y" << p.second
           << "\nG1 Z0\nG1 Z1\n";
      }
    }
  }
  nc << "G0 Z1\nG0 X0 Y0\n";
}

bool saveimages() {
  if (not to_png("_pntlsm_prob.png", probability)) {
    cout << "failed to save prob image" << endl;
    return false;
  }

  array2d<uint8_t> inverted = samples;
  for (size_t i = 0; i < inverted.w * inverted.h; i++) {
    inverted(i) = 255 - inverted(i);
  }
  if (not to_png("_pntlsm_samples.png", inverted)) {
    cout << "failed to save sample image" << endl;
    return false;
  }

  if (sortimage) {
    if (not to_png("_pntlsm_output.png", sortimage)) {
      cout << "failed to save sortimage image" << endl;
      return false;
    }
  }
  return true;
}

int main(int argc, char *argv[]) {
  if (argc < 2) {
    cout << "missing filename" << endl;
    return EXIT_FAILURE;
  }

  if (not load(argv[1])) return EXIT_FAILURE;

  genprob();

#if defined(SCATTER)
  scatter_sample();
  arcsort(); // must come before saveimages!
#elif defined(FEEDER)
  linefeeder_sample();
  linefeeder();
#else
#error no sampling method specified
#endif

  if (not saveimages()) return EXIT_FAILURE;
  return EXIT_SUCCESS;
}