#include <pycuda-complex.hpp>
typedef pycuda::complex<float> complexf;
#define TWOPI 6.2831853071795862f

/** Second stage of VSNR operation: compute
* FT(X) = [F(-lambda_x+beta*y_x)*psi_x.conj() + F(-lambda_y+beta*y_y)*psi_y.conj()] / fphi
* where fphi = 1 + beta * (|psi_x|^2 + |psi_y|^2)
* and then storing FT(X)*FT(Psi_x) and FT(X)*FT(Psi_y) in yx and yy arrays
*
* \param yx, yy: on input, holding the r2c FT of -lambda_x+beta*y_x and -lambda_y+beta*y_y
* \param psif: the real part of the r2c FT of the filter to apply (assumed to be centrosymmetric)
* \param ny: vertical size of the images
* \param nxh: horizontal size of the half-Hermitian array
* \param beta: beta ADMM parameter
*/
__device__ void vsnr2(const int i, complexf *yx, complexf *yy, float* psif, const int ny, const int nxh, const float beta)
{
    const int nx = (nxh-1)*2;  // arrays are half-hermitian
    const int iz = i / (nxh*ny);
    const int ixy = i - iz * nxh * ny;  // index in one 2D image
    const int ix = ixy % nxh;
    const int iy = ixy / nxh;

    // Compute psif_x and psif_y
    float s,c;
    const float psi = psif[ixy];  // psif is real as original Psi is centrosymmetric

    __sincosf((float)(TWOPI * ix)/(float)nx, &s,&c);
    const complexf psifx(psi * (1.0f-c), -psi * s);

    __sincosf((float)(TWOPI * iy)/(float)ny, &s,&c);
    const complexf psify(psi * (1.0f-c), -psi * s);

    // compute fphi = 1 + beta * (psifx**2 + psify**2)
    const float fphi = 1 + beta * (psifx.real() * psifx.real() + psifx.imag() * psifx.imag() + psify.real() * psify.real() + psify.imag() * psify.imag());

    // Compute FT(X) - not stored
    const complexf yx0 = yx[i];
    const complexf yy0 = yy[i];
    const complexf x = complexf(yx0.real()*psifx.real() +  yx0.imag()*psifx.imag() + yy0.real()*psify.real() +  yy0.imag()*psify.imag(),
                               yx0.imag()*psifx.real() -  yx0.real()*psifx.imag() + yy0.imag()*psify.real() -  yy0.real()*psify.imag()) / fphi;
    // Store FT(X) * psif_x in yx [same for y]
    yx[i] = complexf(x.real() * psifx.real() - x.imag() * psifx.imag(), x.real() * psifx.imag() + x.imag() * psifx.real());
    yy[i] = complexf(x.real() * psify.real() - x.imag() * psify.imag(), x.real() * psify.imag() + x.imag() * psify.real());
}

/** Third stage of VSNR operation: compute
* [F(-lambda_x+beta*y_x)*psi_x.conj() + F(-lambda_y+beta*y_y)*psi_y.conj()] / fphi
* where fphi = 1 + beta * (|psi_x|^2 + |psi_y|^2)
*
* \param u0: the stack of 2D images to filter- shape (nz, ny, nx)
* \param lx, ly: the lambda_x and _y arrays, with the same shape as u0
* \param yx, yy: on input, holding IFT[FT(X)*FT(Psi_x)] and IFT[FT(X)*FT(Psi_y)]. On output,
*    the new iteration for lambda_x and lambda_y
* \param psif: the real part of the r2c FT of the Psi filter to apply (Psi assumed to be centrosymmetric,
*    hence the real values)
* \param ny: vertical size of the images
* \param nx: horizontal size of the images
* \param beta: beta ADMM parameter
*/
__device__ void vsnr3(const int i, float *u0, float *lx, float *ly, float *yx, float *yy,
                      const int ny, const int nx, const float beta)
{
   const int iz = i / (nx*ny);
   const int ixy = i - iz * nx * ny;  // index in one 2D image
   const int ix = ixy % nx;
   const int iy = ixy / nx;

   // u0 gradient
   const float u = u0[i];
   const float u0x = u - u0[ i + 1 - nx * (ix==(nx-1))];
   const float u0y = u - u0[ i + nx * (1 - ny * (iy ==(ny-1)))];

   // read AXx and AXy (stored in Y arrays), lambda_x,y
   const float axx = yx[i];
   const float axy = yy[i];
   const float lx0 = lx[i];
   const float ly0 = ly[i];

   // tmpx and y
   const float tmpx = u0x - (axx + lx0 / beta);
   const float tmpy = u0y - (axy + ly0 / beta);

   // norm
   const float ng = fmaxf(sqrtf(tmpx * tmpx + tmpy * tmpy), 1e-6f);

   // Update y
   const float yx1 = u0x - (float)(ng > 1/beta) * tmpx * (1 - 1 / (beta*ng));
   const float yy1 = u0y - (float)(ng > 1/beta) * tmpy * (1 - 1 / (beta*ng));
   yx[i] = yx1;
   yy[i] = yy1;

   // Update lambda
   lx[i] = lx0 + beta * (axx - yx1);
   ly[i] = ly0 + beta * (axy - yy1);
}

