#include <spng.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include "write_png.h"
#include "common.h"
#include "quick_select.h"

static inline double linear_to_srgb(double x){
  if (x <= 0.0) return 0.0;
  if (x >= 1.0) return 1.0;
  if (x <= 0.0031308) return 12.92 * x;
  return 1.055 * pow(x, 1.0/2.4) - 0.055;
}

int buffer_normalize_srgb(double *buffer, int W, int H) {
  double max = 0;
  //for (int j = 0; j < H; j++) {
  //  for (int i = 0; i < W; i++) {
  //    for (int c = 0; c < 3; c++) {
  //      if (max < buffer[color_index(i, j, c)]) max = buffer[color_index(i, j, c)];
  //    }
  //  }
  //}
  max = percentile(buffer, W*H*3, cutperc);
  printf("max in the buffer: %g\n", max);
  if (max == 0) return 1;
  #pragma omp parallel for
  for (int j = 0; j < H; j++) {
    for (int i = 0; i < W; i++) {
      for (int c = 0; c < 3; c++) {
        buffer[color_index(i, j, c)] /= max;
        buffer[color_index(i, j, c)] = linear_to_srgb(buffer[color_index(i, j, c)]);
      }
    }
  }
  return 0;
}

int write_png(char *filename, double *buffer, int W, int H) {
  buffer_normalize_srgb(buffer, W, H);
  size_t bufsize = W*H*3;
  //printf("bufsize: %d\n", bufsize);
  uint8_t *img = (uint8_t *)malloc(bufsize*sizeof(uint8_t));
  if (!img) {return 1;}

  for (int j = 0; j < H; j++) {
    for (int i = 0; i < W; i++) {
      for (int c = 0; c < 3; c++) {
        int v = (int)lrint(buffer[color_index(i,j,c)]*255);
        if (v < 0) v = 0;
        if (v > 255) v = 255;
        img[color_index(i, j, c)] = (uint8_t) v;
        //printf("i: %d, j: %d, c:%d, buffer=%g, v=%d\n", i,j,c,buffer[color_index(i, j, c)],v);
      }
    }
  }

  int flag = 0;
  spng_ctx *ctx = spng_ctx_new(SPNG_CTX_ENCODER);
  if(!ctx){ fprintf(stderr, "spng_ctx_new failed\n"); free(img); return 1; }

  FILE *f = fopen(filename, "wb");
  if(!f){ fprintf(stderr, "fopen failed\n"); spng_ctx_free(ctx); free(img); return 1; }
  spng_set_png_file(ctx, f);

  struct spng_ihdr ihdr = {0};
  ihdr.width = W;
  ihdr.height = H;
  ihdr.bit_depth = 8;
  ihdr.color_type = SPNG_COLOR_TYPE_TRUECOLOR; // RGB
  ihdr.interlace_method = SPNG_INTERLACE_NONE;
  ihdr.compression_method = 0;
  ihdr.filter_method = 0;
  if( (flag= spng_set_ihdr(ctx, &ihdr)) ){
    fprintf(stderr, "spng_set_ihdr: %s\n", spng_strerror(flag));
    fclose(f); spng_ctx_free(ctx); free(img); return 1;
  }

  #ifndef SPNG_SRGB_INTENT_PERCEPTUAL
  #define SPNG_SRGB_INTENT_PERCEPTUAL 0
  #endif
  spng_set_srgb(ctx, SPNG_SRGB_INTENT_PERCEPTUAL);

  flag = spng_encode_image(ctx, img, bufsize, SPNG_FMT_PNG, SPNG_ENCODE_FINALIZE);
  if(flag){
    fprintf(stderr, "spng_encode_image: %s\n", spng_strerror(flag));
  }
  fclose(f);
  spng_ctx_free(ctx);
  free(img);
  return flag? 1 : 0;
}