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

#include "PixelRefine.h"

#include <ceres/ceres.h>
#include <ceres/rotation.h>


struct PixelResidual {
    // Assume that Itrue and Ibkg are already corrected with solid angle and polarization correction
    PixelResidual(double x, double y,
                  double Itrue, double Iobs,
                  double Ibkg, double Ibkg_sigma,
                  double lambda,
                  double pixel_size,
                  double angle_rad,
                  double exp_h, double exp_k,
                  double exp_l,
                  gemmi::CrystalSystem symmetry)
        : Itrue(Itrue), Iobs(Iobs),
          Ibkg(Ibkg), Ibkg_sigma(Ibkg_sigma), obs_x(x),
          obs_y(y),
          inv_lambda(1.0 / lambda),
          pixel_size(pixel_size),
          exp_h(exp_h),
          exp_k(exp_k),
          exp_l(exp_l),
          angle_rad(angle_rad),
          symmetry(symmetry) {
        if (std::fabs(lambda) < 1e-6)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Lambda cannot be close to zero");
    }

    template<typename T>
    bool operator()(const T *const beam,
                    const T *const distance_mm,
                    const T *const detector_rot,
                    const T *const rotation_axis,
                    const T *const p0,
                    const T *const p1,
                    const T *const p2,
                    const T *const scale_factor,
                    const T *const B,
                    const T *const R,
                    T *residual) const {
        // PyFAI convention (left-handed for rot1/rot2):
        // poni_rot = Rz(-rot3) * Rx(-rot2) * Ry(+rot1)
        // detector_rot[0] = rot1, detector_rot[1] = rot2 (rot3 = 0 assumed)

        const T rot1 = detector_rot[0];
        const T rot2 = detector_rot[1];

        // Ry(+rot1): rotation around Y-axis
        const T c1 = ceres::cos(rot1);
        const T s1 = ceres::sin(rot1);

        // Rx(-rot2): rotation around X-axis with inverted sign (PyFAI left-handed)
        const T c2 = ceres::cos(rot2);
        const T s2 = ceres::sin(rot2);

        // Detector coordinates in mm
        const T det_x = (T(obs_x) - beam[0]) * T(pixel_size);
        const T det_y = (T(obs_y) - beam[1]) * T(pixel_size);
        const T det_z = T(distance_mm[0]);

        // Apply Ry(rot1) first: rotate around Y
        const T t1_x = c1 * det_x + s1 * det_z;
        const T t1_y = det_y;
        const T t1_z = -s1 * det_x + c1 * det_z;

        // Then apply Rx(-rot2): rotate around X
        const T x = t1_x;
        const T y = c2 * t1_y + s2 * t1_z;
        const T z = -s2 * t1_y + c2 * t1_z;

        // convert to recip space
        const T lab_norm = ceres::sqrt(x * x + y * y + z * z);
        const T inv_norm = T(1) / lab_norm;

        T recip_raw[3];
        recip_raw[0] = x * inv_norm * T(inv_lambda);
        recip_raw[1] = y * inv_norm * T(inv_lambda);
        recip_raw[2] = (z * inv_norm - T(1.0)) * T(inv_lambda);

        // Apply goniometer "back-to-start" rotation:
        // brings observed reciprocal from image orientation into reference crystal frame
        const T aa_back[3] = {
            T(angle_rad) * rotation_axis[0],
            T(angle_rad) * rotation_axis[1],
            T(angle_rad) * rotation_axis[2]
        };

        T recip_obs[3];
        ceres::AngleAxisRotatePoint(aa_back, recip_raw, recip_obs);

        const Eigen::Matrix<T, 3, 1> e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]);

        // Build unit cell lengths and B (convention: columns are a, b, c prior to global rotation)
        Eigen::Matrix<T, 3, 1> e_uc_len = Eigen::Matrix<T, 3, 1>::Zero();
        Eigen::Matrix<T, 3, 3> B = Eigen::Matrix<T, 3, 3>::Identity();

        if (symmetry == gemmi::CrystalSystem::Hexagonal) {
            e_uc_len << p1[0], p1[0], p1[2];
            B(0, 1) = T(-0.5); // cos(120)
            B(1, 1) = T(sqrt(3.0) / 2.0); // sin(120)
        } else if (symmetry == gemmi::CrystalSystem::Orthorhombic) {
            e_uc_len << p1[0], p1[1], p1[2];
        } else if (symmetry == gemmi::CrystalSystem::Tetragonal) {
            e_uc_len << p1[0], p1[0], p1[2];
        } else if (symmetry == gemmi::CrystalSystem::Cubic) {
            e_uc_len << p1[0], p1[0], p1[0];
        } else if (symmetry == gemmi::CrystalSystem::Monoclinic) {
            // Unique axis b: alpha = gamma = 90°, beta free (angle between a and c)
            e_uc_len << p1[0], p1[1], p1[2];
            B(0, 2) = ceres::cos(p2[0]);
            B(2, 2) = ceres::sin(p2[0]);
        } else {
            // Triclinic: p1 = (a,b,c), p2 = (alpha, beta, gamma) in radians
            const T ca = ceres::cos(p2[0]);
            const T cb = ceres::cos(p2[1]);
            const T cg = ceres::cos(p2[2]);
            const T sg = ceres::sin(p2[2]);

            e_uc_len << p1[0], p1[1], p1[2];

            B(0, 0) = T(1);
            B(1, 0) = T(0);
            B(2, 0) = T(0);
            B(0, 1) = cg;
            B(1, 1) = sg;
            B(2, 1) = T(0);

            // c vector components:
            const T cx = cb;
            const T cy = (ca - cb * cg) / sg;
            const T v = T(1) - cx * cx - cy * cy;
            const T cz = (v >= T(0)) ? ceres::sqrt(v) : T(0);

            B(0, 2) = cx;
            B(1, 2) = cy;
            B(2, 2) = cz;
        }

        // Build unrotated direct lattice columns: (B * D), then rotate them by p0.
        // This avoids AngleAxisToRotationMatrix + matrix multiplications.
        const T L0 = e_uc_len[0];
        const T L1 = e_uc_len[1];
        const T L2 = e_uc_len[2];

        T col0_unrot[3] = {B(0, 0) * L0, B(1, 0) * L0, B(2, 0) * L0};
        T col1_unrot[3] = {B(0, 1) * L1, B(1, 1) * L1, B(2, 1) * L1};
        T col2_unrot[3] = {B(0, 2) * L2, B(1, 2) * L2, B(2, 2) * L2};

        T col0_rot[3], col1_rot[3], col2_rot[3];
        ceres::AngleAxisRotatePoint(p0, col0_unrot, col0_rot);
        ceres::AngleAxisRotatePoint(p0, col1_unrot, col1_rot);
        ceres::AngleAxisRotatePoint(p0, col2_unrot, col2_rot);

        const Eigen::Matrix<T, 3, 1> A(col0_rot[0], col0_rot[1], col0_rot[2]);
        const Eigen::Matrix<T, 3, 1> Bv(col1_rot[0], col1_rot[1], col1_rot[2]);
        const Eigen::Matrix<T, 3, 1> C(col2_rot[0], col2_rot[1], col2_rot[2]);

        const Eigen::Matrix<T, 3, 1> BxC = Bv.cross(C);
        const Eigen::Matrix<T, 3, 1> CxA = C.cross(A);
        const Eigen::Matrix<T, 3, 1> AxB = A.cross(Bv);

        const T V = A.dot(BxC);
        const T invV = T(1) / V;

        const Eigen::Matrix<T, 3, 1> Astar = BxC * invV;
        const Eigen::Matrix<T, 3, 1> Bstar = CxA * invV;
        const Eigen::Matrix<T, 3, 1> Cstar = AxB * invV;

        const T h = T(exp_h);
        const T k = T(exp_k);
        const T l = T(exp_l);

        const Eigen::Matrix<T, 3, 1> e_pred_recip = Astar * h + Bstar * k + Cstar * l;

        // Ewald sphere centre is at -k_i = (0, 0, -inv_lambda)
        // Radial direction: outward normal at g_hkl
        const Eigen::Matrix<T, 3, 1> S_pred(
            e_pred_recip[0],
            e_pred_recip[1],
            e_pred_recip[2] + T(inv_lambda)   // g_hkl + k_i
        );
        const T S_pred_norm = S_pred.norm();
        if (S_pred_norm < T(1e-10))
            return T(0);

        const Eigen::Matrix<T, 3, 1> n_radial = S_pred / S_pred_norm;

        const Eigen::Matrix<T, 3, 1> delta_q = e_obs_recip - e_pred_recip;
        const T eps_radial        = delta_q.dot(n_radial);
        const Eigen::Matrix<T, 3, 1> dq_tang = delta_q - eps_radial * n_radial;
        const T eps_tangential_sq = dq_tang.squaredNorm();   // guaranteed ≥ 0
        // ─────────────────────────────────────────────────────────────

        const T B_term = ceres::exp(- B[0] * e_pred_recip.squaredNorm() / 4.0);

        // Need to normalize by R[0] and R[1]
        const T partiality = ceres::exp(- eps_radial * eps_radial / (R[0] * R[0]) - eps_tangential_sq / (R[1] * R[1]));

        const T Ipred = partiality * Itrue * scale_factor[0] * B_term - Ibkg;

        // Need to weight by sigma
        // I would like to use sigma based on Ipred and Ibkg_sigma - need to come up with a better approach
        residual[0] = (Ipred - Iobs) / Ibkg_sigma;

        return true;
    }

    const double Itrue, Iobs, Ibkg, Ibkg_sigma;
    const double obs_x, obs_y;
    const double inv_lambda;
    const double pixel_size;
    const double exp_h;
    const double exp_k;
    const double exp_l;
    const double angle_rad;
    gemmi::CrystalSystem symmetry;
};

