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

#include "ScaleOnTheFly.h"

#include <future>
#include <ceres/ceres.h>
#include <ceres/rotation.h>

namespace {
    double SafeInv(double x, double fallback) {
        if (!std::isfinite(x) || x == 0.0)
            return fallback;
        return 1.0 / x;
    }

    float RotationPartiality( double delta_phi_deg,
                     double zeta,
                     double mosaicity_deg,
                     double wedge_deg) {
        const double half_wedge = wedge_deg / 2.0;
        const double c1 = zeta / std::sqrt(2.0);
        const double arg_plus = (delta_phi_deg + half_wedge) * c1 / mosaicity_deg;
        const double arg_minus = (delta_phi_deg - half_wedge) * c1 / mosaicity_deg;
        return static_cast<float>((std::erf(arg_plus) - std::erf(arg_minus)) / 2.0);
    }


    struct RegularizationResidual {
        RegularizationResidual(const double weight, const double expected)
            : weight(weight),
              expected(expected) {
        }

        template<typename T>
        bool operator()(const T *parameter, T *residual) const {
            residual[0] = T(weight) * (parameter[0] - T(expected));
            return true;
        }

    private:
        const double weight;
        const double expected;
    };

    class ScalingResidual {
    protected:
        const double Iobs;
        const double Itrue;
        const double weight;
        const double lp;
        const double b_resolution_coeff;

        ScalingResidual(const Reflection &r, double Itrue, double sigma)
            : Iobs(r.I),
              Itrue(Itrue),
              weight(SafeInv(sigma, 1.0)),
              lp(SafeInv(r.rlp, 1.0)),
              b_resolution_coeff(-SafeInv(4.0 * r.d * r.d, 0.0)) {
        }
    };

    struct ScalingRotationResidual : public ScalingResidual {
        ScalingRotationResidual(const Reflection &r, double Itrue, double sigma)
            : ScalingResidual(r, Itrue, sigma),
              delta_phi_deg(r.delta_phi_deg),
              c1(r.zeta / std::sqrt(2.0)) {
        }

        template<typename T>
        bool operator()(const T *const G,
                        const T *const B,
                        const T *const mosaicity,
                        const T *const wedge,
                        T *residual) const {
            if (mosaicity[0] < 1e-6)
                return false;

            const T half_wedge = wedge[0] / T(2.0);
            const T arg_plus = (T(delta_phi_deg) + half_wedge) * T(c1) / mosaicity[0];
            const T arg_minus = (T(delta_phi_deg) - half_wedge) * T(c1) / mosaicity[0];
            const T partiality = (ceres::erf(arg_plus) - ceres::erf(arg_minus)) / T(2.0);
            const T B_term = ceres::exp(B[0] * T(b_resolution_coeff));
            residual[0] = (G[0] * partiality * B_term * T(lp) * T(Itrue) - T(Iobs)) * T(weight);
            return true;
        }

        double delta_phi_deg;
        double c1;
    };

    struct IntensityFixedResidual : public ScalingResidual {
        IntensityFixedResidual(const Reflection &r, double Itrue, double sigma, double partiality)
            : ScalingResidual(r, Itrue, sigma),
              partiality(partiality) {
        }

        template<typename T>
        bool operator()(const T *const G, const T *const B, T *residual) const {
            const T B_term = ceres::exp(T(B[0]) * b_resolution_coeff);
            residual[0] = (G[0] * T(partiality) * B_term * T(lp) * Itrue - T(Iobs)) * T(weight);
            return true;
        }

        double partiality;
    };

struct ScalingPostRefResidual : public ScalingResidual {
    // In this algorithm at the point of post-refinement we don't anymore care for where maximum of
    // the reflection was located and if it fits the observed position.
    // This reflections was already integrated and we cannot integrate it better at this point.
    // But, we could adjust partiality to indicate that this reflection was wrongly predicted.
    // I.e., integrated position was far away from true reflection, so partiality must be low.
    // This is an empiric model and need to see if this will work in practice at all.
    // I hope it will allow the model to find that reflections were misindexed.
    // We assume at this point that initial indexing was done properly and integration was generally OK => most low resolution reflections fit correctly
    // Yet we know, that small errors in indexing are inducing misalignment at high resolution - sometimes it is visible that high-resolution reflections
    // are away from shoe-boxes observed in the image, if we can catch this at post-refinement/scaling step, this would be great.
    // Next logical step is to do this pixel-wise - for each pixel refine partiality and merge pixels
    // This should work for per-image scaling, or even, maybe, for full rotation datasets (3600 images)
    // Then we could properly take into account misalignment of shoe-box center vs. partiality and also remove most pixels
    // in the shoe-box that don't really contribute to the reflection.
    // But...it could also drift to downweighting partiality for all high resolution reflections to make loss function "fake happy".

