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

#include "JFJochViewerReciprocalSpaceWindow.h"

#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QPushButton>
#include <QSurfaceFormat>
#include <QtMath>
#include <algorithm>

// ============================================================================
//  Shaders
// ============================================================================

// One shader handles both spots (GL_POINTS) and lines (GL_LINES).
// Per-vertex color comes straight from the VBO; MVP is a single uniform.

static const char *kVertSrc = R"(
    #version 330 core
    layout(location = 0) in vec3 aPos;
    layout(location = 1) in vec3 aColor;
    out vec3 vColor;
    uniform mat4 uMVP;
    void main() {
        vColor       = aColor;
        gl_Position  = uMVP * vec4(aPos, 1.0);
        gl_PointSize = 6.0;
    }
)";

static const char *kFragSrc = R"(
    #version 330 core
    in  vec3 vColor;
    out vec4 fragColor;
    uniform bool uPointShape;
    void main() {
        if (uPointShape) {
            // Discard corners to draw a circle instead of a square point.
            vec2 c = gl_PointCoord - vec2(0.5);
            if (dot(c, c) > 0.25) discard;
        }
        fragColor = vec4(vColor, 1.0);
    }
)";

// ============================================================================
//  ReciprocalSpaceGLView
// ============================================================================

ReciprocalSpaceGLView::ReciprocalSpaceGLView(QWidget *parent)
    : QOpenGLWidget(parent)
    , spotsVBO_(QOpenGLBuffer::VertexBuffer)
    , linesVBO_(QOpenGLBuffer::VertexBuffer)
{
    // Request OpenGL 3.3 Core Profile (Linux/NVIDIA, Mesa, macOS all support this)
    QSurfaceFormat fmt;
    fmt.setVersion(3, 3);
    fmt.setProfile(QSurfaceFormat::CoreProfile);
    fmt.setDepthBufferSize(24);
    setFormat(fmt);
    setFocusPolicy(Qt::StrongFocus);
}

ReciprocalSpaceGLView::~ReciprocalSpaceGLView() {
    makeCurrent();
    spotsVAO_.destroy();
    spotsVBO_.destroy();
    linesVAO_.destroy();
    linesVBO_.destroy();
    doneCurrent();
}

void ReciprocalSpaceGLView::initializeGL() {
    initializeOpenGLFunctions();

    glClearColor(20.f/255, 20.f/255, 25.f/255, 1.0f);
    glEnable(GL_DEPTH_TEST);
    glEnable(GL_PROGRAM_POINT_SIZE);   // lets the vertex shader set gl_PointSize

    shader_.addShaderFromSourceCode(QOpenGLShader::Vertex,   kVertSrc);
    shader_.addShaderFromSourceCode(QOpenGLShader::Fragment, kFragSrc);
    shader_.link();

    // Create empty VAO/VBOs so they are valid before the first data arrives.
    setupVAO(spotsVAO_, spotsVBO_);
    setupVAO(linesVAO_, linesVBO_);

    glReady_ = true;

    uploadBuffer(spotsVBO_, spotsVAO_, pendingSpots_);
    spotsCount_ = static_cast<int>(pendingSpots_.size());

    uploadBuffer(linesVBO_, linesVAO_, pendingLines_);
    linesCount_ = static_cast<int>(pendingLines_.size());
}

void ReciprocalSpaceGLView::resizeGL(int w, int h) {
    proj_.setToIdentity();
    proj_.perspective(45.0f, float(w) / float(h ? h : 1), 0.1f, 10000.0f);
}

QMatrix4x4 ReciprocalSpaceGLView::currentViewMatrix() const {
    QMatrix4x4 view;
    view.translate(0, 0, -zoom_);
    view.rotate(pitch_, 1, 0, 0);
    view.rotate(yaw_,   0, 1, 0);
    view.translate(-target_);
    return view;
}

QMatrix4x4 ReciprocalSpaceGLView::currentMvpMatrix() const {
    return proj_ * currentViewMatrix();
}

void ReciprocalSpaceGLView::paintGL() {
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    const QMatrix4x4 mvp = currentMvpMatrix();

    shader_.bind();
    shader_.setUniformValue("uMVP", mvp);

    // Draw axes + cell vectors
    if (linesCount_ > 0) {
        shader_.setUniformValue("uPointShape", false);
        linesVAO_.bind();
        glDrawArrays(GL_LINES, 0, linesCount_);
        linesVAO_.release();
    }

    // Draw spots
    if (spotsCount_ > 0) {
        shader_.setUniformValue("uPointShape", true);
        spotsVAO_.bind();
        glDrawArrays(GL_POINTS, 0, spotsCount_);
        spotsVAO_.release();
    }

    shader_.release();
}

