#include "SpectrumDisplay.h"

namespace
{
    constexpr float minDb = -66.0f;
    constexpr float topDb = -6.0f;
    constexpr float dbRange = topDb - minDb;

    const juce::Colour wetLineCol (0xDD9DF5E9);
    const juce::Colour wetFillCol (0xE63CE0C8);
    const juce::Colour wetPeakCol (0x883CE0C8);
    const juce::Colour dryLineCol (0xCCF0A03C);
    const juce::Colour dryFillCol (0x22F0A03C);
    const juce::Colour dryPeakCol (0x55E89B3A);

    juce::String formatFreq (float hz)
    {
        if (hz >= 1000.0f)
            return juce::String (juce::roundToInt (hz / 1000.0f)) + "k";
        return juce::String (juce::roundToInt (hz));
    }
}

SpectrumDisplay::SpectrumDisplay (HorizontAudioProcessor& p)
    : processor (p)
{
    window.resize ((size_t) fftSize);

    for (int i = 0; i < fftSize; ++i)
        window[(size_t) i] = 0.5f * (1.0f - std::cos (2.0f * juce::MathConstants<float>::pi * (float) i / (float) (fftSize - 1)));

    startTimerHz (30);
}

SpectrumDisplay::~SpectrumDisplay()
{
    stopTimer();
}

void SpectrumDisplay::setScaleFactor (float scale)
{
    scaleFactor = juce::jlimit (0.6f, 2.0f, scale);
    repaint();
}

void SpectrumDisplay::computeLevels (const float* samples, int got, std::array<float, numBins>& out)
{
    for (auto& l : out)
        l = -200.0f;

    if (got <= 0)
        return;

    for (int i = 0; i < got; ++i)
        fftData[(size_t) i] = samples[i] * window[(size_t) i];
    for (int i = got; i < fftSize; ++i)
        fftData[(size_t) i] = 0.0f;

    std::fill (fftData.begin() + fftSize, fftData.end(), 0.0f);
    fft.performRealOnlyForwardTransform (fftData.data(), true);

    const float logNyquist = std::log ((float) (fftSize / 2));
    const float binScale = 2.0f / (float) fftSize;

    for (int i = 0; i < numBins; ++i)
    {
        const int kLo = jmax (1, (int) std::floor (std::exp (logNyquist * (float) i / (float) numBins)));
        const int kHi = jmin (fftSize / 2 - 1, (int) std::ceil (std::exp (logNyquist * (float) (i + 1) / (float) numBins)));

        for (int k = kLo; k <= kHi; ++k)
        {
            const float re = fftData[(size_t) (2 * k)];
            const float im = fftData[(size_t) (2 * k + 1)];
            const float mag = std::sqrt (re * re + im * im) * binScale;
            const float db = 20.0f * std::log10 (mag + 1e-6f);

            out[(size_t) i] = jmax (out[(size_t) i], db);
        }
    }
}

void SpectrumDisplay::timerCallback()
{
    juce::AudioBuffer<float> wetBuf (1, fftSize);
    juce::AudioBuffer<float> dryBuf (1, fftSize);
    const int got = processor.readWetSamples (wetBuf.getWritePointer (0), fftSize);
    processor.copyRecentDrySamples (dryBuf.getWritePointer (0), fftSize);

    computeLevels (wetBuf.getReadPointer (0), got, level);
    computeLevels (dryBuf.getReadPointer (0), fftSize, dryLevel);

    for (int i = 0; i < numBins; ++i)
    {
        const int lo = jmax (0, i - 1);
        const int hi = jmin (numBins - 1, i + 1);

        const float wetSmooth = 0.25f * level[(size_t) lo] + 0.5f * level[(size_t) i] + 0.25f * level[(size_t) hi];
        const float wetNorm = (juce::jlimit (minDb, topDb, wetSmooth) - minDb) / dbRange;
        float& wd = display[(size_t) i];
        wd += (wetNorm - wd) * (wetNorm > wd ? 0.45f : 0.15f);
        peak[(size_t) i] = jmax (peak[(size_t) i] * 0.93f, display[(size_t) i]);

        const float drySmooth = 0.25f * dryLevel[(size_t) lo] + 0.5f * dryLevel[(size_t) i] + 0.25f * dryLevel[(size_t) hi];
        const float dryNorm = (juce::jlimit (minDb, topDb, drySmooth) - minDb) / dbRange;
        float& dd = dryDisplay[(size_t) i];
        dd += (dryNorm - dd) * (dryNorm > dd ? 0.45f : 0.15f);
        dryPeak[(size_t) i] = jmax (dryPeak[(size_t) i] * 0.93f, dryDisplay[(size_t) i]);
    }

    repaint();
}

