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

#include "ScaleAll.h"

#include <ceres/ceres.h>

#include <algorithm>
#include <cmath>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <thread>
#include "HKLKey.h"

namespace {
    struct ScaleObs {
        const Reflection *r = nullptr;
        int image = 0;
        int hkl = -1;
        double sigma = 1.0;
    };

    double SafeSigma(double sigma, double min_sigma) {
        if (!std::isfinite(sigma) || sigma <= 0.0)
            return min_sigma;
        return std::max(sigma, min_sigma);
    }

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

    bool AcceptReflectionForScaling(const Reflection &r, const ScaleMergeOptions &opt) {
        if (!AcceptReflection(r, opt.d_min_limit_A))
            return false;

        switch (opt.partiality_model) {
            case ScaleMergeOptions::PartialityModel::Rotation:
                return std::isfinite(r.zeta) && r.zeta > 0.0f;

            case ScaleMergeOptions::PartialityModel::Still:
                return std::isfinite(r.dist_ewald);

            case ScaleMergeOptions::PartialityModel::Fixed:
            case ScaleMergeOptions::PartialityModel::Unity:
                return true;
        }

        return true;
    }

    struct IntensityFixedResidual {
        IntensityFixedResidual(const Reflection &r, double sigma, double partiality)
            : Iobs(r.I),
              weight(SafeInv(sigma, 1.0)),
              correction(partiality * SafeInv(r.rlp, 1.0)) {
        }

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

        double Iobs;
        double weight;
        double correction;
    };

    struct IntensityRotationResidual {
        IntensityRotationResidual(const Reflection &r, double sigma)
            : Iobs(r.I),
              weight(SafeInv(sigma, 1.0)),
              delta_phi_deg(r.delta_phi_deg),
              lp(SafeInv(r.rlp, 1.0)),
              c1(r.zeta / std::sqrt(2.0)) {
        }

        template<typename T>
        bool operator()(const T *const G,
                        const T *const mosaicity,
                        const T *const Itrue,
                        const T *const wedge,
                        T *residual) const {
            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);

            residual[0] = (G[0] * partiality * T(lp) * Itrue[0] - T(Iobs)) * T(weight);
            return true;
        }

