#include "FourierTransform_pffft.h"

#include "pffft.h"

namespace staffpad::audio {
FourierTransform::FourierTransform(int32_t newBlockSize)
    : _blockSize{newBlockSize}
{
    _pffft_scratch = (float*)pffft_aligned_malloc(_blockSize * sizeof(float));
    realFftSpec = pffft_new_setup(_blockSize, PFFFT_REAL);
}

FourierTransform::~FourierTransform()
{
    if (_pffft_scratch) {
        pffft_aligned_free(_pffft_scratch);
        _pffft_scratch = nullptr;
    }
    if (realFftSpec) {
        pffft_destroy_setup(realFftSpec);
        realFftSpec = nullptr;
    }
}

void FourierTransform::forwardReal(const SamplesReal& t, SamplesComplex& c)
{
    assert(t.getNumSamples() == _blockSize);

    for (auto ch = 0; ch < t.getNumChannels(); ++ch) {
        auto* spec = c.getPtr(ch); // interleaved complex numbers, size _blockSize + 2
        auto* cpx_flt = (float*)spec;
        pffft_transform_ordered(realFftSpec, t.getPtr(ch), cpx_flt, _pffft_scratch, PFFFT_FORWARD);
        // pffft combines dc and nyq values into the first complex value,
        // adjust to CCS format.
        auto dc = cpx_flt[0];
        auto nyq = cpx_flt[1];
        spec[0] = { dc, 0.f };
        spec[c.getNumSamples() - 1] = { nyq, 0.f };
    }
}

void FourierTransform::inverseReal(const SamplesComplex& c, SamplesReal& t)
{
    assert(c.getNumSamples() == _blockSize / 2 + 1);

    for (auto ch = 0; ch < c.getNumChannels(); ++ch) {
        auto* spec = c.getPtr(ch);
        // Use t to convert in-place from CCS to pffft format
        t.assignSamples(ch, (float*)spec);
        auto* ts = t.getPtr(ch);
        ts[0] = spec[0].real();
        ts[1] = spec[c.getNumSamples() - 1].real();
        pffft_transform_ordered(realFftSpec, ts, ts, _pffft_scratch, PFFFT_BACKWARD);
    }
}
} // namespace  staffpad::audio