#pragma once
#include <JuceHeader.h>
#include <vector>
#include <cmath>

// ============================== Design =======================================
// - Bank with F frames, each frame is a single-cycle table of N samples.
// - For each frame, we create L mip-levels: level 0 = full bandwidth,
//   level l halves the permitted harmonics (spectral truncation).
// - Runtime chooses level from note frequency and sampleRate, then morphs
//   between adjacent frames and crossfades between the two nearest levels.
// - Table read uses linear interpolation (cheap and good enough with N>=2048).

namespace WT
{
    // Utility: complex array wrapper for JUCE FFT (interleaved real/imag floats)
    struct ComplexBuf
    {
        std::vector<float> data; // size = 2 * N
        explicit ComplexBuf(size_t N = 0) { resize(N); }
        void resize(size_t N) { data.assign(2 * N, 0.0f); }
        juce::dsp::Complex<float>* asComplex() { return reinterpret_cast<juce::dsp::Complex<float>*>(data.data()); }
    };

    // =======================================================================
    // WavetableBank: holds raw frames + mipmapped versions
    // =======================================================================
    class Bank
    {
    public:
        // N = table length (must be power-of-two for FFT), frames = number of morph frames
        // mipLevels = how many spectral levels (>=1). 5 ~ 6 is plenty for synth use.
        Bank(size_t N = 2048, int frames = 16, int mipLevels = 6)
            : tableSize(N), numFrames(frames), numLevels(mipLevels),
              fft((int)std::log2((double)N))
        {
            jassert(juce::isPowerOfTwo((int)N));
            tables.resize((size_t)numLevels);
            for (int l = 0; l < numLevels; ++l)
                tables[(size_t)l].resize((size_t)numFrames, std::vector<float>(tableSize, 0.0f));
        }

        size_t getSize()    const { return tableSize; }
        int    getFrames()  const { return numFrames; }
        int    getLevels()  const { return numLevels; }

        // Provide raw “design” frames (time-domain single-cycle) then call buildMipmaps().
        // framesRaw.size() must equal numFrames, each frame length must equal tableSize.
        void setRawFrames(const std::vector<std::vector<float>>& framesRaw)
        {
            jassert((int)framesRaw.size() == numFrames);
            for (const auto& f : framesRaw) jassert(f.size() == tableSize);
            raw = framesRaw;
        }

        // Convenience: generate 16-frame bank morphing Sine -> Saw -> Square -> Triangle
        void generateDefaultMorph()
        {
            std::vector<std::vector<float>> frames;
            frames.resize((size_t)numFrames, std::vector<float>(tableSize, 0.0f));

            auto fill = [&](int idx, auto func)
            {
                auto& t = frames[(size_t)idx];
                for (size_t n = 0; n < tableSize; ++n)
                {
                    const float ph = (float) (juce::MathConstants<double>::twoPi * (double)n / (double)tableSize);
                    t[n] = func(ph);
                }
                normalise(t);
            };

            // helper waves
            auto sine = [](float ph) { return std::sin(ph); };
            auto saw  = [](float ph) { return (float)(2.0 * (ph / juce::MathConstants<float>::twoPi) - 1.0); };
            auto sq   = [](float ph) { return ph < juce::MathConstants<float>::pi ? 1.0f : -1.0f; };
            auto tri  = [](float ph) {
                float v = (float)(2.0 * std::abs(2.0 * (ph / juce::MathConstants<float>::twoPi) - 1.0) - 1.0);
                return v;
            };

            // 0..5: sine->saw, 6..10: saw->square, 11..15: square->triangle
            const int F = numFrames;
            for (int i = 0; i < F; ++i)
            {
                const float t = (float) i / (float) juce::jmax(1, F - 1);
                std::function<float(float)> a, b;
                float mix = 0.0f;

                if (i <= 5)      { a = sine; b = saw;    mix = (float)i / 5.0f; }
                else if (i <=10) { a = saw;  b = sq;     mix = (float)(i - 6) / 4.0f; }
                else             { a = sq;   b = tri;    mix = (float)(i - 11) / 4.0f; }

                fill(i, [=](float ph){ return (1.0f - mix) * a(ph) + mix * b(ph); });
            }

            setRawFrames(frames);
        }

