// 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/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/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("");
    logger.Info("  Scaling and merging");
    logger.Info("   -S<num>    Space group number");
    logger.Info("   -D<num>    High resolution limit for scaling/merging (default: 0.0; no limit)");
    logger.Info("   -P<txt>    Partiality refinement fixed|rot|unity (default: fixed)");
    logger.Info("   -A         Anomalous mode (don't merge Friedel pairs)");
    logger.Info("   -B         Refine per image B-factor");
    logger.Info("   -w<num>    Refine image wedge during scaling with starting wedge value");
    logger.Info("   -p<num>    Minimum partiality to accept reflection (default: 0.02)");
    logger.Info("   -q<num>    Per-image CC limit in percent (default: no limit)");
    logger.Info("   -i<num>    Number of scaling iterations with no reference data (default: 3) ");
}

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(logger);
        exit(EXIT_FAILURE);
    }

    int opt;
    while ((opt = getopt(argc, argv, "o:z:N:s:e:p:q:Bw::vD:S:A:P:")) != -1) {
        switch (opt) {
            case 'o':
                output_prefix = optarg;
                break;
            case 'z':
                ref_mtz = optarg;
                break;
            case 'N':
                nthreads = atoi(optarg);
                break;
            case 's':
                start_image = atoi(optarg);
                break;
            case 'e':
                end_image = atoi(optarg);
                break;
            case 'p':
                min_partiality = std::stod(optarg);
                break;
            case 'q':
                min_image_cc = std::stod(optarg);
                break;
            case 'B':
                refine_bfactor = true;
                break;
            case 'w':
                refine_wedge = true;
                if (optarg)
                    wedge_for_scaling = std::stod(optarg);
                break;
            case 'v':
                verbose = true;
                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 'S':
                space_group_number = atoi(optarg);
                break;
            case 'A':
                anomalous_mode = true;
                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(logger);
                    exit(EXIT_FAILURE);
                }
                break;
            case 'i':
                scaling_iter = atoi(optarg);
                if (scaling_iter <= 0) {
                    logger.Error("Invalid scaling iteration count: {}", scaling_iter);
                    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);

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

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

    // 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.ImagesPerTrigger(images_to_process);
    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);

    logger.Info("Running scaling (mosaicity refinement) ...");
    auto reflections = reader.ReadReflections(start_image, end_image);
    std::vector<float> mosaicity(end_image - start_image + 1);
    for (int i = 0; i < end_image - start_image + 1; i++) {
        mosaicity[i] = dataset->mosaicity_deg[start_image + i];
    }

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

    auto merge_mask = CalcMergeMask(experiment, scale_result);
    auto merge_result = MergeAll(experiment, reflections, merge_mask);
    auto merge_stats = MergeStats(experiment, merge_result, reflections, merge_mask);

    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
        merge_stats.Print(logger);

        {
            const std::string hkl_path = output_prefix + "_intensities.hkl";
            std::ofstream hkl_file(hkl_path);
            if (!hkl_file) {
                logger.Error("Cannot open {} for writing", hkl_path);
            } else {
                for (const auto &r: merge_result) {
                    hkl_file << r.h << " " << r.k << " " << r.l << " "
                            << r.I << " " << r.sigma
                            << "\n";
                }
                hkl_file.close();
                logger.Info("Wrote {} reflections to {}", merge_result.size(), hkl_path);
            }
        }

        try {
            const std::string cif_path = output_prefix + "_intensities.cif";

            MmcifMetadata cif_meta;

            cif_meta.Fill(experiment);
            cif_meta.data_block_name = output_prefix;

            WriteMmcifReflections(cif_path, merge_result, cif_meta);
            logger.Info("Wrote mmCIF reflections to {}", cif_path);
        } catch (const std::exception &e) {
            logger.Error("Failed to write mmCIF: {}", e.what());
        }
    }
}