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

#ifndef JUNGFRAUJOCH_SENSORABSORPTION_H
#define JUNGFRAUJOCH_SENSORABSORPTION_H

#include <cmath>
#include <string>

#include "../common/BraggIntegrationSettings.h"   // FlightPathMedium

// How much of the beam a flat sensor stops, as a function of the angle at which the beam enters it.
//
// A photon arriving at incidence angle alpha to the sensor normal crosses t/cos(alpha) of sensor
// instead of t, so more of it is absorbed: the detector is MORE efficient at high angle than at
// normal incidence, and an uncorrected reflection out at the detector edge reads high. Dividing by
// that ratio is the whole correction.
//
//   a(alpha) = t / (L cos alpha)          crossing length in attenuation lengths
//   QE(alpha) = 1 - exp(-a)               absorbed fraction
//   correction to I = QE(0) / QE(alpha)   normalised so normal incidence is untouched
//
// Two limits, both physical and both reached in this corpus:
//   - L << t (long wavelength): every photon stops in the entrance skin, QE = 1 at every angle and
//     the correction is exactly 1. This is why the low-energy case is inert rather than gated off.
//   - L >> t (thin sensor, hard X-rays): absorption is proportional to path length and the
//     correction tends to cos(alpha) - a factor of 2 at 2theta = 60 degrees. This is the regime
//     small-molecule work at 18-25 keV lives in.
//
// alpha is the angle to the DETECTOR NORMAL, not the scattering angle. The two coincide only on an
// untilted detector. On a tilted one the correction acquires an azimuthal dependence at fixed
// resolution, which is the only part of it that is not degenerate with an overall Wilson B.
//
// Attenuation coefficients: NIST X-Ray Mass Attenuation Coefficients (Hubbell & Seltzer), total
// mu/rho with coherent scattering. Photoelectric absorption dominates over the whole range used
// here, and because the correction is a RATIO of two absorbed fractions the photoelectric branching
// ratio cancels exactly; only the attenuation length enters. Coherent scattering, which attenuates
// without converting, is a few per cent of mu/rho in Si below 20 keV and is kept in the total - it
// biases L slightly short, and the correction depends on L only through t/L.
//
// NOT modelled, and negligible where this is used: K-fluorescence escape from the sensor (in CdTe
// above the Cd K edge at 26.71 keV a fluorescence photon can leave the sensor, so the charge is
// recorded in the wrong pixel or not at all); charge sharing between pixels; and the obliquity of
// the entrance window. Above the Cd/Te K edges the escape term is large and this model should not
// be trusted as it stands.
namespace sensor_absorption {

// NIST mass attenuation coefficient mu/rho [cm^2/g] against photon energy [keV]. Absorption edges
// are the repeated energies - the table is walked in order and interpolated log-log, which is how
// these tables are meant to be read. Si has only the K edge at 1.839 keV, far below anything used.
struct MuRow { double E_keV, mu_rho; };

inline constexpr MuRow SI_TABLE[] = {
    {1.0, 1570.0}, {1.5, 535.5}, {1.8389, 309.2}, {1.8389, 3192.0}, {2.0, 2777.0},
    {3.0, 978.4}, {4.0, 452.9}, {5.0, 245.0}, {6.0, 147.0}, {8.0, 64.68}, {10.0, 33.89},
    {15.0, 10.34}, {20.0, 4.464}, {30.0, 1.436}, {40.0, 0.7012}, {50.0, 0.4385}};

inline constexpr MuRow CD_TABLE[] = {
    {1.0, 7350.0}, {1.5, 2931.0}, {2.0, 1473.0}, {3.0, 541.4}, {3.5375, 357.5},
    {3.5375, 1152.0}, {3.63101, 1083.0}, {3.727, 1013.0}, {3.727, 1389.0}, {4.0, 1170.0},
    {4.018, 1157.0}, {4.018, 1324.0}, {5.0, 768.5}, {6.0, 479.3}, {8.0, 225.4}, {10.0, 124.4},
    {15.0, 41.78}, {20.0, 19.20}, {26.7112, 8.809}, {26.7112, 50.65}, {30.0, 37.65},
    {40.0, 17.78}, {50.0, 9.779}};

inline constexpr MuRow TE_TABLE[] = {
    {1.0, 8434.0}, {1.5, 3608.0}, {2.0, 1832.0}, {3.0, 679.2}, {4.0, 329.7}, {4.3414, 267.8},
    {4.3414, 788.2}, {4.47465, 750.4}, {4.612, 699.5}, {4.612, 944.5}, {4.7728, 878.2},
    {4.9392, 806.2}, {4.9392, 929.2}, {5.0, 901.4}, {6.0, 572.1}, {8.0, 270.2}, {10.0, 150.1},
    {15.0, 50.78}, {20.0, 23.41}, {30.0, 7.878}, {31.8138, 6.738}, {31.8138, 37.19},
    {40.0, 20.64}, {50.0, 11.45}};

// Dry air near sea level, for the sample-to-detector flight path below. The repeated energy is the
// argon K edge: argon is only 1.28% of air by mass but dominates its absorption just above 3 keV,
// which is inside the long-wavelength range this table is read at.
inline constexpr MuRow AIR_TABLE[] = {
    {1.0, 3606.0}, {1.5, 1191.0}, {2.0, 527.9}, {3.0, 162.5}, {3.2029, 134.0}, {3.2029, 148.5},
    {4.0, 77.88}, {5.0, 40.27}, {6.0, 23.41}, {8.0, 9.921}, {10.0, 5.120}, {15.0, 1.614},
    {20.0, 0.7779}, {30.0, 0.3538}, {40.0, 0.2485}, {50.0, 0.2080}};

// Helium, for a helium flight path. It attenuates about 1/600 of what air does at 3.8 keV - two
// electrons per atom against air's nitrogen and oxygen, and a seventh of the density - which is why
// long-wavelength stations use it. Not zero, though, so it is modelled rather than treated as vacuum.
inline constexpr MuRow HE_TABLE[] = {
    {1.0, 60.84}, {1.5, 16.76}, {2.0, 6.863}, {3.0, 2.007}, {4.0, 0.9329}, {5.0, 0.5766},
    {6.0, 0.4195}, {8.0, 0.2933}, {10.0, 0.2476}, {15.0, 0.2092}, {20.0, 0.1960},
    {30.0, 0.1838}, {40.0, 0.1763}, {50.0, 0.1703}};

// Log-log interpolation, clamped at both ends of the table.
template <std::size_t N>
double InterpMuRho(const MuRow (&tab)[N], double E_keV) {
    if (E_keV <= tab[0].E_keV)
        return tab[0].mu_rho;
    if (E_keV >= tab[N - 1].E_keV)
        return tab[N - 1].mu_rho;
    std::size_t i = 1;
    while (i < N - 1 && tab[i].E_keV < E_keV)
        i++;
    const double e0 = tab[i - 1].E_keV, e1 = tab[i].E_keV;
    if (!(e1 > e0))                    // the two rows of an absorption edge: take the upper side
        return tab[i].mu_rho;
    const double f = (std::log(E_keV) - std::log(e0)) / (std::log(e1) - std::log(e0));
    return std::exp(std::log(tab[i - 1].mu_rho) + f * (std::log(tab[i].mu_rho) - std::log(tab[i - 1].mu_rho)));
}

// Attenuation length 1/mu [um] of a sensor material at a given wavelength. Si and CdTe are the
// sensors in use; an unrecognised material is treated as silicon, which is what DetectorSetup
// defaults to anyway.
inline double AttenuationLength_um(const std::string &material, double lambda_A) {
    if (!(lambda_A > 0.0))
        return 0.0;
    const double E_keV = 12.39842 / lambda_A;
    double mu_rho_cm2_g, rho_g_cm3;
    if (material == "CdTe") {
        // Mass fractions from the atomic weights, Cd 112.414 and Te 127.60.
        constexpr double w_cd = 112.414 / 240.014, w_te = 127.60 / 240.014;
        mu_rho_cm2_g = w_cd * InterpMuRho(CD_TABLE, E_keV) + w_te * InterpMuRho(TE_TABLE, E_keV);
        rho_g_cm3 = 5.85;
    } else {
        mu_rho_cm2_g = InterpMuRho(SI_TABLE, E_keV);
        rho_g_cm3 = 2.3290;
    }
    const double mu_cm = mu_rho_cm2_g * rho_g_cm3;
    return mu_cm > 0.0 ? 1e4 / mu_cm : 0.0;
}

// Everything the per-reflection correction needs, reduced to the two numbers that are constant for
// a dataset. Built once on the host; the GPU predictor takes the two floats.
struct SensorQE {
    float a0 = 0.0f;        // t / L, the optical thickness at normal incidence
    float qe0 = 1.0f;       // absorbed fraction at normal incidence
    bool active = false;    // false where the sensor is opaque and the correction is exactly 1

