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


/** Average the probe over multiple distances.
* This should be called for the first probe and first mode
*/
__device__ void AveProbeZ(const int i, complexf *probe, complexf*pr0, float *f,
                          const int nx, const int ny, const int nz, const int nprobe)
{
  const int nxy = nx*ny;
  // Coordinates in iobs array (centered on array)
  const int imode = i / nxy;  // If multiple probe modes
  const float ix = i % nx;
  const float iy = (i - imode * nxy) / nx;

  // We interpolate differently for each iz
  for(int iz0=0 ;iz0<nz ;iz0++)
  {
    for(int imode=0; imode<nprobe; imode++)
    {
      complexf pr = pr0[i + nxy * (imode + nprobe * iz0)];
      for(int iz=0 ;iz<nz ;iz++)
      {
        if(iz!=iz0)
        {
          const float m = f[iz] / f[iz0]; // Magnification factor between probes
          // Note that we assume nx, ny are multiples of 2
          // TODO: check actual center used when zooming
          const float ixzoom = (ix - nx/2) * m + nx/2;
          const float iyzoom = (iy - ny/2) * m + ny/2;
          // Interpolation is warped across boundaries
          pr += bilinear(pr0, ixzoom, iyzoom, imode + nprobe * iz, nx, ny, true, false);
        }
      }
      probe[i + nxy * (imode + nprobe * iz0)] = pr / float(nz);
    }
  }
}
