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

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

JFPedestalCalc::JFPedestalCalc(const DiffractionExperiment &experiment)
: min_images(experiment.GetPedestalMinImageCount()),
currPedestal(RAW_MODULE_SIZE, 0), currPedestal2(RAW_MODULE_SIZE, 0), wrongCount(RAW_MODULE_SIZE, 0) {

    switch (experiment.GetDetectorMode()) {
        case DetectorMode::PedestalG1:
            gain_level = 1;
            break;
        case DetectorMode::PedestalG2:
            gain_level = 2;
            break;
        default:
            gain_level = 0;
            break;
    }
}

template <unsigned int GAIN_BIT> void JFPedestalCalc::AnalyzeImage(const uint16_t *raw_image) {
    for (int i = 0; i < RAW_MODULE_SIZE; i++) {
        uint16_t adc = raw_image[i] & 0x3FFF;
        uint16_t gain = raw_image[i] & 0xC000;

        if (gain != GAIN_BIT)
            wrongCount[i] ++;
        currPedestal[i] += adc;
        currPedestal2[i] += adc * adc;
    }
    image_number++;
}

void JFPedestalCalc::AnalyzeImage(const uint16_t *raw_image) {
    std::unique_lock ul(m);

    switch (gain_level) {
        case 0:
            AnalyzeImage<0>(raw_image);
            break;
        case 1:
            AnalyzeImage<0x4000>(raw_image);
            break;
        case 2:
            AnalyzeImage<0xc000>(raw_image);
            break;
        default:
            break;
    }
}
void JFPedestalCalc::Export(JFModulePedestal &calibration, size_t allowed_wrong_gains) const {
    std::unique_lock ul(m);

    auto calib_pedestal = calibration.GetPedestal();
    auto calib_rms = calibration.GetPedestalRMS();

    for (int i = 0; i < RAW_MODULE_SIZE; i++) {
        float tmp = static_cast<float>(currPedestal[i]) / static_cast<float>(image_number);
        float variance = static_cast<float>(currPedestal2[i]) / static_cast<float>(image_number) - tmp * tmp;
        float rms = sqrtf(variance);

        // Minimally, full window size needs to be recorded to give viable pedestal
        if ((image_number < min_images) || (wrongCount[i] > allowed_wrong_gains)) {
            calib_pedestal[i] = PEDESTAL_WRONG;
            calib_rms[i] = UINT16_MAX;
        } else {
            calib_pedestal[i] = static_cast<uint16_t>(std::lroundf(tmp));
            if (std::lroundf(rms) > UINT16_MAX)
                calib_rms[i] = UINT16_MAX;
            else
                calib_rms[i] = static_cast<uint16_t>(std::lroundf(rms));
        }
    }
    calibration.SetFrameCount(image_number);
}

JFPedestalCalc &JFPedestalCalc::operator+=(const JFPedestalCalc &other) {
    std::unique_lock ul1(m), ul2(other.m);

    if (this != &other) {
        if (this->gain_level != other.gain_level) {
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Gain levels do not match between the two JFPedestalCalc objects.");
        }

        for (size_t i = 0; i < currPedestal.size(); ++i) {
            this->currPedestal[i] += other.currPedestal[i];
            this->currPedestal2[i] += other.currPedestal2[i];
            this->wrongCount[i] += other.wrongCount[i];
        }
        this->image_number += other.image_number;
    }
    return *this;
}