#include "lattice.h"
#include "consts.h"
#include <cassert>
#include <cmath>
#include <iostream>

namespace {
double mixed_prod(const Eigen::Matrix3d& lat_vec)
{
    Eigen::Vector3d a0 = lat_vec.row(0);
    Eigen::Vector3d a1 = lat_vec.row(1);
    Eigen::Vector3d a2 = lat_vec.row(2);
    double product = a0.cross(a1).dot(a2);
    assert(std::abs(product) > tbplas::base::INFINITESIMAL);
    return product;
}
} // namespace

namespace tbplas::base {

Eigen::Matrix3d gen_lattice_vectors(double a, double b, double c,
    double alpha, double beta, double gamma)
{
    Eigen::VectorXd test { { a, b, c, alpha, beta, gamma } };
    for (const auto& d : test) {
        assert(d > tbplas::base::INFINITESIMAL);
    }
    double alpha_rad = alpha / 180 * tbplas::base::PI;
    double beta_rad = beta / 180 * tbplas::base::PI;
    double gamma_rad = gamma / 180 * tbplas::base::PI;
    double ca = cos(alpha_rad), cb = cos(beta_rad), cg = cos(gamma_rad);
    double sg = sin(gamma_rad);
    Eigen::Matrix3d lat_vec = Eigen::Matrix3d::Zero();
    lat_vec.row(0) = Eigen::Vector3d(a, 0, 0);
    lat_vec.row(1) = Eigen::Vector3d(b * cg, b * sg, 0);
    lat_vec(2, 0) = c * cb;
    lat_vec(2, 1) = c * (ca - cb * cg) / sg;
    lat_vec(2, 2) = c * sqrt(1 + 2 * ca * cb * cg - pow(ca, 2) - pow(cb, 2) - pow(cg, 2)) / sg;
    assert(get_lattice_volume(lat_vec) > tbplas::base::INFINITESIMAL);
    return lat_vec;
}

Eigen::Matrix3d gen_reciprocal_vectors(const Eigen::Matrix3d& lat_vec)
{
    ::mixed_prod(lat_vec);
    Eigen::Matrix3d recip_vec = Eigen::Matrix3d::Zero();
    Eigen::Matrix3d product = tbplas::base::TWOPI * Eigen::Matrix3d::Identity(3, 3);
    for (int idx = 0; idx < 3; idx++) {
        recip_vec.row(idx) = lat_vec.colPivHouseholderQr().solve(product.col(idx));
    }
    return recip_vec;
}

Eigen::Matrix3d gen_reciprocal_vectors2(const Eigen::Matrix3d& lat_vec)
{
    Eigen::Matrix3d recip_vec = Eigen::Matrix3d::Zero();
    Eigen::Vector3d a0 = lat_vec.row(0);
    Eigen::Vector3d a1 = lat_vec.row(1);
    Eigen::Vector3d a2 = lat_vec.row(2);
    double product = ::mixed_prod(lat_vec);
    recip_vec.row(0) = tbplas::base::TWOPI / product * a1.cross(a2);
    recip_vec.row(1) = tbplas::base::TWOPI / product * a2.cross(a0);
    recip_vec.row(2) = tbplas::base::TWOPI / product * a0.cross(a1);
    return recip_vec;
}

Eigen::Matrix3Xd frac2cart(const Eigen::Matrix3d& lat_vec,
    const Eigen::Matrix3Xd& frac_coord, const Eigen::Vector3d& origin)
{
#ifndef NDEBUG
    // WARNING: It is likely that the caller may send a 3x3 matrix in row-wise order
    // to this function. However, it will not trigger compilation-time errors, but
    // waste the developers hours debugging. So, we check the number of columns of
    // incoming matrix. DO NOT REMOVE THIS PART OR YOU WILL BE TO BLAME.
    if (frac_coord.cols() == 3) {
        std::cout << "INFO: You have passed a 3x3 matrix to " << __func__ << "\n"
                  << "Please make sure the matrix is in column-wise order.\n";
    }
#endif
    assert(get_lattice_volume(lat_vec) > tbplas::base::INFINITESIMAL);
    long num_col = frac_coord.cols();
    Eigen::Matrix3Xd cart_coord(3, num_col);
    for (int idx = 0; idx < num_col; idx++) {
        cart_coord.col(idx) = (frac_coord.col(idx).transpose() * lat_vec).transpose();
    }
    cart_coord.colwise() += origin;
    return cart_coord;
}

Eigen::Matrix3Xd cart2frac(const Eigen::Matrix3d& lat_vec,
    const Eigen::Matrix3Xd& cart_coord, const Eigen::Vector3d& origin)
{
#ifndef NDEBUG
    // WARNING: It is likely that the caller may send a 3x3 matrix in row-wise order
    // to this function. However, it will not trigger compilation-time errors, but
    // waste the developers hours debugging. So, we check the number of columns of
    // incoming matrix. DO NOT REMOVE THIS PART OR YOU WILL BE TO BLAME.
    if (cart_coord.cols() == 3) {
        std::cout << "INFO: You have passed a 3x3 matrix to " << __func__ << "\n"
                  << "Please make sure the matrix is in column-wise order.\n";
    }
#endif
    assert(get_lattice_volume(lat_vec) > tbplas::base::INFINITESIMAL);
    long num_col = cart_coord.cols();
    Eigen::Matrix3Xd frac_coord(3, num_col);
    Eigen::Matrix3Xd cart_coord_shift = cart_coord.colwise() - origin;
    Eigen::Matrix3d conv_mat = lat_vec.inverse();
    for (int idx = 0; idx < num_col; idx++) {
        frac_coord.col(idx) = (cart_coord_shift.col(idx).transpose() * conv_mat).transpose();
    }
    return frac_coord;
}

double get_lattice_area(const Eigen::Matrix3d& lat_vec, const std::string& direction)
{
    assert((direction == "a") || (direction == "b") || (direction == "c"));
    Eigen::Vector3d a0, a1;
    if (direction == "a") {
        a0 = lat_vec.row(1);
        a1 = lat_vec.row(2);
    } else if (direction == "b") {
        a0 = lat_vec.row(2);
        a1 = lat_vec.row(0);
    } else {
        a0 = lat_vec.row(0);
        a1 = lat_vec.row(1);
    }
    double area = a0.cross(a1).norm();
    assert(area > tbplas::base::INFINITESIMAL);
    return area;
}

double get_lattice_volume(const Eigen::Matrix3d& lat_vec)
{
    return std::abs(::mixed_prod(lat_vec));
}

Eigen::Matrix3Xd rotate_coord(const Eigen::Matrix3Xd& coord, const double& angle,
    const std::string& axis, const Eigen::Vector3d& center)
{
#ifndef NDEBUG
    // WARNING: It is likely that the caller may send a 3x3 matrix in row-wise order
    // to this function. However, it will not trigger compilation-time errors, but
    // waste the developers hours debugging. So, we check the number of columns of
    // incoming matrix. DO NOT REMOVE THIS PART OR YOU WILL BE TO BLAME.
    if (coord.cols() == 3) {
        std::cout << "INFO: You have passed a 3x3 matrix to " << __func__ << "\n"
                  << "Please make sure the matrix is in column-wise order.\n";
    }
#endif
    assert((axis == "x") || (axis == "y") || (axis == "z"));

    // Determine rotation matrix
    Eigen::Vector3d u;
    if (axis == "x") {
        u = Eigen::Vector3d(1, 0, 0);
    } else if (axis == "y") {
        u = Eigen::Vector3d(0, 1, 0);
    } else {
        u = Eigen::Vector3d(0, 0, 1);
    }
    double ux = u(0), uy = u(1), uz = u(2);
    Eigen::Matrix3d u_prod { { 0, -uz, uy }, { uz, 0, -ux }, { -uy, ux, 0 } };
    Eigen::Matrix3d u_tens = u * u.transpose();
    double cos_ang = cos(angle), sin_ang = sin(angle);
    Eigen::Matrix3d rot_mat = cos_ang * Eigen::Matrix3d::Identity() + sin_ang * u_prod + (1 - cos_ang) * u_tens;

    // Rotate coordinates
    // CAUTION: rot_mat is for column-wise vectors only!
    long num_col = coord.cols();
    Eigen::Matrix3Xd coord_rot(3, num_col);
    for (int idx = 0; idx < num_col; idx++) {
        coord_rot.col(idx) = rot_mat * (coord.col(idx) - center) + center;
    }
    return coord_rot;
}
} // namespace tbplas::base
