// Example 5.2.  See page 37 of the FG tutorial.

#include "FG.h"
#include <complex>
#include <iostream>
#include <cmath>
using namespace std;

typedef complex<double> complex_double;

// Class to package together the parameters needed to compute FFTs.
// By making it a class, the destructor takes care of deallocating the
// array of level numbers.
class FFT_macro_params {
public:
  FFT_macro_params(int logn);   // constructor, takes number of levels
  ~FFT_macro_params();          // destructor
  int logn;                     // number of levels
  int *levels;                  // array of level numbers
};

// Constructor.  Stores the number of levels and creates the array of
// level numbers.
FFT_macro_params::FFT_macro_params(int logn)
{
  this->logn = logn;            // save the number of levels

  // Create and initialize the array of level numbers with 0, 1, ...,
  // log n - 1.
  levels = new int [logn];
  for (int j = 0; j < logn; j++)
    levels[j] = j;
}

// Destructor.  Deletes the array of level numbers.
FFT_macro_params::~FFT_macro_params()
{
  delete [] levels;
}

// Function prototypes.
FG_macro make_fft_macro(FG_pipeline_thread_helper fft_thread,
                        FFT_macro_params &FFT_params);
int log2(int n);
int bit_rev_perm_func(FG_stage my_stage, FG_stage_params my_params);
int level_func(FG_stage my_stage, FG_stage_params my_params);
int spike_func(FG_stage my_stage, FG_stage_params my_params);
int check_func(FG_stage my_stage, FG_stage_params my_params);

int main()
{
  const int N = 8;
  int logn = log2(N);
  int size = N * sizeof(complex_double);

  // Declare the threads, and store them in a NULL-terminated array.
  FG_pipeline_thread_helper
    spike_thread = new FG_pipeline_thread_helper_info(),
    fft_thread = new FG_pipeline_thread_helper_info(),
    check_thread = new FG_pipeline_thread_helper_info();

  FG_pipeline_thread_helper threads[] = {spike_thread, fft_thread,
                                         check_thread, NULL};

  // Declare the stages.
  FG_pipeline_stage_helper
    spike_stage = new FG_pipeline_stage_helper_info(spike_thread),
    check_stage = new FG_pipeline_stage_helper_info(check_thread);

  // Create the parameters and macro for FFT.  The level numbers, as
  // members of the FFT_params object, are created before the
  // FG_macro.  The addresses of the level numbers are passed to the
  // stages as parameters.  Once the pipeline finishes, it's OK to
  // destroy the FFT_params object.
  FFT_macro_params FFT_params(logn);
  FG_macro FFT_macro = make_fft_macro(fft_thread, FFT_params);

  // Store the stages in a NULL-terminated array of FG_node.
  FG_node stages[] = {spike_stage, FFT_macro, check_stage, NULL};

  // Set the stage functions.
  spike_stage->set_func(spike_func, NULL);
  check_stage->set_func(check_func, NULL);

  // Instantiate the pipeline.
  FG_pipeline my_pipeline = FG_pipeline_info::create(stages, threads);

  // Set the pipeline variables.
  my_pipeline->set_buff_size(size);
  my_pipeline->set_num_rounds(N);
  my_pipeline->set_num_buffs(logn+3);

  // Fix the pipeline, which also runs it, and store the return code.
  int error = my_pipeline->fix();
  
  // Display the return code.
  cout << "FINISHED PIPELINE WITH ERROR CODE " << error << endl;

  // Dismantle the pipeline.
  my_pipeline->dismantle();

  return 0;
}

// Make the macro for the FFT code.  This function is given a thread
// and an FFT_params object (already initialized), and it creates all
// the stages needed to perform an FFT, packages them into an
// FG_macro, and returns this FG_macro to the caller.
FG_macro make_fft_macro(FG_pipeline_thread_helper fft_thread,
                        FFT_macro_params &FFT_params)
{
  int logn = FFT_params.logn;   // store logn into a local variable

  // Allocate space for threads.
  FG_pipeline_thread_helper threads[logn+1];

  // Store threads in an array.
  for (int j = 0; j <= logn; j++)
    threads[j] = new FG_pipeline_thread_helper_info();

  // Allocate space for stages.
  FG_node *stages = new FG_node[logn+2];

  // Define the stages and store them in a NULL-terminated array of FG_node.
  for (int j = 0; j <= logn; j++)
    stages[j] = new FG_pipeline_stage_helper_info(threads[j]);
  stages[logn+1] = NULL;

  // Define the stage functions.
  ((FG_pipeline_stage_helper) stages[0])->set_func(bit_rev_perm_func,
                                                   &FFT_params.logn);
  for (int j = 1; j <= logn; j++)
    ((FG_pipeline_stage_helper) stages[j])->
      set_func(level_func, &FFT_params.levels[j-1]);

  // Now we're ready to make the macro.
  FG_macro FFT_macro = new FG_macro_info(fft_thread);

  // Add the stages we just made to the FFT macro.
  FFT_macro->add_children(stages);

  // Return the macro we've created.
  return FFT_macro;
} 

// Return the log-base-2 of n.
int log2(int n)
{
  int result = -1;

  while (n > 0) {
    n >>= 1;
    result++;
  }

  return result;
}

