#include <iostream>
#include <random>
#include <string>
#include <vector>
// Command-line parsing
#include <CLI11.hpp>

// Image filtering and I/O
#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <SimpleProgressBar.hpp>
#include <stb_image_write.h>

// Global flags
bool silent = false;
bool test_energy = false;
int max_step = 1;

// Get index for any table indexed by [width*(i : height) + (j : width)], but a
// : dim_long, b : dim_large
#define im_index(a, b) (vertical ? (width * a + b) : (width * b + a))

bool nearly_equal(float a, float b) {
  return std::nextafter(a, std::numeric_limits<float>::lowest()) <= b &&
         std::nextafter(a, std::numeric_limits<float>::max()) >= b;
}

std::pair<int, float> operator+(std::pair<int, float>& p1, std::pair<int, float>& p2) {
  return {
    // If one of the two pixels is "protected" (ie 2), we want to prevent this line removing
    // Else, we want to keep the information of "is there a pixel to remove in this seam" (ie 0)
    (p1.first==2) || (p2.first==2) ? 2 : std::min(p1.first, p2.first),
    p1.second+p2.second
  };
}

void operator+=(std::pair<int, float>& p1, std::pair<int, float>& p2) {
  p1 = p1+p2;
}

namespace limits {
  struct max_energy {
      template<class T> operator T() {
        return std::numeric_limits<T>::max();
      }
      operator std::pair<int, float>() {
        return {3, std::numeric_limits<float>::max()};
      }
  };  
}


void export_image(const char *filename, const void *data, int width, int height,
                  int nbChannels) {
  if (!silent)
    std::cout << "Exporting to \"" << filename << "\".." << std::endl;
  int errcode = stbi_write_png(filename, width, height, nbChannels, data,
                               nbChannels * width);
  if (!errcode) {
    std::cerr << "Error while exporting the resulting image." << std::endl;
    exit(errcode);
  }
}

#define compute_energy_for_pixel(source, width, height, i, j, nbChannels,      \
                                 nbColorChannels, dest)                        \
  auto indexPixel = (nbChannels) * (width * (j) + (i));                        \
  auto indexPixel_up =                                                         \
      ((j) - 1 > 0) ? (nbChannels) * (width * ((j) - 1) + (i)) : indexPixel;   \
  auto indexPixel_down = ((j) + 1 < height)                                    \
                             ? (nbChannels) * (width * ((j) + 1) + (i))        \
                             : indexPixel;                                     \
  auto indexPixel_left =                                                       \
      ((i) - 1 > 0) ? (nbChannels) * (width * (j) + ((i) - 1)) : indexPixel;   \
  auto indexPixel_right = ((i) + 1 < width)                                    \
                              ? (nbChannels) * (width * (j) + ((i) + 1))       \
                              : indexPixel;                                    \
  dest = 0;                                                                    \
  for (auto ch = 0; ch < (nbColorChannels); ch++) {                            \
    dest +=                                                                    \
        (fabs((float)source[indexPixel_up + ch] - source[indexPixel + ch]) +   \
         fabs((float)source[indexPixel_down + ch] - source[indexPixel + ch]) + \
         fabs((float)source[indexPixel_left + ch] - source[indexPixel + ch]) + \
         fabs((float)source[indexPixel_right + ch] -                           \
              source[indexPixel + ch]));                                       \
  }                                                                            \
// Le alpha n'est pas pris en compte dans l'énergie
// Here, we use this /ugly/ macro to avoid defining a function that would be way
// longer...

/** e_1 energy, energy is always normalized between 0 and 1 */
std::vector<float> energy_e1(std::vector<unsigned char> source, int width,
                             int height, int nbChannels) {
  int nbColorChannels =
      nbChannels > 3 ? 3 : nbChannels; // nombre de canaux, excepté le alpha
  std::vector<float> energy(width * height);
  float max_energy = 0;
  for (auto i = 0; i < width; i++) {
    for (auto j = 0; j < height; j++) {
      compute_energy_for_pixel(
        source, width, height,
        i, j, nbChannels, nbColorChannels, energy[width*j+i]
      );
    }
  }
  return energy;
}