PixelRefine::PixelRefine(const DiffractionExperiment &experiment,
                         const AzimuthalIntegrationMapping &mapping,
                         const std::vector<MergedReflection> &reference,
                         BraggPrediction &prediction)
    : prediction(prediction),
      mapping(mapping),
      xpixel(experiment.GetXPixelsNum()),
      ypixel(experiment.GetYPixelsNum()),
      experiment(experiment),
      hkl_key_generator(experiment.GetScalingSettings().GetMergeFriedel(),
                        experiment.GetSpaceGroupNumber().value_or(1)) {
    for (const auto &ref: reference)
        reference_data[hkl_key_generator(ref)] = ref.I;
}

template<class T>
void PixelRefine::Run(const T *image,
         const AzimuthalIntegrationProfile &profile,
         PixelRefineData &data) {

    ceres::Problem problem;

    // We predict reflections based on initial geometry and default settings
    // To be tuned later
    const BraggPredictionSettings settings_prediction{
        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
        .max_hkl = 100,
        .centering = data.centering
    };
    prediction.Calc(experiment, data.latt, settings_prediction);

    auto azim_result = profile.GetResult();
    auto azim_std = profile.GetStd();

    // For each reflection we select some area (3-5 pixels around it)
    const int radius = 3;

    for (const auto &refl : prediction.GetReflections()) {
        auto hkl = hkl_key_generator(refl);
        // We only handle reflections that are present in the reference set
        if (!reference_data.contains(hkl))
            continue;
        const double I_true = reference_data[hkl];

        int min_y = std::max<int>(refl.predicted_y - radius, 0);
        int max_y = std::min<int>(refl.predicted_y + radius, ypixel - 1);
        int min_x = std::max<int>(refl.predicted_x - radius, 0);
        int max_x = std::min<int>(refl.predicted_x + radius, xpixel - 1);

        for (int y = min_y; y <= max_y; ++y) {
            for (int x = min_x; x <= max_x; ++x) {
                const size_t npixel = xpixel * y + x;
                int azim_bin = mapping.GetPixelToBin()[npixel];
                // If pixel is not mapped to azimuthal bin
                // or pixel has special value (lowest/highest integer)
                // it should be ignored for the purpose of this try
                // We should check if pixel mask is needed, but for most workflows it is already applied

                if (azim_bin >= mapping.GetAzimuthalBinCount())
                    continue;
                if (image[npixel] == std::numeric_limits<T>::max())
                    continue;
                if (std::is_signed_v<T>() && (image[npixel] == std::numeric_limits<T>::min()))
                    continue;

                // Get per-pixel polarization and solid angle correction for the pixel from the AzimuthalIntegrationMapping
                // Warning! this is missing Lorentz correction, but we don't worry at the moment about it
                // Important is -> this correction is also applied to background, so we must be consistent here
                float correction = mapping.Corrections()[npixel];

                // Get mean pixel value for background in the azimuthal bin + sigma
                float bkg_value = azim_result[azim_bin];
                float bkg_sigma = azim_std[azim_bin];
                float pixel_value = image[npixel];


            }
        }
    }

}