    // Opaque beyond this: exp(-20) = 2e-9, so QE(0)/QE(alpha) rounds to exactly 1.0f for every
    // incidence angle and the correction is bit-identical to not applying it. This is what makes
    // the long-wavelength case inert without a flag or a threshold anyone has to choose.
    static constexpr float OPAQUE_A0 = 20.0f;

    static SensorQE Build(const std::string &material, double thickness_um, double lambda_A) {
        SensorQE q;
        const double L = AttenuationLength_um(material, lambda_A);
        if (!(thickness_um > 0.0) || !(L > 0.0))
            return q;
        q.a0 = static_cast<float>(thickness_um / L);
        q.qe0 = static_cast<float>(1.0 - std::exp(-q.a0));
        q.active = q.a0 < OPAQUE_A0;
        return q;
    }

    // The multiplicative correction to an intensity recorded at incidence angle alpha: QE(0)/QE(alpha).
    // Always <= 1, because a sensor is more efficient off-normal than head-on.
    [[nodiscard]] float Factor(float cos_alpha) const {
        if (!active || !(cos_alpha > 1e-3f))
            return 1.0f;
        const float qe = 1.0f - std::exp(-a0 / cos_alpha);
        return qe > 0.0f ? qe0 / qe : 1.0f;
    }
};


// What the diffracted beam crosses between the sample and the sensor - the one other term that
// carries the same cos(alpha) dependence as the sensor crossing above, and carries it with the
// opposite sign. A reflection leaving the sample at incidence angle alpha to the detector normal
// reaches its pixel after D/cos(alpha) of flight rather than D, so it crosses more of the medium and
// arrives attenuated, where the sensor above makes it read high.
//
//   T(alpha) = exp(-D / (L cos alpha))     Beer-Lambert; L = 1/mu is the attenuation length
//   correction to I = T(0) / T(alpha)      normalised so normal incidence is untouched
//                   = exp(d_over_L * (1/cos alpha - 1))
//
// Normalising at alpha = 0 divides out exp(-D/L), a constant for the dataset that the fitted
// per-image scale absorbs; what is left is the only part of the flight path that is not degenerate
// with that scale. The whole term is COMPUTED, never fitted: mu comes from the NIST tables above,
// and the distance and the wavelength are both stated by the file.
//
// The attenuation length of air falls steeply toward low energy - 8.2 m at 18 keV, 3.0 m at
// 12.4 keV, but only 8.9 cm at 3.8 keV - so in air the term is a few tenths of a per cent at hard
// X-rays and several-fold at long wavelength. That is why the medium is a user choice: a station
// working at 3.8 keV puts the beam in helium precisely because air at that energy is unusable.
//
// The medium is NOT detected, and that is a conclusion rather than an omission. Nothing in the
// files states it, and the implied transmission does not separate the cases either - in this corpus
// a confirmed helium station sits at 51% implied air transmission and a confirmed air station at
// 63%. Any rule dividing those would be a threshold fitted between two points, so there is none:
// the medium is declared (rugnux --flight-path), defaulted to air, and printed in the report.
struct FlightPathAttenuation {
    float d_over_L = 0.0f;   // normal-incidence flight path in attenuation lengths; 0 = vacuum

