// Demonstration that inverse real FFT is giving incorrect results.
//
// June 25, 2018
// HIP version: 1.5.18151
// rocfft version 0.8.1
//
// Can be compiled with:
//
//  hipcc -o rocfft_real_inverse rocfft_real_inverse.cpp -I/opt/rocmfft/include -L/opt/rocfft/lib -lrocfft
//

#include <iostream>
#include <hip/hip_runtime.h>
#include <hip/hip_vector_types.h>
#include "rocfft/rocfft.h"
#include <algorithm>
#include <cmath>


// Parameters
const int width = 12;   // innermost
const int height = 4;

// Derived_parameters.
const int N = width * height;
const int hermitian_width = width/2 + 1;
const int N_hermitian = hermitian_width * height;


#define FFT_SAFE(x) { \
  rocfft_status status = x; \
  if ( status != rocfft_status_success ) { \
    std::cerr << "Error " << status << std::endl; \
    exit(-1); \
  }}

// -------------------------------------------------------------------------

void inverse_real_fft(
  float* in,
  float* out )
{
  // Create rocFFT plan
  rocfft_plan inverse_plan = nullptr;
  size_t length[2] = { width, height };
  FFT_SAFE( rocfft_plan_create(
              &inverse_plan,
              rocfft_placement_notinplace,
              rocfft_transform_type_real_inverse,
              rocfft_precision_single,
              2,
              length,
              1,
              nullptr ) );
  size_t device_workspace_size;
  rocfft_plan_get_work_buffer_size( inverse_plan, &device_workspace_size );
  char* device_workspace = nullptr;
  rocfft_execution_info info = nullptr;
  hipMalloc(&device_workspace, device_workspace_size);
  rocfft_execution_info_create(&info);
  rocfft_execution_info_set_work_buffer(info, device_workspace, device_workspace_size);

  // Create HIP device buffer for K-space data (input to inverse FFT)
  float *d_in;
  hipMalloc( &d_in, 2*sizeof(float)*N_hermitian );
  // Create HIP device buffer for real-space data (output to inverse FFT)
  float *d_out;
  hipMalloc( &d_out, sizeof(float)*N );

  // Copy data to device
  hipMemcpy( d_in, in, 2*sizeof(float)*N_hermitian, hipMemcpyHostToDevice );

  // Execute inverse plan
  FFT_SAFE( rocfft_execute( inverse_plan, (void**)&d_in, (void**)&d_out, info) );

  // Wait for execution to finish
   hipDeviceSynchronize();

  // Copy data back to host
  hipMemcpy( out, d_out, sizeof(float)*N, hipMemcpyDeviceToHost );

  // Clean-up
  hipFree( device_workspace );
  hipFree( d_out );
  hipFree( d_in );
  FFT_SAFE( rocfft_execution_info_destroy(info) );
  FFT_SAFE( rocfft_plan_destroy(inverse_plan) );
}

// -------------------------------------------------------------------------

void forward_real_fft(
  float* in,
  float* out )
{
  // Create rocFFT plan
  rocfft_plan forward_plan = nullptr;
  size_t length[2] = { width, height };
  FFT_SAFE( rocfft_plan_create(
              &forward_plan,
              rocfft_placement_notinplace,
              rocfft_transform_type_real_forward,
              rocfft_precision_single,
              2,
              length,
              1,
              nullptr ) );
  size_t device_workspace_size;
  rocfft_plan_get_work_buffer_size( forward_plan, &device_workspace_size );
  char* device_workspace = nullptr;
  rocfft_execution_info info = nullptr;
  hipMalloc(&device_workspace, device_workspace_size);
  rocfft_execution_info_create(&info);
  rocfft_execution_info_set_work_buffer(info, device_workspace, device_workspace_size);

  // Create HIP device buffer for real-space data (input to forward FFT)
  float *d_in;
  hipMalloc( &d_in, sizeof(float)*N );
  // Create HIP device buffer K-space data (output to forward FFT)
  float *d_out;
  hipMalloc( &d_out, 2*sizeof(float)*N_hermitian );

  // Copy data to device
  hipMemcpy( d_in, in, sizeof(float)*N, hipMemcpyHostToDevice );

  // Execute forward plan
  FFT_SAFE( rocfft_execute( forward_plan, (void**)&d_in, (void**)&d_out, info) );

  // Wait for execution to finish
  hipDeviceSynchronize();

  // Copy data back to host
  hipMemcpy( out, d_out, 2*sizeof(float)*N_hermitian, hipMemcpyDeviceToHost );

  // Clean-up
  hipFree( device_workspace );
  hipFree( d_out );
  hipFree( d_in );
  FFT_SAFE( rocfft_execution_info_destroy(info) );
  FFT_SAFE( rocfft_plan_destroy(forward_plan) );
}

