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

#ifndef CHECKIMAGEOUTPUT_H
#define CHECKIMAGEOUTPUT_H

#include <iostream>
#include <fstream>
#include "../common/DiffractionExperiment.h"
#include "../common/RawToConvertedGeometry.h"
#include "../common/JFJochException.h"
#include "../jungfrau/JFCalibration.h"
#include "../jungfrau/JFModuleGainCalibration.h"
#include "../jungfrau/JFConversionFloatingPoint.h"

template <class T> void LoadBinaryFile(const std::string &filename, T* output, size_t size) {
    std::fstream file(filename.c_str(), std::fstream::in | std::fstream::binary);
    if (file.is_open())
        file.read((char *) output, size * sizeof(T));
    else
        throw JFJochException(JFJochExceptionCategory::GainFileOpenError, "Gain file cannot be opened");
}

template <class T> double Compare(T *source, std::vector<double> &reference, int32_t npixel) {
    double result = 0;

    for (size_t i = 0; i < npixel; i++) {
        double diff = reference[i] - source[i];
        result += diff * diff;
    }

    return sqrt(result / npixel);
}

template <class T> double MaxErrorOnConversion(T *source, std::vector<double> &reference, int32_t npixel) {
    double ret = 0;

    for (size_t i = 0; i < npixel; i++) {
        double val = abs(reference[i] - source[i]);
        if (val > ret) ret = val;
    }
    return ret;
}

inline std::vector<double> ConversionRef(const DiffractionExperiment &experiment,
                                  const JFCalibration &calib,
                                  const uint16_t *raw,
                                  size_t nmodules,
                                  size_t storage_cell) {
    std::vector<double> conversion_ref(nmodules * RAW_MODULE_SIZE);

    if (experiment.IsJungfrauConvPhotonCnt()) {
        JFConversionFloatingPoint conversion(experiment);
        for (int m = 0; m < nmodules; m++) {
            conversion.Setup(calib.GainCalibration(m),
                             calib.Pedestal(m, 0, storage_cell),
                             calib.Pedestal(m, 1, storage_cell),
                             calib.Pedestal(m, 2, storage_cell),
                             experiment.GetIncidentEnergy_keV(),
                             experiment.IsUsingGainHG0());
            conversion.ConvertFP(conversion_ref.data() + m * RAW_MODULE_SIZE,
                                 raw + m * RAW_MODULE_SIZE);
        }
    } else {
        if (experiment.GetPixelValueLowThreshold()) {
            int64_t thr = experiment.GetPixelValueLowThreshold().value();
            for (int i = 0; i < nmodules * RAW_MODULE_SIZE; i++) {
                if (raw[i] < thr)
                    conversion_ref[i] = 0;
                else
                    conversion_ref[i] = experiment.GetSummation() * raw[i];
            }
        } else {
            for (int i = 0; i < nmodules * RAW_MODULE_SIZE; i++)
                conversion_ref[i] = experiment.GetSummation() * raw[i];
        }
    }

    if (experiment.GetLossyCompressionPoisson()) {
        for (int i = 0; i < nmodules * RAW_MODULE_SIZE; i++) {
            if (((int64_t)conversion_ref[i] >= experiment.GetSaturationLimit())
                || ((int64_t)conversion_ref[i] == experiment.GetUnderflow()));
            else if (conversion_ref[i] < 0.0)
                conversion_ref[i] = 0.0;
            else {
                conversion_ref[i] = std::round(experiment.GetLossyCompressionPoisson().value() *
                                               std::sqrt(conversion_ref[i]));
            }
        }
    }

    if (!experiment.IsJungfrauConvPhotonCnt()) {
        for (int i = 0; i < nmodules * RAW_MODULE_SIZE; i++) {
            if (conversion_ref[i] >= experiment.GetSaturationLimit()) {
                conversion_ref[i] = experiment.GetSaturationLimit();
            }
        }
    }
    return conversion_ref;
}

template <class T> double CheckImageOutput(const DiffractionExperiment &experiment,
                                           const JFCalibration &calib,
                                           const uint16_t *raw, T *converted,
                                           size_t storage_cell = 0) {
    auto conversion_ref = ConversionRef(experiment, calib, raw, experiment.GetModulesNum(), storage_cell);

    if (!experiment.IsGeometryTransformed())
        return Compare(converted, conversion_ref, RAW_MODULE_SIZE);

    std::vector<double> conversion_ref_transformed(experiment.GetPixelsNum(), experiment.GetImageFillValue());

    RawToConvertedGeometryAdjustMultipixels(experiment, conversion_ref_transformed.data(), conversion_ref.data() );
    return Compare(converted, conversion_ref_transformed, experiment.GetPixelsNum());
}

#endif //CHECKIMAGEOUTPUT_H