using System;
using System.Collections.Generic;
using System.Collections.Concurrent;
using System.Linq;
using System.Text;
using System.Numerics;
using System.Threading.Tasks;
using System.Threading;
using System.Diagnostics;

namespace WpfApplication1
{
    public class timefreq
    {
        public float[][] timeFreqData;
        public int wSamp;
        public Complex[] twiddles;

        public timefreq(float[] x, int windowSamp, bool parallelFlag)
        {
            int ii;
            double pi = 3.14159265;
            Complex i = Complex.ImaginaryOne;
            this.wSamp = windowSamp;
            twiddles = new Complex[wSamp];

            //Calculate twiddle factors, calculated once based on N aka wSamp
            //TDoesn't take too much time, lots of pow and float math
            //Should be easily parallelised, but unnecessary
            for (ii = 0; ii < wSamp; ii++)
            {
                double a = 2 * pi * ii / (double)wSamp;
                twiddles[ii] = Complex.Pow(Complex.Exp(-i), (float)a);
            }

            timeFreqData = new float[wSamp/2][];

            //Quantizes length and samples
            int nearest = (int)Math.Ceiling((double)x.Length / (double)wSamp);
            nearest = nearest * wSamp;

            Complex[] compX = new Complex[nearest];
            for (int kk = 0; kk < nearest; kk++)
            {
                if (kk < x.Length)
                {
                    compX[kk] = x[kk];
                }
                else
                {
                    compX[kk] = Complex.Zero;
                }
            }

            //Wait this squares the data off
            //When there are double the samples there are half the columns
            //Sum total always equals nearest (2383872)
            int cols = 2 * nearest /wSamp;

            for (int jj = 0; jj < wSamp / 2; jj++)
            {
                timeFreqData[jj] = new float[cols];
            }

            //Also need to check to apply parallel to everything before this
            if (parallelFlag)
            {
                timeFreqData = stftPar(compX, wSamp);
            }
            else
            {
                timeFreqData = stftSeq(compX, wSamp);
            }


        }

        //Short-time Fourier Transform
        //Performs FFT on short, overlapping time segments
        float[][] stftSeq(Complex[] x, int wSamp)
        {
            int ii = 0;
            int jj = 0;
            int kk = 0;
            int ll = 0;
            int N = x.Length;

            //Stores the absolute maximum intensity of any frequency
            //Presumably to determine a colour scale
            float fftMax = 0;

            //Array of arrays to store the maximum frequency intensities at each time slice
            float[][] Y = new float[wSamp / 2][];
            for (ll = 0; ll < wSamp / 2; ll++)
            {
                //Each row is a frequency
                Y[ll] = new float[2 * (int)Math.Floor((double)N / (double)wSamp)];
            }

            Complex[] temp = new Complex[wSamp];
            Complex[] tempFFT = new Complex[wSamp];

            //Investigate this loop
            //This is where most time is spent
            //Takes a slice of the wave file "temp" (array of doubles)
            //Runs FFT on it and returns the intensities of each frequency "tempFFT" (array of complexes)
            //Stores the absolute intensity values of each frequency [kk] into slice [ii] of "Y" (array of float arrays)

            for (ii = 0; ii < 2 * Math.Floor((double)N / (double)wSamp) - 1; ii++)
            //Parallel.For(0, Convert.ToInt64(2 * Math.Floor((double)N / (double)wSamp) - 1),
            //ii =>
            {
                //Take a slice of the wave file
                //Overlaps by 50%?
                //Misses the very start?
                for (jj = 0; jj < wSamp; jj++)
                {
                    temp[jj] = x[ii * (wSamp / 2) + jj];
                }

                //Recursive FFT algorithm, takes up all the time
                tempFFT = fft2(temp);

                //Stores the values and checks to see if any are higher than the max already encountered
                for (kk = 0; kk < wSamp / 2; kk++)
                {
                    Y[kk][ii] = (float)Complex.Abs(tempFFT[kk]);

                    //Needs to be broken out in the final version, fftMax is a shared variable
                    //Possibly keep locally within the thread and find max later
                    if (Y[kk][ii] > fftMax)
                    {
                        fftMax = Y[kk][ii];
                    }
                }

            }

            for (ii = 0; ii < 2 * Math.Floor((double)N / (double)wSamp) - 1; ii++)
            {
                for (kk = 0; kk < wSamp / 2; kk++)
                {
                    Y[kk][ii] /= fftMax;
                }
            }

            return Y;
        }

