// 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"
#include "../common/hkl_key.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;
        double xcal, ycal, zcal, rlp, peak, corr;
        int32_t maxc;

        if (!(iss >> l >> k >> h >> I >> sigma >> xcal >> ycal >> zcal >> rlp >> peak >> corr >> maxc)) {
            continue;
        }

        HKLData entry;
        entry.h = -h;
        entry.k = k;
        entry.l = l;
        entry.I = I;
        entry.sigma = sigma;
        entry.rlp = rlp;
        entry.image_number = zcal;

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

        result[hkl_key(-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;

        for (const auto &i_x: xds_list) {
            for (const auto &i_o: ours_list) {
                if (std::fabs(i_x.image_number - i_o.image_number) > 30.0)
                    continue;

                int64_t h = i_x.h;
                int64_t k = i_x.k;
                int64_t l = i_x.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_o.I;
                sum_y[idx] += i_x.I;
                sum_x2[idx] += i_o.I * i_o.I;
                sum_y2[idx] += i_x.I * i_x.I;
                sum_xy[idx] += i_o.I * i_x.I;
            }
        }
    }

    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;
}