// Copyright (2019-2023) Paul Scherrer Institute

#include "RadialIntegration.h"
#include "../common/JFJochException.h"

RadialIntegration::RadialIntegration(const std::vector<uint16_t>& in_mapping, uint32_t in_nbins, uint32_t in_pixel_split) :
        pixel_to_bin(in_mapping), nbins(in_nbins), sum(in_nbins, 0), count(in_nbins, 0),
        pixel_split(in_pixel_split)
{

    if (pixel_split == 4) {
        if (pixel_to_bin.size() % 4 != 0)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "With pixel split of 4 input array must be of size multiple of 4");
    } else if (pixel_split != 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Only pixel split of 1 and 4 allowed at the moment for radial integration");

    coeff = (float *) std::aligned_alloc(64, in_mapping.size() * sizeof(float));
    if (coeff == nullptr)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Memory allocation error");

    for (int i = 0; i < in_mapping.size(); i++)
        coeff[i] = 1.0f;
}

RadialIntegration::RadialIntegration(const RadialIntegrationMapping &mapping, uint32_t in_pixel_split) :
        RadialIntegration(mapping.GetPixelToBinMapping(), mapping.GetBinNumber(), in_pixel_split)
{}

RadialIntegration::~RadialIntegration() {
    std::free(coeff);
}

void RadialIntegration::LoadRadialIntegrationCorr(const std::vector<float> &v) {
    if (v.size() != pixel_to_bin.size() * pixel_split)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in size of pixel-to-bin mapping and correction");
    memcpy(coeff, v.data(), pixel_to_bin.size() * pixel_split * sizeof(float));
}

void RadialIntegration::Clear() {
    for (auto &i : sum)
        i = 0;

    for (auto &i : count)
        i = 0;
}

void RadialIntegration::Process(const int16_t *__restrict data, size_t npixel) {
    if (npixel != pixel_to_bin.size() / pixel_split)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in size of pixel-to-bin mapping and image");

    const auto coeff_aligned = std::assume_aligned<64>(coeff);

    if (pixel_split == 1) {
        for (int i = 0; i < npixel; i++) {
            auto value = data[i];
            if ((value > INT16_MIN + 4) && (value < INT16_MAX - 4)) {
                auto bin = pixel_to_bin[i];
                if (bin < nbins) {
                    sum[bin] += coeff_aligned[i] * value;
                    count[bin] += 1.0f;
                }
            }
        }
    } else if (pixel_split == 4) {
        for (int i = 0; i < npixel; i++) {
            auto value = data[i];
            if ((value > INT16_MIN + 4) && (value < INT16_MAX - 4)) {
                for (int p = 0; p < 4; p++) {
                    auto bin = pixel_to_bin[i * pixel_split + p];
                    if (bin < nbins) {
                        sum[bin] += coeff_aligned[i * pixel_split + p] * value * 0.25f;
                        count[bin] += 0.25f;
                    }
                }
            }
        }
    }
}

void RadialIntegration::ProcessOneImage(const int16_t *data, size_t npixel) {
    Clear();
    Process(data, npixel);
}

void RadialIntegration::GetResult(std::vector<float> &result) const {
    result.resize(nbins);

    for (int i = 0; i < nbins; i++) {
        if (count[i] > 0)
            result[i] = static_cast<float>(sum[i]) / static_cast<float>(count[i]);
        else
            result[i] = 0;
    }
}

const std::vector<float> &RadialIntegration::GetSum() const {
    return sum;
}

const std::vector<float> &RadialIntegration::GetCount() const {
    return count;
}

float RadialIntegration::GetRangeValue(uint32_t min_bin, uint32_t max_bin) {
    float ret_sum = 0;
    float ret_count = 0;

    for (int i = std::min(nbins,min_bin); i <= std::min(nbins-1,max_bin); i++) {
        ret_sum += sum[i];
        ret_count += count[i];
    }
    if (ret_count == 0)
        return 0;
    else
        return static_cast<float>(ret_sum) / static_cast<float>(ret_count);
}