#include "Spectrum.h"
#include "Curve_Sampling.h"
#include "Plottable_Real_Time_Data.h"
#include "../render/Blend2D_Cache.h"
#include "../renderable/Render_Partition.h"
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <sstream>
namespace renderive::detail {
namespace {
struct Spectrum_Frame {
    std::vector<double> samples;
    std::vector<double> maxima;
    std::vector<double> minima;
};
struct Spectrum_Interaction_Base {};
struct Spectrum_Interaction {
    std::vector<double> markers;
    int selected_marker = -1;
    Hover_Tooltip_Runtime tooltip;
};
using Spectrum_Interaction_State = Double_State_Strategy<Spectrum_Interaction_Base, Spectrum_Interaction>;
struct Prepared_Curve {
    std::vector<PointF> points;
    std::vector<PointF> fill;
};
struct Spectrum_Partition_Buffer {
    RectF clip;
    Prepared_Curve maximum;
    Prepared_Curve minimum;
    Prepared_Curve current;
};
enum class Spectrum_Marker_Style {
    middle,
    marker,
    selected
};
struct Spectrum_Marker_Buffer {
    PointF first;
    PointF second;
    Spectrum_Marker_Style style{};
};
struct Spectrum_Extreme_Buffer {
    PointF point;
    bool maximum{};
};
struct Spectrum_Prepare_Buffer {
    Spectrum_Properties properties;
    Spectrum_Frame frame;
    Spectrum_Interaction interaction;
    Axis_Transform frequency_axis;
    Axis_Transform power_axis;
    RectF content;
    RectF sweep_region;
    std::vector<Spectrum_Partition_Buffer> partitions;
    std::vector<Spectrum_Marker_Buffer> markers;
    std::vector<Spectrum_Extreme_Buffer> extremes;
    RectF tooltip_box;
    std::string tooltip_text;
    int active_partitions{1};
    std::size_t work_size{};
    bool valid{};
};

struct Curve_Partition {
    std::span<const double> values;
    Range domain;
    Range clip_domain;
};

Curve_Partition curve_partition(std::span<const double> values, Range domain,
                                int partition_index, int partition_count) {
    if (values.size() < 2 || partition_count <= 0)
        return {};
    const std::size_t segment_count = values.size() - 1;
    const auto core = render_partition_range(segment_count, partition_index, partition_count);
    if (core.first == core.last)
        return {};
    const std::size_t first = core.first == 0 ? 0 : core.first - 1;
    const std::size_t last = std::min(segment_count, core.last + 1);
    const auto coordinate = [domain, segment_count](std::size_t index) {
        return domain.origin + domain.length() * static_cast<double>(index) /
                                   static_cast<double>(segment_count);
    };
    return {
        values.subspan(first, last - first + 1),
        {coordinate(first), coordinate(last)},
        {coordinate(core.first), coordinate(core.last)}
    };
}
Prepared_Curve prepare_curve(std::span<const double> values, Range domain,
                             const Axis_Transform& frequency_axis,
                             const Axis_Transform& power_axis, bool visible_only,
                             Line_Interpolation_Mode interpolation) {
    Prepared_Curve result;
    result.points = curve_points(values, domain, frequency_axis, power_axis,
                                 visible_only, interpolation);
    if(result.points.size() < 2)
        return result;
    result.fill.reserve(result.points.size() + 2);
    PointF first_baseline = result.points.front();
    PointF last_baseline = result.points.back();
    const double baseline =
        power_axis.coord_to_pixel(power_axis.coordinate_range.target);
    if (power_axis.orientation == Orientation::Horizontal) {
        first_baseline.x = baseline;
        last_baseline.x = baseline;
    } else {
        first_baseline.y = baseline;
        last_baseline.y = baseline;
    }
    result.fill.push_back(first_baseline);
    result.fill.insert(result.fill.end(), result.points.begin(), result.points.end());
    result.fill.push_back(last_baseline);
    return result;
}
void paint_curve(Painter& painter, const Prepared_Curve& curve, const Pen& pen,
                 const Brush& brush) {
    if(curve.points.size() < 2)
        return;
    if(brush.enabled() && !curve.fill.empty())
        painter.polygon(curve.fill, Pen{.style = Line_Style::None}, brush);
    painter.polyline(curve.points, pen);
}
double spectrum_power_at(const Spectrum_Properties& properties, const Spectrum_Frame& frame, double frequency, bool& ok) {
    ok = false;
    if(frame.samples.empty() || !properties.frequency_range.contains(frequency) || properties.frequency_range.length() == 0.0)
        return 0.0;
    const double normalized = (frequency - properties.frequency_range.origin) / properties.frequency_range.length();
    const double position = std::clamp(normalized, 0.0, 1.0) * static_cast<double>(frame.samples.size() - 1);
    const auto lower = static_cast<std::size_t>(std::floor(position));
    const auto upper = std::min(lower + 1, frame.samples.size() - 1);
    const double fraction = position - static_cast<double>(lower);
    ok = true;
    return frame.samples[lower] * (1.0 - fraction) + frame.samples[upper] * fraction;
}
}
struct Spectrum_Control::Impl {
    Impl(Spectrum_Control& owner, renderive_Owner<Frequency_Axis> frequency, renderive_Owner<Axis> power)
        : frequency_axis(std::move(frequency)), power_axis(std::move(power)), frame(owner) {}
    renderive_Owner<Frequency_Axis> frequency_axis;
    renderive_Owner<Axis> power_axis;
    Plottable_Latest_Real_Time_Data<Spectrum_Frame> frame;
    Spectrum_Interaction_State interaction;
    std::mutex frame_update_mutex;
    Adaptive_Render_Partitioner partitioner;
    Spectrum_Prepare_Buffer prepare_buffer;
};
Spectrum_Control::Spectrum_Control(const Spectrum_Properties& properties, renderive_Owner<Frequency_Axis> frequency_axis, renderive_Owner<Axis> power_axis)
    : Plottable_State(properties), impl_(std::make_unique<Impl>(*this, std::move(frequency_axis), std::move(power_axis))) {}
Spectrum_Control::~Spectrum_Control() = default;
void Spectrum_Control::update_samples(std::span<const double> values) {
    if(get<&Spectrum_Properties::frequency_point_size>() <= 0)
        set<&Spectrum_Properties::frequency_point_size>(static_cast<int>(values.size()));
    std::lock_guard lock(impl_->frame_update_mutex);
    Spectrum_Frame frame = impl_->frame.snapshot().value_or(Spectrum_Frame{});
    frame.samples.assign(values.begin(), values.end());
    if(frame.maxima.size() != values.size())
        frame.maxima.assign(values.begin(), values.end());
    else
        for(std::size_t index = 0; index < values.size(); ++index)
            frame.maxima[index] = std::max(frame.maxima[index], values[index]);
    if(frame.minima.size() != values.size())
        frame.minima.assign(values.begin(), values.end());
    else
        for(std::size_t index = 0; index < values.size(); ++index)
            frame.minima[index] = std::min(frame.minima[index], values[index]);
    impl_->frame.update(std::move(frame));
    render_graph_changed();
}
void Spectrum_Control::update_samples(std::pmr::vector<double>&& values) {
    update_samples(std::span<const double>(values.data(), values.size()));
}
std::size_t Spectrum_Control::sample_count() const {
    const auto frame = impl_->frame.snapshot();
    return frame ? frame->samples.size() : 0;
}
std::size_t Spectrum_Control::rendered_point_count() const {
    const auto state = properties();
    const auto frame = impl_->frame.snapshot();
    return frame ? curve_points(frame->samples, state.frequency_range, impl_->frequency_axis->transform(), impl_->power_axis->transform(), state.visible_range_only, state.interpolation_mode).size() : 0;
}
double Spectrum_Control::power_at(double frequency, bool& ok) const {
    const auto state = properties();
    const auto frame = impl_->frame.snapshot();
    return frame ? spectrum_power_at(state, *frame, frequency, ok) : (ok = false, 0.0);
}
void Spectrum_Control::add_custom_marker(double frequency) {
    add_custom_line_marker(frequency);
}
void Spectrum_Control::add_custom_line_marker(double frequency) {
    impl_->interaction.update([frequency](Spectrum_Interaction& interaction) { interaction.markers.push_back(frequency); });
    changed();
}
void Spectrum_Control::remove_custom_marker(double frequency) {
    bool removed{};
    impl_->interaction.update([&](Spectrum_Interaction& interaction) {
        if(interaction.markers.empty())
            return;
        auto closest = std::min_element(interaction.markers.begin(), interaction.markers.end(), [frequency](double left, double right) { return std::abs(left - frequency) < std::abs(right - frequency); });
        const int removed_index = static_cast<int>(std::distance(interaction.markers.begin(), closest));
        interaction.markers.erase(closest);
        if(interaction.selected_marker == removed_index)
            interaction.selected_marker = -1;
        else if(interaction.selected_marker > removed_index)
            --interaction.selected_marker;
        removed = true;
    });
    if(removed)
        changed();
}
void Spectrum_Control::remove_selected_marker() {
    bool removed{};
    impl_->interaction.update([&](Spectrum_Interaction& interaction) {
        if(interaction.selected_marker < 0 || interaction.selected_marker >= static_cast<int>(interaction.markers.size()))
            return;
        interaction.markers.erase(interaction.markers.begin() + interaction.selected_marker);
        interaction.selected_marker = -1;
        removed = true;
    });
    if(removed)
        changed();
}
void Spectrum_Control::clear_custom_markers() {
    impl_->interaction.update([](Spectrum_Interaction& interaction) {
        interaction.markers.clear();
        interaction.selected_marker = -1;
    });
    changed();
}
int Spectrum_Control::selectable_line_marker_count() const {
    return impl_->interaction.read([](const Spectrum_Interaction& interaction) { return static_cast<int>(interaction.markers.size()); });
}
int Spectrum_Control::selected_marker_index() const {
    return impl_->interaction.get<&Spectrum_Interaction::selected_marker>();
}
void Spectrum_Control::set_selected_marker_index(int index) {
    impl_->interaction.update([index](Spectrum_Interaction& interaction) { interaction.selected_marker = index >= 0 && index < static_cast<int>(interaction.markers.size()) ? index : -1; });
    changed();
}
void Spectrum_Control::select_next_marker() {
    impl_->interaction.update([](Spectrum_Interaction& interaction) {
        if(interaction.markers.empty())
            interaction.selected_marker = -1;
        else
            interaction.selected_marker = (interaction.selected_marker + 1) % static_cast<int>(interaction.markers.size());
    });
    changed();
}
void Spectrum_Control::select_previous_marker() {
    impl_->interaction.update([](Spectrum_Interaction& interaction) {
        if(interaction.markers.empty())
            interaction.selected_marker = -1;
        else
            interaction.selected_marker = (interaction.selected_marker <= 0 ? static_cast<int>(interaction.markers.size()) : interaction.selected_marker) - 1;
    });
    changed();
}
void Spectrum_Control::clear_marker_selection() {
    set_selected_marker_index(-1);
}
double Spectrum_Control::marker_frequency(int index) const {
    return impl_->interaction.read([index](const Spectrum_Interaction& interaction) { return index >= 0 && index < static_cast<int>(interaction.markers.size()) ? interaction.markers[index] : 0.0; });
}
void Spectrum_Control::set_marker_frequency(int index, double frequency) {
    bool updated{};
    impl_->interaction.update([&](Spectrum_Interaction& interaction) {
        if(index < 0 || index >= static_cast<int>(interaction.markers.size()))
            return;
        interaction.markers[index] = frequency;
        updated = true;
    });
    if(updated)
        changed();
}
void Spectrum_Control::set_current_marker_frequency(double frequency) {
    const int index = selected_marker_index();
    if(index >= 0)
        set_marker_frequency(index, frequency);
}
void Spectrum_Control::handle_event(const Event& event) {
    bool updated{};
    impl_->interaction.update([&](Spectrum_Interaction& interaction) { updated = update_hover_tooltip(interaction.tooltip, event); });
    if(updated)
        changed();
}
void Spectrum_Control::publish() {
    publish_properties();
    impl_->interaction.publish();
}

void Spectrum_Control::build_prepare_graph(Renderable_Graph_Builder& builder) {
    const auto view = render_state_view();
    const auto state = properties();
    const auto& published_frame = view.get(impl_->frame);
    const std::size_t work_size = published_frame && published_frame->samples.size() > 1
        ? published_frame->samples.size() - 1
        : 0;
    const int partition_count = impl_->partitioner.graph_partition_count(
        state.partition_mode, state.partition_count.get(),
        static_cast<int>(Scene_Base::task_executor_worker_count()), work_size, 128);
    const auto prepare = add_prepare_task(
        builder, "prepare", "Prepare Spectrum",
        [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,
                                     impl_->prepare_buffer.frame.samples.size());
                context.metrics->set(Node_Metric_Kind::chunk_size,
                                     impl_->prepare_buffer.work_size /
                                         std::max(1, impl_->prepare_buffer.active_partitions));
            }
        });
    for (int index = 0; index < partition_count; ++index) {
        const auto partition = builder.emplace(
            "chunk_prepare:" + std::to_string(index),
            "Spectrum Chunk " + std::to_string(index + 1) + " Prepare",
            [this, index](const Prepare_Render_Context& context) {
                prepare_partition(index);
                if (context.metrics) {
                    const auto range = render_partition_range(
                        impl_->prepare_buffer.work_size, index,
                        impl_->prepare_buffer.active_partitions);
                    context.metrics->set(Node_Metric_Kind::prepared_cells,
                                         range.last - range.first);
                }
            });
        builder.precede(prepare, partition);
    }
}

