// Copyright (2019-2024) Paul Scherrer Institute

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

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

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

void PixelMask::LoadMask(const std::vector<uint32_t> &input_mask, uint8_t bit) {
    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);
        else
            mask[i] &= ~(1<<bit);
    }
}

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

std::vector<uint32_t> PixelMask::GetMask(const DiffractionExperiment &experiment, bool conv) const {
    std::vector<uint32_t> mask_out = mask;

    // 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 (experiment.GetMaskModuleEdges()) {
                    if ((line == 0)
                        || (line == RAW_MODULE_LINES - 1)
                        || (col == 0)
                        || (col == RAW_MODULE_COLS - 1))
                        mask_out[pixel] |= (1 << ModuleEdgePixelBit);
                }

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

    if (conv && 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_out.data());
        return tmp;
    } else
        return mask_out;
}

std::vector<uint32_t> PixelMask::GetUserMask(const DiffractionExperiment &experiment, bool conv) const {
    std::vector<uint32_t> ret = GetMask(experiment, conv);
    for (auto &i: ret)
        i = ((i & (1 << UserMaskedPixelBit)) != 0) ? 1 : 0;
    return ret;
}

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) {
        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());
        LoadMask(raw_mask, UserMaskedPixelBit);
    } else
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Size of input user mask invalid");
}