        // Build mip-levels by FFT → spectral truncation → IFFT
        void buildMipmaps()
        {
            jassert(!raw.empty());
            ComplexBuf freq(tableSize);
            ComplexBuf time(tableSize);

            for (int f = 0; f < numFrames; ++f)
            {
                // Forward FFT of raw frame
                std::fill(freq.data.begin(), freq.data.end(), 0.0f);
                for (size_t n = 0; n < tableSize; ++n)
                {
                    time.data[2 * n + 0] = raw[(size_t)f][n];
                    time.data[2 * n + 1] = 0.0f;
                }
                fft.performRealOnlyForwardTransform(time.data.data());
                // After JUCE real FFT, bins are laid out as: Re[0], Re[N/2], Re[1], Im[1], Re[2], Im[2], ...
                // We'll reconstruct complex bins for easy masking.

                // Helper to zero all harmonics above kMax (inclusive index in [0..N/2])
                auto maskAndIFFT = [&](int level, int kMax)
                {
                    // Copy time.data into working complex bins
                    auto* bins = freq.asComplex();
                    // DC & Nyquist are purely real in real-FFT
                    bins[0].real (time.data[0]);
                    bins[0].imag (0.0f);
                    bins[tableSize/2].real (time.data[1]);
                    bins[tableSize/2].imag (0.0f);

                    // Rebuild the rest (Re[k], Im[k]) packed starting at index 2
                    for (size_t k = 1; k < tableSize/2; ++k)
                    {
                        bins[k].real (time.data[2 * k + 0]);
                        bins[k].imag (time.data[2 * k + 1]);
                    }

                    // Mask
                    for (size_t k = (size_t)kMax + 1; k < tableSize/2; ++k)
                        bins[k] = { 0.0f, 0.0f };

                    // Pack back into real-FFT layout for inverse
                    time.data[0] = bins[0].real();  // DC
                    time.data[1] = bins[tableSize/2].real(); // Nyquist
                    for (size_t k = 1; k < tableSize/2; ++k)
                    {
                        time.data[2 * k + 0] = bins[k].real();
                        time.data[2 * k + 1] = bins[k].imag();
                    }

                    // IFFT
                    fft.performRealOnlyInverseTransform(time.data.data());

                    // Copy, normalise a little (scale JUCE inverse divides by N already)
                    auto& dst = tables[(size_t)level][(size_t)f];
                    for (size_t n = 0; n < tableSize; ++n)
                        dst[n] = time.data[2 * n + 0];

                    normalise(dst);
                };

                // Level 0 → all harmonics available up to N/2 - 1
                for (int l = 0; l < numLevels; ++l)
                {
                    const int maxH = (int)((tableSize / 2) >> l);       // halve per level
                    const int kMax = juce::jmax(1, juce::jmin(maxH, (int)tableSize/2 - 1));
                    maskAndIFFT(l, kMax);
                }
            }
        }

        // sample at (frame, level, phase in [0,1))
        inline float lookup (float frameIdx, int level, float phase) const noexcept
        {
            const int f0 = juce::jlimit(0, numFrames - 1, (int)std::floor(frameIdx));
            const int f1 = juce::jlimit(0, numFrames - 1, f0 + 1);
            const float t = juce::jlimit(0.0f, 1.0f, frameIdx - (float)f0);

            const auto& T0 = tables[(size_t)level][(size_t)f0];
            const auto& T1 = tables[(size_t)level][(size_t)f1];

            const float pos = phase * (float)tableSize;
            const int   i0  = (int) std::floor(pos) & (int)(tableSize - 1);
            const int   i1  = (i0 + 1) & (int)(tableSize - 1);
            const float a   = pos - (float) std::floor(pos);

            const float s0 = juce::jmap(a, T0[(size_t)i0], T0[(size_t)i1]);
            const float s1 = juce::jmap(a, T1[(size_t)i0], T1[(size_t)i1]);
            return juce::jmap(t, s0, s1);
        }

        // choose mip-level for given frequency (Hz) & sampleRate
        inline int chooseLevel (float freq, double sampleRate) const noexcept
        {
            // permitted harmonics at this pitch:
            const float maxH = (float) (0.5 * sampleRate / juce::jmax(1.0f, freq));
            // level so that harmonic budget of level >= maxH, i.e. l = ceil(log2((N/2)/maxH))
            const float base = (float)(tableSize * 0.5);
            const float ratio = base / juce::jmax(1.0f, maxH);
            int l = (int) std::ceil (std::log2 (ratio));
            return juce::jlimit (0, numLevels - 1, l);
        }

        static void normalise (std::vector<float>& t)
        {
            float mx = 0.0f;
            for (float v : t) mx = juce::jmax(mx, std::abs(v));
            if (mx < 1.0e-6f) return;
            for (float& v : t) v /= mx;
        }

    private:
        size_t tableSize;
        int    numFrames;
        int    numLevels;

        juce::dsp::FFT fft;
        std::vector<std::vector<float>> raw;
        // [level][frame][sample]
        std::vector<std::vector<std::vector<float>>> tables;
    };

    // =======================================================================
    // Wavetable Oscillator
    // =======================================================================
    class Osc
    {
    public:
        void prepare (double sr) { sampleRate = sr; }
        void setBank (std::shared_ptr<Bank> b) { bank = std::move(b); }
        void setFrequency (float f) { freq = juce::jmax(0.0f, f); phaseInc = freq / (float)sampleRate; }
        void setMorph (float m) { morph = m; }      // 0..frames-1 (continuous)
        void resetPhase (float p = 0.0f) { phase = juce::jlimit(0.0f, 1.0f, p); }

        float process()
        {
            if (!bank) return 0.0f;

            const int l0 = bank->chooseLevel(freq, sampleRate);
            const int l1 = juce::jmin(l0 + 1, bank->getLevels() - 1);
            const float preferL0 = 1.0f - juce::jlimit(0.0f, 1.0f,
                                   (float)l0 - (float)bank->chooseLevel(freq * 0.99f, sampleRate));

            const float s0 = bank->lookup(morph, l0, phase);
            const float s1 = bank->lookup(morph, l1, phase);
            const float out = juce::jmap(preferL0, s1, s0); // simple crossfade

            phase += phaseInc;
            while (phase >= 1.0f) phase -= 1.0f;
            return out;
        }

    private:
        std::shared_ptr<Bank> bank;
        double sampleRate { 44100.0 };
        float  freq { 0.0f };
        float  morph { 0.0f };     // 0..frames-1
        float  phase { 0.0f };
        float  phaseInc { 0.0f };
    };
} // namespace WT