// 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 <getopt.h>

#include "../reader/JFJochHDF5Reader.h"
#include "../common/Logger.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../common/AzimuthalIntegrationMapping.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/LoadFCalcFromMtz.h"
#include "../image_analysis/scale_merge/Merge.h"
#include "../image_analysis/scale_merge/SearchSpaceGroup.h"
#include "../image_analysis/WriteReflections.h"
#include "../image_analysis/UpdateReflectionResolution.h"

void print_usage() {
    std::cout << "Usage ./jfjoch_scale {<options>} <input.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 << "   -v, --verbose                         Verbose output" << std::endl;
    std::cout << "" << std::endl;
    std::cout << "  Scaling and merging" << std::endl;
    std::cout << "   -S, --space-group <num>               Space group number" << std::endl;
    std::cout <<
            "   -D, --scaling-high-resolution <num>   High resolution limit for scaling/merging (default: 0.0; no limit)"
            << std::endl;
    std::cout << "   -P, --partiality <txt>                Partiality refinement fixed|rot|unity (default: fixed)" <<
            std::endl;
    std::cout << "   -A, --anomalous                       Anomalous mode (don't merge Friedel pairs)" << std::endl;
    std::cout << "   -B, --refine-bfactor                  Refine per image B-factor" << std::endl;
    std::cout << "   -w, --wedge[=num]                     Refine image wedge during scaling with starting wedge value"
            << std::endl;
    std::cout << "   --min-partiality <num>            Minimum partiality to accept reflection (default: 0.02)" <<
            std::endl;
    std::cout << "   --min-image-cc <num>              Per-image CC limit in percent (default: no limit)" << std::endl;
    std::cout << "   --scaling-iterations <num>        Number of scaling iterations with no reference data (default: 3)"
            << std::endl;
    std::cout << "   -z, --reference-mtz <file>            Reference MTZ file" << std::endl;
}

