// 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>

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);
    }

    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;
    };
}

ScaleOnTheFly::ScaleOnTheFly(const std::vector<MergedReflection> &ref, const DiffractionExperiment &x)
    : 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) {
    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;
}

ScaleOnTheFlyResult ScaleOnTheFly::Scale(std::vector<Reflection> &reflections, std::optional<double> 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;

    if (s.GetRefineB()) {
        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.scaling_correction = static_cast<float>(r.rlp / denom);
            } else {
                r.scaling_correction = NAN;
            }
    }

    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<double> &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<double> 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;
        }
    } 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<double> 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;
                    i = curr_image.fetch_add(1);
                }
            }));

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

    return result;
}