#include <math.h>
#include <stdio.h>
#include <time.h>
#include <stdlib.h>

double LambertW(double z);
double LambertWm1(double z);
double ScaledWoodsSaxon(double x, double XA);
double GetPmax(double XA);
double ScaledWoodsSaxon(double x, double XA);
double ScaledWoodsSaxonEnv(double x, double XA, double Pmax, double x1, double x2);
double GetX1(double Pmax);
double GetX2(double XA, double Pmax);
double GenRandWS_r(double XA, double Pmax, double x1, double x2);
int binning(double x, double xi, double xf, double *bins, int num_bins);

double genrand_real(); // Uses C's own random number generator. Replace this with a better one.

/*
Generates random numbers according to the Woods-Saxon distribution
rho_WS = rho_0/(1 + exp((r-R_A)/a))
*/

// Test driver
int main(void)
{
 double third, aA, RA, XA, Pmax, x1, x2, A, x, Dx;
 double *bins;
 int iter, itermax, num_bins;
 FILE *output;

 fprintf(stderr, "RAND_MAX = %ld\n", RAND_MAX);
 
 third = 1.0/3.0;

 printf("A = ");
 scanf("%lf", &A);

 aA = 0.55; // in fm
 RA = 1.12*pow(A, third); // in fm
 
 fprintf(stdout, "RA = %e fm\n", RA);

 XA = RA/aA;
 
 fprintf(stdout, "XA = %e fm\n", XA);

 Pmax = GetPmax(XA);
 x1 = GetX1(Pmax);
 x2 = GetX2(XA, Pmax);

 fprintf(stdout, "Pmax = %e\n", Pmax);
 fprintf(stdout, "x1 = %e\n", x1);
 fprintf(stdout, "x2 = %e\n", x2);
 fprintf(stdout, "R1 = %e\n", x1*aA);
 fprintf(stdout, "R2 = %e\n", x2*aA);

 itermax = 10000;

 num_bins = 25;
 bins = (double *) malloc(sizeof(double)*num_bins);
 
 for(iter = 0; iter < itermax; iter++)
  {
   x = GenRandWS_r(XA, Pmax, x1, x2);  
   binning(x, 0, 2.0*XA, bins, num_bins);
  }// iter
 
 output = fopen("WS_dist.dat","w");
 Dx = 2.0*XA/((double) num_bins);
 for(iter = 0; iter < num_bins; iter++)
  {
   x = (iter + 0.5)*Dx;
   fprintf(output, "%e  %e\n", x*aA, bins[iter]/Dx/aA);
  }// iter

 fclose(output); 
 return 1;
}// main


// generate random numbers according to P(x) = x^2/(1 + exp(x-XA))

double GenRandWS_r(double XA, double Pmax, double x1, double x2)
{
 double x, y, z, z1, z2, z3, z_arg, ratio;
 double third = 1.0/3.0;

 z1 = pow(x1, 3.0)/3.0;
 z2 = z1 + Pmax*(x2 - x1);
 z3 = z2 + (2.0+x2)*(2.0+x2)*exp(-x2+XA);

do
{
 z = z3*genrand_real();

 if( (0 <= z) && (z < z1) ) // zone 1
  {
   x = pow(3.0*z, third);
  }
 else if( (z1 <= z) && (z < z2) ) // zone 2
  {
   x = x1 + (z-z1)/Pmax;
  }
 else // zone 3
  {
   z_arg = -0.5*sqrt(z3 - z)*exp(-(XA+2.0)/2.0);
   x = -2.0*(LambertWm1(z_arg) + 1.0); 
  }

 y = genrand_real();
 ratio = ScaledWoodsSaxon(x, XA)/ScaledWoodsSaxonEnv(x, XA, Pmax, x1, x2);
} while(ratio < y);

 return x;
}// GenRandWS_r


double GetX1(double Pmax)
{
 return sqrt(Pmax);
}// GetX1


// Solve  f(x) = x(x+2) e^(-x+XA) = Pmax 
// using Newton's method
// x(n) = x(n-1) + (Pmax - f(n-1))/f'(n-1) 
// Derivative = (2 - x*x)exp(-x+XA)

double GetX2(double XA, double Pmax)
{
 double x_now, x_next, num, denom, err, TOL;
 int iter;

// initial value
 x_now = sqrt(Pmax);
 iter = 0;
 TOL = 1.0e-14;

 do
  {
   iter++;
   num = Pmax - x_now*(x_now + 2.0)*exp(-x_now+XA);
   denom = (2.0 - x_now*x_now)*exp(-x_now+XA);
   x_next = x_now + num/denom;

   err = fabs(x_next - x_now)/fabs(x_next + x_now);
   x_now = x_next;
  } while(err > TOL);
 
 return x_now;
}// GetX2


double GetPmax(double XA)
{
 double f, x, y;

 y = 2.0*exp(XA-2.0);
 
 x = LambertW(y) + 2.0;

 f = ScaledWoodsSaxon(x, XA);
 
 return f;
}// GetPmax


double ScaledWoodsSaxon(double x, double XA)
{
 double f;

 f = x*x/(1.0 + exp(x - XA));

 return f;
}// ScaledWoodsSaxon


