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

#include "LatticeVisualizerWidget.h"

#include <QVBoxLayout>
#include <QVector3D>
#include <QQuaternion>
#include <Qt3DExtras/QCylinderMesh>
#include <Qt3DRender/QCamera>
#include <Qt3DExtras/QOrbitCameraController>
#include <Qt3DRender/QPointLight>
#include <Qt3DExtras/QForwardRenderer>
#include <Qt3DRender/QRenderSettings>

LatticeVisualizerWidget::LatticeVisualizerWidget(QWidget *parent) : QWidget(parent) {
    auto *layout = new QVBoxLayout(this);
    layout->setContentsMargins(0, 0, 0, 0);

    view = new Qt3DExtras::Qt3DWindow();
    view->defaultFrameGraph()->setClearColor(Qt::white);

    // Optimize rendering to be on demand (only when scene changes or camera moves)
    // This prevents the widget from consuming 100% of a CPU core for a static image
    view->renderSettings()->setRenderPolicy(Qt3DRender::QRenderSettings::OnDemand);

    QWidget *container = QWidget::createWindowContainer(view, this);
    container->setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Expanding);
    layout->addWidget(container);


    rootEntity = new Qt3DCore::QEntity();
    view->setRootEntity(rootEntity);

    // Camera setup
    Qt3DRender::QCamera *camera = view->camera();
    camera->lens()->setOrthographicProjection(-50.0f, 50.0f, -50.0f, 50.0f, -200.0f, 200.0f);
    camera->setPosition(QVector3D(0.0f, 0.0f, 50.0f));
    camera->setViewCenter(QVector3D(0.0f, 0.0f, 0.0f));

    auto *camController = new Qt3DExtras::QOrbitCameraController(rootEntity);
    camController->setCamera(camera);
    camController->setLinearSpeed(0.0f); // Disable translation to keep the fixed scale view
    camController->setLookSpeed(180.0f);

    // Light
    auto *lightEntity = new Qt3DCore::QEntity(rootEntity);
    auto *light = new Qt3DRender::QPointLight(lightEntity);
    light->setColor("white");
    light->setIntensity(1.0f);
    lightEntity->addComponent(light);

    // Attach light to camera position or a fixed position?
    // Attaching to a transform that follows camera is complex without parenting.
    // For simple viz, a fixed light + a light attached to camera usually works best.
    // Let's put a simple light at viewer position (approx)
    auto *lightTransform = new Qt3DCore::QTransform(lightEntity);
    lightTransform->setTranslation(QVector3D(0.0f, 0.0f, 50.0f));
    lightEntity->addComponent(lightTransform);

    // Create bars with color-blind friendly palette
    // Vector a: Vermilion (D55E00)
    createBar(0, QColor(213, 94, 0));
    // Vector b: Sky Blue (56B4E9)
    createBar(1, QColor(86, 180, 233));
    // Vector c: Bluish Green (009E73)
    createBar(2, QColor(0, 158, 115));
}

void LatticeVisualizerWidget::createBar(int index, const QColor &color) {
    bars[index].entity = new Qt3DCore::QEntity(rootEntity);

    auto *mesh = new Qt3DExtras::QCylinderMesh();
    mesh->setRadius(1.0f); // Base radius, scaled later
    mesh->setLength(1.0f); // Base length, scaled later
    mesh->setRings(20);
    mesh->setSlices(20);

    bars[index].transform = new Qt3DCore::QTransform();

    bars[index].material = new Qt3DExtras::QPhongMaterial();
    bars[index].material->setAmbient(color);
    bars[index].material->setDiffuse(color);
    bars[index].material->setSpecular(Qt::white);
    bars[index].material->setShininess(50.0f);

    bars[index].entity->addComponent(mesh);
    bars[index].entity->addComponent(bars[index].transform);
    bars[index].entity->addComponent(bars[index].material);
}

void LatticeVisualizerWidget::updateBar(int index, float x, float y, float z) {
    QVector3D vec(x, y, z);
    float length = vec.length();

    if (length < 1e-3f) {
        bars[index].entity->setEnabled(false);
        return;
    }
    bars[index].entity->setEnabled(true);

    // Cylinder is Y-up by default.
    QVector3D yAxis(0.0f, 1.0f, 0.0f);
    QQuaternion rotation = QQuaternion::rotationTo(yAxis, vec.normalized());

    // Center of cylinder is at origin, we need to move it to vec/2
    QVector3D translation = vec / 2.0f;

    // Scale: X/Z is thickness, Y is length
    float thickness = 1.0f; // 40 Angstroms thick
    bars[index].transform->setScale3D(QVector3D(thickness, length, thickness));
    bars[index].transform->setRotation(rotation);
    bars[index].transform->setTranslation(translation);
}

void LatticeVisualizerWidget::setLattice(const float matrix[9]) {
    // Assuming row-major:
    // a = (matrix[0], matrix[1], matrix[2])
    // b = (matrix[3], matrix[4], matrix[5])
    // c = (matrix[6], matrix[7], matrix[8])
    updateBar(0, matrix[0], matrix[1], matrix[2]);
    updateBar(1, matrix[3], matrix[4], matrix[5]);
    updateBar(2, matrix[6], matrix[7], matrix[8]);

    // Ensure the vertical scale is strictly enforced as requested
    Qt3DRender::QCamera *camera = view->camera();
    if (camera) {
        // Adjust aspect ratio
        float aspect = (float)width() / (float)height();
        if (std::isnan(aspect) || height() == 0) aspect = 1.0f;

        // Height is fixed +/- 200
        float h = 200.0f;
        float w = h * aspect;

        camera->setTop(h);
        camera->setBottom(-h);
        camera->setLeft(-w);
        camera->setRight(w);
    }
}

void LatticeVisualizerWidget::loadImage(std::shared_ptr<const JFJochReaderImage> image) {
    if (image && image->ImageData().indexing_lattice) {
        setLattice(image->ImageData().indexing_lattice->GetVector().data());
    } else {
        float zero[9] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
        setLattice(zero);
    }
}