//
// Created by Lenn on 2025/10/17.
//

#ifndef TOUCHSENSOR_HEATMAP_IMPL_HH
#define TOUCHSENSOR_HEATMAP_IMPL_HH

#include "heatmap.hh"
#include "creeper-qt/utility/theme/theme.hh"
#include "creeper-qt/widget/sliders.hh"
#include "qcustomplot/qcustomplot.h"
#include <algorithm>
#include <memory>
#include <optional>
#include <qcolor.h>
#include <qdebug.h>
#include <qfont.h>
#include <vector>
using namespace creeper::plot_widget::internal;

struct BasicPlot::Impl {
    explicit Impl(BasicPlot& self) noexcept
        : matrix_size(QSize{3, 4})
        , color_min(0.0)
        , color_max(800.0)
        , show_labels(false)
        , initialized(false)
        , self{self} { }

public:
    std::optional<creeper::ColorScheme> scheme;
    auto set_xlabel_text(const QString& text) -> void {
        xlabel = text;
        if (initialized) {
            self.xAxis->setLabel(text);
            self.replot();
        }
    }

    auto set_ylabel_text(const QString& text) -> void {
        ylabel = text;
        if (initialized) {
            self.yAxis->setLabel(text);
            self.replot();
        }
    }

    auto set_matrix_size(const QSize& size) -> void {
        matrix_size = size;
        const int expected = std::max(0, matrix_size.width() * matrix_size.height());
        last_values.assign(static_cast<std::size_t>(expected), 0.0);
        if (initialized) {
            reset_plot();
            if (!data_points.isEmpty()) {
                set_data(data_points);
            } else if (show_labels) {
                sync_labels(last_values);
                self.replot();
            }
        }
    }

    auto load_theme_manager(ThemeManager& mgr) -> void {
        mgr.append_handler(&self, [this](const ThemeManager& mgr) {
            scheme = mgr.color_scheme();
            apply_color_scheme();
            if (initialized) {
                self.replot();
            }
        });
    }

    auto set_color_gradient_range(const double& min, const double& max) -> void {
        if (initialized && color_map) {
            color_map->setDataRange(QCPRange(min, max));
            self.replot();
        }
        color_min = min;
        color_max = max;
    }

    auto set_data(const QVector<PointData>& data) -> void {
        data_points = data;
        if (initialized && color_map) {
            update_plot_data();
        }
    }

    auto set_labels_visible(bool visible) -> void {
        show_labels = visible;
        if (initialized) {
            sync_labels(last_values);
            self.replot();
        }
    }

    auto labels_visible() const -> bool {
        return show_labels;
    }

    auto initialize_plot() -> void {
        if (initialized) return;

        color_map = new QCPColorMap(self.xAxis, self.yAxis);
        auto* cpmp = color_map;
        cpmp->data()->setSize(matrix_size.width(), matrix_size.height());
        cpmp->data()->setRange(QCPRange(0.5, matrix_size.width() - 0.5), 
                             QCPRange(0.5, matrix_size.height() - 0.5));
        
        QSharedPointer<QCPAxisTickerText> xticker(new QCPAxisTickerText);
        QSharedPointer<QCPAxisTickerText> yticker(new QCPAxisTickerText);
        xticker->setSubTickCount(1);
        yticker->setSubTickCount(1);
        self.xAxis->setVisible(false);
        self.yAxis->setVisible(false);
        self.xAxis->setTicker(xticker);
        self.yAxis->setTicker(yticker);

        self.xAxis->grid()->setPen(Qt::NoPen);
        self.yAxis->grid()->setPen(Qt::NoPen);
        self.xAxis->grid()->setSubGridVisible(true);
        self.yAxis->grid()->setSubGridVisible(true);
        self.xAxis->setSubTicks(true);
        self.yAxis->setSubTicks(true);
        self.xAxis->setTickLength(0);
        self.yAxis->setTickLength(0);
        self.xAxis->setSubTickLength(6);
        self.yAxis->setSubTickLength(6);
        

        self.xAxis->setRange(0, matrix_size.width());
        self.yAxis->setRange(0, matrix_size.height());

        if (!xlabel.isEmpty()) self.xAxis->setLabel(xlabel);
        if (!ylabel.isEmpty()) self.yAxis->setLabel(ylabel);
        
        if (!color_scale) {
            color_scale = new QCPColorScale(&self);
            color_scale->setType(QCPAxis::atBottom);
        }
        if (self.plotLayout()) {
            auto existing = self.plotLayout()->element(1, 0);
            if (existing && existing != color_scale) {
                self.plotLayout()->take(existing);
            }
            if (!existing || existing != color_scale) {
                self.plotLayout()->addElement(1, 0, color_scale);
            }
        }
        cpmp->setColorScale(color_scale);

        QCPColorGradient gradient;
        gradient.setColorStopAt(0.0, QColor(0, 176, 80));
        gradient.setColorStopAt(0.5, QColor(255, 214, 10));
        gradient.setColorStopAt(1.0, QColor(204, 0, 0));
        cpmp->setGradient(gradient);

        cpmp->setDataRange(QCPRange(color_min, color_max));
        cpmp->setInterpolate(false);
        
        QCPMarginGroup *margin_group = new QCPMarginGroup(&self);
        self.axisRect()->setMarginGroup(QCP::msLeft | QCP::msRight, margin_group);
        color_scale->setMarginGroup(QCP::msLeft | QCP::msRight, margin_group);

        initialized = true;
        apply_color_scheme();

        if (!data_points.isEmpty()) {
            update_plot_data();
        }
    }

