#include "CustomWaveform.hpp"

#include "PerFrameContext.hpp"
#include "PresetFileParser.hpp"

#include <Renderer/BlendMode.hpp>

#include <algorithm>
#include <cmath>

namespace libprojectM {
namespace MilkdropPreset {

static constexpr int CustomWaveformMaxSamples = std::max(Audio::WaveformSamples, Audio::SpectrumSamples);

CustomWaveform::CustomWaveform(PresetState& presetState)
    : m_presetState(presetState)
    , m_perFrameContext(presetState.globalMemory, &presetState.globalRegisters)
    , m_perPointContext(presetState.globalMemory, &presetState.globalRegisters)
    , m_mesh(Renderer::VertexBufferUsage::StreamDraw, true, false)
{
    m_perFrameContext.RegisterBuiltinVariables();
    m_perPointContext.RegisterBuiltinVariables();

    // Allocate space for max number of vertices possible, so we won't have to resize the vertex
    // buffers, which may change on each frame.
    m_mesh.SetVertexCount(std::max(Audio::SpectrumSamples, Audio::WaveformSamples) * 2 + 2);
}

void CustomWaveform::Initialize(PresetFileParser& parsedFile, int index)
{
    std::string const wavecodePrefix = "wavecode_" + std::to_string(index) + "_";
    std::string const wavePrefix = "wave_" + std::to_string(index) + "_";

    m_index = index;
    m_enabled = parsedFile.GetBool(wavecodePrefix + "enabled", m_enabled);
    m_samples = parsedFile.GetInt(wavecodePrefix + "samples", m_samples);
    m_sep = parsedFile.GetInt(wavecodePrefix + "sep", m_sep);
    m_spectrum = parsedFile.GetBool(wavecodePrefix + "bSpectrum", m_spectrum);
    m_useDots = parsedFile.GetBool(wavecodePrefix + "bUseDots", m_useDots);
    m_drawThick = parsedFile.GetBool(wavecodePrefix + "bDrawThick", m_drawThick);
    m_additive = parsedFile.GetBool(wavecodePrefix + "bAdditive", m_additive);
    m_scaling = parsedFile.GetFloat(wavecodePrefix + "scaling", m_scaling);
    m_smoothing = parsedFile.GetFloat(wavecodePrefix + "smoothing", m_smoothing);
    m_r = parsedFile.GetFloat(wavecodePrefix + "r", m_r);
    m_g = parsedFile.GetFloat(wavecodePrefix + "g", m_g);
    m_b = parsedFile.GetFloat(wavecodePrefix + "b", m_b);
    m_a = parsedFile.GetFloat(wavecodePrefix + "a", m_a);

    m_mesh.SetRenderPrimitiveType(m_useDots ? Renderer::Mesh::PrimitiveType::Points : Renderer::Mesh::PrimitiveType::LineStrip);
}

void CustomWaveform::CompileCodeAndRunInitExpressions(const PerFrameContext& presetPerFrameContext)
{
    m_perFrameContext.LoadStateVariables(m_presetState, presetPerFrameContext, *this);
    m_perFrameContext.EvaluateInitCode(m_presetState.customWaveInitCode[m_index], *this);

    for (int t = 0; t < TVarCount; t++)
    {
        m_tValuesAfterInitCode[t] = *m_perFrameContext.t_vars[t];
    }

    m_perFrameContext.CompilePerFrameCode(m_presetState.customWavePerFrameCode[m_index], *this);
    m_perPointContext.CompilePerPointCode(m_presetState.customWavePerPointCode[m_index], *this);
}

void CustomWaveform::Draw(const PerFrameContext& presetPerFrameContext)
{
    static_assert(Audio::WaveformSamples <= WaveformMaxPoints, "WaveformMaxPoints is larger than WaveformSamples");
    static_assert(Audio::SpectrumSamples <= WaveformMaxPoints, "WaveformMaxPoints is larger than SpectrumSamples");

    if (!m_enabled)
    {
        return;
    }

    int const maxSampleCount{m_spectrum ? Audio::SpectrumSamples : Audio::WaveformSamples};

    int sampleCount = std::min(maxSampleCount, static_cast<int>(*m_perFrameContext.samples));
    sampleCount -= m_sep;

    // Initialize and execute per-frame code
    LoadPerFrameEvaluationVariables(presetPerFrameContext);
    m_perFrameContext.ExecutePerFrameCode();

    // Copy Q and T vars to per-point context
    InitPerPointEvaluationVariables();

    sampleCount = std::min(maxSampleCount, static_cast<int>(*m_perFrameContext.samples));

    // If there aren't enough samples to draw a single line or dot, skip drawing the waveform.
    if ((m_useDots && sampleCount < 1) || sampleCount < 2)
    {
        return;
    }

    const auto* pcmL = m_spectrum
                           ? m_presetState.audioData.spectrumLeft.data()
                           : m_presetState.audioData.waveformLeft.data();
    const auto* pcmR = m_spectrum
                           ? m_presetState.audioData.spectrumRight.data()
                           : m_presetState.audioData.waveformRight.data();

    const float mult = m_scaling * m_presetState.waveScale * (m_spectrum ? 0.15f : 0.004f);

    // PCM data smoothing
    const int offset1 = m_spectrum ? 0 : (maxSampleCount - sampleCount) / 2 - m_sep / 2;
    const int offset2 = m_spectrum ? 0 : (maxSampleCount - sampleCount) / 2 + m_sep / 2;
    const float t = m_spectrum ? static_cast<float>(maxSampleCount - m_sep) / static_cast<float>(sampleCount) : 1.0f;
    const float mix1 = std::pow(m_smoothing * 0.98f, 0.5f);
    const float mix2 = 1.0f - mix1;

    std::array<float, CustomWaveformMaxSamples> sampleDataL{};
    std::array<float, CustomWaveformMaxSamples> sampleDataR{};

    sampleDataL[0] = pcmL[offset1];
    sampleDataR[0] = pcmR[offset2];

    // Smooth forward
    for (int sample = 1; sample < sampleCount; sample++)
    {
        sampleDataL[sample] = pcmL[static_cast<int>(sample * t) + offset1] * mix2 + sampleDataL[sample - 1] * mix1;
        sampleDataR[sample] = pcmR[static_cast<int>(sample * t) + offset2] * mix2 + sampleDataR[sample - 1] * mix1;
    }

    // Smooth backwards (this fixes the asymmetry of the beginning & end)
    for (int sample = sampleCount - 2; sample >= 0; sample--)
    {
        sampleDataL[sample] = sampleDataL[sample] * mix2 + sampleDataL[sample + 1] * mix1;
        sampleDataR[sample] = sampleDataR[sample] * mix2 + sampleDataR[sample + 1] * mix1;
    }

    // Scale waveform to final size
    for (int sample = 0; sample < sampleCount; sample++)
    {
        sampleDataL[sample] *= mult;
        sampleDataR[sample] *= mult;
    }

    std::vector<Renderer::Point> points(sampleCount);
    std::vector<Renderer::Color> colors(sampleCount);

    float const sampleMultiplicator = sampleCount > 1 ? 1.0f / static_cast<float>(sampleCount - 1) : 0.0f;
    for (int sample = 0; sample < sampleCount; sample++)
    {
        float const sampleIndex = static_cast<float>(sample) * sampleMultiplicator;
        LoadPerPointEvaluationVariables(sampleIndex, sampleDataL[sample], sampleDataR[sample]);

        m_perPointContext.ExecutePerPointCode();

        points[sample] = Renderer::Point(static_cast<float>((*m_perPointContext.x * 2.0 - 1.0) * m_presetState.renderContext.invAspectX),
                                         static_cast<float>((*m_perPointContext.y * -2.0 + 1.0) * m_presetState.renderContext.invAspectY));

        colors[sample] = Renderer::Color::Modulo(Renderer::Color(static_cast<float>(*m_perPointContext.r),
                                                                 static_cast<float>(*m_perPointContext.g),
                                                                 static_cast<float>(*m_perPointContext.b),
                                                                 static_cast<float>(*m_perPointContext.a)));
    }

    SmoothWave(points, colors);

#ifndef USE_GLES
    glDisable(GL_LINE_SMOOTH);
#endif
    glLineWidth(1);

    // Additive wave drawing (vice overwrite)
    if (m_additive)
    {
        Renderer::BlendMode::Set(true, Renderer::BlendMode::Function::SourceAlpha, Renderer::BlendMode::Function::One);
    }
    else
    {
        Renderer::BlendMode::Set(true, Renderer::BlendMode::Function::SourceAlpha, Renderer::BlendMode::Function::OneMinusSourceAlpha);
    }

    auto shader = m_presetState.untexturedShader.lock();
    shader->Bind();
    shader->SetUniformMat4x4("vertex_transformation", PresetState::orthogonalProjection);
    shader->SetUniformFloat("vertex_point_size", m_drawThick ? 2.0f : 1.0f);

    auto iterations = (m_drawThick && !m_useDots) ? 4 : 1;

    // Need to use +/- 1.0 here instead of 2.0 used in Milkdrop to achieve the same rendering result.
    auto incrementX = 1.0f / static_cast<float>(m_presetState.renderContext.viewportSizeX);
    auto incrementY = 1.0f / static_cast<float>(m_presetState.renderContext.viewportSizeX);

    size_t smoothedVertexCount = m_mesh.Indices().Size();

    auto& vertices = m_mesh.Vertices();

    // If thick outline is used, draw the shape four times with slight offsets
    // (top left, top right, bottom right, bottom left).
    for (auto iteration = 0; iteration < iterations; iteration++)
    {
        switch (iteration)
        {
            case 0:
            default:
                break;

            case 1:
                for (size_t j = 0; j < smoothedVertexCount; j++)
                {
                    vertices[j].SetX(vertices[j].X() + incrementX);
                }
                break;

            case 2:
                for (size_t j = 0; j < smoothedVertexCount; j++)
                {
                    vertices[j].SetY(vertices[j].Y() + incrementY);
                }
                break;

            case 3:
                for (size_t j = 0; j < smoothedVertexCount; j++)
                {
                    vertices[j].SetX(vertices[j].X() - incrementX);
                }
                break;
        }

        m_mesh.Update();
        m_mesh.Draw();
    }

    m_mesh.Unbind();
    Renderer::Shader::Unbind();
    Renderer::BlendMode::SetBlendActive(false);
}

void CustomWaveform::LoadPerFrameEvaluationVariables(const PerFrameContext& presetPerFrameContext)
{
    m_perFrameContext.LoadStateVariables(m_presetState, presetPerFrameContext, *this);
    m_perPointContext.LoadReadOnlyStateVariables(presetPerFrameContext);
}

void CustomWaveform::InitPerPointEvaluationVariables()
{
    for (int q = 0; q < QVarCount; q++)
    {
        *m_perPointContext.q_vars[q] = *m_perFrameContext.q_vars[q];
    }
    for (int t = 0; t < TVarCount; t++)
    {
        *m_perPointContext.t_vars[t] = *m_perFrameContext.t_vars[t];
    }
}

void CustomWaveform::LoadPerPointEvaluationVariables(float sample, float value1, float value2)
{
    *m_perPointContext.sample = static_cast<double>(sample);
    *m_perPointContext.value1 = static_cast<double>(value1);
    *m_perPointContext.value2 = static_cast<double>(value2);
    *m_perPointContext.x = static_cast<double>(0.5f + value1);
    *m_perPointContext.y = static_cast<double>(0.5f + value2);
    *m_perPointContext.r = *m_perFrameContext.r;
    *m_perPointContext.g = *m_perFrameContext.g;
    *m_perPointContext.b = *m_perFrameContext.b;
    *m_perPointContext.a = *m_perFrameContext.a;
}

void CustomWaveform::SmoothWave(const std::vector<Renderer::Point>& points, const std::vector<Renderer::Color>& colors)
{
    constexpr float c1{-0.15f};
    constexpr float c2{1.15f};
    constexpr float c3{1.15f};
    constexpr float c4{-0.15f};
    constexpr float inverseSum{1.0f / (c1 + c2 + c3 + c4)};

    size_t outputIndex = 0;
    size_t iBelow = 0;
    size_t iAbove2 = 1;

    size_t vertexCount = points.size();

    auto& outVertices = m_mesh.Vertices();
    auto& outColors = m_mesh.Colors();

    for (size_t inputIndex = 0; inputIndex < vertexCount - 1; inputIndex++)
    {
        size_t const iAbove = iAbove2;
        iAbove2 = std::min(vertexCount - 1, inputIndex + 2);
        outVertices[outputIndex] = points[inputIndex];
        outColors[outputIndex] = colors[inputIndex];
        outColors[outputIndex + 1] = colors[inputIndex];
        auto& smoothedPoint = outVertices[outputIndex + 1];
        smoothedPoint = points[inputIndex];
        smoothedPoint.SetX((c1 * points[iBelow].X() + c2 * points[inputIndex].X() + c3 * points[iAbove].X() + c4 * points[iAbove2].X()) * inverseSum);
        smoothedPoint.SetY((c1 * points[iBelow].Y() + c2 * points[inputIndex].Y() + c3 * points[iAbove].Y() + c4 * points[iAbove2].Y()) * inverseSum);
        iBelow = inputIndex;
        outputIndex += 2;
    }

    outVertices[outputIndex] = points[vertexCount - 1];
    outColors[outputIndex] = colors[vertexCount - 1];

    auto& indices = m_mesh.Indices();
    indices.Resize(outputIndex + 1);
    indices.MakeContinuous();
}

} // namespace MilkdropPreset
} // namespace libprojectM