#ifndef TBPLAS_ATOM_H
#define TBPLAS_ATOM_H

#include <cmath>
#include <cstdint>

namespace atom {

/// @brief Evaluate the factorial of n.
inline double factorial(int n)
{
    double y = 1.0;
    for (int i = 1; i <= n; ++i) {
        y *= i;
    }
    return y;
}

/// @brief Evaluate the renormalization factor of Rnl(r).
inline double rnl_factor(int z, int n, int l)
{
    double a = 2.0 / (pow(n, 2) * factorial(2 * l + 1));
    double b = sqrt(factorial(n + l) / factorial(n - l - 1));
    double c = sqrt(pow(z, 3));
    return a * b * c;
}

/// @brief Evaluate the Kummer function in Rnl(r).
inline double kummer(int alpha, int gamma, double ita)
{
    double sum = 1.0;
    double ck = 1.0;
    for (int k = 1; k <= -alpha; ++k) {
        ck = ck * (alpha + k - 1) / (gamma + k - 1) / k;
        sum = sum + ck * pow(ita, k);
    }
    return sum;
}

/// @brief Evaluate the un-normalized Rnl(r).
inline double radial(int z, int n, int l, double r)
{
    double ita = 2.0 * z / n * r;
    double a = exp(-0.5 * ita);
    double b = pow(ita, l);
    double c = kummer(-n + l + 1, 2 * l + 2, ita);
    return a * b * c;
}

/// @brief Evaluate the coefficient term in ylm(x, y, z).
inline double coeff(int a, int b, int c, int d)
{
    constexpr double pi = 3.14159265358979323846;
    return static_cast<double>(a) / b * sqrt(c / (d * pi));
}

/// @brief Evaluate the renormalization factor of ylm(x, y, z).
inline double ylm_factor(int l, int m)
{
    double factor = 0.0;
    if (l == 0) {
        factor = coeff(1, 2, 1, 1);
    }
    else if (l == 1) {
        switch (m) {
        case -1:
            factor = coeff(1, 1, 3, 4);
            break; // py
        case 0:
            factor = coeff(1, 1, 3, 4);
            break; // pz
        case 1:
            factor = coeff(1, 1, 3, 4);
            break; // px
        default:
            factor = 0.0;
            break;
        }
    } else if (l == 2) {
        switch (m) {
        case -2:
            factor = coeff(1, 2, 15, 1);
            break; // dxy
        case -1:
            factor = coeff(1, 2, 15, 1);
            break; // dyz
        case 0:
            factor = coeff(1, 4, 5, 1);
            break; // dz^2
        case 1:
            factor = coeff(1, 2, 15, 1);
            break; // dxz
        case 2:
            factor = coeff(1, 4, 15, 1);
            break; // d(x^2-y^2)
        default:
            factor = 0.0;
            break;
        }
    } else if (l == 3) {
        switch (m) {
        case -3:
            factor = coeff(1, 4, 35, 2);
            break; // fy(3x^2-y^2)
        case -2:
            factor = coeff(1, 2, 105, 1);
            break; // fxyz
        case -1:
            factor = coeff(1, 4, 21, 2);
            break; // fyz^2
        case 0:
            factor = coeff(1, 4, 7, 1);
            break; // fz^3
        case 1:
            factor = coeff(1, 4, 21, 2);
            break; // fxz^2
        case 2:
            factor = coeff(1, 4, 105, 1);
            break; // fz(x^2-y^2)
        case 3:
            factor = coeff(1, 4, 35, 2);
            break; // fx(x^2-3y^2)
        default:
            factor = 0.0;
            break;
        }
    } else {
        factor = 0.0;
    }
    return factor;
}

/// @brief Evaluate the un-normalized real spherical harmonics ylm(x, y, z).
/// @note Reference:
/// https://en.wikipedia.org/wiki/Table_of_spherical_harmonics#Real_spherical_harmonics
inline double ylm(int l, int m, double x, double y, double z, double r)
{
    double ylm = 1.0;
    if (l == 0) {
        ylm = 1.0;
    } else if (l == 1) {
        double denorm = r;
        switch (m) {
        case -1:
            ylm = 1.0 / denorm * y;
            break; // py
        case 0:
            ylm = 1.0 / denorm * z;
            break; // pz
        case 1:
            ylm = 1.0 / denorm * x;
            break; // px
        default:
            ylm = 0.0;
            break;
        }
    } else if (l == 2) {
        double denorm = pow(r, 2.0);
        switch (m) {
        case -2:
            ylm = 1.0 / denorm * (x * y);
            break; // dxy
        case -1:
            ylm = 1.0 / denorm * (y * z);
            break; // dyz
        case 0:
            ylm = 1.0 / denorm * (2 * z * z - x * x - y * y);
            break; // dz^2
        case 1:
            ylm = 1.0 / denorm * (z * x);
            break; // dxz
        case 2:
            ylm = 1.0 / denorm * (x * x - y * y);
            break; // d(x^2-y^2)
        default:
            ylm = 0.0;
            break;
        }
    } else if (l == 3) {
        double denorm = pow(r, 3.0);
        switch (m) {
        case -3:
            ylm = 1.0 / denorm * y * (3 * x * x - y * y);
            break; // fy(3x^2-y^2)
        case -2:
            ylm = 1.0 / denorm * (x * y * z);
            break; // fxyz
        case -1:
            ylm = 1.0 / denorm * y * (4 * z * z - x * x - y * y);
            break; // fyz^2
        case 0:
            ylm = 1.0 / denorm * z * (2 * z * z - 3 * x * x - 3 * y * y);
            break; // fz^3
        case 1:
            ylm = 1.0 / denorm * x * (4 * z * z - x * x - y * y);
            break; // fxz^2
        case 2:
            ylm = 1.0 / denorm * z * (x * x - y * y);
            break; // fz(x^2-y^2)
        case 3:
            ylm = 1.0 / denorm * x * (x * x - 3 * y * y);
            break; // fx(x^2-3y^2)
        default:
            ylm = 0.0;
            break;
        }
    } else {
        ylm = 0.0;
    }
    return ylm;
}

inline void set_cube_cxx(
    const double* c0,
    const int32_t* qn,
    const double* cube_origin,
    const int32_t* num_grid,
    const double& res,
    double* cube)
{
    double x_min = cube_origin[0];
    double y_min = cube_origin[1];
    double z_min = cube_origin[2];
    int32_t grid_x = num_grid[0];
    int32_t grid_y = num_grid[1];
    int32_t grid_z = num_grid[2];
    int32_t grid_yz = grid_y * grid_z;
    int32_t z = qn[0];
    int32_t n = qn[1];
    int32_t l = qn[2];
    int32_t m = qn[3];

    double factor = rnl_factor(z, n, l) * ylm_factor(l, m);
    int32_t i = 0;
    double dx = 0.0;
    double dx2 = 0.0;
    int64_t offset_x = 0;
    int32_t j = 0;
    double dy = 0.0;
    double dx2y2 = 0.0;
    int64_t offset_y = 0;
    int32_t k = 0;
    double dz = 0.0;
    double r = 0.0;
    double grid_val = 0.0;

#pragma omp parallel for private(dx, dx2, offset_x, j, dy, dx2y2, offset_y, k, dz, r, grid_val)
    for (i = 0; i < grid_x; ++i) {
        dx = x_min + res * i - c0[0];
        dx2 = dx * dx;
        offset_x = i * grid_yz;

        for (j = 0; j < grid_y; ++j) {
            dy = y_min + res * j - c0[1];
            dx2y2 = dx2 + dy * dy;
            offset_y = offset_x + j * grid_z;

            for (k = 0; k < grid_z; ++k) {
                dz = z_min + res * k - c0[2];
                r = sqrt(dx2y2 + dz * dz);
                grid_val = 0.0;
                if (r > 1.0e-7) {
                    grid_val = radial(z, n, l, r) * ylm(l, m, dx, dy, dz, r) * factor;
                }
                cube[offset_y + k] = grid_val;
            }
        }
    }
}
} // namespace atom
#endif