        //Short-time Fourier Transform
        //Performs FFT on short, overlapping time segments
        float[][] stftPar(Complex[] x, int wSamp)
        {
            int N = x.Length;
            int timeSlices = (2 * (int)Math.Floor((double)N / (double)wSamp));
            int frequencies = wSamp / 2;

            //Stores the absolute maximum intensity of any frequency
            //Presumably to determine a colour scale
            float fftMax = 0;

            //Array of arrays to store the maximum frequency intensities at each time slice
            float[][] Y = new float[timeSlices][];
            for (int ii = 0; ii < Y.Length; ii++)
            {
                //Each row is the frequency values at that timeslice
                Y[ii] = new float[frequencies];
            }

            //Investigate this loop
            //This is where most time is spent
            //Takes a slice of the wave file "temp" (array of doubles)
            //Runs FFT on it and returns the intensities of each frequency "tempFFT" (array of complexes)
            //Stores the absolute intensity values of each frequency [kk] into slice [ii] of "Y" (array of float arrays)

            Object lockObject = new Object();
            ParallelOptions pOptions = new ParallelOptions();
            pOptions.MaxDegreeOfParallelism = -1;

            Parallel.For(0, timeSlices - 1, pOptions, () => new float(),
            (index, loop, localMax) =>
            {
                long ii = index;

                Complex[] timeSlice = new Complex[wSamp];
                Complex[] tempFFT = new Complex[wSamp];

                var fftMag = Y[ii];

                //Take a slice of the wave file
                //Overlaps by 50%?
                //Misses the very start?
                for (int jj = 0; jj < wSamp; jj++)
                {
                    timeSlice[jj] = x[ii * (frequencies) + jj];
                }

                //Recursive FFT algorithm, takes up all the time
                tempFFT = fft2(timeSlice);

                //Stores the values and checks to see if any are higher than the max already encountered
                for (int kk = 0; kk < frequencies; kk++)
                {
                    fftMag[kk] = (float)Complex.Abs(tempFFT[kk]);

                    //Needs to be broken out in the final version, fftMax is a shared variable
                    //Possibly keep locally within the thread and find max later
                    if (fftMag[kk] > localMax)
                    {
                        localMax = fftMag[kk];
                    }

                }
                Buffer.BlockCopy(fftMag, 0, Y[ii], 0, fftMag.Length);
                return localMax;
            },
                (lMax) =>
                {
                    lock (lockObject)
                    {
                        if (lMax > fftMax)
                        {
                            fftMax = lMax;
                        }
                    }
                }
            );

            var Z = new float[Y[0].Length][];
            for (int ii = 0; ii < Z.Length; ii++) Z[ii] = new float[Y.Length];

            //Normalise all to fftMax
            //for (int ii = 0; ii < 2 * Math.Floor((double)N / (double)wSamp) - 1; ii++)
            Parallel.For(0, Y.Length - 1,
                index =>
            {
                long ii = index;
                for (int kk = 0; kk < Y[0].Length; kk++)
                {
                    Z[kk][ii] = Y[ii][kk] / fftMax;
                }
            });
            return Z;
        }


        //Could increase the radix from 2 to explore caching optimisations
        Complex[] fft(Complex[] x)
        {
            int ii = 0;
            int kk = 0;
            int N = x.Length;

            Complex[] Y = new Complex[N];

            // NEED TO MEMSET TO ZERO?

            if (N == 1)
            {
                Y[0] = x[0]; return Y;
            }

            //So much memory access, could reduce by simply calling fft with O(x) or E(X) inline
            //Maybe, need to profile though
            Complex[] E = new Complex[N/2];
            Complex[] O = new Complex[N/2];
            Complex[] even = new Complex[N/2];
            Complex[] odd = new Complex[N/2];

            //Splits N and x into even and odd sets
            //Could be optimised by replacing % with bit shift hack
            //Obviously parallel
            //Or could just eliminated completely by calculating in the call to fft() later
            for (ii = 0; ii < N; ii++)
            {

                if (ii % 2 == 0)
                {
                    even[ii / 2] = x[ii];
                }
                if (ii % 2 == 1)
                {
                    odd[(ii - 1) / 2] = x[ii];
                }
            }

            //RECURSION
            //Algorithm splits the sample space in half and abuses the properties of
            //the DFT to use intermediate results in later calculations (called a butterfly)
            //This allows recursion up to Log2(N) (e.g. 11 for 2048) along with the final loop for N = 1
            //Do this for N samples and you have a depth first graph search N wide and Log2(N+1) deep
            //Results in O(N log N) complexity
            //This could possibly be done in distinct threads, absolutely hilariously
            //Apparently this isn't hilarious and is actually possible
            E = fft(even);
            O = fft(odd);

            //Should be able to be parallelised
            //Might not be worth the overhead depending on N
            //Could also compute O*twiddles in a single loop for SSE reasons (lots of float multiplication)
            for (kk = 0; kk < N; kk++)
            {
                //Might be able to fix up % with bit shifting
                Y[kk] = E[(kk % (N / 2))] + O[(kk % (N / 2))] * twiddles[kk * wSamp / N];
            }


           return Y;
        }


        Complex[] fft2(Complex[] x)
        {
            int bits = (int)Math.Log(x.Length, 2);
            for(int i = 1; i < x.Length; i++)
            {
                int swapPosition = bitReversal(i, bits);
                if(swapPosition > i)
                {
                    var temp = x[i];
                    x[i] = x[swapPosition];
                    x[i] = temp;
                }
            }

            for(int n = 2; n <= x.Length; n = n << 1)
            {
                for(int i = 0; i < x.Length; i += n)
                {
                    for(int j = 0; j < (n/2); j++)
                    {
                        int evenIndex = i + j;
                        int oddIndex = i + j + (n/2);
                        var even = x[evenIndex];
                        var odd = x[oddIndex];

                        Complex exp = twiddles[j * wSamp / n] * odd;

                        x[evenIndex] = even + exp;
                        x[oddIndex] = even - exp;
                    }
                }
            }

            return x;
        }

        static int bitReversal(int n, int bits)
        {
            int reversed = n;
            int count = bits - 1;

            n = n >> 1;
            while(n > 0)
            {
                reversed = (reversed << 1) | (n & 1);
                count--;
                n = n >> 1;
            }

            return (reversed << count) & (1 << bits - 1);
        }
    }
}