#include "datascience/mc_random_walk_2d_gpu.h"
#include "gfx/gl_loader.h"
#include "gfx/gpu_compute_program.h"
#include "gfx/gpu_dispatch.h"
#include "gfx/gpu_rng_glsl.h"

#include <cstdio>
#include <vector>

namespace fn::ds {

constexpr int kSeedBinding = 9;
constexpr int kUserParams  = 16;

// trace_xy se almacena como float[2*N] xy-interleaved — directamente
// consumible por gpu_histogram_2d.
static const char* k_body_template = R"glsl(
layout(std430, binding = 0) buffer Trace        { float trace_xy[];      };
layout(std430, binding = 1) buffer AcceptCounts { uint  accept_counts[]; };
layout(std430, binding = 2) buffer CurrXY       { float curr_xy[];       };

uniform uint  u_n_walkers;
uniform uint  u_n_steps;
uniform float u_proposal_sigma;

void main() {
    uint wid = gl_GlobalInvocationID.x;
    if (wid >= u_n_walkers) return;

    uint  s   = rng_seeds[wid];
    vec2  xy  = vec2(curr_xy[2u * wid + 0u], curr_xy[2u * wid + 1u]);
    float lp  = target_log_pdf(xy);
    uint  acc = 0u;

    uint base = wid * u_n_steps * 2u;
    for (uint i = 0u; i < u_n_steps; ++i) {
        vec2 prop = xy + u_proposal_sigma * vec2(rng_normal(s), rng_normal(s));
        float lp_p = target_log_pdf(prop);
        float la = lp_p - lp;
        bool accept = false;
        if (la >= 0.0) {
            accept = true;
        } else {
            if (rng_uniform(s) < exp(la)) accept = true;
        }
        if (accept) { xy = prop; lp = lp_p; acc += 1u; }
        trace_xy[base + 2u * i + 0u] = xy.x;
        trace_xy[base + 2u * i + 1u] = xy.y;
    }

    rng_seeds[wid]              = s;
    curr_xy[2u * wid + 0u]      = xy.x;
    curr_xy[2u * wid + 1u]      = xy.y;
    accept_counts[wid]         += acc;
}
)glsl";

McRandomWalk2DGpu mc_rw2d_gpu_create(int m_walkers, int n_steps_per_run,
                                     const std::string& target_log_pdf_glsl) {
    McRandomWalk2DGpu s{};
    if (m_walkers <= 0 || n_steps_per_run <= 0) return s;

    auto rng = fn::gfx::glsl_rng_preamble(kSeedBinding);
    std::string preamble;
    preamble += rng;
    char buf[64];
    std::snprintf(buf, sizeof(buf), "uniform float u_user[%d];\n", kUserParams);
    preamble += buf;
    preamble += target_log_pdf_glsl;
    if (!preamble.empty() && preamble.back() != '\n') preamble += '\n';

    auto r = fn::gfx::compile_compute(k_body_template, 64, preamble);
    if (!r.ok) {
        std::fprintf(stderr, "[mc_rw2d_gpu] compile error: %s\n", r.err_msg.c_str());
        return s;
    }
    s.program = r.program;
    s.m_walkers = m_walkers;
    s.n_steps_per_run = n_steps_per_run;

    s.loc_n_walkers = static_cast<unsigned int>(glGetUniformLocation(s.program, "u_n_walkers"));
    s.loc_n_steps   = static_cast<unsigned int>(glGetUniformLocation(s.program, "u_n_steps"));
    s.loc_sigma     = static_cast<unsigned int>(glGetUniformLocation(s.program, "u_proposal_sigma"));

    s.trace_xy = fn::gfx::ssbo_create(
        static_cast<std::size_t>(m_walkers) *
        static_cast<std::size_t>(n_steps_per_run) * 2u * sizeof(float),
        nullptr, GL_DYNAMIC_COPY);
    s.accept_counts = fn::gfx::ssbo_create(
        static_cast<std::size_t>(m_walkers) * sizeof(unsigned int),
        nullptr, GL_DYNAMIC_COPY);
    s.curr_xy = fn::gfx::ssbo_create(
        static_cast<std::size_t>(m_walkers) * 2u * sizeof(float),
        nullptr, GL_DYNAMIC_COPY);
    s.rng_seeds = fn::gfx::ssbo_create(
        static_cast<std::size_t>(m_walkers) * sizeof(unsigned int),
        nullptr, GL_DYNAMIC_COPY);
    return s;
}