/** Final stage of VSNR operation: same as vsnr2, but on output storing FT(X)*FT(Psi) in yx
*
* \param yx, yy: on input, holding the r2c FT of -lambda_x+beta*y_x and -lambda_y+beta*y_y
* \param psif: the real part of the r2c FT of the filter to apply (assumed to be centrosymmetric)
* \param ny: vertical size of the images
* \param nxh: horizontal size of the half-Hermitian array
* \param beta: beta ADMM parameter
*/
__device__ void vsnr4(const int i, complexf *yx, complexf *yy, float* psif, const int ny, const int nxh, const float beta)
{
   const int nx = (nxh-1) * 2;  // arrays are half-hermitian
   const int iz = i / (nxh*ny);
   const int ixy = i - iz * nxh * ny;  // index in one 2D image
   const int ix = ixy % nxh;
   const int iy = ixy / nxh;

   // Compute psif_x and psif_y
   float s,c;
   const float psi = psif[ixy];  // psif is real as original Psi is centrosymmetric

   __sincosf((float)(TWOPI * ix)/(float)nx, &s,&c);
   const complexf psifx(psi * (1.0f-c), -psi * s);

   __sincosf((float)(TWOPI * iy)/(float)ny, &s,&c);
   const complexf psify(psi * (1.0f-c), -psi * s);

   // compute fphi = 1 + beta * (psifx**2 + psify**2)
   const float fphi = 1 + beta * (psifx.real() * psifx.real() + psifx.imag() * psifx.imag() + psify.real() * psify.real() + psify.imag() * psify.imag());

   // Compute FT(X) - not stored
   const complexf yx0 = yx[i];
   const complexf yy0 = yy[i];
   const complexf x = complexf(yx0.real()*psifx.real() +  yx0.imag()*psifx.imag() + yy0.real()*psify.real() +  yy0.imag()*psify.imag(),
                               yx0.imag()*psifx.real() -  yx0.real()*psifx.imag() + yy0.imag()*psify.real() -  yy0.real()*psify.imag()) / fphi;

   // Store FT(X) * psif in yx
   yx[i] = x * psi;
}

/** Apodisation function
*/
__device__ void apod(const int i, float *y, const int w, const int ny, const int nx)
{
    const int ix = i % nx;
    const int iy = (i % (nx * ny)) / nx;
    float cpad = 1.0f;
    if(ix<w)            cpad *= 0.5f * (1 - __cosf(   (float)      ix  * 3.1415926535897932f / (float)w));
    else if(ix>=(nx-w)) cpad *= 0.5f * (1 - __cosf(   (float)(nx-1-ix) * 3.1415926535897932f / (float)w));
    if(iy<w)            cpad *= 0.5f * (1 - __cosf(   (float)      iy  * 3.1415926535897932f / (float)w));
    else if(iy>=(ny-w)) cpad *= 0.5f * (1 - __cosf(   (float)(ny-1-iy) * 3.1415926535897932f / (float)w));
    if(cpad != 1.0f) y[i] *= cpad;
}
