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

#include "Merge.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <map>
#include <stdexcept>
#include <tuple>
#include <unordered_set>

#include "../../common/ResolutionShells.h"
#include "HKLKey.h"

namespace {
    struct Obs {
        const Reflection *r = nullptr;
        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;
    }

    std::vector<Obs> BuildObservations(const std::vector<std::vector<Reflection>> &observations,
                                       const ScaleMergeOptions &opt,
                                       std::vector<HKLKey> &slot_to_hkl) {
        std::map<HKLKey, int> hkl_to_slot;
        std::vector<Obs> out;

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

        for (const auto &image: observations) {
            for (const auto &r: image) {

                if (r.scaling_correction <= 0.0 || !std::isfinite(r.scaling_correction))
                    continue;
                if (!AcceptReflection(r, opt.d_min_limit_A))
                    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>(slot_to_hkl.size());
                    it = hkl_to_slot.emplace(key, slot).first;
                    slot_to_hkl.push_back(key);
                }

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

        return out;
    }

    ScaleMergeResult InitResult(const std::vector<HKLKey> &slot_to_hkl,
                                const std::vector<Obs> &obs) {
        ScaleMergeResult out;
        out.merged.resize(slot_to_hkl.size());

        for (int i = 0; i < static_cast<int>(slot_to_hkl.size()); ++i) {
            out.merged[i].h = slot_to_hkl[i].h;
            out.merged[i].k = slot_to_hkl[i].k;
            out.merged[i].l = slot_to_hkl[i].l;
            out.merged[i].I = 0.0;
            out.merged[i].sigma = 0.0;
            out.merged[i].d = 0.0;
        }

        std::vector<std::vector<double>> d_values(slot_to_hkl.size());
        for (const auto &o: obs) {
            if (std::isfinite(o.r->d) && o.r->d > 0.0f)
                d_values[o.hkl].push_back(o.r->d);
        }

        for (int h = 0; h < static_cast<int>(d_values.size()); ++h) {
            auto &v = d_values[h];
            if (v.empty())
                continue;

            std::nth_element(v.begin(), v.begin() + static_cast<long>(v.size() / 2), v.end());
            out.merged[h].d = v[v.size() / 2];
        }

        return out;
    }

    void Merge(size_t nhkl, ScaleMergeResult &out, const std::vector<Obs> &obs) {
        struct Accum {
            double sum_wI = 0.0;
            double sum_w = 0.0;
            double sum_wsigma2 = 0.0;
        };

        std::vector<Accum> acc(nhkl);

        for (const auto &o: obs) {
            const double I_corr = static_cast<double>(o.r->I) * o.r->scaling_correction;
            const double sigma_corr = o.sigma * o.r->scaling_correction;

            if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0)
                continue;

            // Extra factor o.r->scaling_correction down-weights weak images / low partiality observations.
            const double w = o.r->scaling_correction / (sigma_corr * sigma_corr);

            auto &a = acc[o.hkl];
            a.sum_wI += w * I_corr;
            a.sum_w += w;
            a.sum_wsigma2 += w * w * sigma_corr * sigma_corr;
        }

        for (int h = 0; h < static_cast<int>(nhkl); ++h) {
            const auto &a = acc[h];
            if (a.sum_w <= 0.0)
                continue;

            out.merged[h].I = a.sum_wI / a.sum_w;
            out.merged[h].sigma = std::sqrt(a.sum_wsigma2) / a.sum_w;
        }
    }

