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

#include "JFJochGridScanImage.h"

#include <QGraphicsItem>

JFJochGridScanImage::JFJochGridScanImage(QWidget *parent) : JFJochImage(parent) {}

void JFJochGridScanImage::clear() {
    W = 0;
    H = 0;
    this->settings = std::nullopt;
    composite = false;
    obj_map.clear();
    protein_map.clear();
    ice_map.clear();
    crystals.clear();
    if (scene())
        scene()->clear();
    resetScenePointers();
}

void JFJochGridScanImage::SetupGrid(const GridScanSettings &in_settings) {
    this->settings = in_settings;

    W = in_settings.GetGridSizeX_step();
    H = in_settings.GetGridSizeY_step();
    // Draw a cell in the proportion of the area it covers, so a scan with a coarse x step and a
    // fine y step looks like the sample region it was taken over, not like a square grid.
    pixel_aspect_ = std::fabs(in_settings.GetGridStepY_um() / in_settings.GetGridStepX_um());

    std::vector<int64_t> indices(static_cast<size_t>(W) * H);
    std::iota(indices.begin(), indices.end(), 0);
    image_index = in_settings.Rearrange(indices, -1);
}

void JFJochGridScanImage::loadData(const std::vector<float> &data, const GridScanSettings &settings, bool in_one_over_d2) {
    if (data.empty()) {
        clear();
        return;
    }
    composite = false;
    SetupGrid(settings);
    image_fp = settings.Rearrange(data, NAN);

    one_over_d2 = in_one_over_d2;
    if (one_over_d2) {
        for (int i = 0; i < image_fp.size(); i++)
            if (!std::isnan(image_fp[i]))
                image_fp[i] = 1.0f / (image_fp[i] * image_fp[i]);
    }

    float minv = std::numeric_limits<float>::infinity();
    float maxv = -std::numeric_limits<float>::infinity();
    for (float v : image_fp) {
        if (std::isfinite(v)) {
            if (v < minv) minv = v;
            if (v > maxv) maxv = v;
        }
    }
    if (!std::isfinite(minv) || !std::isfinite(maxv)) {
        minv = 0.0f;
        maxv = 1.0f;
    }
    if (!(maxv > minv)) {
        maxv = minv + 1.0f;
    }
    background = minv;
    foreground = maxv;

    RenderImage();
    Redraw();
}

void JFJochGridScanImage::loadComposite(const std::vector<float> &bkg, const std::vector<float> &protein,
                                        const std::vector<float> &ice,
                                        const std::vector<GridScanCrystal> &in_crystals,
                                        const GridScanSettings &settings) {
    if (bkg.empty() && protein.empty() && ice.empty()) {
        clear();
        return;
    }
    composite = true;
    one_over_d2 = false;
    SetupGrid(settings);

    obj_map = settings.Rearrange(grid_scan::NormaliseBackground(bkg), NAN);
    // A score that has not been measured is not a score of zero, but for the colouring the two are
    // the same thing: nothing painted over the grey.
    protein_map = settings.Rearrange(protein, 0.0f);
    ice_map = settings.Rearrange(ice, 0.0f);
    crystals = in_crystals;

    image_fp = obj_map;
    background = 0.0f;
    foreground = 1.0f;

    RenderImage();
    Redraw();
}

void JFJochGridScanImage::ColorRow(size_t y, const PixelColorMap &map, QRgb *out) const {
    if (!composite) {
        JFJochImage::ColorRow(y, map, out);
        return;
    }

    for (size_t x = 0; x < W; ++x) {
        const size_t idx = y * W + x;
        if (std::isnan(obj_map[idx])) {
            out[x] = qRgb(map.gap.r, map.gap.g, map.gap.b);
            continue;
        }
        const rgb c = grid_scan::CompositeColor(obj_map[idx], protein_map[idx], ice_map[idx]);
        out[x] = qRgb(c.r, c.g, c.b);
    }
}

QString JFJochGridScanImage::PixelLabel(int x, int y) const {
    if (!composite)
        return JFJochImage::PixelLabel(x, y);

    const size_t idx = static_cast<size_t>(y) * W + x;
    if (std::isnan(obj_map[idx]))
        return QStringLiteral("-");
    return QString("%1/%2").arg(QString::number(protein_map[idx], 'f', 2),
                                QString::number(ice_map[idx], 'f', 2));
}