void Spectrum_Control::build_paint_graph(Renderable_Graph_Builder& builder) {
    const auto view = render_state_view();
    const auto state = properties();
    const auto& published_frame = view.get(impl_->frame);
    const std::size_t work_size = published_frame && published_frame->samples.size() > 1
        ? published_frame->samples.size() - 1
        : 0;
    const int partition_count = impl_->partitioner.graph_partition_count(
        state.partition_mode, state.partition_count.get(),
        static_cast<int>(Scene_Base::task_executor_worker_count()), work_size, 128);
    const auto paint = add_paint_task(
        builder, "paint", "Paint Spectrum",
        [this, partition_count](Painter& painter,
                                const Paint_Render_Context& context) {
            paint_background(painter, context.frame.render_state);
            for (int index = 0; index < partition_count; ++index)
                paint_partition(painter, index, context.frame.render_state);
            paint_overlay(painter, context.frame.render_state);
            if (context.metrics)
                context.metrics->set(Node_Metric_Kind::primitive_count,
                                     impl_->prepare_buffer.work_size);
        });
    if (partition_count == 0)
        builder.precede(builder.find("prepare"), paint);
    else
        for (int index = 0; index < partition_count; ++index)
            builder.precede(builder.find("chunk_prepare:" + std::to_string(index)),
                            paint);
}

