#include "Waterfall.h"
#include "Heatmap_Utils.h"
#include "Plottable_Real_Time_Data.h"
#include "../renderable/Render_Partition.h"
#include "../renderable/Renderable_p.h"
#include <algorithm>
#include <deque>
#include <iomanip>
#include <optional>
#include <sstream>
namespace renderive::detail {
namespace {
struct Waterfall_Row {
    int tick{};
    std::vector<double> values;
};
struct Waterfall_Interaction_Base {};
struct Waterfall_Interaction {
    Hover_Tooltip_Runtime tooltip;
};
using Waterfall_History = Plottable_History_Real_Time_Data<Waterfall_Row, std::deque<Waterfall_Row>>;
using Waterfall_Interaction_State = Double_State_Strategy<Waterfall_Interaction_Base, Waterfall_Interaction>;
struct Waterfall_Prepare_Buffer {
    Waterfall_Properties properties;
    std::deque<Waterfall_Row> rows;
    Waterfall_Interaction interaction;
    Axis_Raster_Layout layout;
    Frequency_Columns columns;
    Axis_Transform frequency_axis;
    std::vector<Pixel> pixels;
    RectF tooltip_box;
    std::string tooltip_text;
    int source_width{};
    int source_height{};
    int active_partitions{1};
    std::size_t work_size{};
    bool valid{};
};
std::size_t waterfall_work_size(const Waterfall_Properties& state,
                                const std::deque<Waterfall_Row>& rows,
                                const Axis_Transform& frequency_axis) {
    if (rows.empty())
        return 0;
    const auto shortest = std::min_element(
        rows.begin(), rows.end(),
        [](const auto& left, const auto& right) {
            return left.values.size() < right.values.size();
        });
    const int source_width = std::max(
        0, std::min(state.frequency_bin_count.get(),
                    static_cast<int>(shortest->values.size())));
    const auto columns = frequency_columns(state.frequency_range,
                                           frequency_axis.coordinate_range,
                                           source_width,
                                           state.visible_range_only);
    return columns
        ? static_cast<std::size_t>(columns->last - columns->first + 1) * rows.size()
        : 0;
}
}
struct Waterfall_Control::Impl : Renderable::Impl {
    Impl(renderive_Owner<Frequency_Axis> frequency,
         renderive_Owner<Time_Axis> time)
        : frequency_axis(std::move(frequency)), time_axis(std::move(time)) {}
    renderive_Owner<Frequency_Axis> frequency_axis;
    renderive_Owner<Time_Axis> time_axis;
    std::optional<Waterfall_History> rows;
    Waterfall_Interaction_State interaction;
    Adaptive_Render_Partitioner partitioner;
    Waterfall_Prepare_Buffer prepare_buffer;
};
Waterfall_Control::Waterfall_Control(const Waterfall_Properties& properties, renderive_Owner<Frequency_Axis> frequency_axis, renderive_Owner<Time_Axis> time_axis)
    : Plottable_State(properties, std::make_unique<Impl>(
          std::move(frequency_axis), std::move(time_axis))) {
    d_func<Impl>().rows.emplace(*this);
}
Waterfall_Control::~Waterfall_Control() = default;
void Waterfall_Control::append_row(int tick, std::span<const double> values) {
    const std::size_t limit = static_cast<std::size_t>(
        std::max(2, d_func<Impl>().time_axis->get<&Time_Axis_Properties::visible_count>()));
    if(get<&Waterfall_Properties::frequency_bin_count>() <= 0)
        set<&Waterfall_Properties::frequency_bin_count>(static_cast<int>(values.size()));
    d_func<Impl>().rows->update({tick, {values.begin(), values.end()}}, limit);
    render_graph_changed();
}
void Waterfall_Control::append_row(int tick, std::pmr::vector<double>&& values) {
    append_row(tick, std::span<const double>(values.data(), values.size()));
}
void Waterfall_Control::append_row(Time_Of_Day time, std::span<const double> values) {
    append_row(d_func<Impl>().time_axis->append_time(time), values);
}
void Waterfall_Control::append_row(Time_Of_Day time, std::pmr::vector<double>&& values) {
    append_row(time, std::span<const double>(values.data(), values.size()));
}
std::size_t Waterfall_Control::row_count() const {
    return d_func<Impl>().rows->size();
}
std::size_t Waterfall_Control::stored_point_count() const {
    const auto rows = d_func<Impl>().rows->snapshot();
    std::size_t count{};
    for(const auto& row : rows)
        count += row.values.size();
    return count;
}
std::size_t Waterfall_Control::rendered_cell_count() const {
    const auto state = properties();
    const auto rows = d_func<Impl>().rows->snapshot();
    if(rows.empty())
        return 0;
    const int source_width = std::min(state.frequency_bin_count.get(), static_cast<int>(std::min_element(rows.begin(), rows.end(), [](const auto& left, const auto& right) { return left.values.size() < right.values.size(); })->values.size()));
    if(source_width <= 0)
        return 0;
    const auto columns = frequency_columns(
        state.frequency_range,
        d_func<Impl>().frequency_axis->get<&Axis_Properties::coordinates>(),
        source_width,
        state.visible_range_only);
    return columns ? static_cast<std::size_t>(columns->last - columns->first + 1) * rows.size() : 0;
}
void Waterfall_Control::handle_event(const Event& event) {
    bool updated{};
    d_func<Impl>().interaction.update([&](Waterfall_Interaction& interaction) { updated = update_hover_tooltip(interaction.tooltip, event); });
    if(updated)
        changed();
}
void Waterfall_Control::publish() {
    publish_properties();
    d_func<Impl>().interaction.publish();
}

void Waterfall_Control::build_prepare_graph(Renderable_Graph_Builder& builder) {
    const auto view = d_func().render_state_view();
    const auto state = properties();
    const auto& rows = view.get(*d_func<Impl>().rows);
    const std::size_t work_size = waterfall_work_size(
        state, rows, d_func<Impl>().frequency_axis->transform(view));
    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 Waterfall",
        [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.rows.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> partitions;
    partitions.reserve(static_cast<std::size_t>(partition_count));
    for (int index = 0; index < partition_count; ++index) {
        const auto partition = builder.emplace(
            "chunk_prepare:" + std::to_string(index),
            "Waterfall Chunk " + std::to_string(index + 1) + " Prepare",
            [this, index](const Prepare_Render_Context& context) {
                render_partition(index);
                if (context.metrics) {
                    const auto range = render_partition_range(
                        d_func<Impl>().prepare_buffer.work_size, index,
                        d_func<Impl>().prepare_buffer.active_partitions);
                    context.metrics->set(Node_Metric_Kind::prepared_cells,
                                         range.last - range.first);
                }
            });
        builder.precede(prepare, partition);
        partitions.push_back(partition);
    }
}
void Waterfall_Control::build_paint_graph(Renderable_Graph_Builder& builder) {
    const auto paint_image = add_paint_task(
        builder, "paint", "Paint Waterfall",
        [this](Painter& painter, const Paint_Render_Context& context) {
            paint_render_frame(painter, context.frame.render_state);
            if (context.metrics)
                context.metrics->set(Node_Metric_Kind::pixel_count,
                                     d_func<Impl>().prepare_buffer.work_size);
        });
    const auto view = d_func().render_state_view();
    const auto state = properties();
    const auto& rows = view.get(*d_func<Impl>().rows);
    const std::size_t work_size = waterfall_work_size(
        state, rows, d_func<Impl>().frequency_axis->transform(view));
    const int 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);
    if (count == 0) {
        builder.precede(builder.find("prepare"), paint_image);
    } else {
        for (int index = 0; index < count; ++index)
            builder.precede(builder.find("chunk_prepare:" + std::to_string(index)),
                            paint_image);
    }
}

void Waterfall_Control::prepare_render_frame(const Render_State_View& view,
                                             int graph_partition_count) {
    const auto& state = render_properties(view);
    const auto& rows = view.get(*d_func<Impl>().rows);
    auto& output = d_func<Impl>().prepare_buffer;
    output = {};
    output.properties = state;
    output.rows.assign(rows.begin(), rows.end());
    output.interaction = view.get(d_func<Impl>().interaction);
    if (rows.empty())
        return;
    const int source_width = std::min(state.frequency_bin_count.get(), static_cast<int>(std::min_element(rows.begin(), rows.end(), [](const auto& left, const auto& right) { return left.values.size() < right.values.size(); })->values.size()));
    const int height = static_cast<int>(rows.size());
    if (source_width <= 0 || height <= 0)
        return;
    const Axis_Transform frequency_axis = d_func<Impl>().frequency_axis->transform(view);
    const Axis_Transform time_axis = d_func<Impl>().time_axis->transform(view);
    const auto columns = frequency_columns(state.frequency_range, frequency_axis.coordinate_range,
                                           source_width, state.visible_range_only);
    if (!columns)
        return;
    const int width = columns->last - columns->first + 1;
    const Range time_range = rows.size() == 1
                                 ? time_axis.coordinate_range
                                 : Range{static_cast<double>(rows.front().tick),
                                         static_cast<double>(rows.back().tick)};
    const auto layout = axis_raster_layout(frequency_axis, time_axis, columns->range,
                                           time_range, width, height);
    if (!layout.valid())
        return;
    output.layout = layout;
    output.columns = *columns;
    output.frequency_axis = frequency_axis;
    output.source_width = width;
    output.source_height = height;
    output.work_size = static_cast<std::size_t>(width) * height;
    output.pixels.resize(output.work_size);
    output.active_partitions = d_func<Impl>().partitioner.begin(graph_partition_count,
                                                        output.work_size);
    if (state.tooltip_enabled && output.interaction.tooltip.active &&
        output.layout.target.contains(output.interaction.tooltip.position)) {
        const double frequency = frequency_axis.point_to_coord(
            output.interaction.tooltip.position);
        std::ostringstream text;
        text << std::fixed << std::setprecision(2) << frequency << " Hz";
        output.tooltip_text = text.str();
        output.tooltip_box = {output.interaction.tooltip.position.x + 8.0,
                              output.interaction.tooltip.position.y + 8.0,
                              110.0, 24.0};
    }
    output.valid = true;
}

void Waterfall_Control::render_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& rows = output.rows;
    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));
        const auto& row = rows[static_cast<std::size_t>(y)].values;
        output.pixels[output.layout.index(x, y, output.source_width, output.source_height)] =
            state.color_map.at_normalized(normalized_value(
                row[static_cast<std::size_t>(output.columns.first) + static_cast<std::size_t>(x)],
                    state.power_range));
    }
}

void Waterfall_Control::paint_render_frame(Painter& painter,
                                           const Render_State_View& view) {
    const auto& output = d_func<Impl>().prepare_buffer;
    if (!output.valid)
        return;
    const auto& paint_state = render_properties(view);
    painter.heatmap(output.layout.target, output.layout.width, output.layout.height,
                    output.pixels, paint_state.interpolation_mode);
    if (!output.tooltip_text.empty()) {
        painter.rect(output.tooltip_box, Pen{paint_state.tooltip_text_pen.color},
                     paint_state.tooltip_background_brush);
        painter.text({output.tooltip_box.x + 4.0, output.tooltip_box.y + 3.0},
                     output.tooltip_text, paint_state.tooltip_font,
                     paint_state.tooltip_text_pen);
    }
}

void Waterfall_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 Waterfall_Control::prepare_frame(const Prepare_Render_Context& context) {
    prepare_render_frame(context.frame.render_state, 1);
    render_partition(0);
}
void Waterfall_Control::paint(Painter& painter,
                              const Paint_Render_Context& context) {
    paint_render_frame(painter, context.frame.render_state);
}
}