// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "JFJochAzIntImageView.h"
#include <QVBoxLayout>

JFJochAzIntImageView::JFJochAzIntImageView(QWidget *parent)
    : QWidget(parent), label(new QLabel(this)) {
    label->setAlignment(Qt::AlignCenter);
    label->setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Expanding);
    auto *layout = new QVBoxLayout(this);
    layout->setContentsMargins(0, 0, 0, 0);
    layout->addWidget(label);
    setLayout(layout);
    color_scale.Select(ColorScaleEnum::Viridis);
}

void JFJochAzIntImageView::Clear() {
    img = QImage();
    label->clear();
}

// data: size = azimuthal_bins * q_bins
// Layout: q varies fastest (i % q_bins == q index)
void JFJochAzIntImageView::SetData(const std::vector<float> &data, int azimuthal_bins, int q_bins) {
    if (azimuthal_bins <= 0 || q_bins <= 0 ||
            data.size() != static_cast<size_t>(azimuthal_bins * q_bins)) {
        Clear();
        return;
            }

    float local_min = range_min, local_max = range_max;
    if (auto_range) {
        local_min = std::numeric_limits<float>::infinity();
        local_max = -std::numeric_limits<float>::infinity();
        for (float v : data) {
            if (std::isfinite(v)) {
                if (v < local_min) local_min = v;
                if (v > local_max) local_max = v;
            }
        }
        if (!std::isfinite(local_min) || !std::isfinite(local_max) || local_max <= local_min) {
            Clear();
            return;
        }
    } else {
        if (!(std::isfinite(local_min) && std::isfinite(local_max) && local_max > local_min)) {
            Clear();
            return;
        }
    }

    // Build an RGB image (q along X, azimuth along Y) using ColorScale
    img = QImage(q_bins, azimuthal_bins, QImage::Format_RGB32);
    for (int az = 0; az < azimuthal_bins; ++az) {
        QRgb* scan = reinterpret_cast<QRgb*>(img.scanLine(az));
        for (int q = 0; q < q_bins; ++q) {
            float v = data[az * q_bins + q];
            if (!std::isfinite(v)) {
                auto c = color_scale.Apply(ColorScaleSpecial::BadPixel);
                scan[q] = qRgb(c.r, c.g, c.b);
            } else {
                auto c = color_scale.Apply(v, local_min, local_max);
                scan[q] = qRgb(c.r, c.g, c.b);
            }
        }
    }

    updateLabelPixmap();

}

void JFJochAzIntImageView::SetColorScale(ColorScaleEnum map) {
    color_scale.Select(map);
    updateLabelPixmap();
}

void JFJochAzIntImageView::SetRangeAuto() {
    auto_range = true;
    updateLabelPixmap();
}

void JFJochAzIntImageView::SetRange(float min_val, float max_val) {
    range_min = min_val;
    range_max = max_val;
    auto_range = false;
    updateLabelPixmap();
}


void JFJochAzIntImageView::resizeEvent(QResizeEvent *e) {
    QWidget::resizeEvent(e);
    updateLabelPixmap();
}


void JFJochAzIntImageView::updateLabelPixmap() {
    if (img.isNull()) {
        label->clear();
        return;
    }
    // Keep aspect ratio; smooth transform for readability
    auto scaled = QPixmap::fromImage(img).scaled(label->size(), Qt::KeepAspectRatio, Qt::SmoothTransformation);
    label->setPixmap(scaled);
}