double ScaledWoodsSaxonEnv(double x, double XA, double Pmax, double x1, double x2)
{
 double f;

 if( (0 <= x) && (x < x1) ) // region 1
  {
   return x*x;
  }
 else if( (x1 <= x) && (x < x2) ) // region 2
  {
   return Pmax;
  }
 else // region 3
  {
   return x*(x+2.0)*exp(-x+XA);
  }
 
 return f;
}// ScaledWoodsSaxonEnv


// Solves y = x e^x
// This is the principal branch
double LambertW(double z)
{
 double w_new, w_old, ratio, e_old, tol, logz;
 double a1, a2, a3, a4, b1, b2, b3, b4;
 int n;

 /* iteration */

 tol = 1.0e-14;

 if(z <= -exp(-1.0))
  {
   fprintf(stderr, "LambertW is not defined for z = %e\n", z);
   fprintf(stderr, "z needs to be bigger than %e\n", -exp(-1.0));
   exit(0);
  }

// approximate values to initiate iteration process

 if(z > 3.0)
  {
   logz = log(z);
   w_old = logz - log(logz)*(1.0-1.0/logz);
  }
 else if( (0 <= z) && (z <= 3.0) )
  {
   w_old = 0.88921*atan(0.74423*z);
  }
 else 
  {
   w_old = 1.5*( sqrt(z + exp(-1.0))-exp(-0.5) );
  }

 w_new = 0.0;
 ratio = 1.0;
 n = 0;
 while(fabs(ratio) > tol) 
  {
   e_old = exp(w_old);
   // Halley iteration 
   ratio = w_old*e_old - z;
   ratio /= ( e_old*(w_old + 1.0) - (w_old+2.0)*(w_old*e_old-z)/(2.0*w_old + 2.0) );

   w_new = w_old - ratio;
   w_old = w_new;
   n++;
   if(n > 99) 
    {
     fprintf(stderr, "LambertW is not converging after 100 iterations.\n");
     fprintf(stderr, "LambertW: z = %e\n", z);
     fprintf(stderr, "LambertW: w_old = %e\n", w_old);
     fprintf(stderr, "LambertW: w_new = %e\n", w_new);
     fprintf(stderr, "LambertW: ratio = %e\n", ratio);
     exit(0);
    }
  }

 return w_new;
}// LambertW


// Solves y = x e^x
// This is the second branch for -1/e < z < 0
double LambertWm1(double z)
{
 double w_new, w_old, ratio, e_old, tol, logz;
 double a0, a1, a2, b1, b2, b3, b4, b5;
 int n;

 /* iteration */

 tol = 1.0e-14;

 if( (z <= -exp(-1.0)) || (z > 0.0) )
  {
   fprintf(stderr, "LambertWm1 is not defined for z = %e\n", z);
   fprintf(stderr, "z needs to be between 0 and %e\n", -exp(-1.0));
   exit(0);
  }

// approximate value to initiate iteration process
// from 1003.1628v2, Darko Veberic

 a0 = -7.81417672390744;
 a1 = 253.88810188892484;
 a2 = 657.9493176902304;

 b1 = -60.43958713690808;
 b2 = 99.9856708310761;
 b3 = 682.6073999909428;
 b4 = 962.1784396969866;
 b5 = 1477.9341280760887;

 // starting point of the iteration
 w_old = a0 + a1*z + a2*z*z;
 w_old /= 1.0 + b1*z + b2*z*z + b3*z*z*z + b4*z*z*z*z + b5*z*z*z*z*z;

 w_new = 0.0;
 ratio = 1.0;
 n = 0;
 while(fabs(ratio) > tol) 
  {
   e_old = exp(w_old);
   // Halley iteration 
   ratio = w_old*e_old - z;
   ratio /= ( e_old*(w_old + 1.0) - (w_old+2.0)*(w_old*e_old-z)/(2.0*w_old + 2.0) );

   w_new = w_old - ratio;
   w_old = w_new;
   n++;
   if(n > 99) 
    {
     fprintf(stderr, "LambertWm1 is not converging after 100 iterations.\n");
     fprintf(stderr, "LambertWm1: z = %e\n", z);
     fprintf(stderr, "LambertWm1: w_old = %e\n", w_old);
     fprintf(stderr, "LambertWm1: w_new = %e\n", w_new);
     fprintf(stderr, "LambertWm1: ratio = %e\n", ratio);
     exit(0);
    }
  }

 return w_new;
}// LambertWm1


// Bins x in the interval (xi, xf)
int binning(double x, double xi, double xf, double *bins, int num_bins)
{
  double dx;
  int i;
  dx = (xf-xi)/(num_bins);
   
  // for instance if dx = 1
  // using floor means all 0.*** goes into the i = 0 bin
  // all 1.*** goes into the i = 1 bin, etc
  i = (int) floor((x-xi)/dx); 

  if( (0 <= i) && (i < num_bins) ) 
   { bins[i] += 1.0;
    return i;}
  else // outside the range 
   {
    return -1; 
   }
}// binning


// Uniform random number generator 
// Uses C's rand
double genrand_real()
{
 static int ind = 0;
 int r;
 double x;

 if(ind == 0) {srand(time(NULL));}
 ind++;
 
 r = rand();
 x = ((double) r)/((double) RAND_MAX);

 return x;
}// genrand_real