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

#include <iostream>
#include <string>
#include <vector>
#include <future>
#include <atomic>
#include <chrono>
#include <optional>
#include <algorithm>
#include <getopt.h>

#include "../reader/JFJochHDF5Reader.h"
#include "../common/Logger.h"
#include "../common/Definitions.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/AzimuthalIntegrationProfile.h"
#include "../common/time_utc.h"
#include "../common/print_license.h"
#include "../image_analysis/azint/AzIntEngineCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBuffer.h"
#include "../receiver/JFJochReceiverPlots.h"
#include "../writer/FileWriter.h"

void print_usage() {
    std::cout << "Usage ./jfjoch_azint {<options>} <input.h5>" << std::endl;
    std::cout << "Runs CPU azimuthal integration on a Jungfraujoch HDF5 file and writes <prefix>_process.h5" << std::endl;
    std::cout << "Options:" << std::endl;
    std::cout << "   -o, --output-prefix <txt>        Output file prefix (default: output)" << std::endl;
    std::cout << "   -N, --threads <num>              Number of threads (default: 1)" << std::endl;
    std::cout << "   -s, --start-image <num>          Start image number (default: 0)" << std::endl;
    std::cout << "   -e, --end-image <num>            End image number (default: all)" << std::endl;
    std::cout << "   -t, --stride <num>               Image stride (default: 1)" << std::endl;
    std::cout << "   -v, --verbose                    Verbose output" << std::endl;
    std::cout << std::endl;

    std::cout << "  Azimuthal integration (defaults taken from the input file)" << std::endl;
    std::cout << "   --min-q <num>                    Minimum Q for integration (1/A)" << std::endl;
    std::cout << "   --max-q <num>                    Maximum Q for integration (1/A)" << std::endl;
    std::cout << "   --q-spacing <num>                Q bin spacing (1/A)" << std::endl;
    std::cout << "   --azimuthal-bins <num>           Number of azimuthal bins (default: 1)" << std::endl;
    std::cout << "   --polarization-correction <on|off>  Enable/disable polarization correction" << std::endl;
    std::cout << "   --solid-angle-correction <on|off>   Enable/disable solid angle correction" << std::endl;
    std::cout << std::endl;

    std::cout << "  Geometry overrides (defaults taken from the input file)" << std::endl;
    std::cout << "   --beam-x <num>                   Beam center X (pixel)" << std::endl;
    std::cout << "   --beam-y <num>                   Beam center Y (pixel)" << std::endl;
    std::cout << "   --detector-distance <num>        Detector distance (mm)" << std::endl;
    std::cout << "   --wavelength <num>               Wavelength (A)" << std::endl;
    std::cout << "   --rot1 <num>                     PONI rotation 1 (rad)" << std::endl;
    std::cout << "   --rot2 <num>                     PONI rotation 2 (rad)" << std::endl;
    std::cout << "   --polarization <num>             Polarization factor" << std::endl;
}

enum {
    OPT_MIN_Q = 1000,
    OPT_MAX_Q,
    OPT_Q_SPACING,
    OPT_AZIMUTHAL_BINS,
    OPT_POLARIZATION_CORRECTION,
    OPT_SOLID_ANGLE_CORRECTION,
    OPT_BEAM_X,
    OPT_BEAM_Y,
    OPT_DETECTOR_DISTANCE,
    OPT_WAVELENGTH,
    OPT_ROT1,
    OPT_ROT2,
    OPT_POLARIZATION
};

static option long_options[] = {
    {"verbose", no_argument, nullptr, 'v'},
    {"output-prefix", required_argument, nullptr, 'o'},
    {"threads", required_argument, nullptr, 'N'},
    {"start-image", required_argument, nullptr, 's'},
    {"end-image", required_argument, nullptr, 'e'},
    {"stride", required_argument, nullptr, 't'},

    {"min-q", required_argument, nullptr, OPT_MIN_Q},
    {"max-q", required_argument, nullptr, OPT_MAX_Q},
    {"q-spacing", required_argument, nullptr, OPT_Q_SPACING},
    {"azimuthal-bins", required_argument, nullptr, OPT_AZIMUTHAL_BINS},
    {"polarization-correction", required_argument, nullptr, OPT_POLARIZATION_CORRECTION},
    {"solid-angle-correction", required_argument, nullptr, OPT_SOLID_ANGLE_CORRECTION},

    {"beam-x", required_argument, nullptr, OPT_BEAM_X},
    {"beam-y", required_argument, nullptr, OPT_BEAM_Y},
    {"detector-distance", required_argument, nullptr, OPT_DETECTOR_DISTANCE},
    {"wavelength", required_argument, nullptr, OPT_WAVELENGTH},
    {"rot1", required_argument, nullptr, OPT_ROT1},
    {"rot2", required_argument, nullptr, OPT_ROT2},
    {"polarization", required_argument, nullptr, OPT_POLARIZATION},
    {nullptr, 0, nullptr, 0}
};

