// 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() = default;

PixelMask::PixelMask(size_t width, size_t height)
        : mask(width*height, 0) {}

PixelMask::PixelMask(const DiffractionExperiment &experiment)
        : PixelMask(experiment.GetXPixelsNumConv(),
                    experiment.GetYPixelsNumConv()) {
    CalcEdgePixels(experiment);
}

PixelMask::PixelMask(const std::vector<uint32_t> &in_mask) : mask(in_mask) {
    statistics.user_mask = 0;
    statistics.wrong_gain = 0;
    statistics.too_high_pedestal_rms = 0;
    for (unsigned int i : mask) {
        if (i & (1 << UserMaskedPixelBit))
            statistics.user_mask++;
        if (i & (1 << ErrorPixelBit))
            statistics.wrong_gain++;
        if (i & (1 << TooHighPedestalRMSPixelBit))
            statistics.too_high_pedestal_rms++;
    }
}

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::CalcEdgePixels(const DiffractionExperiment &experiment)  {
    if (experiment.GetDetectorType() == DetectorType::DECTRIS)
        return;

    size_t nmodules = experiment.GetModulesNum();
    auto settings = experiment.GetImageFormatSettings();

    // Set module gaps to 1
    std::vector<uint32_t> module_gaps(nmodules * RAW_MODULE_SIZE, 0);
    std::vector<uint32_t> module_gaps_conv(experiment.GetPixelsNumConv(), 1);
    RawToConvertedGeometry(experiment, module_gaps_conv.data(), module_gaps.data());
    LoadMask(module_gaps_conv, ModuleGapPixelBit);

    // Calculate module edges and chip edges
    std::vector<uint32_t> module_edge(nmodules * RAW_MODULE_SIZE, 0);
    std::vector<uint32_t> chip_edge(nmodules * RAW_MODULE_SIZE, 0);
    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 ((line == 0)
                    || (line == RAW_MODULE_LINES - 1)
                    || (col == 0)
                    || (col == RAW_MODULE_COLS - 1))
                    module_edge[pixel] = 1;

                if ((col == 255) || (col == 256)
                    || (col == 511) || (col == 512)
                    || (col == 767) || (col == 768)
                    || (line == 255) || (line == 256))
                    chip_edge[pixel] = 1;
            }
        }
    }

    std::vector<uint32_t> module_edge_conv(experiment.GetPixelsNumConv(), 0);
    if (experiment.GetMaskModuleEdges())
        RawToConvertedGeometry(experiment, module_edge_conv.data(), module_edge.data());
    LoadMask(module_edge_conv, ModuleEdgePixelBit);

    std::vector<uint32_t> chip_edge_conv(experiment.GetPixelsNumConv(), 0);
    if (experiment.GetMaskChipEdges())
        RawToConvertedGeometry(experiment, chip_edge_conv.data(), chip_edge.data());
    LoadMask(chip_edge_conv, ChipGapPixelBit);
}

std::vector<uint32_t> PixelMask::GetMaskRaw(const DiffractionExperiment &experiment) const {
    std::vector<uint32_t> ret(experiment.GetModulesNum() * RAW_MODULE_SIZE, 0);
    ConvertedToRawGeometry(experiment, ret.data(), mask.data());
    return ret;
}

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

std::vector<uint32_t> PixelMask::GetMask(const DiffractionExperiment& experiment) const {
    if (experiment.IsGeometryTransformed())
        return GetMask();
    else
        return GetMaskRaw(experiment);
}

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

std::vector<uint32_t> PixelMask::GetUserMask(const DiffractionExperiment& experiment) const {
    if (experiment.IsGeometryTransformed())
        return GetUserMask();
    else {
        std::vector<uint32_t> tmp = GetUserMask();
        std::vector<uint32_t> ret(experiment.GetModulesNum() * RAW_MODULE_SIZE, 0);
        ConvertedToRawGeometry(experiment, ret.data(), tmp.data());
        return ret;
    }
}


void PixelMask::LoadDetectorBadPixelMask(const DiffractionExperiment &experiment, const JFCalibration *calib) {
    if (experiment.GetDetectorType() == DetectorType::DECTRIS) {
        statistics.wrong_gain = 0;
        statistics.too_high_pedestal_rms = 0;
        return;
    }

    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) {
        for (int sc = 0; sc < experiment.GetStorageCellNumber(); sc++) {
            // For multiple SC PixelMask is logical sum of all image masks
            // (this can be too much, but better than too little)
            auto pedestal_g0 = calib->GetPedestal(0, sc);
            auto pedestal_g0_rms = calib->GetPedestalRMS(0, sc);
            auto pedestal_g1 = calib->GetPedestal(1, sc);
            auto pedestal_g2 = calib->GetPedestal(2, sc);

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

    std::vector<uint32_t> input_mask_conv(experiment.GetPixelsNumConv(), 0);
    RawToConvertedGeometry(experiment, input_mask_conv.data(), input_mask.data());
    statistics.wrong_gain = LoadMask(input_mask_conv, ErrorPixelBit);

    std::vector<uint32_t> input_mask_rms_conv(experiment.GetPixelsNumConv(), 0);
    RawToConvertedGeometry(experiment, input_mask_rms_conv.data(), input_mask_rms.data());
    statistics.too_high_pedestal_rms = LoadMask(input_mask_rms_conv, TooHighPedestalRMSPixelBit);

    CalcEdgePixels(experiment);
}


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

void PixelMask::LoadUserMask(const DiffractionExperiment& experiment, const std::vector<uint32_t> &in_mask) {
    if (in_mask.size() == mask.size()) {
        statistics.user_mask = LoadMask(in_mask, UserMaskedPixelBit);
    } else if (in_mask.size() == experiment.GetModulesNum() * RAW_MODULE_SIZE) {
        std::vector<uint32_t> tmp(experiment.GetPixelsNumConv(), 0);
        RawToConvertedGeometry(experiment, tmp.data(), in_mask. data());
        statistics.user_mask = LoadMask(tmp, UserMaskedPixelBit);
    } else
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Size of input user mask invalid");
}

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

    uint32_t user_bitmask = (1 << UserMaskedPixelBit);
    uint32_t det_bitmask = ~user_bitmask;
    uint32_t bad_pixel_bitmask = ~((1 << UserMaskedPixelBit) | (1 << ModuleGapPixelBit) | (1 << ChipGapPixelBit));

    statistics.user_mask = 0;
    statistics.wrong_gain = 0;
    statistics.too_high_pedestal_rms = 0;

    for (int i = 0; i < mask.size(); i++) {
        uint32_t user_mask = (mask[i] & user_bitmask);
        if (user_mask != 0)
            ++statistics.user_mask;
        uint32_t detector_mask = (input_mask[i] & det_bitmask);
        uint32_t bad_pixel = (input_mask[i] & bad_pixel_bitmask);
        if (bad_pixel != 0)
            ++statistics.wrong_gain;

        mask[i] = user_mask | detector_mask;
    }
}