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

void 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(s.pos) *
            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 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(1), .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);

    auto shader = ShaderManager::get(kShader::BrushStroke);

    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;
}
StrokeSample Stroke::randomize_sample(const glm::vec3& pos, float pressure, float dir_angle)
{
    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::vec3(cosf(rad), sinf(rad), 0); };  // normalized direction vector
    auto rnd_bneg = [&] { return prng() % 2 == 0 ? -1.f : 1.f; }; // -1 or 1

    float size_dyn = m_brush->m_tip_size_pressure ? pressure : 1.f;
    float flow_dyn = m_brush->m_tip_flow_pressure ? pressure : 1.f;
    float opacity_dyn = m_brush->m_tip_opacity_pressure ? pressure : 1.f;
    float size = glm::min(m_brush->m_tip_size / glm::tan(glm::radians(m_camera.fov * 0.5f)), m_max_size);
    float randflipx = m_brush->m_tip_randflipx ? rnd_bneg() : 1.f;
    float randflipy = m_brush->m_tip_randflipy ? rnd_bneg() : 1.f;
    
    glm::vec2 scatter_axis = m_brush->m_jitter_scatter_bothaxis ? glm::vec2(1.f, 1.f) : glm::vec2(0.f, 1.f);
    auto scatter_scale = glm::vec3(scatter_axis * glm::orientate2(-dir_angle), 1.f);

    float aspect_jitter = m_brush->m_jitter_aspect_bothaxis ?
        rnd_neg() * 0.5f * m_brush->m_jitter_aspect :
        rnd_nor() * 0.5f * m_brush->m_jitter_aspect;
    float aspect = glm::clamp(m_brush->m_tip_aspect + aspect_jitter, 0.1f, 0.9f);
    glm::vec2 aspect_scale = {
        (aspect <= 0.5 ? aspect * 2.f : 1.f),
        (aspect > 0.5 ? 1.f - (aspect - .5f) * 2.f : 1.f)
    };

    StrokeSample s;
    s.origin = pos;
    s.scale = m_brush->m_tip_scale * randflipx * (m_brush->m_tip_flipx ? -1.f : 1.f) * aspect_scale;
    s.angle = (m_brush->m_tip_angle + rnd_neg() * m_brush->m_jitter_angle) * (float)(M_PI * 2.0);
    s.size = size * (1.f - rnd_nor() * m_brush->m_jitter_scale) * size_dyn;
    s.pos = pos + (scatter_scale * rnd_vec() * m_brush->m_jitter_scatter * s.size * 0.5f); // 0.5 because PS scatters by half size
    s.flow = m_brush->m_tip_flow * (1.f - rnd_nor() * m_brush->m_jitter_flow) * flow_dyn;
    s.opacity = m_brush->m_tip_opacity * (1.f - rnd_nor() * m_brush->m_jitter_opacity) * opacity_dyn;
    auto hsv = convert_rgb2hsv(m_brush->m_tip_color);
    hsv.x = glm::clamp(glm::mix(hsv.x, (pressure - 0.5f) * 2.0f, m_brush->m_tip_hue) + (rnd_nor() - 0.5f) * m_brush->m_jitter_hue, 0.f, 1.f);
    hsv.y = glm::clamp(glm::mix(hsv.y, (1.f - pressure - 0.5f) * 2.0f, m_brush->m_tip_sat) + (rnd_nor() - 0.5f) * m_brush->m_jitter_sat, 0.f, 1.f);
    hsv.z = glm::clamp(glm::mix(hsv.z, (pressure - 0.5f) * 2.0f, m_brush->m_tip_val) + (rnd_nor() - 0.5f) * m_brush->m_jitter_val, 0.f, 1.f);
    m_hsv_jitter.add(hsv);
    s.col = convert_hsv2rgb(m_hsv_jitter.average());
    return s;
}
std::vector<StrokeSample> 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;
        m_dir_dist += m_step;
        int old_kp = m_last_kp;
        while (m_dist > m_keypoints[m_last_kp + 1].dist)
            m_last_kp++;
        const auto& A = m_keypoints[old_kp];
        const auto& B = m_keypoints[m_last_kp == old_kp ? m_last_kp + 1 : m_last_kp]; // 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);

        if (m_dir_dist > m_dir_step && m_last_kp != m_dir_kp)
        {
            glm::vec2 v = glm::normalize(m_keypoints[m_last_kp].pos - m_keypoints[m_dir_kp].pos);
            m_dir_angle = -glm::orientedAngle(v, m_dir_ref);
            if (m_brush->m_tip_angle_smooth > 0 && (glm::abs(m_dir_angle) > glm::radians(30.f) || !m_dir_valid))
            {
                if (glm::abs(m_dir_angle) > glm::radians(100.f))
                {
                    //LOG("BIG ANGLE");
                    m_direction.clear();
                }

                auto old_dir = m_dir_ref;
                m_dir_ref = v;
                m_dir_ref_angle = -glm::orientedAngle(m_dir_ref, { 1, 0 });
                m_dir_angle = 0;
                auto angle_diff = -glm::orientedAngle(m_dir_ref, old_dir);
                for (int i = 0; i < m_direction.m_count; i++)
                    m_direction.m_vec[i] -= angle_diff;
            }
            m_dir_kp = m_last_kp;
            m_dir_dist = 0;
            if (!m_dir_valid)
            {
                m_dir_init = m_dir_angle + m_dir_ref_angle;
                m_dir_valid = true;
            }
        }

        bool need_dir = false;
        need_dir |= m_brush->m_tip_angle_follow;
        need_dir |= m_brush->m_tip_angle_init;
        need_dir |= m_brush->m_jitter_angle > 0;

        // angle is not ready yet
        if (need_dir && !m_dir_valid)
            continue;

        auto s = randomize_sample(pos, pressure, m_dir_angle + m_dir_ref_angle);
        if (s.valid())
        {
            if (m_brush->m_tip_angle_follow)
            {
                m_direction.add(m_dir_angle);
                s.angle += m_direction.average() + m_dir_ref_angle;
            }
            else if (m_brush->m_tip_angle_init)
            {
                s.angle += m_dir_init;
            }

            m_prev_sample = s;
            samples.push_back(s);
        }
        else
        {
            LOG("Invalid sample");
        }
    }
    return samples;
}
bool 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 Stroke::reset(bool clear_keypoints /*= false*/)
{
    m_dir_kp = 0;
    m_dir_angle = 0;
    m_dir_valid = false;
    m_dir_dist = 0;
    m_last_kp = 0;
    m_dist = 0.f;
    if (clear_keypoints)
        m_keypoints.clear();
}
void Stroke::add_point(glm::vec3 pos, float pressure)
{
#ifdef __IOS__
    m_curve = glm::min(m_curve + 0.1f, 1.f);
    //pressure = pressure * glm::pow(m_curve, 2.f);
    if (m_filter_points && m_hold_points.size() < 5)
    {
        m_hold_points.push_back({pos, pressure});
        return;
    }
#endif // __IOS__
    //m_pressure_buff.add(pressure);
    //pressure = m_pressure_buff.average();
    
    if (m_brush->m_tip_size_pressure)
    {
        float aspect_width = glm::min(1.f, glm::clamp(m_brush->m_tip_aspect, .1f, .9f) * 2.f);
        float raw_size = glm::clamp(m_brush->m_tip_size / glm::tan(glm::radians(m_camera.fov * 0.5f)), 1.f, m_max_size);
        float size = aspect_width * glm::min(m_brush->m_tip_scale.x, m_brush->m_tip_scale.y) * raw_size;
        m_step = glm::max(0.5f, m_brush->m_tip_spacing * size);
    }

    float dist = m_keypoints.empty() ? m_step :
        m_keypoints.back().dist + glm::distance(m_keypoints.back().pos, pos);
    if (m_keypoints.empty())
        m_prev_sample = randomize_sample(pos, pressure, 0);
    else if (m_keypoints.back().pos == pos)
        return; // skip same point, leading to black samples (NaN values)
    Keypoint kp;
    kp.pos = pos;
    kp.pressure = pressure;
    kp.dist = dist;
    m_keypoints.push_back(kp);
}
void Stroke::start(const std::shared_ptr<Brush>& brush)
{
    m_hold_points.clear();
    m_curve = 0.f;
    m_direction.clear();
    m_pressure_buff.clear();
    m_hsv_jitter.clear();
    m_last_kp = 0;
    m_dist = 0.f;
    m_dir_kp = 0;
    m_dir_angle = 0;
    m_dir_valid = false;
    m_dir_dist = 0;
    m_brush = brush;
    
    float aspect_width = glm::min(1.f, glm::clamp(m_brush->m_tip_aspect, .1f, .9f) * 2.f);
    float raw_size = glm::clamp(m_brush->m_tip_size / glm::tan(glm::radians(m_camera.fov * 0.5f)), 1.f, m_max_size);
    float size = aspect_width * glm::min(m_brush->m_tip_scale.x, m_brush->m_tip_scale.y) * raw_size;
    m_step = glm::max(0.5f, m_brush->m_tip_spacing * size);
    
    m_direction.resize(std::max<int>(1, m_brush->m_tip_angle_smooth * 200.f / m_step));
    prng.seed(0);
}

