#define TINYEXR_IMPLEMENTATION
#include "tinyexr.h"

#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Weverything"
#endif

#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include "stb_image_resize.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#include "cxxopts.hpp"

#ifdef __clang__
#pragma clang diagnostic pop
#endif

namespace {

static void vnormalize(float v[3]) {
  const float d2 = v[0] * v[0] + v[1] * v[1] + v[2] * v[2];
  if (d2 > 1.0e-6f) {
    const float inv_d = 1.0f / std::sqrt(d2);
    v[0] *= inv_d;
    v[1] *= inv_d;
    v[2] *= inv_d;
  }
  return;
}


template<typename T>
static inline T clamp(const T v, const T min_v, const T max_v)
{
  return std::max(min_v, std::min(max_v, v));
}

//
// Compute gradient from scalar field.
// dx = (x + 1, y    ) - (x, y) 
// dy = (x    , y + 1) - (x, y) 
//
// TODO(syoyo): Use central difference with texel filtering.
//
static void Gradient(
  const std::vector<float> &src,
  const size_t width,
  const size_t height,
  const size_t x, const size_t y,
  const float bumpness,
  float dir[3])
{
  const size_t x1 = clamp(x + 1, size_t(0), width - 1);
  const size_t y1 = clamp(y + 1, size_t(0), height - 1);


  float v00 = src[y * width + x];
  float v01 = src[y * width + x1];
  float v11 = src[y1 * width + x];

  
  float dx = bumpness * (v01 - v00);
  float dy = bumpness * (v11 - v00);

  dir[0] = dx;
  dir[1] = dy;
  dir[2] = 0.0f;
    
}

///
/// Convert image(bump map for single channel, vector displacement map for 3 channels input) to normal map.
/// @param[in] base Base value fo
///
///
static void ToNormalMap(
  const std::vector<float> &src,
  const size_t width,
  const size_t height,
  const size_t channels,
  const float strength,
  std::vector<float> *dst)
{
  assert((channels == 1) || (channels == 3) || (channels == 4));

  dst->resize(width * height * 3);

  if (channels == 1) {
    // bump map
    for (size_t y = 0; y < height; y++) {
      for (size_t x = 0; x < width; x++) {
        float d[3];
        Gradient(src, width, height, x, y, strength, d);

        (*dst)[3 * (y * width + x) + 0] = d[0];
        (*dst)[3 * (y * width + x) + 1] = d[1];
        (*dst)[3 * (y * width + x) + 2] = d[2];

      }
    }

  } else {
    // vector displacement map

    for (size_t y = 0; y < height; y++) {
      for (size_t x = 0; x < width; x++) {

        float v[3];
        v[0] = src[channels * (y * width + x) + 0];
        v[1] = src[channels * (y * width + x) + 1];
        v[2] = src[channels * (y * width + x) + 2];

        v[0] *= strength;
        v[1] *= strength;
        v[2] *= strength;

        // Add (0, 0, 1)
        v[2] += 1.0f;

        // TODO(syoyo): Add option to not normalize.
        vnormalize(v);

        (*dst)[3 * (y * width + x) + 0] = 0.5f * v[0] + 0.5f;
        (*dst)[3 * (y * width + x) + 1] = 0.5f * v[1] + 0.5f;
        (*dst)[3 * (y * width + x) + 2] = 0.5f * v[2] + 0.5f;

      }
    }

  }

}

inline unsigned char ftouc(float f)
{
  int i = static_cast<int>(f * 255.0f);
  if (i > 255) i = 255;
  if (i < 0) i = 0;

  return static_cast<unsigned char>(i);
}

bool SaveImage(const char* filename, const float* rgb, int width, int height) {

  std::vector<unsigned char> dst(width * height * 3);

  for (size_t i = 0; i < width * height; i++) {
      dst[i * 3 + 0] = ftouc(rgb[i * 3 + 0]);
      dst[i * 3 + 1] = ftouc(rgb[i * 3 + 1]);
      dst[i * 3 + 2] = ftouc(rgb[i * 3 + 2]);
  }

  int ret = stbi_write_png(filename, width, height, 3, static_cast<const void*>(dst.data()), width * 3);

  return (ret > 0);
}

std::string GetFileExtension(const std::string &filename) {
  if (filename.find_last_of(".") != std::string::npos)
    return filename.substr(filename.find_last_of(".") + 1);
  return "";
}

} // namespace

int main(int argc, char **argv)
{
  cxxopts::Options options("normalmap", "help");
  options.add_options()
    ("s,strength", "Strength(scaling) for normal value", cxxopts::value<float>())
    ("i,input", "Input filename", cxxopts::value<std::string>())
    ("o,output", "Output filename", cxxopts::value<std::string>())
    ("r,resize", "Resize image. 0.5 = 50%%, 0.1 = 10%%", cxxopts::value<float>())
    ;

  auto result = options.parse(argc, argv);

  if (result.count("input") == 0) {
    std::cerr << "input filename missing" << std::endl;
    return EXIT_FAILURE;
  }

  if (result.count("output") == 0) {
    std::cerr << "output filename missing" << std::endl;
    return EXIT_FAILURE;
  }

  float strength = 1.0f;
  if (result.count("strength")) {
    strength = result["strength"].as<float>();
  }

  float resize = 1.0f;
  if (result.count("resize")) {
    resize = result["resize"].as<float>();
  }

  std::string input_filename = result["input"].as<std::string>();
  std::string output_filename = result["output"].as<std::string>();

  std::vector<float> src;
  size_t src_width;
  size_t src_height;

  {
    float *rgba = nullptr;
    int width, height;
    const char *err = nullptr;
    int ret = LoadEXR(&rgba, &width, &height, input_filename.c_str(), &err);
    if (TINYEXR_SUCCESS != ret) {
      std::cerr << "Failed to load EXR file [" << input_filename << "] code = " << ret << std::endl;
      if (err) {
        std::cerr << err << std::endl;
        FreeEXRErrorMessage(err);
      }
      
      return EXIT_FAILURE;
    }

    std::cout << "loaded EXR. width x height = " << width << "x" << height << std::endl;
    src.resize(size_t(width * height * 3));

    // ignore alpha for now
    for (size_t i = 0; i < size_t(width * height); i++) {
      src[3 * i + 0] = rgba[4 * i + 0];
      src[3 * i + 1] = rgba[4 * i + 1];
      src[3 * i + 2] = rgba[4 * i + 2];
    }

    src_width  = size_t(width);
    src_height = size_t(height);

    free(rgba);
  }

  std::cout << "strength = " << strength << std::endl;

  std::vector<float> dst;
  ToNormalMap(src, src_width, src_height, 3, strength, &dst);


  std::string ext = GetFileExtension(output_filename);
  if ((ext.compare("png") == 0) ||
      (ext.compare("PNG") == 0)) {
    // Save as LDR image.
    // Do not apply sRGB conversion for PNG(LDR) image.
    if (!SaveImage(output_filename.c_str(), dst.data(), int(src_width), int(src_height))) {
      std::cerr << "Failed to write a file : " << output_filename << std::endl;
      return EXIT_FAILURE;
    } 
  } else {
    // assume EXR.
    float *rgba = nullptr;
    int width, height;
    int ret = SaveEXR(dst.data(), int(src_width), int(src_height), /* component */3, /* fp16 */0, output_filename.c_str(), nullptr);
    if (TINYEXR_SUCCESS != ret) {
      std::cerr << "Failed to save EXR file [" << input_filename << "] code = " << ret << std::endl;
      return EXIT_FAILURE;
    }
  }

  return EXIT_SUCCESS;
}