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

#include "GlobalScale.h"

#include <ceres/ceres.h>

#include <thread>
#include <iostream>
#include <algorithm>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <tuple>
#include <unordered_map>
#include <vector>

namespace {
    struct HKLKey {
        int64_t h = 0;
        int64_t k = 0;
        int64_t l = 0;
        bool is_positive = true; // only relevant if opt.merge_friedel == false

        bool operator==(const HKLKey &o) const noexcept {
            return h == o.h && k == o.k && l == o.l && is_positive == o.is_positive;
        }
    };

    struct HKLKeyHash {
        size_t operator()(const HKLKey &key) const noexcept {
            auto mix = [](uint64_t x) {
                x ^= x >> 33;
                x *= 0xff51afd7ed558ccdULL;
                x ^= x >> 33;
                x *= 0xc4ceb9fe1a85ec53ULL;
                x ^= x >> 33;
                return x;
            };
            const uint64_t a = static_cast<uint64_t>(key.h);
            const uint64_t b = static_cast<uint64_t>(key.k);
            const uint64_t c = static_cast<uint64_t>(key.l);
            const uint64_t d = static_cast<uint64_t>(key.is_positive ? 1 : 0);
            return static_cast<size_t>(mix(a) ^ (mix(b) << 1) ^ (mix(c) << 2) ^ (mix(d) << 3));
        }
    };

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

    inline double SafeD(double d) {
        if (!std::isfinite(d) || d <= 0.0)
            return std::numeric_limits<double>::quiet_NaN();
        return d;
    }

    inline int SafeToInt(int64_t x) {
        if (x < std::numeric_limits<int>::min() || x > std::numeric_limits<int>::max())
            throw std::out_of_range("HKL index out of int range for Gemmi");
        return static_cast<int>(x);
    }

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

    inline HKLKey CanonicalizeHKLKey(const Reflection &r, const GlobalScaleOptions &opt) {
        HKLKey key{};
        key.h = r.h;
        key.k = r.k;
        key.l = r.l;
        key.is_positive = true;

        if (!opt.space_group.has_value()) {
            if (!opt.merge_friedel) {
                const HKLKey neg{-r.h, -r.k, -r.l, true};
                const bool pos = std::tie(key.h, key.k, key.l) >= std::tie(neg.h, neg.k, neg.l);
                if (!pos) {
                    key.h = -key.h;
                    key.k = -key.k;
                    key.l = -key.l;
                    key.is_positive = false;
                }
            }
            return key;
        }

        const gemmi::SpaceGroup &sg = *opt.space_group;
        const gemmi::GroupOps gops = sg.operations();
        const gemmi::ReciprocalAsu rasu(&sg);

        const gemmi::Op::Miller in{{SafeToInt(r.h), SafeToInt(r.k), SafeToInt(r.l)}};
        const auto [asu_hkl, sign_plus] = rasu.to_asu_sign(in, gops);

        key.h = asu_hkl[0];
        key.k = asu_hkl[1];
        key.l = asu_hkl[2];
        key.is_positive = opt.merge_friedel ? true : sign_plus;
        return key;
    }

    struct IntensityRotResidual {
        IntensityRotResidual(const Reflection &r, double sigma_obs, double wedge_deg)
            : Iobs_(static_cast<double>(r.I)),
              weight_(SafeInv(sigma_obs, 1.0)),
              delta_phi_(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 {
            T partiality;
            if (mosaicity[0] >= 0.0) {
                const T half_wedge = wedge[0] / T(2.0);
                const T arg_plus = T(delta_phi_ + half_wedge) * T(c1_) / mosaicity[0];
                const T arg_minus = T(delta_phi_ - half_wedge) * T(c1_) / mosaicity[0];
                partiality = (ceres::erf(arg_plus) - ceres::erf(arg_minus)) / T(2.0);
            } else
                partiality = T(1.0);

            const T Ipred = G[0] * partiality * T(lp_) * Itrue[0];
            residual[0] = (Ipred - T(Iobs_)) * T(weight_);
            return true;
        }

        double Iobs_;
        double weight_;
        double delta_phi_;
        double lp_;
        double c1_;
    };

    struct IntensityStillResidual {
        IntensityStillResidual(const Reflection &r, double sigma_obs)
            : Iobs_(static_cast<double>(r.I)),
              weight_(SafeInv(sigma_obs, 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,
                        const T *const Itrue,
                        T *residual) const {
            const T partiality = ceres::exp(-T(dist_ewald_sq_)/R[0]);
            const T Ipred = G[0] * partiality * T(lp_) * Itrue[0];
            residual[0] = (Ipred - T(Iobs_)) * T(weight_);
            return true;
        }

        double Iobs_;
        double weight_;
        double lp_;
        double dist_ewald_sq_;
    };

    struct IntensityFixedResidual {
        IntensityFixedResidual(const Reflection &r, double sigma_obs, double partiality)
            : Iobs_(static_cast<double>(r.I)),
              weight_(SafeInv(sigma_obs, 1.0)),
              corr_(partiality * SafeInv(r.rlp, 1.0))
        {}