void Spectrum_Control::prepare_render_frame(const Render_State_View& view,
                                            int graph_partition_count) {
    const auto& published_frame = view.get(impl_->frame);
    auto& output = impl_->prepare_buffer;
    output = {};
    output.properties = render_properties(view);
    if (published_frame)
        output.frame = *published_frame;
    output.interaction = view.get(impl_->interaction);
    output.frequency_axis = impl_->frequency_axis->transform(view);
    output.power_axis = impl_->power_axis->transform(view);
    output.valid = axes_are_orthogonal(output.frequency_axis, output.power_axis);
    output.work_size = published_frame && published_frame->samples.size() > 1
        ? published_frame->samples.size() - 1
        : 0;
    output.active_partitions = impl_->partitioner.begin(graph_partition_count,
                                                        output.work_size);
    output.partitions.resize(static_cast<std::size_t>(graph_partition_count));
    if (!output.valid)
        return;
    output.content = mapped_rect(output.frequency_axis, output.power_axis,
                                 output.frequency_axis.coordinate_range,
                                 output.power_axis.coordinate_range);
    if (output.content.empty()) {
        output.valid = false;
        return;
    }
    const auto& state = output.properties;
    if (state.sweep_region_visible)
        output.sweep_region = mapped_rect(output.frequency_axis, output.power_axis,
                                          state.sweep_frequency_range,
                                          output.power_axis.coordinate_range);
    output.markers.push_back({
        mapped_point(output.frequency_axis, state.center_frequency,
                     output.power_axis, output.power_axis.coordinate_range.origin),
        mapped_point(output.frequency_axis, state.center_frequency,
                     output.power_axis, output.power_axis.coordinate_range.target),
        Spectrum_Marker_Style::middle});
    for (std::size_t index = 0; index < output.interaction.markers.size(); ++index) {
        const double frequency = output.interaction.markers[index];
        output.markers.push_back({
            mapped_point(output.frequency_axis, frequency, output.power_axis,
                         output.power_axis.coordinate_range.origin),
            mapped_point(output.frequency_axis, frequency, output.power_axis,
                         output.power_axis.coordinate_range.target),
            static_cast<int>(index) == output.interaction.selected_marker
                ? Spectrum_Marker_Style::selected
                : Spectrum_Marker_Style::marker});
    }
    if (!output.frame.samples.empty() &&
        (state.max_marker_visible || state.use_min_marker)) {
        const auto prepare_extreme = [&](bool maximum) {
            const auto iterator = maximum
                ? std::max_element(output.frame.samples.begin(), output.frame.samples.end())
                : std::min_element(output.frame.samples.begin(), output.frame.samples.end());
            const std::size_t index = static_cast<std::size_t>(
                std::distance(output.frame.samples.begin(), iterator));
            const double denominator = output.frame.samples.size() > 1
                ? static_cast<double>(output.frame.samples.size() - 1)
                : 1.0;
            const double frequency = state.frequency_range.origin +
                state.frequency_range.length() * static_cast<double>(index) / denominator;
            output.extremes.push_back({
                mapped_point(output.frequency_axis, frequency, output.power_axis, *iterator),
                maximum});
        };
        if (state.max_marker_visible)
            prepare_extreme(true);
        if (state.use_min_marker)
            prepare_extreme(false);
    }
    if (state.tooltip_enabled && output.interaction.tooltip.active &&
        output.content.contains(output.interaction.tooltip.position)) {
        const double frequency = output.frequency_axis.point_to_coord(
            output.interaction.tooltip.position);
        bool ok{};
        const double power = spectrum_power_at(state, output.frame, frequency, ok);
        if (ok) {
            std::ostringstream text;
            text << std::fixed << std::setprecision(2) << frequency << " Hz  " << power;
            output.tooltip_text = text.str();
            output.tooltip_box = {output.interaction.tooltip.position.x + 8.0,
                                  output.interaction.tooltip.position.y + 8.0,
                                  170.0, 24.0};
        }
    }
}

