// #include "cuda_fp16.h"
#define twopi 6.2831853071795862f

/** Put observed intensities in Psi array, and normalise with empty_beam image
*
* This should be called for the complete iobs array of shape: (nb_proj, nbz, ny, nx)
* Psi has a shape: (nb_proj, nb_obj=1, nb_probe=1, ny, nx)
* The iobs_empty amplitude data is stored in the real part of the probe.
*/
__device__ void Iobs2Psi(const int i, float *iobs, complexf* probe, complexf *psi,
                         float* background, const int nx, const int ny)
{
  // coordinates in probe array (not fft-shifted like psi and iobs)
  const int prx = i % nx;
  const int pry = (i % (nx * ny)) / nx;
  const int iy = pry - ny/2 + ny * (pry<(ny/2));
  const int ix = prx - nx/2 + nx * (prx<(nx/2));
  const int ipr  = ix + iy * nx ;

  float obs = iobs[i];
  if(obs<0) obs = -(obs+1);  // padded / masked intensities stored as -1-I_interp

  float b = background[prx + nx * pry];
  if(b<0) b = -(b+1);  // same for the background (dark)

  psi[i] = complexf((obs-b) / fmaxf(probe[ipr].real() * probe[ipr].real() -b, 1e-2f), 0.0f);
}

// Paganin filter in Fourier space.
// This should be called for the whole psi array
__device__ void paganin_fourier(const int i, complexf *psi, float alpha, const float px,
                                const int nx, const int ny)
{
  // Coordinates in psi (no modes so same size)
  const int nxy = nx * ny;
  const int ix = i % nx;
  const int iy = (i % nxy) / nx;

  // Assumes ny, nx are multiples of 2
  const float ky = (iy - (int)ny *(int)(iy >= ((int)ny / 2))) * twopi / (px * (float)ny) ;
  const float kx = (ix - (int)nx *(int)(ix >= ((int)nx / 2))) * twopi / (px * (float)nx) ;

  // Paganin original method
  const complexf ps = psi[i];
  const float a = 1.0f + 0.5f * alpha * (kx*kx + ky*ky);
  //psi[ipsi] = psi[ipsi] / float(1.0f + 0.5f * alpha[iz] * (kx*kx + ky*ky));
  psi[i] = complexf(ps.real()/a, ps.imag()/a) ;
}

/** This function should be called for the first psi frame,
* and will compute the new object
* There are no object or probe modes, so:
*   Object has 2 dimensions: ny, nx.
*   Psi stack has dimensions: npsi, ny, nx
*
*
*/
__device__ void paganin2obj(const int i, complexf *psi, float *obj_mut, float* obj_norm,
                            float* cx,  float* cy, const int nx, const int ny,
                            const int nxo, const int nyo, const int npsi, const int padding,
                            const int padding_window)
{
  // Coordinates from the corner of the array (probe)
  const int prx = i % nx;
  const int pry = i / nx;
  const int interp=1; // no interpolation for paganin

  // Coordinates, fft-shifted in psi. Assumes nx ny are multiple of 2
  const int iy1 = pry - ny/2 + ny * (pry<(ny/2));
  const int ix1 = prx - nx/2 + nx * (prx<(nx/2));

  // Use a Tukey window for padded data in near field
  float cpad = 1.0;
  if(padding>0)
  { // Tukey dampening
    const int p0 = padding-padding_window;
    const int p1 = padding+padding_window;
    if(prx<p0)            cpad = 0;
    else if(prx<p1)       cpad *= 0.5 * (1 - __cosf(   (prx-p0) * 1.57079632679f / padding_window));
    if(prx>=(nx-p0))      cpad =0;
    else if(prx>=(nx-p1)) cpad *= 0.5 * (1 - __cosf((nx-p0-prx) * 1.57079632679f / padding_window));
    if(pry<p0)            cpad =0;
    else if(pry<p1)       cpad *= 0.5 * (1 - __cosf(   (pry-p0) * 1.57079632679f / padding_window));
    if(pry>=(ny-p0))      cpad = 0;
    else if(pry>=(ny-p1)) cpad *= 0.5 * (1 - __cosf((ny-p0-pry) * 1.57079632679f / padding_window));
  }

  for(int j=0;j<npsi;j++)
  {
    if(cx[j]<1e8)
    {
      complexf o=0;
      const float mut = -logf(abs(psi[ix1 + nx * (iy1 + nx * j ) ]));

      // Distribute the computed mu*t on the 4 corners of the interpolated object
      bilinear_atomic_add_f(obj_mut, mut * cpad, cx[j] + prx, cy[j] + pry, 0, nxo, nyo, interp);
      bilinear_atomic_add_f(obj_norm,      cpad, cx[j] + prx, cy[j] + pry, 0, nxo, nyo, interp);
    }
  }
}

// Compute average mu*t after Paganin, and save complex object
__device__ void PaganinObjNorm(const int i, float* obj_norm, float* obj_mut, complexf *obj,
                               const float delta_beta)
{
  const float mut = obj_mut[i] / fmaxf(obj_norm[i], 1e-6f);
   // Use approximations if absorption or phase shift is small
  float a = 0.5 * mut;
  if(fabs(a) < 1e-4)
      a = 1 - a ;
  else
    a = expf(-a);

  const float ph = -0.5f * mut * delta_beta;
  if(fabs(ph)<1e-4)
    obj[i] = complexf(a * (1-ph*ph), a * ph);
  else
    obj[i] = complexf(a * cos(ph), a * sin(ph));
}
