#include "Plottables.h"

#include "../plot/Plot_Core.h"
#include "../render/Blend2D_Cache.h"

#include <algorithm>
#include <cmath>
#include <iomanip>
#include <optional>
#include <sstream>

namespace renderive {
namespace {

RectF mapped_rect(const Axis_Transform& horizontal, const Axis_Transform& vertical,
                  Range horizontal_range, Range vertical_range) {
    const double x1 = horizontal.coord_to_pixel(horizontal_range.origin);
    const double x2 = horizontal.coord_to_pixel(horizontal_range.target);
    const double y1 = vertical.coord_to_pixel(vertical_range.origin);
    const double y2 = vertical.coord_to_pixel(vertical_range.target);
    return {std::min(x1, x2), std::min(y1, y2), std::abs(x2 - x1), std::abs(y2 - y1)};
}

double normalized_value(double value, Range range) {
    if (range.length() == 0.0)
        return 0.0;
    return std::clamp((value - range.origin) / range.length(), 0.0, 1.0);
}

struct Frequency_Columns {
    int first{};
    int last{};
    Range range;
};

std::optional<Frequency_Columns> frequency_columns(Range data_range,
                                                   Range visible_range,
                                                   int column_count,
                                                   bool visible_only) {
    if (column_count <= 0)
        return std::nullopt;
    if (!visible_only || column_count == 1 || data_range.length() == 0.0)
        return Frequency_Columns{0, column_count - 1, data_range};

    const auto [data_low, data_high] = std::minmax(data_range.origin, data_range.target);
    const auto [visible_low, visible_high] = std::minmax(visible_range.origin,
                                                        visible_range.target);
    const double clipped_low = std::max(data_low, visible_low);
    const double clipped_high = std::min(data_high, visible_high);
    if (clipped_low > clipped_high)
        return std::nullopt;

    const auto position = [data_range, column_count](double coordinate) {
        return (coordinate - data_range.origin) / data_range.length() *
               static_cast<double>(column_count - 1);
    };
    const auto [position_low, position_high] =
        std::minmax(position(clipped_low), position(clipped_high));
    int first = std::clamp(static_cast<int>(std::floor(position_low)),
                           0, column_count - 1);
    int last = std::clamp(static_cast<int>(std::ceil(position_high)),
                          first, column_count - 1);
    if (first == last) {
        if (last + 1 < column_count)
            ++last;
        else if (first > 0)
            --first;
    }
    const auto coordinate = [data_range, column_count](int index) {
        return data_range.origin + data_range.length() * static_cast<double>(index) /
                                       static_cast<double>(column_count - 1);
    };
    return Frequency_Columns{first, last, {coordinate(first), coordinate(last)}};
}

} // namespace

Waterfall::Waterfall(Plot_Core& plot,
                     std::shared_ptr<Frequency_Axis> frequency_axis,
                     std::shared_ptr<Time_Axis> time_axis)
    : Renderable(plot) {
    state_.frequency_axis = std::move(frequency_axis);
    state_.time_axis = std::move(time_axis);
}

void Waterfall::append_row(int tick, std::span<const double> values) {
    std::shared_ptr<Time_Axis> axis;
    {
        std::lock_guard lock(mutex_);
        state_.rows.push_back(Row{tick, {values.begin(), values.end()}});
        if (state_.bin_count <= 0)
            state_.bin_count = static_cast<int>(values.size());
        axis = state_.time_axis;
    }
    const std::size_t limit = axis
                                  ? static_cast<std::size_t>(std::max(2, axis->visible_time_point_count()))
                                  : 256;
    {
        std::lock_guard lock(mutex_);
        while (state_.rows.size() > limit)
            state_.rows.pop_front();
    }
    changed();
}

void Waterfall::append_row(int tick, std::pmr::vector<double>&& values) {
    append_row(tick, std::span<const double>(values.data(), values.size()));
}

void Waterfall::append_row(Time_Of_Day time, std::span<const double> values) {
    auto axis = time_axis();
    if (!axis)
        return;
    append_row(axis->append_time(time), values);
}

void Waterfall::append_row(Time_Of_Day time, std::pmr::vector<double>&& values) {
    append_row(time, std::span<const double>(values.data(), values.size()));
}

std::shared_ptr<Frequency_Axis> Waterfall::frequency_axis() const { std::lock_guard lock(mutex_); return state_.frequency_axis; }
std::shared_ptr<Time_Axis> Waterfall::time_axis() const { std::lock_guard lock(mutex_); return state_.time_axis; }

#define RENDERIVE_WATERFALL_PROPERTY(Type, Method, Field) \
    Type Waterfall::Method() const { std::lock_guard lock(mutex_); return state_.Field; } \
    void Waterfall::set_##Method(Type value) { { std::lock_guard lock(mutex_); state_.Field = std::move(value); } changed(); }

RENDERIVE_WATERFALL_PROPERTY(Range, frequency_range, frequency_range)
RENDERIVE_WATERFALL_PROPERTY(Range, power_range, power_range)
RENDERIVE_WATERFALL_PROPERTY(bool, visible_range_only, visible_only)
RENDERIVE_WATERFALL_PROPERTY(Image_Interpolation_Mode, interpolation_mode, interpolation)

#undef RENDERIVE_WATERFALL_PROPERTY

int Waterfall::frequency_bin_count() const {
    std::lock_guard lock(mutex_);
    return state_.bin_count;
}

void Waterfall::set_frequency_bin_count(int value) {
    {
        std::lock_guard lock(mutex_);
        state_.bin_count = std::max(0, value);
    }
    changed();
}

std::size_t Waterfall::row_count() const {
    std::lock_guard lock(mutex_);
    return state_.rows.size();
}

std::size_t Waterfall::stored_point_count() const {
    std::lock_guard lock(mutex_);
    std::size_t count{};
    for (const auto& row : state_.rows)
        count += row.values.size();
    return count;
}

std::size_t Waterfall::rendered_cell_count() const {
    const State state = state_snapshot();
    if (!state.frequency_axis || state.rows.empty())
        return 0;
    const int source_width = std::min(
        state.bin_count,
        static_cast<int>(std::min_element(
            state.rows.begin(), state.rows.end(),
            [](const Row& left, const Row& right) {
                return left.values.size() < right.values.size();
            })->values.size()));
    if (source_width <= 0)
        return 0;
    const auto columns = frequency_columns(state.frequency_range,
                                           state.frequency_axis->coord_range(),
                                           source_width, state.visible_only);
    return columns ? static_cast<std::size_t>(columns->last - columns->first + 1) *
                         state.rows.size()
                   : 0;
}

Waterfall::State Waterfall::state_snapshot() const { std::lock_guard lock(mutex_); return state_; }

void Waterfall::handle_event(const Event& event) { update_hover(event); }

void Waterfall::paint(detail::Painter& painter) {
    const State state = state_snapshot();
    if (!state.frequency_axis || !state.time_axis || state.rows.empty())
        return;
    const int source_width = std::min(state.bin_count,
                                      static_cast<int>(std::min_element(
                                          state.rows.begin(), state.rows.end(),
                                          [](const Row& left, const Row& right) {
                                              return left.values.size() < right.values.size();
                                          })->values.size()));
    const int height = static_cast<int>(state.rows.size());
    if (source_width <= 0 || height <= 0)
        return;
    const Axis_Transform horizontal = state.frequency_axis->transform();
    const auto columns = frequency_columns(state.frequency_range,
                                           horizontal.coordinate_range,
                                           source_width,
                                           state.visible_only);
    if (!columns)
        return;
    const int width = columns->last - columns->first + 1;
    std::vector<Pixel> pixels(static_cast<std::size_t>(width) * height);
    for (int y = 0; y < height; ++y) {
        const auto& row = state.rows[static_cast<std::size_t>(y)].values;
        for (int x = 0; x < width; ++x) {
            pixels[static_cast<std::size_t>(y) * width + x] =
                state.color_map.at_normalized(normalized_value(
                    row[static_cast<std::size_t>(columns->first + x)],
                    state.power_range));
        }
    }
    const Axis_Transform vertical = state.time_axis->transform();
    const Range time_range{static_cast<double>(state.rows.front().tick),
                           static_cast<double>(state.rows.back().tick)};
    RectF target = mapped_rect(horizontal, vertical, columns->range, time_range);
    if (target.height < 1.0)
        target.height = std::max(1.0, static_cast<double>(state.time_axis->pixel_length()));
    painter.heatmap(target, width, height, pixels, state.interpolation);

    const Hover_Tooltip_Snapshot tooltip = tooltip_snapshot();
    if (tooltip.enabled && tooltip.active && target.contains(tooltip.position)) {
        const double frequency = horizontal.pixel_to_coord(tooltip.position.x);
        std::ostringstream text;
        text << std::fixed << std::setprecision(2) << frequency << " Hz";
        const RectF box{tooltip.position.x + 8.0, tooltip.position.y + 8.0, 110.0, 24.0};
        painter.rect(box, Pen{tooltip.text_pen.color}, tooltip.background);
        painter.text({box.x + 4.0, box.y + 3.0}, text.str(), tooltip.font, tooltip.text_pen);
    }
}

Waterfall::Builder::Builder(std::shared_ptr<Renderable> parent,
                            std::shared_ptr<Frequency_Axis> frequency_axis,
                            std::shared_ptr<Time_Axis> time_axis)
    : parent_(std::move(parent)), frequency_axis_(std::move(frequency_axis)),
      time_axis_(std::move(time_axis)) {}

std::shared_ptr<Waterfall> Waterfall::Builder::build() {
    if (!parent_ || !frequency_axis_ || !time_axis_)
        return {};
    auto result = parent_->plot().make_renderable<Waterfall>(parent_, frequency_axis_, time_axis_);
    result->set_frequency_range(frequency_range_);
    result->set_power_range(power_range_);
    result->set_frequency_bin_count(bin_count_);
    result->set_visible_range_only(visible_only_);
    result->set_interpolation_mode(interpolation_);
    {
        std::lock_guard lock(result->mutex_);
        result->state_.color_map = color_map_;
    }
    return result;
}

Frequency_Trace::Frequency_Trace(Plot_Core& plot,
                                 std::shared_ptr<Time_Axis> time_axis,
                                 std::shared_ptr<Axis> value_axis)
    : Renderable(plot), time_axis_(std::move(time_axis)), value_axis_(std::move(value_axis)) {}

void Frequency_Trace::append_sample(int tick, double value) {
    int limit = 512;
    {
        std::lock_guard lock(mutex_);
        samples_.emplace_back(tick, value);
        if (time_axis_)
            limit = std::max(2, time_axis_->visible_time_point_count());
        while (samples_.size() > static_cast<std::size_t>(limit))
            samples_.pop_front();
    }
    changed();
}

void Frequency_Trace::append_sample(Time_Of_Day time, double value) {
    auto axis = time_axis();
    if (axis)
        append_sample(axis->append_time(time), value);
}

std::shared_ptr<Time_Axis> Frequency_Trace::time_axis() const { std::lock_guard lock(mutex_); return time_axis_; }
std::shared_ptr<Axis> Frequency_Trace::value_axis() const { std::lock_guard lock(mutex_); return value_axis_; }
std::size_t Frequency_Trace::sample_count() const { std::lock_guard lock(mutex_); return samples_.size(); }
std::size_t Frequency_Trace::rendered_point_count() const {
    std::lock_guard lock(mutex_);
    return time_axis_ && value_axis_ && samples_.size() >= 2 ? samples_.size() : 0;
}

void Frequency_Trace::paint(detail::Painter& painter) {
    std::shared_ptr<Time_Axis> time_axis;
    std::shared_ptr<Axis> value_axis;
    std::deque<std::pair<int, double>> samples;
    Pen pen;
    {
        std::lock_guard lock(mutex_);
        time_axis = time_axis_;
        value_axis = value_axis_;
        samples = samples_;
        pen = pen_;
    }
    if (!time_axis || !value_axis || samples.size() < 2)
        return;
    const Axis_Transform x = time_axis->transform();
    const Axis_Transform y = value_axis->transform();
    std::vector<PointF> points;
    points.reserve(samples.size());
    for (const auto& [tick, value] : samples)
        points.push_back({x.coord_to_pixel(tick), y.coord_to_pixel(value)});
    painter.polyline(points, pen);
}

Frequency_Trace::Builder::Builder(std::shared_ptr<Renderable> parent,
                                  std::shared_ptr<Time_Axis> time_axis,
                                  std::shared_ptr<Axis> value_axis)
    : parent_(std::move(parent)), time_axis_(std::move(time_axis)),
      value_axis_(std::move(value_axis)) {}

std::shared_ptr<Frequency_Trace> Frequency_Trace::Builder::build() {
    if (!parent_ || !time_axis_ || !value_axis_)
        return {};
    auto result = parent_->plot().make_renderable<Frequency_Trace>(parent_, time_axis_, value_axis_);
    { std::lock_guard lock(result->mutex_); result->pen_ = pen_; }
    return result;
}

Afterglow::Afterglow(Plot_Core& plot,
                     std::shared_ptr<Frequency_Axis> frequency_axis,
                     std::shared_ptr<Axis> power_axis)
    : Renderable(plot) {
    state_.frequency_axis = std::move(frequency_axis);
    state_.power_axis = std::move(power_axis);
}

std::shared_ptr<Frequency_Axis> Afterglow::frequency_axis() const { std::lock_guard lock(mutex_); return state_.frequency_axis; }
std::shared_ptr<Axis> Afterglow::power_axis() const { std::lock_guard lock(mutex_); return state_.power_axis; }

#define RENDERIVE_AFTERGLOW_PROPERTY(Type, Method, Field) \
    Type Afterglow::Method() const { std::lock_guard lock(mutex_); return state_.Field; } \
    void Afterglow::set_##Method(Type value) { { std::lock_guard lock(mutex_); state_.Field = std::move(value); } changed(); }

RENDERIVE_AFTERGLOW_PROPERTY(Range, frequency_range, frequency_range)
RENDERIVE_AFTERGLOW_PROPERTY(Range, power_range, power_range)
RENDERIVE_AFTERGLOW_PROPERTY(bool, interpolate, interpolate)

#undef RENDERIVE_AFTERGLOW_PROPERTY

int Afterglow::frequency_point_size() const {
    std::lock_guard lock(mutex_);
    return state_.frequency_count;
}

void Afterglow::set_frequency_point_size(int value) {
    { std::lock_guard lock(mutex_); state_.frequency_count = std::max(0, value); }
    changed();
}

int Afterglow::power_point_size() const {
    std::lock_guard lock(mutex_);
    return state_.power_count;
}

void Afterglow::set_power_point_size(int value) {
    { std::lock_guard lock(mutex_); state_.power_count = std::max(0, value); }
    changed();
}

double Afterglow::attenuation_rate() const {
    std::lock_guard lock(mutex_);
    return state_.attenuation;
}

void Afterglow::set_attenuation_rate(double value) {
    if (!std::isfinite(value))
        value = 0.0;
    { std::lock_guard lock(mutex_); state_.attenuation = std::clamp(value, 0.0, 1.0); }
    changed();
}

std::size_t Afterglow::history_count() const {
    std::lock_guard lock(mutex_);
    return state_.history.size();
}

std::size_t Afterglow::latest_spectrum_point_count() const {
    std::lock_guard lock(mutex_);
    return state_.history.empty() ? 0 : state_.history.back().size();
}

std::size_t Afterglow::rendered_cell_count() const {
    const State state = state_snapshot();
    if (!state.power_axis || state.history.empty())
        return 0;
    const int width = std::min(state.frequency_count,
                               static_cast<int>(state.history.back().size()));
    const int height = state.power_count > 0
                           ? state.power_count
                           : std::max(1, static_cast<int>(state.power_axis->pixel_length()));
    return width > 0 && height > 0
               ? static_cast<std::size_t>(width) * static_cast<std::size_t>(height)
               : 0;
}

void Afterglow::append_spectrum(std::span<const double> values) {
    {
        std::lock_guard lock(mutex_);
        state_.history.emplace_back(values.begin(), values.end());
        if (state_.frequency_count <= 0)
            state_.frequency_count = static_cast<int>(values.size());
        while (state_.history.size() > 64)
            state_.history.pop_front();
    }
    changed();
}

void Afterglow::append_spectrum(std::pmr::vector<double>&& values) {
    append_spectrum(std::span<const double>(values.data(), values.size()));
}

Afterglow::State Afterglow::state_snapshot() const { std::lock_guard lock(mutex_); return state_; }

void Afterglow::paint(detail::Painter& painter) {
    const State state = state_snapshot();
    if (!state.frequency_axis || !state.power_axis || state.history.empty())
        return;
    const int width = std::min(state.frequency_count,
                               static_cast<int>(state.history.back().size()));
    const int height = state.power_count > 0
                           ? state.power_count
                           : std::max(1, static_cast<int>(state.power_axis->pixel_length()));
    if (width <= 0 || height <= 0)
        return;
    std::vector<double> intensity(static_cast<std::size_t>(width) * height);
    double weight = 1.0;
    const double decay = 1.0 - state.attenuation;
    for (auto iterator = state.history.rbegin(); iterator != state.history.rend(); ++iterator) {
        const auto& spectrum = *iterator;
        const int count = std::min(width, static_cast<int>(spectrum.size()));
        for (int x = 0; x < count; ++x) {
            const double normalized = normalized_value(spectrum[static_cast<std::size_t>(x)],
                                                       state.power_range);
            const int y = std::clamp(height - 1 - static_cast<int>(normalized * (height - 1)),
                                     0, height - 1);
            intensity[static_cast<std::size_t>(y) * width + x] += weight;
            if (state.interpolate && y + 1 < height)
                intensity[static_cast<std::size_t>(y + 1) * width + x] += weight * 0.35;
        }
        weight *= decay;
        if (weight < 0.01)
            break;
    }
    const double maximum = std::max(1.0, *std::max_element(intensity.begin(), intensity.end()));
    std::vector<Pixel> pixels(intensity.size());
    for (std::size_t index = 0; index < pixels.size(); ++index)
        pixels[index] = state.color_map.at_normalized(intensity[index] / maximum);
    const RectF target = mapped_rect(state.frequency_axis->transform(), state.power_axis->transform(),
                                     state.frequency_range, state.power_range);
    painter.heatmap(target, width, height, pixels, Image_Interpolation_Mode::Bilinear);
}

Afterglow::Builder::Builder(std::shared_ptr<Renderable> parent,
                            std::shared_ptr<Frequency_Axis> frequency_axis,
                            std::shared_ptr<Axis> power_axis)
    : parent_(std::move(parent)), frequency_axis_(std::move(frequency_axis)),
      power_axis_(std::move(power_axis)) {}

std::shared_ptr<Afterglow> Afterglow::Builder::build() {
    if (!parent_ || !frequency_axis_ || !power_axis_)
        return {};
    auto result = parent_->plot().make_renderable<Afterglow>(parent_, frequency_axis_, power_axis_);
    result->set_frequency_range(frequency_range_);
    result->set_power_range(power_range_);
    result->set_frequency_point_size(frequency_count_);
    result->set_power_point_size(power_count_);
    result->set_interpolate(interpolate_);
    result->set_attenuation_rate(attenuation_);
    {
        std::lock_guard lock(result->mutex_);
        result->state_.color_map = color_map_;
    }
    return result;
}

} // namespace renderive