// ---------------------------------------------------------------------------

void ReciprocalSpaceGLView::setSpots(const std::vector<Vertex> &spots) {
    pendingSpots_ = spots;
    spotsCount_ = static_cast<int>(pendingSpots_.size());

    if (glReady_) {
        makeCurrent();
        uploadBuffer(spotsVBO_, spotsVAO_, pendingSpots_);
        doneCurrent();
    }

    update();
}

void ReciprocalSpaceGLView::setLines(const std::vector<Vertex> &lines) {
    pendingLines_ = lines;
    linesCount_ = static_cast<int>(pendingLines_.size());

    if (glReady_) {
        makeCurrent();
        uploadBuffer(linesVBO_, linesVAO_, pendingLines_);
        doneCurrent();
    }

    update();
}

// ---------------------------------------------------------------------------

void ReciprocalSpaceGLView::setupVAO(QOpenGLVertexArrayObject &vao,
                                      QOpenGLBuffer &vbo) {
    vao.create();
    vao.bind();
    vbo.create();
    vbo.bind();

    // location 0: position (xyz)
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE,
                          sizeof(Vertex),
                          reinterpret_cast<void *>(offsetof(Vertex, x)));
    // location 1: color (rgb)
    glEnableVertexAttribArray(1);
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE,
                          sizeof(Vertex),
                          reinterpret_cast<void *>(offsetof(Vertex, r)));

    vbo.release();
    vao.release();
}

void ReciprocalSpaceGLView::uploadBuffer(QOpenGLBuffer &vbo,
                                          QOpenGLVertexArrayObject &vao,
                                          const std::vector<Vertex> &data) {
    vao.bind();
    vbo.bind();
    if (data.empty()) {
        vbo.allocate(nullptr, 0);
    } else {
        vbo.allocate(data.data(),
                     static_cast<int>(data.size() * sizeof(Vertex)));
    }
    vbo.release();
    vao.release();
}

// ---------------------------------------------------------------------------
// Mouse / wheel -> orbit camera
// ---------------------------------------------------------------------------

void ReciprocalSpaceGLView::mousePressEvent(QMouseEvent *e) {
    lastMousePos_ = e->pos();
}

void ReciprocalSpaceGLView::mouseDoubleClickEvent(QMouseEvent *e) {
    if (e->button() == Qt::LeftButton && focusNearestSpot(e->pos())) {
        e->accept();
        return;
    }

    QOpenGLWidget::mouseDoubleClickEvent(e);
}

void ReciprocalSpaceGLView::mouseMoveEvent(QMouseEvent *e) {
    const QPoint delta = e->pos() - lastMousePos_;
    lastMousePos_ = e->pos();

    if (e->buttons() & Qt::LeftButton) {
        yaw_   += delta.x() * 0.5f;
        pitch_ += delta.y() * 0.5f;
        pitch_  = qBound(-89.0f, pitch_, 89.0f);
        update();
    } else if (e->buttons() & Qt::MiddleButton) {
        const float panScale = zoom_ * 0.0015f;

        QMatrix4x4 rot;
        rot.rotate(pitch_, 1, 0, 0);
        rot.rotate(yaw_,   0, 1, 0);

        const QVector3D right = rot.inverted().mapVector(QVector3D(1, 0, 0));
        const QVector3D up    = rot.inverted().mapVector(QVector3D(0, 1, 0));

        target_ -= right * float(delta.x()) * panScale;
        target_ += up    * float(delta.y()) * panScale;
        update();
    }
}

void ReciprocalSpaceGLView::wheelEvent(QWheelEvent *e) {
    zoom_ *= (e->angleDelta().y() > 0) ? 0.9f : 1.1f;
    zoom_  = qBound(1.0f, zoom_, 5000.0f);
    update();
}

