// Class for managing the Discrete FFT transformation
// George F. Riley
// ECE3090 Project 4

#include <iostream>
#include <fstream>
#include <string>
#include <vector>

#include "complex.h"
#include "fft.h"

using namespace std;

// Constructors
FFT::FFT(const vector<Complex>& inputs) 
{ // Construct an FFT object with data from the inputs vector
  // Your code here
  N = inputs.size();
  h = new Complex[N];
  W = new Complex[N/2];
  // Load the initial H array
  for (unsigned i = 0; i < N; ++i) h[i] = inputs[i];
}

void FFT::DumpH(ostream& os)
{
  for (unsigned i = 0; i < N; ++i)
    {
      os << i << " - " << h[i] << endl;
      // match matlab output
      //printf("   %9.7e\n", h[i].real);
    }
}

void FFT::DumpW(ostream& os)
{
  for (unsigned i = 0; i < N/2; ++i)
    {
      os << i << " - " << W[i] << endl;
    }
}

void FFT::LoadWeights()
{
  for (unsigned i = 0; i < N/2; ++i)
    { // Complex value representing e exp(2pi*i/N*n)
      double theta = M_PI * 2 * i/ (double)N;
      W[i] = Complex(cos(theta), -sin(theta));
    }
}

void FFT::BitReverseElements()
{ // Shuffle the h elements per the Cooley-Tukey bit-reverse algorithm
  for (unsigned i = 0; i < N; ++i)
    {
      unsigned j = ReverseBits(i);
      if (i < j)
        { // Swap the elements
          Complex t = h[i];
          h[i] = h[j];
          h[j] = t;
        }
    }
}

void FFT::TransformSimple()
{
  Complex* H = new Complex[N]; // Results
  for (unsigned n = 0; n < N; ++n)
    {
      H[n] = Complex(0.0);
      for (unsigned k = 0; k < N; ++k)
        {
          double theta = M_PI * 2 * k * n/ (double)N;
          H[n] = H[n] + h[k] * Complex(cos(theta), -sin(theta));
        }
    }
  // Write the results
  for (unsigned i = 0; i < N; ++i)
    {
      if (fabs(H[i].real) < 1e-12) H[i].real = 0;
      if (fabs(H[i].imag) < 1e-12) H[i].imag = 0;
      cout << i << " - " << H[i] << endl;
    }
}

void FFT::Transform()
{ // Do the transform
  //cout << "Initial Data Points" << endl;
  //DumpH(cout);
  
  BitReverseElements(); // Shuffle the elements with the bit reversal
  LoadWeights();        // Calculate the W weight factors
  //cout << "After bit reverse" << endl;
  //DumpH(cout);
  //cout << "Weights are" << endl;
  //DumpW(cout);
  for (unsigned nPoints = 2; nPoints <= N; nPoints *= 2)
    { // Perform each sub-transform
      unsigned delta = nPoints / 2; // Divides into even half and odd half
      for (unsigned i = 0; i < N; i+=nPoints)
        {
          for (unsigned j = 0; j < nPoints/2; ++j)
            {
              Complex hTemp = h[i+j];
              h[i+j] = h[i+j] + W[j*N/nPoints]*h[i+j+delta]; 
              h[i+j+delta] = hTemp - W[j*N/nPoints]*h[i+j+delta];
            }
        }
      //cout << "After transform length " << nPoints << endl;
      //DumpH(cout);
    }
  DumpH(cout);
}

   
// Private functions
unsigned FFT::ReverseBits(unsigned v)
{ //  Provided to students
  unsigned n = N; // Size of array (which is even 2 power k value)
  unsigned r = 0; // Return value
   
  for (--n; n > 0; n >>= 1)
    {
      r <<= 1;        // Shift return value
      r |= (v & 0x1); // Merge in next bit
      v >>= 1;        // Shift reversal value
    }
  return r;
}

bool FFT::ReadInputs(char* fn, vector<Complex>& data)
{ // Read inputs from a file, store in the "data" vector
  ifstream f(fn);
  if (!f)
    {
      cout << "Can't open input data file " << fn << endl;
      return false;
    }
  while(f)
    {
      string line;
      getline(f, line);
      if (line.length() == 0) continue; // empty
      if (line[0] == '#')     continue; // comment
      data.push_back(Complex(atof(line.c_str()))); // Add to data vector
    }
  return true; // Read correctly
}

int main(int argc, char** argv)
{
  if (argc < 2)
    {
      cout << "Usage: fft filename" << endl;
      exit(1);
    }
  bool simple = false;
  if (argc > 2) simple = true;
  vector<Complex> inputs;            // Holds the time domain samples
  FFT::ReadInputs(argv[1], inputs);  // Reads the samples from the input file
  cout << "Found " << inputs.size() << " input samples" << endl;
  FFT f(inputs);                     // Creates the FFT object
  if (simple)
    {
      f.TransformSimple();           // Do the transform (simple, inefficient)
    }
  else
    {
      f.Transform();                  // Do the transform
    }
}