#include <gsl/gsl_matrix_double.h>
#include <stdint.h>
#include <gsl/gsl_odeiv2.h>
#include <gsl/gsl_matrix.h>
#include <spng.h>

#include "common.h"

// ODE
// du/dphi = v
// dv/dphi = 3Mu^2 - u
int func(double phi, const double y[], double f[], void *params){
  f[0] = y[1];
  f[1] = -y[0] + 3.0*M*y[0]*y[0];
  return GSL_SUCCESS;
}

int jac(double phi, const double y[], double *dfdy, double dfdphi[], void *params){
  gsl_matrix_view dfdy_mat = gsl_matrix_view_array(dfdy, 2, 2);
  gsl_matrix *m = &dfdy_mat.matrix;
  gsl_matrix_set(m, 0, 0, 0.0);
  gsl_matrix_set(m, 0, 1, 1.0);
  gsl_matrix_set(m, 1, 0, -3*M*y[0]-1.0);
  gsl_matrix_set(m, 1, 1, 0.0);

  dfdphi[0] = 0.0;
  dfdphi[0] = 0.0;

  return GSL_SUCCESS;
}

double find_root(double *y, double phi){
  //printf("phi: %g, y:{%g, %g}\n", phi, y[0], y[1]);
  return phi - y[0]/y[1];
}

double find_root_4ord(double u1, double u2, double v1, double v2, double phi1, double phi2){
  double Du = u1 - u2;
  double Dphi = phi2 - phi1;
  return -(u2*v1*Du*Du*Du - v1*v2*Du*Du*Du*phi2 + u2*u2*Du*(2*v1*Du + v2*Du + 3*v1*v2*Dphi) + u2*u2*u2*(v2*Du + v1*(Du + 2*v2*Dphi)))/(v1*v2*Du*Du*Du);
}
inline double find_root_5ord(double u1, double u2, double v1, double v2, double phi1, double phi2){
  return (-(phi1*pow(u2,5)*v1*v2) + 
     u1*pow(u2,4)*(u2 - phi1*v1)*
      v2 + pow(u1,2)*pow(u2,3)*
      (u2 + 8*phi1*v1)*v2 + 
     pow(u1,4)*u2*v1*
      (-u2 + phi2*v2) + 
     2*pow(u1,3)*pow(u2,2)*
      (u2*(v1 - v2) - 4*phi2*v1*v2)
       + pow(u1,5)*
      (-(u2*v1) + phi2*v1*v2))/
   (pow(u1 - u2,3)*
     (pow(u1,2) + 4*u1*u2 + 
       pow(u2,2))*v1*v2);
}

// \chi(b) is the totol deflection angle with parameter b
double chi(double b){
  double x;
  gsl_odeiv2_system sys = {func, jac, 2, &x};
  gsl_odeiv2_driver *d = gsl_odeiv2_driver_alloc_y_new(&sys, gsl_odeiv2_step_rk8pd, 1e-3, 1e-12, 1e-10);
  
  double u0 = 1.0/R0;
  double y[2] = {
    u0,
    sqrt(1.0/(b*b) - u0*u0 + 2.0*M*u0*u0*u0)
  };
  double phi = 0;
  double phi_max = 8*PI;
  double y_previous[2] = {0.0, 0.0};
  double phi_previous=phi;

  while (y[0] > 0 && phi < phi_max){
    double phi_next = phi + 0.1;
    phi_previous = phi;
    y_previous[0] = y[0];
    y_previous[1] = y[1];
    int status = gsl_odeiv2_driver_apply(d, &phi, phi_next, y);
    if (status != GSL_SUCCESS){
      printf("Error: GSL driver returns %d\n", status);
      break;
    }
    //printf("phi: %g, y[0]: %g, y[1]: %g\n",phi, y[0], y[1]);
  }

  //printf("%g, %g \n",find_root(y_previous, phi_previous),find_root(y, phi));
  //printf("%.16g\n",find_root_4ord(y_previous[0], y[0], y_previous[1], y[1], phi_previous, phi));
  //printf("%.16g\n",find_root_5ord(y_previous[0], y[0], y_previous[1], y[1], phi_previous, phi));
  //return (find_root(y_previous, phi_previous) + find_root(y, phi))/2.0;

  gsl_odeiv2_driver_free(d);
  if (phi >= phi_max){
    return phi_max;
  }
  return find_root_4ord(y_previous[0], y[0], y_previous[1], y[1], phi_previous, phi);
}

// (x,y) on screen to far away angle
double phi(double x, double y){
  return atan2(y,x);
}

double theta(double x, double y){
  //double psi = atan(sqrt(x*x+y*y));
  double tanpsi2 = x*x+y*y;
  double cotpsi2 = 1/tanpsi2;
  double b = R0/sqrt(f(R0)*(1+cotpsi2));

  if (b < bmin){
    return THETAERROR;
  }
  return chi(b);
}

