/* -*- 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
*/

/** Amplitude projection: apply the observed intensity to the calculated complex amplitude.
* This should be called for the first frame of a stack (i.e. cu_obs[0]) and will apply to all valid frames
*
*/
__device__ void ProjectionAmplitude(const int i, float *iobs, complexf *dcalc, float *background,
                                    const unsigned int nbmode, const unsigned int nxy,
                                    const int nxystack, const int npsi, const float scale_in,
                                    const float scale_out)
{
  const float s2 = scale_in * scale_in;
  const float sio = scale_in * scale_out;
  const float b = background[i];
  for(int j=0;j<npsi;j++)
  {
    const float obs = iobs[i + j * nxy];
    if(obs >= 0)
    {
      float dc2=0;
      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        // TODO: use local memory to avoid double-reading of dcalc !
        // Would require a __local memory array with the size=number of modes
        //dc[mode] = dcalc[i + mode*nxystack];
        //dc2 += dot(dc[mode],dc[mode]);
        const complexf dc = dcalc[i + j * nxy + mode * nxystack];
        dc2 += s2 * dot(dc,dc);
      }

      // Normalization to observed amplitude, taking into account all modes
      dc2 = fmaxf(dc2,1e-12f); // TODO: KLUDGE ? 1e-12f is arbitrary

      // Flip the sign if obs - background < 0 ?
      const float diff = obs - b;
      const float d = copysignf(sqrtf(fmaxf(diff, 0) / dc2) * sio, diff);

      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        //dcalc[i + mode*nxystack] = complexf(d*dc[mode].real() , d*dc[mode].imag());
        dcalc[i + j * nxy + mode * nxystack] *= d;
      }
    }
    else if(sio != 1.0f)
    {
      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        dcalc[i + j * nxy + mode * nxystack] *= sio;
      }
    }
  }
}


/** Same as ProjectionAmplitude, but with the calculated array fft-shifted relatively
* to the observed data. This applies only to a single frame
* This is only used for the probe in near-field ptycho
*/
__device__ void ProjectionAmplitudeShift1(const int i, float *iobs, complexf *dcalc, float *background,
                                          const unsigned int nx, const unsigned int ny, const float scale)
{
  const float b = background[i];
  const float obs = iobs[i];
  // Coordinates in the psi array (fft-shifted)
  const int ix = i % nx;
  const int iy = i / nx;

  // FFT-shifted coordinate
  const int i1 = ix - nx/2 + nx * (ix<(nx/2)) + nx * (iy - ny/2 + ny * (iy<(ny/2)));

  const complexf dc = dcalc[i1];
  const float dc2 = fmaxf(scale * dot(dc,dc), 1e-12f);

  // Flip the sign if obs - background < 0 ?
  const float diff = obs - b;
  const float d = copysignf(sqrtf(fmaxf(diff, 0) / dc2) * scale, diff);

  dcalc[i1] *= d;
}


/** Amplitude projection: apply the observed intensity to the calculated complex amplitude.
* This should be called for the first frame of a stack (i.e. cu_obs[0]) and will apply to all valid frames
*
* This version is specific to near field ptychography, and allows to use a smooth window function
* around the padding boundaries, to avoid an abrupt change which would diffract and create
* artefacts from the border.
* The smoothing is done by progressively going from using the calculated amplitude to the observed one.
*/
__device__ void ProjectionAmplitudePadWindow(const int i, float *iobs, complexf *dcalc, float *background,
                                    const unsigned int nbmode, const unsigned int nx, const unsigned int ny,
                                    const int nxystack, const int npsi, const float scale_in,
                                    const float scale_out, const int padding, const int padding_window)
{
  if(padding_window==0)
    return ProjectionAmplitude(i, iobs, dcalc, background, nbmode, nx*ny, nxystack, npsi, scale_in, scale_out);

  const unsigned int nxy = nx*ny;
  const float s2 = scale_in * scale_in;
  const float sio = scale_in * scale_out;
  const float b = background[i];

  // Coordinates in the psi array (fft-shifted)
  const int ix = i % nx;
  const int iy = i / nx;

  // Coordinates starting from the corner of the non-shifted Iobs data.
  // Assume nx ny are multiple of 2
  const int pry = iy - ny/2 + ny * (iy<(ny/2));
  const int prx = ix - nx/2 + nx * (ix<(nx/2));

  // Use a window (Tukey or erfc) for padded data in near field
  float cpad = 1.0f;

  // Tukey window
  const int p0 = padding-padding_window; // window is 50% at the boundary
  const int p1 = padding+padding_window;
  //  const int p0 = padding;  // Alternative: start windowing inside the actually observed data
  //  const int p1 = padding+2*padding_window;
  // Padding factor goes from 0 on the border to 1 at 2*padding pixels from the border
  if(prx<p0)            cpad = 0;
  else if(prx<p1)       cpad *= 0.5f * (1.0f - __cosf(   (prx-p0) * 1.57079632679f / padding_window));
  if(prx>=(nx-p0))      cpad =0;
  else if(prx>=(nx-p1)) cpad *= 0.5f * (1.0f - __cosf((nx-p0-prx) * 1.57079632679f / padding_window));
  if(pry<p0)            cpad =0;
  else if(pry<p1)       cpad *= 0.5f * (1.0f - __cosf(   (pry-p0) * 1.57079632679f / padding_window));
  if(pry>=(ny-p0))      cpad = 0;
  else if(pry>=(ny-p1)) cpad *= 0.5f * (1.0f - __cosf((ny-p0-pry) * 1.57079632679f / padding_window));

  for(int j=0;j<npsi;j++)
  {
    if(cpad==0)
    {
      // In the padding area, outside the smoothed part - keep the calculated values, scaled
      if(sio != 1.0f)
      {
        for(unsigned int mode=0 ; mode<nbmode ; mode++)
          dcalc[i + j * nxy + mode * nxystack] *= sio;
      }
    }
    else
    {
      float obs = iobs[i + j * nxy];

      if((obs<0) && (prx>=padding) && (prx<(nx-padding)) && (pry>=padding) && (pry<(ny-padding)))
      { // Masked pixel, inside the observed data array => keep calculated value
        if(sio != 1.0f)
        {
          for(unsigned int mode=0 ; mode<nbmode ; mode++)
            dcalc[i + j * nxy + mode * nxystack] *= sio;
        }
      }
      else
      {
        // Un-mask extrapolated data
        if(obs<0) obs = -(obs+1);

        // Calculated intensity from all modes
        float dc2=0;
        for(unsigned int mode=0 ; mode<nbmode ; mode++)
        {
          const complexf dc = dcalc[i + j * nxy + mode * nxystack];
          dc2 += s2 * dot(dc,dc);
        }
        dc2 = fmaxf(dc2,1e-12f); // TODO: KLUDGE ? 1e-12f is arbitrary

        // Smoothing part: switch from calculated to observed depending on cpad
        const float icalc = cpad * (obs -b) + (1-cpad) * dc2;

        // Flip the sign if obs - background < 0 ?
        const float d = copysignf(sqrtf(fmaxf(icalc, 0) / dc2) * sio, icalc);

        for(unsigned int mode=0 ; mode<nbmode ; mode++)
        {
          //dcalc[i + mode*nxystack] = complexf(d*dc[mode].real() , d*dc[mode].imag());
          dcalc[i + j * nxy + mode * nxystack] *= d;
        }
      }
    }
  }
}