bool ReciprocalSpaceGLView::focusNearestSpot(const QPoint &screenPos) {
    if (pendingSpots_.empty() || width() <= 0 || height() <= 0)
        return false;

    const QMatrix4x4 mvp = currentMvpMatrix();

    float bestDist2 = std::numeric_limits<float>::max();
    QVector3D bestPos;
    bool found = false;

    for (const auto &v : pendingSpots_) {
        const QVector4D clip = mvp * QVector4D(v.x, v.y, v.z, 1.0f);
        if (clip.w() <= 0.0f)
            continue;

        const QVector3D ndc = clip.toVector3DAffine();
        if (ndc.z() < -1.0f || ndc.z() > 1.0f)
            continue;

        const float sx = (ndc.x() * 0.5f + 0.5f) * float(width());
        const float sy = (0.5f - ndc.y() * 0.5f) * float(height());

        const float dx = sx - float(screenPos.x());
        const float dy = sy - float(screenPos.y());
        const float dist2 = dx * dx + dy * dy;

        if (dist2 < bestDist2) {
            bestDist2 = dist2;
            bestPos = QVector3D(v.x, v.y, v.z);
            found = true;
        }
    }

    // Roughly 15 px picking radius.
    if (!found || bestDist2 > 15.0f * 15.0f)
        return false;

    target_ = bestPos;
    zoom_ = qMax(5.0f, zoom_ * 0.35f);
    update();
    return true;
}

// ============================================================================
//  JFJochViewerReciprocalSpaceWindow
// ============================================================================

JFJochViewerReciprocalSpaceWindow::JFJochViewerReciprocalSpaceWindow(QWidget *parent)
    : JFJochHelperWindow(parent)
{
    setWindowTitle("Reciprocal space");
    resize(800, 800);

    auto *central = new QWidget(this);
    auto *layout  = new QVBoxLayout(central);

    glView_ = new ReciprocalSpaceGLView(central);
    glView_->setMinimumSize(400, 400);
    layout->addWidget(glView_, 1);

    auto *controls = new QHBoxLayout();
    crystalFrameCheck = new QCheckBox("Crystal frame (angle = 0)", central);
    crystalFrameCheck->setChecked(false);
    crystalFrameCheck->setEnabled(false);
    showCellCheck = new QCheckBox("Show reciprocal cell", central);
    showCellCheck->setChecked(true);
    accumulateCheck = new QCheckBox("Accumulate", central);
    accumulateCheck->setChecked(false);
    auto *clearButton = new QPushButton("Clear accumulated", central);

    controls->addWidget(crystalFrameCheck);
    controls->addWidget(showCellCheck);
    controls->addWidget(accumulateCheck);
    controls->addStretch(1);
    controls->addWidget(clearButton);
    layout->addLayout(controls);
    setCentralWidget(central);

    connect(crystalFrameCheck, &QCheckBox::toggled,
            this, &JFJochViewerReciprocalSpaceWindow::rebuildGL);
    connect(showCellCheck, &QCheckBox::toggled,
            this, &JFJochViewerReciprocalSpaceWindow::rebuildGL);
    connect(clearButton, &QPushButton::clicked, this, [this] {
        spots_.clear();
        rebuildGL();
    });

    rebuildGL();
}

// ---------------------------------------------------------------------------
// Public slots
// ---------------------------------------------------------------------------

void JFJochViewerReciprocalSpaceWindow::datasetLoaded(
        std::shared_ptr<const JFJochReaderDataset> in_dataset)
{
    spots_.clear();
    indexed_lattice_.reset();
    has_rotation_ = false;

    if (in_dataset)
        has_rotation_ = in_dataset->experiment.GetGoniometer().has_value();

    crystalFrameCheck->setEnabled(has_rotation_);
    if (!has_rotation_)
        crystalFrameCheck->setChecked(false);

    rebuildGL();
}

void JFJochViewerReciprocalSpaceWindow::imageLoaded(std::shared_ptr<const JFJochReaderImage> image) {
    if (!accumulateCheck->isChecked()) {
        spots_.clear();
        indexed_lattice_.reset();
    }

    if (!image) {
        rebuildGL();
        return;
    }

    const auto &dataset        = image->Dataset();
    const auto  geom           = dataset.experiment.GetDiffractionGeometry();
    const auto  axis           = dataset.experiment.GetGoniometer();
    const int64_t image_number = image->ImageData().number;

    std::optional<RotMatrix> back_rot;
    if (axis) {
        const float angle_deg = axis->GetAngle_deg(static_cast<float>(image_number))
                                + axis->GetWedge_deg() / 2.0f;
        back_rot = axis->GetTransformationAngle(angle_deg);
    }

    spots_.reserve(spots_.size() + image->ImageData().spots.size());
    for (const auto &s : image->ImageData().spots) {
        AccumulatedSpot acc;
        acc.recip_lab     = s.ReciprocalCoord(geom);
        acc.recip_crystal = back_rot ? (*back_rot * acc.recip_lab) : acc.recip_lab;
        acc.indexed       = s.indexed;
        acc.ice_ring      = s.ice_ring;
        spots_.emplace_back(acc);
    }

    // Soft cap: no per-spot Qt objects here, so 200k is fine.
    constexpr size_t max_accumulated = 200000;
    if (spots_.size() > max_accumulated)
        spots_.erase(spots_.begin(),
                     spots_.begin() +
                     static_cast<std::ptrdiff_t>(spots_.size() - max_accumulated));

    if (image->ImageData().indexing_lattice)
        indexed_lattice_ = image->ImageData().indexing_lattice;

    rebuildGL();
}

