Jungfraujoch/tools/jfjoch_process.cpp

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

#include <iostream>
#include <vector>
#include <string>
#include <unistd.h>
#include <future>
#include <mutex>
#include <atomic>
#include <chrono>
#include <fstream>
#include <sstream>

#include "../reader/JFJochHDF5Reader.h"
#include "../common/Logger.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../common/AzimuthalIntegration.h"
#include "../common/time_utc.h"
#include "../common/print_license.h"
#include "../image_analysis/MXAnalysisWithoutFPGA.h"
#include "../image_analysis/indexing/IndexerFactory.h"
#include "../writer/FileWriter.h"
#include "../image_analysis/IndexAndRefine.h"
#include "../receiver/JFJochReceiverPlots.h"
#include "../compression/JFJochCompressor.h"
#include "../image_analysis/scale_merge/FrenchWilson.h"
#include "../image_analysis/scale_merge/SearchSpaceGroup.h"
#include "../image_analysis/WriteMmcif.h"

void print_usage(Logger &logger) {
    logger.Info("Usage ./jfjoch_analysis {<options>} <input.h5>");
    logger.Info("Options:");
    logger.Info("   -o<txt>    Output file prefix (default: output)");
    logger.Info("   -N<num>    Number of threads (default: 1)");
    logger.Info("   -s<num>    Start image number (default: 0)");
    logger.Info("   -e<num>    End image number (default: all)");
    logger.Info("   -v         Verbose output");
    logger.Info("   -R[num]    Rotation indexing (optional: min angular range deg)");
    logger.Info("   -F         Use FFT indexing algorithm (shortcut for -XFFT)");
    logger.Info("   -X<txt>    Indexing algorithm (FFBIDX|FFT|FFTW|Auto|None)");
    logger.Info("   -x         No least-square beam center refinement");
    logger.Info("   -d<num>    High resolution limit for spot finding (default: 1.5)");
    logger.Info("   -D<num>    High resolution limit for scaling/merging (default: 0.0; no limit)");
    logger.Info("   -S<num>    Space group number");
    logger.Info("   -M         Scale and merge (refine mosaicity) and write scaled.hkl + image.dat");
    logger.Info("   -P<txt>    Partiality refinement fixed|rot|unity (default: fixed)");
    logger.Info("   -A         Anomalous mode (don't merge Friedel pairs)");
    logger.Info("   -C<cell>   Fix reference unit cell: -C\"a,b,c,alpha,beta,gamma\" (comma-separated, no spaces; quotes optional)");
    logger.Info("   -c<num>    Max spot count (default: 250)");
    logger.Info("   -W         HDF5 file with analysis results is written");
    logger.Info("   -T<num>    Noise sigma level for spot finding (default: 3.0)");
    logger.Info("   -t<num>    Photon count threshold for spot finding (default: 10)");
}

void trim_in_place(std::string& t) {
    size_t b = 0;
    while (b < t.size() && std::isspace(static_cast<unsigned char>(t[b]))) b++;
    size_t e = t.size();
    while (e > b && std::isspace(static_cast<unsigned char>(t[e - 1]))) e--;
    t = t.substr(b, e - b);
};

