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

#include "XdsIntegrateParser.h"

#include <fstream>
#include <sstream>
#include <stdexcept>

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

IntegrateMap ParseXdsIntegrateHkl(const std::string& filename) {
    std::ifstream in(filename);
    if (!in.is_open()) {
        throw std::runtime_error("Failed to open INTEGRATE.HKL file: " + filename);
    }

    IntegrateMap result;
    std::string line;

    while (std::getline(in, line)) {
        if (line.empty()) continue;
        char c0 = line.front();
        if (c0 == '!' || c0 == '#') continue;

        std::istringstream iss(line);
        int32_t h, k, l;
        double I, sigma;

        if (!(iss >> l >> k >> h >> I >> sigma)) {
            continue;
        }

        HKLData entry;
        entry.h = h;
        entry.k = k;
        entry.l = l;
        entry.I = I;
        entry.sigma = sigma;

        double v = 0.0;
        while (iss >> v) {
            entry.tail.push_back(v);
        }

        result[hkl_key_16(h, k, l)].push_back(std::move(entry));
    }

    return result;
}


namespace {
double SumIntensity(const std::vector<HKLData>& v) {
    double s = 0.0;
    for (const auto& e : v)
        s += e.I;
    return s;
}
} // namespace

CcHalfByResolutionResult ComputeCcByResolution(
    const CrystalLattice& lattice,
    const IntegrateMap& ours,
    const IntegrateMap& xds,
    float d_min,
    float d_max,
    int32_t nshells) {

    ResolutionShells shells(d_min, d_max, nshells);

    std::vector<double> sum_x(nshells, 0.0);
    std::vector<double> sum_y(nshells, 0.0);
    std::vector<double> sum_x2(nshells, 0.0);
    std::vector<double> sum_y2(nshells, 0.0);
    std::vector<double> sum_xy(nshells, 0.0);
    std::vector<int32_t> n(nshells, 0);

    const Coord astar = lattice.Astar();
    const Coord bstar = lattice.Bstar();
    const Coord cstar = lattice.Cstar();

    for (const auto& [key, ours_list] : ours) {
        auto it = xds.find(key);
        if (it == xds.end())
            continue;

        const auto& xds_list = it->second;

        double i_ours = SumIntensity(ours_list);
        double i_xds = SumIntensity(xds_list);

        int64_t h = ours_list.front().h;
        int64_t k = ours_list.front().k;
        int64_t l = ours_list.front().l;

        Coord recip = astar * h + bstar * k + cstar * l;
        double recip_len = std::sqrt(recip.x * recip.x + recip.y * recip.y + recip.z * recip.z);
        if (recip_len <= 0.0)
            continue;

        float d = static_cast<float>(1.0 / recip_len);
        auto shell = shells.GetShell(d);
        if (!shell)
            continue;

        int idx = shell.value();
        n[idx] += 1;
        sum_x[idx] += i_ours;
        sum_y[idx] += i_xds;
        sum_x2[idx] += i_ours * i_ours;
        sum_y2[idx] += i_xds * i_xds;
        sum_xy[idx] += i_ours * i_xds;
    }

    CcHalfByResolutionResult result;
    result.cc.resize(nshells, std::numeric_limits<double>::quiet_NaN());
    result.pairs = n;
    result.shell_mean_one_over_d2 = shells.GetShellMeanOneOverResSq();

    for (int i = 0; i < nshells; ++i) {
        if (n[i] < 2)
            continue;

        double denom_x = n[i] * sum_x2[i] - sum_x[i] * sum_x[i];
        double denom_y = n[i] * sum_y2[i] - sum_y[i] * sum_y[i];
        double denom = std::sqrt(denom_x * denom_y);
        if (denom <= 0.0)
            continue;

        result.cc[i] = (n[i] * sum_xy[i] - sum_x[i] * sum_y[i]) / denom;
    }

    return result;
}