// Copyright (2019-2023) Paul Scherrer Institute

#include "JFCalibration.h"

#include <cstring>

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),
        mask(in_nmodules * RAW_MODULE_SIZE) {
    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;
}

uint32_t &JFCalibration::Mask(size_t id) {
    return mask.at(id);
}

const uint32_t &JFCalibration::Mask(size_t id) const {
    return mask.at(id);
}

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

std::vector<uint32_t> JFCalibration::CalculateMask(const DiffractionExperiment &experiment, size_t storage_cell) const {
    if (experiment.GetModulesNum() != nmodules)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in module size");

    std::vector<uint32_t> tmp(RAW_MODULE_SIZE * nmodules, 0);

    for (int64_t module = 0; module < nmodules; module++) {
        auto mask_g0 = Pedestal(module, 0, storage_cell).GetPedestalMask();
        auto mask_g1 = Pedestal(module, 1, storage_cell).GetPedestalMask();
        auto mask_g2 = Pedestal(module, 2, storage_cell).GetPedestalMask();

        for (int64_t line = 0; line < RAW_MODULE_LINES; line++) {
            for (int64_t col = 0; col < RAW_MODULE_COLS; col++) {
                int64_t pixel = module * RAW_MODULE_SIZE + line * RAW_MODULE_COLS + col;

                if (experiment.GetMaskModuleEdges()) {
                    if ((line == 0)
                        || (line == RAW_MODULE_LINES - 1)
                        || (col == 0)
                        || (col == RAW_MODULE_COLS - 1))
                        tmp[pixel] |= (1 << 30);
                }

                if (experiment.GetMaskChipEdges()) {
                    if ((col == 255) || (col == 256)
                        || (col == 511) || (col == 512)
                        || (col == 767) || (col == 768)
                        || (line == 255) || (line== 256))
                        tmp[pixel] |= (1 << 31);
                }

                tmp[pixel] |= mask[pixel];
                tmp[pixel] |= mask_g0[line * RAW_MODULE_COLS + col];
                tmp[pixel] |= mask_g1[line * RAW_MODULE_COLS + col];
                tmp[pixel] |= mask_g2[line * RAW_MODULE_COLS + col];
            }
        }
    }

    return tmp;
}

std::vector<uint32_t> JFCalibration::CalculateNexusMask(const DiffractionExperiment &experiment, size_t storage_cell) const {
    auto mask_raw = CalculateMask(experiment, storage_cell);

    if (experiment.GetDetectorMode() == DetectorMode::Conversion) {
        std::vector<uint32_t> tmp(experiment.GetPixelsNum(), 1);
        RawToConvertedGeometry<uint32_t, uint32_t>(experiment, tmp.data(), mask_raw.data());
        return tmp;
    } else
        return mask_raw;
}

std::vector<uint8_t> JFCalibration::CalculateOneByteMask(const DiffractionExperiment &experiment, size_t storage_cell) const {
    auto nexus_mask = CalculateNexusMask(experiment, storage_cell);

    std::vector<uint8_t> tmp(experiment.GetPixelsNum());
    for (int i = 0; i < experiment.GetPixelsNum(); i++)
        tmp[i] = (nexus_mask[i] == 0);
    return tmp;
}

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

    for (int i = 0; i < RAW_MODULE_SIZE; i++) {
        if ((mask[RAW_MODULE_SIZE * module_number+i] != 0) || (mask_g0[i] != 0) || (mask_g1[i] != 0) || (mask_g2[i] != 0))
            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.masked_pixels = (CountMaskedPixels(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);
}