/*
 * Copyright (C) 2020 Linux Studio Plugins Project <https://lsp-plug.in/>
 *           (C) 2020 Vladimir Sadovnikov <sadko4u@gmail.com>
 *
 * This file is part of lsp-plugins
 * Created on: 14 авг. 2016 г.
 *
 * lsp-plugins is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * any later version.
 *
 * lsp-plugins 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 Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with lsp-plugins. If not, see <https://www.gnu.org/licenses/>.
 */

#include <lsp-plug.in/dsp-units/util/Analyzer.h>
#include <lsp-plug.in/dsp-units/misc/envelope.h>
#include <lsp-plug.in/dsp-units/misc/windows.h>
#include <lsp-plug.in/dsp-units/units.h>
#include <lsp-plug.in/stdlib/math.h>
#include <lsp-plug.in/common/alloc.h>
#include <lsp-plug.in/dsp/dsp.h>

namespace lsp
{
    namespace dspu
    {
        Analyzer::Analyzer()
        {
            construct();
        }

        Analyzer::~Analyzer()
        {
            destroy();
        }

        void Analyzer::construct()
        {
            nChannels       = 0;
            nMaxRank        = 0;
            nRank           = 0;
            nSampleRate     = 0;
            nMaxSampleRate  = 0;
            nBufSize        = 0;
            nCounter        = 0;
            nPeriod         = 0;
            nStep           = 0;
            nHead           = 0;
            fReactivity     = 0.0f;
            fTau            = 1.0f;
            fRate           = 1.0f;
            fMinRate        = 1.0f;
            fShift          = 1.0f;
            nReconfigure    = 0;
            nEnvelope       = envelope::PINK_NOISE;
            nWindow         = windows::HANN;
            bActive         = true;

            vChannels       = NULL;
            vData           = NULL;
            vSigRe          = NULL;
            vFftReIm        = NULL;
            vWindow         = NULL;
            vEnvelope       = NULL;
        }

        void Analyzer::destroy()
        {
            if (vChannels != NULL)
            {
                delete [] vChannels;
                vChannels   = NULL;
            }

            free_aligned(vData);
        }

        bool Analyzer::init(size_t channels, size_t max_rank, size_t max_sr, float min_rate)
        {
            destroy();

            size_t fft_size         = 1 << max_rank;
            nBufSize                = align_size(fft_size + size_t(float(max_sr * 2) / min_rate) + DEFAULT_ALIGN, DEFAULT_ALIGN);
            size_t allocate         = 5 * fft_size +                // vSigRe, vFftReIm (re + im), vWindow, vEnvelope
                                      channels * nBufSize +         // c->vBuffer
                                      channels * fft_size +         // c->vAmp
                                      channels * fft_size;          // c->vData

            // Allocate data
            float *abuf         = alloc_aligned<float>(vData, allocate);
            if (abuf == NULL)
                return false;

            // Allocate channels
            channel_t *clist    = new channel_t[channels];
            if (clist == NULL)
            {
                delete [] abuf;
                return false;
            }

            nChannels           = channels;
            nMaxRank            = max_rank;
            nRank               = max_rank;
            nMaxSampleRate      = max_sr;
            fMinRate            = min_rate;

            // Clear buffers
            dsp::fill_zero(abuf, allocate);

            // Initialize buffers
            vSigRe              = abuf;
            abuf               += fft_size;
            vFftReIm            = abuf;
            abuf               += fft_size * 2;
            vWindow             = abuf;
            abuf               += fft_size;
            vEnvelope           = abuf;
            abuf               += fft_size;

            // Initialize channels
            vChannels           = clist;
            for (size_t i=0; i<channels; ++i)
            {
                channel_t *c        = &vChannels[i];

                // FFT buffers
                c->vBuffer          = abuf;
                abuf               += nBufSize;
                c->vAmp             = abuf;
                abuf               += fft_size;
                c->vData            = abuf;
                abuf               += fft_size;

                // Counters
                c->nDelay           = 0;
                c->bFreeze          = false;
                c->bActive          = true;
            }

            // Set reconfiguration flags
            nReconfigure        = R_ALL;

            return true;
        }

        void Analyzer::set_sample_rate(size_t sr)
        {
            sr              = lsp_min(sr, nMaxSampleRate);
            if (nSampleRate == sr)
                return;

            nSampleRate     = sr;
            nReconfigure   |= R_ALL;
        }

        void Analyzer::set_rate(float rate)
        {
            rate        = lsp_max(fMinRate, rate);
            if (fRate == rate)
                return;

            fRate           = rate;
            nReconfigure   |= R_COUNTERS;
        }

        void Analyzer::set_window(size_t window)
        {
            if (nWindow == window)
                return;

            nWindow         = window;
            nReconfigure   |= R_WINDOW;
        }

        void Analyzer::set_envelope(size_t envelope)
        {
            if (nEnvelope == envelope)
                return;

            nEnvelope       = envelope;
            nReconfigure   |= R_ENVELOPE;
        }

        void Analyzer::set_shift(float shift)
        {
            if (fShift == shift)
                return;

            fShift          = shift;
            nReconfigure   |= R_ENVELOPE;
        }

