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

#include "../../common/JFJochMath.h"
#include "LatticeReduction.h"

#include <Eigen/Dense>

void LatticeToRodriguesAndLengths_GS(const CrystalLattice &latt,
                                            double rod[3],
                                            double lengths[3]) {
    // Load lattice columns
    const Coord a = latt.Vec0();
    const Coord b = latt.Vec1();
    const Coord c = latt.Vec2();

    Eigen::Vector3d A(a[0], a[1], a[2]);
    Eigen::Vector3d B(b[0], b[1], b[2]);
    Eigen::Vector3d C(c[0], c[1], c[2]);

    // Lengths = column norms (orthorhombic assumption)
    lengths[0] = A.norm();
    lengths[1] = B.norm();
    lengths[2] = C.norm();

    auto safe_unit = [](const Eigen::Vector3d &v, double eps = 1e-15) -> Eigen::Vector3d {
        double n = v.norm();
        return (n > eps) ? (v / n) : Eigen::Vector3d(1.0, 0.0, 0.0);
    };

    // Gram–Schmidt with original order: x from A, y from B orthogonalized vs x
    Eigen::Vector3d e1 = safe_unit(A);
    Eigen::Vector3d y = B - (e1.dot(B)) * e1;
    Eigen::Vector3d e2 = safe_unit(y);

    // z from cross to ensure right-handed basis
    Eigen::Vector3d e3 = e1.cross(e2);
    double n3 = e3.norm();
    if (n3 < 1e-15) {
        // Degenerate case: B nearly collinear with A → use C instead
        y = C - (e1.dot(C)) * e1;
        e2 = safe_unit(y);
        e3 = e1.cross(e2);
        n3 = e3.norm();
        if (n3 < 1e-15) {
            // Still degenerate: pick any perpendicular to e1
            e2 = safe_unit((std::abs(e1.x()) < 0.9)
                               ? Eigen::Vector3d::UnitX().cross(e1)
                               : Eigen::Vector3d::UnitY().cross(e1));
            e3 = e1.cross(e2);
        }
    } else {
        e3 /= n3;
    }

    Eigen::Matrix3d R;
    R.col(0) = e1;
    R.col(1) = e2;
    R.col(2) = e3;

    // Convert rotation to Rodrigues (axis * angle)
    Eigen::AngleAxisd aa(R);
    Eigen::Vector3d r = aa.angle() * aa.axis();
    rod[0] = r.x();
    rod[1] = r.y();
    rod[2] = r.z();
}

void LatticeToRodriguesAndLengths_Hex(const CrystalLattice &latt, double rod[3], double ac[3]) {
    const Coord a = latt.Vec0();
    const Coord b = latt.Vec1();
    const Coord c = latt.Vec2();

    Eigen::Vector3d A(a[0], a[1], a[2]);
    Eigen::Vector3d B(b[0], b[1], b[2]);
    Eigen::Vector3d C(c[0], c[1], c[2]);

    const double a_len = A.norm();
    const double b_len = B.norm();
    const double c_len = C.norm();

    ac[0] = (a_len + b_len) / 2.0;
    ac[1] = (a_len + b_len) / 2.0;
    ac[2] = c_len;

    Eigen::Vector3d e1;
    Eigen::Vector3d e3;

    if (a_len > 0.0)
        e1 = A / a_len;
    else
        e1 = Eigen::Vector3d::UnitX();

    if (c_len > 0.0)
        e3 = C / c_len;
    else
        e3 = Eigen::Vector3d::UnitZ();

    Eigen::Vector3d e2 = e3.cross(e1);
    if (e2.norm() < 1e-15) {
        e2 = (std::abs(e1.x()) < 0.9)
                 ? Eigen::Vector3d::UnitX().cross(e1)
                 : Eigen::Vector3d::UnitY().cross(e1);
    }
    e2.normalize();
    e3 = e1.cross(e2).normalized();

    Eigen::Matrix3d R;
    R.col(0) = e1;
    R.col(1) = e2;
    R.col(2) = e3;

    Eigen::AngleAxisd aa(R);
    Eigen::Vector3d r = aa.angle() * aa.axis();
    rod[0] = r.x();
    rod[1] = r.y();
    rod[2] = r.z();
}