        template<typename T>
        bool operator()(const T *const G,
                        const T *const Itrue,
                        T *residual) const {
            const T Ipred = T(corr_) * G[0] * Itrue[0];
            residual[0] = (Ipred - T(Iobs_)) * T(weight_);
            return true;
        }

        double Iobs_;
        double weight_;
        double corr_;
    };

    struct ScaleRegularizationResidual {
        explicit ScaleRegularizationResidual(double sigma_k)
            : inv_sigma_(SafeInv(sigma_k, 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 SmoothnessRegularizationResidual {
        explicit SmoothnessRegularizationResidual(double sigma)
            : inv_sigma_(SafeInv(sigma, 1.0)) {
        }

        template<typename T>
        bool operator()(const T *const k0,
                        const T *const k1,
                        const T *const k2,
                        T *residual) const {
            residual[0] = (ceres::log(k0[0]) + ceres::log(k2[0]) - T(2.0) * ceres::log(k1[0])) * T(inv_sigma_);
            return true;
        }

        double inv_sigma_;
    };

    struct ObsRef {
        const Reflection *r = nullptr;
        int img_id = 0;
        int hkl_slot = -1;
        double sigma = 0.0;
    };

    void scale(const GlobalScaleOptions &opt,
               std::vector<double> &g,
               std::vector<double> &mosaicity,
               std::vector<double> &R_sq,
               const std::vector<uint8_t> &image_slot_used,
               bool rotation_crystallography,
               size_t nhkl,
               const std::vector<ObsRef> &obs) {
        ceres::Problem problem;

        std::vector<double> Itrue(nhkl, 0.0);

        // Initialize Itrue from per-HKL median of observed intensities
        {
            std::vector<std::vector<double> > per_hkl_I(nhkl);
            for (const auto &o: obs) {
                per_hkl_I[o.hkl_slot].push_back(static_cast<double>(o.r->I));
            }
            for (int h = 0; h < nhkl; ++h) {
                auto &v = per_hkl_I[h];
                if (v.empty()) {
                    Itrue[h] = std::max(opt.min_sigma, 1e-6);
                    continue;
                }
                std::nth_element(v.begin(), v.begin() + static_cast<long>(v.size() / 2), v.end());
                double med = v[v.size() / 2];
                if (!std::isfinite(med) || med <= opt.min_sigma)
                    med = opt.min_sigma;
                Itrue[h] = med;
            }
        }

        double wedge = opt.wedge_deg.value_or(0.0);

        for (const auto &o: obs) {
            switch (opt.partiality_model) {
                case GlobalScaleOptions::PartialityModel::Rotation: {
                    auto *cost = new ceres::AutoDiffCostFunction<IntensityRotResidual, 1, 1, 1, 1, 1>(
                        new IntensityRotResidual(*o.r, o.sigma, opt.wedge_deg.value_or(0.0)));
                    problem.AddResidualBlock(cost,
                                             nullptr,
                                             &g[o.img_id],
                                             &mosaicity[o.img_id],
                                             &Itrue[o.hkl_slot],
                                             &wedge);
                }
                break;
                case GlobalScaleOptions::PartialityModel::Still: {
                    auto *cost = new ceres::AutoDiffCostFunction<IntensityStillResidual, 1, 1, 1, 1>(
                        new IntensityStillResidual(*o.r, o.sigma));
                    problem.AddResidualBlock(cost,
                                             nullptr,
                                             &g[o.img_id],
                                             &R_sq[o.img_id],
                                             &Itrue[o.hkl_slot]);
                }
                break;
                case GlobalScaleOptions::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.img_id],
                                             &Itrue[o.hkl_slot]);
                }
                break;
                case GlobalScaleOptions::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.img_id],
                                             &Itrue[o.hkl_slot]);
                }
                break;
            }
        }

        for (int i = 0; i < g.size(); ++i) {
            if (image_slot_used[i]) {
                auto *cost = new ceres::AutoDiffCostFunction<ScaleRegularizationResidual, 1, 1>(
                    new ScaleRegularizationResidual(0.05));
                problem.AddResidualBlock(cost, nullptr, &g[i]);
            }
        }

        if (rotation_crystallography) {
            if (opt.smoothen_g) {
                for (int i = 0; i < g.size() - 2; ++i) {
                    if (image_slot_used[i] && image_slot_used[i + 1] && image_slot_used[i + 2]) {
                        auto *cost = new ceres::AutoDiffCostFunction<SmoothnessRegularizationResidual, 1, 1, 1, 1>(
                            new SmoothnessRegularizationResidual(0.05));

                        problem.AddResidualBlock(cost, nullptr, &g[i], &g[i + 1], &g[i + 2]);
                    }
                }
            }

            if (opt.smoothen_mos && opt.partiality_model == GlobalScaleOptions::PartialityModel::Rotation) {
                for (int i = 0; i < mosaicity.size() - 2; ++i) {
                    if (image_slot_used[i] && image_slot_used[i + 1] && image_slot_used[i + 2]) {
                        auto *cost = new ceres::AutoDiffCostFunction<SmoothnessRegularizationResidual, 1, 1, 1, 1>(
                            new SmoothnessRegularizationResidual(0.05));
                        problem.AddResidualBlock(cost, nullptr, &mosaicity[i], &mosaicity[i + 1], &mosaicity[i + 2]);
                    }
                }
            }
        }