bool Brush::load_tip(const std::string& path, const std::string& thumb)
{
    m_tip_texture = std::make_shared<Texture2D>();
    if (!m_tip_texture->load(path))
    {
        m_tip_texture = nullptr;
        return false;
    }
    m_tip_texture->create_mipmaps();
    m_tip_texture->auto_destroy = true;
    m_brush_path = path;
    m_brush_thumb_path = thumb;
    auto sz = m_tip_texture->size();
    m_tip_scale = sz.y > sz.x ? glm::vec2(sz.x / sz.y, 1.f) : glm::vec2(1.f, sz.y / sz.x);
    return true;
}

bool Brush::load_dual(const std::string& path, const std::string& thumb)
{
    m_dual_texture = std::make_shared<Texture2D>();
    if (!m_dual_texture->load(path))
    {
        m_dual_texture = nullptr;
        return false;
    }
    m_dual_texture->create_mipmaps();
    m_dual_texture->auto_destroy = true;
    m_dual_path = path;
    m_dual_thumb_path = thumb;
    return true;
}

bool Brush::load_pattern(const std::string& path, const std::string& thumb)
{
    m_pattern_texture = std::make_shared<Texture2D>();
    if (!m_pattern_texture->load(path))
    {
        m_pattern_texture = nullptr;
        return false;
    }
    m_pattern_texture->create_mipmaps();
    m_pattern_texture->auto_destroy = true;
    m_pattern_path = path;
    m_pattern_thumb_path = thumb;
    return true;
}

