#include "iteration_functions.h"

void compute_stokes(double ene, double* k_lab, double* polvect_cart, stokes_parameters* ptr_stokes, double* array_ubounds, int flag)

{
  int FlagFound;
  int ii, jj;
  double Qs;
  double Us;
  double cos_theta;

  static double Qp_cart[3];
  static double Qm_cart[3];
  static double Up_cart[3];
  static double Um_cart[3];

  cos_theta = k_lab[2];

  /* =====================================================================# */
  /* Loop on energies */
  /* =====================================================================# */

  FlagFound = FALSE;

  if (flag == 0)
  {
    for (ii = 0; ii < POLAR_NBOUNDS - 1; ii++)
    {
      if ((ene >= ene_bound[ii]) && (ene <= ene_bound[ii + 1]))
      {
        /* =====================================================================# */
        /* Loop on u array, the cosine of the polar angle. Indeed, k[2]=u */
        /* =====================================================================# */

        for (jj = 0; jj < nstepangles; jj++)
        {
          if ((cos_theta <= array_ubounds[jj]) && (cos_theta >= array_ubounds[jj + 1]))
          {

            detector_axis(k_lab, Qp_cart, Qm_cart, Up_cart, Um_cart);

            Qs = pow(dot_prod(Qp_cart, polvect_cart), 2) - pow(dot_prod(Qm_cart, polvect_cart), 2);
            Us = pow(dot_prod(Up_cart, polvect_cart), 2) - pow(dot_prod(Um_cart, polvect_cart), 2);


            if (isnan(Qs))
            {
              printf("Error: Qs is Nan\n");
              exit(1);
            }

            if (isnan(Us))
            {
              printf("Error: Us is Nan\n");
              exit(1);
            }

            ptr_stokes[ii].array_I[jj] = ptr_stokes[ii].array_I[jj] + 1;
            ptr_stokes[ii].array_Q[jj] = ptr_stokes[ii].array_Q[jj] + Qs;
            ptr_stokes[ii].array_U[jj] = ptr_stokes[ii].array_U[jj] + Us;
            ptr_stokes[ii].counter[jj]++;

            FlagFound = TRUE;

            break;
          }
        }
      }
    }

    if (FlagFound == FALSE)
    {
      double zio = 1;
      printf("Warning: no correspondence for Ene=%5.4f k_z=%5.4f\n", ene, k_lab[2]);
      // exit(1);
    }
  }
  else
  {
    double u_center;
    double Qtot;
    double Utot;
    double Itot_u1;
    double Itot;
    int Itot_allbands = 0;

    double* Pdeg = (double*)malloc(nstepangles * sizeof(double));
    dat_integral_polar = fopen(integral, "w");

    fprintf(dat_integral_polar, "!Col1: cos(theta)\n");
    fprintf(dat_integral_polar, "!Col2: Pdeg\n");
    fprintf(dat_integral_polar, "!Col3: Q\n");
    fprintf(dat_integral_polar, "!Col4: U\n");
    fprintf(dat_integral_polar, "!Col5: Itot/I(u=1)*1/u\n\n");

    for (jj = 0; jj < nstepangles; jj++)
    {
      u_center = (array_ubounds[jj] + array_ubounds[jj + 1]) / 2;

      Qtot = 0;
      Utot = 0;
      Itot = 0;

      for (ii = 0; ii < POLAR_NBOUNDS - 1; ii++)
      {
        Itot_allbands = Itot_allbands + ptr_stokes[ii].array_I[jj];

        Qtot = Qtot + ptr_stokes[ii].array_Q[jj];
        Utot = Utot + ptr_stokes[ii].array_U[jj];
        Itot = Itot + ptr_stokes[ii].array_I[jj];

        if (isnan(Qtot))
        {
          printf("Qtot is Nan at jj=%d ii=%d\n", jj, ii);
          exit(1);
        }

        if (isnan(Utot))
        {
          printf("Utot is Nan at jj=%d ii=%d\n", jj, ii);
          exit(1);
        }

        if (isnan((double)Itot))
        {
          printf("Itot is Nan at jj=%d ii=%d\n", jj, ii);
          exit(1);
        }
      }

      Pdeg[jj] = sqrt(Qtot * Qtot + Utot * Utot) / Itot;

      if (jj == 0)
      {
        Itot_u1 = Itot;
      }

      fprintf(dat_integral_polar, "%lf %lf %lf %lf %5.4f\n", u_center, Pdeg[jj], Qtot / Itot,
              Utot / Itot, 1/u_center*(Itot / Itot_u1));
    }

    fclose(dat_integral_polar);

    printf("Total counts all bands %d\n", Itot_allbands);
  }
}