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

#include "XtalRefine.h"

#include <array>

#include "Dual.h"
#include "../../common/ParallelFor.h"

namespace {
    // Ambient layout of the parameter vector and the order of the blocks in it.
    constexpr int OFF_BEAM = 0, OFF_ROT = 2, OFF_AXIS = 4, OFF_P0 = 7, OFF_LEN = 10, OFF_ANG = 13, N_AMBIENT = 16;

    // Derivative lanes. The observed half of a residual depends on beam, detector angles and spindle
    // (OBS lanes), the predicted half on orientation and cell (PRED lanes); each half is carried on a
    // dual number of its own width and the two are put side by side only in the sums.
    constexpr int OBS = 6, PRED = 9, LANES = OBS + PRED;
    using DO = Dual<OBS>;
    using DP = Dual<PRED>;

    // The reduced problem: beam (2) and orientation (3) only.
    constexpr int R_OBS = 2, R_PRED = 3, R_LANES = R_OBS + R_PRED;

    // Sums over one block of residuals, in lane coordinates.
    template<int L>
    struct Sums {
        double cost = 0.0;
        double g[L] = {};
        double H[L][L] = {};   // upper triangle

        void Add(double r, const double *J, double w2) {
            cost += 0.5 * w2 * r * r;
            for (int i = 0; i < L; i++) {
                const double wj = w2 * J[i];
                g[i] += wj * r;
                for (int j = i; j < L; j++)
                    H[i][j] += wj * J[j];
            }
        }

        void Add(const Sums &o) {
            cost += o.cost;
            for (int i = 0; i < L; i++) {
                g[i] += o.g[i];
                for (int j = i; j < L; j++)
                    H[i][j] += o.H[i][j];
            }
        }
    };

    std::vector<LMBlock> MakeBlocks(const XtalRefineProblem &p) {
        std::vector<LMBlock> blocks(6);
        blocks[0].offset = OFF_BEAM;
        blocks[0].size = 2;
        blocks[0].constant = p.beam_constant;

        blocks[1].offset = OFF_ROT;
        blocks[1].size = 2;
        blocks[1].constant = p.beam_and_orientation_only || p.detector_rot_constant;
        for (int j = 0; j < 2; j++) {
            blocks[1].lower[j] = p.detector_rot_lower[j];
            blocks[1].upper[j] = p.detector_rot_upper[j];
        }

        blocks[2].offset = OFF_AXIS;
        blocks[2].size = 3;
        blocks[2].constant = p.beam_and_orientation_only || p.rot_vec_constant;
        blocks[2].sphere = true;

        blocks[3].offset = OFF_P0;
        blocks[3].size = 3;

        blocks[4].offset = OFF_LEN;
        blocks[4].size = 3;
        blocks[4].constant = p.beam_and_orientation_only || p.latt_vec1_constant;

        blocks[5].offset = OFF_ANG;
        blocks[5].size = 3;
        blocks[5].constant = p.beam_and_orientation_only || p.latt_vec2_constant;

        for (int j = 0; j < 3; j++) {
            blocks[4].lower[j] = p.latt_vec1_lower[j];
            blocks[4].upper[j] = p.latt_vec1_upper[j];
            blocks[5].lower[j] = p.latt_vec2_lower[j];
            blocks[5].upper[j] = p.latt_vec2_upper[j];
        }
        return blocks;
    }

    // Where each lane lands among the tangent coordinates of the free blocks; -1 for a held one.
    template<size_t L>
    std::array<int, L> LaneToTangent(const std::vector<LMBlock> &blocks, const std::array<int, L> &lane_block,
                                     const std::array<int, L> &lane_index) {
        std::array<int, L> map{};
        std::array<int, 6> tangent_offset{};
        int t = 0;
        for (size_t b = 0; b < blocks.size(); b++) {
            tangent_offset[b] = t;
            t += blocks[b].TangentSize();
        }
        for (size_t l = 0; l < L; l++)
            map[l] = blocks[lane_block[l]].constant ? -1 : tangent_offset[lane_block[l]] + lane_index[l];
        return map;
    }