    // We assume rot3 == 0. Rot3 is not really helping much in crystallography (other than fixing polarization correction)
    ScalingPostRefResidual(const Reflection &r, double Itrue, double sigma,
                           const DiffractionGeometry &geometry,
                           const CrystalLattice &lattice,
                           double exp_h, double exp_k,
                           double exp_l)
    : ScalingResidual(r, Itrue, sigma),
          integration_center_x(r.predicted_x),
          integration_center_y(r.predicted_y),
          inv_lambda(SafeInv(geometry.GetWavelength_A(), 0.0)),
          pixel_size(geometry.GetPixelSize_mm()),
          det_dist_mm(geometry.GetDetectorDistance_mm()),
          beam_x(geometry.GetBeamX_pxl()),
          beam_y(geometry.GetBeamY_pxl()),
          exp_h(exp_h),
          exp_k(exp_k),
          exp_l(exp_l),
          Astar(lattice.Astar()),
          Bstar(lattice.Bstar()),
          Cstar(lattice.Cstar()),
          c1(std::cos(geometry.GetPoniRot1_rad())),
          s1(std::sin(geometry.GetPoniRot1_rad())),
          c2(std::cos(geometry.GetPoniRot2_rad())),
          s2(std::sin(geometry.GetPoniRot2_rad())) {
    }

    template <typename T>
    T CalcPartiality(const T *const R_radial,
                    const T *const R_tangential,
                    const T *const beam_corr,
                    const T *const p0) const {
        // Detector coordinates in mm
        const T det_x = (T(integration_center_x) - beam_x - beam_corr[0]) * T(pixel_size);
        const T det_y = (T(integration_center_y) - beam_y - beam_corr[1]) * T(pixel_size);
        const T det_z = T(det_dist_mm);

        // Apply Ry(rot1) first: rotate around Y
        const T t1_x = T(c1) * det_x + T(s1) * det_z;
        const T t1_y = det_y;
        const T t1_z = T(-s1) * det_x + T(c1) * det_z;

        // Then apply Rx(-rot2): rotate around X
        const T x = t1_x;
        const T y = T(c2) * t1_y + T(s2) * t1_z;
        const T z = - T(s2) * t1_y + T(c2) * t1_z;

        // convert to recip space
        const T lab_norm = ceres::sqrt(x * x + y * y + z * z);
        const T inv_norm = T(1) / lab_norm;

        T recip_obs[3];
        recip_obs[0] = x * inv_norm * inv_lambda;
        recip_obs[1] = y * inv_norm * inv_lambda;
        recip_obs[2] = (z * inv_norm - T(1.0)) * inv_lambda;

        const Eigen::Matrix<T, 3, 1> e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]);

        const T astar_unrot[3] = {T(Astar.x), T(Astar.y), T(Astar.z)};
        const T bstar_unrot[3] = {T(Bstar.x), T(Bstar.y), T(Bstar.z)};
        const T cstar_unrot[3] = {T(Cstar.x), T(Cstar.y), T(Cstar.z)};

        T astar_rot[3], bstar_rot[3], cstar_rot[3];

        ceres::AngleAxisRotatePoint(p0, astar_unrot, astar_rot);
        ceres::AngleAxisRotatePoint(p0, bstar_unrot, bstar_rot);
        ceres::AngleAxisRotatePoint(p0, cstar_unrot, cstar_rot);

        const Eigen::Matrix<T, 3, 1> e_pred_recip(T(exp_h) * astar_rot[0] + T(exp_k) * bstar_rot[0] + T(exp_l) * cstar_rot[0],
            T(exp_h) * astar_rot[1] + T(exp_k) * bstar_rot[1] + T(exp_l) * cstar_rot[1],
            T(exp_h) * astar_rot[2] + T(exp_k) * bstar_rot[2] + T(exp_l) * cstar_rot[2]
        );