void JFJochGridScanImage::mouseHover(const QPointF &pt, Qt::KeyboardModifiers modifiers) {
    // Map mouse position to image pixel if inside bounds
    if (W == 0 || H == 0 || image_index.empty())
        return;

    // Convert view coordinates to image pixel by truncation
    int x = static_cast<int>(pt.x());
    int y = static_cast<int>(pt.y());

    if (x < 0 || x >= W || y < 0 || y >= H)
        return;

    int idx = y * W + x;
    if (idx < 0 || idx >= static_cast<int>(image_index.size()))
        return;
    int64_t image_id = image_index[idx];
    if (image_id >= 0) {
        if (modifiers & Qt::ShiftModifier)
            emit imageSelected(image_id);

        if (composite) {
            emit writeStatusBar(QString("Image %1 x %2 y %3 protein %4 ice %5")
                                .arg(image_id)
                                .arg(x)
                                .arg(y)
                                .arg(QString::number(protein_map[idx], 'f', 2))
                                .arg(QString::number(ice_map[idx], 'f', 2)),
                                6000);
        } else if (one_over_d2) {
            if (std::isnan(image_fp[idx]))
                emit writeStatusBar(QString("Image %1 x %2 y %3 no resolution estimation")
                                    .arg(image_id)
                                    .arg(x)
                                    .arg(y),
                                    6000);
            else
                emit writeStatusBar(QString("Image %1 x %2 y %3 d = %4 Å")
                                    .arg(image_id)
                                    .arg(x)
                                    .arg(y)
                                    .arg(QString::number(1 / std::sqrt(image_fp[idx]), 'f', 2)),
                                    6000);
        } else {
            emit writeStatusBar(QString("Image %1 x %2 y %3 value %4")
                                .arg(image_id)
                                .arg(x)
                                .arg(y)
                                .arg(QString::number(image_fp[idx], 'f', 3)),
                                6000);
        }
    }
}

void JFJochGridScanImage::loadImage(QMouseEvent *event) {
    // Map mouse position to image pixel if inside bounds
    if (W == 0 || H == 0 || image_index.empty())
        return;

    const QPointF pt = mapToScene(event->pos());

    // Convert view coordinates to image pixel by truncation
    int x = static_cast<int>(pt.x());
    int y = static_cast<int>(pt.y());

    if (x < 0 || x >= W || y < 0 || y >= H)
        return;

    int idx = y * W + x;
    if (idx < 0 || idx >= static_cast<int>(image_index.size()))
        return;
    int64_t image_id = image_index[idx];
    if (image_id >= 0)
        emit imageSelected(image_id);
}

void JFJochGridScanImage::mouseDoubleClickEvent(QMouseEvent *event) {
    // Map mouse position to image pixel if inside bounds
    if (W == 0 || H == 0 || image_index.empty())
        return;

    const QPointF pt = mapToScene(event->pos());

    // Convert view coordinates to image pixel by truncation
    int x = static_cast<int>(pt.x());
    int y = static_cast<int>(pt.y());

    if (x < 0 || x >= W || y < 0 || y >= H)
        return;

    int idx = y * W + x;
    if (idx < 0 || idx >= static_cast<int>(image_index.size()))
        return;
    int64_t image_id = image_index[idx];
    if (image_id >= 0)
        emit imageSelected(image_id);
    JFJochImage::mouseDoubleClickEvent(event);
}

void JFJochGridScanImage::setImage(int64_t val) {
    if (settings) {
        current_image_W = settings->GetElementPosX_step(val);
        current_image_H = settings->GetElementPosY_step(val);
    } else {
        current_image_W = -1;
        current_image_H = -1;
    }
    Redraw();
}

void JFJochGridScanImage::drawCrystalFrames() {
    // White over black rather than a colour from the palette: the frame has to stay legible over
    // the pale plate, the grey loop, cyan, orange and everything in between, and no single hue
    // does that.
    QPen halo(Qt::black, 5);
    halo.setCosmetic(true);
    QPen line(Qt::white, 2);
    line.setCosmetic(true);

    const float step_x = settings->GetGridStepX_um();
    const float step_y = settings->GetGridStepY_um();

    for (const auto &c : crystals) {
        if (grid_scan::IsRoundBlob(c.major_um, c.minor_um)) {
            // No orientation to show, so no oriented box - a circle of the mean extent.
            const float r_um = 0.25f * (c.major_um + c.minor_um);
            const QRectF box(c.nx + 0.5 - r_um / step_x, c.ny + 0.5 - r_um / step_y,
                             2 * r_um / step_x, 2 * r_um / step_y);
            addOverlayItem(scene()->addEllipse(box, halo));
            addOverlayItem(scene()->addEllipse(box, line));
        } else {
            QPolygonF poly;
            for (const auto &corner : grid_scan::CrystalFrameCorners(c.nx, c.ny, c.major_um, c.minor_um,
                                                                     c.angle_deg, step_x, step_y))
                poly << QPointF(corner.x, corner.y);
            addOverlayItem(scene()->addPolygon(poly, halo));
            addOverlayItem(scene()->addPolygon(poly, line));
        }
    }
}

void JFJochGridScanImage::addCustomOverlay() {
    if (composite && settings)
        drawCrystalFrames();

    if (current_image_W < 0 || current_image_H < 0 || current_image_W >= W || current_image_H >= H)
        return;

    QPen pen(feature_color, 3);
    pen.setCosmetic(true);

    auto *rect = scene()->addRect(current_image_W, current_image_H, 1, 1, pen);
    addOverlayItem(rect);
}