#include "iteration_functions.h"

#define NSC_MAX 1000

extern double albedobase;

/*==================================================================*/
/*Function prototypes*/
/*==================================================================*/

void comptonization(double E_lab,
                    double* k_lab,
                    double* Efield_lab,
                    double* Magfield_lab,
                    double* E_new_lab,
                    double* k_new_lab,
                    double* Efield_lab_new,
                    double* Magfield_lab_new);


double distance_bottom(double z0, double kz);
double distance_top(double z0, double kz, double H);

/*==================================================================*/
/*Datatypes defined in the header globals.h*/
/*==================================================================*/

//double obsmindeg;
//double obsmaxdeg;


//double* array_angles;
//double* array_ubounds;

//double** array_gb;
//double** ptr_2colfile;


const gsl_root_fsolver_type* T_phi;
gsl_root_fsolver* s_phi;

//const gsl_root_fsolver_type* T_disk;
//gsl_root_fsolver* s_disk;

const gsl_root_fsolver_type* T_maxw;
gsl_root_fsolver* s_maxw;

/*==================================================================*/

int slab_mc(int nphot, int seed)
{
  int ii;
  int jj;
  int n_sc;
  int n_zero_sc;
  int nph_esc;
  int Fval_max;
  int FlagLocation;
  int photon_diffusion[NSC_MAX];
  int status;


  double E_lab;
  double E_new_lab;
  double k_lab[3];
  double k_new_lab[3];
  double Efield_lab[3];
  double Efield_lab_new[3];
  double Magfield_lab[3];
  double Magfield_lab_new[3];
  double pos[3];
  double cos_theta;
  double sin_theta;
  double phi;
  double NeDisk;
  double r_units;
  double Hphot;
  double L_top;
  double tau;
  //double eps;
  double l_int;
  double csi;
  double Pdeg;

  T_disk = gsl_root_fsolver_bisection;
  s_disk = gsl_root_fsolver_alloc(T_disk);

  T_phi = gsl_root_fsolver_bisection;
  s_phi = gsl_root_fsolver_alloc(T_phi);

  /*This is for sampling beta from subrelativistics maxwellian*/

  T_maxw = gsl_root_fsolver_bisection;
  s_maxw = gsl_root_fsolver_alloc(T_maxw);

  nstepangles = 40;
  obsmindeg = 0.;
  obsmaxdeg = 90.;

  ptr_gb = read_array(
      stringconcat(getenv("MC_INSTDIR"), "/electron_distributions/maxwellian/temp_random.txt"));

  matrix_gb =
      stringconcat(getenv("MC_INSTDIR"), "/electron_distributions/maxwellian/gammabeta_lab.fits");
  array_gb = read_matrix(matrix_gb);

  double array_ene[NSTEP_ENE];
  make_array_energy(array_ene, 0.01, 20, NSTEP_ENE);

  double** array_counts;
  array_counts = dmatrix(0, NSTEP_ENE - 1, 0, nstepangles - 1);

  double** reduce_array_counts;
  reduce_array_counts = dmatrix(0, NSTEP_ENE - 1, 0, nstepangles - 1);

  for (ii = 0; ii < NSTEP_ENE - 1; ii++)
  {
    for (jj = 0; jj < nstepangles; jj++)
    {
      array_counts[ii][jj] = 0;
      reduce_array_counts[ii][jj] = 0;
    }
  }

  /*=======================================================================*/

  Hphot = 0.5;
  r_units = NS_MASS * R_SCW_SUN;
  NeDisk = disktau / (SIGMA_T * r_units * Hphot);

  /*=======================================================================*/

  if (DiskPolarizationFlag == FROMFILE)
  {
    radiation_field_t* radiation_field = (radiation_field_t*)malloc(sizeof(radiation_field_t));

    read_resultsfile(polarinput, radiation_field, MC_SLAB);

    setup_gsl_objects(radiation_field, MC_SLAB);
  }


  Pdeg=0;
  nph_esc = 0;
  cos_theta = 0;

  for (ii = 0; ii < NSC_MAX; ii++)
  {
    photon_diffusion[ii] = 0;
  }
  /*===================================================================*/

  array_angles = (double*)malloc((nstepangles) * sizeof(double));
  array_ubounds = (double*)malloc((nstepangles + 1) * sizeof(double));
  make_array_angles(array_angles, array_ubounds);

  stokes_parameters struct_stokes[POLAR_NBOUNDS - 1];
  allocate_stokes_pointer(struct_stokes);

  stokes_parameters struct_stokes_average[POLAR_NBOUNDS - 1];
  allocate_stokes_pointer(struct_stokes_average);

  n_zero_sc=0;
  /*=============================================================*/
  /*Here starts loop over photon */
  /*=============================================================*/

  for (ii = 0; ii < nphot; ii++)
  {
    //======================================================================== */
    // New unpolarized photon */
    // ===================================================================== */

    if (rank == 0 && ii % 5000 == 0)
    {
      printf("======================\n");
      printf("Photon %d\n", ii);
    }

    n_sc = 0;
    E_lab = bb_spec(ktbb);
    E_new_lab = 0;

    pos[0] = 0;
    pos[1] = 0;

    if (seed == 1)
    {
      pos[2] = Hphot / 2;
      cos_theta = 2 * clock_random() - 1;
    }
    else if (seed == 2)
    {
      pos[2] = clock_random() * Hphot;
      cos_theta = 2 * clock_random() - 1;
    }
    else
    {
      pos[2] = 0;

      if (DiskPolarizationFlag == FROMFILE)
      {
        csi = clock_random();
        status = gsl_spline_eval_e(Spline_sample_Iu, csi, xacc_1d, &cos_theta);

        if (status)
        {
          printf("Error\n");
          printf("File %s\n", __FILE__);
        }

        printf("sampling %lf\n", cos_theta);
      }
      else
      {
        cos_theta = clock_random();
      }
    }

    sin_theta = sqrt(1 - cos_theta * cos_theta);
    phi = 2 * PI * clock_random();

    k_lab[0] = sin_theta * cos(phi);
    k_lab[1] = sin_theta * sin(phi);
    k_lab[2] = cos_theta;


    if (DiskPolarizationFlag == FROMFILE)
    {
      electricfield_fromfile(k_lab, Efield_lab, Magfield_lab, &Pdeg);
    }
    else
    {
      unpolarized_seedphot(k_lab, Efield_lab, Magfield_lab);
    }

    /*====================================================================*/

    FlagLocation = 1;

    while (FlagLocation)
    {
      if (k_lab[2] >= 0)
      {
        L_top = distance_top(pos[2], k_lab[2], Hphot);
      }
      else
      {
        L_top = distance_bottom(pos[2], k_lab[2]);
      }

      /*====================================================================*/
      /*Compute the optical depth from top or base of the atmosphere*/
      /*=====================================================================*/

      tau = NeDisk * SIGMA_T * r_units * L_top;
      eps = -log(clock_random());

      if (tau > eps)
      {
        l_int = eps / (NeDisk * SIGMA_T * r_units);

        pos[0] = pos[0] + k_lab[0] * l_int;
        pos[1] = pos[1] + k_lab[1] * l_int;
        pos[2] = pos[2] + k_lab[2] * l_int;

        //printf("Photon %d before scattering %d  Pdeg = %5.4f\n", ii, n_sc, Pdeg);

        comptonization(E_lab, k_lab, Efield_lab, Magfield_lab,  &E_new_lab, k_new_lab,
                       Efield_lab_new, Magfield_lab_new);

        //printf("After scattering Pdeg = %5.4f\n\n", Pdeg);

        n_sc++;
        E_lab = E_new_lab;

        for (jj=0; jj<3; jj++)
        {
        	k_lab[jj] = k_new_lab[jj];
        	Efield_lab[jj] = Efield_lab_new[jj];
        	Magfield_lab[jj] = Magfield_lab_new[jj];
        }

      }
      else

      /*Photon potentially escaping from the slab	  */
      {
        if (k_lab[2] > 0)
        {
          // printf("Photon %d escaping after %d scattering\n", ii, n_sc);

          if (n_sc==0) n_zero_sc++;

          compute_stokes(E_lab, k_lab, Efield_lab, struct_stokes, array_ubounds, 0);
          photon_diffusion[n_sc]++;
          collect_photons(E_new_lab, k_lab, array_ene, array_counts, 0);

          nph_esc++;
          FlagLocation = 0;
        }
        else
        {
          csi = clock_random();

          /*Photon absorbed and newly re-emitted with probability A=albedobase*/
          if (csi <= albedobase)
          {
            // printf("Photon %d reflected\n", ii);

            cos_theta = clock_random();
            sin_theta = sqrt(1 - cos_theta * cos_theta);
            phi = 2 * PI * clock_random();

            k_lab[0] = sin_theta * cos(phi);
            k_lab[1] = sin_theta * sin(phi);
            k_lab[2] = cos_theta;

            pos[0] = 0;
            pos[1] = 0;
            pos[2] = 0;

            unpolarized_seedphot(k_lab, Efield_lab, Magfield_lab);
            E_lab = bb_spec(ktbb);
          }

          /*Photon absorbed*/
          else
          {
            FlagLocation = 0;
            break;
          }
        } // Close the else with condition k_lab[2] < 0

      } // End of confition  if (tau > eps)... else...

    } // end of while (FlagLocation)
  }

  /*==================================================================*/

  for (ii = 0; ii < POLAR_NBOUNDS - 1; ii++)
  {
    MPI_Reduce(struct_stokes[ii].array_Is, struct_stokes_average[ii].array_Is, nstepangles,
               MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);

    MPI_Reduce(struct_stokes[ii].array_Qs, struct_stokes_average[ii].array_Qs, nstepangles,
               MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);

    MPI_Reduce(struct_stokes[ii].array_Us, struct_stokes_average[ii].array_Us, nstepangles,
               MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);

    MPI_Reduce(struct_stokes[ii].counter, struct_stokes_average[ii].counter, nstepangles, MPI_INT,
               MPI_SUM, 0, MPI_COMM_WORLD);
  }

  for (ii = 0; ii < NSTEP_ENE; ii++)
  {
    for (jj = 0; jj < nstepangles; jj++)
    {
      MPI_Reduce(&array_counts[ii][jj], &reduce_array_counts[ii][jj], 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    }
  }

  // Get the number of processes

   MPI_Comm_size(MPI_COMM_WORLD, &nproc);


  /*==================================================================*/

  if (rank == 0)
  {
    compute_stokes(1, k_lab, Efield_lab, struct_stokes_average, array_ubounds, 1);

    printf("Percentage of escaping photons from the top layer %4.3f\n", 1. * nph_esc / nphot);
    printf("Percentage of escaping non scattered photons from the top layer %4.3f\n", (double) n_zero_sc/ nphot);
    /*============================================================*/

    dat_diffusion = fopen(diffusion, "w");
    fprintf(dat_diffusion, "!Percentage of escaping photons %4.3f\n", 1. * nph_esc / nphot);

    Fval_max=-10;

    for (ii = 0; ii < 300; ii++)
    {
      if (photon_diffusion[ii] > Fval_max)
      {
        Fval_max = photon_diffusion[ii];
      }
    }

    for (ii = 0; ii < 300; ii++)
    {
      fprintf(dat_diffusion, " %d  %3.2f\n", ii, (double)photon_diffusion[ii]/Fval_max);
    }

    fclose(dat_diffusion);
  }

  collect_photons(0, k_lab, array_ene, reduce_array_counts, 1);

  dat_energy_polar = fopen(polarfile, "w");

  write_polardata(struct_stokes_average, ene_bound, dat_energy_polar);

  fclose(dat_energy_polar);

  gsl_root_fsolver_free(s_phi);
  gsl_root_fsolver_free(s_disk);

  return 0;

} // End of function


void comptonization(double E_lab,
                    double* k_lab,
                    double* Efield_lab,
                    double* Magfield_lab,
                    double* E_new_lab,
                    double* k_new_lab,
                    double* Efield_lab_new,
                    double* Magfield_lab_new)

{

  int jj;
  int status;
  int PolarFlag;
  double Pdeg;
  double phik;
  double E_erf;
  double E_new_erf;
  double csi;
  double costheta_erf;
  double sintheta_erf;
  double beta_el;

  double k_erf[3];
  double k_new_erf[3];
  double el_vel_vect[3];
  double ElectField_new_erf[3];
  double ElectField_old_erf[3];
  double MagField_Erf[3];


  PolarFlag=0;
  Pdeg=0;

  //================================================================== */
  // Find the K' components in the ERF */
  // If also polarization is computed, transform also the electric */
  // field vector to ERF */
  //================================================================== */

  beta_el = sample_electron(E_lab, k_lab, kte, ACCRETION_DISK, el_vel_vect);

  versor_boost(k_lab, el_vel_vect, k_erf, 1);

  elmag_lorentztransform(Efield_lab, Magfield_lab, el_vel_vect, ElectField_old_erf, MagField_Erf, 1);

  E_erf = energy_lorentz_boost(E_lab, el_vel_vect, k_lab, 1);


  csi = clock_random();

  costheta_erf = (-1 + pow(-2 + 4 * csi + sqrt(5 - 16 * csi + 16 * pow(csi, 2)), 2 / 3.)) /
                 pow(-2 + 4 * csi + sqrt(5 - 16 * csi + 16 * pow(csi, 2)), 1 / 3.);

  sintheta_erf = sqrt(1 - costheta_erf * costheta_erf);

  // ==================================================================
  // Find the new energy photon in the erf after scattering
  //==================================================================

  E_new_erf = new_energy_erf(E_erf, costheta_erf);


  phik = sample_azim_angle(E_erf, E_new_erf, costheta_erf, 1);

  //================================================================== */
  /*Find the components of the photon versor after scattering*/
  /*Use Equation(2) from Peplow, 1999*/
  /*Nuclear Science and Engineering, 131, 132*/
  //================================================================== */

  for (jj = 0; jj < 3; jj++)
  {
    k_new_erf[jj] = k_erf[jj] * costheta_erf + ElectField_old_erf[jj] * cos(phik) * sintheta_erf +
                    MagField_Erf[jj] * sin(phik) * sintheta_erf;
  }

  //================================================================== */
  // Check the probability for depolarization*/
  // and compute the new electric field vector
  //================================================================== */

  PolarFlag=1;
  status = depolarization_probability(E_erf, E_new_erf, costheta_erf, phik, &Pdeg, &PolarFlag);
  

  if (status == 1)
  {
    exit(1);
  }

  electfield_as_new(k_new_erf, ElectField_old_erf, ElectField_new_erf, PolarFlag);

  vect_prod(k_new_erf, ElectField_new_erf, MagField_Erf);

  // =========================================================================== */
  // Back convert from (K1)' in ERF to lab*/
  // =========================================================================== */

  versor_boost(k_new_erf, el_vel_vect, k_new_lab, 2);

  // ================================================================== */
  // Back-transform electric field versor to the lab frame */
  // ================================================================== */

  elmag_lorentztransform(ElectField_new_erf, MagField_Erf, el_vel_vect, Efield_lab_new, Magfield_lab_new, 2);

  // ================================================================== */
  // Back convert electron energy from ERF to lab */
  // ================================================================== */

  *E_new_lab = energy_lorentz_boost(E_new_erf, el_vel_vect, k_new_erf, 2);

  if (isnan(*E_new_lab) || *E_new_lab < 0)
  {
    printf("Error in photon energy in lab frame after scattering: value %5.4f\n", *E_new_lab);
    exit(1);
  }
}

/*======================================================================*/

double distance_top(double z_pos, double kz, double H)
{
  double t;

  t = (H - z_pos) / kz;

  return t;
}

/*======================================================================*/

double distance_bottom(double z_pos, double kz)
{
  double t;

  t = -z_pos / kz;

  return t;
}