/** Given the energy value, returns the optimal seam */
template <typename T> std::vector<int> optimal_seam(std::vector<T> energy, int width, int height,
                              bool vertical) {

  // dyn_energy is indexed by [dim_large*(i : dim_long) + (j : dim_large)]
  std::vector<T> dyn_energy(width * height);

  int dim_large = vertical ? width : height;
  int dim_long = vertical ? height : width; // Number of elements in the seam

  //* Find an end of the minimal connected vertical/horizontal seam
  for (auto i = 0; i < dim_large; i++) {
    dyn_energy[i] = energy[i];
  }

  for (auto i = 1; i < dim_long; i++) { // Propagate dyn_energy
    for (auto j = 0; j < dim_large; j++) {
      dyn_energy[dim_large * i + j] = limits::max_energy();
      int lower_bound = std::max(j - max_step, 0);
      int upper_bound = std::min(j + max_step, dim_large - 1);

      for (auto k = lower_bound; k <= upper_bound;
           k++) { // Compute energy based on predecessors
        dyn_energy[dim_large * i + j] = std::min(
            dyn_energy[dim_large * i + j], dyn_energy[dim_large * (i - 1) + k]
        );
      }
      dyn_energy[dim_large * i + j] += energy[im_index(i, j)];
    }
  }

  std::vector<int> result(dim_long);
  // Find the seam end
  int min_idx = -1;
  T min_val = limits::max_energy();

  for (auto j = 0; j < dim_large; j++) {
    if (min_val > dyn_energy[dim_large * (dim_long - 1) + j]) {
      min_idx = j;
      min_val = dyn_energy[dim_large * (dim_long - 1) + j];
    }
  }
  result[dim_long - 1] = min_idx;

  //* Backtracking to find the path
  for (auto i = dim_long - 1; i > 0; i--) {
    // We want to find either (bot_l, bot_c, bot_r) with dyn_energy[.] = min_val -
    // energy[cur]

    // Idea : float next_energy = min_val - energy[width*i + min_idx];
    //! With floats, we don't always have x + y - y == x, so we check is x+y == x+y
    // This define is a bit ugly but 200x faster than using a lambda function
    #define is_next_idx(idx)                                                       \
      (dyn_energy[(i - 1) * dim_large + idx] + energy[im_index(i, min_idx)] ==     \
      min_val)

    // This is not the nicest way to do this but we want to check in priority
    // at the center to have straight seams
    bool found = false;
    if (is_next_idx(min_idx)) {

      min_val = dyn_energy[(i - 1) * dim_large + min_idx];
      found = true;
    }
    for (auto k = 1; !found && k <= max_step; k++) {
      if (min_idx + k < dim_large && is_next_idx(min_idx + k)) {
        min_val = dyn_energy[(i - 1) * dim_large + min_idx + k];
        min_idx = min_idx + k;

        found = true;
      } else if (min_idx - k >= 0 && is_next_idx(min_idx - k)) {
        min_val = dyn_energy[(i - 1) * dim_large + min_idx - k];
        min_idx = min_idx - k;

        found = true;
      }
    }
    if (!found) {

      std::cerr << "Unable to backtrack path !" << std::endl;
      exit(1);
    }
    result[i - 1] = min_idx;
  }

  return result;
}

/** Carves an image by one seam. Returns the optimal seam used */
template <typename T>
void remove_seam(const std::vector<T> source, std::vector<T> &output, int width,
                 int height, int nbChannels, bool vertical,
                 const std::vector<int> seam) {
  // remove the given seam from the image, the result is in output
  // the output must have at least the right size!
  int dim_large = vertical ? width : height;
  int dim_long = vertical ? height : width;
  for (auto i = 0; i < dim_long; i++) {
    int cur_j = 0;
    for (auto j = 0; cur_j < dim_large - 1 && j < dim_large; j++) {
      if (seam[i] != j) {
        int out_pixelIndex = nbChannels * (vertical ? ((width - 1) * i + cur_j)
                                                    : (width * cur_j + i));
        int src_pixelIndex = nbChannels * im_index(i, j);
        for (auto ch = 0; ch < nbChannels; ch++)
          output[out_pixelIndex + ch] = source[src_pixelIndex + ch];
        cur_j++;
      }
    }
  }
}

// It would be preferable to use templates only for the value assignation but this is in fact far less efficient
void recompute_energy_along_seam(
    std::vector<unsigned char> carved_img, std::vector<float> &output_energy, std::vector<int> opt_seam,
    int width, int height,  int nbChannels, int nbColorChannels, bool vertical
  ) {
  int dim_large = vertical ? width : height;
  int dim_long = vertical ? height : width;

  int newWidth  = vertical ? width-1 : width;
  int newHeight = vertical ? height : height-1;

  for (auto j = 0; j < dim_long; j++) {
    for (auto i = -1; i < 2; i++) {
      int x = vertical ? (opt_seam[j] + i) : j;
      int y = vertical ? j : (opt_seam[j] + i);
      if ((0 < (opt_seam[j] + i)) && ((opt_seam[j] + i) < dim_large - 1)) {
        compute_energy_for_pixel(
          carved_img, newWidth, newHeight,
          x, y, nbChannels,
          nbColorChannels, output_energy[width*y+x]
        );
      }
    }
  }
}