enum {
    OPT_MIN_PARTIALITY = 1000,
    OPT_MIN_IMAGE_CC,
    OPT_SCALING_ITERATIONS
};


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'},
    {"reference-mtz", required_argument, nullptr, 'z'},
    {"space-group", required_argument, nullptr, 'S'},
    {"partiality", required_argument, nullptr, 'P'},
    {"anomalous", no_argument, nullptr, 'A'},
    {"refine-bfactor", no_argument, nullptr, 'B'},
    {"wedge", optional_argument, nullptr, 'w'},
    {"scaling-high-resolution", required_argument, nullptr, 'D'},
    {"min-partiality", required_argument, nullptr, OPT_MIN_PARTIALITY},
    {"min-image-cc", required_argument, nullptr, OPT_MIN_IMAGE_CC},
    {"scaling-iterations", required_argument, nullptr, OPT_SCALING_ITERATIONS},
    {nullptr, 0, nullptr, 0}
};

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

    print_license("jfjoch_scale");

    Logger logger("jfjoch_scale");

    std::string input_file;
    std::string output_prefix = "output";
    int nthreads = 0;
    int start_image = 0;
    int end_image = -1; // -1 indicates process until end
    bool verbose = false;
    bool anomalous_mode = false;
    std::optional<int64_t> space_group_number;
    bool refine_bfactor = false;
    bool refine_wedge = false;
    std::optional<double> wedge_for_scaling;
    std::string ref_mtz;
    double min_partiality = 0.02;
    double min_image_cc = 0.0;
    int64_t scaling_iter = 3;

    PartialityModel partiality_model = PartialityModel::Fixed;

    std::optional<float> d_min_scale_merge;

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

    int opt;
    int option_index = 0;
    const char *short_opts = "vo:N:s:e:z:S:P:ABw::D:p:q:i:";

    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 'z':
                ref_mtz = optarg;
                break;
            case 'S':
                space_group_number = atoi(optarg);
                break;
            case 'P':
                if (strcmp(optarg, "unity") == 0)
                    partiality_model = PartialityModel::Unity;
                else if (strcmp(optarg, "fixed") == 0)
                    partiality_model = PartialityModel::Fixed;
                else if (strcmp(optarg, "rot") == 0)
                    partiality_model = PartialityModel::Rotation;
                else {
                    logger.Error("Invalid partiality mode: {}", optarg);
                    print_usage();
                    exit(EXIT_FAILURE);
                }
                break;
            case 'A':
                anomalous_mode = true;
                break;
            case 'B':
                refine_bfactor = true;
                break;
            case 'w':
                refine_wedge = true;
                if (optarg)
                    wedge_for_scaling = std::stod(optarg);
                break;
            case 'D':
                d_min_scale_merge = atof(optarg);
                logger.Info("High resolution limit for scaling/merging set to {:.2f} A", d_min_scale_merge.value());
                break;
            case OPT_MIN_PARTIALITY:
                min_partiality = std::stod(optarg);
                break;
            case OPT_MIN_IMAGE_CC:
                min_image_cc = std::stod(optarg);
                break;
            case OPT_SCALING_ITERATIONS:
                scaling_iter = atoi(optarg);
                if (scaling_iter <= 0) {
                    logger.Error("Invalid scaling iteration count: {}", scaling_iter);
                    exit(EXIT_FAILURE);
                }
                break;
        }
    }

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

    std::vector<std::vector<Reflection>> reflections;
    std::vector<float> mosaicity;

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

    logger.Verbose(verbose);
    logger.Info("Loading reflections");

    // 1. Read Input files
    JFJochHDF5Reader reader;
    std::shared_ptr<const JFJochReaderDataset> dataset;
    for (int i = optind; i < argc; i++) {
        input_file = argv[optind];
        try {
            reader.ReadFile(input_file);
            auto tmp_dataset = reader.GetDataset();
            if (i == optind)
                dataset = tmp_dataset;

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

            size_t mosaicity_size = mosaicity.size();
            mosaicity.resize(mosaicity_size + end_image - start_image + 1);
            for (int i = 0; i < end_image - start_image + 1; i++)
                mosaicity[mosaicity_size + i] = dataset->mosaicity_deg[start_image + i];

            reader.ReadReflections(reflections, start_image, end_image);
            logger.Info("Loaded dataset from {}", input_file);
        } catch (const std::exception &e) {
            logger.Error("Error reading input file: {}", e.what());
            // Hard exit is only for the first file
            if (i == optind)
                exit(EXIT_FAILURE);
        }
    }


    std::vector<MergedReflection> reference_data;
    if (!ref_mtz.empty()) {
        reference_data = LoadFCalcFromMtz(ref_mtz);
        logger.Info("Loaded {} reflections from {} MTZ file", reference_data.size(), ref_mtz);
    }

    logger.Info("Starting analysis of {} images", reflections.size());

    // 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);
    experiment.SpaceGroupNumber(space_group_number);
    experiment.NumTriggers(1);

    ScalingSettings scaling_settings;
    scaling_settings.SetPartialityModel(partiality_model);
    if (d_min_scale_merge)
        scaling_settings.HighResolutionLimit_A(d_min_scale_merge.value());
    scaling_settings.MergeFriedel(!anomalous_mode);
    scaling_settings.RefineB(refine_bfactor);
    scaling_settings.RefineRotationWedge(refine_wedge);
    if (wedge_for_scaling.has_value())
        scaling_settings.RotationWedgeForScaling(wedge_for_scaling);
    scaling_settings.MinPartiality(min_partiality);
    scaling_settings.MinCCForImage(min_image_cc);

    experiment.ImportScalingSettings(scaling_settings);

    if (!experiment.GetUnitCell()) {
        logger.Error("Experiment unit cell not found, cannot update reflection resolution");
        exit(EXIT_FAILURE);
    }

    auto refl_stats = UpdateReflectionResolution(experiment.GetUnitCell().value(), reflections);

    logger.Info("Read {} reflections from {} images", refl_stats.n_reflections, refl_stats.n_images);

    experiment.ImagesPerTrigger(refl_stats.n_images);

    ScalingResult scale_result(0);

    auto scale_start = std::chrono::steady_clock::now();
    for (int i = 0; i < scaling_iter; i++) {
        auto iter_start = std::chrono::steady_clock::now();

        if (reference_data.empty()) {
            auto merged = MergeAll(experiment, reflections);
            ScaleOnTheFly scaling(experiment, merged);
            scale_result = scaling.Scale(reflections, mosaicity, nthreads);
        } else {
            ScaleOnTheFly scaling(experiment, reference_data);
            scale_result = scaling.Scale(reflections, mosaicity, nthreads);
        }

        scale_result.SaveToFile(output_prefix + "_iter" + std::to_string(i) + "_scale.dat");
        auto iter_end = std::chrono::steady_clock::now();
        double iter_time = std::chrono::duration<double>(iter_end - iter_start).count();
        logger.Info("Scaling iteration {} took {:.3f} seconds", i, iter_time);
    }

    auto scale_end = std::chrono::steady_clock::now();
    double scale_time = std::chrono::duration<double>(scale_end - scale_start).count();
    logger.Info("Scaling completed in {:.2f} s", scale_time);

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

    std::vector<uint8_t> merge_mask(reflections.size(), 1);
    auto rejected = CalcMergeMaskCC(experiment, scale_result.image_cc, merge_mask);
    if (rejected > 0)
        logger.Info("Rejected {} images due to low CC with reference", rejected);

    auto merge_result = MergeAll(experiment, reflections, merge_mask);
    auto merge_stats = MergeStats(experiment, merge_result, reflections, experiment.GetUnitCell().value(), merge_mask,
                                  reference_data);

    auto merge_end = std::chrono::steady_clock::now();
    double merge_time = std::chrono::duration<double>(merge_end - merge_start).count();

    logger.Info("Merge completed in {:.2f} s ({} unique reflections)", merge_time,
                merge_result.size());

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

    if (!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 = experiment.GetScalingSettings().GetMergeFriedel();
        sg_opts.d_min_limit_A = experiment.GetScalingSettings().GetHighResolutionLimit_A().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(merge_result, 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("");
    }

    // Print resolution-shell statistics table
    std::cout << merge_stats;

    if (!output_prefix.empty())
        WriteReflections(merge_result, *experiment.GetUnitCell(), experiment, output_prefix);
}