/*
 * Copyright (C) 2013 Darth Affe <http://wyrez.org> and contributors
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 * 
 * This code is based on a tutorial by http://www.badlogicgames.com (as of 16.05.2013).
 * The whole tutorial can be found here:
 * http://www.badlogicgames.com/wordpress/?cat=18&paged=3
 */
package org.wyrez.audio;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import org.wyrez.audio.analysis.FFT;
import org.wyrez.audio.decoder.Decoder;
import org.wyrez.audio.util.Band;
import org.wyrez.audio.util.SampleBuffer;
import org.wyrez.audio.util.SampleHelper;

/**
 * Processor to create analysis-data of audio samples
 *
 * @author Darth Affe
 */
public class AudioProcessor {

    private Band band;
    private int bufferSize;
    private SampleBuffer sampleBuffer;
    private float[] spectrum;
    private float[] spectralFlux;
    private float[] threshold;
    private float[] prunedSpectralFlux;
    private float[] peaks;
    private float peakAverage;
    private float samplingRate;
    private float bpm;

    /**
     * Constructor, initializes an AudioProcessor for the given Decoder.
     *
     * @param decoder The decoder to process
     */
    public AudioProcessor(Decoder decoder) {
        this(SampleHelper.createSampleBuffer(decoder));
    }

    /**
     * Constructor, initializes an AudioProcessor for the given Decoder and
     * band.
     *
     * @param decoder The decoder to process
     * @param band The band to process
     */
    public AudioProcessor(Decoder decoder, Band band) {
        this(SampleHelper.createSampleBuffer(decoder), band);
    }

    public AudioProcessor(SampleBuffer sampleBuffer) {
        this(sampleBuffer, null);
    }

    public AudioProcessor(SampleBuffer sampleBuffer, Band band) {
        this.sampleBuffer = sampleBuffer;
        this.samplingRate = sampleBuffer.samplingRate;
        this.band = band;
    }

    /**
     * Removes cached objects to save some memory.
     */
    public void clean() {
        spectrum = null;
        spectralFlux = null;
        threshold = null;
        prunedSpectralFlux = null;
        peaks = null;
        bpm = 0f;
        bufferSize = 0;
    }

    /**
     * Calculates a full analysis of the samples data using default values:
     * BufferSize = 1024 ThresholdWindowSize = 10 ThresholdMultiplier = 1.5
     */
    public void calculate() {
        calculate(1024, 10, 1.5f);
    }

    /**
     * Calculates a full analysis of the samples data using a given threshold.
     *
     * @param thresholdWindowSize The window -size used to calculate the
     * threshold
     * @param thresholdMultiplier The multiplier used to calculate the threshold
     */
    public void calculate(int thresholdWindowSize, float thresholdMultiplier) {
        calculate(1024, thresholdWindowSize, thresholdMultiplier);
    }

    /**
     * Calculates a full analysis of the samples data using a given threshold
     * and a custom buffer size.
     *
     * @param bufferSize The size of the buffer
     * @param thresholdWindowSize The window -size used to calculate the
     * threshold
     * @param thresholdMultiplier The multiplier used to calculate the threshold
     */
    public void calculate(int bufferSize, int thresholdWindowSize, float thresholdMultiplier) {
        calculate(bufferSize, thresholdWindowSize, thresholdMultiplier, -1f, -1);
    }

    /**
     * Calculates a full analysis of the samples data using a given
     * threshold-multiplier and cuts fast peaks.
     *
     * @param thresholdWindowSize The window -size used to calculate the
     * threshold
     * @param minPeakDiff The minimum difference in samples between to peaks
     */
    public void calculate(int thresholdWindowSize, float thresholdMultiplier, int minPeakDiff) {
        calculate(bufferSize, thresholdWindowSize, thresholdMultiplier, -1f, -1);
    }

    /**
     * Calculates a full analysis of the samples data using a given
     * threshold-multiplier and cuts low peaks.
     *
     * @param thresholdWindowSize The window -size used to calculate the
     * threshold
     * @param percentageToKeep The minimum amplitude in percent of the avarage
     * amplitude
     */
    public void calculate(int thresholdWindowSize, float thresholdMultiplier, float percentageToKeep) {
        calculate(bufferSize, thresholdWindowSize, thresholdMultiplier, percentageToKeep, -1);
    }