        if (opt.partiality_model == GlobalScaleOptions::PartialityModel::Still) {
            for (int i = 0; i < R_sq.size(); ++i) {
                if (image_slot_used[i]) {
                    problem.SetParameterLowerBound(&R_sq[i], 0, 1e-9);
                    problem.SetParameterUpperBound(&R_sq[i], 0, 1.0);
                }
            }
        }

        // Scaling factors must be always positive
        for (int i = 0; i < g.size(); i++) {
            if (image_slot_used[i])
                problem.SetParameterLowerBound(&g[i], 0, 1e-12);
        }

        // Mosaicity refinement + bounds
        if (opt.partiality_model == GlobalScaleOptions::PartialityModel::Rotation) {
            for (int i = 0; i < mosaicity.size(); ++i) {
                if (image_slot_used[i]) {
                    problem.SetParameterLowerBound(&mosaicity[i], 0, opt.mosaicity_min_deg);
                    problem.SetParameterUpperBound(&mosaicity[i], 0, opt.mosaicity_max_deg);
                }
            }
            if (!opt.refine_wedge)
                problem.SetParameterBlockConstant(&wedge);
            else
                problem.SetParameterLowerBound(&wedge, 0, 0.0);
        }

        // use all available threads
        unsigned int hw = std::thread::hardware_concurrency();
        if (hw == 0)
            hw = 1; // fallback

        ceres::Solver::Options options;

        options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;
        options.minimizer_progress_to_stdout = false;
        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.BriefReport() << std::endl;
    }

    void proc_obs(const std::vector<std::vector<Reflection> > &observations,
        const GlobalScaleOptions &opt,
        std::vector<uint8_t> &image_slot_used,
        std::vector<ObsRef> &obs,
        std::unordered_map<HKLKey, int, HKLKeyHash> &hklToSlot
        ) {
        for (int i = 0; i < observations.size(); i++) {
            for (const auto &r: observations[i]) {
                const double d = SafeD(r.d);
                if (!std::isfinite(d))
                    continue;
                if (!std::isfinite(r.I))
                    continue;

                if (opt.d_min_limit_A > 0.0 && d < opt.d_min_limit_A)
                    continue;

                if (!std::isfinite(r.zeta) || r.zeta <= 0.0f)
                    continue;
                if (!std::isfinite(r.rlp) || r.rlp == 0.0f)
                    continue;

                const double sigma = SafeSigma(r.sigma, opt.min_sigma);

                const int img_id = i / opt.image_cluster;
                image_slot_used[img_id] = 1;

                int hkl_slot;
                try {
                    const HKLKey key = CanonicalizeHKLKey(r, opt);
                    auto it = hklToSlot.find(key);
                    if (it == hklToSlot.end()) {
                        hkl_slot = static_cast<int>(hklToSlot.size());
                        hklToSlot.emplace(key, hkl_slot);
                    } else {
                        hkl_slot = it->second;
                    }
                } catch (...) {
                    continue;
                }

                ObsRef o;
                o.r = &r;
                o.img_id = img_id;
                o.hkl_slot = hkl_slot;
                o.sigma = sigma;
                obs.push_back(o);
            }
        }
    }
} // namespace

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

    const bool rotation_crystallography = opt.wedge_deg.has_value();

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

    std::vector<ObsRef> obs;
    obs.reserve(nrefl);

    std::unordered_map<HKLKey, int, HKLKeyHash> hklToSlot;
    hklToSlot.reserve(nrefl);

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

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

    proc_obs(observations, opt, image_slot_used, obs, hklToSlot);

    const int nhkl = static_cast<int>(hklToSlot.size());

    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 < n_image_slots; i++) {
        if (!image_slot_used[i]) {
            mosaicity[i] = NAN;
            g[i] = NAN;
            R_sq[i] = NAN;
        } else if (opt.mosaicity_init_deg_vec.size() > i && std::isfinite(opt.mosaicity_init_deg_vec[i])) {
            mosaicity[i] = opt.mosaicity_init_deg_vec[i];
        }
    }

    scale(opt, g, mosaicity, R_sq, image_slot_used, rotation_crystallography, nhkl, obs);

    GlobalScaleResult out;
    out.image_scale_g.resize(observations.size(), NAN);
    out.mosaicity_deg.resize(observations.size(), NAN);
    for (int i = 0; i < observations.size(); i++) {
        size_t img_slot = i / opt.image_cluster;
        if (image_slot_used[img_slot]) {
            out.image_scale_g[i] = static_cast<float>(g[img_slot]);
            out.mosaicity_deg[i] = static_cast<float>(mosaicity[img_slot]);
        }
    }
    return out;
}