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

/** Reverse the order of an array along multiple dimensions.
* This can be used to perform a circular shift using 2 successive
* calls.
*
*  This should be called as an elementwise operation with the full array
* as first argument
*
* To perform a circular shift of an array with shifts (dx, dy, dz), two calls
* are necessary:
*  reverse_array(a, nx, ny, nz, nx, ny, nz);
*  reverse_array(a, nx, ny, nz, dx, dy, dz);
*
* \param nx, ny, nz: shape of the array a (nx being the fast axis)
* \param dx, dy, dz: boundaries for the inner blocks to be reversed independently.
*  the array is reversed by blocks, e.g. for a 2D array, a flip is done within
* four blocks:
*    0<=ix<dx, 0<=iy<dy
*    0<=ix<dx, dy<=iy<ny
*    dx<=ix<nx, 0<=iy<dy
*    dx<=ix<nx,  dy<=iy<ny
*
* TODO: for each block of total size N, only N/2 do something. Can this be improved ?
*  In principle it should be possible to call for only half the size of the array
*  (precisely N//2), and then split the N//2 among the different half-blocks...
*  moderately tedious (and obfuscating), not worth it if operation is bandwidth-limited (likely).
*
* Equivalent numpy code in 2D:
import numpy as np
ny, nx = 5, 5
dy, dx = 2, 1
v = np.char.array([v for v in "ABCDEFGHIJKLMNOPQRSTUVWXY"]).reshape((ny,nx))
print(v)
# Flip the whole array
v = v[::-1, ::-1]
# Flip the 4 blocks
v[:dy, :dx] = v[:dy, :dx][::-1, ::-1]
v[dy:, :dx] = v[dy:, :dx][::-1, ::-1]
v[:dy, dx:] = v[:dy, dx:][::-1, ::-1]
v[dy:, dx:] = v[dy:, dx:][::-1, ::-1]
print()
print(v)
*/