        // Ewald sphere centre is at -k_i = (0, 0, -inv_lambda)
        // Radial direction: outward normal at g_hkl
        const Eigen::Matrix<T, 3, 1> S_pred(
            e_pred_recip[0],
            e_pred_recip[1],
            e_pred_recip[2] + T(inv_lambda)   // g_hkl + k_i
        );
        const T S_pred_norm = S_pred.norm();
        if (S_pred_norm < T(1e-10))
            return T(0);

        const Eigen::Matrix<T, 3, 1> n_radial = S_pred / S_pred_norm;

        const Eigen::Matrix<T, 3, 1> delta_q = e_obs_recip - e_pred_recip;
        const T eps_radial        = delta_q.dot(n_radial);
        const Eigen::Matrix<T, 3, 1> dq_tang = delta_q - eps_radial * n_radial;
        const T eps_tangential_sq = dq_tang.squaredNorm();   // guaranteed ≥ 0
        // ─────────────────────────────────────────────────────────────

        return ceres::exp(
            - eps_radial * eps_radial / (R_radial[0]    * R_radial[0])
            - eps_tangential_sq       / (R_tangential[0] * R_tangential[0])
        );
    }

    template<typename T>
    bool operator()(const T *const G,
                    const T *const B,
                    const T *const R_radial,
                    const T *const R_tangential,
                    const T *const beam_corr,
                    const T *const p0,
                    T *residual) const {
        if (R_radial[0] < T(1e-10) || R_tangential[0] < T(1e-10))
            return false;

        const T B_term = ceres::exp(B[0] * T(b_resolution_coeff));

        const T partiality = CalcPartiality(R_radial, R_tangential, beam_corr, p0);
        residual[0] = (G[0] * partiality * B_term * T(lp) * T(Itrue)
                       - T(Iobs)) * T(weight);
        return true;
    }

    const double integration_center_x, integration_center_y;
    const double inv_lambda;
    const double pixel_size;
    const double det_dist_mm;
    const double beam_x, beam_y;
    const double exp_h;
    const double exp_k;
    const double exp_l;
    const Coord Astar, Bstar, Cstar;
    const double c1,s1,c2,s2;
};
}

ScaleOnTheFly::ScaleOnTheFly(const DiffractionExperiment &x, const std::vector<MergedReflection> &ref)
    : sg(x.GetGemmiSpaceGroup()),
      model(x.GetPartialityModel()),
      s(x.GetScalingSettings()),
      rot_wedge_deg(x.GetRotationWedgeForScaling()),
      refine_rot_wedge(x.GetRefineRotationWedgeInScaling()),
      hkl_key_generator(s.GetMergeFriedel(), x.GetSpaceGroupNumber().value_or(1)) {
    for (const auto &r: ref) {
        const auto key = hkl_key_generator(r);
        reference_data[key] = r.I;
    }
}

bool ScaleOnTheFly::Accept(const Reflection &r) const {
    if (!AcceptReflection(r, s.GetHighResolutionLimit_A()))
        return false;

    switch (model) {
        case PartialityModel::Rotation:
            return std::isfinite(r.zeta) && r.zeta > 0.0f;
        case PartialityModel::Fixed:
        case PartialityModel::Unity:
            return true;
    }

    return true;
}

std::pair<double, size_t> ScaleOnTheFly::CalculateGlobalCC(const std::vector<Reflection> &reflections) const {
    double sum_x = 0.0;
    double sum_y = 0.0;
    double sum_x2 = 0.0;
    double sum_y2 = 0.0;
    double sum_xy = 0.0;
    size_t n = 0;

    for (const auto &r: reflections) {
        if (!AcceptReflection(r, s.GetHighResolutionLimit_A()))
            continue;
        if (r.partiality < s.GetMinPartiality())
            continue;
        if (!std::isfinite(r.I) || !std::isfinite(r.image_scale_corr) || r.image_scale_corr <= 0.0f)
            continue;
        if (!std::isfinite(r.sigma) || r.sigma <= 0.0f)
            continue;

        const HKLKey key = hkl_key_generator(r);
        const auto it = reference_data.find(key);
        if (it == reference_data.end())
            continue;

        const double image_i = static_cast<double>(r.I) * static_cast<double>(r.image_scale_corr);
        const double ref_i = it->second;

        if (!std::isfinite(image_i) || !std::isfinite(ref_i))
            continue;

        sum_x += image_i;
        sum_y += ref_i;
        sum_x2 += image_i * image_i;
        sum_y2 += ref_i * ref_i;
        sum_xy += image_i * ref_i;
        ++n;
    }

    if (n < MIN_REFLECTIONS)
        return {NAN, n};

    const double nd = static_cast<double>(n);
    const double cov = sum_xy - sum_x * sum_y / nd;
    const double var_x = sum_x2 - sum_x * sum_x / nd;
    const double var_y = sum_y2 - sum_y * sum_y / nd;

    if (!(var_x > 0.0 && var_y > 0.0))
        return {NAN, n};

    return {cov / std::sqrt(var_x * var_y), n};
}

