// 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 {
    // Assumes planar detector, 90 deg towards beam
    return {(x - beam_x_pxl) * pixel_size_mm ,
            (y - beam_y_pxl) * pixel_size_mm ,
            det_distance_mm};
}

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 = recip + GetScatteringVector();
    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::CosTwoTheta(float x, float y) const {
    auto lab = LabCoord(x, y);
    return det_distance_mm / lab.Length();
}

float DiffractionGeometry::Phi(float x, float y) const {
    auto lab = LabCoord(x, y);
    return atan2f(lab.y, lab.x);
}

float DiffractionGeometry::PxlToRes(float x, float y) const {
    float cos_2theta = CosTwoTheta( x, y);
    if (cos_2theta == 1.0f)
        return std::numeric_limits<float>::infinity();

    // cos(2theta) = cos(theta)^2 - sin(theta)^2
    // cos(2theta) = 1 - 2*sin(theta)^2
    // Technically two solutions for two theta, but it makes sense only to take positive one in this case
    float sin_theta = sqrtf((1 - cos_2theta)/2);
    return wavelength_A / (2 * sin_theta);
}

float DiffractionGeometry::PxlToRes(float dist_pxl) const {
    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 = CosTwoTheta(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 = CosTwoTheta(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;
}