int xy_to_angle_BH(double *xy, double *angle){
  angle[1] = phi(xy[0], xy[1]);
  angle[0] = theta(xy[0], xy[1]);
  int flag = (angle[0]==THETAERROR);
  //int i = ((int) floor(angle[1]/PI))&1;
  //if (i==1){
  //  angle[1] = 2*PI - angle[1];
  //}
  //printf("xy=(%g,%g), thetaphi=(%g,%g)\n",xy[0],xy[1],angle[0],angle[1]);
  return flag;
}
int xy_to_angle_flat(double *xy, double *angle){
  angle[1] = atan2(xy[1], xy[0]);
  angle[0] = PI - atan(sqrt(xy[0]*xy[0] + xy[1]*xy[1]));
  return 0;
}

// angle to pixel
int test(double *thetaphi, double *rgb){
  int i = floor(8.0*thetaphi[0]/PI);
  int j = floor(8.0*(thetaphi[1]+PI)/PI);
  int k = (i+j)&1;
  int l = floor(2*thetaphi[1]/PI);
  rgb[0] = k;
  rgb[1] = k;
  rgb[2] = k;
  switch (l%4) {
    case 1:
      rgb[1] = rgb[2] = 0;
      break;
    case 2:
      rgb[0] = rgb[2] = 0;
      break;
    case 3:
      rgb[0] = rgb[1] = 0;
      break;
    default:
      break;
  }
  return 0;
}

int test2(double *thetaphi, double *rgb){
  double tantheta = tan(PI-thetaphi[0]);
  int i = (int)floor(4*tantheta*cos(thetaphi[1]));
  int j = (int)floor(4*tantheta*sin(thetaphi[1]));
  int k = (i+j)&1;
  rgb[0] = 1-k;
  rgb[1] = 1-k;
  rgb[2] = 1-k;
  return 0;
}

int test3(double *thetaphi, double *rgb){
  double x = sin(thetaphi[0])*sin(thetaphi[1]);
  double y = sin(thetaphi[0])*cos(thetaphi[1]);
  double z = cos(thetaphi[0]);
  double theta1 = atan2(hypot(z, x), y);
  double phi1 = atan2(-z, x);
  int i = (int)floor(16*theta1/PI);
  int j = (int)floor(16*phi1/PI);
  int k = (i+j)&1;
  rgb[0] = 1-k;
  rgb[1] = 1-k;
  rgb[2] = 1-k;
  return 0;
}

int testtheta(double *thetaphi, double *rgb){
  rgb[0] = thetaphi[0]/PI;
  rgb[0] -= floor(rgb[0]);
  rgb[0] *= rgb[0];
  rgb[1] = 0;
  rgb[2] = 0;
  return 0;
}

int testphi(double *thetaphi, double *rgb){
  if (thetaphi[1] >=0 && thetaphi[1] <= PI/3){
    rgb[0] = 1;
  }
  else {
    rgb[0] = 0;
  }
  rgb[1] = 0;
  rgb[2] = 0;
  return 0;
}

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 render(char *filename, int W, int H, double w, int angle_to_pixel(double *, double *), int xy_to_angle(double *, double*)){
  size_t bufsize = W*H*3;
  uint8_t *img = (uint8_t *)malloc(bufsize);
  if (!img) {return 1;}

  double xy[2]={0};
  double angle[2] = {0};
  int flag = 0;
  double rgb_lin[3] = {0};

  for (int i = 0; i < W; i++){
    xy[0] = (w/W)*(i + 0.5 - W/2.0);
    for(int j=0; j<H;j++){
      xy[1] = -(w/W)*(j + 0.5 - H/2.0);
      flag = xy_to_angle(xy, angle);
      if (flag == 0){
        flag = angle_to_pixel(angle, rgb_lin);
        if (flag != 0){
          rgb_lin[0] = 1;
          rgb_lin[1] = 0;
          rgb_lin[2] = 0;
        }
      }
      else{
        rgb_lin[0] = 0;
        rgb_lin[1] = 1;
        rgb_lin[2] = 0;
      }
      //printf("i=%d,j=%d,xy=(%g,%g),angle=(%g,%g),rgb=(%g,%g,%g)\n",i,j,xy[0],xy[1],angle[0],angle[1],rgb_lin[0],rgb_lin[1],rgb_lin[2]);
      for(int c=0; c<3; ++c){
        double sr = linear_to_srgb(rgb_lin[c]);
        int v = (int)lrint(sr * 255.0);
        if(v<0) v=0;
        if(v>255) v=255;
        img[(size_t)j*W*3+ (size_t)i*3 + c] = (uint8_t)v;
      }
    }
  }

  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;
}

// Main
int main(){
  int W = 3840;
  int H = 2160;
  double w = 2.56;
  //xy_to_angle_flat(xy, angle);
  //printf("xy=(%g,%g), angle=(%g,%g)\n", xy[0], xy[1], angle[0], angle[1]);
  render("test_flat.png", W, H, w, test, xy_to_angle_flat);
  render("test_BH.png", W, H, w, test, xy_to_angle_BH);
  //render("test2_flat.png", W, H, w, test2, xy_to_angle_flat);
  //render("test2_BH.png", W, H, w, test2, xy_to_angle_BH);
  //render("test3_flat.png", W, H, w, test3, xy_to_angle_flat);
  //render("test3_BH.png", W, H, w, test3, xy_to_angle_BH);
  //render("testtheta_flat.png", W, H, w, testtheta, xy_to_angle_flat);
  //render("testtheta_BH.png", W, H, w, testtheta, xy_to_angle_BH);
}