// Extract rotation (Rodrigues), lengths (a,b,c) and beta (rad) for monoclinic (unique axis b).
// Frame choice: e2 aligned with b; e1 from a projected orthogonal to e2; e3 = e1 x e2.
void LatticeToRodriguesLengthsBeta_Mono(const CrystalLattice &latt,
                                        double rod[3],
                                        double lengths[3],
                                        double &beta_rad) {
    const Coord a = latt.Vec0();
    const Coord b = latt.Vec1();
    const Coord c = latt.Vec2();

    const Eigen::Vector3d A(a[0], a[1], a[2]);
    const Eigen::Vector3d Bv(b[0], b[1], b[2]);
    const Eigen::Vector3d C(c[0], c[1], c[2]);

    // Unit cell lengths
    const double a_len = A.norm();
    const double b_len = Bv.norm();
    const double c_len = C.norm();

    lengths[0] = a_len;
    lengths[1] = b_len;
    lengths[2] = c_len;

    // Monoclinic beta = angle(a, c)
    double cos_beta = 0.0;
    if (a_len > 1e-15 && c_len > 1e-15) {
        cos_beta = A.dot(C) / (a_len * c_len);
        cos_beta = std::clamp(cos_beta, -1.0, 1.0);
    }

    beta_rad = std::acos(cos_beta);

    // Protect against singular construction
    const double sin_beta = std::max(std::abs(std::sin(beta_rad)), 1e-12);

    // Canonical monoclinic basis:
    //
    // B =
    // [ 1   0   cos(beta) ]
    // [ 0   1        0     ]
    // [ 0   0   sin(beta) ]
    //
    Eigen::Matrix3d Bmono = Eigen::Matrix3d::Zero();
    Bmono(0,0) = 1.0;
    Bmono(1,1) = 1.0;
    Bmono(0,2) = std::cos(beta_rad);
    Bmono(2,2) = sin_beta;

    // Scale by lengths
    Eigen::DiagonalMatrix<double,3> D(a_len, b_len, c_len);

    // Ideal body-frame lattice
    const Eigen::Matrix3d M = Bmono * D;

    // Observed lattice
    Eigen::Matrix3d L;
    L.col(0) = A;
    L.col(1) = Bv;
    L.col(2) = C;

    // Estimate rotation:
    // R ≈ L * M^{-1}
    Eigen::Matrix3d R_est = L * M.inverse();

    // Project to nearest proper rotation matrix
    Eigen::JacobiSVD<Eigen::Matrix3d> svd(R_est, Eigen::ComputeFullU | Eigen::ComputeFullV);

    Eigen::Matrix3d R = svd.matrixU() * svd.matrixV().transpose();

    // Enforce det(R)=+1
    if (R.determinant() < 0.0) {
        Eigen::Matrix3d U = svd.matrixU();
        U.col(2) *= -1.0;
        R = U * svd.matrixV().transpose();
    }

    // Rodrigues vector
    Eigen::AngleAxisd aa(R);
    const Eigen::Vector3d r = aa.angle() * aa.axis();

    rod[0] = r.x();
    rod[1] = r.y();
    rod[2] = r.z();
}

static inline Eigen::Matrix3d B_from_angles(double alpha_rad, double beta_rad, double gamma_rad) {
    const double ca = std::cos(alpha_rad);
    const double cb = std::cos(beta_rad);
    const double cg = std::cos(gamma_rad);
    const double sg = std::sin(gamma_rad);

    Eigen::Matrix3d B = Eigen::Matrix3d::Identity();

    // a along x, b in x-y, c general
    B(0, 0) = 1.0;  B(1, 0) = 0.0;  B(2, 0) = 0.0;
    B(0, 1) = cg;   B(1, 1) = sg;   B(2, 1) = 0.0;

    // c vector components (standard crystallography construction)
    const double cx = cb;
    const double cy = (ca - cb * cg) / sg;
    const double cz = std::sqrt(std::max(0.0, 1.0 - cx * cx - cy * cy));

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

    return B;
}

CrystalLattice AngleAxisAndCellToLattice(const double rod[3],
                                         const double lengths[3],
                                         double alpha_rad,
                                         double beta_rad,
                                         double gamma_rad) {
    const Eigen::Vector3d r(rod[0], rod[1], rod[2]);
    const double angle = r.norm();

    Eigen::Matrix3d R = Eigen::Matrix3d::Identity();
    if (angle > 0.0)
        R = Eigen::AngleAxisd(angle, r / angle).toRotationMatrix();

    const Eigen::DiagonalMatrix<double, 3> D(lengths[0], lengths[1], lengths[2]);
    const Eigen::Matrix3d B = B_from_angles(alpha_rad, beta_rad, gamma_rad);

    // IMPORTANT convention: L = R * B * D (scale columns by lengths)
    const Eigen::Matrix3d latt = R * B * D;

    return CrystalLattice(Coord(latt(0, 0), latt(1, 0), latt(2, 0)),
                          Coord(latt(0, 1), latt(1, 1), latt(2, 1)),
                          Coord(latt(0, 2), latt(1, 2), latt(2, 2)));
}


CrystalLattice AngleAxisAndLengthsToLattice(const double rod[3], const double lengths[3], bool hex) {
    return AngleAxisAndCellToLattice(rod, lengths,
                                     /*alpha=*/PI / 2.0,
                                     /*beta =*/PI / 2.0,
                                     /*gamma=*/hex ? 2.0 * PI / 3.0 : PI / 2.0);
}