#include "Afterglow.h"
#include "Heatmap_Utils.h"
#include "Plottable_Real_Time_Data.h"
#include "../render/Blend2D_Cache.h"
#include "../renderable/Render_Partition.h"
#include "../renderable/Renderable_p.h"
#include <algorithm>
#include <deque>
#include <optional>
namespace renderive::detail {
namespace {
using Afterglow_History = Plottable_History_Real_Time_Data<std::vector<double>, std::deque<std::vector<double>>>;
struct Afterglow_Prepare_Buffer {
    Afterglow_Properties properties;
    std::deque<std::vector<double>> history;
    Axis_Raster_Layout layout;
    std::vector<double> intensity;
    std::vector<Pixel> pixels;
    int source_width{};
    int source_height{};
    int active_partitions{1};
    double maximum{1.0};
    std::size_t work_size{};
    bool valid{};
};
}
struct Afterglow_Control::Impl : Renderable::Impl {
    Impl(renderive_Owner<Frequency_Axis> frequency,
         renderive_Owner<Axis> power)
        : frequency_axis(std::move(frequency)),
          power_axis(std::move(power)) {}
    renderive_Owner<Frequency_Axis> frequency_axis;
    renderive_Owner<Axis> power_axis;
    std::optional<Afterglow_History> history;
    Adaptive_Render_Partitioner partitioner;
    Afterglow_Prepare_Buffer prepare_buffer;
};
Afterglow_Control::Afterglow_Control(const Afterglow_Properties& properties, renderive_Owner<Frequency_Axis> frequency_axis, renderive_Owner<Axis> power_axis)
    : Plottable_State(properties, std::make_unique<Impl>(
          std::move(frequency_axis), std::move(power_axis))) {
    d_func<Impl>().history.emplace(*this);
}
Afterglow_Control::~Afterglow_Control() = default;
std::size_t Afterglow_Control::history_count() const {
    return d_func<Impl>().history->size();
}
std::size_t Afterglow_Control::latest_spectrum_point_count() const {
    const auto history = d_func<Impl>().history->snapshot();
    return history.empty() ? 0 : history.back().size();
}
std::size_t Afterglow_Control::rendered_cell_count() const {
    const auto state = properties();
    const auto history = d_func<Impl>().history->snapshot();
    if (history.empty())
        return 0;
    const int width = std::min(state.frequency_point_size.get(), static_cast<int>(history.back().size()));
    const int height = state.power_point_size.get() > 0
                           ? state.power_point_size.get()
                           : std::max(1, static_cast<int>(d_func<Impl>().power_axis->transform().pixel_length));
    return width > 0 && height > 0 ? static_cast<std::size_t>(width) * static_cast<std::size_t>(height) : 0;
}
void Afterglow_Control::append_spectrum(std::span<const double> values) {
    if (get<&Afterglow_Properties::frequency_point_size>() <= 0)
        set<&Afterglow_Properties::frequency_point_size>(static_cast<int>(values.size()));
    d_func<Impl>().history->update({values.begin(), values.end()}, 64);
    render_graph_changed();
}
void Afterglow_Control::append_spectrum(std::pmr::vector<double>&& values) {
    append_spectrum(std::span<const double>(values.data(), values.size()));
}
void Afterglow_Control::publish() {
    publish_properties();
}

void Afterglow_Control::build_prepare_graph(Renderable_Graph_Builder& builder) {
    const auto view = d_func().render_state_view();
    const auto state = properties();
    const auto& history = view.get(*d_func<Impl>().history);
    const int width = history.empty()
        ? 0
        : std::min(state.frequency_point_size.get(),
                   static_cast<int>(history.back().size()));
    const int height = state.power_point_size.get() > 0
        ? state.power_point_size.get()
        : std::max(1, static_cast<int>(d_func<Impl>().power_axis->transform(view).pixel_length));
    const std::size_t work_size = width > 0 && height > 0
        ? static_cast<std::size_t>(width) * static_cast<std::size_t>(height)
        : 0;
    const int partition_count = d_func<Impl>().partitioner.graph_partition_count(
        state.partition_mode, state.partition_count.get(),
        static_cast<int>(Scene_Base::task_executor_worker_count()), work_size, 4096);
    const auto prepare = add_prepare_task(
        builder, "prepare", "Prepare Afterglow",
        [this, partition_count](const Prepare_Render_Context& context) {
            prepare_render_frame(context.frame.render_state, partition_count);
            if (context.metrics) {
                context.metrics->set(Node_Metric_Kind::input_count,
                                     d_func<Impl>().prepare_buffer.history.size());
                context.metrics->set(Node_Metric_Kind::chunk_size,
                                     d_func<Impl>().prepare_buffer.work_size /
                                         std::max(1, d_func<Impl>().prepare_buffer.active_partitions));
            }
        });
    std::vector<Renderable_Graph_Builder::Task> accumulation;
    accumulation.reserve(static_cast<std::size_t>(partition_count));
    for (int index = 0; index < partition_count; ++index) {
        const auto task = builder.emplace(
            "accumulate:" + std::to_string(index),
            "Afterglow Chunk " + std::to_string(index + 1) + " Accumulate",
            [this, index](const Prepare_Render_Context& context) {
                accumulate_partition(index);
                if (context.metrics) {
                    const auto range = render_partition_range(
                        static_cast<std::size_t>(d_func<Impl>().prepare_buffer.source_width),
                        index, d_func<Impl>().prepare_buffer.active_partitions);
                    context.metrics->set(Node_Metric_Kind::prepared_cells,
                                         (range.last - range.first) *
                                             d_func<Impl>().prepare_buffer.source_height);
                }
            });
        builder.precede(prepare, task);
        accumulation.push_back(task);
    }
    if (partition_count == 0)
        return;
    const auto normalize = builder.emplace(
        "normalize", "Normalize Afterglow",
        [this](const Prepare_Render_Context&) { normalize_render_frame(); });
    for (const auto task : accumulation)
        builder.precede(task, normalize);
    for (int index = 0; index < partition_count; ++index) {
        const auto task = builder.emplace(
            "color:" + std::to_string(index),
            "Afterglow Chunk " + std::to_string(index + 1) + " Color",
            [this, index](const Prepare_Render_Context&) { color_partition(index); });
        builder.precede(normalize, task);
    }
}

void Afterglow_Control::build_paint_graph(Renderable_Graph_Builder& builder) {
    const auto view = d_func().render_state_view();
    const auto state = properties();
    const auto& history = view.get(*d_func<Impl>().history);
    const int width = history.empty()
        ? 0
        : std::min(state.frequency_point_size.get(),
                   static_cast<int>(history.back().size()));
    const int height = state.power_point_size.get() > 0
        ? state.power_point_size.get()
        : std::max(1, static_cast<int>(d_func<Impl>().power_axis->transform(view).pixel_length));
    const std::size_t work_size = width > 0 && height > 0
        ? static_cast<std::size_t>(width) * static_cast<std::size_t>(height)
        : 0;
    const int partition_count = d_func<Impl>().partitioner.graph_partition_count(
        state.partition_mode, state.partition_count.get(),
        static_cast<int>(Scene_Base::task_executor_worker_count()), work_size, 4096);
    const auto paint_image = add_paint_task(
        builder, "paint", "Paint Afterglow",
        [this](Painter& painter, const Paint_Render_Context& context) {
            paint_render_frame(painter);
            if (context.metrics)
                context.metrics->set(Node_Metric_Kind::pixel_count,
                                     d_func<Impl>().prepare_buffer.work_size);
        });
    if (partition_count == 0) {
        builder.precede(builder.find("prepare"), paint_image);
    } else {
        for (int index = 0; index < partition_count; ++index)
            builder.precede(builder.find("color:" + std::to_string(index)), paint_image);
    }
}

void Afterglow_Control::prepare_render_frame(const Render_State_View& view,
                                             int graph_partition_count) {
    const auto& state = render_properties(view);
    const auto& history = view.get(*d_func<Impl>().history);
    auto& output = d_func<Impl>().prepare_buffer;
    output = {};
    output.properties = state;
    output.history.assign(history.begin(), history.end());
    if (history.empty())
        return;
    const int width = std::min(state.frequency_point_size.get(), static_cast<int>(history.back().size()));
    const int height = state.power_point_size.get() > 0
                           ? state.power_point_size.get()
                           : std::max(1, static_cast<int>(d_func<Impl>().power_axis->transform(view).pixel_length));
    if (width <= 0 || height <= 0)
        return;
    const auto layout = axis_raster_layout(d_func<Impl>().frequency_axis->transform(view),
                                           d_func<Impl>().power_axis->transform(view),
                                           state.frequency_range, state.power_range,
                                           width, height);
    if (!layout.valid())
        return;
    output.layout = layout;
    output.source_width = width;
    output.source_height = height;
    output.work_size = static_cast<std::size_t>(width) * height;
    output.intensity.resize(output.work_size);
    output.pixels.resize(output.work_size);
    output.active_partitions = d_func<Impl>().partitioner.begin(graph_partition_count,
                                                        output.work_size);
    output.valid = true;
}

void Afterglow_Control::accumulate_partition(int partition_index) {
    auto& output = d_func<Impl>().prepare_buffer;
    if (!output.valid || partition_index >= output.active_partitions)
        return;
    const auto& state = output.properties;
    const auto& history = output.history;
    const auto columns = render_partition_range(static_cast<std::size_t>(output.source_width),
                                                partition_index, output.active_partitions);
    for (std::size_t x = columns.first; x < columns.last; ++x)
        for (int y = 0; y < output.source_height; ++y)
            output.intensity[static_cast<std::size_t>(y) * output.source_width + x] = 0.0;
    double weight = 1.0;
    const double decay = 1.0 - state.attenuation_rate.get();
    for (auto iterator = history.rbegin(); iterator != history.rend(); ++iterator) {
        const std::size_t count = std::min<std::size_t>(output.source_width, iterator->size());
        for (std::size_t x = columns.first; x < std::min(columns.last, count); ++x) {
            const double normalized = normalized_value((*iterator)[x], state.power_range);
            const int y = std::clamp(static_cast<int>(normalized * (output.source_height - 1)),
                                     0, output.source_height - 1);
            output.intensity[static_cast<std::size_t>(y) * output.source_width + x] += weight;
            if (state.interpolate && y + 1 < output.source_height)
                output.intensity[static_cast<std::size_t>(y + 1) * output.source_width + x] += weight * 0.35;
        }
        weight *= decay;
        if (weight < 0.01)
            break;
    }
}

void Afterglow_Control::normalize_render_frame() {
    auto& output = d_func<Impl>().prepare_buffer;
    if (output.valid)
        output.maximum = std::max(1.0, *std::max_element(output.intensity.begin(),
                                                        output.intensity.end()));
}

void Afterglow_Control::color_partition(int partition_index) {
    auto& output = d_func<Impl>().prepare_buffer;
    if (!output.valid || partition_index >= output.active_partitions)
        return;
    const auto& state = output.properties;
    const auto range = render_partition_range(output.work_size, partition_index,
                                              output.active_partitions);
    for (std::size_t cell = range.first; cell < range.last; ++cell) {
        const int y = static_cast<int>(cell / static_cast<std::size_t>(output.source_width));
        const int x = static_cast<int>(cell % static_cast<std::size_t>(output.source_width));
        output.pixels[output.layout.index(x, y, output.source_width, output.source_height)] =
            state.color_map.at_normalized(output.intensity[cell] / output.maximum);
    }
}

void Afterglow_Control::paint_render_frame(Painter& painter) {
    const auto& output = d_func<Impl>().prepare_buffer;
    if (!output.valid)
        return;
    painter.heatmap(output.layout.target, output.layout.width, output.layout.height,
                    output.pixels, Image_Interpolation_Mode::Bilinear);
}

void Afterglow_Control::render_frame_completed(
    std::uint64_t target_interval_ns) {
    const auto& output = d_func<Impl>().prepare_buffer;
    if (!output.valid || !is_visible())
        return;
    const auto& state = output.properties;
    if (d_func<Impl>().partitioner.finish(
            state.partition_mode, output.active_partitions, target_interval_ns,
            static_cast<int>(Scene_Base::task_executor_worker_count()),
            output.work_size, 4096))
        render_graph_changed();
}

void Afterglow_Control::prepare_frame(const Prepare_Render_Context& context) {
    prepare_render_frame(context.frame.render_state, 1);
    accumulate_partition(0);
    normalize_render_frame();
    color_partition(0);
}

void Afterglow_Control::paint(Painter& painter,
                              const Paint_Render_Context&) {
    paint_render_frame(painter);
}
}