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

#include "PixelMask.h"
#include "RawToConvertedGeometry.h"
#include "JFJochException.h"
#include "JFJochCompressor.h"

PixelMask::PixelMask(const DetectorSetup &detector)
    : nmodules(detector.GetModulesNum()),
      mask(detector.GetModulesNum() * RAW_MODULE_SIZE, 0) {
}

PixelMask::PixelMask(const DiffractionExperiment &experiment)
    : PixelMask(experiment.GetDetectorSetup()) {
    Update(experiment.GetImageFormatSettings());
}

uint32_t PixelMask::LoadMask(const std::vector<uint32_t> &input_mask, uint8_t bit) {
    uint32_t ret = 0;
    if (input_mask.size() != mask.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Input match doesn't fit the detector ");

    for (int i = 0; i < mask.size(); i++) {
        if (input_mask[i] != 0) {
            mask[i] |= (1 << bit);
            ret++;
        } else
            mask[i] &= ~(1 << bit);
    }
    return ret;
}

void PixelMask::LoadDetectorBadPixelMask(const std::vector<uint32_t> &input_mask) {
    LoadMask(input_mask, ErrorPixelBit);
}

const std::vector<uint32_t> &PixelMask::GetMaskRaw() const {
    return mask;
}

std::vector<uint32_t> PixelMask::GetMask(const DiffractionExperiment &experiment) const {
    if (experiment.IsGeometryTransformed()) {
        std::vector<uint32_t> tmp(experiment.GetPixelsNum(), 1 << ModuleGapPixelBit);
        // nonfunctional areas (i.e. gaps) are filled with 1
        RawToConvertedGeometry<uint32_t, uint32_t>(experiment, tmp.data(), mask.data());
        return tmp;
    } else
        return mask;
}

std::vector<uint8_t> PixelMask::GetCompressedMask(const DiffractionExperiment &experiment) const {
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    return compressor.Compress(GetMask(experiment));
}

std::vector<uint32_t> PixelMask::GetUserMask(const DiffractionExperiment &experiment, bool conv) const {
    std::vector<uint32_t> ret;

    if (conv)
        ret = GetMask(experiment);
    else
        ret = GetMaskRaw();

    for (auto &i: ret)
        i = ((i & (1 << UserMaskedPixelBit)) != 0) ? 1 : 0;
    return ret;
}

PixelMaskStatistics PixelMask::GetStatistics() const {
    return statistics;
}

void PixelMask::LoadUserMask(const DiffractionExperiment &experiment,
                             const std::vector<uint32_t> &in_mask) {
    DiffractionExperiment tmp_experiment = experiment;
    tmp_experiment.GeometryTransformation(true);

    if (tmp_experiment.GetModulesNum() != nmodules)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in module size");

    if (in_mask.size() == nmodules * RAW_MODULE_SIZE) {
        statistics.user_mask = LoadMask(in_mask, UserMaskedPixelBit);
    } else if (in_mask.size() == tmp_experiment.GetPixelsNum()) {
        std::vector<uint32_t> raw_mask(nmodules * RAW_MODULE_SIZE);
        ConvertedToRawGeometry(tmp_experiment, raw_mask.data(), in_mask.data());
        statistics.user_mask = LoadMask(raw_mask, UserMaskedPixelBit);
    } else
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Size of input user mask invalid");
}

void PixelMask::LoadDetectorBadPixelMask(const DiffractionExperiment &experiment, const JFCalibration *calib) {
    std::vector<uint32_t> input_mask(experiment.GetModulesNum() * RAW_MODULE_SIZE, 0);
    std::vector<uint32_t> input_mask_rms(experiment.GetModulesNum() * RAW_MODULE_SIZE, 0);

    if (calib != nullptr && experiment.GetStorageCellNumber() == 1) {
        // with SC no point to calculate detector pixel_mask
        auto pedestal_g0 = calib->GetPedestal(0, 0);
        auto pedestal_g0_rms = calib->GetPedestalRMS(0, 0);
        auto pedestal_g1 = calib->GetPedestal(1, 0);
        auto pedestal_g2 = calib->GetPedestal(2, 0);

        for (int i = 0; i < experiment.GetModulesNum() * RAW_MODULE_SIZE; i++) {
            if (pedestal_g1[i] > 16383)
                input_mask[i] = 1;
            if (!experiment.IsFixedGainG1()) {
                if (pedestal_g0[i] >= 16383) {
                    if (experiment.IsMaskPixelsWithoutG0())
                        input_mask[i] = 1;
                } else if (pedestal_g0_rms[i] > experiment.GetImageFormatSettings().GetPedestalG0RMSLimit())
                    input_mask_rms[i] = 1;

                if (pedestal_g2[i] >= 16383)
                    input_mask[i] = 1;
            }
        }
    }

    statistics.wrong_gain = LoadMask(input_mask, ErrorPixelBit);
    statistics.too_high_pedestal_rms = LoadMask(input_mask_rms, TooHighPedestalRMSPixelBit);
    Update(experiment.GetImageFormatSettings());
}

void PixelMask::Update(const ImageFormatSettings &settings) {
    std::vector<uint32_t> module_edge(nmodules * RAW_MODULE_SIZE, 0);
    std::vector<uint32_t> chip_edge(nmodules * RAW_MODULE_SIZE, 0);

    // apply edge mask
    for (int64_t module = 0; module < nmodules; module++) {
        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 (settings.IsMaskModuleEdges())
                    if ((line == 0)
                        || (line == RAW_MODULE_LINES - 1)
                        || (col == 0)
                        || (col == RAW_MODULE_COLS - 1))
                        module_edge[pixel] = 1;

                if (settings.IsMaskChipEdges())
                    if ((col == 255) || (col == 256)
                        || (col == 511) || (col == 512)
                        || (col == 767) || (col == 768)
                        || (line == 255) || (line == 256))
                        chip_edge[pixel] = 1;
            }
        }
    }
    LoadMask(module_edge, ModuleEdgePixelBit);
    LoadMask(chip_edge, ChipGapPixelBit);
}