__device__ void ProjectionAmplitudeBackground(const int i, float *iobs, complexf *dcalc, float *background,
                                              float *vd, float *vd2, float *vz2, float *vdz2,
                                              const unsigned int nbmode, const unsigned int nxy,
                                              const int nxystack, const int npsi, const char first_pass,
                                              const float scale_in, const float scale_out)
{
  const float s2 = scale_in * scale_in;
  const float sio = scale_in * scale_out;
  const float b = background[i];
  // For the background update
  float psi2 = 0;
  float dz2 = 0;
  float d2 = 0;
  float d = 0;

  for(int j=0;j<npsi;j++)
  {
    const float obs = iobs[i + j * nxy];
    if(obs >= 0)
    {
      float dc2=0;
      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        // TODO: use local memory to avoid double-reading of dcalc !
        // Would require a __local memory array with the size=number of modes
        //dc[mode] = dcalc[i + mode*nxystack];
        //dc2 += dot(dc[mode],dc[mode]);
        const complexf dc = dcalc[i + j * nxy + mode * nxystack];
        dc2 += s2 * dot(dc,dc);
      }

      // Normalization to observed amplitude, taking into account all modes
      dc2 = fmaxf(dc2,1e-12f); // TODO: KLUDGE ? 1e-12f is arbitrary

      // Background update
      const float dd = obs - b;
      psi2 += dc2;
      dz2 += dd * dc2;
      d2 += dd * dd;
      d += dd;

      const float d = sqrtf(fmaxf(obs-b, 0) / dc2) * sio;

      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        //dcalc[i + mode*nxystack] = complexf(d*dc[mode].real() , d*dc[mode].imag());
        dcalc[i + j * nxy + mode * nxystack] *= d;
      }
    }
    else if(sio != 1.0f)
    {
      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        dcalc[i + j * nxy + mode * nxystack] *= sio;
      }
    }
  }
  if(first_pass)
  {
    vd  [i] = d;
    vd2 [i] = d2 ;
    vz2 [i] = psi2;
    vdz2[i] = dz2;
  }
  else
  {
    vd  [i] += d;
    vd2 [i] += d2 ;
    vz2 [i] += psi2;
    vdz2[i] += dz2
    ;
  }
}

/** Amplitude projection with background update using a mode approach.
*
*/
__device__ void ProjectionAmplitudeBackgroundMode(const int i, float *iobs, complexf *dcalc, float *background,
                                              float *background_new,
                                              const unsigned int nbmode, const unsigned int nxy,
                                              const int nxystack, const int npsi, const char first_pass,
                                              const float scale_in, const float scale_out)
{
  const float s2 = scale_in * scale_in;
  const float sio = scale_in * scale_out;
  const float b = background[i];
  float db=0;
  for(int j=0;j<npsi;j++)
  {
    const float obs = iobs[i + j * nxy];
    if(obs >= 0)
    {
      float dc2=b;
      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        // TODO: use local memory to avoid double-reading of dcalc !
        // Would require a __local memory array with the size=number of modes
        //dc[mode] = dcalc[i + mode*nxystack];
        //dc2 += dot(dc[mode],dc[mode]);
        const complexf dc = dcalc[i + j * nxy + mode * nxystack];
        dc2 += s2 * 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 / dc2) * sio;

      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        //dcalc[i + mode*nxystack] = complexf(d*dc[mode].real() , d*dc[mode].imag());
        dcalc[i + j * nxy + mode * nxystack] *= d;
      }
      // Background update as an incoherent mode
      db += b * obs / dc2;
    }
    else if(scale_in * scale_out != 1.0f)
    {
      for(unsigned int mode=0 ; mode<nbmode ; mode++)
      {
        dcalc[i + j * nxy + mode * nxystack] *= sio;
      }
    }
  }
  if(first_pass) background_new[i] = db;
  else background_new[i] += db;
}
