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

#include "GeometryRefiner.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <vector>

#include "../../common/JFJochMath.h"   // PI
#include "LatticeReduction.h"
#include "XtalResidual.h"
#include "ceres/ceres.h"

namespace {

// Trigonal is described with the hexagonal cell + B, same as the per-image XtalOptimizer.
gemmi::CrystalSystem MapSystem(gemmi::CrystalSystem s) {
    return (s == gemmi::CrystalSystem::Trigonal) ? gemmi::CrystalSystem::Hexagonal : s;
}

// The cell angles (radians) that XtalResidual bakes into its B matrix for a given system. Only
// monoclinic (beta) and triclinic (all three) read p2; the rest use fixed symmetry angles. We keep the
// angles constant in v1, so these are set once from the input cell.
void InitCellAngles(gemmi::CrystalSystem system, const UnitCell &cell, double ang[3]) {
    const double d2r = PI / 180.0;
    switch (system) {
        case gemmi::CrystalSystem::Hexagonal:
            ang[0] = PI / 2.0; ang[1] = PI / 2.0; ang[2] = 2.0 * PI / 3.0;
            break;
        case gemmi::CrystalSystem::Monoclinic:
            // XtalResidual reads p2[0] as beta.
            ang[0] = cell.beta * d2r; ang[1] = PI / 2.0; ang[2] = PI / 2.0;
            break;
        case gemmi::CrystalSystem::Triclinic:
            ang[0] = cell.alpha * d2r; ang[1] = cell.beta * d2r; ang[2] = cell.gamma * d2r;
            break;
        default:  // orthorhombic / tetragonal / cubic
            ang[0] = PI / 2.0; ang[1] = PI / 2.0; ang[2] = PI / 2.0;
            break;
    }
}

// Effective (lengths, alpha,beta,gamma) that reproduce the lattice XtalResidual builds from (len, ang),
// so a frame's lattice can be reconstructed with AngleAxisAndCellToLattice for HKL re-assignment.
void EffectiveCell(gemmi::CrystalSystem system, const double len[3], const double ang[3],
                   double lengths[3], double &alpha, double &beta, double &gamma) {
    switch (system) {
        case gemmi::CrystalSystem::Hexagonal:
            lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[2];
            alpha = PI / 2.0; beta = PI / 2.0; gamma = 2.0 * PI / 3.0;
            break;
        case gemmi::CrystalSystem::Tetragonal:
            lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[2];
            alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
            break;
        case gemmi::CrystalSystem::Cubic:
            lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[0];
            alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
            break;
        case gemmi::CrystalSystem::Monoclinic:
            lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
            alpha = PI / 2.0; beta = ang[0]; gamma = PI / 2.0;
            break;
        case gemmi::CrystalSystem::Triclinic:
            lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
            alpha = ang[0]; beta = ang[1]; gamma = ang[2];
            break;
        default:  // orthorhombic
            lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
            alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
            break;
    }
}

// Orientation seed (angle-axis) for a frame's lattice, using the per-system decomposition the per-image
// optimizer uses so it matches XtalResidual's B convention.
void OrientationSeed(gemmi::CrystalSystem system, const CrystalLattice &latt, double rod[3]) {
    double scratch[3];
    if (system == gemmi::CrystalSystem::Hexagonal) {
        LatticeToRodriguesAndLengths_Hex(latt, rod, scratch);
    } else if (system == gemmi::CrystalSystem::Monoclinic) {
        double beta;
        LatticeToRodriguesLengthsBeta_Mono(latt, rod, scratch, beta);
    } else {
        LatticeToRodriguesAndLengths_GS(latt, rod, scratch);
    }
}

UnitCell BuildUnitCell(gemmi::CrystalSystem system, const double len[3], const double ang[3]) {
    const double r2d = 180.0 / PI;
    UnitCell uc{};
    switch (system) {
        case gemmi::CrystalSystem::Hexagonal:
            uc = {(float)len[0], (float)len[0], (float)len[2], 90.0f, 90.0f, 120.0f};
            break;
        case gemmi::CrystalSystem::Tetragonal:
            uc = {(float)len[0], (float)len[0], (float)len[2], 90.0f, 90.0f, 90.0f};
            break;
        case gemmi::CrystalSystem::Cubic:
            uc = {(float)len[0], (float)len[0], (float)len[0], 90.0f, 90.0f, 90.0f};
            break;
        case gemmi::CrystalSystem::Monoclinic:
            uc = {(float)len[0], (float)len[1], (float)len[2], 90.0f, (float)(ang[0] * r2d), 90.0f};
            break;
        case gemmi::CrystalSystem::Triclinic:
            uc = {(float)len[0], (float)len[1], (float)len[2],
                  (float)(ang[0] * r2d), (float)(ang[1] * r2d), (float)(ang[2] * r2d)};
            break;
        default:  // orthorhombic
            uc = {(float)len[0], (float)len[1], (float)len[2], 90.0f, 90.0f, 90.0f};
            break;
    }
    return uc;
}

// Anchors the cell lengths to the input cell. Penalises each length's deviation from its input value
// with a large weight, breaking the low-resolution distance<->cell-scale degeneracy (position ~
// distance / cell_scale) so distance is determined by the data with the cell pinned.
struct CellLengthRegularizer {
    double w0, w1, w2;
    double l0, l1, l2;
    template <typename T>
    bool operator()(const T *const len, T *r) const {
        r[0] = T(w0) * (len[0] - T(l0));
        r[1] = T(w1) * (len[1] - T(l1));
        r[2] = T(w2) * (len[2] - T(l2));
        return true;
    }
};

}   // namespace

GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_geom,
                                           const UnitCell &input_cell,
                                           const std::vector<GeomRefineFrame> &frames,
                                           const GeometryRefinerSettings &settings) {
    GeometryRefinerResult result;
    try {
        const gemmi::CrystalSystem system = MapSystem(settings.crystal_system);

        // Shared geometry parameter blocks (refined). Detector tilt and the rotation axis are constant.
        double beam[2] = {nominal_geom.GetBeamX_pxl(), nominal_geom.GetBeamY_pxl()};
        double distance_mm = nominal_geom.GetDetectorDistance_mm();
        double detector_rot[2] = {nominal_geom.GetPoniRot1_rad(), nominal_geom.GetPoniRot2_rad()};
        double rot_vec[3] = {1.0, 0.0, 0.0};
        const double rot3 = nominal_geom.GetPoniRot3_rad();
        const double lambda = nominal_geom.GetWavelength_A();
        const double pixel = nominal_geom.GetPixelSize_mm();

        // Shared cell blocks. Lengths refined (regularized); angles constant.
        double cell_len[3] = {input_cell.a, input_cell.b, input_cell.c};
        const double cell_len0[3] = {input_cell.a, input_cell.b, input_cell.c};
        double cell_ang[3];
        InitCellAngles(system, input_cell, cell_ang);

        double eff_len[3], eff_alpha, eff_beta, eff_gamma;   // recomputed each round from cell_len/ang

        // Per-frame orientation seeds.
        std::vector<std::array<double, 3>> orient(frames.size());
        for (size_t i = 0; i < frames.size(); ++i)
            OrientationSeed(system, frames[i].lattice, orient[i].data());

        // Reciprocal length of one detector pixel at the nominal geometry, used to scale the robust
        // loss (~3 px) into the reciprocal-space residual units.
        const double recip_per_px = pixel / (distance_mm * lambda);
        const double loss_scale = 3.0 * recip_per_px;

        // Tolerance schedule (fractional-hkl norm), tightening each round like the per-image optimizer.
        auto round_tol = [&](int r) {
            if (settings.rounds <= 1) return 0.2;
            const double t = static_cast<double>(r) / (settings.rounds - 1);
            return 0.3 - (0.3 - 0.1) * t;
        };

        for (int round = 0; round < settings.rounds; ++round) {
            EffectiveCell(system, cell_len, cell_ang, eff_len, eff_alpha, eff_beta, eff_gamma);

            DiffractionGeometry cur = nominal_geom;
            cur.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]).DetectorDistance_mm(distance_mm);

            const double tol = round_tol(round);
            const double tol_sq = tol * tol;

            ceres::Problem problem;
            auto ordering = std::make_shared<ceres::ParameterBlockOrdering>();

            int spots_used = 0, frames_used = 0;
            for (size_t i = 0; i < frames.size(); ++i) {
                const CrystalLattice latt = AngleAxisAndCellToLattice(
                    orient[i].data(), eff_len, eff_alpha, eff_beta, eff_gamma);
                const Coord v0 = latt.Vec0(), v1 = latt.Vec1(), v2 = latt.Vec2();

                int used_in_frame = 0;
                for (const auto &s : frames[i].spots) {
                    const Coord recip = cur.DetectorToRecip(s.x, s.y);
                    const double hf = recip * v0, kf = recip * v1, lf = recip * v2;
                    const double h = std::round(hf), k = std::round(kf), l = std::round(lf);
                    const double dn = (hf - h) * (hf - h) + (kf - k) * (kf - k) + (lf - l) * (lf - l);
                    if (dn > tol_sq)
                        continue;

                    problem.AddResidualBlock(
                        new ceres::AutoDiffCostFunction<XtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
                            new XtalResidual(s.x, s.y, lambda, pixel, rot3, 0.0, h, k, l, system)),
                        new ceres::CauchyLoss(loss_scale),
                        beam, &distance_mm, detector_rot, rot_vec, orient[i].data(), cell_len, cell_ang);
                    ++used_in_frame;
                    ++spots_used;
                }
                if (used_in_frame > 0) {
                    ordering->AddElementToGroup(orient[i].data(), 0);   // eliminate orientations first
                    ++frames_used;
                }
            }

            if (spots_used < settings.min_spots_total) {
                if (round == 0)
                    return result;   // ok stays false
                break;               // keep the previous round's solution
            }

            // Cell-length regularizer (weight ~ coeff * sqrt(N_obs) / L0^2 -> effectively pins the cell).
            const double sqrtN = std::sqrt(static_cast<double>(spots_used));
            CellLengthRegularizer *reg = new CellLengthRegularizer{
                settings.cell_reg_coeff * sqrtN / (cell_len0[0] * cell_len0[0]),
                settings.cell_reg_coeff * sqrtN / (cell_len0[1] * cell_len0[1]),
                settings.cell_reg_coeff * sqrtN / (cell_len0[2] * cell_len0[2]),
                cell_len0[0], cell_len0[1], cell_len0[2]};
            problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<CellLengthRegularizer, 3, 3>(reg), nullptr, cell_len);

            // Shared blocks go in the second elimination group; constant ones are stripped by Ceres.
            for (double *b : {beam, &distance_mm, detector_rot, rot_vec, cell_len, cell_ang})
                ordering->AddElementToGroup(b, 1);

            // Beam + distance bounds around the nominal value; cell lengths loosely bounded (the
            // regularizer does the anchoring).
            problem.SetParameterLowerBound(beam, 0, nominal_geom.GetBeamX_pxl() - settings.beam_range_px);
            problem.SetParameterUpperBound(beam, 0, nominal_geom.GetBeamX_pxl() + settings.beam_range_px);
            problem.SetParameterLowerBound(beam, 1, nominal_geom.GetBeamY_pxl() - settings.beam_range_px);
            problem.SetParameterUpperBound(beam, 1, nominal_geom.GetBeamY_pxl() + settings.beam_range_px);
            problem.SetParameterLowerBound(&distance_mm, 0,
                std::max(1.0, nominal_geom.GetDetectorDistance_mm() - settings.distance_range_mm));
            problem.SetParameterUpperBound(&distance_mm, 0,
                nominal_geom.GetDetectorDistance_mm() + settings.distance_range_mm);
            for (int j = 0; j < 3; ++j) {
                problem.SetParameterLowerBound(cell_len, j, 0.7 * cell_len0[j]);
                problem.SetParameterUpperBound(cell_len, j, 1.3 * cell_len0[j]);
            }

            problem.SetParameterBlockConstant(detector_rot);   // tilt not refined
            problem.SetParameterBlockConstant(rot_vec);
            problem.SetParameterBlockConstant(cell_ang);       // angles not refined in v1

            ceres::Solver::Options options;
            options.linear_solver_type = ceres::DENSE_SCHUR;
            options.linear_solver_ordering = ordering;
            options.num_threads = std::max(1, settings.num_threads);
            options.max_num_iterations = 100;
            options.max_solver_time_in_seconds = 60.0;
            options.logging_type = ceres::LoggingType::SILENT;
            options.minimizer_progress_to_stdout = false;

            ceres::Solver::Summary summary;
            ceres::Solve(options, &problem, &summary);

            result.frames_used = frames_used;
            result.spots_used = spots_used;
        }

        // Final fit quality: median detector residual of the retained spots at the refined geometry.
        EffectiveCell(system, cell_len, cell_ang, eff_len, eff_alpha, eff_beta, eff_gamma);
        DiffractionGeometry cur = nominal_geom;
        cur.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]).DetectorDistance_mm(distance_mm);
        std::vector<double> residuals;
        for (size_t i = 0; i < frames.size(); ++i) {
            const CrystalLattice latt = AngleAxisAndCellToLattice(
                orient[i].data(), eff_len, eff_alpha, eff_beta, eff_gamma);
            const Coord a = latt.Astar(), b = latt.Bstar(), c = latt.Cstar();
            const Coord v0 = latt.Vec0(), v1 = latt.Vec1(), v2 = latt.Vec2();
            for (const auto &s : frames[i].spots) {
                const Coord recip = cur.DetectorToRecip(s.x, s.y);
                const double hf = recip * v0, kf = recip * v1, lf = recip * v2;
                const double h = std::round(hf), k = std::round(kf), l = std::round(lf);
                const double dn = (hf - h) * (hf - h) + (kf - k) * (kf - k) + (lf - l) * (lf - l);
                if (dn > 0.1 * 0.1)
                    continue;
                const Coord recip_pred = a * (float)h + b * (float)k + c * (float)l;
                const auto [px, py] = cur.RecipToDetector(recip_pred);
                if (std::isfinite(px) && std::isfinite(py))
                    residuals.push_back(std::hypot(px - s.x, py - s.y));
            }
        }
        if (!residuals.empty()) {
            std::nth_element(residuals.begin(), residuals.begin() + residuals.size() / 2, residuals.end());
            result.median_residual_px = residuals[residuals.size() / 2];
        }

        result.ok = result.spots_used >= settings.min_spots_total && result.frames_used > 0;
        result.beam_x_px = beam[0];
        result.beam_y_px = beam[1];
        result.distance_mm = distance_mm;
        result.cell = BuildUnitCell(system, cell_len, cell_ang);
        return result;
    } catch (...) {
        result.ok = false;
        return result;
    }
}