    template<int L>
    void ToTangent(const Sums<L> &s, const std::array<int, L> &map, Eigen::VectorXd &g, Eigen::MatrixXd &H) {
        for (int i = 0; i < L; i++) {
            if (map[i] < 0)
                continue;
            g[map[i]] += s.g[i];
            for (int j = i; j < L; j++) {
                if (map[j] < 0)
                    continue;
                H(map[i], map[j]) += s.H[i][j];
                if (map[i] != map[j])
                    H(map[j], map[i]) += s.H[i][j];
            }
        }
    }

    void AddPriors(const XtalRefineProblem &p, const double *x, double &cost, Eigen::VectorXd *g,
                   Eigen::MatrixXd *H, int beam_tangent, int rot_tangent) {
        for (const auto &prior: p.priors) {
            const bool on_beam = prior.block == XtalRefinePrior::Block::Beam;
            const double *v = x + (on_beam ? OFF_BEAM : OFF_ROT);
            const double r = prior.weight * (prior.gx * v[0] + prior.gy * v[1] - prior.p0);
            cost += 0.5 * r * r;
            const int t = on_beam ? beam_tangent : rot_tangent;
            if (!g || t < 0)
                continue;
            const double J[2] = {prior.weight * prior.gx, prior.weight * prior.gy};
            for (int i = 0; i < 2; i++) {
                (*g)[t + i] += J[i] * r;
                for (int j = 0; j < 2; j++)
                    (*H)(t + i, t + j) += J[i] * J[j];
            }
        }
    }

    template<class D>
    D Seed(double value, int lane, bool free) {
        return free ? D::Variable(value, lane) : D(value);
    }

    // Residual blocks of at least this many residuals; the cut depends on the count alone.
    constexpr int MIN_RESIDUALS_PER_BLOCK = 256;

    template<int L, class Fn>
    Sums<L> SumResiduals(const XtalRefineProblem &p, int num_threads, Fn &&residual) {
        const int n = static_cast<int>(p.residuals.size());
        std::vector<Sums<L>> partial(ReductionBlocks(n, MIN_RESIDUALS_PER_BLOCK));
        ParallelBlocks(n, std::max(1, num_threads), [&](int b, int lo, int hi) {
            for (int i = lo; i < hi; i++)
                residual(i, partial[b]);
        }, MIN_RESIDUALS_PER_BLOCK);
        Sums<L> total;
        for (const auto &s: partial)
            total.Add(s);
        return total;
    }

    double WeightSq(const XtalRefineProblem &p, int i) {
        return p.weight_sq.empty() ? 1.0 : p.weight_sq[i];
    }

    bool AllFinite(double cost, const Eigen::VectorXd *g, const Eigen::MatrixXd *H) {
        return std::isfinite(cost) && (!g || (g->allFinite() && H->allFinite()));
    }