ScaleOnTheFlyResult ScaleOnTheFly::Scale(std::vector<Reflection> &reflections,
    std::optional<float> mosaicity_deg) {
    auto start = std::chrono::steady_clock::now();

    ceres::Problem problem;

    ScaleOnTheFlyResult result{
        .B = 0.0,
        .G = 1.0
    };

    if (model == PartialityModel::Rotation) {
        if (mosaicity_deg && std::isfinite(*mosaicity_deg) && *mosaicity_deg > 0.0)
            result.mos = *mosaicity_deg;
        else
            result.mos = s.GetDefaultMosaicity();
        result.wedge = rot_wedge_deg.value_or(0.0);
    } else {
        result.mos = NAN;
        result.wedge = NAN;
    }

    size_t n_reflections = 0;
    for (const auto &r: reflections) {
        if (!Accept(r))
            continue;

        const HKLKey key = hkl_key_generator(r);

        if (!reference_data.contains(key))
            continue;

        ++n_reflections;

        const double Itrue = reference_data.at(key);
        const double sigma = r.sigma;

        switch (model) {
            case PartialityModel::Fixed: {
                auto *cost = new ceres::AutoDiffCostFunction<IntensityFixedResidual, 1, 1, 1>(
                    new IntensityFixedResidual(r, Itrue, sigma, r.partiality));
                problem.AddResidualBlock(cost, nullptr, &result.G, &result.B);
            }
            break;
            case PartialityModel::Unity: {
                auto *cost = new ceres::AutoDiffCostFunction<IntensityFixedResidual, 1, 1, 1>(
                    new IntensityFixedResidual(r, Itrue, sigma, 1.0));
                problem.AddResidualBlock(cost, nullptr, &result.G, &result.B);
            }
            break;
            case PartialityModel::Rotation: {
                auto *cost = new ceres::AutoDiffCostFunction<ScalingRotationResidual, 1, 1, 1, 1, 1>(
                    new ScalingRotationResidual(r, Itrue, sigma));
                problem.AddResidualBlock(cost, nullptr, &result.G, &result.B, &result.mos,
                    &result.wedge);
            }
                break;
            default:
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                      "Not supported partiality model");
        }
    }

    if (n_reflections < MIN_REFLECTIONS) {
        result.succesful = false;
        return result;
    }
    result.succesful = true;

    constexpr double SIGMA_SCALE_FACTOR = 2; // We assume scale factor is +/- 2.0
    const double scale_factor_regularization_weight = std::sqrt(static_cast<double>(n_reflections) / SIGMA_SCALE_FACTOR);

    if (s.GetScalingRegularize()) {
        auto *scale_reg_cost = new ceres::AutoDiffCostFunction<RegularizationResidual, 1, 1>(
            new RegularizationResidual(scale_factor_regularization_weight, 1.0));
        problem.AddResidualBlock(scale_reg_cost, nullptr, &result.G);
    }

    problem.SetParameterLowerBound(&result.G, 0, 0.0);

    if (s.GetRefineB()) {
        constexpr double SIGMA_B = 10.0; // in A^2
        const double b_regularization_weight = std::sqrt(static_cast<double>(n_reflections) / SIGMA_B);

        if (s.GetScalingRegularize()) {
            auto *b_reg_cost = new ceres::AutoDiffCostFunction<RegularizationResidual, 1, 1>(
                new RegularizationResidual(b_regularization_weight, 0.0));
            problem.AddResidualBlock(b_reg_cost, nullptr, &result.B);
        }

        problem.SetParameterLowerBound(&result.B, 0, s.GetMinB());
        problem.SetParameterUpperBound(&result.B, 0, s.GetMaxB());
    } else {
        problem.SetParameterBlockConstant(&result.B);
    }

    if (model == PartialityModel::Rotation) {
        if (refine_rot_wedge) {
            problem.SetParameterLowerBound(&result.wedge, 0, s.GetMinWedge());
            problem.SetParameterUpperBound(&result.wedge, 0, s.GetMaxWedge());
        } else {
            problem.SetParameterBlockConstant(&result.wedge);
        }
        problem.SetParameterLowerBound(&result.mos, 0, s.GetMinMosaicity());
        problem.SetParameterUpperBound(&result.mos, 0, s.GetMaxMosaicity());
    }

    ceres::Solver::Options options;
    options.linear_solver_type = ceres::DENSE_QR;
    options.minimizer_progress_to_stdout = false;
    options.num_threads = 1;

    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);

    for (auto &r: reflections) {
        const double B_term = exp(result.B * -SafeInv(4.0 * r.d * r.d, 0.0));

        switch (model) {
            case PartialityModel::Unity:
                r.partiality = 1.0;
                break;
            case PartialityModel::Rotation: {
                if (std::isfinite(r.delta_phi_deg) && std::isfinite(r.zeta) && result.mos > 1e-6)
                    r.partiality = RotationPartiality(r.delta_phi_deg, r.zeta, result.mos, result.wedge);
                break;
            }
            default:
                // For fixed partiality there is no need to change anything
                break;
        }
        const double denom = B_term * r.partiality * result.G;
        if (std::isfinite(r.rlp) &&
            std::isfinite(denom) &&
            denom > 0.0) {
                r.image_scale_corr = static_cast<float>(r.rlp / denom);
            } else {
                r.image_scale_corr = NAN;
            }
    }

    const auto [cc, cc_n] = CalculateGlobalCC(reflections);
    result.cc = cc;
    result.cc_n = cc_n;

    auto end = std::chrono::steady_clock::now();
    result.time_s = std::chrono::duration<float>(end - start).count();
    return result;
}