        void Analyzer::set_reactivity(float reactivity)
        {
            if (fReactivity == reactivity)
                return;

            fReactivity     = reactivity;
            nReconfigure   |= R_TAU;
        }

        bool Analyzer::set_rank(size_t rank)
        {
            if ((rank < 2) || (rank > nMaxRank))
                return false;
            else if (nRank == rank)
                return true;
            nRank           = rank;
            nReconfigure   |= R_ALL;
            return true;
        }

        bool Analyzer::freeze_channel(size_t channel, bool freeze)
        {
            if (channel >= nChannels)
                return false;
            vChannels[channel].bFreeze      = freeze;
            return true;
        }

        bool Analyzer::enable_channel(size_t channel, bool enable)
        {
            if (channel >= nChannels)
                return false;
            if (vChannels[channel].bActive == enable)
                return false;

            vChannels[channel].bActive      = enable;
            nReconfigure   |= R_COUNTERS;
            return true;
        }

        void Analyzer::reconfigure()
        {
            if (!nReconfigure)
                return;

            size_t fft_size     = 1 << nRank;
            size_t fft_period   = float(nSampleRate) / fRate;
            nStep               = fft_period / nChannels;
            nPeriod             = nStep * nChannels;

            // Update envelope
            if (nReconfigure & R_ENVELOPE)
            {
                envelope::reverse_noise(vEnvelope, fft_size, envelope::envelope_t(nEnvelope));
                dsp::mul_k2(vEnvelope, fShift / fft_size, fft_size);
            }

            // Clear analysis
            if (nReconfigure & R_ANALYSIS)
            {
                for (size_t i=0; i<nChannels; ++i)
                {
                    dsp::fill_zero(vChannels[i].vAmp, fft_size);
                    dsp::fill_zero(vChannels[i].vData, fft_size);
                }
            }
            // Update window
            if (nReconfigure & R_WINDOW)
                windows::window(vWindow, fft_size, windows::window_t(nWindow));
            // Update reactivity
            if (nReconfigure & R_TAU)
                fTau    = 1.0f - expf(logf(1.0f - M_SQRT1_2) / seconds_to_samples(float(nSampleRate) / float(nPeriod), fReactivity));
            // Update counters
            if (nReconfigure & R_COUNTERS)
            {
                for (size_t i=0; i<nChannels; ++i)
                    vChannels[i].nDelay     = i*nStep;
            }

            // Clear reconfiguration flag and update strobe signal
            nReconfigure    = 0;
        }

        void Analyzer::process(const float * const *in, size_t samples)
        {
            if (vChannels == NULL)
                return;

            // Auto-apply reconfiguration
            reconfigure();

            // Do main processing
            channel_t *c;
            ssize_t fft_size    = 1 << nRank;
            ssize_t fft_csize   = (fft_size >> 1) + 1;

            for (size_t offset = 0; offset < samples; )
            {
                // Determine the actual position
                size_t channel  = nCounter / nStep;
                size_t off      = nCounter % nStep;

                // Need to do FFT transform/sync?
                if (off == 0)
                {
                    if (nCounter == 0)
                    {
                        // Strobe trigger, copy buffers
                        for (size_t i=0; i<nChannels; ++i)
                            dsp::copy(vChannels[i].vData, vChannels[i].vAmp, fft_size);
                    }

                    // Regular channel, need to perform FFT if
                    c   = &vChannels[channel];

                    // Perform FFT only for active channels
                    if (!c->bFreeze)
                    {
                        if ((bActive) && (c->bActive))
                        {
                            // Get the time mark to start from
                            ssize_t doff    = nHead - (fft_size + c->nDelay);

    //                        lsp_trace("channel=%d, head=%d, delay=%d, offset=%d, buf_size=%d, samples=%d",
    //                                int(channel), int(nHead), int(c->nDelay), int(doff), int(nBufSize), int(samples));

                            if (doff < 0)
                                doff           += nBufSize;

                            // Prepare the real buffer
                            ssize_t count   = nBufSize - doff;
                            if (count < fft_size)
                            {
                                dsp::mul3(vSigRe, &c->vBuffer[doff], vWindow, count);
                                dsp::mul3(&vSigRe[count], c->vBuffer, &vWindow[count], fft_size - count);
                            }
                            else
                                dsp::mul3(vSigRe, &c->vBuffer[doff], vWindow, fft_size);

                            // Do Real->complex conversion and FFT
                            dsp::pcomplex_r2c(vFftReIm, vSigRe, fft_size);
                            dsp::packed_direct_fft(vFftReIm, vFftReIm, nRank);
                            // Get complex argument
                            dsp::pcomplex_mod(vFftReIm, vFftReIm, fft_csize);
                            // Mix with the previous value
                            dsp::mix2(c->vAmp, vFftReIm, 1.0 - fTau, fTau, fft_csize);
                        }
                        else
                            dsp::fill_zero(c->vAmp, fft_size);
                    } // c->bFreeze
                } // off == 0

                // How many samples to process?
                size_t to_process   = lsp_min(samples - offset, nStep - off);

                // Commit data to delay buffers for each channel
                size_t tail         = nBufSize - nHead;
                for (size_t i=0; i<nChannels; ++i)
                {
                    c                   = &vChannels[i];
                    const float *src    = (in != NULL) ? in[i] : NULL;

                    if (tail < to_process)
                    {
                        if (src != NULL)
                        {
                            dsp::copy(&c->vBuffer[nHead], &src[offset], tail);
                            dsp::copy(c->vBuffer, &src[offset + tail], to_process - tail);
                        }
                        else
                        {
                            dsp::fill_zero(&c->vBuffer[nHead], tail);
                            dsp::fill_zero(c->vBuffer, to_process - tail);
                        }
                    }
                    else
                    {
                        if (src != NULL)
                            dsp::copy(&c->vBuffer[nHead], &src[offset], to_process);
                        else
                            dsp::fill_zero(&c->vBuffer[nHead], to_process);
                    }
                }

                // Update positions
                offset     += to_process;
                nCounter   += to_process;
                if (nCounter >= nPeriod)
                    nCounter   -= nPeriod;
                nHead      += to_process;
                if (nHead >= nBufSize)
                    nHead      -= nBufSize;
            }
        }