    /**
     * Calculates a full analysis of the samples data.
     *
     * @param bufferSize The size of the buffer
     * @param thresholdWindowSize The window -size used to calculate the
     * threshold
     * @param thresholdMultiplier The multiplier used to calculate the threshold
     * @param percentageToKeep The minimum amplitude in percent of the avarage
     * amplitude
     * @param minPeakDiff The minimum difference in samples between to peaks
     */
    public void calculate(int bufferSize, int thresholdWindowSize, float thresholdMultiplier, float percentageToKeep, int minPeakDiff) {
        if ((bufferSize & (bufferSize - 1)) != 0) {
            throw new IllegalArgumentException("FFT: bufferSize must be a power of two.");
        }
        this.bufferSize = bufferSize;
        calculateSpectralFlux();
        calculateThreshold(thresholdWindowSize, thresholdMultiplier);
        calculatePrunedSpectralFlux();
        calculatePeaks();
        if (percentageToKeep > 0f) {
            cutLowPeaks(percentageToKeep);
        }
        if (minPeakDiff > 0) {
            cutFastPeaks(minPeakDiff);
        }
    }

    private void calculateSpectralFlux() {
        FFT fft = new FFT(bufferSize, samplingRate);
        fft.window(FFT.HAMMING);
        float[] buffer;
        float[] spectrumBuffer = new float[fft.getBandSize(band)];
        float[] lastSpectrumBuffer = new float[spectrumBuffer.length];
        List<Float> spectrum = new ArrayList<Float>();
        List<Float> spectralFlux = new ArrayList<Float>();

        for (int i = 0; i < sampleBuffer.samples.length; i += bufferSize) {
            buffer = SampleHelper.splitSampleArray(sampleBuffer.samples, i, bufferSize);
            fft.forward(buffer);
            System.arraycopy(spectrumBuffer, 0, lastSpectrumBuffer, 0, spectrumBuffer.length);
            System.arraycopy(fft.getSpectrum(band), 0, spectrumBuffer, 0, spectrumBuffer.length);

            float spectrumAvg = 0;
            float flux = 0;
            for (int j = 0; j < spectrumBuffer.length; j++) {
                spectrumAvg += spectrumBuffer[j];
                float value = (spectrumBuffer[j] - lastSpectrumBuffer[j]);
                flux += value < 0 ? 0 : value;
            }
            spectrumAvg /= spectrumBuffer.length;
            spectrum.add(spectrumAvg);
            spectralFlux.add(flux);
        }
        this.spectralFlux = SampleHelper.convertToFloatArray(spectralFlux);
        this.spectrum = SampleHelper.convertToFloatArray(spectrum);
    }

    private void calculateThreshold(int thresholdWindowSize, float multiplier) {
        threshold = new float[spectralFlux.length];
        for (int i = 0; i < spectralFlux.length; i++) {
            int start = Math.max(0, i - thresholdWindowSize);
            int end = Math.min(spectralFlux.length - 1, i + thresholdWindowSize);
            float mean = 0;
            for (int j = start; j <= end; j++) {
                mean += spectralFlux[j];
            }
            mean /= (end - start);
            threshold[i] = mean * multiplier;
        }
    }

    private void calculatePrunedSpectralFlux() {
        prunedSpectralFlux = new float[threshold.length];
        for (int i = 0; i < threshold.length; i++) {
            prunedSpectralFlux[i] = (threshold[i] <= spectralFlux[i]) ? spectralFlux[i] - threshold[i] : 0f;
        }
    }

    private void calculatePeaks() {
        peaks = new float[prunedSpectralFlux.length - 1];
        for (int i = 0; i < prunedSpectralFlux.length - 1; i++) {
            peaks[i] = (prunedSpectralFlux[i] > prunedSpectralFlux[i + 1]) ? prunedSpectralFlux[i] : 0f;
        }
        calculatePeakData();
    }

    private void calculatePeakData() {
        calculatePeakAverage();
        calculateBPM();
    }

    private void calculatePeakAverage() {
        int counter = 0;
        float sum = 0f;
        for (int i = 0; i < peaks.length; i++) {
            if (peaks[i] > 0f) {
                sum += peaks[i];
                counter++;
            }
        }
        peakAverage = sum / counter;
    }

