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

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

std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x, const std::vector<std::vector<Reflection> > &reflections) {
    auto scaling_settings = x.GetScalingSettings();
    HKLKeyGenerator key_generator(scaling_settings.GetMergeFriedel(), x.GetSpaceGroupNumber().value_or(1));
    const std::optional<double> high_resolution_limit = scaling_settings.GetHighResolutionLimit_A();

    struct Accum {
        // Keep anomalous + / - together, but separate
        int32_t h = 0;
        int32_t k = 0;
        int32_t l = 0;
        float d = NAN;
        double sum_wI = 0.0;
        double sum_w = 0.0;
        double sum_wI_anom[2] = {0.0, 0.0};
        double sum_w_anom[2] = {0.0, 0.0};
        bool present[2] = {false, false};
    };

    std::unordered_map<uint64_t, Accum> acc;

    for (const auto &image: reflections) {
        for (const auto &r: image) {
            if (r.scaling_correction <= 0.0 || !std::isfinite(r.scaling_correction))
                continue;
            if (!AcceptReflection(r, high_resolution_limit))
                continue;

            const float I_corr = r.I * r.scaling_correction;
            const float sigma_corr = r.sigma * r.scaling_correction;
            if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0)
                continue;

            auto hkl = key_generator(r);
            auto hkl_key = hkl.pack_no_anom();

            auto it = acc.find(hkl_key);
            if (it == acc.end())
                it = acc.emplace(hkl_key, Accum{
                .h = hkl.h,
                .k = hkl.k,
                .l = hkl.l
                }).first;

            int solution = hkl.plus ? 0 : 1;

            const float w = 1.0f / (sigma_corr * sigma_corr);
            it->second.sum_wI += w * I_corr;
            it->second.sum_w += w;
            it->second.sum_wI_anom[solution] += w * I_corr;
            it->second.sum_w_anom[solution] += w;

            it->second.present[solution] = true;
            if (!std::isfinite(it->second.d) && std::isfinite(r.d) && r.d > 0.0f)
                it->second.d = r.d;
        }
    }

    std::vector<MergedReflection> out;
    out.reserve(acc.size());
    for (const auto &[key, accum]: acc) {
        if (accum.sum_w <= 0.0)
            continue;

        float I_anom[2] = {NAN, NAN};
        float sigma_anom[2] = {NAN, NAN};

        for (int i = 0; i < 2; ++i) {
            if (accum.present[i] && accum.sum_w_anom[i] > 0.0) {
                I_anom[i] = static_cast<float>(accum.sum_wI_anom[i] / accum.sum_w_anom[i]);
                sigma_anom[i] = 1.0f / std::sqrt(static_cast<float>(accum.sum_w_anom[i]));
            }
        }
        out.emplace_back(MergedReflection{
            .h = accum.h,
            .k = accum.k,
            .l = accum.l,
            .I = static_cast<float>(accum.sum_wI / accum.sum_w),
            .sigma = 1.0f / std::sqrt(static_cast<float>(accum.sum_w)),
            .I_anom = {I_anom[0], I_anom[1]},
            .sigma_anom = {sigma_anom[0], sigma_anom[1]},
            .d = accum.d
        });
    }
    return out;
}

MergeStatistics MergeStats(const DiffractionExperiment &x,
    const std::vector<MergedReflection> &merged,
    const std::vector<std::vector<Reflection> > &reflections) {

    constexpr int n_shells = 10;

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

    auto d_min_limit_A = x.GetScalingSettings().GetHighResolutionLimit_A();
    for (const auto &m: merged) {
        if (!std::isfinite(m.d) || m.d <= 0.0f)
            continue;
        if (d_min_limit_A && m.d < d_min_limit_A)
            continue;

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

    if (!(d_min < d_max && d_min > 0.0f))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Error in resolution calculation");

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

    struct ShellAccum {
        int total_obs = 0;
        int unique = 0;
        double sum_i_over_sigma = 0.0;
        int n_i_over_sigma = 0;
    };

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

    for (const auto &m: merged) {
        const auto shell = shells.GetShell(m.d);
        if (!shell.has_value())
            continue;

        const int s = *shell;
        if (s >= 0 && s < n_shells) {
            if (std::isfinite(m.I) && std::isfinite(m.sigma) && m.sigma > 0.0) {
                acc[s].unique++;
                acc[s].sum_i_over_sigma += m.I / m.sigma;
                ++acc[s].n_i_over_sigma;
            }

        }
    }

    for (const auto &image: reflections) {
        for (const auto &r: image) {
            if (r.scaling_correction <= 0.0 || !std::isfinite(r.scaling_correction))
                continue;

            if (!AcceptReflection(r, d_min_limit_A))
                continue;

            const auto shell = shells.GetShell(r.d);
            if (!shell.has_value())
                continue;
            const int s = *shell;
            if (s >= 0 && s < n_shells)
                acc[s].total_obs++;
        }
    }

    MergeStatistics out;
    out.shells.resize(n_shells);

    for (int s = 0; s < n_shells; ++s) {
        const auto &sa = acc[s];
        auto &ss = out.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 = sa.unique;
        ss.mean_i_over_sigma = sa.n_i_over_sigma > 0
                                   ? sa.sum_i_over_sigma / sa.n_i_over_sigma
                                   : 0.0;
    }

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

    int all_unique = 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 += sa.unique;
        sum_i_over_sigma += sa.sum_i_over_sigma;
        n_i_over_sigma += sa.n_i_over_sigma;
    }

    overall.unique_reflections = all_unique;
    overall.mean_i_over_sigma = n_i_over_sigma > 0 ? sum_i_over_sigma / n_i_over_sigma : 0.0;

    return out;
}

void MergeStatistics::Print(Logger &logger) const {
    logger.Info("");
    logger.Info("  {:>8s} {:>8s} {:>8s} {:>8s}", "d_min", "N_obs", "N_uniq", "<I/sig>");
    logger.Info("  {:->8s} {:->8s} {:->8s} {:->8s}", "", "", "", "");
    for (const auto &sh: shells) {
        if (sh.unique_reflections == 0)
            continue;
        logger.Info("  {:8.2f} {:8d} {:8d} {:8.1f}",
                    sh.d_min, sh.total_observations, sh.unique_reflections,
                    sh.mean_i_over_sigma);
    }
    {
        const auto &ov = overall;
        logger.Info("  {:->8s} {:->8s} {:->8s} {:->8s}", "", "", "", "");
        logger.Info("  {:>8s} {:8d} {:8d} {:8.1f}",
                    "Overall", ov.total_observations, ov.unique_reflections,
                    ov.mean_i_over_sigma);
    }
    logger.Info("");
}