// Copyright (2019-2023) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#ifndef JUNGFRAUJOCH_DIFFRACTIONGEOMETRY_H
#define JUNGFRAUJOCH_DIFFRACTIONGEOMETRY_H

#include "DiffractionExperiment.h"
#include <cmath>

inline Coord DetectorToRecip(const DiffractionExperiment &experiment, float x, float y) {
    return experiment.LabCoord(x, y).Normalize() / experiment.GetWavelength_A() - experiment.GetScatteringVector();
}

inline std::pair<float, float> RecipToDector(const DiffractionExperiment &experiment, const Coord &recip) {
    auto S = recip + experiment.GetScatteringVector();
    float coeff = experiment.GetDetectorDistance_mm() / (S.z * experiment.GetPixelSize_mm());
    float x = experiment.GetBeamX_pxl() + S.x * coeff;
    float y = experiment.GetBeamY_pxl() + S.y * coeff;
    return {x, y};
}

inline float CosTwoTheta(const DiffractionExperiment& experiment, float x, float y) {
    auto lab = experiment.LabCoord(x, y);
    return experiment.GetDetectorDistance_mm() / lab.Length();
}

inline float Phi(const DiffractionExperiment& experiment, float x, float y) {
    auto lab = experiment.LabCoord(x, y);
    return atan2f(lab.y, lab.x);
}

inline float PxlToRes(const DiffractionExperiment& experiment, float x, float y) {
    float cos_2theta = CosTwoTheta(experiment, 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 experiment.GetWavelength_A() / (2 * sin_theta);
}

inline float ResToPxl(const DiffractionExperiment& experiment, float d) {
    if (d == 0)
        return INFINITY;

    float sin_theta = experiment.GetWavelength_A() / (2 * d);
    float theta = asinf(sin_theta);
    float tan_2theta = tanf(2 * theta);
    return tan_2theta * experiment.GetDetectorDistance_mm() / experiment.GetPixelSize_mm();
}

inline float DistFromEwaldSphere(const DiffractionExperiment& experiment, const Coord& recip) {
    auto S = recip + experiment.GetScatteringVector();
    return fabsf(S.Length() - (1.0f/experiment.GetWavelength_A()));
}

inline float CalcRadIntSolidAngleCorr(const DiffractionExperiment& experiment, float q) {
    float sin_theta = q * experiment.GetWavelength_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;
}

inline float CalcRadIntSolidAngleCorr(const DiffractionExperiment& experiment, float x, float y) {
    float cos_2theta = CosTwoTheta(experiment, x, y);
    float cos_2theta_3 = cos_2theta * cos_2theta * cos_2theta;
    return cos_2theta_3;
}

inline float CalcRadIntPolarizationCorr(const DiffractionExperiment& experiment, float x, float y) {
    auto cos_2theta = CosTwoTheta(experiment, x, y);
    float cos_2theta_2 = cos_2theta * cos_2theta;
    float cos_2phi = cosf(2.0f * Phi(experiment, x, y));
    return 0.5f * (1.0f + cos_2theta_2 - experiment.GetPolarizationFactor() * cos_2phi * (1.0f - cos_2theta_2));
}

inline std::vector<float> CalcRadIntCorr(const DiffractionExperiment& experiment) {
    std::vector<float> corr(experiment.GetPixelsNum(), 1.0);
    auto xpixels = experiment.GetXPixelsNum();
    auto ypixels = experiment.GetYPixelsNum();

    for (int y = 0; y < ypixels; y++) {
        for (int x = 0; x < xpixels; x++) {
            if (experiment.GetApplySolidAngleCorr())
                corr[y * xpixels + x] *= CalcRadIntSolidAngleCorr(experiment,
                                                                  static_cast<float>(x),
                                                                  static_cast<float>(y));
            if (experiment.GetApplyPolarizationCorr())
                corr[y * xpixels + x] *= CalcRadIntPolarizationCorr(experiment,
                                                                          static_cast<float>(x),
                                                                          static_cast<float>(y));
        }
    }

    return corr;
}

#endif //JUNGFRAUJOCH_DIFFRACTIONGEOMETRY_H