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

#include <atomic>
#include <csignal>
#include <iostream>
#include <string>
#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/print_license.h"
#include "../process/JFJochProcess.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;
}

namespace {
    std::atomic<JFJochProcess *> g_active_process{nullptr};
    void handle_sigint(int) {
        if (auto *p = g_active_process.load())
            p->Cancel();
    }
}

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

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

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

    // 2. Build experiment: defaults from the input file, overridden by command line. Output and
    //    runtime invariants are set inside JFJochProcess; here we only configure the geometry and
    //    azimuthal-integration settings.
    DiffractionExperiment experiment(dataset->experiment);
    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");

    // 3. Run the shared azimuthal-integration workflow.
    ProcessConfig config;
    config.mode = ProcessMode::AzimuthalIntegration;
    config.start_image = start_image;
    config.end_image = end_image;
    config.stride = image_stride;
    config.nthreads = nthreads;
    config.output_prefix = output_prefix;

    JFJochProcess process(reader, experiment, dataset->pixel_mask, config);

    g_active_process = &process;
    std::signal(SIGINT, handle_sigint);

    ProcessResult result;
    try {
        result = process.Run();
    } catch (const std::exception &e) {
        logger.Error("Processing failed: {}", e.what());
        exit(EXIT_FAILURE);
    }
    g_active_process = nullptr;

    // 4. Report statistics
    std::cout << fmt::format("Processing time:  {:.2f} s", result.processing_time_s) << std::endl;
    std::cout << fmt::format("Frame rate:       {:.2f} Hz", result.frame_rate_hz) << std::endl;
    std::cout << fmt::format("Total throughput: {:.2f} MB/s", result.throughput_MBs) << std::endl;
    if (result.cancelled)
        logger.Warning("Processing was cancelled after {} images", result.images_processed);
}