void mc_rw2d_gpu_reset(McRandomWalk2DGpu& s,
                       unsigned long long master_seed,
                       float x0, float y0,
                       const float* initial_xys) {
    if (s.m_walkers <= 0) return;

    std::vector<unsigned int> seeds(s.m_walkers);
    fn::gfx::seed_walkers_init(master_seed, seeds.data(), s.m_walkers);
    fn::gfx::ssbo_upload(s.rng_seeds, 0,
                         static_cast<std::size_t>(s.m_walkers) * sizeof(unsigned int),
                         seeds.data());

    std::vector<float> xy(static_cast<std::size_t>(s.m_walkers) * 2u);
    if (initial_xys) {
        for (std::size_t k = 0; k < xy.size(); ++k) xy[k] = initial_xys[k];
    } else {
        for (int j = 0; j < s.m_walkers; ++j) {
            xy[2 * j + 0] = x0;
            xy[2 * j + 1] = y0;
        }
    }
    fn::gfx::ssbo_upload(s.curr_xy, 0,
                         xy.size() * sizeof(float), xy.data());

    std::vector<unsigned int> zeros(s.m_walkers, 0u);
    fn::gfx::ssbo_upload(s.accept_counts, 0,
                         static_cast<std::size_t>(s.m_walkers) * sizeof(unsigned int),
                         zeros.data());
}

void mc_rw2d_gpu_run(McRandomWalk2DGpu& s,
                     float proposal_sigma,
                     const float* user_params, int n_user_params) {
    if (s.program == 0 || s.m_walkers <= 0) return;

    glUseProgram(s.program);
    fn::gfx::ssbo_bind(s.trace_xy,      0);
    fn::gfx::ssbo_bind(s.accept_counts, 1);
    fn::gfx::ssbo_bind(s.curr_xy,       2);
    fn::gfx::ssbo_bind(s.rng_seeds,     kSeedBinding);

    glUniform1ui(static_cast<GLint>(s.loc_n_walkers), static_cast<GLuint>(s.m_walkers));
    glUniform1ui(static_cast<GLint>(s.loc_n_steps),   static_cast<GLuint>(s.n_steps_per_run));
    glUniform1f (static_cast<GLint>(s.loc_sigma),     proposal_sigma);

    if (user_params && n_user_params > 0) {
        if (n_user_params > kUserParams) n_user_params = kUserParams;
        GLint loc = glGetUniformLocation(s.program, "u_user");
        if (loc >= 0) glUniform1fv(loc, n_user_params, user_params);
    }

    fn::gfx::dispatch_1d(s.m_walkers, 64);
    // No usamos barrier_buffer_update porque tipicamente lo siguiente es
    // alimentar gpu_histogram_2d (otro compute que lee este SSBO).
    fn::gfx::barrier_storage();
}

void mc_rw2d_gpu_readback_trace(const McRandomWalk2DGpu& s, float* out) {
    if (s.trace_xy.id == 0 || out == nullptr) return;
    fn::gfx::barrier_buffer_update();
    fn::gfx::ssbo_readback(
        s.trace_xy, 0,
        static_cast<std::size_t>(s.m_walkers) *
        static_cast<std::size_t>(s.n_steps_per_run) * 2u * sizeof(float),
        out);
}

void mc_rw2d_gpu_readback_accepts(const McRandomWalk2DGpu& s,
                                  unsigned int* out) {
    if (s.accept_counts.id == 0 || out == nullptr) return;
    fn::gfx::barrier_buffer_update();
    fn::gfx::ssbo_readback(
        s.accept_counts, 0,
        static_cast<std::size_t>(s.m_walkers) * sizeof(unsigned int),
        out);
}

const fn::gfx::Ssbo& mc_rw2d_gpu_trace_ssbo(const McRandomWalk2DGpu& s) {
    return s.trace_xy;
}

void mc_rw2d_gpu_destroy(McRandomWalk2DGpu& s) {
    fn::gfx::delete_compute_program(s.program);
    s.program = 0;
    fn::gfx::ssbo_destroy(s.trace_xy);
    fn::gfx::ssbo_destroy(s.accept_counts);
    fn::gfx::ssbo_destroy(s.curr_xy);
    fn::gfx::ssbo_destroy(s.rng_seeds);
    s.m_walkers = 0;
    s.n_steps_per_run = 0;
}

} // namespace fn::ds