// Return the bit reversal of x, to n bits.
int bitRevOf(int x, int n)
{
  int result = 0;
  int bits;
  static char byteRev[5][16] = {
    { 0x0 },
    { 0x0, 0x1 },
    { 0x0, 0x2, 0x1, 0x3 },
    { 0x0, 0x4, 0x2, 0x6, 0x1, 0x5, 0x3, 0x7 },
    { 0x0, 0x8, 0x4, 0xC, 0x2, 0xA, 0x6, 0xE,
      0x1, 0x9, 0x5, 0xD, 0x3, 0xB, 0x7, 0xF }
  };

  for ( ; n > 0; n -= 4, x >>= 4) {
    bits = n > 4 ? 4 : n;
    result = (result << bits) | byteRev[bits][x & 0xF];
  }

  return result;
}

// Perform a bit-reversal permutation on a buffer.
int bit_rev_perm_func(FG_stage my_stage, FG_stage_params my_params)
{
  // Accept a thumbnail from the preceding stage, and save the
  // location of its buffer.
  FG_pipeline_thumbnail thumb = my_stage->accept_thumbnail();
  complex_double *buff = (complex_double *) thumb->get_address();

  // The number of elements in the buffer is in the thumbnail, but the
  // log of the number of elements is the stage parameter.
  int n = thumb->get_size() / sizeof(complex_double);
  int logn = * (int *) my_params;

  // Do the bit-reversal permutation on the buffer.
  for (int j = 0; j < n; j++) {
    int revj = bitRevOf(j, logn);
    if (j < revj) {
      complex_double t = buff[j];
      buff[j] = buff[revj];
      buff[revj] = t;
    }
  }

  // Convey the thumbnail to the next stage.
  my_stage->convey_thumbnail(thumb);
  
  return 0;
}

// Perform one level of the butterfly network.
int level_func(FG_stage my_stage, FG_stage_params my_params)
{
  // Accept a thumbnail from the preceding stage, and save the
  // location of its buffer.
  FG_pipeline_thumbnail thumb = my_stage->accept_thumbnail();
  complex_double *buff = (complex_double *) thumb->get_address();

  // The number of elements in the buffer is in the thumbnail.
  int n = thumb->get_size() / sizeof(complex_double);

  // The level is the stage parameter.
  int level = * (int *) my_params;

  int gap = 1 << level;         // gap between butterfly partners
  int groups = n / (2 * gap);   // number of groups in this level
  double nth = 2 * M_PI / (1 << (level + 1)); // 2 * PI / n

  complex_double omega_n(cos(nth), sin(nth));

  for (int group = 0; group < groups; group++) {
    complex_double omega(1, 0);
    int group_start = 2 * group * gap; // index of first position in this group

    for (int j = 0; j < gap; j++) {
      complex_double t = omega * buff[group_start + j + gap];
      complex_double u = buff[group_start + j];
      buff[group_start + j] = u + t;
      buff[group_start + j + gap] = u - t;
      omega *= omega_n;
    }
  }

  // Convey the thumbnail to the next stage.
  my_stage->convey_thumbnail(thumb);
  
  return 0;
}

// Generate a spike vector with the spike in the position given by the
// round number.
int spike_func(FG_stage my_stage, FG_stage_params my_params)
{
  // Accept a thumbnail from the preceding stage, and save the
  // location of its buffer.
  FG_pipeline_thumbnail thumb = my_stage->accept_thumbnail();
  complex_double *buff = (complex_double *) thumb->get_address();

  memset(buff, 0, thumb->get_size());

  buff[thumb->get_round()] = complex_double(1, 0);

  // Convey the thumbnail to the next stage.
  my_stage->convey_thumbnail(thumb);
  
  return 0;
}

// Print the contents of a buffer, with its round number.
int check_func(FG_stage my_stage, FG_stage_params my_params)
{
  // Accept a thumbnail from the preceding stage, and save the
  // location of its buffer.
  FG_pipeline_thumbnail thumb = my_stage->accept_thumbnail();
  complex_double *buff = (complex_double *) thumb->get_address();

  int spike = (int)thumb->get_round();
  const double TOLERANCE = 1.0e-9;

  // The number of elements in the buffer is in the thumbnail.
  int n = thumb->get_size() / sizeof(complex_double);

  bool badness = false;         // no error until we find one
  for (int j = 0; j < n; j++) {
    // The target value for position j is e^{j * spike * 2 * pi / n}.
    double t_real = cos(spike * j * 2 * M_PI / n);
    double t_imag = sin(spike * j * 2 * M_PI / n);
    complex_double target(t_real, t_imag);

    // Badness if either the real or the imaginary part of the result
    // is not within the tolerance.
    if (fabs(t_real - buff[j].real()) > TOLERANCE ||
        fabs(t_imag - buff[j].imag()) > TOLERANCE) {
        cout << "error in position " << j
             << ": target = " << target
             << ", actual = " << buff[j] << endl;
        badness = true;
      }
    }

  cout << "Done checking for round " << thumb->get_round();
  if (badness)
    cout << " and there are errors.\n";
  else
    cout << " and all is well.\n";

  // Convey the thumbnail to the next stage.
  my_stage->convey_thumbnail(thumb);
  
  return 0;
}