juce::Path SpectrumDisplay::makeCurve (const juce::Rectangle<float>& graph, const std::array<float, numBins>& data) const
{
    juce::Path p;
    for (int i = 0; i < numBins; ++i)
    {
        const float x = graph.getX() + graph.getWidth() * (float) (i + 0.5f) / (float) numBins;
        const float y = graph.getBottom() - data[(size_t) i] * graph.getHeight();
        if (i == 0)
            p.startNewSubPath (x, y);
        else
            p.lineTo (x, y);
    }
    return p;
}
void SpectrumDisplay::paint (juce::Graphics& g)
{
    const auto bounds = getLocalBounds().toFloat();
    const float s = scaleFactor;

    auto shRect = bounds.expanded (6.0f * s, 4.0f * s).translated (0.0f, 5.0f * s);
    auto shadowGrad = juce::ColourGradient (juce::Colour (0x00000000), { 0.0f, shRect.getY() },
                                            juce::Colour (0x00000000), { 0.0f, shRect.getBottom() },
                                            false);
    shadowGrad.addColour (0.45f, juce::Colour (0x70000000));
    g.setGradientFill (shadowGrad);
    g.fillRoundedRectangle (shRect, 14.0f * s);

    auto bodyGrad = juce::ColourGradient (juce::Colour (0xFF24323A), { 0.0f, bounds.getY() },
                                          juce::Colour (0xFF0B1216), { 0.0f, bounds.getBottom() },
                                          false);
    g.setGradientFill (bodyGrad);
    g.fillRoundedRectangle (bounds, 12.0f * s);

    g.setColour (juce::Colour (0x3A3CE0C8));
    g.drawRoundedRectangle (bounds.expanded (1.0f * s), 13.0f * s, 1.0f * s);

    g.setColour (juce::Colour (0x38FFFFFF));
    g.drawHorizontalLine (juce::roundToInt (bounds.getY() + 1.0f * s), juce::roundToInt (bounds.getX() + 3.0f * s), juce::roundToInt (bounds.getRight() - 3.0f * s));
    g.drawVerticalLine (juce::roundToInt (bounds.getX() + 1.0f * s), juce::roundToInt (bounds.getY() + 3.0f * s), juce::roundToInt (bounds.getBottom() - 3.0f * s));
    g.setColour (juce::Colour (0x50000000));
    g.drawHorizontalLine (juce::roundToInt (bounds.getBottom() - 1.0f * s), juce::roundToInt (bounds.getX() + 3.0f * s), juce::roundToInt (bounds.getRight() - 3.0f * s));
    g.drawVerticalLine (juce::roundToInt (bounds.getRight() - 1.0f * s), juce::roundToInt (bounds.getY() + 3.0f * s), juce::roundToInt (bounds.getBottom() - 3.0f * s));

    g.saveState();
    g.reduceClipRegion (bounds.toNearestInt().reduced (juce::roundToInt (1.0f * s)));
    auto refl = juce::ColourGradient (juce::Colour (0x28FFFFFF), { 0.0f, bounds.getY() },
                                      juce::Colour (0x00FFFFFF), { 0.0f, bounds.getY() + bounds.getHeight() * 0.45f },
                                      false);
    g.setGradientFill (refl);
    g.fillRect (bounds);
    g.setColour (juce::Colour (0x45FFFFFF));
    g.fillRect (bounds.getX() + 2.0f * s, bounds.getY() + 1.0f * s, bounds.getWidth() - 4.0f * s, 1.0f * s);
    g.restoreState();

    const auto area = bounds.reduced (14.0f * s, 8.0f * s);

    const float dbAxisW = 30.0f * s;
    const float graphAreaBottom = area.getBottom() - 16.0f * s;
    const float graphH = graphAreaBottom - area.getY();
    const auto graph = area.withHeight (graphH).withLeft (area.getX() + dbAxisW);

    g.setFont (juce::Font (juce::FontOptions (9.0f * s)));

    for (int i = 0; i < 4; ++i)
    {
        const float db = -6.0f - (float) i * 15.0f;
        const float y = graph.getBottom() - ((db - minDb) / dbRange) * graph.getHeight();
        g.setColour (juce::Colour (0x223A4A4D));
        g.drawHorizontalLine (juce::roundToInt (y), juce::roundToInt (graph.getX()), juce::roundToInt (graph.getRight()));
        g.setColour (juce::Colour (0x55FFFFFF));
        g.drawText (juce::String (juce::roundToInt (db)), juce::Rectangle<float> (graph.getX() - dbAxisW, y - 6.0f * s, dbAxisW - 6.0f * s, 12.0f * s),
                    juce::Justification::right, false);
    }

    const float sr = juce::jmax (1000.0f, (float) processor.getSampleRate());
    const float logNyquist = std::log ((float) (fftSize / 2));

    const float freqs[] = { 50.0f, 100.0f, 200.0f, 500.0f, 1000.0f, 2000.0f, 5000.0f, 10000.0f, 20000.0f };
    for (float f : freqs)
    {
        const float k = f * (float) fftSize / sr;
        if (k < 1.0f || k > (float) (fftSize / 2))
            continue;

        const float x = graph.getX() + (std::log (k) / logNyquist) * graph.getWidth();
        g.setColour (juce::Colour (0x223A4A4D));
        g.drawVerticalLine (juce::roundToInt (x), juce::roundToInt (graph.getY()), juce::roundToInt (graph.getBottom()));
        g.setColour (juce::Colour (0x55FFFFFF));
        g.drawText (formatFreq (f), juce::Rectangle<float> (x - 14.0f * s, graph.getBottom() + 3.0f * s, 28.0f * s, 12.0f * s),
                    juce::Justification::centred, false);
    }

    g.setColour (juce::Colour (0x33FFFFFF));
    g.drawRect (graph, 0.6f * s);

    if (graph.getWidth() > 2.0f * s && graph.getHeight() > 2.0f * s)
    {
        auto curve = makeCurve (graph, display);
        auto dry = makeCurve (graph, dryDisplay);

        auto wetFill = juce::Path (curve);
        wetFill.lineTo (graph.getRight(), graph.getBottom());
        wetFill.lineTo (graph.getX(), graph.getBottom());
        wetFill.closeSubPath();

        auto dryFill = juce::Path (dry);
        dryFill.lineTo (graph.getRight(), graph.getBottom());
        dryFill.lineTo (graph.getX(), graph.getBottom());
        dryFill.closeSubPath();

        auto wetGrad = juce::ColourGradient (wetFillCol, { 0.0f, graph.getY() },
                                             juce::Colour (0x063CE0C8), { 0.0f, graph.getBottom() },
                                             false);
        g.setGradientFill (wetGrad);
        g.fillPath (wetFill);

        auto dryGrad = juce::ColourGradient (dryFillCol, { 0.0f, graph.getY() },
                                             juce::Colour (0x00FFFFFF), { 0.0f, graph.getBottom() },
                                             false);
        g.setGradientFill (dryGrad);
        g.fillPath (dryFill);

        g.setColour (wetPeakCol);
        g.strokePath (makeCurve (graph, peak), juce::PathStrokeType (1.0f * s));

        g.setColour (dryLineCol);
        g.strokePath (dry, juce::PathStrokeType (1.2f * s,
                                                juce::PathStrokeType::JointStyle::curved,
                                                juce::PathStrokeType::EndCapStyle::rounded));

        g.setColour (dryPeakCol);
        g.strokePath (makeCurve (graph, dryPeak), juce::PathStrokeType (0.8f * s));

        g.setColour (wetLineCol);
        g.strokePath (curve, juce::PathStrokeType (1.4f * s,
                                                  juce::PathStrokeType::JointStyle::curved,
                                                  juce::PathStrokeType::EndCapStyle::rounded));
    }

    g.setColour (juce::Colour (0x66FFFFFF));
    g.setFont (juce::Font (juce::FontOptions (10.0f * s, juce::Font::bold)));
    g.drawText ("REVERB SPECTRUM", graph.withTop (graph.getY() - 2.0f * s).withBottom (graph.getY() + 14.0f * s),
                juce::Justification::bottomLeft, false);

    const float legendY = graph.getY() + 4.0f * s;
    const float legendX = graph.getRight() - 2.0f * s;
    const auto legendRow = juce::Rectangle<float> (legendX - 84.0f * s, legendY - 9.0f * s, 84.0f * s, 12.0f * s);

    g.setFont (juce::Font (juce::FontOptions (9.0f * s)));
    g.setColour (wetPeakCol);
    g.drawHorizontalLine (juce::roundToInt (legendRow.getY() + 8.0f * s), juce::roundToInt (legendRow.getX()), juce::roundToInt (legendRow.getX() + 11.0f * s));
    g.setColour (juce::Colour (0x88FFFFFF));
    g.drawText ("WET", legendRow.withX (legendRow.getX() + 14.0f * s).withWidth (28.0f * s),
                juce::Justification::left, false);

    g.setColour (dryPeakCol);
    g.drawHorizontalLine (juce::roundToInt (legendRow.getY() + 8.0f * s), juce::roundToInt (legendRow.getX() + 42.0f * s), juce::roundToInt (legendRow.getX() + 53.0f * s));
    g.setColour (juce::Colour (0x88FFFFFF));
    g.drawText ("DRY", legendRow.withX (legendRow.getX() + 56.0f * s).withWidth (28.0f * s),
                juce::Justification::left, false);
}