# cython: language_level=3
# cython: warn.undeclared=True
# cython: warn.unreachable=True
# cython: warn.maybe_uninitialized=True
# cython: warn.unused=True
# cython: warn.unused_arg=True
# cython: warn.unused_result=True
# cython: warn.multiple_declarators=True
# distutils: language = c++

import cython
from libc.stdint cimport int32_t
import numpy as np


cdef extern from "atom.h" namespace "atom":
    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)


@cython.boundscheck(False)
@cython.wraparound(False)
def set_cube(
    double [::1] c0,
    int32_t[::1] qn,
    double[::1] cube_origin,
    int32_t [::1] num_grid,
    double res,
    double [:,:,::1] cube):
    """
    Write atomic wavefunction to cube.

    Parameters
    ----------
    c0: (3,) float64 array
        Cartesian coordinate of the nuclei in bohr
    qn: (4,) int32 array
        z, n, l, m of the atom
    cube_origin: (3,) float64 array
        Cartesian coordinate of origin of cube in bohr
    num_grid: (3,) int32 array
        number of grid points along x, y, and z directions
    res: float64
        resolution of cube
    cube: (nx, ny, nz) float64 array
        cube to store the atomic wave function
    """
    set_cube_cxx(
        &c0[0],
        &qn[0],
        &cube_origin[0],
        &num_grid[0],
        res,
        &cube[0][0][0])