        double Iobs;
        double weight;
        double delta_phi_deg;
        double lp;
        double c1;
    };

    struct IntensityStillResidual {
        IntensityStillResidual(const Reflection &r, double sigma)
            : Iobs(r.I),
              weight(SafeInv(sigma, 1.0)),
              lp(SafeInv(r.rlp, 1.0)),
              dist_ewald_sq(r.dist_ewald * r.dist_ewald) {
        }

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

        double Iobs;
        double weight;
        double lp;
        double dist_ewald_sq;
    };

    struct ScaleRegularizationResidual {
        explicit ScaleRegularizationResidual(double sigma)
            : inv_sigma(SafeInv(sigma, 1.0)) {
        }

        template<typename T>
        bool operator()(const T *const k, T *residual) const {
            residual[0] = (k[0] - T(1.0)) * T(inv_sigma);
            return true;
        }

        double inv_sigma;
    };

    struct SmoothnessResidual {
        explicit SmoothnessResidual(double sigma)
            : inv_sigma(SafeInv(sigma, 1.0)) {
        }

        template<typename T>
        bool operator()(const T *const x0,
                        const T *const x1,
                        const T *const x2,
                        T *residual) const {
            residual[0] = (ceres::log(x0[0]) + ceres::log(x2[0]) - T(2.0) * ceres::log(x1[0])) * T(inv_sigma);
            return true;
        }

        double inv_sigma;
    };

    std::vector<ScaleObs> BuildScaleObs(const std::vector<std::vector<Reflection>> &observations,
                                        const ScaleMergeOptions &opt,
                                        std::vector<uint8_t> &image_used,
                                        int &nhkl) {
        std::map<HKLKey, int> hkl_to_slot;
        std::vector<ScaleObs> obs;

        size_t nrefl = 0;
        for (const auto &image: observations)
            nrefl += image.size();
        obs.reserve(nrefl);

        for (int image = 0; image < static_cast<int>(observations.size()); ++image) {
            const int image_slot = image / static_cast<int>(opt.image_cluster);

            for (const auto &r: observations[image]) {
                if (!AcceptReflectionForScaling(r, opt))
                    continue;

                HKLKey key;
                try {
                    key = CanonicalHKL(r, opt.merge_friedel, opt.space_group);
                } catch (...) {
                    continue;
                }
                auto it = hkl_to_slot.find(key);
                if (it == hkl_to_slot.end()) {
                    const int slot = static_cast<int>(hkl_to_slot.size());
                    it = hkl_to_slot.emplace(key, slot).first;
                }

                image_used[image_slot] = 1;

                obs.push_back({
                    .r = &r,
                    .image = image_slot,
                    .hkl = it->second,
                    .sigma = SafeSigma(r.sigma, opt.min_sigma)
                });
            }
        }

        nhkl = static_cast<int>(hkl_to_slot.size());
        return obs;
    }

    std::vector<double> InitialIntensities(int nhkl,
                                           const ScaleMergeOptions &opt,
                                           const std::vector<ScaleObs> &obs) {
        std::vector<std::vector<double>> values(nhkl);

        for (const auto &o: obs)
            values[o.hkl].push_back(o.r->I);

        std::vector<double> Itrue(nhkl, opt.min_sigma);

        for (int h = 0; h < nhkl; ++h) {
            auto &v = values[h];
            if (v.empty())
                continue;

            std::nth_element(v.begin(), v.begin() + static_cast<long>(v.size() / 2), v.end());

            Itrue[h] = v[v.size() / 2];
            if (!std::isfinite(Itrue[h]) || Itrue[h] <= opt.min_sigma)
                Itrue[h] = opt.min_sigma;
        }

        return Itrue;
    }

    void Scale(const ScaleMergeOptions &opt,
               const std::vector<ScaleObs> &obs,
               const std::vector<uint8_t> &image_used,
               int nhkl,
               std::vector<double> &G,
               std::vector<double> &mosaicity,
               std::vector<double> &R_sq) {
        ceres::Problem problem;

        auto Itrue = InitialIntensities(nhkl, opt, obs);
        double wedge = opt.wedge_deg.value_or(0.0);

        for (const auto &o: obs) {
            switch (opt.partiality_model) {
                case ScaleMergeOptions::PartialityModel::Rotation: {
                    auto *cost = new ceres::AutoDiffCostFunction<IntensityRotationResidual, 1, 1, 1, 1, 1>(
                        new IntensityRotationResidual(*o.r, o.sigma));
                    problem.AddResidualBlock(cost, nullptr, &G[o.image], &mosaicity[o.image], &Itrue[o.hkl], &wedge);
                    break;
                }

                case ScaleMergeOptions::PartialityModel::Still: {
                    auto *cost = new ceres::AutoDiffCostFunction<IntensityStillResidual, 1, 1, 1, 1>(
                        new IntensityStillResidual(*o.r, o.sigma));
                    problem.AddResidualBlock(cost, nullptr, &G[o.image], &R_sq[o.image], &Itrue[o.hkl]);
                    break;
                }

                case ScaleMergeOptions::PartialityModel::Unity: {
                    auto *cost = new ceres::AutoDiffCostFunction<IntensityFixedResidual, 1, 1, 1>(
                        new IntensityFixedResidual(*o.r, o.sigma, 1.0));
                    problem.AddResidualBlock(cost, nullptr, &G[o.image], &Itrue[o.hkl]);
                    break;
                }

                case ScaleMergeOptions::PartialityModel::Fixed: {
                    auto *cost = new ceres::AutoDiffCostFunction<IntensityFixedResidual, 1, 1, 1>(
                        new IntensityFixedResidual(*o.r, o.sigma, o.r->partiality));
                    problem.AddResidualBlock(cost, nullptr, &G[o.image], &Itrue[o.hkl]);
                    break;
                }
            }
        }

        for (int i = 0; i < static_cast<int>(G.size()); ++i) {
            if (!image_used[i])
                continue;

            problem.SetParameterLowerBound(&G[i], 0, 1e-12);

            if (opt.regularize_scale_to_one) {
                auto *cost = new ceres::AutoDiffCostFunction<ScaleRegularizationResidual, 1, 1>(
                    new ScaleRegularizationResidual(opt.scale_regularization_sigma));
                problem.AddResidualBlock(cost, nullptr, &G[i]);
            }
        }

        if (opt.smoothen_g) {
            for (int i = 0; i + 2 < static_cast<int>(G.size()); ++i) {
                if (!(image_used[i] && image_used[i + 1] && image_used[i + 2]))
                    continue;

                auto *cost = new ceres::AutoDiffCostFunction<SmoothnessResidual, 1, 1, 1, 1>(
                    new SmoothnessResidual(0.05));
                problem.AddResidualBlock(cost, nullptr, &G[i], &G[i + 1], &G[i + 2]);
            }
        }

        if (opt.partiality_model == ScaleMergeOptions::PartialityModel::Rotation) {
            for (int i = 0; i < static_cast<int>(mosaicity.size()); ++i) {
                if (!image_used[i])
                    continue;

                problem.SetParameterLowerBound(&mosaicity[i], 0, opt.mosaicity_min_deg);
                problem.SetParameterUpperBound(&mosaicity[i], 0, opt.mosaicity_max_deg);
            }

            if (opt.smoothen_mos) {
                for (int i = 0; i + 2 < static_cast<int>(mosaicity.size()); ++i) {
                    if (!(image_used[i] && image_used[i + 1] && image_used[i + 2]))
                        continue;

                    auto *cost = new ceres::AutoDiffCostFunction<SmoothnessResidual, 1, 1, 1, 1>(
                        new SmoothnessResidual(0.05));
                    problem.AddResidualBlock(cost, nullptr, &mosaicity[i], &mosaicity[i + 1], &mosaicity[i + 2]);
                }
            }

            if (!opt.refine_wedge)
                problem.SetParameterBlockConstant(&wedge);
            else
                problem.SetParameterLowerBound(&wedge, 0, 0.0);
        }

        if (opt.partiality_model == ScaleMergeOptions::PartialityModel::Still) {
            for (int i = 0; i < static_cast<int>(R_sq.size()); ++i) {
                if (!image_used[i])
                    continue;

                problem.SetParameterLowerBound(&R_sq[i], 0, 1e-9);
                problem.SetParameterUpperBound(&R_sq[i], 0, 1.0);
            }
        }

        unsigned int hw = std::thread::hardware_concurrency();
        if (hw == 0)
            hw = 1;

        ceres::Solver::Options options;
        options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;
        options.minimizer_progress_to_stdout = true;
        options.max_num_iterations = opt.max_num_iterations;
        options.max_solver_time_in_seconds = opt.max_solver_time_s;
        options.num_threads = static_cast<int>(hw);
        options.function_tolerance = 1e-4;

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

        std::cout << summary.FullReport() << std::endl;
    }

    double Partiality(const Reflection &r,
                      const ScaleMergeOptions &opt,
                      int image_slot,
                      const std::vector<double> &mosaicity,
                      const std::vector<double> &R_sq) {
        switch (opt.partiality_model) {
            case ScaleMergeOptions::PartialityModel::Fixed:
                return r.partiality;

            case ScaleMergeOptions::PartialityModel::Unity:
                return 1.0;

            case ScaleMergeOptions::PartialityModel::Rotation: {
                const double half_wedge = opt.wedge_deg.value_or(0.0) / 2.0;
                const double c1 = r.zeta / std::sqrt(2.0);
                const double arg_plus = (r.delta_phi_deg + half_wedge) * c1 / mosaicity[image_slot];
                const double arg_minus = (r.delta_phi_deg - half_wedge) * c1 / mosaicity[image_slot];

                return (std::erf(arg_plus) - std::erf(arg_minus)) / 2.0;
            }

            case ScaleMergeOptions::PartialityModel::Still:
                return std::exp(-r.dist_ewald * r.dist_ewald / R_sq[image_slot]);
        }

        return 1.0;
    }

    void CalcCorrections(std::vector<std::vector<Reflection>> &observations,
                                          const ScaleMergeOptions &opt,
                                          const std::vector<double> &G,
                                          const std::vector<double> &mosaicity,
                                          const std::vector<double> &R_sq) {

        size_t nrefl = 0;
        for (const auto &image: observations)
            nrefl += image.size();

        for (int image = 0; image < static_cast<int>(observations.size()); ++image) {
            const int image_slot = image / static_cast<int>(opt.image_cluster);

            for (auto &r: observations[image]) {
                if (!AcceptReflectionForScaling(r, opt)) {
                    r.scaling_correction = 0.0;
                    continue;
                }

                const double partiality = Partiality(r, opt, image_slot, mosaicity, R_sq);
                if (!std::isfinite(partiality) || partiality < 0.01) {
                    r.partiality = 0.0;
                    r.scaling_correction = 0.0;
                } else {
                    const double denom = G[image_slot] * partiality;
                    const double correction = denom > 0.0 ? r.rlp / denom : 0.0;
                    r.partiality = partiality;
                    r.scaling_correction = std::isfinite(correction) ? correction : 0.0;
                }
            }
        }
    }
}

