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

#include <fstream>

#include <getopt.h>

#include "../common/DiffractionExperiment.h"
#include "../writer/FileWriter.h"
#include "../common/Logger.h"
#include "../common/RawToConvertedGeometry.h"
#include "../receiver/FrameTransformation.h"
#include "../image_analysis/MXAnalyzer.h"
#include "../preview/PreviewImage.h"
#include "../receiver/LossyFilter.h"
#include "../preview/JFJochTIFF.h"
#include "../common/print_license.h"

Logger logger("jfjoch_offline_process");

bool write_debug = false;

void print_usage() {
    logger.Info("Usage ./jfjoch_offline_process {options} <name of master dataset> <name of data dataset>");
    logger.Info("");
    logger.Info("Available options:");
    logger.Info("-s<float>     SNR threshold");
    logger.Info("-c<int>       count threshold");
    logger.Info("-d<float>     high resolution limit");
    logger.Info("-D<float>     low resolution limit");
    logger.Info("-p<int>       min pixels per spot");
    logger.Info("-S<int>       max spots saved per image");
    logger.Info("-o<string>    output prefix");
    logger.Info("-i<int>       image number");
    logger.Info("-I            enable indexing");
    logger.Info("-t<float>     indexing tolerance (0.0-1.0)");
    logger.Info("-l<float>     lossy filter probability (0.0-1.0)");
    logger.Info("-r{<int>}     filter powder rings (min spots)");
    logger.Info("-g            write debug images");
    logger.Info("");
}

void GetGeometry(DiffractionExperiment &x, HDF5Object &master_file) {
    x.BeamX_pxl(
            HDF5DataSet(master_file, "/entry/instrument/detector/beam_center_x").ReadScalar<double>());
    x.BeamY_pxl(
            HDF5DataSet(master_file, "/entry/instrument/detector/beam_center_y").ReadScalar<double>());
    x.DetectorDistance_mm(
            HDF5DataSet(master_file, "/entry/instrument/detector/detector_distance").ReadScalar<double>() *1e3);
    x.IncidentEnergy_keV(WVL_1A_IN_KEV /
                       HDF5DataSet(master_file, "/entry/instrument/beam/incident_wavelength").ReadScalar<double>());

    try {
        std::vector<float> v;
        HDF5DataSet(master_file, "/entry/sample/unit_cell").ReadVector(v);
        if (v.size() == 6) {
            logger.Info("Unit cell {} {} {} {} {} {}", v[0], v[1], v[2], v[3], v[4], v[5]);
            x.SetUnitCell(UnitCell(v[0], v[1], v[2], v[3], v[4], v[5]));
        }
    } catch (...) {}
}

int parse_options(DiffractionExperiment& experiment,
                  SpotFindingSettings& settings,
                  int argc, char **argv) {

    int c;
    static struct option long_options[] = {
            {"count",required_argument,       0, 'c'},
            {"snr",  required_argument,       0, 's'},
            {"minpix",  required_argument,       0, 'm'},
            {"d_min",  required_argument,       0, 'd'},
            {"d_max",  required_argument,       0, 'D'},
            {"output",    required_argument, 0, 'o'},
            {"help",    no_argument, 0, 'h'},
            {"max-spots",    required_argument, 0, 'S'},
            {"indexing",    no_argument, 0, 'I'},
            {"indexing_tolerance",    required_argument, 0, 't'},
            {"nimages",    required_argument, 0, 'n'},
            {"lossy",    required_argument, 0, 'l'},
            {"filter_rings",    optional_argument, 0, 'r'},
            {"debug",    no_argument, 0, 'g'},
            {0, 0, 0, 0}
    };

    int option_index = 0;
    int opt;
    while ((opt = getopt_long(argc, argv, "c:s:o:p:d:D:?hS:Ii:t:l:r::g",long_options, &option_index)) != -1 ) {
        switch (opt) {
            case 'c':
                settings.photon_count_threshold = atol(optarg);
                break;
            case 's':
                settings.signal_to_noise_threshold = atof(optarg);
                break;
            case 'd':
                settings.high_resolution_limit = atof(optarg);
                break;
            case 'D':
                settings.low_resolution_limit = atof(optarg);
                break;
            case 'p':
                settings.min_pix_per_spot = atol(optarg);
                break;
            case 'o':
                experiment.FilePrefix(optarg);
                break;
            case 'S':
                experiment.MaxSpotCount(atol(optarg));
                break;
            case 'I':
                settings.indexing = true;
                break;
            case 'i':
                experiment.ImagesPerTrigger(atol(optarg));
                break;
            case 't':
                settings.indexing_tolerance = atof(optarg);
                break;
            case 'l':
                experiment.LossyCompressionSerialMX(atof(optarg));
                break;
            case 'r':
                settings.filter_spots_powder_ring = true;
                if (optarg != nullptr)
                    settings.min_spot_count_powder_ring = atol(optarg);
                break;
            case 'g':
                write_debug = true;
                break;
            case 'h':
                print_usage();
                exit(EXIT_SUCCESS);
            default:
                logger.Error("Unknown option {}", opt);
                exit(EXIT_FAILURE);
        }
    }
    return optind;
}