void Spectrum_Control::prepare_partition(int partition_index) {
    auto& output = impl_->prepare_buffer;
    if (!output.valid || partition_index >= output.active_partitions)
        return;
    const auto& state = output.properties;
    const auto& frame = output.frame;
    const auto current = curve_partition(frame.samples, state.frequency_range,
                                         partition_index, output.active_partitions);
    if (current.values.empty())
        return;
    auto& partition = output.partitions[static_cast<std::size_t>(partition_index)];
    partition.clip = mapped_rect(output.frequency_axis, output.power_axis,
                                 current.clip_domain,
                                 output.power_axis.coordinate_range);
    if (state.max_hold_visible) {
        const auto maximum = curve_partition(frame.maxima, state.frequency_range,
                                             partition_index, output.active_partitions);
        partition.maximum = prepare_curve(
            maximum.values, maximum.domain, output.frequency_axis, output.power_axis,
            state.visible_range_only, state.interpolation_mode);
    }
    if (state.min_hold_visible) {
        const auto minimum = curve_partition(frame.minima, state.frequency_range,
                                             partition_index, output.active_partitions);
        partition.minimum = prepare_curve(
            minimum.values, minimum.domain, output.frequency_axis, output.power_axis,
            state.visible_range_only, state.interpolation_mode);
    }
    partition.current = prepare_curve(
        current.values, current.domain, output.frequency_axis, output.power_axis,
        state.visible_range_only, state.interpolation_mode);
}

