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

#include "JFCalibration.h"

#include <cstring>
#include "../preview/JFJochTIFF.h"

JFCalibration::JFCalibration(size_t in_nmodules, size_t in_nstorage_cells) :
        nmodules(in_nmodules),
        nstorage_cells(in_nstorage_cells),
        pedestal(in_nmodules * in_nstorage_cells * 3),
        gain_calibration(in_nmodules) {
    if (in_nmodules * in_nstorage_cells == 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Size of JFCalibration cannot be 0");
}

JFCalibration::JFCalibration(const DiffractionExperiment &experiment) :
        JFCalibration(experiment.GetModulesNum(), experiment.GetStorageCellNumber()) {}

size_t JFCalibration::GetModulesNum() const {
    return nmodules;
}

size_t JFCalibration::GetStorageCellNum() const {
    return nstorage_cells;
}

JFModulePedestal &JFCalibration::Pedestal(size_t module_number, size_t gain_level, size_t storage_cell) {
    if (gain_level >= 3)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Gain level must be in range 0-2");
    if (module_number >= nmodules)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Module out of bounds");
    if (storage_cell >= nstorage_cells)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Storage cell " + std::to_string(storage_cell) + " out of bounds");
    return pedestal.at((storage_cell * nmodules + module_number) * 3 + gain_level);
}

const JFModulePedestal &JFCalibration::Pedestal(size_t module_number, size_t gain_level, size_t storage_cell) const {
    if (gain_level >= 3)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Gain level must be in range 0-2");
    if (module_number >= nmodules)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Module out of bounds");
    if (storage_cell >= nstorage_cells)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Storage cell " + std::to_string(storage_cell) + " out of bounds");
    return pedestal.at((storage_cell * nmodules + module_number) * 3 + gain_level);
}

int64_t JFCalibration::CountBadPixels(size_t module_number, size_t storage_cell) const {
    int64_t ret = 0;
    auto mask_g0 = Pedestal(module_number, 0, storage_cell).GetPedestal();
    auto mask_g1 = Pedestal(module_number, 1, storage_cell).GetPedestal();
    auto mask_g2 = Pedestal(module_number, 2, storage_cell).GetPedestal();

    for (int i = 0; i < RAW_MODULE_SIZE; i++) {
        if ((mask_g0[i] == UINT16_MAX) || (mask_g1[i] == UINT16_MAX) || (mask_g2[i] == UINT16_MAX))
            ret++;
    }
    return ret;
}

JFCalibrationModuleStatistics JFCalibration::GetModuleStatistics(size_t module_number, size_t storage_cell) const {
    JFCalibrationModuleStatistics ret{};
    ret.module_number = module_number;
    ret.storage_cell_number = storage_cell;
    ret.pedestal_g0_mean = std::round(Pedestal(module_number, 0, storage_cell).Mean());
    ret.pedestal_g1_mean = std::round(Pedestal(module_number, 1, storage_cell).Mean());
    ret.pedestal_g2_mean = std::round(Pedestal(module_number, 2, storage_cell).Mean());
    ret.gain_g0_mean = GainCalibration(module_number).GetG0Mean();
    ret.gain_g1_mean = GainCalibration(module_number).GetG1Mean();
    ret.gain_g2_mean = GainCalibration(module_number).GetG2Mean();
    ret.bad_pixels = (CountBadPixels(module_number, storage_cell));
    return ret;
}

std::vector<JFCalibrationModuleStatistics> JFCalibration::GetModuleStatistics(size_t storage_cell) const {
    std::vector<JFCalibrationModuleStatistics> output;
    for (int i = 0; i < nmodules; i++)
        output.push_back(GetModuleStatistics(i, storage_cell));
    return output;
}

std::vector<JFCalibrationModuleStatistics>JFCalibration::GetModuleStatistics() const {
    std::vector<JFCalibrationModuleStatistics> output;
    for (int s = 0; s < nstorage_cells; s++) {
        for (int i = 0; i < nmodules; i++)
            output.push_back(GetModuleStatistics(i, s));
    }
    return output;
}

std::vector<uint16_t> JFCalibration::GetPedestal(size_t gain_level, size_t storage_cell) const {
    std::vector<uint16_t> ret(nmodules * RAW_MODULE_SIZE);
    for (int m = 0; m < nmodules; m++) {
        memcpy(ret.data() + m * RAW_MODULE_SIZE, Pedestal(m, gain_level, storage_cell).GetPedestal(),
               RAW_MODULE_SIZE * sizeof(uint16_t));
    }
    return ret;
}

JFModuleGainCalibration &JFCalibration::GainCalibration(size_t module_num) {
    if (module_num >= nmodules)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Module out of bounds");
    return gain_calibration.at(module_num);
}

const JFModuleGainCalibration &JFCalibration::GainCalibration(size_t module_num) const {
    if (module_num >= nmodules)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Module out of bounds");
    return gain_calibration.at(module_num);
}

std::vector<uint16_t> JFCalibration::GetPedestalRMS(size_t gain_level, size_t storage_cell) const {
    std::vector<uint16_t> ret(nmodules * RAW_MODULE_SIZE);
    for (int m = 0; m < nmodules; m++) {
        memcpy(ret.data() + m * RAW_MODULE_SIZE, Pedestal(m, gain_level, storage_cell).GetPedestalRMS(),
               RAW_MODULE_SIZE * sizeof(uint16_t));
    }
    return ret;
}