int main(int argc, char **argv) {

    RegisterHDF5Filter();

    print_license("jfjoch_offline_process");

    DiffractionExperiment x;
    x.FilePrefix("output").ImagesPerTrigger(100000).Summation(1).NumTriggers(1);

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.indexing = false;
    int first_argc = parse_options(x, settings, argc, argv);

    try {
        DiffractionExperiment::CheckDataProcessingSettings(settings);
    } catch (const std::exception &e) {
        logger.ErrorException(e);
        exit(EXIT_FAILURE);
    }

    if (argc - first_argc != 2) {
        print_usage();
        exit(EXIT_FAILURE);
    }

    HDF5ReadOnlyFile master_file(argv[first_argc]);
    HDF5ReadOnlyFile data_file(argv[first_argc + 1]);
    HDF5DataSet dataset(data_file, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    if (file_space.GetNumOfDimensions() != 3) {
        logger.Error("/entry/data/data must be 3D");
        exit(EXIT_FAILURE);
    }

    if ((file_space.GetDimensions()[1] == 2164) && (file_space.GetDimensions()[2] == 2068)) {
        logger.Info("JF4M with gaps detected (2068 x 2164)");
        x.Detector(DetJF4M());
    } else if ((file_space.GetDimensions()[1] == 3264) && (file_space.GetDimensions()[2] == 3106)) {
        logger.Info("JF9M with gaps detected (3106x3264)");
        x.Detector(DetJF9M());
    } else {
        logger.Error("Unknown geometry - exiting");
        exit(EXIT_FAILURE);
    }

    GetGeometry(x, master_file);

    uint64_t nimages = std::min<uint64_t>(file_space.GetDimensions()[0], x.GetImageNum());
    logger.Info("Number of images in the original dataset: " + std::to_string(nimages));
    x.Mode(DetectorMode::Conversion).ImagesPerTrigger(nimages);

    logger.Info("Beam center {} {} pxl",x.GetBeamX_pxl(), x.GetBeamY_pxl());
    logger.Info("Detector distance {} mm", x.GetDetectorDistance_mm());
    logger.Info("Energy {} keV", x.GetIncidentEnergy_keV());
    logger.Info("Output file prefix {}", x.GetFilePrefix());
    if (x.GetLossyCompressionSerialMX() != 1.0)
        logger.Info("MX reduction factor {:0.2f}", x.GetLossyCompressionSerialMX());
    if (settings.indexing)
        logger.Info("Indexing tolerance {:0.4f}", settings.indexing_tolerance);

    logger.Info("Spot finding parameters SNR {:.2f} count {} minpix {} d {:.2f} - {:.2f}",
                settings.signal_to_noise_threshold,
                settings.photon_count_threshold,
                settings.min_pix_per_spot,
                settings.high_resolution_limit,
                settings.low_resolution_limit);

    StartMessage start_message;
    x.FillMessage(start_message);

    // Master & calibration files are written outside of timing routine
    auto fileset = std::make_unique<FileWriter>(start_message);

    std::vector<int16_t> original_image(x.GetPixelsNum());
    std::vector<int16_t> image_in_raw_coordinates(x.GetModulesNum() * RAW_MODULE_SIZE);
    FrameTransformation transformation(x);

    MXAnalyzer analyzer(x);

    PixelMask mask(x);

    PreviewImage preview;
    preview.Configure(x, mask);

    uint64_t indexed_images = 0;

    LossyFilter filter(x);

    for (int i = 0; i < nimages; i++) {
        std::vector<hsize_t> start = {(hsize_t) i,0,0};
        std::vector<hsize_t> size = {1, (hsize_t) x.GetYPixelsNum(), (hsize_t) x. GetXPixelsNum()};

        dataset.ReadVector(original_image, start, size);
        ConvertedToRawGeometry(x,
                               image_in_raw_coordinates.data(),
                               original_image.data());
        for (int m = 0; m < x.GetModulesNum(); m++) {
            transformation.ProcessModule(image_in_raw_coordinates.data() + m * RAW_MODULE_SIZE,
                                         m,
                                         0);
            analyzer.ReadFromCPU(image_in_raw_coordinates.data() + m * RAW_MODULE_SIZE,
                                 settings,
                                 m);
        }

        auto image = transformation.GetCompressedImage();

        DataMessage message{};
        message.image = image;
        message.number = i;
        message.original_number = i;

        analyzer.Process(message, settings);
        if (message.indexing_result)
            indexed_images++;

        if (filter.ApplyFilter(message))
            fileset->WriteHDF5(message);

        if ((i > 0) && (i % 1000 == 0))
            logger.Info("Images {:10d}   indexing rate: {:0.1f}%", i,
                        static_cast<float>(indexed_images) / static_cast<float>(i+1) * 100.0f);
    }
    fileset.reset();

    if (write_debug) {
        std::vector<float> convg(x.GetPixelsNum());
        std::vector<int32_t> convg_u32(x.GetPixelsNum());

        if (!analyzer.GetCPUMean().empty()) {
            RawToConvertedGeometry(x, convg.data(), analyzer.GetCPUMean().data());
            for (int i = 0; i < convg.size(); i++)
                convg_u32[i] = std::lround(convg[i] * 10.0f);
            WriteTIFFToFile("spot_finding_mean.tiff", convg_u32.data(), x.GetXPixelsNum(), x.GetYPixelsNum(), 4, true);
        }

        if (!analyzer.GetCPUStdDev().empty()) {
            RawToConvertedGeometry(x, convg.data(), analyzer.GetCPUStdDev().data());
            for (int i = 0; i < convg.size(); i++)
                convg_u32[i] = std::lround(convg[i] * 10.0f);
            WriteTIFFToFile("spot_finding_stddev.tiff", convg_u32.data(), x.GetXPixelsNum(), x.GetYPixelsNum(), 4, true);
        }

        if (!analyzer.GetCPUStrongPixel().empty()) {
            RawToConvertedGeometry(x, convg_u32.data(), analyzer.GetCPUStrongPixel().data());
            WriteTIFFToFile("spot_finding_strong_pixel.tiff", convg_u32.data(), x.GetXPixelsNum(), x.GetYPixelsNum(), 4, true);
        }

        WriteTIFFToFile("spot_finding_original_image.tiff", original_image.data(), x.GetXPixelsNum(), x.GetYPixelsNum(), 2, true);
    }

    if (settings.indexing)
        logger.Info("Indexing rate: {:0.1f}%", static_cast<float>(indexed_images) / static_cast<float>(nimages) * 100.0f);
}