void Spectrum_Control::paint_partition(Painter& painter, int partition_index,
                                       const Render_State_View& view) {
    const auto& output = impl_->prepare_buffer;
    if (!output.valid || partition_index >= output.active_partitions)
        return;
    const auto& partition = output.partitions[static_cast<std::size_t>(partition_index)];
    const auto& state = render_properties(view);
    const auto clip = painter.scoped_clip(partition.clip);
    paint_curve(painter, partition.maximum, state.max_pen, state.max_brush);
    paint_curve(painter, partition.minimum, state.min_pen, state.min_brush);
    paint_curve(painter, partition.current, state.current_pen, state.current_brush);
}

void Spectrum_Control::paint_background(Painter& painter,
                                        const Render_State_View& view) {
    const auto& output = impl_->prepare_buffer;
    if (!output.valid || output.sweep_region.empty())
        return;
    painter.rect(output.sweep_region, Pen{.style = Line_Style::None},
                 render_properties(view).sweep_region_brush);
}

void Spectrum_Control::paint_overlay(Painter& painter,
                                     const Render_State_View& view) {
    const auto& output = impl_->prepare_buffer;
    if (!output.valid)
        return;
    const auto& paint_state = render_properties(view);
    for (const auto& marker : output.markers) {
        const Pen& pen = marker.style == Spectrum_Marker_Style::middle
            ? paint_state.middle_frequency_pen
            : marker.style == Spectrum_Marker_Style::selected
                ? paint_state.selected_marker_pen
                : paint_state.marker_pen;
        if (pen.enabled())
            painter.line(marker.first, marker.second, pen);
    }
    for (const auto& extreme : output.extremes) {
        const Pen& pen = extreme.maximum ? paint_state.max_pen : paint_state.min_pen;
        painter.circle(extreme.point, 3.0, pen,
                       Brush{pen.color, Brush_Style::Solid});
    }
    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 Spectrum_Control::render_frame_completed(
    std::uint64_t target_interval_ns) {
    const auto& output = impl_->prepare_buffer;
    if (!output.valid || !is_visible())
        return;
    const auto& state = output.properties;
    if (impl_->partitioner.finish(
            state.partition_mode, output.active_partitions, target_interval_ns,
            static_cast<int>(Scene_Base::task_executor_worker_count()),
            output.work_size, 128))
        render_graph_changed();
}

void Spectrum_Control::prepare_frame(const Prepare_Render_Context& context) {
    prepare_render_frame(context.frame.render_state, 1);
    prepare_partition(0);
}

void Spectrum_Control::paint(Painter& painter,
                             const Paint_Render_Context& context) {
    paint_background(painter, context.frame.render_state);
    paint_partition(painter, 0, context.frame.render_state);
    paint_overlay(painter, context.frame.render_state);
}
}