#include "pch.h"
#include "log.h"
#include "brush.h"

void ui::BrushMesh::draw(const std::vector<StrokeSample>& samples, const glm::mat4& proj)
{
    std::vector<instance_t> attributes;
    attributes.reserve(samples.size());
    for (const auto& s : samples)
    {
        auto mvp = proj *
            glm::translate(glm::vec3(s.pos, 0)) *
            glm::scale(glm::vec3(s.size, s.size, 1)) *
            glm::eulerAngleZ(s.angle);
        attributes.emplace_back(instance_t{ mvp, s.flow });
    }
#ifdef USE_VBO
    glBindBuffer(GL_ARRAY_BUFFER, buffers[2]);
    glBufferData(GL_ARRAY_BUFFER, (int)(sizeof(instance_t) * attributes.size()), attributes.data(), GL_STATIC_DRAW);
    glBindBuffer(GL_ARRAY_BUFFER, 0);

    glBindVertexArray(vao);
    glDrawElementsInstanced(GL_TRIANGLES, 6, GL_UNSIGNED_SHORT, 0, (int)samples.size());
    glBindVertexArray(0);
#else
    glEnableVertexAttribArray(0);
    glEnableVertexAttribArray(1);
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, buffers[1]);
    glBindBuffer(GL_ARRAY_BUFFER, buffers[0]);
    glVertexAttribPointer(0, 4, GL_FLOAT, GL_FALSE, sizeof(vertex_t), (GLvoid*)offsetof(vertex_t, pos));
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, sizeof(vertex_t), (GLvoid*)offsetof(vertex_t, uvs));

    // Likewise, we can do the same with the model matrix. Note that a
    // matrix input to the vertex shader consumes N consecutive input
    // locations, where N is the number of columns in the matrix. So...
    // we have four vertex attributes to set up.
    glBindBuffer(GL_ARRAY_BUFFER, buffers[2]);
    glBufferData(GL_ARRAY_BUFFER, (int)(sizeof(instance_t) * attributes.size()), attributes.data(), GL_STATIC_DRAW);
    // Loop over each column of the matrix...
    for (int i = 0; i < 4; i++)
    {
        // Set up the vertex attribute
        glVertexAttribPointer(loc_mvp + i, 4, GL_FLOAT, GL_FALSE, sizeof(instance_t),
            (GLvoid*)(offsetof(instance_t, mvp) + sizeof(glm::vec4) * i));
        // Enable it
        glEnableVertexAttribArray(loc_mvp + i);
        // Make it instanced
        glVertexAttribDivisor(loc_mvp + i, 1);
    }
    glEnableVertexAttribArray(loc_flow);
    glVertexAttribPointer(loc_flow, 1, GL_FLOAT, GL_FALSE, sizeof(instance_t),
        (GLvoid*)offsetof(instance_t, flow));
    glVertexAttribDivisor(loc_flow, 1);

    glDrawElementsInstanced(GL_TRIANGLES, 6, GL_UNSIGNED_SHORT, 0, (int)samples.size());
    //glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_SHORT, 0);

    glDisableVertexAttribArray(0);
    glDisableVertexAttribArray(1);
    for (int i = 0; i < 4; i++)
        glDisableVertexAttribArray(loc_mvp + i);
    glDisableVertexAttribArray(loc_flow);

    glBindBuffer(GL_ARRAY_BUFFER, 0);
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, 0);
#endif // USE_VBO
}
bool ui::BrushMesh::create()
{
    static GLushort idx[6]{ 0, 1, 2, 0, 2, 3 };
    static vertex_t vertices[4]{
        { { -.5f, -.5f, 0, 1 }, { 0, 0 } }, // A       B----C
        { { -.5f,  .5f, 0, 1 }, { 0, 1 } }, // B  --\  |    |
        { {  .5f,  .5f, 0, 1 }, { 1, 1 } }, // C  --/  |    |
        { {  .5f, -.5f, 0, 1 }, { 1, 0 } }, // D       A----D
    };
    glGenBuffers(3, buffers);
    if (!(buffers[0] && buffers[1] && buffers[2]))
        return false;

    static instance_t inst{ glm::mat4(), .1f };
    glBindBuffer(GL_ARRAY_BUFFER, buffers[2]);
    glBufferData(GL_ARRAY_BUFFER, sizeof(instance_t), &inst, GL_STATIC_DRAW);

    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, buffers[1]);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(idx), idx, GL_STATIC_DRAW);
    glBindBuffer(GL_ARRAY_BUFFER, buffers[0]);
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, 0);
    glBindBuffer(GL_ARRAY_BUFFER, 0);

    // STROKE - INSTANCED
    static const char* shader_stroke_inst_v =
        SHADER_VERSION
        "in vec4 pos;"
        "in vec2 uvs;"
        "in mat4 a_mvp;"
        "in float a_flow;"
        "out vec3 uv;"
        "out float alpha;"
        "void main(){"
        "   uv = vec3(uvs, pos.w);"
        "   alpha = a_flow;"
        "   gl_Position = a_mvp * vec4(pos.xyz, 1.0);"
        "}";
    static const char* shader_stroke_inst_f =
        SHADER_VERSION
        "uniform mediump sampler2D tex;"
        "uniform mediump vec4 col;"
        "in mediump float alpha;"
        "in mediump vec3 uv;"
        "out mediump vec4 frag;"
        "void main(){"
        "   mediump float a = (1.0 - texture(tex, uv.xy).r) * alpha;"
        "   frag = vec4(col.rgb, a);"
        "}";


    if (!shader.create(shader_stroke_inst_v, shader_stroke_inst_f))
        LOG("Failed to create shader BrushMesh Stroke");

    loc_flow = shader.GetAttribLocation("a_flow");
    loc_mvp = shader.GetAttribLocation("a_mvp");