        bool Analyzer::read_frequencies(float *frq, float start, float stop, size_t count, size_t flags)
        {
            if ((vChannels == NULL) || (count == 0))
                return false;
            else if (count == 1)
            {
                *frq            = start;
                return true;
            }

            // Analyze flags
            if (flags == FRQA_SCALE_LOGARITHMIC)
            {
                // Initialize list of frequencies
                float norm          = logf(stop/start) / (--count);

                for (size_t i=0; i<count; ++i)
                    frq[i]              = start * expf(i * norm);
            }
            else if (flags == FRQA_SCALE_LINEAR)
            {
                float norm          = (stop - start) / (--count);
                for (size_t i=0; i<count; ++i)
                    frq[i]              = start + i * norm;
            }
            else
                return false;

            frq[count]          = stop;
            return true;
        }

        bool Analyzer::get_spectrum(size_t channel, float *out, const uint32_t *idx, size_t count)
        {
            if ((vChannels == NULL) || (channel >= nChannels))
                return false;

            channel_t *c        = &vChannels[channel];
            for (size_t i=0; i<count; ++i)
            {
                size_t j            = idx[i];
                out[i]              = c->vData[j] * vEnvelope[j];
            }

            return true;
        }

        float Analyzer::get_level(size_t channel, const uint32_t idx)
        {
            if ((vChannels == NULL) || (channel >= nChannels))
                return 0.0f;

            return vChannels[channel].vData[idx] * vEnvelope[idx];
        }

        void Analyzer::get_frequencies(float *frq, uint32_t *idx, float start, float stop, size_t count)
        {
            size_t fft_size     = 1 << nRank;
            size_t fft_csize    = (fft_size >> 1) + 1;
            float scale         = float(fft_size) / float(nSampleRate);

            // Initialize list of frequencies
            float norm          = logf(stop/start) / (count - 1);

            for (size_t i=0; i<count; ++i)
            {
                float f         = start * expf(i * norm);
                size_t ix       = scale * f;
                if (ix > fft_csize)
                    ix                  = fft_csize;

                frq[i]          = f;
                idx[i]          = ix;
            }
        }

        void Analyzer::dump(IStateDumper *v) const
        {
            v->write("nChannels", nChannels);
            v->write("nMaxRank", nMaxRank);
            v->write("nRank", nRank);
            v->write("nSampleRate", nSampleRate);
            v->write("nMaxSampleRate", nMaxSampleRate);
            v->write("nBufSize", nBufSize);
            v->write("nCounter", nCounter);
            v->write("nPeriod", nPeriod);
            v->write("nStep", nStep);
            v->write("nHead", nHead);
            v->write("fReactivity", fReactivity);
            v->write("fTau", fTau);
            v->write("fRate", fRate);
            v->write("fMinRate", fMinRate);
            v->write("fShift", fShift);
            v->write("nReconfigure", nReconfigure);
            v->write("nEnvelope", nEnvelope);
            v->write("nWindow", nWindow);
            v->write("bActive", bActive);

            v->begin_array("vChannels", vChannels, nChannels);
            for (size_t i=0; i<nChannels; ++i)
            {
                const channel_t *c = &vChannels[i];
                v->begin_object(c, sizeof(channel_t));
                {
                    v->write("vBuffer", c->vBuffer);
                    v->write("vAmp", c->vAmp);
                    v->write("vData", c->vData);
                    v->write("nDelay", c->nDelay);
                    v->write("bFreeze", c->bFreeze);
                    v->write("bActive", c->bActive);
                }
                v->end_object();
            }
            v->end_array();

            v->write("vData", vData);
            v->write("vSigRe", vSigRe);
            v->write("vFftReIm", vFftReIm);
            v->write("vWindow", vWindow);
            v->write("vEnvelope", vEnvelope);
        }
    }
} /* namespace lsp */
