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

#pragma once

#include "JFJochException.h"
#include "Coord.h"
#include "DetectorOrientation.h"
#include "GoniometerAxis.h"

// The two directions of the PONI convention, as pure functions, so the conversion can be exercised
// on its own. rot_matrix = Rz(-rot3) * Rx(-rot2) * Ry(+rot1) in the internal frame (x = column,
// y = row downward, z = beam); its columns are the lab directions of a +1 column step, a +1 row step
// and the sample->PONI vector.
RotMatrix PoniRotMatrix(float rot1, float rot2, float rot3);

// The inverse. rot2 comes back in [-pi/2, pi/2] and rot1, rot3 in (-pi, pi], which is the canonical
// branch: on it the round trip is the identity. At rot2 = +-pi/2 only rot1 +- rot3 is determined, and
// the convention is to put it all into rot1 and leave rot3 at zero.
void PoniAnglesFromMatrix(const RotMatrix &rot_matrix, float &rot1, float &rot2, float &rot3);

class DiffractionGeometry {
    float beam_x_pxl = 0.0;
    float beam_y_pxl = 0.0;
    float det_distance_mm = 100.0;
    float pixel_size_mm = 0.075;
    float wavelength_A = 1.0;
    float poni_rot_1 = 0.0f;
    float poni_rot_2 = 0.0f;
    float poni_rot_3 = 0.0f;
    DetectorOrientation orientation;
    // The full detector orientation: the PONI rotation composed with the discrete image orientation.
    // Its columns are the fast, slow and normal axes. Orthogonal, but improper when the image is
    // mirrored, so transpose() is still its inverse.
    RotMatrix det_matrix;

    std::optional<GoniometerAxis> axis;
    void UpdateDetectorMatrix();
public:
    DiffractionGeometry &BeamX_pxl(float input);
    DiffractionGeometry &BeamY_pxl(float input);
    DiffractionGeometry &DetectorDistance_mm(float input);
    DiffractionGeometry &PixelSize_mm(float input);
    DiffractionGeometry &Wavelength_A(float input);
    DiffractionGeometry &PoniRot1_rad(float input);
    DiffractionGeometry &PoniRot2_rad(float input);
    DiffractionGeometry &PoniRot3_rad(float input);
    DiffractionGeometry &Orientation(const DetectorOrientation &input);
    // Sets the detector plane from its two axis vectors (unit, orthogonal). The discrete orientation
    // is left as it is - it says how the image is stored, which two vectors cannot - and the PONI
    // angles are re-derived so that the two views stay in step.
    DiffractionGeometry &DetectorAxes(const Coord &fast, const Coord &slow);
    DiffractionGeometry &Rotation(const std::optional<GoniometerAxis> &input);

    [[nodiscard]] float GetBeamX_pxl() const;
    [[nodiscard]] float GetBeamY_pxl() const;
    [[nodiscard]] float GetDetectorDistance_mm() const;
    [[nodiscard]] float GetPixelSize_mm() const;
    [[nodiscard]] float GetWavelength_A() const;
    [[nodiscard]] Coord GetScatteringVector() const;
    [[nodiscard]] float GetPoniRot1_rad() const;
    [[nodiscard]] float GetPoniRot2_rad() const;
    [[nodiscard]] float GetPoniRot3_rad() const;
    [[nodiscard]] DetectorOrientation GetOrientation() const;
    [[nodiscard]] Coord GetFastAxis() const;     // lab direction of a +1 column step
    [[nodiscard]] Coord GetSlowAxis() const;     // lab direction of a +1 row step
    [[nodiscard]] Coord GetNormalAxis() const;   // sample -> PONI direction
    [[nodiscard]] std::pair<float, float> GetDirectBeam_pxl() const;
    [[nodiscard]] std::optional<GoniometerAxis> GetRotation() const;

    [[nodiscard]] Coord LabCoord(float x, float y) const;
    [[nodiscard]] Coord DetectorToRecip(float x, float y) const;
    [[nodiscard]] std::pair<float, float> RecipToDetector(const Coord &recip) const;
    [[nodiscard]] float TwoTheta_rad(float x, float y) const;
    [[nodiscard]] float Phi_rad(float x, float y) const;
    [[nodiscard]] float PxlToRes(float x, float y) const;
    [[nodiscard]] float PxlToQ(float x, float y) const;
    [[nodiscard]] float PxlToRes(float dist_pxl) const;
    [[nodiscard]] float ResToPxl(float d_A) const;
    [[nodiscard]] Coord ResToPxl(float d_A, float phi) const;
    [[nodiscard]] float DistFromEwaldSphere(const Coord& recip) const;
    [[nodiscard]] float CalcAzIntSolidAngleCorr(float x, float y) const;
    [[nodiscard]] float CalcAzIntPolarizationCorr(float x, float y, float coeff) const;
    [[nodiscard]] std::pair<float, float> ResPhiToPxl(float d_A, float phi_rad) const;
    [[nodiscard]] Coord ProjectToEwaldSphere(const Coord &p0) const;

    // eq. 18 in https://journals.iucr.org/d/issues/2014/08/00/dz5332/index.html
    [[nodiscard]] float AngleFromEwaldSphere_deg(const Coord &p0) const;

    [[nodiscard]] const RotMatrix& GetDetectorMatrix() const;
};