#if USE_VBO
    glGenVertexArrays(1, &vao);
    if (!vao)
        return false;
    glBindVertexArray(vao);
    glEnableVertexAttribArray(0);
    glEnableVertexAttribArray(1);
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, buffers[1]);
    glBindBuffer(GL_ARRAY_BUFFER, buffers[0]);
    glVertexAttribPointer(0, 4, GL_FLOAT, GL_FALSE, sizeof(vertex_t), (GLvoid*)offsetof(vertex_t, pos));
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, sizeof(vertex_t), (GLvoid*)offsetof(vertex_t, uvs));
    glBindBuffer(GL_ARRAY_BUFFER, buffers[2]);
    // Loop over each column of the matrix...
    for (int i = 0; i < 4; i++)
    {
        // Set up the vertex attribute
        glVertexAttribPointer(loc_mvp + i, 4, GL_FLOAT, GL_FALSE, sizeof(instance_t),
            (GLvoid*)(offsetof(instance_t, mvp) + sizeof(glm::vec4) * i));
        // Enable it
        glEnableVertexAttribArray(loc_mvp + i);
        // Make it instanced
        glVertexAttribDivisor(loc_mvp + i, 1);
    }
    glEnableVertexAttribArray(loc_flow);
    glVertexAttribPointer(loc_flow, 1, GL_FLOAT, GL_FALSE, sizeof(instance_t),
        (GLvoid*)offsetof(instance_t, flow));
    glVertexAttribDivisor(loc_flow, 1);
    glBindVertexArray(0);
#endif

    return true;
}
ui::StrokeSample ui::Stroke::randomize_sample(const glm::vec2& pos, float pressure)
{
    auto rnd_nor = [&] { return float((double)prng() / (double)prng.max()); };              // normalized [0, +1]
    //auto rnd_neg = [&] { return float((double)prng() / (double)prng.max() * 2.0 - 1.0); };  // normalized [-1, +1]
    auto rnd_rad = [&] { return float((double)prng() / (double)prng.max() * M_PI * 2.0); }; // normalized [0, 2pi]
    auto rnd_vec = [&] { float rad = rnd_rad(); return glm::vec2(cosf(rad), sinf(rad)); };  // normalized direction vector

    StrokeSample s;
    s.angle = (m_brush.m_tip_angle + rnd_nor() * m_brush.m_jitter_angle) * (float)(M_PI * 2.0);
    s.pos = pos + (rnd_vec() * m_brush.m_jitter_spread * 100.f);
    s.size = 100.f * m_brush.m_tip_size * (1.f - rnd_nor() * m_brush.m_jitter_scale);
    s.flow = m_brush.m_tip_flow * (1.f - rnd_nor() * m_brush.m_jitter_flow) * pressure;
    return s;
}
std::vector<ui::StrokeSample> ui::Stroke::compute_samples()
{
    if (m_keypoints.empty()) return {};
    int nsamples = (int)glm::floor((m_keypoints.back().dist - m_dist) / m_step);
    std::vector<StrokeSample> samples;
    samples.reserve(nsamples); // preallocate the estimate number of samples
    while (m_keypoints.back().dist > (m_dist + m_step))
    {
        m_dist += m_step;
        while (m_dist > m_keypoints[m_last_kp + 1].dist)
            m_last_kp++;
        const auto& A = m_keypoints[m_last_kp];
        const auto& B = m_keypoints[m_last_kp + 1]; // NOTE: this should be true when while is true
        float t = (m_dist - A.dist) / (B.dist - A.dist); // NOTE: must be A != B
        auto pos = glm::lerp(A.pos, B.pos, t);
        float pressure = glm::lerp(A.pressure, B.pressure, t);
        auto s = randomize_sample(pos, pressure);
        samples.push_back(s);
    }
    return std::move(samples);
}
bool ui::Stroke::has_sample()
{
    return m_keypoints.empty() ? false :                // no keypoints
        (m_keypoints.back().dist > (m_dist + m_step));  // check if next kp is closer than spacing
}
void ui::Stroke::reset(bool clear_keypoints /*= false*/)
{
    m_last_kp = 0;
    m_dist = 0.f;
    if (clear_keypoints)
        m_keypoints.clear();
}
void ui::Stroke::add_point(glm::vec2 pos, float pressure)
{
    float dist = m_keypoints.empty() ? 0.f :
        m_keypoints.back().dist + glm::distance(m_keypoints.back().pos, pos);
    m_keypoints.emplace_back();
    m_keypoints.back().pos = pos;
    m_keypoints.back().pressure = pressure;
    m_keypoints.back().dist = dist;
}
void ui::Stroke::start(const ui::Brush& brush)
{
    m_last_kp = 0;
    m_dist = 0.f;
    m_step = glm::max(brush.m_tip_spacing * brush.m_tip_size * 30, 0.1f);
    m_brush = brush;
    prng.seed(0);
}