    private void calculateBPM() {
        HashMap<Integer, Integer> peakDiff = new HashMap<Integer, Integer>();
        int counter = 0;
        boolean first = true;
        for (int i = 0; i < peaks.length; i++) {
            if (peaks[i] > 0f) {
                if (!first) {
                    Integer value = peakDiff.get(counter);
                    if (value == null) {
                        value = 1;
                    } else {
                        value += 1;
                    }
                    peakDiff.put(counter, value);
                } else {
                    first = false;
                }
                counter = 1;
            } else {
                counter++;
            }
        }

        int most = -1;
        int highestCount = -1;
        for (Integer i : peakDiff.keySet()) {
            if (peakDiff.get(i) > highestCount) {
                highestCount = peakDiff.get(i);
                most = i;
            }
        }
        float timePerPeak = (1 / samplingRate) * bufferSize;
        bpm = 1f / (timePerPeak * most) * 60f;
    }

    /**
     * Cuts the last calculated peaks. Removes every peak with a amplitude lower
     * than peakAverage*percentageToKeep.
     */
    public void cutLowPeaks(float percentageToKeep) {
        if (peaks == null) {
            return;
        }
        float valueToKeep = peakAverage * percentageToKeep;
        for (int i = 0; i < peaks.length; i++) {
            if (peaks[i] < valueToKeep) {
                peaks[i] = 0f;
            }
        }
        calculatePeakData();
    }

    /**
     * Cuts the last calculated peaks. Removes peaks witch a sample difference
     * lower than minPeakDiff. (This algorithm does NOT analyse the whole data
     * therfore the result may be a bit inaccurately)
     */
    public void cutFastPeaks(int minPeakDiff) {
        if (peaks == null) {
            return;
        }
        int counter = 1;
        boolean first = true;
        for (int i = 0; i < peaks.length; i++) {
            if (peaks[i] > 0f) {
                if (!first) {
                    if (counter < minPeakDiff) {
                        peaks[i] = 0f;
                    } else {
                        counter = 1;
                    }
                } else {
                    first = false;
                }
            } else {
                counter++;
            }
        }
        calculatePeakData();
    }

    /**
     * Returns the samples which are analysed.
     *
     * @returns A SampleBuffer containing the samples
     */
    public SampleBuffer getSampleBuffer() {
        return sampleBuffer;
    }

    /**
     * Returns the last calculated fft-spectrum.
     *
     * @returns A float array containing the spectrum data
     */
    public float[] getSpectrum() {
        return spectrum;
    }

    /**
     * Returns the last calculated spectral flux.
     *
     * @returns A float array containing the spectral flux data
     */
    public float[] getSpectralFlux() {
        return spectralFlux;
    }

    /**
     * Returns the last calculated threshold.
     *
     * @returns A float array containing the threshold data
     */
    public float[] getThreshold() {
        return threshold;
    }

    /**
     * Returns the last calculated pruned spectral flux.
     *
     * @returns A float array containing the pruned spectral flux data
     */
    public float[] getPrunedSpectralFlux() {
        return prunedSpectralFlux;
    }

    /**
     * Returns the last calculated peaks.
     *
     * @returns A float array containing the peaks
     */
    public float[] getPeaks() {
        return peaks;
    }

    /**
     * Returns the last calculated avarage peak amplitude.
     *
     * @returns The avarage peak amplitude as float
     */
    public float getPeakAverage() {
        return peakAverage;
    }

    /**
     * Returns the sampling rate of the used samples.
     *
     * @return The sampling rate as float
     */
    public float getSamplingRate() {
        return samplingRate;
    }

    /**
     * Returns the last calculated bpm.
     *
     * @returns The bpm as float
     */
    public float getBpm() {
        return bpm;
    }

    /**
     * Returns the buffer size used for the last calculation.
     *
     * @returns The buffer size as int
     */
    public int getBufferSize() {
        return bufferSize;
    }

    /**
     * Returns the band used for the last calculation.
     *
     * @returns The band
     */
    public Band getBand() {
        return band;
    }

    /**
     * Sets the band.
     *
     * @param band Band to use
     */
    public void setBand(Band band) {
        this.band = band;
    }
}