ScaleMergeResult ScaleAndMergeReflectionsCeres(std::vector<std::vector<Reflection>> &observations,
                                               const ScaleMergeOptions &opt) {
    if (opt.image_cluster <= 0)
        throw std::invalid_argument("image_cluster must be positive");

    const size_t n_image_slots = observations.size() / opt.image_cluster +
                                 (observations.size() % opt.image_cluster > 0 ? 1 : 0);

    std::vector<uint8_t> image_used(n_image_slots, 0);

    int nhkl = 0;
    auto scale_obs = BuildScaleObs(observations, opt, image_used, nhkl);

    std::vector<double> G(n_image_slots, 1.0);
    std::vector<double> mosaicity(n_image_slots, opt.mosaicity_init_deg);
    std::vector<double> R_sq(n_image_slots, 0.001 * 0.001);

    for (int i = 0; i < static_cast<int>(n_image_slots); ++i) {
        if (!image_used[i]) {
            G[i] = NAN;
            mosaicity[i] = NAN;
            R_sq[i] = NAN;
        } else if (opt.mosaicity_init_deg_vec.size() > static_cast<size_t>(i) &&
                   std::isfinite(opt.mosaicity_init_deg_vec[i])) {
            mosaicity[i] = opt.mosaicity_init_deg_vec[i];
        }
    }

    Scale(opt, scale_obs, image_used, nhkl, G, mosaicity, R_sq);

    CalcCorrections(observations, opt, G, mosaicity, R_sq);

    auto out = MergeReflections(observations, opt);

    out.image_scale_g.resize(observations.size(), NAN);
    out.mosaicity_deg.resize(observations.size(), NAN);

    for (int image = 0; image < static_cast<int>(observations.size()); ++image) {
        const int image_slot = image / static_cast<int>(opt.image_cluster);

        if (image_slot < static_cast<int>(image_used.size()) && image_used[image_slot]) {
            out.image_scale_g[image] = G[image_slot];
            out.mosaicity_deg[image] = mosaicity[image_slot];
        }
    }

    return out;
}