void JFJochViewerReciprocalSpaceWindow::setSpotColor(QColor input) {
    if (!input.isValid()) return;
    spot_color = input;
    rebuildGL();
}

void JFJochViewerReciprocalSpaceWindow::setFeatureColor(QColor input) {
    if (!input.isValid()) return;
    indexed_color = input;
    rebuildGL();
}

// ---------------------------------------------------------------------------
// Private
// ---------------------------------------------------------------------------

QColor JFJochViewerReciprocalSpaceWindow::spotColorFor(bool indexed,
                                                        bool ice_ring) const {
    if (indexed)  return indexed_color;
    if (ice_ring) return ice_ring_color;
    return spot_color;
}

void JFJochViewerReciprocalSpaceWindow::rebuildGL() {
    const bool crystal_frame = crystalFrameCheck->isChecked() && has_rotation_;

    // --- Spots ---------------------------------------------------------------
    std::vector<ReciprocalSpaceGLView::Vertex> spotVerts;
    spotVerts.reserve(spots_.size() + 1);

    // Precompute the three possible colors as floats — avoids QColor in the loop.
    auto toF = [](const QColor &c, float &r, float &g, float &b) {
        r = float(c.redF()); g = float(c.greenF()); b = float(c.blueF());
    };
    float sr, sg, sb, ir, ig, ib, xr, xg, xb;
    toF(spot_color,     sr, sg, sb);
    toF(indexed_color,  ir, ig, ib);
    toF(ice_ring_color, xr, xg, xb);

    for (const auto &spot : spots_) {
        const Coord &c = crystal_frame ? spot.recip_crystal : spot.recip_lab;
        float r, g, b;
        if      (spot.indexed)  { r = ir; g = ig; b = ib; }
        else if (spot.ice_ring) { r = xr; g = xg; b = xb; }
        else                    { r = sr; g = sg; b = sb; }
        spotVerts.push_back({ c.x * scene_scale_,
                              c.y * scene_scale_,
                              c.z * scene_scale_,
                              r, g, b });
    }

    spotVerts.push_back({ 0, 0, 0, 1.0, 1.0, 1.0 });

    // --- Lines: axes + optional cell vectors ---------------------------------
    // Each segment = 2 vertices (GL_LINES draws pairs).
    std::vector<ReciprocalSpaceGLView::Vertex> lineVerts;

    auto addLine = [&](QVector3D a, QVector3D b, QColor col) {
        const float r = float(col.redF()),
                    g = float(col.greenF()),
                    bv = float(col.blueF());
        lineVerts.push_back({ a.x(), a.y(), a.z(), r, g, bv });
        lineVerts.push_back({ b.x(), b.y(), b.z(), r, g, bv });
    };

    // Reference axes
    const float len = 20.0f;
    addLine({0,0,0}, {len,0,0}, QColor(255, 60,  60));
    addLine({0,0,0}, {0,len,0}, QColor(60,  255, 60));
    addLine({0,0,0}, {0,0,len}, QColor(60,  60,  255));

    // Reciprocal cell vectors
    if (showCellCheck->isChecked() && indexed_lattice_) {
        auto addCell = [&](const Coord &v, QColor col) {
            addLine({0, 0, 0},
                    { v.x * scene_scale_,
                      v.y * scene_scale_,
                      v.z * scene_scale_ },
                    col);
        };
        addCell(indexed_lattice_->Astar(), Qt::red);
        addCell(indexed_lattice_->Bstar(), Qt::green);
        addCell(indexed_lattice_->Cstar(), Qt::blue);
    }

    glView_->setSpots(spotVerts);
    glView_->setLines(lineVerts);
}