bool parse_on_off(const char *arg, bool &out) {
    std::string s = arg ? arg : "";
    std::transform(s.begin(), s.end(), s.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    if (s == "on" || s == "1" || s == "true" || s == "yes") {
        out = true;
        return true;
    }
    if (s == "off" || s == "0" || s == "false" || s == "no") {
        out = false;
        return true;
    }
    return false;
}

int main(int argc, char **argv) {
    for (int i = 0; i < argc; i++)
        std::cout << argv[i] << " ";
    std::cout << std::endl << std::endl;

    RegisterHDF5Filter();
    print_license("jfjoch_azint");

    Logger logger("jfjoch_azint");

    std::string output_prefix = "output";
    int nthreads = 1;
    int start_image = 0;
    int end_image = -1; // -1 indicates process until end
    int image_stride = 1;
    bool verbose = false;

    // Azimuthal integration overrides (default: keep value from input file)
    std::optional<float> min_q;
    std::optional<float> max_q;
    std::optional<float> q_spacing;
    std::optional<int32_t> azimuthal_bins;
    std::optional<bool> polarization_correction;
    std::optional<bool> solid_angle_correction;

    // Geometry overrides (default: keep value from input file)
    std::optional<float> beam_x;
    std::optional<float> beam_y;
    std::optional<float> detector_distance_mm;
    std::optional<float> wavelength_A;
    std::optional<float> rot1_rad;
    std::optional<float> rot2_rad;
    std::optional<float> polarization_factor;

    if (argc == 1) {
        print_usage();
        exit(EXIT_FAILURE);
    }

    int opt;
    int option_index = 0;
    const char *short_opts = "vo:N:s:e:t:";

    while ((opt = getopt_long(argc, argv, short_opts, long_options, &option_index)) != -1) {
        switch (opt) {
            case 'o': output_prefix = optarg; break;
            case 'v': verbose = true; break;
            case 'N': nthreads = atoi(optarg); break;
            case 's': start_image = atoi(optarg); break;
            case 'e': end_image = atoi(optarg); break;
            case 't': image_stride = atoi(optarg); break;

            case OPT_MIN_Q: min_q = atof(optarg); break;
            case OPT_MAX_Q: max_q = atof(optarg); break;
            case OPT_Q_SPACING: q_spacing = atof(optarg); break;
            case OPT_AZIMUTHAL_BINS: azimuthal_bins = atoi(optarg); break;
            case OPT_POLARIZATION_CORRECTION: {
                bool value;
                if (!parse_on_off(optarg, value)) {
                    logger.Error("Invalid polarization correction value (expected on|off): {}", optarg);
                    exit(EXIT_FAILURE);
                }
                polarization_correction = value;
                break;
            }
            case OPT_SOLID_ANGLE_CORRECTION: {
                bool value;
                if (!parse_on_off(optarg, value)) {
                    logger.Error("Invalid solid angle correction value (expected on|off): {}", optarg);
                    exit(EXIT_FAILURE);
                }
                solid_angle_correction = value;
                break;
            }

            case OPT_BEAM_X: beam_x = atof(optarg); break;
            case OPT_BEAM_Y: beam_y = atof(optarg); break;
            case OPT_DETECTOR_DISTANCE: detector_distance_mm = atof(optarg); break;
            case OPT_WAVELENGTH: wavelength_A = atof(optarg); break;
            case OPT_ROT1: rot1_rad = atof(optarg); break;
            case OPT_ROT2: rot2_rad = atof(optarg); break;
            case OPT_POLARIZATION: polarization_factor = atof(optarg); break;

            default:
                print_usage();
                exit(EXIT_FAILURE);
        }
    }

    if (optind != argc - 1) {
        logger.Error("Input file not specified");
        print_usage();
        exit(EXIT_FAILURE);
    }

    const std::string input_file = argv[optind];
    logger.Verbose(verbose);

    // 1. Read input file
    JFJochHDF5Reader reader;
    try {
        reader.ReadFile(input_file);
    } catch (const std::exception &e) {
        logger.Error("Error reading input file: {}", e.what());
        exit(EXIT_FAILURE);
    }

    const auto dataset = reader.GetDataset();
    if (!dataset) {
        logger.Error("No experiment dataset found in the input file");
        exit(EXIT_FAILURE);
    }
    logger.Info("Loaded dataset from {}", input_file);

    if (image_stride <= 0) {
        logger.Error("Image stride must be positive");
        exit(EXIT_FAILURE);
    }

    const uint64_t total_images_in_file = reader.GetNumberOfImages();
    if (end_image < 0 || end_image > total_images_in_file)
        end_image = total_images_in_file;

    const int images_to_process = (end_image - start_image) / image_stride;
    if (images_to_process <= 0) {
        logger.Warning("No images to process (Start: {}, End: {}, Stride: {}, Total: {})",
                       start_image, end_image, image_stride, total_images_in_file);
        return 0;
    }

    // 2. Set up experiment - defaults come from the input file, overridden by command line
    DiffractionExperiment experiment(dataset->experiment);
    experiment.BitDepthImage(32).PixelSigned(true);
    experiment.Mode(DetectorMode::Standard);
    experiment.FilePrefix(output_prefix);
    experiment.OverwriteExistingFiles(true);
    experiment.SetFileWriterFormat(FileWriterFormat::NXmxLegacy);
    experiment.ImagesPerTrigger(images_to_process);
    experiment.NumTriggers(1);

    if (beam_x.has_value()) experiment.BeamX_pxl(beam_x.value());
    if (beam_y.has_value()) experiment.BeamY_pxl(beam_y.value());
    if (detector_distance_mm.has_value()) experiment.DetectorDistance_mm(detector_distance_mm.value());
    if (wavelength_A.has_value()) experiment.IncidentEnergy_keV(WVL_1A_IN_KEV / wavelength_A.value());
    if (rot1_rad.has_value()) experiment.PoniRot1_rad(rot1_rad.value());
    if (rot2_rad.has_value()) experiment.PoniRot2_rad(rot2_rad.value());
    if (polarization_factor.has_value()) experiment.PolarizationFactor(polarization_factor.value());

    AzimuthalIntegrationSettings azint_settings = experiment.GetAzimuthalIntegrationSettings();
    if (min_q.has_value() || max_q.has_value())
        azint_settings.QRange_recipA(min_q.value_or(azint_settings.GetLowQ_recipA()),
                                     max_q.value_or(azint_settings.GetHighQ_recipA()));
    if (q_spacing.has_value()) azint_settings.QSpacing_recipA(q_spacing.value());
    if (azimuthal_bins.has_value()) azint_settings.AzimuthalBinCount(azimuthal_bins.value());
    if (polarization_correction.has_value()) azint_settings.PolarizationCorrection(polarization_correction.value());
    if (solid_angle_correction.has_value()) azint_settings.SolidAngleCorrection(solid_angle_correction.value());
    experiment.ImportAzimuthalIntegrationSettings(azint_settings);

    logger.Info("Geometry: beam ({:.2f}, {:.2f}) pxl, distance {:.2f} mm, wavelength {:.5f} A",
                experiment.GetBeamX_pxl(), experiment.GetBeamY_pxl(),
                experiment.GetDetectorDistance_mm(), experiment.GetWavelength_A());
    logger.Info("Azimuthal integration: Q range [{:.4f}, {:.4f}] 1/A, spacing {:.4f} 1/A, {} Q bins x {} azimuthal bins",
                azint_settings.GetLowQ_recipA(), azint_settings.GetHighQ_recipA(),
                azint_settings.GetQSpacing_recipA(), azint_settings.GetQBinCount(),
                azint_settings.GetAzimuthalBinCount());
    logger.Info("Corrections: polarization {}, solid angle {}",
                azint_settings.IsPolarizationCorrection() ? "on" : "off",
                azint_settings.IsSolidAngleCorrection() ? "on" : "off");

    PixelMask pixel_mask = dataset->pixel_mask;
    AzimuthalIntegrationMapping mapping(experiment, pixel_mask);

    JFJochReceiverPlots plots;
    plots.Setup(experiment, mapping);

    // 3. Set up the output file
    StartMessage start_message;
    experiment.FillMessage(start_message);
    start_message.arm_date = dataset->arm_date;
    start_message.az_int_bin_to_q = mapping.GetBinToQ();
    start_message.az_int_bin_to_two_theta = mapping.GetBinToTwoTheta();
    start_message.az_int_q_bin_count = mapping.GetQBinCount();
    start_message.az_int_phi_bin_count = mapping.GetAzimuthalBinCount();
    if (mapping.GetAzimuthalBinCount() > 1)
        start_message.az_int_bin_to_phi = mapping.GetBinToPhi();
    start_message.pixel_mask["default"] = pixel_mask.GetMask(experiment);
    start_message.master_suffix = "process";
    start_message.file_format = FileWriterFormat::NXmxIntegrated;
    start_message.write_master_file = true;
    start_message.write_images = false;
    start_message.hdf5_source_data = reader.GetHDF5DataSource(start_image, images_to_process);

    std::unique_ptr<FileWriter> writer;
    try {
        if (!output_prefix.empty())
            writer = std::make_unique<FileWriter>(start_message);
    } catch (const std::exception &e) {
        logger.Error("Failed to initialize file writer: {}", e.what());
        exit(EXIT_FAILURE);
    }

    logger.Info("Starting azimuthal integration of {} images (range {}-{}) using {} threads",
                images_to_process, start_image, end_image, nthreads);

    // 4. Process images on N CPU threads
    std::atomic<int> next_image = 0;
    std::atomic<int> finished_count = 0;
    std::atomic<uint64_t> total_uncompressed_bytes = 0;

    const auto start_time = std::chrono::steady_clock::now();

    auto worker = [&]() {
        std::vector<uint8_t> decompression_buffer;
        ImagePreprocessorCPU preprocessor(experiment, pixel_mask);
        ImagePreprocessorBuffer buffer(experiment.GetPixelsNum());
        AzIntEngineCPU azint(mapping);
        AzimuthalIntegrationProfile profile(mapping);

        while (true) {
            const int image_ordinal = next_image.fetch_add(1);
            const int image_idx = start_image + image_ordinal * image_stride;
            if (image_idx >= end_image) break;

            std::shared_ptr<JFJochReaderRawImage> img;
            try {
                img = reader.GetRawImage(image_idx);
            } catch (const std::exception &e) {
                logger.Error("Failed to load image {}: {}", image_idx, e.what());
                continue;
            }
            if (!img) continue;

            DataMessage msg{};
            msg.image = img->image;
            msg.number = image_ordinal;
            msg.original_number = image_idx;
            if (dataset->efficiency.size() > image_idx)
                msg.image_collection_efficiency = dataset->efficiency[image_idx];
            total_uncompressed_bytes += msg.image.GetUncompressedSize();

            const auto azint_start = std::chrono::steady_clock::now();
            try {
                const uint8_t *image_ptr = msg.image.GetUncompressedPtr(decompression_buffer);
                preprocessor.Analyze(buffer, image_ptr, msg.image.GetMode());
                azint.Run(buffer, profile);
            } catch (const std::exception &e) {
                logger.Error("Error integrating image {}: {}", image_idx, e.what());
                continue;
            }
            const auto azint_end = std::chrono::steady_clock::now();

            msg.azint_time_s = std::chrono::duration<float>(azint_end - azint_start).count();
            msg.processing_time_s = msg.azint_time_s;
            msg.az_int_profile = profile.GetResult();
            msg.az_int_profile_count = profile.GetPixelCount();
            msg.az_int_profile_std = profile.GetStd();
            msg.bkg_estimate = profile.GetBkgEstimate(mapping.Settings());
            msg.run_number = experiment.GetRunNumber();
            msg.run_name = experiment.GetRunName();

            plots.Add(msg, profile);
            if (writer)
                writer->Write(msg);

            const int done = finished_count.fetch_add(1) + 1;
            if (done % 1000 == 0) {
                const double elapsed_s = std::chrono::duration<double>(
                    std::chrono::steady_clock::now() - start_time).count();
                const double frame_rate_hz = (elapsed_s > 0.0) ? (done / elapsed_s) : 0.0;
                logger.Info("Integrated {} / {} images ({:.1f} Hz)", done, images_to_process, frame_rate_hz);
            }
        }
    };

    std::vector<std::future<void> > futures;
    futures.reserve(nthreads);
    for (int i = 0; i < nthreads; ++i)
        futures.push_back(std::async(std::launch::async, worker));
    for (auto &f: futures)
        f.get();

    const auto end_time = std::chrono::steady_clock::now();

    // 5. Finalize output file
    EndMessage end_msg;
    end_msg.max_image_number = images_to_process;
    end_msg.images_collected_count = images_to_process;
    end_msg.images_sent_to_write_count = images_to_process;
    end_msg.end_date = time_UTC(std::chrono::system_clock::now());
    end_msg.run_number = experiment.GetRunNumber();
    end_msg.run_name = experiment.GetRunName();
    end_msg.bkg_estimate = plots.GetBkgEstimate();
    end_msg.az_int_result["dataset"] = plots.GetAzIntProfile();

    if (writer) {
        writer->WriteHDF5(end_msg);
        writer->Finalize();
    }

    // 6. Report statistics
    const double processing_time = std::chrono::duration<double>(end_time - start_time).count();
    const double throughput_MBs = static_cast<double>(total_uncompressed_bytes) / (processing_time * 1e6);
    const double frame_rate = static_cast<double>(finished_count.load()) / processing_time;

    std::cout << fmt::format("Processing time:  {:.2f} s", processing_time) << std::endl;
    std::cout << fmt::format("Frame rate:       {:.2f} Hz", frame_rate) << std::endl;
    std::cout << fmt::format("Total throughput: {:.2f} MB/s", throughput_MBs) << std::endl;
}