#include "Afterglow.h"
#include "Heatmap_Utils.h"
#include "Plottable_Real_Time_Data.h"
#include "../render/Blend2D_Cache.h"
#include "../renderable/Render_Partition.h"
#include <algorithm>
#include <deque>
namespace renderive {
namespace detail {
namespace {
using Afterglow_History = Plottable_History_Real_Time_Data<std::vector<double>, std::deque<std::vector<double>>>;
struct Afterglow_Render_Frame {
    Adaptive_Render_Partitioner partitioner;
    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 {
    Impl(Afterglow_Control& owner, std::shared_ptr<Frequency_Axis> frequency, std::shared_ptr<Axis> power)
        : frequency_axis(std::move(frequency)), power_axis(std::move(power)), history(owner) {}
    std::shared_ptr<Frequency_Axis> frequency_axis;
    std::shared_ptr<Axis> power_axis;
    Afterglow_History history;
    Afterglow_Render_Frame render_frame;
};
Afterglow_Control::Afterglow_Control(Plot_Core& plot, const Afterglow_Properties& properties, std::shared_ptr<Frequency_Axis> frequency_axis, std::shared_ptr<Axis> power_axis)
    : Plottable_State(plot, properties), impl_(std::make_unique<Impl>(*this, std::move(frequency_axis), std::move(power_axis))) {}
Afterglow_Control::~Afterglow_Control() = default;
std::size_t Afterglow_Control::history_count() const {
    return impl_->history.size();
}
std::size_t Afterglow_Control::latest_spectrum_point_count() const {
    const auto history = 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 = 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>(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()));
    impl_->history.update({values.begin(), values.end()}, 64);
    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_paint_task_graph(Renderable_Task_Graph& graph) {
    auto& output = impl_->render_frame;
    const auto view = render_state_view();
    const auto& state = render_properties(view);
    const auto& history = view.get(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>(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 = output.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 = graph.emplace([this, partition_count](const Scene_Render_Context&) {
        prepare_render_frame(render_state_view(), partition_count);
    }, "prepare afterglow");
    std::vector<Renderable_Task_Graph::Task> accumulation;
    accumulation.reserve(static_cast<std::size_t>(partition_count));
    for (int index = 0; index < partition_count; ++index) {
        const auto task = graph.emplace(
            [this, index](const Scene_Render_Context&) {
                accumulate_partition(render_state_view(), index);
            },
            "accumulate afterglow partition " + std::to_string(index + 1));
        graph.precede(prepare, task);
        accumulation.push_back(task);
    }
    const auto normalize = graph.emplace([this](const Scene_Render_Context&) {
        normalize_render_frame();
    }, "normalize afterglow");
    for (const auto task : accumulation)
        graph.precede(task, normalize);
    std::vector<Renderable_Task_Graph::Task> coloring;
    coloring.reserve(static_cast<std::size_t>(partition_count));
    for (int index = 0; index < partition_count; ++index) {
        const auto task = graph.emplace(
            [this, index](const Scene_Render_Context&) {
                color_partition(render_state_view(), index);
            },
            "color afterglow partition " + std::to_string(index + 1));
        graph.precede(normalize, task);
        coloring.push_back(task);
    }
    const auto paint_image = add_paint_task(
        graph, "paint afterglow shared image",
        [this](Painter& painter, const Render_State_View& frame_view,
               const Scene_Render_Context& context) {
            paint_render_frame(painter, frame_view,
                               context.frame_control_state.next_refresh_interval_ns);
        });
    for (const auto task : coloring)
        graph.precede(task, 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(impl_->history);
    auto& output = impl_->render_frame;
    output.valid = false;
    output.work_size = 0;
    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>(impl_->power_axis->transform(view).pixel_length));
    if (width <= 0 || height <= 0)
        return;
    const auto layout = axis_raster_layout(impl_->frequency_axis->transform(view),
                                           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 = output.partitioner.begin(graph_partition_count,
                                                        output.work_size);
    output.valid = true;
}

void Afterglow_Control::accumulate_partition(const Render_State_View& view,
                                              int partition_index) {
    auto& output = impl_->render_frame;
    if (!output.valid || partition_index >= output.active_partitions)
        return;
    const auto& state = render_properties(view);
    const auto& history = view.get(impl_->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 = impl_->render_frame;
    if (output.valid)
        output.maximum = std::max(1.0, *std::max_element(output.intensity.begin(),
                                                        output.intensity.end()));
}

void Afterglow_Control::color_partition(const Render_State_View& view,
                                        int partition_index) {
    auto& output = impl_->render_frame;
    if (!output.valid || partition_index >= output.active_partitions)
        return;
    const auto& state = render_properties(view);
    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 Render_State_View& view,
                                           std::uint64_t frame_interval_ns) {
    auto& output = impl_->render_frame;
    if (!output.valid)
        return;
    painter.heatmap(output.layout.target, output.layout.width, output.layout.height,
                    output.pixels, Image_Interpolation_Mode::Bilinear);
    const auto& state = render_properties(view);
    if (output.partitioner.finish(
            state.partition_mode, output.active_partitions, frame_interval_ns,
            static_cast<int>(Scene_Base::task_executor_worker_count()),
            output.work_size, 4096))
        task_graph_changed();
}

void Afterglow_Control::paint(Painter& painter, const Render_State_View& view) {
    prepare_render_frame(view, 1);
    accumulate_partition(view, 0);
    normalize_render_frame();
    color_partition(view, 0);
    paint_render_frame(painter, view, 0);
}
}
}