    void Stats(const ScaleMergeOptions &opt, ScaleMergeResult &out, const std::vector<Obs> &obs) {
        constexpr int n_shells = 10;

        float d_min = std::numeric_limits<float>::max();
        float d_max = 0.0f;

        for (const auto &m: out.merged) {
            const auto d = static_cast<float>(m.d);
            if (!std::isfinite(d) || d <= 0.0f)
                continue;
            if (opt.d_min_limit_A > 0.0 && d < static_cast<float>(opt.d_min_limit_A))
                continue;

            d_min = std::min(d_min, d);
            d_max = std::max(d_max, d);
        }

        if (!(d_min < d_max && d_min > 0.0f))
            return;

        const float d_min_pad = d_min * 0.999f;
        const float d_max_pad = d_max * 1.001f;

        ResolutionShells shells(d_min_pad, d_max_pad, n_shells);
        const auto shell_mean_1_d2 = shells.GetShellMeanOneOverResSq();
        const auto shell_min_res = shells.GetShellMinRes();

        std::vector<int> hkl_shell(out.merged.size(), -1);
        for (int h = 0; h < static_cast<int>(out.merged.size()); ++h) {
            auto s = shells.GetShell(out.merged[h].d);
            if (s)
                hkl_shell[h] = *s;
        }

        struct PerHKL {
            double sum_I = 0.0;
            std::vector<double> I;
        };

        std::vector<PerHKL> per_hkl(out.merged.size());

        for (const auto &o: obs) {
            if (o.hkl < 0 || o.hkl >= static_cast<int>(per_hkl.size()))
                continue;
            if (hkl_shell[o.hkl] < 0)
                continue;

            const double I_corr = static_cast<double>(o.r->I) * o.r->scaling_correction;
            if (!std::isfinite(I_corr))
                continue;

            per_hkl[o.hkl].sum_I += I_corr;
            per_hkl[o.hkl].I.push_back(I_corr);
        }

        struct ShellAccum {
            int total_obs = 0;
            std::unordered_set<int> unique;
            double rmeas_num = 0.0;
            double rmeas_den = 0.0;
            double sum_i_over_sigma = 0.0;
            int n_i_over_sigma = 0;
        };

        std::vector<ShellAccum> acc(n_shells);

        for (int h = 0; h < static_cast<int>(per_hkl.size()); ++h) {
            const int s = hkl_shell[h];
            if (s < 0 || per_hkl[h].I.empty())
                continue;

            auto &sa = acc[s];
            const auto &ph = per_hkl[h];
            const int n = static_cast<int>(ph.I.size());
            const double mean_I = ph.sum_I / n;

            sa.unique.insert(h);
            sa.total_obs += n;

            if (n >= 2) {
                double sum_abs_dev = 0.0;
                for (double I: ph.I)
                    sum_abs_dev += std::abs(I - mean_I);

                sa.rmeas_num += std::sqrt(static_cast<double>(n) / (n - 1.0)) * sum_abs_dev;
            }

            for (double I: ph.I)
                sa.rmeas_den += std::abs(I);

            if (out.merged[h].sigma > 0.0) {
                sa.sum_i_over_sigma += out.merged[h].I / out.merged[h].sigma;
                ++sa.n_i_over_sigma;
            }
        }

        out.statistics.shells.resize(n_shells);

        for (int s = 0; s < n_shells; ++s) {
            const auto &sa = acc[s];
            auto &ss = out.statistics.shells[s];

            ss.mean_one_over_d2 = shell_mean_1_d2[s];
            ss.d_min = shell_min_res[s];
            ss.d_max = s == 0 ? d_max_pad : shell_min_res[s - 1];
            ss.total_observations = sa.total_obs;
            ss.unique_reflections = static_cast<int>(sa.unique.size());
            ss.rmeas = sa.rmeas_den > 0.0 ? sa.rmeas_num / sa.rmeas_den : 0.0;
            ss.mean_i_over_sigma = sa.n_i_over_sigma > 0
                                        ? sa.sum_i_over_sigma / sa.n_i_over_sigma
                                        : 0.0;
            ss.completeness = 0.0;
            ss.possible_reflections = 0;
        }

        auto &overall = out.statistics.overall;
        overall.d_min = d_min;
        overall.d_max = d_max;

        std::unordered_set<int> all_unique;
        double rmeas_num = 0.0;
        double rmeas_den = 0.0;
        double sum_i_over_sigma = 0.0;
        int n_i_over_sigma = 0;

        for (const auto &sa: acc) {
            overall.total_observations += sa.total_obs;
            all_unique.insert(sa.unique.begin(), sa.unique.end());
            rmeas_num += sa.rmeas_num;
            rmeas_den += sa.rmeas_den;
            sum_i_over_sigma += sa.sum_i_over_sigma;
            n_i_over_sigma += sa.n_i_over_sigma;
        }

        overall.unique_reflections = static_cast<int>(all_unique.size());
        overall.rmeas = rmeas_den > 0.0 ? rmeas_num / rmeas_den : 0.0;
        overall.mean_i_over_sigma = n_i_over_sigma > 0 ? sum_i_over_sigma / n_i_over_sigma : 0.0;
        overall.completeness = 0.0;
        overall.possible_reflections = 0;
    }
}

ScaleMergeResult MergeReflections(const std::vector<std::vector<Reflection>> &observations,
                                  const ScaleMergeOptions &opt) {
    std::vector<HKLKey> slot_to_hkl;
    auto obs = BuildObservations(observations, opt, slot_to_hkl);

    auto out = InitResult(slot_to_hkl, obs);

    Merge(slot_to_hkl.size(), out, obs);
    Stats(opt, out, obs);

    return out;
}