std::optional<UnitCell> parse_unit_cell_arg(const char* arg)  {
    if (!arg)
        return std::nullopt;

    std::string s(arg);


    trim_in_place(s);

    if (s.size() >= 2 && ((s.front() == '"' && s.back() == '"') || (s.front() == '\'' && s.back() == '\''))) {
        s = s.substr(1, s.size() - 2);
        trim_in_place(s);
    }

    std::vector<std::string> parts;
    parts.reserve(6);
    size_t start = 0;
    while (true) {
        size_t pos = s.find(',', start);
        if (pos == std::string::npos) {
            parts.push_back(s.substr(start));
            break;
        }
        parts.push_back(s.substr(start, pos - start));
        start = pos + 1;
    }

    if (parts.size() != 6)
        return std::nullopt;

    auto parse_float_strict = [](const std::string& t, float& out) -> bool {
        try {
            size_t idx = 0;
            out = std::stof(t, &idx);
            return idx == t.size();
        } catch (...) {
            return false;
        }
    };

    UnitCell uc{};
    if (!parse_float_strict(parts[0], uc.a)) return std::nullopt;
    if (!parse_float_strict(parts[1], uc.b)) return std::nullopt;
    if (!parse_float_strict(parts[2], uc.c)) return std::nullopt;
    if (!parse_float_strict(parts[3], uc.alpha)) return std::nullopt;
    if (!parse_float_strict(parts[4], uc.beta)) return std::nullopt;
    if (!parse_float_strict(parts[5], uc.gamma)) return std::nullopt;

    return uc;
};

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

    print_license("jfjoch_analysis");

    Logger logger("jfjoch_analysis");

    std::string input_file;
    std::string output_prefix = "output";
    int nthreads = 1;
    int start_image = 0;
    int end_image = -1; // -1 indicates process until end
    bool verbose = false;
    bool rotation_indexing = false;
    std::optional<float> rotation_indexing_range;
    bool refine_beam_center = true;
    bool run_scaling = false;
    bool anomalous_mode = false;
    std::optional<int> space_group_number;
    std::optional<UnitCell> fixed_reference_unit_cell;
    bool write_output = false;
    std::optional<int64_t> max_spot_count_override;
    float sigma_spot_finding = 3.0;
    int64_t photon_count_threshold_spot_finding = 10;

    IndexingAlgorithmEnum indexing_algorithm = IndexingAlgorithmEnum::Auto;

    ScaleMergeOptions::PartialityModel partiality_model = ScaleMergeOptions::PartialityModel::Fixed;

    float d_min_spot_finding = 1.5;
    std::optional<float> d_min_scale_merge;

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

    int opt;
    while ((opt = getopt(argc, argv, "o:N:s:e:vc:R::FX:xd:S:MP:AD:C:T:t:W")) != -1) {
        switch (opt) {
            case 'o':
                output_prefix = optarg;
                break;
            case 'N':
                nthreads = atoi(optarg);
                break;
            case 's':
                start_image = atoi(optarg);
                break;
            case 'e':
                end_image = atoi(optarg);
                break;
            case 'W':
                write_output = true;
                break;
            case 'v':
                verbose = true;
                break;
            case 'd':
                d_min_spot_finding = atof(optarg);
                break;
            case 'c':
                max_spot_count_override = atoll(optarg);
                break;
            case 'R':
                rotation_indexing = true;
                if (optarg) rotation_indexing_range = atof(optarg);
                break;
            case 'F':
                indexing_algorithm = IndexingAlgorithmEnum::FFT;
                break;
            case 'X': {
                std::string alg = optarg ? optarg : "";
                std::transform(alg.begin(), alg.end(), alg.begin(),
                           [](unsigned char c) { return static_cast<char>(std::tolower(c)); });

                if (alg == "ffbidx")
                    indexing_algorithm = IndexingAlgorithmEnum::FFBIDX;
                else if (alg == "fft")
                    indexing_algorithm = IndexingAlgorithmEnum::FFT;
                else if (alg == "fftw")
                    indexing_algorithm = IndexingAlgorithmEnum::FFTW;
                else if (alg == "auto")
                    indexing_algorithm = IndexingAlgorithmEnum::Auto;
                else if (alg == "none")
                    indexing_algorithm = IndexingAlgorithmEnum::None;
                else {
                    logger.Error("Invalid indexing algorithm: {}", alg);
                    print_usage(logger);
                    exit(EXIT_FAILURE);
                }
                break;
            }
            case 'x':
                refine_beam_center = false;
                break;
            case 'D':
                d_min_scale_merge = atof(optarg);
                logger.Info("High resolution limit for scaling/merging set to {:.2f} A", d_min_spot_finding);
                break;
            case 'S':
                space_group_number = atoi(optarg);
                break;
            case 'M':
                run_scaling = true;
                break;
            case 'A':
                anomalous_mode = true;
                break;
            case 'T':
                sigma_spot_finding = atof(optarg);
                logger.Info("Noise threshold level for spot finding set to {:.2f} sigma", sigma_spot_finding);
                break;
            case 't':
                photon_count_threshold_spot_finding = atoi(optarg);
                logger.Info("Photon-count threshold level for spot finding set to {:d}", photon_count_threshold_spot_finding);
                break;
            case 'C': {
                auto uc = parse_unit_cell_arg(optarg);
                if (!uc.has_value()) {
                    logger.Error("Invalid -C unit cell. Expected: -C\"a,b,c,alpha,beta,gamma\" (6 floats, comma-separated, no spaces). Got: {}", optarg ? optarg : "<null>");
                    print_usage(logger);
                    exit(EXIT_FAILURE);
                }
                fixed_reference_unit_cell = uc;
                logger.Info("Fixed reference unit cell set: a={:.3f} b={:.3f} c={:.3f} alpha={:.3f} beta={:.3f} gamma={:.3f}",
                            uc->a, uc->b, uc->c, uc->alpha, uc->beta, uc->gamma);
                break;
            }
            case 'P':
                if (strcmp(optarg, "unity") == 0)
                    partiality_model = ScaleMergeOptions::PartialityModel::Unity;
                else if (strcmp(optarg, "fixed") == 0)
                    partiality_model = ScaleMergeOptions::PartialityModel::Fixed;
                else if (strcmp(optarg, "rot") == 0)
                    partiality_model = ScaleMergeOptions::PartialityModel::Rotation;
                else if (strcmp(optarg, "still") == 0)
                    partiality_model = ScaleMergeOptions::PartialityModel::Still;
                else {
                    logger.Error("Invalid partiality mode: {}", optarg);
                    print_usage(logger);
                    exit(EXIT_FAILURE);
                }
                break;
            default:
                print_usage(logger);
                exit(EXIT_FAILURE);
        }
    }

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

    input_file = argv[optind];
    logger.Verbose(verbose);

    // Validate space group number early
    const gemmi::SpaceGroup *space_group = nullptr;
    if (space_group_number.has_value()) {
        space_group = gemmi::find_spacegroup_by_number(space_group_number.value());
        if (!space_group) {
            logger.Error("Unknown space group number {}", space_group_number.value());
            exit(EXIT_FAILURE);
        }
        logger.Info("Using space group {} (number {})", space_group->hm, space_group_number.value());
    }

    // 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. Setup Experiment & Components
    DiffractionExperiment experiment(dataset->experiment);
    experiment.BitDepthImage(32).Compression(CompressionAlgorithm::BSHUF_LZ4);
    experiment.FilePrefix(output_prefix);
    experiment.Mode(DetectorMode::Standard); // Ensure full image analysis
    experiment.PixelSigned(true);
    experiment.OverwriteExistingFiles(true);
    experiment.PolarizationFactor(0.99);
    experiment.SetFileWriterFormat(FileWriterFormat::NXmxLegacy);

    if (fixed_reference_unit_cell.has_value())
        experiment.SetUnitCell(*fixed_reference_unit_cell);

    if (max_spot_count_override.has_value()) {
        experiment.MaxSpotCount(max_spot_count_override.value());
        logger.Info("Max spot count overridden to {}", max_spot_count_override.value());
    }

    // Configure Indexing
    IndexingSettings indexing_settings;
    indexing_settings.Algorithm(indexing_algorithm);
    indexing_settings.RotationIndexing(rotation_indexing);
    if (rotation_indexing_range.has_value())
        indexing_settings.RotationIndexingMinAngularRange_deg(rotation_indexing_range.value());

    if (refine_beam_center)
        indexing_settings.GeomRefinementAlgorithm(GeomRefinementAlgorithmEnum::BeamCenter);
    else
        indexing_settings.GeomRefinementAlgorithm(GeomRefinementAlgorithmEnum::None);
    experiment.ImportIndexingSettings(indexing_settings);

    SpotFindingSettings spot_settings;
    spot_settings.enable = true;
    spot_settings.indexing = true;
    spot_settings.high_resolution_limit = d_min_spot_finding;
    spot_settings.signal_to_noise_threshold = sigma_spot_finding;
    spot_settings.photon_count_threshold = photon_count_threshold_spot_finding;
    if (d_min_scale_merge > 0)
        spot_settings.high_resolution_limit = d_min_spot_finding;

    // Initialize Analysis Components
    PixelMask pixel_mask = dataset->pixel_mask;

    // If dataset has a mask you wish to use, you might need to load it into pixel_mask here
    // e.g. pixel_mask.LoadUserMask(dataset->pixel_mask, ...);

    AzimuthalIntegration mapping(experiment, pixel_mask);
    IndexerThreadPool indexer_pool(experiment.GetIndexingSettings());

    // Statistics collector
    JFJochReceiverPlots plots;
    plots.Setup(experiment, mapping);

    // 3. Setup FileWriter
    StartMessage start_message;
    experiment.FillMessage(start_message);
    start_message.arm_date = dataset->arm_date; // Use original arm date
    start_message.az_int_bin_to_q = mapping.GetBinToQ();
    start_message.az_int_q_bin_count = mapping.GetQBinCount();
    if (mapping.GetAzimuthalBinCount() > 1)
        start_message.az_int_bin_to_phi = mapping.GetBinToPhi();

    start_message.pixel_mask["default"] = pixel_mask.GetMask(experiment);
    start_message.max_spot_count = experiment.GetMaxSpotCount();

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

    // 4. Processing Setup
    int total_images_in_file = reader.GetNumberOfImages();
    if (end_image < 0 || end_image > total_images_in_file)
        end_image = total_images_in_file;

    int images_to_process = end_image - start_image;

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

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

    std::atomic<int> processed_count = 0;
    std::atomic<uint64_t> total_uncompressed_bytes = 0;
    std::atomic<uint64_t> max_image_number_sent = 0;

    // Mimic JFJochReceiver lattice handling (IndexAndRefine handles the logic per thread,
    // but we need a central accumulator or use the pool's functionality if IndexAndRefine wraps it)
    // Here we will use per-thread IndexAndRefine which uses the shared thread pool.

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

    IndexAndRefine indexer(experiment, &indexer_pool);

    std::atomic<int> finished_count = 0;

    auto worker = [&](int thread_id) {
        // Thread-local analysis resources
        MXAnalysisWithoutFPGA analysis(experiment, mapping, pixel_mask, indexer);

        AzimuthalIntegrationProfile profile(mapping);

        while (true) {
            int current_idx_offset = processed_count.fetch_add(1);
            int image_idx = start_image + current_idx_offset;

            if (image_idx >= end_image) break;

            // Load Image
            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_idx;
            msg.image_collection_efficiency = dataset->efficiency[image_idx];

            total_uncompressed_bytes += msg.image.GetUncompressedSize();

            auto image_start_time = std::chrono::high_resolution_clock::now();

            // Analyze
            try {
                analysis.Analyze(msg, profile, spot_settings);
            } catch (const std::exception &e) {
                logger.Error("Error analyzing image {}: {}", image_idx, e.what());
                continue;
            }

            auto image_end_time = std::chrono::high_resolution_clock::now();
            std::chrono::duration<float> image_duration = image_end_time - image_start_time;

            msg.processing_time_s = image_duration.count();
            msg.original_number = msg.number;
            msg.run_number = experiment.GetRunNumber();
            msg.run_name = experiment.GetRunName();

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

            // Update max sent tracking
            uint64_t current_max = max_image_number_sent.load();
            while (static_cast<uint64_t>(msg.number) > current_max) {
                if (max_image_number_sent.compare_exchange_weak(current_max, static_cast<uint64_t>(msg.number)))
                    break;
            }

            finished_count.fetch_add(1);

            // Progress log
            if (current_idx_offset > 0 && current_idx_offset % 100 == 0) {
                std::optional<float> indexing_rate = plots.GetIndexingRate();

                const auto now = std::chrono::steady_clock::now();
                const double elapsed_s = std::chrono::duration<double>(now - start_time).count();
                const int processed_images = finished_count.load();
                const double frame_rate_hz = (elapsed_s > 0.0) ? (processed_images / elapsed_s) : 0.0;

                if (indexing_rate.has_value()) {
                    logger.Info("Processed {} / {} images ({:.2f} Hz, indexing rate {:.1f}%)",
                                current_idx_offset, images_to_process,
                                frame_rate_hz,
                                indexing_rate.value() * 100.0f);
                } else {
                    logger.Info("Processed {} / {} images ({:.2f} Hz, indexing rate N/A)",
                                current_idx_offset, images_to_process,
                                frame_rate_hz);
                }
            }
        }

        // Finalize per-thread indexing (if any per-thread aggregation is needed, though pool handles most)
        // IndexAndRefine doesn't have a per-thread finalize that returns a lattice,
        // the main lattice determination is usually done on the aggregated results in the pool or main thread
    };

    // Launch threads
    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, i));
    }

    // Wait for completion
    for (auto &f: futures) {
        f.get();
    }

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

    // 5. Finalize Statistics and Write EndMessage
    EndMessage end_msg;
    end_msg.max_image_number = max_image_number_sent;
    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();

    // Gather statistics from plots
    end_msg.bkg_estimate = plots.GetBkgEstimate();
    end_msg.indexing_rate = plots.GetIndexingRate();
    end_msg.az_int_result["dataset"] = plots.GetAzIntProfile();

    // Finalize Indexing (Global) to get rotation lattice
    // We create a temporary IndexAndRefine to call Finalize() which aggregates pool results
    const auto rotation_indexer_ret = indexer.Finalize();

    if (rotation_indexer_ret.has_value()) {
        end_msg.rotation_lattice = rotation_indexer_ret->lattice;
        end_msg.rotation_lattice_type = LatticeMessage{
            .centering = rotation_indexer_ret->search_result.centering,
            .niggli_class = rotation_indexer_ret->search_result.niggli_class,
            .crystal_system = rotation_indexer_ret->search_result.system
        };
        logger.Info("Rotation Indexing found lattice");
    }

    // --- Optional: run scaling (mosaicity refinement) on accumulated reflections ---
    // --- Optional: run scaling (mosaicity refinement) on accumulated reflections ---
    if (run_scaling) {
        logger.Info("Running scaling (mosaicity refinement) ...");

        ScaleMergeOptions scale_opts;
        scale_opts.partiality_model = partiality_model;
        scale_opts.max_num_iterations = 500;
        scale_opts.max_solver_time_s = 240.0; // generous cutoff for now
        scale_opts.merge_friedel = !anomalous_mode;
        scale_opts.d_min_limit_A = d_min_scale_merge.value_or(0.0);

        const bool fixed_space_group = space_group || experiment.GetGemmiSpaceGroup().has_value();

        if (space_group)
            scale_opts.space_group = *space_group;
        else
            scale_opts.space_group = experiment.GetGemmiSpaceGroup();

        auto scale_start = std::chrono::steady_clock::now();
        auto scale_result = indexer.ScaleRotationData(scale_opts);
        auto scale_end = std::chrono::steady_clock::now();
        double scale_time = std::chrono::duration<double>(scale_end - scale_start).count();

        if (scale_result && !fixed_space_group) {
            logger.Info("Searching for space group from P1-merged reflections ...");

            SearchSpaceGroupOptions sg_opts;
            sg_opts.crystal_system.reset();
            sg_opts.centering = '\0';
            sg_opts.merge_friedel = !anomalous_mode;
            sg_opts.d_min_limit_A = d_min_scale_merge.value_or(0.0);
            sg_opts.min_i_over_sigma = 0.0;
            sg_opts.min_operator_cc = 0.80;
            sg_opts.min_pairs_per_operator = 20;
            sg_opts.min_total_compared = 100;
            sg_opts.test_systematic_absences = true;

            const auto sg_search = SearchSpaceGroup(scale_result->merged, sg_opts);

            logger.Info("");
            {
                std::istringstream iss(SearchSpaceGroupResultToText(sg_search));
                std::string line;
                while (std::getline(iss, line)) {
                    if (!line.empty())
                        logger.Info("{}", line);
                }
            }
            logger.Info("");

            if (sg_search.best_space_group.has_value()) {
                logger.Info("Re-running scaling in detected space group {}", sg_search.best_space_group->short_name());

                scale_opts.space_group = *sg_search.best_space_group;

                auto rescale_start = std::chrono::steady_clock::now();
                auto refined_scale_result = indexer.ScaleRotationData(scale_opts);
                auto rescale_end = std::chrono::steady_clock::now();

                if (refined_scale_result) {
                    scale_result = std::move(refined_scale_result);
                    scale_time += std::chrono::duration<double>(rescale_end - rescale_start).count();
                }
            } else {
                logger.Warning("No space group accepted; keeping P1-merged result");
            }
        }

        if (scale_result) {
            end_msg.scale_factor = scale_result->image_scale_g;

            logger.Info("Scaling completed in {:.2f} s  ({} unique reflections)",
                        scale_time, scale_result->merged.size());

            // Print resolution-shell statistics table
            {
                const auto &stats = scale_result->statistics;
                logger.Info("");
                logger.Info("  {:>8s} {:>8s} {:>8s} {:>8s} {:>8s} {:>10s}",
                            "d_min", "N_obs", "N_uniq", "Rmeas", "<I/sig>", "Complete");
                logger.Info("  {:->8s} {:->8s} {:->8s} {:->8s} {:->8s} {:->10s}",
                            "", "", "", "", "", "");
                for (const auto &sh: stats.shells) {
                    if (sh.unique_reflections == 0)
                        continue;
                    std::string compl_str = (sh.completeness > 0.0)
                                                ? fmt::format("{:8.1f}%", sh.completeness * 100.0)
                                                : "      N/A";
                    logger.Info("  {:8.2f} {:8d} {:8d} {:8.3f}% {:8.1f} {:>10s}",
                                sh.d_min, sh.total_observations, sh.unique_reflections,
                                sh.rmeas * 100, sh.mean_i_over_sigma, compl_str);
                }
                {
                    const auto &ov = stats.overall;
                    logger.Info("  {:->8s} {:->8s} {:->8s} {:->8s} {:->8s} {:->10s}",
                                "", "", "", "", "", "");
                    std::string compl_str = (ov.completeness > 0.0)
                                                ? fmt::format("{:8.1f}%", ov.completeness * 100.0)
                                                : "      N/A";
                    logger.Info("  {:>8s} {:8d} {:8d} {:8.3f}% {:8.1f} {:>10s}",
                                "Overall", ov.total_observations, ov.unique_reflections,
                                ov.rmeas * 100, ov.mean_i_over_sigma, compl_str);
                }
                logger.Info("");
            }

            {
                const std::string img_path = output_prefix + "_image.dat";
                std::ofstream img_file(img_path);
                if (!img_file) {
                    logger.Error("Cannot open {} for writing", img_path);
                } else {
                    img_file << "# image_id  mosaicity_deg  K\n";
                    for (size_t i = 0; i < scale_result->mosaicity_deg.size(); ++i) {
                        img_file << i << " " << scale_result->mosaicity_deg[i] << " " << scale_result->image_scale_g[i]
                                << "\n";
                    }
                    img_file.close();
                }
            }

            {
                FrenchWilsonOptions fw_opts;
                fw_opts.acentric = true; // typical for MX
                fw_opts.num_shells = 20;

                auto fw = FrenchWilson(scale_result->merged, fw_opts);
                {
                    {
                        const std::string hkl_path = output_prefix + "_amplitudes.hkl";
                        std::ofstream hkl_file(hkl_path);
                        if (!hkl_file) {
                            logger.Error("Cannot open {} for writing", hkl_path);
                        } else {
                            for (const auto &r: fw) {
                                hkl_file << r.h << " " << r.k << " " << r.l << " "
                                        << r.F << " " << r.sigmaF << " "
                                        << r.I << " " << r.sigmaI
                                        << "\n";
                            }
                            hkl_file.close();
                            logger.Info("Wrote {} reflections to {}", scale_result->merged.size(), hkl_path);
                        }
                    }
                    MmcifMetadata cif_meta;

                    if (rotation_indexer_ret.has_value()) {
                        cif_meta.unit_cell = rotation_indexer_ret->lattice.GetUnitCell();
                    } else if (experiment.GetUnitCell().has_value()) {
                        cif_meta.unit_cell = experiment.GetUnitCell().value();
                    }

                    if (scale_opts.space_group.has_value()) {
                        cif_meta.space_group_name = scale_opts.space_group->hm;
                        cif_meta.space_group_number = scale_opts.space_group->number;
                    } else if (auto sg = experiment.GetGemmiSpaceGroup(); sg.has_value()) {
                        cif_meta.space_group_name = sg->hm;
                        cif_meta.space_group_number = sg->number;
                    }

                    cif_meta.detector_name = experiment.GetDetectorDescription();
                    cif_meta.wavelength_A = experiment.GetWavelength_A();
                    cif_meta.detector_distance_mm = experiment.GetDetectorDistance_mm();
                    cif_meta.sample_temperature_K = experiment.GetSampleTemperature_K();
                    cif_meta.sample_name = experiment.GetSampleName();
                    cif_meta.data_block_name = output_prefix;

                    cif_meta.beamline = experiment.GetInstrumentName();
                    cif_meta.source = experiment.GetSourceName();

                    const std::string cif_path = output_prefix + "_amplitudes.cif";
                    try {
                        WriteMmcifReflections(cif_path, fw, cif_meta);
                        logger.Info("Wrote mmCIF reflections to {}", cif_path);
                    } catch (const std::exception &e) {
                        logger.Error("Failed to write mmCIF: {}", e.what());
                    }
                }
            }
        } else {
            logger.Warning("Scaling skipped — too few reflections accumulated (need >= 20)");
            logger.Info("Scaling wall-clock time: {:.2f} s", scale_time);
        }
    }

    // Write End Message
    if (writer) {
        writer->WriteHDF5(end_msg);
        auto stats = writer->Finalize();
    }

    // 6. Report Statistics to Console
    double processing_time = std::chrono::duration<double>(end_time - start_time).count();
    double throughput_MBs = static_cast<double>(total_uncompressed_bytes) / (processing_time * 1e6);
    double frame_rate = static_cast<double>(images_to_process) / processing_time;

    logger.Info("");
    logger.Info("Processing time: {:.2f} s", processing_time);
    logger.Info("Frame rate:      {:.2f} Hz", frame_rate);
    logger.Info("Total throughput:{:.2f} MB/s", throughput_MBs);

    // Print extended stats similar to Receiver
    if (!end_msg.indexing_rate.has_value()) {
        logger.Info("Indexing rate:   {:.2f}%", end_msg.indexing_rate.value() * 100.0);
    }

    auto image_mean_time = plots.GetMeanProcessingTime();
    logger.Info("Per-image time: (mean; microseconds): decompress {:.0f} preprocess {:.0f} azint {:.0f} spot finding {:.0f} indexing {:.0f} refinement {:.0f} prediction {:.0f} integration {:.0f} total {:.0f}",
                image_mean_time.compression * 1e6,
                image_mean_time.preprocessing * 1e6,
                image_mean_time.azint * 1e6,
                image_mean_time.spot_finding * 1e6,
                image_mean_time.indexing * 1e6,
                image_mean_time.refinement * 1e6,
                image_mean_time.bragg_prediction * 1e6,
                image_mean_time.integration * 1e6,
                image_mean_time.processing * 1e6);

    if (rotation_indexer_ret.has_value()) {
        auto latt = rotation_indexer_ret->lattice;
        if (auto axis_ = experiment.GetGoniometer()) {
            const float angle_deg = axis_->GetAngle_deg(0);
            const auto rot = axis_->GetTransformationAngle(angle_deg);
            latt = latt.Multiply(rot);
        }

        auto vec0 = rotation_indexer_ret->lattice.Vec0();
        auto vec1 = rotation_indexer_ret->lattice.Vec1();
        auto vec2 = rotation_indexer_ret->lattice.Vec2();
        auto uc = rotation_indexer_ret->lattice.GetUnitCell();
        logger.Info("Lattice vec0: {:.3f}, {:.3f}, {:.3f}", vec0.x, vec0.y, vec0.z);
        logger.Info("Lattice vec1: {:.3f}, {:.3f}, {:.3f}", vec1.x, vec1.y, vec1.z);
        logger.Info("Lattice vec2: {:.3f}, {:.3f}, {:.3f}", vec2.x, vec2.y, vec2.z);
        logger.Info("Lattice:         a={:.2f} b={:.2f} c={:.2f} alpha={:.2f} beta={:.2f} gamma={:.2f}",
                    uc.a, uc.b, uc.c, uc.alpha, uc.beta, uc.gamma);
    }
}