// -------------------------------------------------------------------------

void inverse_complex_fft(
  float* in,
  float* out )
{
  // Create rocFFT plan
  rocfft_plan inverse_plan = nullptr;
  size_t length[2] = { width, height };
  FFT_SAFE( rocfft_plan_create(
              &inverse_plan,
              rocfft_placement_notinplace,
              rocfft_transform_type_complex_inverse,
              rocfft_precision_single,
              2,
              length,
              1,
              nullptr ) );
  size_t device_workspace_size;
  rocfft_plan_get_work_buffer_size( inverse_plan, &device_workspace_size );
  char* device_workspace = nullptr;
  rocfft_execution_info info = nullptr;
  hipMalloc(&device_workspace, device_workspace_size);
  rocfft_execution_info_create(&info);
  rocfft_execution_info_set_work_buffer(info, device_workspace, device_workspace_size);

  // Create HIP device buffer for K-space data (input to inverse FFT)
  float *d_in;
  hipMalloc( &d_in, 2*sizeof(float)*N );
  // Create HIP device buffer for real-space data (output to inverse FFT)
  float *d_out;
  hipMalloc( &d_out, 2*sizeof(float)*N );

  // Copy data to device
  hipMemcpy( d_in, in, 2*sizeof(float)*N, hipMemcpyHostToDevice );

  // Execute inverse plan
  FFT_SAFE( rocfft_execute( inverse_plan, (void**)&d_in, (void**)&d_out, info) );

  // Wait for execution to finish
  hipDeviceSynchronize();

  // Copy data to back to host
  hipMemcpy( out, d_out, 2*sizeof(float)*N, hipMemcpyDeviceToHost );

  // Clean-up
  hipFree( device_workspace );
  hipFree( d_out );
  hipFree( d_in );
  FFT_SAFE( rocfft_execution_info_destroy(info) );
  FFT_SAFE( rocfft_plan_destroy(inverse_plan) );
}

// -------------------------------------------------------------------------

// Utility function to print maximum amplitude (over both arrays),
// and the maximum difference (amplitude) between the arrays.
void compare_complex_to_real( float* c, float* r, int count )
{
  float max_amplitude = 0;
  float max_difference = 0;
  for ( int i=0; i<count; ++i ) {
    max_amplitude = std::max( max_amplitude, r[i] );
    float c_real = c[2*i];
    float c_imag = c[2*i+1];
    max_amplitude = std::max( max_amplitude, std::sqrt(c_real*c_real + c_imag*c_imag) );
    float d_real = r[i] - c_real;
    float d_imag = c_imag;
    max_difference = std::max( max_difference, std::sqrt(d_real*d_real + d_imag*d_imag) );
  }
  std::cout << "  max_amplitude  = " << max_amplitude
          << "\n  max_difference = " << max_difference << std::endl;
}

// -------------------------------------------------------------------------

// Utility function to print maximum amplitude (over both arrays),
// and the maximum difference (amplitude) between the arrays.
void compare_complex_to_complex( float* c1, float* c2, int count )
{
  float max_amplitude = 0;
  float max_difference = 0;
  for ( int i=0; i<count; ++i ) {
    float c1_real = c1[2*i];
    float c1_imag = c1[2*i+1];
    max_amplitude = std::max( max_amplitude, std::sqrt(c1_real*c1_real + c1_imag*c1_imag) );
    float c2_real = c2[2*i];
    float c2_imag = c2[2*i+1];
    max_amplitude = std::max( max_amplitude, std::sqrt(c2_real*c2_real + c2_imag*c2_imag) );
    float d_real = c1_real - c2_real;
    float d_imag = c1_imag - c2_imag;
    max_difference = std::max( max_difference, std::sqrt(d_real*d_real + d_imag*d_imag) );
  }
  std::cout << "  max_amplitude  = " << max_amplitude
          << "\n  max_difference = " << max_difference << std::endl;
}

