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

#include <fstream>

#include "XdsIntegrateParser.h"
#include "../reader/JFJochHDF5Reader.h"
#include "../common/print_license.h"
#include "../common/Logger.h"
#include "../common/hkl_key.h"

void print_usage(Logger &logger) {
    logger.Info("Usage ./jfjoch_extract_hkl {<options>} <path to master file>");
    logger.Info("Options:");
    logger.Info("   -I<num>     Number of images");
    logger.Info("   -o<file>    Output filename");
    logger.Info("   -x<num>     Max dist from Ewald sphere");
    logger.Info("   -R{<file>}  Sum same HKL across neighboring images (image +/- 1) - optional reference INTEGRATE.HKL file name");
}


int main(int argc, char **argv) {
    int64_t image_number = 0;
    float max_dist_ewald_sphere = 1.0;
    bool sum_neighboring = false;
    std::string ref_file;

    std::string output_filename = "out.hkl";

    print_license("jfjoch_extract_hkl");

    Logger logger("jfjoch_extract_hkl");
    logger.Verbose(true);
    int opt;
    while ((opt = getopt(argc, argv, "I:x:o:R::")) != -1) {
        switch (opt) {
            case 'I':
                image_number = atol(optarg);
                break;
            case 'x':
                max_dist_ewald_sphere = atof(optarg);
                break;
            case 'o':
                output_filename = optarg;
                break;
            case 'R':
                sum_neighboring = true;
                if (optarg)
                    ref_file = std::string(optarg);
                break;
            default: /* '?' */
                print_usage(logger);
                exit(EXIT_FAILURE);
        }
    }

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

    JFJochHDF5Reader reader;
    reader.ReadFile(argv[optind]);

    auto dataset = reader.LoadImage(image_number);

    if (dataset) {
        if (sum_neighboring) {
            CrystalLattice latt;
            IntegrateMap agg;
            int64_t total_images = reader.GetNumberOfImages();

            IntegrateMap xds_result;
            if (!ref_file.empty())
                xds_result = ParseXdsIntegrateHkl(ref_file);

            for (int i = 0; i < total_images; ++i) {
                auto dataset = reader.LoadImage(i);
                if (!dataset) continue;
                if (dataset->ImageData().indexing_lattice)
                    latt = dataset->ImageData().indexing_lattice.value();
                for (const auto &r: dataset->ImageData().reflections) {
                    int64_t key = hkl_key(r.h, r.k, r.l);
                    auto it = agg.find(key);
                    HKLData data{
                        .h = r.h,
                        .k = r.k,
                        .l = r.l,
                        .I = r.I,
                        .last_image = i,
                        .count = 1,
                        .rlp = r.rlp,
                        .image_number = r.image_number
                    };
                    bool found = false;
                    if (it != agg.end()) {
                        for (auto &val: it->second) {
                            if (val.last_image == i - 1) {
                                val.I += r.I;
                                val.last_image = i;
                                val.count++;
                                found = true;
                                break;
                            }
                        }
                    }
                    if (!found)
                        agg[key].push_back(data);
                }
            }

            std::fstream f(output_filename, std::ios::out);
            for (const auto &[key, val]: agg) {
                if (!val.empty()) {

                    if (xds_result.empty()) {
                        f << val[0].h << " " << val[0].k << " " << val[0].l << " " << val[0].I;
                        f << std::endl;
                    } else {
                        auto xds_it = xds_result.find(key);
                        if (xds_it != xds_result.end() && !xds_it->second.empty()) {
                            f << val[0].h << " " << val[0].k << " " << val[0].l << " " << val[0].I;

                            f << " " << xds_it->second[0].I << " " << val[0].rlp << " " << xds_it->second[0].rlp;
                            f << std::endl;
                        }
                    }


                }
            }

            auto cc_result = ComputeCcByResolution(latt, agg,xds_result, 1.0, 50.0, 25);
            for (int i = 0; i < cc_result.cc.size(); ++i)
                std::cout << 1/ std::sqrt(cc_result.shell_mean_one_over_d2[i]) << " " << cc_result.cc[i]* 100.0 << " " <<  cc_result.pairs[i] << std::endl;

        } else {
            auto geom = dataset->Dataset().experiment.GetDiffractionGeometry();

            std::fstream f(output_filename, std::ios::out);
            if (dataset->ImageData().indexing_lattice) {
                auto latt = dataset->ImageData().indexing_lattice.value();

                Coord astar = latt.Astar();
                Coord bstar = latt.Bstar();
                Coord cstar = latt.Cstar();
                auto uc = latt.GetUnitCell();
                logger.Info("Cell {} {} {} {} {} {}", uc.a, uc.b, uc.c, uc.alpha, uc.beta, uc.gamma);
                logger.Info("A {} {} {}", astar.x, astar.y, astar.z);
                logger.Info("B {} {} {}", bstar.x, bstar.y, bstar.z);
                logger.Info("C {} {} {}", cstar.x, cstar.y, cstar.z);

                logger.Info("{} reflections identified", dataset->ImageData().reflections.size());
                for (const auto &r: dataset->ImageData().reflections) {
                    auto recip = astar * r.h + bstar * r.k + cstar * r.l;
                    auto dist = geom.DistFromEwaldSphere(recip);
                    if (dist < max_dist_ewald_sphere)
                        f << r.l << " " << r.k << " " << r.h << " "
                                << r.I << " " << r.sigma << " " << dist << " "
                                << r.predicted_x << " " << r.predicted_y << " " << r.delta_phi_deg
                                << std::endl;
                }
            }
        }
    }
}