///
/// Functions and classes for coordinates conversion and lattice operations
///
/// @author: Yunhai Li (yhli)
///

#ifndef TBPLAS_BASE_LATTICE_H
#define TBPLAS_BASE_LATTICE_H

#include <eigen3/Eigen/Dense>

namespace tbplas::base {
/// @brief Generate lattice vectors from given lattice parameters
/// @param[in] a lattice constant a in ANY unit
/// @param[in] b lattice constant b in ANY unit
/// @param[in] c lattice constant c in ANY unit
/// @param[in] alpha angle between a2 and a3 in DEGREE
/// @param[in] beta angle between a3 and a1 in DEGREE
/// @param[in] gamma between a1 and a2 in DEGREE
/// @return Cartesian coordinates of lattice vectors in the same unit as a/b/c,
/// with each ROW being a vector.
Eigen::Matrix3d gen_lattice_vectors(double a, double b, double c,
    double alpha, double beta, double gamma);

/// @brief Calculate Cartesian coordinates of reciprocal lattice vectors
/// @param[in] lat_vec Cartesian coordinates of real lattice in ANY unit
/// with each ROW being a vector
/// @return Cartesian coordinates of reciprocal lattice vectors, with each
/// ROW being a vector. The unit is inverse to the unit of real-space lattice
//  vectors.
Eigen::Matrix3d gen_reciprocal_vectors(const Eigen::Matrix3d& lat_vec);

/// @brief Calculate Cartesian coordinates of reciprocal lattice vectors using
/// another algorithm
/// @param[in] lat_vec Cartesian coordinates of real lattice in ANY unit
/// with each ROW being a vector
/// @return Cartesian coordinates of reciprocal lattice vectors, with each
/// ROW being a vector. The unit is inverse to the unit of real-space lattice
//  vectors.
Eigen::Matrix3d gen_reciprocal_vectors2(const Eigen::Matrix3d& lat_vec);

/// @brief Convert fractional coordinates to Cartesian coordinates
/// @param[in] lat_vec Cartesian coordinates of lattice vectors
/// with each ROW being a vector
/// @param[in] frac_coord Fractional coordinates to convert
/// with each COLUMN being a coordinate
/// @param[in] origin Cartesian coordinate of lattice origin
/// @return Cartesian coordinates after conversion with each COLUMN being
/// a coordinate in the same unit as lat_vec
Eigen::Matrix3Xd frac2cart(const Eigen::Matrix3d& lat_vec,
    const Eigen::Matrix3Xd& frac_coord,
    const Eigen::Vector3d& origin = Eigen::VectorXd::Zero(3));

/// @brief Convert Cartesian coordinates to fractional coordinates
/// @param[in] lat_vec Cartesian coordinates of lattice vectors
/// with each ROW being a vector
/// @param[in] cart_coord Cartesian coordinates to convert
/// with each COLUMN being a coordinate
/// @param[in] origin Cartesian coordinate of lattice origin
/// @return Fractional coordinates after conversion with each COLUMN being
/// a coordinate
Eigen::Matrix3Xd cart2frac(const Eigen::Matrix3d& lat_vec,
    const Eigen::Matrix3Xd& cart_coord,
    const Eigen::Vector3d& origin = Eigen::VectorXd::Zero(3));

/// @brief Calculate the area along given direction
/// @param[in] lat_vec Cartesian coordinates of lattice vectors
/// with each ROW being a vector
/// @param[in] direction direction along which the area is evaluated, should be in ("a", "b", "c")
/// @return the area in squared unit of lattice vectors
double get_lattice_area(const Eigen::Matrix3d& lat_vec, const std::string& direction = "c");

/// @brief Calculate lattice volume
/// @param[in] lat_vec Cartesian coordinates of lattice vectors
/// with each ROW being a vector
/// @return the volume in cubic unit of lattice vectors
double get_lattice_volume(const Eigen::Matrix3d& lat_vec);

/// @brief Rotate Cartesian coordinates according to Euler angles.
/// @param[in] coord Cartesian coordinates to rotate, with each COLUMN being a vector
/// @param[in] angle rotation angle in RADIAN, not degrees
/// @param[in] axis axis around which the rotation is performed
/// @param[in] center Cartesian coordinate of rotation center
/// @return rotated Cartesian coordinates
/// @note  the right-hand rule is used, i.e., a positive value means the angle is counter-clockwise.
/// Reference: https://en.wikipedia.org/wiki/Rotation_matrix#Rotation_matrix_from_axis_and_angle
Eigen::Matrix3Xd rotate_coord(const Eigen::Matrix3Xd& coord, const double& angle = 0.0,
    const std::string& axis = "z", const Eigen::Vector3d& center = Eigen::Vector3d::Zero());

/**
 * @brief Convert derivatives of f w.r.t ka/kb/kc to kx/ky/kz.
 *
 * @tparam scalar_t datatype of f
 * @param[in] recip_vec Cartesian coordinates of reciprocal lattice vectors,
 *   with each ROW being a vector
 * @param[in] df_dka partial derivative of f w.r.t ka
 * @param[in] df_dkb partial derivative of f w.r.t kb
 * @param[in] df_dkc partial derivative of f w.r.t kc
 * @param[out] df_dkx partial derivative of f w.r.t kx
 * @param[out] df_dky partial derivative of f w.r.t ky
 * @param[out] df_dkz partial derivative of f w.r.t kz
 */
template <typename scalar_t>
void df_dk_frac2cart(
    const Eigen::Matrix3d& recip_vec,
    const Eigen::MatrixBase<scalar_t>& df_dka,
    const Eigen::MatrixBase<scalar_t>& df_dkb,
    const Eigen::MatrixBase<scalar_t>& df_dkc,
    Eigen::MatrixBase<scalar_t>& df_dkx,
    Eigen::MatrixBase<scalar_t>& df_dky,
    Eigen::MatrixBase<scalar_t>& df_dkz)
{
    // We use the chain rule:
    //   df/dk? = df/dka * dka/dk? + df/dkb * dkb/dk? + df/dkc * dkc/dk?
    // with ? being x, y, and z
    //
    // ka/kb/kc are related to kx/ky/kz via:
    //   ka = t00 * kx + t10 * ky + t20 * kz
    //   kb = t01 * kx + t11 * ky + t21 * kz
    //   kc = t02 * kx + t12 * ky + t22 * kz
    // with t being conv_mat of cart2frac
    Eigen::Matrix3d conv_mat = recip_vec.inverse();
    df_dkx = df_dka * conv_mat(0, 0) + df_dkb * conv_mat(0, 1) + df_dkc * conv_mat(0, 2);
    df_dky = df_dka * conv_mat(1, 0) + df_dkb * conv_mat(1, 1) + df_dkc * conv_mat(1, 2);
    df_dkz = df_dka * conv_mat(2, 0) + df_dkb * conv_mat(2, 1) + df_dkc * conv_mat(2, 2);
}

/**
 * @brief Convert derivatives of f w.r.t kx/ky/kz to ka/kb/kc.
 *
 * @tparam scalar_t datatype of f
 * @param[in] recip_vec Cartesian coordinates of reciprocal lattice vectors,
 *   with each ROW being a vector
 * @param[in] df_dkx partial derivative of f w.r.t kx
 * @param[in] df_dky partial derivative of f w.r.t ky
 * @param[in] df_dkz partial derivative of f w.r.t kz
 * @param[out] df_dka partial derivative of f w.r.t ka
 * @param[out] df_dkb partial derivative of f w.r.t kb
 * @param[out] df_dkc partial derivative of f w.r.t kc
 */
template <typename scalar_t>
void df_dk_cart2frac(
    const Eigen::Matrix3d& recip_vec,
    const Eigen::MatrixBase<scalar_t>& df_dkx,
    const Eigen::MatrixBase<scalar_t>& df_dky,
    const Eigen::MatrixBase<scalar_t>& df_dkz,
    Eigen::MatrixBase<scalar_t>& df_dka,
    Eigen::MatrixBase<scalar_t>& df_dkb,
    Eigen::MatrixBase<scalar_t>& df_dkc)
{
    // Similar as df_dk_frac2cart, but conv_mat being frac2cart
    const Eigen::Matrix3d& conv_mat = recip_vec;
    df_dka = df_dkx * conv_mat(0, 0) + df_dky * conv_mat(0, 1) + df_dkz * conv_mat(0, 2);
    df_dkb = df_dkx * conv_mat(1, 0) + df_dky * conv_mat(1, 1) + df_dkz * conv_mat(1, 2);
    df_dkc = df_dkx * conv_mat(2, 0) + df_dky * conv_mat(2, 1) + df_dkz * conv_mat(2, 2);
}
} // namespace tbplas::base

#endif
