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

#include <cmath>

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

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 poni_rot * detectorCoord;
}

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

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

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

std::pair<float, float> DiffractionGeometry::RecipToDector(const Coord &recip) const {
    auto S_unrotated = recip + GetScatteringVector();
    auto S = poni_rot.transpose() * S_unrotated;

    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(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(float x, float y) const {
    auto lab = LabCoord(x, y);
    auto v = atan2f(-lab.y, lab.x);
    if (v < 0)
        v += 2.0f * M_PI;
    return v;
}

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

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 fabsf(S.Length() - (1.0f/wavelength_A));
}

float DiffractionGeometry::CalcAzIntSolidAngleCorr(float q) const {
    float sin_theta = q * wavelength_A / (4 * static_cast<float>(M_PI));
    float cos_2theta = 1.0f - 2.0f * sin_theta * sin_theta; // cos(2*alpha) = 1 - 2 * sin(alpha)^2
    float cos_2theta_3 = cos_2theta * cos_2theta * cos_2theta;
    return cos_2theta_3;
}

float DiffractionGeometry::CalcAzIntSolidAngleCorr(float x, float y) const {
    float cos_2theta = cosf(TwoTheta(x, y));
    float cos_2theta_3 = cos_2theta * cos_2theta * cos_2theta;
    return cos_2theta_3;
}

float DiffractionGeometry::CalcAzIntPolarizationCorr(float x, float y, float coeff) const {
    auto cos_2theta = cosf(TwoTheta(x, y));
    float cos_2theta_2 = cos_2theta * cos_2theta;
    float cos_2phi = cosf(2.0f * Phi(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(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::UpdatePoniRotMatrix() {
    poni_rot = RotMatrix(-poni_rot_1, {0,1,0})
               * RotMatrix(poni_rot_2, {1,0,0})
               * RotMatrix(poni_rot_3, {0,0,1});
}

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

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

DiffractionGeometry &DiffractionGeometry::PoniRot3_rad(float input) {
    poni_rot_3 = input;
    UpdatePoniRotMatrix();
    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;
}