#include <assert.h>
#include <complex.h>
#include <getopt.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <png.h>
#include <pthread.h>

typedef uint32_t pixel_t;

typedef struct {
  pixel_t *pixels;
  size_t width;
  size_t height;
} bitmap_t;

/* The maximum iteration count. */
const int max_it = 170;

/**
 * Compute the out-coloring based on the iteration counter.
 *
 * Out-coloring means: coloring when a pixel is determined
 * to be outside of the mandelbrot set.
 *
 * @param it iteration counter when diversion was determined.
 *
 * @return pixel color
 */
pixel_t outcolor(int it) {
  return 0x00010001 * ((it * 0xff) / max_it);
}

/**
 * "Compute" the in-coloring.
 *
 * In-coloring means: coloring when a pixel is determined
 * to be inside the mandelbrot set.
 *
 * @return pixel, i.e. color
 */
pixel_t incolor() {
  return 0x00000000;            /* black */
}


/* Mandelbrot bitmap. */
bitmap_t *mbrot;

/**
 * Writes a bitmap to a PNG file.
 *
 * @param bitmap bitmap to save
 * @param path   PNG file to write to
 *
 * @return 0 on success, non-zero on error.
 */
static int save_png_to_file(bitmap_t * bitmap, const char *path) {
  FILE *fp;

  png_structp png_ptr = NULL;
  png_infop info_ptr = NULL;
  size_t x, y;
  png_byte **row_pointers = NULL;

  /* "status" contains the return value of this function. At first it is set to
   * a value which means 'failure'. When the routine has finished its work, it
   * is set to a value which means 'success'. */
  int status = -1;

  int pixel_size = 3;
  int depth = 8;

  fp = fopen(path, "wb");
  if (!fp) {
    goto fopen_failed;
  }

  png_ptr =
    png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  if (png_ptr == NULL) {
    goto png_create_write_struct_failed;
  }

  info_ptr = png_create_info_struct(png_ptr);
  if (info_ptr == NULL) {
    goto png_create_info_struct_failed;
  }

  /* Set up error handling. */
  if (setjmp(png_jmpbuf(png_ptr))) {
    goto png_failure;
  }

  /* Set image attributes. */
  png_set_IHDR(png_ptr,
      info_ptr,
      bitmap->width,
      bitmap->height,
      depth,
      PNG_COLOR_TYPE_RGB,
      PNG_INTERLACE_NONE,
      PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

  /* Initialize rows of PNG. */
  row_pointers = png_malloc(png_ptr, bitmap->height * sizeof(png_byte *));
  for (y = 0; y < bitmap->height; ++y) {
    png_byte *row =
      png_malloc(png_ptr, sizeof(uint8_t) * bitmap->width * pixel_size);
    row_pointers[y] = row;
    for (x = 0; x < bitmap->width; ++x) {
      pixel_t *pixel = bitmap->pixels + (y * bitmap->width + x);
      *row++ = (*pixel & 0x00ff0000) >> 16;
      *row++ = (*pixel & 0x0000ff00) >> 8;
      *row++ = (*pixel & 0x000000ff);
    }
  }

  /* Write the image data to "fp". */
  png_init_io(png_ptr, fp);
  png_set_rows(png_ptr, info_ptr, row_pointers);
  png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);

  /* The routine has successfully written the file, so we set "status" to a
   * value which indicates success. */
  status = 0;

  for (y = 0; y < bitmap->height; y++) {
    png_free(png_ptr, row_pointers[y]);
  }
  png_free(png_ptr, row_pointers);

png_failure:
png_create_info_struct_failed:
  png_destroy_write_struct(&png_ptr, &info_ptr);
png_create_write_struct_failed:
  fclose(fp);
fopen_failed:
  return status;
}

int main() {

  /* Set up mandelbrot bitmap. */
  mbrot = calloc(sizeof(bitmap_t), 1);
  mbrot->width = 1600;
  mbrot->height = 1200;
  mbrot->pixels = calloc(sizeof(pixel_t), mbrot->width * mbrot->height);
  assert(mbrot->pixels != NULL);

  /* Compute. */
  #pragma omp parallel for schedule (dynamic, 50000)
  for (unsigned int y = 0; y < mbrot->height; y++) {
    for (size_t x = 0; x < mbrot->width; x++) {

      float complex c = ((3.0f * x / mbrot->width) - 2.0f)
        + I * ((2.0f * y / mbrot->height) - 1.0f);

      bool diverges = false;
      float complex z = 0;

      int it;
      for (it = 1; it <= max_it; it++) {
        /* z = z² + c */
        z = cpowf(z, 2) + c;

        /* If |z| ever gets greater than 2, it diverges. */
        if (cabsf(z) > 2) {
          diverges = true;
          break;
        }
      }

      uint32_t color;
      if (diverges) {
        color = outcolor(it);
      } else {
        color = incolor();
      }

      size_t i = y * mbrot->width + x;
      mbrot->pixels[i] = color;
    }
  }

  /* Save PNG. */
  char *file = "multibrot-openmp.png";
  int ret = save_png_to_file(mbrot, file);
  assert(ret == 0);

  /* Clean up. */
  free(mbrot->pixels);
  free(mbrot);

  return 0;
}