/*
 * Copyright 2018, Theodore Chabardes
 * <theodore.chabardes@mines-paristech.fr>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include "src/paraws.hpp"
#include "src/io_png_raw.hpp"
#include "src/output_color.hpp"
#include <iostream>
#include <limits>
#include <string>
// Portable version of getopt.
#define OPTPARSE_IMPLEMENTATION
#define OPTPARSE_API static
#include "src/optparse.h"
#ifdef BENCH
#include "src/bench.hpp"
#include "src/debug.hpp"
#include "src/label.hpp"
#endif // BENCH

bool check_compliance_image(const char *path, bool is_uint8, bool is_uint16,
                            bool is_uint32, size_t x, size_t y,
                            std::string &ext) {
  if (path == NULL)
    return false;
  std::string path_str = std::string(path);
  ext = path_str.substr(path_str.find_last_of(".") + 1);
  if (ext == "raw" && (x == 0 || y == 0))
    return false;
  return (is_uint8 + is_uint16 + is_uint32) == 1 &&
         (ext == "png" || ext == "raw");
}

bool check_structuring_element(bool cross, bool hexa, bool square, bool cross3d,
                               bool cube) {
  return (cross + hexa + square + cross3d + cube) == 1;
}

void usage(const char *execName) {
  std::cout
      << execName << " [options]" << std::endl
      << "Options:" << std::endl
#ifdef BENCH
      << "-A | --paraws_bench           Benchs paraws" << std::endl
      << "-B | --hqws_bench             Benchs hqws" << std::endl
#endif // BENCH
      << "-h | --help                   Print this help" << std::endl
      << "-i | --input_uint8            Path to uint8 input image" << std::endl
      << "-I | --input_uint16           Path to uint16 input image" << std::endl
      << "-m | --marker_uint8           Path to uint8 marker image" << std::endl
      << "-M | --marker_uint16          Path to uint16 marker image"
      << std::endl
      << "-N | --marker_uint32          Path to uint32 marker image"
      << std::endl
      << "-o | --output_uint8           Path to uint8 output image" << std::endl
      << "-O | --output_uint16          Path to uint16 output image"
      << std::endl
      << "-P | --output_uint32          Path to uint32 output image"
      << std::endl
      << "-C | --cross                  4-connex structuring element"
      << std::endl
      << "-H | --hexa                   Hexagonal structuring element"
      << std::endl
      << "-S | --square                 8-connex structuring element"
      << std::endl
      << "-R | --cross3d                6-connex structuring element (3d)"
      << std::endl
      << "-U | --cube                   26-connex structuring element (3d)"
      << std::endl
      << "-x | --size_x                 Width dimension (only for raw images)"
      << std::endl
      << "-y | --size_y                 Height dimension (only for raw images)"
      << std::endl
      << "-z | --size_z                 Depth dimension (only for raw images)"
      << std::endl
      << "-w | --without_labels_limit   Ignore errors due to label number "
         "overflow"
      << std::endl
      << "-g | --export_RGB             Export the output image in RGB format."
      << std::endl
      << std::endl
      << "This program needs at least one input image, one output image and a "
         "structuring element. Only png and raw format are supported, and it "
         "should be identical for the input and the ouput. If raw format is "
         "used, dimensions must be specified. An additional markers image can "
         "be used for constrained watershed, whose type should match the "
         "output image.."
      << std::endl;
}

template <class Tin, class Tout>
int exec_paraws(const char *inp, const char *outp, const char *markp,
                std::string ext1, std::string ext2, std::string ext3, size_t &x,
                size_t &y, size_t &z, const structuring_element &se,
                bool ignore_label_limit, bool export_RGB) {
  Tin *im = NULL;
  Tout *out = NULL;

  int result = RES_OK;

  if (ext1 == "png") {
    io_png_read(im, inp, &x, &y);
  } else if (ext1 == "raw") {
    io_raw_read(im, inp, x, y, z);
  }

  if (im == NULL)
    return ERROR_IN;

  out = new Tout[x * y * z];

  if (markp) {
    Tout *markers = NULL;
    if (ext3 == "png") {
      io_png_read(markers, markp, &x, &y);
    } else if (ext3 == "raw") {
      io_raw_read(markers, markp, x, y, z);
    }
    if (markers == NULL)
      return ERROR_IN;
    result = paraws(im, markers, out, x, y, z, se, ignore_label_limit);
    delete[] markers;
  } else {
    result = paraws(im, out, x, y, z, se, ignore_label_limit);
  }

  if(result == RES_OK) {
    if (!export_RGB) {
      if (ext2 == "png") {
        if (io_png_write(outp, out, x, y) != 0)
          result = ERROR_IN;
      } else if (ext2 == "raw") {
        if (io_raw_write(outp, out, x, y, z) != 0)
          result = ERROR_IN;
      }
    } else {
      io_png_write_color (outp, out, x, y);
    }
  }

  delete[] im;
  delete[] out;
  return result;
}

template <class Tin, class Tout>
ERROR benchmark(size_t x, size_t y, size_t z, const structuring_element &se,
                bool ignore_label_limit, bool paraws_bench, bool hqws_bench) {
  Tin *im = new Tin[x * y * z];
  Tout *out = new Tout[x * y * z]();

  random_matrix(im, false, domain(x, y, z));

  ERROR result = RES_OK;

  if (paraws_bench) {
    result = paraws(im, out, x, y, z, se, ignore_label_limit);
    if (result != RES_OK) {
      delete[] im;
      delete[] out;
      return result;
    }
  }
  if (hqws_bench) {
    result = hqws(im, out, x, y, z, se, ignore_label_limit);
  }

  delete[] im;
  delete[] out;

  return result;
}

int main(const int __attribute__((__unused__)) argc, char **argv) {
  // Data types
  bool in8 = false, in16 = false, mark8 = false, mark16 = false, mark32 = false,
       out8 = false, out16 = false, out32 = false;
  // Connectivity
  bool cross = false, hexa = false, square = false, cross3d = false,
       cube = false;
#ifdef BENCH
  // Bench
  bool paraws_bench = false;
  bool hqws_bench = false;
#endif // BENCH
  // paths
  char *inp = NULL, *markp = NULL, *outp = NULL;
  // Dimensions
  size_t x = 0, y = 0, z = 1; // Default value for z.
  // Labels overflow
  bool ignore_label_limit = false;
  // Export the output image in RGB format for visualisation purposes.
  bool export_RGB = false;

  struct optparse_long long_options[] = {
#ifdef BENCH
      {"paraws_bench", 'A', OPTPARSE_NONE},
      {"hqws_bench", 'B', OPTPARSE_NONE},
#endif // BENCH
      {"input_uint8", 'i', OPTPARSE_REQUIRED},
      {"input_uint16", 'I', OPTPARSE_REQUIRED},
      {"marker_uint8", 'm', OPTPARSE_REQUIRED},
      {"marker_uint16", 'M', OPTPARSE_REQUIRED},
      {"marker_uint32", 'N', OPTPARSE_REQUIRED},
      {"output_uint8", 'o', OPTPARSE_REQUIRED},
      {"output_uint16", 'O', OPTPARSE_REQUIRED},
      {"output_uint32", 'P', OPTPARSE_REQUIRED},
      {"cross", 'C', OPTPARSE_NONE},
      {"hexa", 'H', OPTPARSE_NONE},
      {"square", 'S', OPTPARSE_NONE},
      {"cross3d", 'R', OPTPARSE_NONE},
      {"cube", 'U', OPTPARSE_NONE},
      {"size_x", 'x', OPTPARSE_REQUIRED},
      {"size_y", 'y', OPTPARSE_REQUIRED},
      {"size_z", 'z', OPTPARSE_REQUIRED},
      {"without_labels_limit", 'w', OPTPARSE_NONE},
      {"export_RGB", 'g', OPTPARSE_NONE},
      {"help", 'h', OPTPARSE_NONE},
      {0, 0, OPTPARSE_NONE}};

  // optparse variables.
  int option;
  struct optparse options;

  optparse_init(&options, argv);

  while ((option = optparse_long(&options, long_options, NULL)) != -1) {
    switch (option) {
#ifdef BENCH
    case 'A':
      paraws_bench = true;
      break;
    case 'B':
      hqws_bench = true;
      break;
#endif // BENCH
    case 'i':
      inp = options.optarg;
      in8 = true;
      break;
    case 'I':
      inp = options.optarg;
      in16 = true;
      break;
    case 'm':
      markp = options.optarg;
      mark8 = true;
      break;
    case 'M':
      markp = options.optarg;
      mark16 = true;
      break;
    case 'N':
      markp = options.optarg;
      mark32 = true;
      break;
    case 'o':
      outp = options.optarg;
      out8 = true;
      break;
    case 'O':
      outp = options.optarg;
      out16 = true;
      break;
    case 'P':
      outp = options.optarg;
      out32 = true;
      break;
    case 'C':
      cross = true;
      break;
    case 'H':
      hexa = true;
      break;
    case 'S':
      square = true;
      break;
    case 'R':
      cross3d = true;
      break;
    case 'U':
      cube = true;
      break;
    case 'x':
      x = atoi(options.optarg);
      break;
    case 'y':
      y = atoi(options.optarg);
      break;
    case 'z':
      z = atoi(options.optarg);
      break;
    case 'w':
      ignore_label_limit = true;
      break;
    case 'g':
      export_RGB = true;
      break;
    case 'h':
      break;
    case '?':
      std::cerr << argv[0] << ": " << options.errmsg << std::endl;
      exit(EXIT_FAILURE);
    }
  }

  // If arguments remain, then failure.
  if ((optparse_arg(&options))) {
    std::cerr << argv[0] << ": "
              << "Unknown arguments." << std::endl;
  }

  // Checking if everything is correct, before reading...
  bool is_compliant = true;
  is_compliant &= check_structuring_element(cross, hexa, square, cross3d, cube);

  // Loading the appropriate structuring element.
  structuring_element se;
  if (cross) {
    se = get_structuring_element(SE_CROSS);
  } else if (hexa) {
    se = get_structuring_element(SE_HEX);
  } else if (square) {
    se = get_structuring_element(SE_SQUARE);
  } else if (cross3d) {
    se = get_structuring_element(SE_CROSS_3D);
  } else if (cube) {
    se = get_structuring_element(SE_CUBE);
  }

  // For Benchs
#ifdef BENCH
  bool bench = paraws_bench || hqws_bench;
  if (bench) {
    // Something is wrong, print the help message and exit.
    if (!is_compliant) {
      usage(argv[0]);
      exit(EXIT_FAILURE);
    }
    if (x == 0 || y == 0) {
      usage(argv[0]);
      exit(EXIT_FAILURE);
    }
    if (benchmark<uint8_t, uint8_t>(x, y, z, se, ignore_label_limit,
                                    paraws_bench, hqws_bench) == RES_OK)
      exit(EXIT_SUCCESS);
    exit(EXIT_FAILURE);
  }
#endif // BENCH

  size_t x1 = x, y1 = y, z1 = z;
  std::string ext1;
  is_compliant &= check_compliance_image(inp, in8, in16, false, x1, y1, ext1);
  size_t x2 = x, y2 = y, z2 = z;
  std::string ext2;
  is_compliant &=
      check_compliance_image(outp, out8, out16, out32, x2, y2, ext2);
  is_compliant &= x1 == x2 && y1 == y2 && z1 == z2;
  is_compliant &= ext1 == ext2;

  // Color images are only available in png format.
  if (ext2 == "raw" && export_RGB)
      is_compliant = false;

  size_t x3 = x, y3 = y, z3 = z;
  std::string ext3;
  if (mark8 || mark16) {
    is_compliant &=
        check_compliance_image(markp, mark8, mark16, mark32, x3, y3, ext3);
    is_compliant &= ext2 == ext3;
    is_compliant &= (out8 == mark8) && (out16 == mark16) && (out32 == mark32);
    is_compliant &= x1 == x3 && y1 == y3 && z1 == z3;
  }

  // Something is wrong, print the help message and exit.
  if (!is_compliant) {
    usage(argv[0]);
    exit(EXIT_FAILURE);
  }

  int result = EXIT_SUCCESS;

  if (in8) {
    if (out8) {
      result = exec_paraws<uint8_t, uint8_t>(inp, outp, markp, ext1, ext2, ext3,
                                             x, y, z, se, ignore_label_limit, export_RGB);
    } else if (out16) {
      result = exec_paraws<uint8_t, uint16_t>(
          inp, outp, markp, ext1, ext2, ext3, x, y, z, se, ignore_label_limit, export_RGB);
    } else if (out32) {
      result = exec_paraws<uint8_t, uint32_t>(
          inp, outp, markp, ext1, ext2, ext3, x, y, z, se, ignore_label_limit, export_RGB);
    }
  } else if (in16) {
    if (out8) {
      result = exec_paraws<uint16_t, uint8_t>(
          inp, outp, markp, ext1, ext2, ext3, x, y, z, se, ignore_label_limit, export_RGB);
    } else if (out16) {
      result = exec_paraws<uint16_t, uint16_t>(
          inp, outp, markp, ext1, ext2, ext3, x, y, z, se, ignore_label_limit, export_RGB);
    } else if (out32) {
      result = exec_paraws<uint16_t, uint32_t>(
          inp, outp, markp, ext1, ext2, ext3, x, y, z, se, ignore_label_limit, export_RGB);
    }
  }
  if (result != 0) {
    if (result == ERROR_IN) {
      std::cerr << "Error: in/out error..." << std::endl;
      usage(argv[0]);
    } else if (result == ERROR_LABEL) {
      std::cerr << "Error: maximum label value reached..." << std::endl;
    }
    exit(EXIT_FAILURE);
  }
  exit(EXIT_SUCCESS);
}