/* -*- coding: utf-8 -*-
*
* PyNX - Python tools for Nano-structures Crystallography
*   (c) 2017-present : ESRF-European Synchrotron Radiation Facility
*       authors:
*         Vincent Favre-Nicolin, favre@esrf.fr
*/

/** Project the complex calculated psi with the observed magnitude, unless the observed intensity is negative (masked).
*
* This should be called for the complete iobs array of shape: (nb_proj, nbz, ny, nx)
* Psi has a shape: (nb_proj, nbz, nb_obj, nb_probe, ny, nx)
*/
__device__ void ProjectionAmplitude(const int i, float *iobs, complexf *psi,
                                    const int nx, const int ny, const bool unmask,
                                    const bool fftshift)
{
  float obs = iobs[i];
  if(obs < 0)
  {
    if(unmask) obs = -(obs+1);
    else return;
  }

  const int nxy = nx*ny;
  // Coordinates in iobs array (centered on array)
  const int ix = i % nx;
  const int iy = (i % nxy) / nx;

  int ix1, iy1;
  if(fftshift)
  {
    // Coordinates in Psi array (centered in (0,0)). Assumes nx ny are multiple of 2
    iy1 = iy - ny/2 + ny * (iy<(ny/2));
    ix1 = ix - nx/2 + nx * (ix<(nx/2));
  }
  else
  {
    iy1 = iy;
    ix1 = ix;
  }
  // Coordinate of first mode in Psi array
  const int i0 = ix1 + iy1 * nx + (i / nxy) * nxy;

  const complexf dc = psi[i0];
  float dc2 = dot(dc,dc);

  // Normalization to observed amplitude, taking into account all modes
  dc2 = fmaxf(dc2, 1e-12f); // TODO: KLUDGE ? 1e-12f is arbitrary
  const float d = sqrtf(obs) / sqrtf(dc2);
  psi[i0] *= d;
}

__device__ void ProjectionAmplitudeIcalc(const int i, float *iobs, complexf *psi, float* icalc,
                                         const int nx, const int ny, const bool unmask,
                                         const bool fftshift)
{
  float obs = iobs[i];
  if(obs < 0)
  {
    if(unmask) obs = -(obs+1);
    else return;
  }

  const int nxy = nx*ny;
  // Coordinates in iobs array (centered on array)
  const int ix = i % nx;
  const int iz = i / nxy;
  const int iy = (i - nxy * iz) / nx;

  int ix1, iy1;
  if(fftshift)
  {
    // Coordinates in Psi/Icalc array (centered in (0,0)). Assumes nx ny are multiple of 2
    iy1 = iy - ny/2 + ny * (iy<(ny/2));
    ix1 = ix - nx/2 + nx * (ix<(nx/2));
  }
  else
  {
    iy1 = iy;
    ix1 = ix;
  }
  // Intensity from Icalc array with 2 extra columns for in-place rfft
  // Note: intensity can be <0 due to fft convolution
  float dc2 = icalc[ix1 + (nx+2) * (iy1  + ny * iz)];

  // Normalization to observed amplitude ; only accept factors between 1/10 and 10
  const float d = fmaxf(0.1f, fminf(sqrtf(obs / fmaxf(dc2, 1e-12f)), 10.0f));
  psi[ix1 + nx * (iy1  + ny * iz)] *= d;
}


/** Project the complex calculated psi with the observed *normalised* magnitude,
* unless the observed intensity is negative (masked).
*
* This should be called for a single projection iobs array of shape: (nbz, ny, nx),
* and the kernel will loop other all projections.
* Psi has a shape: (nb_proj, nbz, ny, nx)
*
* TODO: handle multiple (interpolated) iobs_empty.
*/
__device__ void ProjectionAmplitudeNormalised(const int i, float *iobs, float *iobs_empty, complexf *psi,
                                              float *iobs_empty_mean, const int nb_proj,
                                              const int nx, const int ny, const int nz)
{
  const int nxyz = nx*ny*nz;
  const float inorm = iobs_empty[i] / iobs_empty_mean[i/(nx*ny)];
  const int nxy = nx*ny;

  // 2D Coordinates in iobs array (centered on array)
  const int ix = i % nx;
  const int iy = (i % nxy) / nx;

  // 2D Coordinates in Psi array (centered in (0,0)). Assumes nx ny are multiple of 2
  const int iy1 = iy - ny/2 + ny * (iy<(ny/2));
  const int ix1 = ix - nx/2 + nx * (ix<(nx/2));

  for(int iproj=0; iproj<nb_proj; iproj++)
  {
    const float obs = iobs[i + nxyz*iproj] / inorm;
    if(obs >= 0)
    {
      // Coordinate pixel in Psi array
      const int i0 = ix1 + iy1 * nx + nxyz * iproj;

      const float dc2=dot(psi[i0], psi[i0]);

      // Normalization to observed amplitude
      psi[i0] *= sqrtf(obs) / sqrtf(fmaxf(dc2, 1e-12f));
    }
  }
}
