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

#include "JFJochMath.h"
#include <algorithm>
#include <cmath>

#include "DiffractionGeometry.h"
#include "RawToConvertedGeometry.h"

RotMatrix PoniRotMatrix(float rot1, float rot2, float rot3) {
    return RotMatrix(-rot3, {0,0,1})
           * RotMatrix(-rot2, {1,0,0})
           * RotMatrix(rot1, {0,1,0});
}

void PoniAnglesFromMatrix(const RotMatrix &rot_matrix, float &rot1, float &rot2, float &rot3) {
    const Coord fast = rot_matrix.Column(0);
    const Coord slow = rot_matrix.Column(1);
    const Coord normal = rot_matrix.Column(2);

    rot2 = asinf(std::clamp(-slow.z, -1.0f, 1.0f));
    if (fabsf(cosf(rot2)) < 1e-6f) {
        // Gimbal lock: only rot1 +- rot3 is determined, so put it all into rot1.
        rot1 = atan2f(normal.x, fast.x);
        rot3 = 0.0f;
    } else {
        rot1 = atan2f(-fast.z, normal.z);
        rot3 = atan2f(slow.x, slow.y);
    }
}

Coord DiffractionGeometry::LabCoord(float x, float y) const {
    Coord detectorCoord = {(x - beam_x_pxl) * pixel_size_mm ,
                           (y - beam_y_pxl) * pixel_size_mm ,
                           det_distance_mm};
    return det_matrix * detectorCoord;
}

std::pair<float, float> DiffractionGeometry::GetDirectBeam_pxl() const {
    return RecipToDetector({0,0,0});
}

Coord DiffractionGeometry::GetScatteringVector() const {
    return {0, 0, 1.0f / wavelength_A};
}

Coord DiffractionGeometry::DetectorToRecip(float x, float y) const {
    return LabCoord(x, y).Normalize() / wavelength_A - GetScatteringVector();
}

std::pair<float, float> DiffractionGeometry::RecipToDetector(const Coord &recip) const {
    auto S_unrotated = recip + GetScatteringVector();
    auto S = det_matrix.transpose() * S_unrotated;
    if (S.z <= 0)
        return {NAN, NAN};

    float coeff = det_distance_mm / (S.z * pixel_size_mm);
    float x = beam_x_pxl + S.x * coeff;
    float y = beam_y_pxl + S.y * coeff;
    return {x, y};
}

float DiffractionGeometry::TwoTheta_rad(float x, float y) const {
    auto lab = LabCoord(x, y);

    float r = sqrtf(lab.x * lab.x + lab.y * lab.y);
    return atan2f(r, lab.z);
}

float DiffractionGeometry::Phi_rad(float x, float y) const {
    auto lab = LabCoord(x, y);
    auto v = atan2f(lab.y, lab.x);
    if (v < 0)
        v += 2.0f * PI;
    return v;
}

float DiffractionGeometry::PxlToRes(float x, float y) const {
    float two_theta = TwoTheta_rad(x, y);
    return wavelength_A / (2.0f * sinf(two_theta/2.0f));
}

float DiffractionGeometry::PxlToQ(float x, float y) const {
    return 2.0f * PI / PxlToRes(x,y);
}

float DiffractionGeometry::PxlToRes(float dist_pxl) const {
    // This is agnostic to detector rotation!!!
    if (dist_pxl == 0)
        return INFINITY;
    float tan_2theta = dist_pxl * pixel_size_mm / det_distance_mm;
    float theta = atanf(tan_2theta) / 2.0;
    float d_A = wavelength_A / (2.0f * sinf(theta));
    return d_A;
}

float DiffractionGeometry::ResToPxl(float d_A) const {
    if (d_A == 0)
        return INFINITY;

    float sin_theta = wavelength_A / (2 * d_A);
    float theta = asinf(sin_theta);
    float tan_2theta = tanf(2 * theta);
    return tan_2theta * det_distance_mm / pixel_size_mm;
}

float DiffractionGeometry::DistFromEwaldSphere(const Coord &recip) const {
    auto S = recip + GetScatteringVector();
    return S.Length() - (1.0f/wavelength_A);
}

float DiffractionGeometry::CalcAzIntSolidAngleCorr(float x, float y) const {
    // The solid angle of a flat pixel depends on the incidence angle to the detector
    // normal, cos(alpha) = det_distance / |detector-frame position|. This is evaluated
    // in the detector's own frame, so it is invariant under detector tilt (rot1/rot2/rot3),
    // matching PyFAI solidAngleArray and MAX IV azint. It reduces to cos^3(2*theta) only
    // for an untilted detector.
    float u = (x - beam_x_pxl) * pixel_size_mm;
    float v = (y - beam_y_pxl) * pixel_size_mm;
    float cos_alpha = det_distance_mm / sqrtf(u * u + v * v + det_distance_mm * det_distance_mm);
    return cos_alpha * cos_alpha * cos_alpha;
}

float DiffractionGeometry::CalcAzIntPolarizationCorr(float x, float y, float coeff) const {
    auto cos_2theta = cosf(TwoTheta_rad(x, y));
    float cos_2theta_2 = cos_2theta * cos_2theta;
    float cos_2phi = cosf(2.0f * Phi_rad(x, y));
    return 0.5f * (1.0f + cos_2theta_2 - coeff * cos_2phi * (1.0f - cos_2theta_2));
}

float DiffractionGeometry::GetBeamX_pxl() const {
    return beam_x_pxl;
}

float DiffractionGeometry::GetBeamY_pxl() const {
    return beam_y_pxl;
}

