/** Multiply object and probe to get a Psi to be propagated.
* Version for DM, Compute (2-alpha) * P * O - (1-alpha) * Psi
* alpha is the coefficient allowing a fraction of AP with DM.
*
* Object has 3 dimensions: projections, y, x.
* Probe has 4 dimensions: z (propagation distances), modes, y, x.
* The 5-dimensional Psi stack is calculated, with dimensions:
*    projections, propagation distances, probe modes, y, x
*
* This kernel function should be called for a 2D array, but will be applied to all z, projections and modes.
*
* This version assumes the probe is different for each z (propagation distances)
*/
__device__ void ObjectProbe2PsiDM1(const int i, complexf* obj, complexf* probe, complexf* psi,
                      float *dx, float *dy,
                      const int nb_proj, const int nb_z, const int nb_probe,
                      const int nx, const int ny, const float alpha)
{
  // Coordinates in object array (origin at array center)
  const int ix = i % nx;
  const int iy = i / nx;

  // Coordinates in Psi array (origin at (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));
  const int ipsi  = ix1 + iy1 * nx ;

  for(int iprobe=0; iprobe<nb_probe; iprobe++)
  {
    for(int iproj=0;iproj<nb_proj;iproj++)
    {
      complexf o;
      for(int iz=0;iz<nb_z;iz++)
      {
        // We have to read the object for each z due to shifts
        o = bilinear(obj, ix+dx[iz + iproj * nb_z], iy+dy[iz + iproj * nb_z], iproj, nx, ny, false, false);

        const complexf p = probe[ix + nx * (iy + ny * (iprobe + nb_probe * iz))];
        const complexf ps = complexf(o.real()*p.real() - o.imag()*p.imag() , o.real()*p.imag() + o.imag()*p.real());

        const int i1 = ipsi + nx * ny * (iprobe + nb_probe * (iz + nb_z * iproj));
        psi[i1] = (2.0f-alpha) * ps - (1.0f-alpha) * psi[i1];
      }
    }
  }
}

/** 2nd stage Psi update for DM, computes Psi(n-1) - P*O + Psi,
* where Psi is the result of magnitude projection in detector space.
*
* Object has 3 dimensions: projections, y, x.
* Probe has 4 dimensions: z (propagation distances), modes, y, x.
* The 5-dimensional Psi stack is calculated, with dimensions:
*    projections, propagation distances, probe modes, y, x
*
* This kernel function should be called for a 2D array, but will be applied to all z, projections and modes.
*
* This version assumes the probe is different for each z (propagation distances)
*/
__device__ void ObjectProbe2PsiDM2(const int i, complexf* obj, complexf* probe, complexf* psi, complexf* psi_copy,
                      float *dx, float *dy,
                      const int nb_proj, const int nb_z, const int nb_probe,
                      const int nx, const int ny, const float alpha)
{
  // Coordinates in object array (origin at array center)
  const int ix = i % nx;
  const int iy = i / nx;

  // Coordinates in Psi array (origin at (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));
  const int ipsi  = ix1 + iy1 * nx ;

  for(int iprobe=0; iprobe<nb_probe; iprobe++)
  {
    for(int iproj=0;iproj<nb_proj;iproj++)
    {
      for(int iz=0;iz<nb_z;iz++)
      {
        const complexf o = bilinear(obj, ix+dx[iz + iproj * nb_z],
                                    iy+dy[iz + iproj * nb_z], iproj, nx, ny, false, false);
        const complexf p = probe[ix + nx * (iy + ny * (iprobe + nb_probe * iz))];
        const complexf ps = complexf(o.real()*p.real() - o.imag()*p.imag() , o.real()*p.imag() + o.imag()*p.real());

        const int i1 = ipsi + nx * ny * (iprobe + nb_probe * (iz + nb_z * iproj));
        // DM mixed with AP: Psi(n+1)= (1-alpha)*(Psi - P_s Psi) + P_F[(2-alpha)P_S Psi - (1-alpha) * Psi]
        psi[i1] += (1.0f-alpha) * (ps - psi_copy[i1]) ;
      }
    }
  }
}