    auto reset_plot() -> void {
        // 清除所有绘图元素
        clear_labels();
        self.clearPlottables();
        self.clearGraphs();
        self.clearItems();
        self.clearFocus();
        color_map = nullptr;
        
        // 重新初始化
        initialized = false;
        initialize_plot();
    }
    auto update_plot_data() -> void {
        if (!initialized || !color_map) return;

        // ensure_labels();

        const int width = matrix_size.width();
        const int height = matrix_size.height();
        const int expected = width * height;
        std::vector<double> values(static_cast<std::size_t>(expected), 0.0);

        // 设置新数据
        for (const auto& item : data_points) {
            if (item.x >= 0 && item.x < matrix_size.width() && 
                item.y >= 0 && item.y < matrix_size.height()) {
                color_map->data()->setCell(item.x, item.y, item.z);
                const int idx = static_cast<int>(item.y) * width + static_cast<int>(item.x);
                if (idx >= 0 && idx < expected) {
                    values[static_cast<std::size_t>(idx)] = item.z;
                }
            }
        }

        last_values = values;
        sync_labels(last_values);
        
        // 重绘
        self.replot();
    }

    auto is_plot_initialized() const -> bool {
        return initialized;
    }

    auto get_matrix_size() const -> QSize {
        return matrix_size;
    }

private:
    QString xlabel;
    QString ylabel;
    QSize matrix_size;
    QVector<PointData> data_points;
    std::vector<double> last_values;
    double color_min = 0.0;
    double color_max = 800.0;
    bool show_labels = false;
    bool initialized;
    BasicPlot& self;
    QCPColorScale* color_scale = nullptr;
    QCPColorMap* color_map = nullptr;
    QVector<QCPItemText*> cell_labels;
    QColor label_text_color = QColor(0, 0, 0);

    void apply_color_scheme() {
        QColor text_color = QColor(30, 30, 30);
        if (scheme.has_value()) {
            if (scheme->on_surface.isValid()) {
                text_color = scheme->on_surface;
            }
        }
        label_text_color = text_color;

        const auto pen = QPen(text_color);

        self.xAxis->setTickLabelColor(text_color);
        self.yAxis->setTickLabelColor(text_color);
        self.xAxis->setLabelColor(text_color);
        self.yAxis->setLabelColor(text_color);
        self.xAxis->setBasePen(pen);
        self.yAxis->setBasePen(pen);
        self.xAxis->setTickPen(pen);
        self.yAxis->setTickPen(pen);
        if (color_scale && color_scale->axis()) {
            color_scale->axis()->setTickLabelColor(text_color);
            color_scale->axis()->setLabelColor(text_color);
            color_scale->axis()->setBasePen(pen);
            color_scale->axis()->setTickPen(pen);
        }

        // 已有标签更新
        for (auto* label : cell_labels) {
            if (!label) continue;
            label->setColor(label_text_color);
        }
    }

    void clear_labels() {
        for (auto* label : cell_labels) {
            if (label) {
                self.removeItem(label);
            }
        }
        cell_labels.clear();
    }

    void ensure_labels() {
        const int width = matrix_size.width();
        const int height = matrix_size.height();
        const int expected = width * height;
        if (expected <= 0) {
            clear_labels();
            return;
        }
        if (cell_labels.size() == expected) {
            return;
        }

        clear_labels();
        cell_labels.reserve(expected);
        for (int y = 0; y < height; ++y) {
            for (int x = 0; x < width; ++x) {
                auto* label = new QCPItemText(&self);
                label->position->setType(QCPItemPosition::ptPlotCoords);
                label->setClipToAxisRect(true);
                label->setClipAxisRect(self.axisRect());
                label->setPositionAlignment(Qt::AlignCenter);
                label->position->setCoords(x + 0.5, y + 0.5);
                label->setBrush(Qt::NoBrush);
                label->setPen(Qt::NoPen);
                QFont font = label->font();
                if (font.pointSize() > 0) {
                    font.setPointSize(std::max(font.pointSize() - 1, 6));
                } else {
                    font.setPointSize(8);
                }
                label->setFont(font);
                label->setColor(label_text_color);
                label->setSelectable(false);
                cell_labels.push_back(label);
            }
        }
    }

    void update_label_values(const std::vector<double>& values) {
        const int width = matrix_size.width();
        const int height = matrix_size.height();
        const int expected = width * height;
        for (int idx = 0; idx < expected && idx < cell_labels.size(); ++idx) {
            auto* label = cell_labels[idx];
            if (!label) {
                continue;
            }
            const double value = values.size() > static_cast<std::size_t>(idx)
                                     ? values[static_cast<std::size_t>(idx)]
                                     : 0.0;
            label->setText(QString::number(value, 'f', 0));
            const int x = idx % width;
            const int y = idx / width;
            label->position->setCoords(x + 0.5, y + 0.5);
        }
    }

    void sync_labels(const std::vector<double>& values) {
        if (!initialized) {
            return;
        }
        if (show_labels) {
            ensure_labels();
            update_label_values(values);
        } else if (!cell_labels.isEmpty()) {
            clear_labels();
        }
    }
};

#endif // TOUCHSENSOR_HEATMAP_IMPL_HH