#include <iostream>
#include <fstream>
#include <iomanip>
#include <cmath>
#include <complex>
#include "AudioFile.h"

/*
  Matlab to convert to image:
  w=importdata("sine-wave_waterfall.dat");
  imagesc(w);
  xlabel("Time/block")
  ylabel("wavelength/AU")

  xmgrace spectrum.dat

*/

using namespace std;

int main()
{
    //---------------------------------------------------------------
    // 1. Set a file path to an audio file on your machine

    const std::string base    = "noise";
    const std::string input   = base+".wav";
    const std::string output1 = base+"_waterfall.dat";
    const std::string output2 = base+"_spectrum-rawz.dat";
    const std::string output3 = base+"_spectrum-norm.dat";

    //---------------------------------------------------------------
    // 2. Create an AudioFile object and load the audio file

    AudioFile<float> a;
    bool loadedOK = a.load(input);
    assert (loadedOK);

    cout << "Read: "              << input                  << endl;
    cout << "Bit Depth: "         << a.getBitDepth()        << endl;
    cout << "Sample Rate: "       << a.getSampleRate()      << endl;
    cout << "Num Channels: "      << a.getNumChannels()     << endl;
    cout << "Length in Seconds: " << a.getLengthInSeconds() << endl;

    double rate=a.getSampleRate();
    int blocksize = 4410;
    int nblocks = a.getNumSamplesPerChannel()/blocksize;

    double fmin=rate/blocksize;
    double fmax=rate/2;

    cout << "Number of blocks: " << nblocks << endl;
    cout << "Period of one block: " << blocksize/rate << "s" << endl;

    cout << "Minimal frequency due to block size " << fmin << " Hz" << endl;
    cout << "Maximal frequency due sampling      " << fmax << " Hz" << endl;
    //---------------------------------------------------------------
    // 3. Let's Fourier transform each of the blocks

    int channel=0;
    vector<complex<float> > amplitudesstore(nblocks*blocksize, 0.0);
    vector<complex<float> > e(blocksize*blocksize);


    //---------------------------------------------------------------
    // 4. Lets precalculate exp(i 2pi j*k / blocksize)

    for (int k=0; k < blocksize ; k++)
        for (int j=0 ; j < blocksize; j++)
        {
            int o=k*blocksize+j;
            complex<double> z(0.0, 2.0*M_PI*j*k/blocksize);
            e[o] = exp(z);
        }

    //---------------------------------------------------------------
    // 5. Lets calculate fourier transform the individual blocks

    for (int b = 0; b < nblocks; b++)        // Block
        for (int k=0; k < blocksize ; k++)     // k frequency
        {
            int offset=b*blocksize;
            complex<double> ak(0.0,0.0);   // amplitude of k'th frequncy

            //  ak = sum          f(j)            exp(i 2pi k j/blocksize)

            for (int j=0 ; j < blocksize; j++)
                ak += a.samples[channel][offset+j]*e[k*blocksize+j];

            // Store ak
            amplitudesstore[offset+k]=ak;
        }

    //---------------------------------------------------------------
    // 6. Waterfall spectrum

    ofstream fo(output1);
    // Weird order, but Matlab is column-row order. Thus we load high frequency to low frequency, then blocks.
    double Deltat = 1.0/rate;
    for (int k=blocksize/10; k >=0 ; k--)
    {
        for (int b = 0; b < nblocks; b++)
        {
            int offset=b*blocksize;
            double Re=amplitudesstore[offset+k].real();
            double Im=amplitudesstore[offset+k].imag();
            double I=Deltat*sqrt(Re*Re+Im*Im);
            fo << I << " ";
        }
        fo << endl;
    }
    fo.close();

    //---------------------------------------------------------------
    // 7. Raw and normalized spectrum for last block

    int block=nblocks-1;                 // Last block in file
    ofstream fo2(output2);               // Complex coeffients
    ofstream fo3(output3);               // Power spectrum.
    for (int k=0; k < blocksize ; k++)
    {
        int offset=block*blocksize;
        double Re=amplitudesstore[offset+k].real();
        double Im=amplitudesstore[offset+k].imag();
        double I=Deltat*sqrt(Re*Re+Im*Im);

        fo2 << k<< " " << amplitudesstore[offset+k].real() << " " << amplitudesstore[offset+k].imag() << endl;

        if (k<blocksize/2) fo3 << k*fmin             << " " << I << endl;
                      else fo3 << (k-blocksize)*fmin << " " << I << endl;
    }

    fo2.close();
    fo3.close();
}