float DiffractionGeometry::GetDetectorDistance_mm() const {
    return det_distance_mm;
}

float DiffractionGeometry::GetPixelSize_mm() const {
    return pixel_size_mm;
}

float DiffractionGeometry::GetWavelength_A() const {
    return wavelength_A;
}

DiffractionGeometry &DiffractionGeometry::BeamX_pxl(float input) {
    beam_x_pxl = input;
    return *this;
}

DiffractionGeometry &DiffractionGeometry::BeamY_pxl(float input) {
    beam_y_pxl = input;
    return *this;
}

DiffractionGeometry &DiffractionGeometry::DetectorDistance_mm(float input) {
    if (input < 1.0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Det distance must be above 1.0 mm ");
    det_distance_mm = input;
    return *this;
}

DiffractionGeometry &DiffractionGeometry::PixelSize_mm(float input) {
    if (input <= 0.0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Pixel size must be positive number");
    pixel_size_mm = input;
    return *this;
}

DiffractionGeometry &DiffractionGeometry::Wavelength_A(float input) {
    if (input <= 0.0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Wavelength must be positive number");
    wavelength_A = input;
    return *this;
}

float DiffractionGeometry::AngleFromEwaldSphere_deg(const Coord &p0) const {
    // https://journals.iucr.org/d/issues/2014/08/00/dz5332/index.html
    Coord S0 = GetScatteringVector();

    const float epsilon = 1e-5f;

    float S0_sq = S0 * S0;
    float p0_sq = p0 * p0;
    float S0_p0 = S0 * p0;

    float val = S0_sq * p0_sq - S0_p0 * S0_p0;

    if (fabsf(val) < epsilon)
        return NAN;

    float A = std::sqrt((S0_sq - 1.0f/4.0f * p0_sq) * p0_sq / val);
    float B = (A * S0_p0 + p0_sq / 2.0f) / S0_sq;

    Coord p_star = A * p0 - B * S0;

    return angle_deg(p_star, p0);
}

void DiffractionGeometry::UpdateDetectorMatrix() {
    det_matrix = PoniRotMatrix(poni_rot_1, poni_rot_2, poni_rot_3) * orientation.Matrix();
}

DiffractionGeometry &DiffractionGeometry::PoniRot1_rad(float input) {
    poni_rot_1 = input;
    UpdateDetectorMatrix();
    return *this;
}

DiffractionGeometry &DiffractionGeometry::PoniRot2_rad(float input) {
    poni_rot_2 = input;
    UpdateDetectorMatrix();
    return *this;
}

DiffractionGeometry &DiffractionGeometry::PoniRot3_rad(float input) {
    poni_rot_3 = input;
    UpdateDetectorMatrix();
    return *this;
}

float DiffractionGeometry::GetPoniRot1_rad() const {
    return poni_rot_1;
}

float DiffractionGeometry::GetPoniRot2_rad() const {
    return poni_rot_2;
}

float DiffractionGeometry::GetPoniRot3_rad() const {
    return poni_rot_3;
}

DiffractionGeometry &DiffractionGeometry::Orientation(const DetectorOrientation &input) {
    orientation = input;
    UpdateDetectorMatrix();
    return *this;
}

DetectorOrientation DiffractionGeometry::GetOrientation() const {
    return orientation;
}

DiffractionGeometry &DiffractionGeometry::DetectorAxes(const Coord &fast, const Coord &slow) {
    const Coord f = fast.Normalize();
    const Coord s = slow.Normalize();
    // The normal is not free: it is the sample->PONI direction, and whether it is +fast x slow or
    // -fast x slow is exactly whether the stored image is mirrored, which the orientation already says.
    const Coord n = orientation.IsMirrorY() ? -(f % s) : (f % s);

    PoniAnglesFromMatrix(RotMatrix(f, s, n) * orientation.Matrix().transpose(),
                         poni_rot_1, poni_rot_2, poni_rot_3);
    UpdateDetectorMatrix();
    return *this;
}

Coord DiffractionGeometry::GetFastAxis() const {
    return det_matrix.Column(0);
}

Coord DiffractionGeometry::GetSlowAxis() const {
    return det_matrix.Column(1);
}

Coord DiffractionGeometry::GetNormalAxis() const {
    return det_matrix.Column(2);
}

std::pair<float, float> DiffractionGeometry::ResPhiToPxl(float d_A, float phi_rad) const {
    // Guard invalid inputs
    if (wavelength_A <= 0.0f || d_A <= wavelength_A / 2.0f)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Resolution to high for a given wavelength");

    float sin_theta = wavelength_A / (2.0f * d_A);
    float theta = asinf(sin_theta);
    float k = 1.0f / wavelength_A;

    float s2t = sinf(2.0f * theta);
    float c2t = cosf(2.0f * theta);

    float cphi = cosf(phi_rad);
    float sphi = sinf(phi_rad);

    return RecipToDetector(Coord{ k * s2t * cphi,k * s2t * sphi,k * (c2t - 1.0f)});
}

Coord DiffractionGeometry::ProjectToEwaldSphere(const Coord &p0) const {
    Coord S0 = GetScatteringVector();
    Coord S = p0 + S0;
    S = S.Normalize() / wavelength_A;
    return S - S0;
}

const RotMatrix &DiffractionGeometry::GetDetectorMatrix() const {
    return det_matrix;
}

std::optional<GoniometerAxis> DiffractionGeometry::GetRotation() const {
    return axis;
}

DiffractionGeometry &DiffractionGeometry::Rotation(const std::optional<GoniometerAxis> &input) {
    axis = input;
    return *this;
}