bool Brush::load()
{
    if (!m_brush_path.empty())
    {
        m_tip_texture = std::make_shared<Texture2D>();
        if (!m_tip_texture->load(m_brush_path))
        {
            LOG("failed to load %s", m_brush_path.c_str());
            m_tip_texture = nullptr;
            return false;
        }
        m_tip_texture->create_mipmaps();
        m_tip_texture->auto_destroy = true;
    }
    if (!m_dual_path.empty())
    {
        m_dual_texture = std::make_shared<Texture2D>();
        if (!m_dual_texture->load(m_dual_path))
        {
            LOG("failed to load %s", m_dual_path.c_str());
            m_tip_texture = nullptr;
            m_dual_texture = nullptr;
            return false;
        }
        m_dual_texture->create_mipmaps();
        m_dual_texture->auto_destroy = true;
    }
    if (!m_pattern_path.empty())
    {
        m_pattern_texture = std::make_shared<Texture2D>();
        if (!m_pattern_texture->load(m_pattern_path))
        {
            LOG("failed to load %s", m_pattern_path.c_str());
            m_tip_texture = nullptr;
            m_dual_texture = nullptr;
            m_pattern_texture = nullptr;
            return false;
        }
        m_pattern_texture->create_mipmaps();
        m_pattern_texture->auto_destroy = true;
    }
    return true;
}