    // Dry air at 20 C and 1 atm is 1.205e-3 g/cm^3; helium at the same conditions 1.663e-4 (NIST).
    static FlightPathAttenuation Build(FlightPathMedium medium, double distance_mm, double lambda_A) {
        FlightPathAttenuation a;
        if (medium == FlightPathMedium::Vacuum || !(distance_mm > 0.0) || !(lambda_A > 0.0))
            return a;
        const double E_keV = 12.39842 / lambda_A;
        const double mu_cm = medium == FlightPathMedium::Helium
                             ? InterpMuRho(HE_TABLE, E_keV) * 1.663e-4
                             : InterpMuRho(AIR_TABLE, E_keV) * 1.205e-3;
        if (mu_cm > 0.0)
            a.d_over_L = static_cast<float>(distance_mm * 0.1 * mu_cm);
        return a;
    }

    // The multiplicative correction to an intensity recorded at incidence angle alpha. Always >= 1,
    // because a reflection that arrives obliquely crossed more of the medium than one arriving
    // head-on.
    [[nodiscard]] float Factor(float cos_alpha) const {
        if (!(d_over_L > 0.0f) || !(cos_alpha > 1e-3f))
            return 1.0f;
        return std::exp(d_over_L * (1.0f / cos_alpha - 1.0f));
    }

    // What assuming this medium is worth, as a shift in the Wilson B the merged data will show.
    //
    // On an untilted detector alpha is the scattering angle, so the correction is a pure function of
    // resolution: it cancels within a resolution shell and cannot move R_meas or CC1/2 there. Its
    // whole effect on merged data is therefore a change of slope in the Wilson plot, and that is a
    // number the report can state. Least-squares slope of ln(factor) against s^2 = (sin theta /
    // lambda)^2 over the resolution range, halved because I falls as exp(-2 B s^2); returned
    // negative, since correcting an attenuation lifts the high-angle data and flattens the fall-off.
    [[nodiscard]] double WilsonBShift_A2(double d_min_A, double d_max_A, double lambda_A) const {
        if (!(d_over_L > 0.0f) || !(d_min_A > 0.0) || !(d_max_A > d_min_A) || !(lambda_A > 0.0))
            return 0.0;
        constexpr int N = 128;
        double sx = 0.0, sy = 0.0, sxx = 0.0, sxy = 0.0;
        int n = 0;
        for (int i = 0; i < N; i++) {
            const double d = d_min_A + (d_max_A - d_min_A) * i / (N - 1.0);
            const double sin_theta = lambda_A / (2.0 * d);
            if (!(sin_theta < 1.0))
                continue;
            const double two_theta = 2.0 * std::asin(sin_theta);
            const double x = (sin_theta / lambda_A) * (sin_theta / lambda_A);
            const double y = d_over_L * (1.0 / std::cos(two_theta) - 1.0);
            sx += x; sy += y; sxx += x * x; sxy += x * y; n++;
        }
        const double den = n * sxx - sx * sx;
        if (n < 2 || !(std::fabs(den) > 0.0))
            return 0.0;
        return -(n * sxy - sx * sy) / den / 2.0;
    }
};

}   // namespace sensor_absorption

#endif   // JUNGFRAUJOCH_SENSORABSORPTION_H