    // Seven-block residual (XtalResidualFixedDistance), with every block that depends on parameters
    // alone - detector-angle sines and cosines, the spindle back-rotation of each frame, the reciprocal
    // basis of the cell, the orientation's rotation - worked out once per evaluation.
    bool EvaluateGeneral(const XtalRefineProblem &p, const std::array<int, LANES> &map, int num_threads,
                         const double *x, double &cost, Eigen::VectorXd *g, Eigen::MatrixXd *H) {
        const bool beam_free = map[0] >= 0, rot_free = map[2] >= 0, axis_free = map[4] >= 0;
        const bool p0_free = map[OBS] >= 0, len_free = map[OBS + 3] >= 0, ang_free = map[OBS + 6] >= 0;

        const DO beam[2] = {Seed<DO>(x[OFF_BEAM], 0, beam_free),
                            Seed<DO>(x[OFF_BEAM + 1], 1, beam_free)};
        const DO rot1 = Seed<DO>(x[OFF_ROT], 2, rot_free);
        const DO rot2 = Seed<DO>(x[OFF_ROT + 1], 3, rot_free);
        const DO c1 = cos(rot1), s1 = sin(rot1), c2 = cos(rot2), s2 = sin(rot2);

        DO axis[3] = {x[OFF_AXIS], x[OFF_AXIS + 1], x[OFF_AXIS + 2]};
        if (axis_free) {
            double J[3][2];
            SpherePlusJacobian3(x + OFF_AXIS, J);
            for (int k = 0; k < 3; k++) {
                axis[k].v[4] = J[k][0];
                axis[k].v[5] = J[k][1];
            }
        }
        std::vector<AngleAxisRotator<DO>> rot_back;
        rot_back.reserve(p.frame_angle_rad.size());
        for (const double angle: p.frame_angle_rad) {
            const DO aa_back[3] = {angle * axis[0], angle * axis[1], angle * axis[2]};
            rot_back.emplace_back(aa_back);
        }

        DP p0[3], len[3], ang[3];
        for (int k = 0; k < 3; k++) {
            p0[k] = Seed<DP>(x[OFF_P0 + k], k, p0_free);
            len[k] = Seed<DP>(x[OFF_LEN + k], 3 + k, len_free);
            ang[k] = Seed<DP>(x[OFF_ANG + k], 6 + k, ang_free);
        }
        Eigen::Matrix<DP, 3, 1> bxc, cxa, axb;
        DP invV;
        XtalResidual::ReciprocalBasis(len, ang, p.crystal_system, bxc, cxa, axb, invV);
        const AngleAxisRotator<DP> rot_p0(p0);

        const Sums<LANES> s = SumResiduals<LANES>(p, num_threads, [&](int i, Sums<LANES> &acc) {
            const XtalResidual &res = p.residuals[i];
            DO obs[3];
            res.ObservedRecipCore(beam, p.distance_mm, c1, s1, c2, s2, rot_back[p.frame[i]], obs);
            DP unrot[3], pred[3];
            res.CombineRecipUnrot(bxc, cxa, axb, invV, unrot);
            rot_p0.Rotate(unrot, pred);
            const double w2 = WeightSq(p, i);
            for (int k = 0; k < 3; k++) {
                double J[LANES];
                for (int l = 0; l < OBS; l++)
                    J[l] = obs[k].v[l];
                for (int l = 0; l < PRED; l++)
                    J[OBS + l] = -pred[k].v[l];
                acc.Add(obs[k].a - pred[k].a, J, w2);
            }
        });

        cost = s.cost;
        if (g)
            ToTangent<LANES>(s, map, *g, *H);
        AddPriors(p, x, cost, g, H, map[0], map[2]);
        return AllFinite(cost, g, H);
    }

    // The reduced residual (XtalResidualBeamOrientation): detector, spindle and cell held, so the
    // back-rotation of each frame and the unrotated prediction of each residual are constants.
    bool EvaluateBeamOrientation(const XtalRefineProblem &p, const std::array<int, R_LANES> &map,
                                 const std::vector<AngleAxisRotator<double>> &rot_back,
                                 const std::vector<std::array<double, 3>> &unrot, double c1, double s1,
                                 double c2, double s2, int num_threads,
                                 const double *x, double &cost, Eigen::VectorXd *g, Eigen::MatrixXd *H) {
        using DB = Dual<R_OBS>;
        using DR = Dual<R_PRED>;
        const bool beam_free = map[0] >= 0;
        const DB beam[2] = {beam_free ? DB::Variable(x[OFF_BEAM], 0) : DB(x[OFF_BEAM]),
                            beam_free ? DB::Variable(x[OFF_BEAM + 1], 1) : DB(x[OFF_BEAM + 1])};
        const DR p0[3] = {DR::Variable(x[OFF_P0], 0), DR::Variable(x[OFF_P0 + 1], 1),
                          DR::Variable(x[OFF_P0 + 2], 2)};
        const AngleAxisRotator<DR> rot_p0(p0);

        const Sums<R_LANES> s = SumResiduals<R_LANES>(p, num_threads, [&](int i, Sums<R_LANES> &acc) {
            const XtalResidual &res = p.residuals[i];
            DB obs[3];
            res.ObservedRecipCore(beam, p.distance_mm, c1, s1, c2, s2, rot_back[p.frame[i]], obs);
            DR pred[3];
            rot_p0.Rotate(unrot[i].data(), pred);
            const double w2 = WeightSq(p, i);
            for (int k = 0; k < 3; k++) {
                double J[R_LANES];
                for (int l = 0; l < R_OBS; l++)
                    J[l] = obs[k].v[l];
                for (int l = 0; l < R_PRED; l++)
                    J[R_OBS + l] = -pred[k].v[l];
                acc.Add(obs[k].a - pred[k].a, J, w2);
            }
        });

        cost = s.cost;
        if (g)
            ToTangent<R_LANES>(s, map, *g, *H);
        AddPriors(p, x, cost, g, H, map[0], -1);
        return AllFinite(cost, g, H);
    }
}