// -------------------------------------------------------------------------

int main( int argc, char **argv )
{
  rocfft_setup();

  // ===== Input data: hermitian =====

  // Create some data in k-space: this data is the usual hermitian half of the complex values
  // Will treat this as interleaved complex.
  float* K_hermitian = (float*)malloc( 2*sizeof(float)*N_hermitian );

  // These strides are in size of floats, since we will use float pointers
  size_t hermitian_strides[2] = { 2, 2*hermitian_width };

  // Set all values to zero except for a single non-zero value.
  for ( int i=0; i<2*N_hermitian; ++i ) {
    K_hermitian[i] = 0;
  }
  // Set element k=3, l=1
  int k = 3;
  int l = 1;
  K_hermitian[ k*hermitian_strides[0] + l*hermitian_strides[1]    ] = 1.0f;   // real part
  K_hermitian[ k*hermitian_strides[0] + l*hermitian_strides[1] +1 ] = 1.0f;   // imag part

  // ===== Input data: complete complex data, not stored in hermitian form =====

  float* K_complex = (float*)malloc( 2*sizeof(float)*N );

  // These strides are in size of floats, since we will use float pointers
  size_t complex_strides[2] = { 2, 2*width };

  // Set all values to zero initially.
  for ( int i=0; i<2*N; ++i ) {
    K_complex[i] = 0;
  }

  // To be equivalent to K_hermitian, must set the same element, plus the
  // one satifying the Hermitian property.
  K_complex[ k*complex_strides[0] + l*complex_strides[1]    ] = 1.0f;   // real part
  K_complex[ k*complex_strides[0] + l*complex_strides[1] +1 ] = 1.0f;   // imag part
  K_complex[ (width-k)*complex_strides[0] + (height-l)*complex_strides[1]    ] =  1.0f;   // real part
  K_complex[ (width-k)*complex_strides[0] + (height-l)*complex_strides[1] +1 ] = -1.0f;   // imag part

  // ===== Perform the transforms =====

  // Allocate for output data
  float *X_real = (float*)malloc( sizeof(float)*N );
  float *X_complex = (float*)malloc( 2*sizeof(float)*N );

  inverse_real_fft( K_hermitian, X_real );
  inverse_complex_fft( K_complex, X_complex );

  std::cout << "Compare inverse real FFT and inverse complex FFT:\n";
  compare_complex_to_real( X_complex, X_real, N );

  // ==== Use forward real-to-complex FFT on X_real to see if we recover the original K_hermitian ====
  // (Answer: no)

  float* K_hermitian_round_trip = (float*)malloc(2*sizeof(float)*N_hermitian);
  forward_real_fft( X_real, K_hermitian_round_trip );
  // normalize
  for ( int i=0; i<2*N_hermitian; ++i ) {
    K_hermitian_round_trip[i] /= N;
  }
  std::cout << "Compare forward real FFT of inverse real FFT to original K-space data:\n";
  compare_complex_to_complex( K_hermitian, K_hermitian_round_trip, N_hermitian );

  // ==== Use forward real-to-complex FFT on X_complex (real part) to see if we recover the original K_hermitian ====
  // (Answer: yes)

  // For X_complex, need to extract just the real part (the imaginary parts are zero: can be checked)
  float *X_real_from_complex = (float*)malloc( sizeof(float)*N );
  for ( int i=0; i<N; ++i ) {
    X_real_from_complex[i] = X_complex[2*i];
  }

  float* K_hermitian_from_complex = (float*)malloc(2*sizeof(float)*N_hermitian);
  forward_real_fft( X_real_from_complex, K_hermitian_from_complex );
  // normalize
  for ( int i=0; i<2*N_hermitian; ++i ) {
    K_hermitian_from_complex[i] /= N;
  }
  std::cout << "Compare forward real FFT of inverse complex FFT to original K-space data:\n";
  compare_complex_to_complex( K_hermitian, K_hermitian_from_complex, N_hermitian );

  // ==== Clean up ====

  free( K_hermitian_round_trip );
  free( X_complex );
  free( X_real );
  free( K_complex );
  free( K_hermitian );

  rocfft_cleanup();
}