void recompute_energy_along_seam(
  std::vector<unsigned char> carved_img, std::vector<std::pair<int, float>> &output_energy, std::vector<int> opt_seam,
  int width, int height,  int nbChannels, int nbColorChannels, bool vertical
  ) {
  int dim_large = vertical ? width : height;
  int dim_long = vertical ? height : width;

  int newWidth  = vertical ? width-1 : width;
  int newHeight = vertical ? height : height-1;

  for (auto j = 0; j < dim_long; j++) {
    for (auto i = -1; i < 2; i++) {
      int x = vertical ? (opt_seam[j] + i) : j;
      int y = vertical ? j : (opt_seam[j] + i);
      if ((0 < (opt_seam[j] + i)) && ((opt_seam[j] + i) < dim_large - 1)) {
        compute_energy_for_pixel(
          carved_img, newWidth, newHeight,
          x, y, nbChannels,
          nbColorChannels, output_energy[width*y+x].first
        );
      }
    }
  }
}


/** Carves an image and its energy by one seam, and recomputes the energy.
    Returns the optimal seam used */
template <typename T>
std::vector<int> carving_step(const std::vector<unsigned char> source_img,
                              std::vector<T> source_energy,
                              std::vector<unsigned char> &output_img,
                              std::vector<T> &output_energy, int width,
                              int height, int nbChannels, int nbColorChannels,
                              bool vertical) {

  std::vector<int> opt_seam =
    optimal_seam(source_energy, width, height, vertical);
  remove_seam(source_img, output_img, width, height, nbChannels, vertical,
              opt_seam);

  remove_seam(source_energy, output_energy, width, height, 1, vertical,
              opt_seam);

  // Recompute the energy along the seam, we need a separate function for templating
  recompute_energy_along_seam(
    source_img, output_energy, opt_seam,
    width, height, nbChannels, nbColorChannels,
    vertical
  );

  return opt_seam;
}

std::vector<std::pair<int, float>> mask_energy(std::vector<float> energy,
                                              int width, int height, unsigned char* mask) {
  std::vector<std::pair<int, float>> output(width*height);

  for (auto i=0; i < width*height; i++) {
    output[i] = {mask[i], energy[i]};
  }

  return output;
}

void seam_carving(unsigned char *source, int width, int height, int nbChannels,
                  const char *out_filename, int nbSeams, bool vertical, unsigned char* mask=nullptr) {
  int nbColorChannels = nbChannels > 3 ? 3 : nbChannels;
  int curWidth = width;
  int curHeight = height;

  int dim_large = vertical ? width : height;
  int dim_long = vertical ? height : width;
  // dim_long=longueur des seam

  std::vector<unsigned char> source_img(width * height * nbChannels);
  // Contains at each step the carved image
  std::vector<float> source_energy(width * height);
  // Contains at each step the carved energy
  std::vector<std::pair<int, float>> masked_energy;
  std::vector<std::pair<int, float>> output_masked_energy(width*height);
  // Source energy with (-1, 0, 1) on first element according to mask value
  std::vector<unsigned char> output_img(width * height * nbChannels);
  // Receives at each step the newly carved image
  std::vector<float> output_energy(width * height);
  // Contains at each step the carved energy
  std::vector<bool> complete_blacklist(width * height);
  // Contains all removed pixels, for "test_energy"
  std::vector<float> ini_energy;
  // Contains the initial energy, only for "test_energy"
  std::vector<unsigned char> test_energy_output(width * height * nbChannels);
  // Final output for "test_energy"

  for (auto i = 0; i < width * height * nbChannels; i++) {
    source_img[i] = source[i];
  }

  source_energy = energy_e1(source_img, width, height, nbChannels);
  if (mask)
    masked_energy = mask_energy(source_energy, width, height, mask);

  if (test_energy) {
    ini_energy = energy_e1(source_img, width, height, nbChannels);
    for (auto k = 0; k < width * height; k++) {
      complete_blacklist[k] = false;
    }

    //* Prepare final output

    float max_energy = __FLT_MIN__;
    for (auto k = 0; k < width * height; k++) {
      max_energy = std::max(max_energy, ini_energy[k]);
    }
    if (max_energy != 0) {
      for (auto k = 0; k < width * height; k++) {
        ini_energy[k] /= max_energy;
      }
    }
    for (auto k = 0; k < width * height; k++) {
      //* Uncomment if you prefer to see darkened source image
      // for (auto i=0; i < nbColorChannels; i++)
      //   output[nbChannels*k+i] = source_img[nbChannels*k+i]/nbChannels;

      for (auto ch = 0; ch < nbColorChannels; ch++) {
        test_energy_output[nbChannels * k + ch] = ini_energy[k] * 255;
      }
      if (nbChannels == 4)
        test_energy_output[nbChannels * k + 3] = source_img[nbChannels * k + 3];

      if (mask) {
        if (mask[k] == 2) // Green
          test_energy_output[nbChannels*k + 1] = 125;
        else if (mask[k] == 0) {// Red
          test_energy_output[nbChannels * k] = 125;
        }
      }
    }
  }

  SimpleProgressBar::ProgressBar bar(nbSeams);
  bar.print();

  for (auto seam_index = 0; seam_index < nbSeams; seam_index++) {
    std::vector<int> opt_seam;
    if (mask) {
      opt_seam = carving_step(
        source_img, masked_energy, output_img, output_masked_energy,
        curWidth, curHeight, nbChannels, nbColorChannels, vertical
      );
    } else {
      opt_seam = carving_step(
        source_img, source_energy, output_img, output_energy,
        curWidth, curHeight, nbChannels, nbColorChannels, vertical
      );
    }
    std::copy(output_img.begin(), output_img.end(), source_img.begin());
    std::copy(output_energy.begin(), output_energy.end(),
              source_energy.begin());
    if (mask)
      std::copy(output_masked_energy.begin(), output_masked_energy.end(),
        masked_energy.begin());

    vertical ? curWidth-- : curHeight--; // We just reduced the dimension

    if (test_energy) { // Update blacklist
      for (auto i = 0; i < dim_long; i++) {
        int j, cur_j = 0; // cur_j is the index relative to the current carved
                          // image. j is absolute in the source image
        for (j = 0; j < dim_large &&
                    (cur_j < opt_seam[i] || complete_blacklist[im_index(i, j)]);
             j++) {
          if (!complete_blacklist[im_index(i, j)]) {
            cur_j++;
          }
        }
        assert(cur_j == opt_seam[i]); // Else, j == width and cur_j is not in
                                      // the source image..
        complete_blacklist[im_index(i, j)] = true;
        test_energy_output[nbChannels * im_index(i, j)] =
            255; // Set carved pixel to red
      }
    }
    bar.increment();
    bar.print();
  }
  std::cout << std::endl; // Add newline after ProgressBar

  if (test_energy) {
    export_image(out_filename, test_energy_output.data(), width, height,
                 nbChannels);
  } else {
    export_image(out_filename, source_img.data(), curWidth, curHeight,
                 nbChannels);
  }
}

