#include "pngfile.h"
#include "../structures/colormap.h"

#include "../structures/image2d.h"

#include <stdexcept>

PngFile::PngFile(const std::string &filename, unsigned width, unsigned height)
    : _filename(filename), _width(width), _height(height), _pixelSize(4) {}

PngFile::~PngFile() {}

void PngFile::BeginWrite() {
  _fp = fopen(_filename.c_str(), "wb");
  if (!_fp) throw std::runtime_error("Can not open file");

  _png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, (png_voidp)NULL,
                                     NULL, NULL);

  if (!_png_ptr) {
    fclose(_fp);
    throw std::runtime_error("Can not create png write structure");
  }

  _info_ptr = png_create_info_struct(_png_ptr);
  if (!_info_ptr) {
    png_destroy_write_struct(&_png_ptr, (png_infopp)NULL);
    fclose(_fp);
    throw std::runtime_error("Can not write info structure to file");
  }

  if (setjmp(png_jmpbuf(_png_ptr))) {
    png_destroy_write_struct(&_png_ptr, &_info_ptr);
    fclose(_fp);
    throw std::runtime_error(
        "Unknown error occured during writing of png file");
  }

  png_init_io(_png_ptr, _fp);

  png_set_IHDR(_png_ptr, _info_ptr, _width, _height, 8,
               PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE,
               PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

  _row_pointers =
      (png_bytep *)png_malloc(_png_ptr, _height * sizeof(png_bytep));

  for (unsigned i = 0; i < _height; i++)
    _row_pointers[i] = (png_bytep)png_malloc(_png_ptr, _width * _pixelSize);
}

void PngFile::Close() {
  png_set_rows(_png_ptr, _info_ptr, _row_pointers);
  png_write_png(_png_ptr, _info_ptr, PNG_TRANSFORM_IDENTITY, NULL);
  png_write_end(_png_ptr, _info_ptr);

  for (unsigned i = 0; i < _height; i++) png_free(_png_ptr, _row_pointers[i]);
  png_free(_png_ptr, _row_pointers);

  png_destroy_write_struct(&_png_ptr, &_info_ptr);
  fclose(_fp);
}

void PngFile::Clear(int colorR, int colorG, int colorB, int colorA) {
  for (unsigned y = 0; y < _height; y++) {
    int xa = 0;
    for (unsigned x = 0; x < _width; x++) {
      _row_pointers[y][xa++] = colorR;
      _row_pointers[y][xa++] = colorG;
      _row_pointers[y][xa++] = colorB;
      _row_pointers[y][xa++] = colorA;
    }
  }
}

void PngFile::PlotDatapoint(unsigned x, unsigned y, int colorR, int colorG,
                            int colorB, int colorA) {
  int width = 3;
  for (int xi = (signed)x - width / 2; xi <= (signed)x + (width - 1) / 2;
       xi++) {
    if (xi >= 0 && xi < (signed)_width) {
      _row_pointers[y][xi * _pixelSize] = colorR;
      _row_pointers[y][xi * _pixelSize + 1] = colorG;
      _row_pointers[y][xi * _pixelSize + 2] = colorB;
      _row_pointers[y][xi * _pixelSize + 3] = colorA;
    }
  }
  for (int yi = (signed)y - width / 2; yi <= (signed)y + (width - 1) / 2;
       yi++) {
    if (yi >= 0 && yi < (signed)_height) {
      _row_pointers[yi][x * _pixelSize] = colorR;
      _row_pointers[yi][x * _pixelSize + 1] = colorG;
      _row_pointers[yi][x * _pixelSize + 2] = colorB;
      _row_pointers[yi][x * _pixelSize + 3] = colorA;
    }
  }
}

void PngFile::SetFromImage(const class Image2D &image,
                           const class ColorMap &colorMap,
                           long double normalizeFactor, long double zeroLevel) {
  png_bytep *row_pointers = RowPointers();
  for (unsigned long y = 0; y < image.Height(); y++) {
    for (unsigned long x = 0; x < image.Width(); x++) {
      int xa = x * PixelSize();
      row_pointers[y][xa] = colorMap.ValueToColorR(
          (image.Value(x, y) - zeroLevel) * normalizeFactor);
      row_pointers[y][xa + 1] = colorMap.ValueToColorG(
          (image.Value(x, y) - zeroLevel) * normalizeFactor);
      row_pointers[y][xa + 2] = colorMap.ValueToColorB(
          (image.Value(x, y) - zeroLevel) * normalizeFactor);
      row_pointers[y][xa + 3] = colorMap.ValueToColorA(
          (image.Value(x, y) - zeroLevel) * normalizeFactor);
    }
  }
}

void PngFile::Save(const Image2D &image, const std::string &filename) {
  std::unique_ptr<ColorMap> colorMap(ColorMap::CreateColorMap("monochrome"));
  Save(image, filename, *colorMap);
}

void PngFile::Save(const Image2D &image, const std::string &filename,
                   const ColorMap &colorMap) {
  Save(image, filename, colorMap, image.GetMaxMinNormalizationFactor());
}

void PngFile::Save(const Image2D &image, const std::string &filename,
                   const ColorMap &colorMap, long double normalizeFactor,
                   long double zeroLevel) {
  PngFile pngFile(filename, image.Width(), image.Height());
  pngFile.BeginWrite();
  pngFile.SetFromImage(image, colorMap, normalizeFactor, zeroLevel);
  pngFile.Close();
}

void PngFile::Save(const Image2D &image, const ColorMap &colorMap) {
  long double normalizeFactor = image.GetMaxMinNormalizationFactor();

  png_bytep *row_pointers = RowPointers();

  for (unsigned long y = 0; y < image.Height(); ++y) {
    for (unsigned long x = 0; x < image.Width(); ++x) {
      int xa = x * PixelSize();
      row_pointers[y][xa] =
          colorMap.ValueToColorR(image.Value(x, y) * normalizeFactor);
      row_pointers[y][xa + 1] =
          colorMap.ValueToColorG(image.Value(x, y) * normalizeFactor);
      row_pointers[y][xa + 2] =
          colorMap.ValueToColorB(image.Value(x, y) * normalizeFactor);
      row_pointers[y][xa + 3] =
          colorMap.ValueToColorA(image.Value(x, y) * normalizeFactor);
    }
  }
}
