// Copyright (2019-2023) Paul Scherrer Institute

#include <iostream>
#include <fstream>

#include <getopt.h>

#include "../common/DiffractionExperiment.h"
#include "../writer/HDF5Writer.h"
#include "../common/Logger.h"
#include "../common/RawToConvertedGeometry.h"
#include "../receiver/FrameTransformation.h"
#include "../image_analysis/MXAnalyzer.h"
#include "../export_images/PreviewImage.h"

Logger logger("jfjoch_spot_finding_test");

bool write_jpeg = false;

void print_usage() {
    logger.Info("Usage ./jfjoch_spot_finding_test {options} <name of master dataset> <name of data dataset>");
    logger.Info("");
    logger.Info("Available options:");
    logger.Info("-j            write JPEGs");
    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("-M<int>       max spots saved per image");
    logger.Info("-o<string>    output prefix");
    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.PhotonEnergy_keV(WVL_1A_IN_KEV /
                       HDF5DataSet(master_file, "/entry/instrument/beam/incident_wavelength").ReadScalar<double>());
}

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'},
            {"jpeg",    no_argument, 0, 'j'},
            {"help",    no_argument, 0, 'h'},
            {"max-spots",    required_argument, 0, 'S'},
            {0, 0, 0, 0}
    };

    int option_index = 0;
    int opt;
    while ((opt = getopt_long(argc, argv, "c:s:o:p:d:D:j?hS:",long_options, &option_index)) != -1 ) {
        switch (opt) {
            case 'j':
                write_jpeg = true;
                break;
            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 '?':
            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();

    DiffractionExperiment x;
    x.FilePrefix("output");

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    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(DetectorGeometry(8, 2, 8, 36));
    } else if ((file_space.GetDimensions()[1] == 3264) && (file_space.GetDimensions()[2] == 3106)) {
        logger.Info("JF9M with gaps detected (3106x3264)");
        x.Detector(DetectorGeometry(18, 3, 8, 36));
    } else {
        logger.Error("Unknown geometry - exiting");
        exit(EXIT_FAILURE);
    }

    GetGeometry(x, master_file);

    uint64_t nimages = file_space.GetDimensions()[0];
    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.GetPhotonEnergyForConversion_keV());
    logger.Info("Output file prefix {}", x.GetFilePrefix());
    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<HDF5Writer>(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);

    PreviewImage preview(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;

        analyzer.Process(message, settings);

        if (write_jpeg) {
            preview.UpdateImage(transformation.GetImage(), message.spots);
            auto preview_image = preview.GenerateJPEG(PreviewJPEGSettings{
                    .saturation_value = 10,
                    .jpeg_quality = 100,
                    .show_spots = true,
                    .show_roi = false
            });

            auto fname = fmt::format("{}_{:06d}.jpeg", x.GetFilePrefix(), i);
            std::ofstream f(fname, std::ios::binary);
            f.write(preview_image.data(), preview_image.size());
        }
        fileset->Write(message);
    }
    fileset.reset();
}