LMSummary SolveXtalRefine(XtalRefineProblem &p, int num_threads) {
    const std::vector<LMBlock> blocks = MakeBlocks(p);

    std::vector<double> x(N_AMBIENT);
    const auto put = [&](int off, const double *v, int n) { for (int i = 0; i < n; i++) x[off + i] = v[i]; };
    put(OFF_BEAM, p.beam, 2);
    put(OFF_ROT, p.detector_rot, 2);
    put(OFF_AXIS, p.rot_vec, 3);
    put(OFF_P0, p.latt_vec0, 3);
    put(OFF_LEN, p.latt_vec1, 3);
    put(OFF_ANG, p.latt_vec2, 3);

    LMSummary summary;
    if (p.beam_and_orientation_only) {
        const std::array<int, R_LANES> lane_block = {0, 0, 3, 3, 3};
        const std::array<int, R_LANES> lane_index = {0, 1, 0, 1, 2};
        const auto map = LaneToTangent<R_LANES>(blocks, lane_block, lane_index);

        std::vector<AngleAxisRotator<double>> rot_back;
        rot_back.reserve(p.frame_angle_rad.size());
        for (const double angle: p.frame_angle_rad)
            rot_back.push_back(XtalFrameConstants::BackRotator(angle, p.rot_vec));
        Eigen::Matrix<double, 3, 1> bxc, cxa, axb;
        double invV;
        XtalResidual::ReciprocalBasis(p.latt_vec1, p.latt_vec2, p.crystal_system, bxc, cxa, axb, invV);
        std::vector<std::array<double, 3>> unrot(p.residuals.size());
        for (size_t i = 0; i < p.residuals.size(); i++)
            p.residuals[i].CombineRecipUnrot(bxc, cxa, axb, invV, unrot[i].data());
        const double c1 = std::cos(p.detector_rot[0]), s1 = std::sin(p.detector_rot[0]);
        const double c2 = std::cos(p.detector_rot[1]), s2 = std::sin(p.detector_rot[1]);

        summary = SolveLM(x, blocks, p.options, [&](const double *at, double &cost, Eigen::VectorXd *g,
                                                    Eigen::MatrixXd *H) {
            return EvaluateBeamOrientation(p, map, rot_back, unrot, c1, s1, c2, s2, num_threads, at, cost, g, H);
        });
    } else {
        const std::array<int, LANES> lane_block = {0, 0, 1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5};
        const std::array<int, LANES> lane_index = {0, 1, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2, 0, 1, 2};
        const auto map = LaneToTangent<LANES>(blocks, lane_block, lane_index);
        summary = SolveLM(x, blocks, p.options, [&](const double *at, double &cost, Eigen::VectorXd *g,
                                                    Eigen::MatrixXd *H) {
            return EvaluateGeneral(p, map, num_threads, at, cost, g, H);
        });
    }

    if (summary.IsSolutionUsable()) {
        const auto get = [&](int off, double *v, int n) { for (int i = 0; i < n; i++) v[i] = x[off + i]; };
        get(OFF_BEAM, p.beam, 2);
        get(OFF_ROT, p.detector_rot, 2);
        get(OFF_AXIS, p.rot_vec, 3);
        get(OFF_P0, p.latt_vec0, 3);
        get(OFF_LEN, p.latt_vec1, 3);
        get(OFF_ANG, p.latt_vec2, 3);
    }
    return summary;
}