ScalingResult ScaleOnTheFly::Scale(std::vector<std::vector<Reflection> > &reflections,
                        const std::vector<float> &mosaicity,
                        size_t nthreads) {
    ScalingResult result(reflections.size());

    if (nthreads == 0)
        nthreads = std::thread::hardware_concurrency();

    if (nthreads <= 1) {
        for (int i = 0; i < reflections.size(); i++) {
            std::optional<float> mos_val;
            if (model == PartialityModel::Rotation
                && mosaicity.size() > i
                && std::isfinite(mosaicity[i])
                && mosaicity[i] > 0.0)
                mos_val = mosaicity[i];

            auto local_result = Scale(reflections[i], mos_val);
            result.mosaicity_deg[i] = local_result.mos;
            result.image_bfactor_Ang2[i] = local_result.B;
            result.image_scale_g[i] = local_result.G;
            result.rotation_wedge_deg[i] = local_result.wedge;
            result.image_cc[i] = local_result.cc;
            result.image_cc_n[i] = local_result.cc_n;
        }
    } else {
        auto local_nthreads = std::min(nthreads, reflections.size());
        std::vector<std::future<void>> futures;
        futures.reserve(local_nthreads);
        std::atomic<size_t> curr_image = 0;

        for (size_t t = 0; t < local_nthreads; ++t)
            futures.emplace_back(std::async(std::launch::async, [&] {
                size_t i = curr_image.fetch_add(1);
                while (i < reflections.size()) {
                    std::optional<float> mos_val;
                    if (model == PartialityModel::Rotation
                        && mosaicity.size() > i
                        && std::isfinite(mosaicity[i])
                        && mosaicity[i] > 0.0)
                        mos_val = mosaicity[i];

                    auto local_result = Scale(reflections[i], mos_val);
                    result.mosaicity_deg[i] = local_result.mos;
                    result.image_bfactor_Ang2[i] = local_result.B;
                    result.image_scale_g[i] = local_result.G;
                    result.rotation_wedge_deg[i] = local_result.wedge;
                    result.image_cc[i] = local_result.cc;
                    result.image_cc_n[i] = local_result.cc_n;
                    i = curr_image.fetch_add(1);
                }
            }));

        for (auto &f: futures)
            f.get();
    }

    return result;
}