int main(int argc, char **argv) {
  CLI::App app{"seam-carving"};
  std::string maskImage;
  std::string sourceImage;
  std::string outputImage = "output.png";
  int nbSeams = 1;
  bool vertical = false;

  app.add_option("-s,--source", sourceImage, "Source image")
      ->required()
      ->check(CLI::ExistingFile);
  app.add_option("-o,--output", outputImage, "Output image")->required();
  app.add_option("-m,--mask", maskImage, "Source image")
      ->check(CLI::ExistingFile);

  app.add_option("-n,--nb-seams", nbSeams, "Number of seams")
      ->check(CLI::Number);
  app.add_option("--max-step", max_step, "Max width of step to find a seam")
      ->check(CLI::Number);

  app.add_flag("--vertical", vertical,
               "Vertical carving (remove vertical seams)");
  app.add_flag("--silent", silent, "No verbose messages");
  app.add_flag("--test-energy", test_energy,
               "Don't resize image, just try the specified energy function");
  CLI11_PARSE(app, argc, argv);

  // Image loading
  int width, height, nbChannels;
  unsigned char *source =
      stbi_load(sourceImage.c_str(), &width, &height, &nbChannels, 0);

  unsigned char* mask = nullptr;
  if (!maskImage.empty()) {
    int maskWidth, maskHeight, maskChannels;
    mask =
      stbi_load(maskImage.c_str(), &maskWidth, &maskHeight, &maskChannels, 0);

      if (maskWidth != width || maskHeight != height) {
        std::cerr << maskImage << " and " << sourceImage
        << " differ in dimension. Please provide a valid mask."
        << std::endl;
        exit(1);
      }
      if (maskChannels < 3) {
        std::cerr << maskImage << " needs to be RGB." << std::endl;
        exit(1);
      }
      unsigned char r, g, b;
      for (auto i=0; i < width*height; i++) {
        r = mask[maskChannels*i];
        g = mask[maskChannels*i+1];
        b = mask[maskChannels*i+2];
        bool positive = (g > r && g > b && g > 100); // Mask images are not always the cleanest
        bool negative = (r > g && r > b && r > 100);

        mask[i] = positive ? 2 : (negative ? 0 : 1);
      }
      //* From now on, mask has the same dimensions as source and one single channel
      //* The values are:
      //* . (2) we want to keep the pixel
      //* . (1) nothing in particular
      //* . (0) we want to remove the pixel
  }

  nbSeams = std::min(nbSeams, vertical ? width-1 : height-1); // We want to keep at least one row or column

  seam_carving(source, width, height, nbChannels, outputImage.c_str(),
                nbSeams, vertical, mask=mask);

  stbi_image_free(source);
  exit(0);
}