// Copyright (2019-2023) Paul Scherrer Institute

#include "JFCalibration.h"

#include <bitshuffle/bitshuffle.h>
#include "../compression/JFJochCompressor.h"
#include <nlohmann/json.hpp>

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) {
        DiffractionExperiment x(experiment);
        x.Binning2x2(false);

        std::vector<uint32_t> tmp(x.GetPixelsNum(), 1);
        RawToConvertedGeometry<uint32_t, uint32_t>(x, tmp.data(), mask_raw.data());
        if (experiment.GetBinning2x2()) {
            std::vector<uint32_t> tmp_binned(experiment.GetPixelsNum());
            Bin2x2_or(tmp_binned.data(), tmp.data(), x.GetXPixelsNum(), experiment.GetYPixelsNum());
            return tmp_binned;
        } else
            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;
}

void
JFCalibration::GetModuleStatistics(JFJochProtoBuf::JFCalibrationStatistics &statistics, size_t storage_cell) const {
    for (int i = 0; i < nmodules; i++) {
        auto tmp = statistics.add_module_statistics();
        tmp->set_module_number(i);
        tmp->set_storage_cell_number(storage_cell);
        tmp->set_pedestal_g0_mean(std::round(Pedestal(i, 0, storage_cell).Mean()));
        tmp->set_pedestal_g1_mean(std::round(Pedestal(i, 1, storage_cell).Mean()));
        tmp->set_pedestal_g2_mean(std::round(Pedestal(i, 2, storage_cell).Mean()));
        tmp->set_gain_g0_mean(GainCalibration(i).GetG0Mean());
        tmp->set_gain_g1_mean(GainCalibration(i).GetG1Mean());
        tmp->set_gain_g2_mean(GainCalibration(i).GetG2Mean());
        tmp->set_masked_pixels(CountMaskedPixels(i, storage_cell));
    }
}

JFJochProtoBuf::JFCalibrationStatistics JFCalibration::GetModuleStatistics(size_t storage_cell) const {
    JFJochProtoBuf::JFCalibrationStatistics output;
    GetModuleStatistics(output, storage_cell);
    return output;
}

JFJochProtoBuf::JFCalibrationStatistics JFCalibration::GetModuleStatistics() const {
    JFJochProtoBuf::JFCalibrationStatistics output;
    for (int s = 0; s < nstorage_cells; s++)
        GetModuleStatistics(output, s);
    return output;
}

JFCalibration::JFCalibration(const JFJochProtoBuf::JFCalibration &input) :
        JFCalibration(input.nmodules(), input.nstorage_cells()) {

    if (input.pedestal_size() != pedestal.size())
        throw JFJochException(JFJochExceptionCategory::ZSTDCompressionError,
                              "Mismatch in size of serialized data (pedestal count)");

    for (int i = 0; i < pedestal.size(); i++)
        pedestal[i] = input.pedestal(i);

    if (input.gain_calibration_size() != gain_calibration.size())
        throw JFJochException(JFJochExceptionCategory::ZSTDCompressionError,
                              "Mismatch in size of serialized data (gain calibration count)");

    for (int i = 0; i < gain_calibration.size(); i++){
        if (input.gain_calibration(i).gain_calibration().size() != 3 * RAW_MODULE_SIZE * sizeof(double))
            throw JFJochException(JFJochExceptionCategory::ZSTDCompressionError,
                                  "Mismatch in size of serialized data (gain calibration size)");
        std::vector<double> tmp(3 * RAW_MODULE_SIZE);
        memcpy(tmp.data(), input.gain_calibration(i).gain_calibration().data(), 3 * RAW_MODULE_SIZE * sizeof(double));
        gain_calibration[i] = JFModuleGainCalibration(tmp);
    }

    memcpy(mask.data(), input.mask().data(), nmodules * RAW_MODULE_SIZE * sizeof(uint32_t));
}

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

JFCalibration::operator JFJochProtoBuf::JFCalibration() const {
    JFJochProtoBuf::JFCalibration output;

    output.set_nmodules(nmodules);
    output.set_nstorage_cells(nstorage_cells);

    for (const auto & i : pedestal) {
        auto tmp = output.add_pedestal();
        *tmp = i;
    }

    for (const auto &i: gain_calibration) {
        auto tmp = output.add_gain_calibration();
        auto &tmp_in = i.GetGainCalibration();
        tmp->set_gain_calibration(tmp_in.data(), tmp_in.size() * sizeof(double));
    }

    output.set_mask(mask.data(), RAW_MODULE_SIZE * nmodules * sizeof(uint32_t));

    return output;
}

JFCalibration::operator JFJochProtoBuf::JFCalibrationStatistics() const {
    return GetModuleStatistics();
}

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