__device__ void reverse_block(const int i, float *a, const int nx, const int ny, const int nz,
                              const int dx0, const int dy0,const int dz0)
{
  // x,y, z indices
  const int ix = i % nx;
  const int iz = i / (ny * nx);
  const int iy = (i - iz * ny * nx) / nx;
  // Make sure dc, dy dz are >=0
  const int dx = (dx0 + nx) % nx;
  const int dy = (dy0 + ny) % ny;
  const int dz = (dz0 + nz) % nz;
  if(ix<dx)
  {
    if(iy<dy)
    {
      if(iz<dz)
      { // ix < dx ; iy < dy ; iz < dz
        if((ix + dx * (iy + dy * iz)) < ((dx*dy*dz)/2))
        { // flip from first half block
          const float v0 = a[i];
          // flipped coordinates
          const int ix1 = dx-ix-1;
          const int iy1 = dy-iy-1;
          const int iz1 = dz-iz-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
      else
      { // ix < dx ; iy < dy ; iz >= dz
        // coordinate in block
        const int iz0 = iz - dz;
        if((ix + dx * (iy + dy * iz0)) < ((dx*dy*(nz-dz))/2))
        {// flip from first half block
          const float v0 = a[i];
          // flipped coordinates
          const int ix1 = dx-ix-1;
          const int iy1 = dy-iy-1;
          const int iz1 = nz-iz0-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
    }
    else
    { //iy>=dy
      if(iz<dz)
      { // ix < dx ; iy >= dy ; iz < dz
        // coordinate in block
        const int iy0 = iy - dy;
        if((ix + dx * (iy0 + (ny-dy) * iz)) < ((dx*(ny-dy)*dz)/2))
        { // flip from first half block
          const float v0 = a[i];
          // flipped coordinates
          const int ix1 = dx-ix-1;
          const int iy1 = ny-iy0-1;
          const int iz1 = dz-iz-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
      else
      { // ix < dx ; iy >= dy ; iz >= dz
        // coordinate in block
        const int iy0 = iy - dy;
        const int iz0 = iz - dz;
        if((ix + dx * (iy0 + (ny-dy) * iz0)) < ((dx*(ny-dy)*(nz-dz))/2))
        {// flip from first half block
          const float v0 = a[i];
          // flipped coordinates
          const int ix1 = dx-ix-1;
          const int iy1 = ny-iy0-1;
          const int iz1 = nz-iz0-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
    }
  }
  else
  { //ix>=dx
    if(iy<dy)
    {
      if(iz<dz)
      { // ix >= dx ; iy < dy ; iz < dz
        // coordinate in block
        const int ix0 = ix - dx;
        if((ix0 + (nx-dx) * (iy + dy * iz)) < (((nx-dx)*dy*dz)/2))
        { // flip from first half block
          const float v0 = a[i];
          // flipped coordinates
          const int ix1 = nx-ix0-1;
          const int iy1 = dy-iy-1;
          const int iz1 = dz-iz-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
      else
      { // ix < dx ; iy < dy ; iz >= dz
        // coordinate in block
        const int ix0 = ix - dx;
        const int iz0 = iz - dz;
        if((ix0 + (nx-dx) * (iy + dy * iz0)) < (((nx-dx)*dy*(nz-dz))/2))
        {// flip from first half block
          const float v0 = a[i];
          // flipped coordinates
          const int ix1 = nx-ix0-1;
          const int iy1 = dy-iy-1;
          const int iz1 = nz-iz0-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
    }
    else
    { //iy>=dy
      if(iz<dz)
      { // ix < dx ; iy >= dy ; iz < dz
        // coordinate in block
        const int ix0 = ix - dx;
        const int iy0 = iy - dy;
        if((ix0 + (nx-dx) * (iy0 + (ny-dy) * iz)) < (((nx-dx)*(ny-dy)*dz)/2))
        { // flip from first half block
          const float v0 = a[i];
          // flipped coordinates
          const int ix1 = nx-ix0-1;
          const int iy1 = ny-iy0-1;
          const int iz1 = dz-iz-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
      else
      { // ix < dx ; iy >= dy ; iz >= dz
        // coordinate in block
        const int ix0 = ix - dx;
        const int iy0 = iy - dy;
        const int iz0 = iz - dz;
        if((ix0 + (nx-dx) * (iy0 + (ny-dy) * iz0)) < (((nx-dx)*(ny-dy)*(nz-dz))/2))
        {// flip from first half block
          const float v0 = a[i];
          // flipped coordinates
          const int ix1 = nx-ix0-1;
          const int iy1 = ny-iy0-1;
          const int iz1 = nz-iz0-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
    }
  }
}

/** Same as reverse_blocks, for complexf arrays
*
*/
__device__ void reverse_block_c(const int i, complexf *a, const int nx, const int ny, const int nz,
                                const int dx0, const int dy0,const int dz0)
{
  // x,y, z indices
  const int ix = i % nx;
  const int iz = i / (ny * nx);
  const int iy = (i - iz * ny * nx) / nx;
  // Make sure dc, dy dz are >=0
  const int dx = (dx0 + nx) % nx;
  const int dy = (dy0 + ny) % ny;
  const int dz = (dz0 + nz) % nz;
  if(ix<dx)
  {
    if(iy<dy)
    {
      if(iz<dz)
      { // ix < dx ; iy < dy ; iz < dz
        if((ix + dx * (iy + dy * iz)) < ((dx*dy*dz)/2))
        { // flip from first half block
          const complexf v0 = a[i];
          // flipped coordinates
          const int ix1 = dx-ix-1;
          const int iy1 = dy-iy-1;
          const int iz1 = dz-iz-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
      else
      { // ix < dx ; iy < dy ; iz >= dz
        // coordinate in block
        const int iz0 = iz - dz;
        if((ix + dx * (iy + dy * iz0)) < ((dx*dy*(nz-dz))/2))
        {// flip from first half block
          const complexf v0 = a[i];
          // flipped coordinates
          const int ix1 = dx-ix-1;
          const int iy1 = dy-iy-1;
          const int iz1 = nz-iz0-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
    }
    else
    { //iy>=dy
      if(iz<dz)
      { // ix < dx ; iy >= dy ; iz < dz
        // coordinate in block
        const int iy0 = iy - dy;
        if((ix + dx * (iy0 + (ny-dy) * iz)) < ((dx*(ny-dy)*dz)/2))
        { // flip from first half block
          const complexf v0 = a[i];
          // flipped coordinates
          const int ix1 = dx-ix-1;
          const int iy1 = ny-iy0-1;
          const int iz1 = dz-iz-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
      else
      { // ix < dx ; iy >= dy ; iz >= dz
        // coordinate in block
        const int iy0 = iy - dy;
        const int iz0 = iz - dz;
        if((ix + dx * (iy0 + (ny-dy) * iz0)) < ((dx*(ny-dy)*(nz-dz))/2))
        {// flip from first half block
          const complexf v0 = a[i];
          // flipped coordinates
          const int ix1 = dx-ix-1;
          const int iy1 = ny-iy0-1;
          const int iz1 = nz-iz0-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
    }
  }
  else
  { //ix>=dx
    if(iy<dy)
    {
      if(iz<dz)
      { // ix >= dx ; iy < dy ; iz < dz
        // coordinate in block
        const int ix0 = ix - dx;
        if((ix0 + (nx-dx) * (iy + dy * iz)) < (((nx-dx)*dy*dz)/2))
        { // flip from first half block
          const complexf v0 = a[i];
          // flipped coordinates
          const int ix1 = nx-ix0-1;
          const int iy1 = dy-iy-1;
          const int iz1 = dz-iz-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
      else
      { // ix < dx ; iy < dy ; iz >= dz
        // coordinate in block
        const int ix0 = ix - dx;
        const int iz0 = iz - dz;
        if((ix0 + (nx-dx) * (iy + dy * iz0)) < (((nx-dx)*dy*(nz-dz))/2))
        {// flip from first half block
          const complexf v0 = a[i];
          // flipped coordinates
          const int ix1 = nx-ix0-1;
          const int iy1 = dy-iy-1;
          const int iz1 = nz-iz0-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
    }
    else
    { //iy>=dy
      if(iz<dz)
      { // ix < dx ; iy >= dy ; iz < dz
        // coordinate in block
        const int ix0 = ix - dx;
        const int iy0 = iy - dy;
        if((ix0 + (nx-dx) * (iy0 + (ny-dy) * iz)) < (((nx-dx)*(ny-dy)*dz)/2))
        { // flip from first half block
          const complexf v0 = a[i];
          // flipped coordinates
          const int ix1 = nx-ix0-1;
          const int iy1 = ny-iy0-1;
          const int iz1 = dz-iz-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
      else
      { // ix < dx ; iy >= dy ; iz >= dz
        // coordinate in block
        const int ix0 = ix - dx;
        const int iy0 = iy - dy;
        const int iz0 = iz - dz;
        if((ix0 + (nx-dx) * (iy0 + (ny-dy) * iz0)) < (((nx-dx)*(ny-dy)*(nz-dz))/2))
        {// flip from first half block
          const complexf v0 = a[i];
          // flipped coordinates
          const int ix1 = nx-ix0-1;
          const int iy1 = ny-iy0-1;
          const int iz1 = nz-iz0-1;
          // exchange values
          a[i] = a[ix1 + nx * (iy1 + ny * iz1)];
          a[ix1 + nx * (iy1 + ny * iz1)] = v0;
        }
      }
    }
  }
}
