# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cuest.bindings as ce
import numpy as np
import argparse
from pathlib import Path
import sys

# Inject the directory where the example helper utilities live
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from helpers.cuda_utils import cuda_malloc, cuda_free, cuda_memcpy_dtoh, cuda_memcpy_htod, WorkspaceDescriptor, Workspace
from helpers.utilities import normalize_shell_coefficients
from helpers.parsers import simple_xyz_parser, simple_gbs_parser
from helpers.grid_utils import build_ahlrichs_radial_quadrature, symbol_to_ahlrichs_radius

def run(
    *,
    xyz_filename,
    gbs_filename,
    ):

    # => Parse User Input <= #

    symbols, xyzs, Zs = simple_xyz_parser(
        filename=xyz_filename,
        to_bohr_scale_factor=1.0 / 0.52917720859,
        )

    shellinfo = simple_gbs_parser(
        filename=gbs_filename,
        symbols=symbols,
        )

    num_atoms = len(symbols)
    sizeof_double = np.dtype('double').itemsize

    # => Set Up cuEST <= #

    def cuest_check(
        title,
        return_code
        ):
        " A simple helper function to call cuEST functions and check the return code. "

        if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
            raise RuntimeError(f"{title} failed with code {return_code}")

    # These are user-provided functions that are responsible for allocating
    # host and device workspace arrays.
    persistent_workspace_descriptor = WorkspaceDescriptor()
    temporary_workspace_descriptor = WorkspaceDescriptor()
    # Allow 2Gb of DRAM for XC evaluation routines
    variable_buffer_size = WorkspaceDescriptor(
        host_buffer_size_in_bytes=0,
        device_buffer_size_in_bytes=2000000000,
        )

    cuest_handle_parameters = ce.cuestHandleParameters()

    cuest_check('Create Handle Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            outParameters=cuest_handle_parameters,
            )
        )

    # Here we allow cuEST to use the default stream, cuBLAS, etc., but those
    # parameters should be set in the handle parameters before this call.
    cuest_handle = ce.cuestHandle()
    cuest_check('Create Cuest Handle',
        ce.cuestCreate(
            parameters=cuest_handle_parameters,
            handle=cuest_handle,
            )
        )

    stream_handle = ce.data_cudaStream_t()
    cuest_check('Query Handle Stream',
        ce.cuestParametersQuery(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            parameters=cuest_handle_parameters,
            attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_CUDASTREAM,
            attributeValue=stream_handle,
            )
        )

    cuest_check('Destroy Handle Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            parameters=cuest_handle_parameters,
            )
        )

    # => Build Gaussian Shells From Basis Parser Info <= #

    aoshell_parameters = ce.cuestAOShellParameters()
    cuest_check('Create AO Shell Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
            outParameters=aoshell_parameters,
            )
        )

    shells = []
    n_shells_per_atom = []
    shell_count = 1
    for atom_shells in shellinfo:
        n_shells_per_atom.append(len(atom_shells))
        for shell in atom_shells:
            L = shell['L']
            exponents = shell['exponents']
            raw_coefficients = shell['coefficients']
            normalized_coefficients = normalize_shell_coefficients(
                coefficients=raw_coefficients,
                exponents=exponents,
                L=L,
                normalization=shell['normalization'],
                )
            aoshell_handle = ce.cuestAOShellHandle()
            cuest_check(f'Create AO Shell {shell_count}',
                ce.cuestAOShellCreate(
                    handle=cuest_handle,
                    isPure=shell['is_pure'],
                    L=L,
                    numPrimitive=len(exponents),
                    exponents=exponents,
                    coefficients=normalized_coefficients,
                    parameters=aoshell_parameters,
                    outShell=aoshell_handle)
                )
            shells.append(aoshell_handle)
            shell_count += 1

    cuest_check('Destroy AO Shell Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
            parameters=aoshell_parameters,
            )
        )

    # => Build Basis Sets From Gaussian Shells <= #

    aobasis_parameters = ce.cuestAOBasisParameters()
    cuest_check('Create AO Basis Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
            outParameters=aobasis_parameters,
            )
        )

    # Orbital Basis
    aobasis_handle = ce.cuestAOBasisHandle()
    cuest_check('Create AOBasis Workspace Query',
        ce.cuestAOBasisCreateWorkspaceQuery(
            handle=cuest_handle,
            numAtoms=num_atoms,
            numShellsPerAtom=n_shells_per_atom,
            shells=shells,
            parameters=aobasis_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outBasis=aobasis_handle,
            )
        )

    aobasis_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    aobasis_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('Create AOBasis',
        ce.cuestAOBasisCreate(
            handle=cuest_handle,
            numAtoms=num_atoms,
            numShellsPerAtom=n_shells_per_atom,
            shells=shells,
            parameters=aobasis_parameters,
            persistentWorkspace=aobasis_persistent_workspace.pointer,
            temporaryWorkspace=aobasis_temporary_workspace.pointer,
            outBasis=aobasis_handle,
            )
        )

    del aobasis_temporary_workspace

    for i, shell in enumerate(shells):
        cuest_check(f'Destroy AO Shell {i+1}',
            ce.cuestAOShellDestroy(
                handle=shell,
                )
            )

    aobasis_num_ao = ce.data_uint64_t()
    cuest_check('Query AO Basis Num AO',
        ce.cuestQuery(
            handle=cuest_handle,
            type=ce.CuestType.CUEST_AOBASIS,
            object=aobasis_handle,
            attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_AO,
            attributeValue=aobasis_num_ao,
            )
        )
    nao = aobasis_num_ao.value

    cuest_check('Destroy AO Basis Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
            parameters=aobasis_parameters,
            )
        )

    # => Build Molecular Integration Grid <= #

    atomgrid_parameters = ce.cuestAtomGridParameters()

    cuest_check('Create AtomGrid Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ATOMGRID_PARAMETERS,
            outParameters=atomgrid_parameters,
            )
        )

    atomgrids = []
    for atom,symbol in enumerate(symbols):
        # For simplicity, build an unpruned 75/302 grid using the Ahlrichs radial quadrature scheme
        ahlrichs_radius = symbol_to_ahlrichs_radius(
            symbol=symbol
            )
        radial_nodes, radial_weights = build_ahlrichs_radial_quadrature(
            npoint=75,
            R=ahlrichs_radius,
            )
        num_angular_points = [302]*len(radial_nodes)
        atomgrid_handle = ce.cuestAtomGridHandle()
        cuest_check(f'Create AtomGrid {atom}',
            ce.cuestAtomGridCreate(
                handle=cuest_handle,
                numRadialPoints=75,
                radialNodes=radial_nodes,
                radialWeights=radial_weights,
                numAngularPoints=num_angular_points,
                parameters=atomgrid_parameters,
                outAtomGrid=atomgrid_handle,
                )
            )
        atomgrids.append(atomgrid_handle)

    cuest_check('Destroy AtomGrid Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ATOMGRID_PARAMETERS,
            parameters=atomgrid_parameters,
            )
        )

    moleculargrid_parameters = ce.cuestMolecularGridParameters()

    cuest_check('Create MolecularGrid Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_MOLECULARGRID_PARAMETERS,
            outParameters=moleculargrid_parameters,
            )
        )

    cuest_check('Create MolecularGrid Workspace Query',
        ce.cuestMolecularGridCreateWorkspaceQuery(
            handle=cuest_handle,
            numAtoms=len(Zs),
            atomGrid=atomgrids,
            xyz=xyzs,
            parameters=moleculargrid_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outGrid=ce.cuestMolecularGridHandle(),
            )
        )

    grid_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    grid_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    moleculargrid=ce.cuestMolecularGridHandle()

    cuest_check('Create MolecularGrid',
        ce.cuestMolecularGridCreate(
            handle=cuest_handle,
            numAtoms=len(Zs),
            atomGrid=atomgrids,
            xyz=xyzs,
            parameters=moleculargrid_parameters,
            persistentWorkspace=grid_persistent_workspace.pointer,
            temporaryWorkspace=grid_temporary_workspace.pointer,
            outGrid=moleculargrid,
            )
        )

    del grid_temporary_workspace

    for atom,grid in enumerate(atomgrids):
        cuest_check(f'Destroy AtomGrid{atom}',
            ce.cuestAtomGridDestroy(
                atomGrid=grid,
                )
            )

    cuest_check('Destroy MolecularGrid Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_MOLECULARGRID_PARAMETERS,
            parameters=moleculargrid_parameters,
            )
        )

    # => Build XC Integral Plan <= #

    xcintplan_parameters = ce.cuestXCIntPlanParameters()
    cuest_check('Create XCIntPlan Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_XCINTPLAN_PARAMETERS,
            outParameters=xcintplan_parameters,
            )
        )

    cuest_check('Create XCIntPlan Workspace Query',
        ce.cuestXCIntPlanCreateWorkspaceQuery(
            handle=cuest_handle,
            basis=aobasis_handle,
            grid=moleculargrid,
            functional=ce.CuestXCIntPlanParametersFunctional.CUEST_XCINTPLAN_PARAMETERS_FUNCTIONAL_PBE,
            parameters=xcintplan_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outPlan=ce.cuestXCIntPlanHandle(),
            )
        )

    xcintplan = ce.cuestXCIntPlanHandle()
    xcintplan_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    xcintplan_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('Create XCIntPlan',
        ce.cuestXCIntPlanCreate(
            handle=cuest_handle,
            basis=aobasis_handle,
            grid=moleculargrid,
            functional=ce.CuestXCIntPlanParametersFunctional.CUEST_XCINTPLAN_PARAMETERS_FUNCTIONAL_PBE,
            parameters=xcintplan_parameters,
            persistentWorkspace=xcintplan_persistent_workspace.pointer,
            temporaryWorkspace=xcintplan_temporary_workspace.pointer,
            outPlan=xcintplan,
            )
        )

    del xcintplan_temporary_workspace

    cuest_check('Destroy XCIntPlan Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_XCINTPLAN_PARAMETERS,
            parameters=xcintplan_parameters,
            )
        )

    # => Configure the Nonlocal Potential <= #

    nonlocal_xc_compute_parameters = ce.cuestNonlocalXCPotentialRKSComputeParameters()

    cuest_check('NonlocalXC Parameters Create',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            outParameters=nonlocal_xc_compute_parameters,
            )
        )

    vv10_b = ce.data_double(6.0)
    cuest_check('VV10 b Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialRKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS_VV10_B,
            attributeValue=vv10_b,
            )
        )

    vv10_C = ce.data_double(0.01)
    cuest_check('VV10 C Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialRKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS_VV10_C,
            attributeValue=vv10_C,
            )
        )

    vv10_scale = ce.data_double(1.0)
    cuest_check('VV10 Scale Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialRKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS_VV10_SCALE,
            attributeValue=vv10_scale,
            )
        )

    # => RKS Potential Matrix Computation <= #

    # => Build Fake Occupied Orbitals <= #

    nocc = np.abs(int(np.sum(Zs)//2))

    Cocc_host = np.random.normal(size=(nocc, nao)).astype(np.double)
    Cocc_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nocc * nao
        )
    Cocc_device_handle = ce.Pointer(Cocc_device_pointer)
    cuda_memcpy_htod(
        device_pointer=Cocc_device_pointer,
        host_pointer=Cocc_host.ctypes.data,
        size_in_bytes=sizeof_double * nocc * nao,
        stream=stream_handle.value,
        )

    # => Call the XC Potential Compute Routines <= #

    Vxc_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nao**2
        )
    Vxc_device_handle = ce.Pointer(Vxc_device_pointer)

    Exc = ce.data_double()
    cuest_check('NonlocalXCPotentialRKSCompute Workspace Query',
        ce.cuestNonlocalXCPotentialRKSComputeWorkspaceQuery(
            handle=cuest_handle,
            plan=xcintplan,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            parameters=nonlocal_xc_compute_parameters,
            numOccupied=nocc,
            coefficientMatrix=Cocc_device_handle,
            outXCEnergy=Exc,
            outXCPotentialMatrix=Vxc_device_handle,
            )
        )

    Vxc_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('NonlocalXCPotentialRKSCompute',
        ce.cuestNonlocalXCPotentialRKSCompute(
            handle=cuest_handle,
            plan=xcintplan,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspace=Vxc_workspace.pointer,
            parameters=nonlocal_xc_compute_parameters,
            numOccupied=nocc,
            coefficientMatrix=Cocc_device_handle,
            outXCEnergy=Exc,
            outXCPotentialMatrix=Vxc_device_handle,
            )
        )

    del Vxc_workspace

    cuda_free(
        array=Cocc_device_pointer
        )
    cuest_check('Destroy NonlocalXC Parameters',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            )
        )

    Vxc_host_array = np.empty(
        (nao * nao),
        dtype=np.double
        )
    cuda_memcpy_dtoh(
        host_pointer=Vxc_host_array.ctypes.data,
        device_pointer=Vxc_device_pointer,
        size_in_bytes=sizeof_double * nao**2,
        stream=stream_handle.value,
        )

    cuda_free(
        array=Vxc_device_pointer
        )

    # => UKS Potential Matrix Computation <= #

    # => Build Fake Occupied Orbitals <= #

    # Assume a cation
    nocca = np.abs(int(np.sum(Zs)//2))
    noccb = nocca - 1

    # Alpha and Beta occupied molecular orbitals
    Cocca_host = np.random.normal(size=(nocca, nao)).astype(np.double)
    Cocca_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nocca * nao
        )
    Cocca_device_handle = ce.Pointer(Cocca_device_pointer)
    cuda_memcpy_htod(
        device_pointer=Cocca_device_pointer,
        host_pointer=Cocca_host.ctypes.data,
        size_in_bytes=sizeof_double * nocca * nao,
        stream=stream_handle.value,
        )

    Coccb_host = np.random.normal(size=(noccb, nao)).astype(np.double)
    Coccb_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * noccb * nao
        )
    Coccb_device_handle = ce.Pointer(Coccb_device_pointer)
    cuda_memcpy_htod(
        device_pointer=Coccb_device_pointer,
        host_pointer=Coccb_host.ctypes.data,
        size_in_bytes=sizeof_double * noccb * nao,
        stream=stream_handle.value,
        )

    # => Call the XC Potential Compute Routines <= #

    Vxc_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nao**2
        )
    Vxc_device_handle = ce.Pointer(Vxc_device_pointer)
    nonlocal_xc_compute_parameters = ce.cuestNonlocalXCPotentialUKSComputeParameters()

    cuest_check('NonlocalXC UKS Parameters Create',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS,
            outParameters=nonlocal_xc_compute_parameters,
            )
        )
    vv10_b = ce.data_double(6.0)
    cuest_check('VV10 b UKS Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialUKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS_VV10_B,
            attributeValue=vv10_b,
            )
        )
    vv10_C = ce.data_double(0.01)
    cuest_check('VV10 C UKS Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameter
# ... truncated ...


run

#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cuest.bindings as ce
import numpy as np
import argparse
from pathlib import Path
import sys

# Inject the directory where the example helper utilities live
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from helpers.cuda_utils import cuda_malloc, cuda_free, cuda_memcpy_dtoh, cuda_memcpy_htod, WorkspaceDescriptor, Workspace
from helpers.utilities import normalize_shell_coefficients
from helpers.parsers import simple_xyz_parser, simple_gbs_parser
from helpers.grid_utils import build_ahlrichs_radial_quadrature, symbol_to_ahlrichs_radius

def run(
    *,
    xyz_filename,
    gbs_filename,
    ):

    # => Parse User Input <= #

    symbols, xyzs, Zs = simple_xyz_parser(
        filename=xyz_filename,
        to_bohr_scale_factor=1.0 / 0.52917720859,
        )

    shellinfo = simple_gbs_parser(
        filename=gbs_filename,
        symbols=symbols,
        )

    num_atoms = len(symbols)
    sizeof_double = np.dtype('double').itemsize

    # => Set Up cuEST <= #

    def cuest_check(
        title,
        return_code
        ):
        " A simple helper function to call cuEST functions and check the return code. "

        if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
            raise RuntimeError(f"{title} failed with code {return_code}")

    # These are user-provided functions that are responsible for allocating
    # host and device workspace arrays.
    persistent_workspace_descriptor = WorkspaceDescriptor()
    temporary_workspace_descriptor = WorkspaceDescriptor()
    # Allow 2Gb of DRAM for XC evaluation routines
    variable_buffer_size = WorkspaceDescriptor(
        host_buffer_size_in_bytes=0,
        device_buffer_size_in_bytes=2000000000,
        )

    cuest_handle_parameters = ce.cuestHandleParameters()

    cuest_check('Create Handle Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            outParameters=cuest_handle_parameters,
            )
        )

    # Here we allow cuEST to use the default stream, cuBLAS, etc., but those
    # parameters should be set in the handle parameters before this call.
    cuest_handle = ce.cuestHandle()
    cuest_check('Create Cuest Handle',
        ce.cuestCreate(
            parameters=cuest_handle_parameters,
            handle=cuest_handle,
            )
        )

    stream_handle = ce.data_cudaStream_t()
    cuest_check('Query Handle Stream',
        ce.cuestParametersQuery(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            parameters=cuest_handle_parameters,
            attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_CUDASTREAM,
            attributeValue=stream_handle,
            )
        )

    cuest_check('Destroy Handle Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            parameters=cuest_handle_parameters,
            )
        )

    # => Build Gaussian Shells From Basis Parser Info <= #

    aoshell_parameters = ce.cuestAOShellParameters()
    cuest_check('Create AO Shell Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
            outParameters=aoshell_parameters,
            )
        )

    shells = []
    n_shells_per_atom = []
    shell_count = 1
    for atom_shells in shellinfo:
        n_shells_per_atom.append(len(atom_shells))
        for shell in atom_shells:
            L = shell['L']
            exponents = shell['exponents']
            raw_coefficients = shell['coefficients']
            normalized_coefficients = normalize_shell_coefficients(
                coefficients=raw_coefficients,
                exponents=exponents,
                L=L,
                normalization=shell['normalization'],
                )
            aoshell_handle = ce.cuestAOShellHandle()
            cuest_check(f'Create AO Shell {shell_count}',
                ce.cuestAOShellCreate(
                    handle=cuest_handle,
                    isPure=shell['is_pure'],
                    L=L,
                    numPrimitive=len(exponents),
                    exponents=exponents,
                    coefficients=normalized_coefficients,
                    parameters=aoshell_parameters,
                    outShell=aoshell_handle)
                )
            shells.append(aoshell_handle)
            shell_count += 1

    cuest_check('Destroy AO Shell Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
            parameters=aoshell_parameters,
            )
        )

    # => Build Basis Sets From Gaussian Shells <= #

    aobasis_parameters = ce.cuestAOBasisParameters()
    cuest_check('Create AO Basis Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
            outParameters=aobasis_parameters,
            )
        )

    # Orbital Basis
    aobasis_handle = ce.cuestAOBasisHandle()
    cuest_check('Create AOBasis Workspace Query',
        ce.cuestAOBasisCreateWorkspaceQuery(
            handle=cuest_handle,
            numAtoms=num_atoms,
            numShellsPerAtom=n_shells_per_atom,
            shells=shells,
            parameters=aobasis_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outBasis=aobasis_handle,
            )
        )

    aobasis_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    aobasis_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('Create AOBasis',
        ce.cuestAOBasisCreate(
            handle=cuest_handle,
            numAtoms=num_atoms,
            numShellsPerAtom=n_shells_per_atom,
            shells=shells,
            parameters=aobasis_parameters,
            persistentWorkspace=aobasis_persistent_workspace.pointer,
            temporaryWorkspace=aobasis_temporary_workspace.pointer,
            outBasis=aobasis_handle,
            )
        )

    del aobasis_temporary_workspace

    for i, shell in enumerate(shells):
        cuest_check(f'Destroy AO Shell {i+1}',
            ce.cuestAOShellDestroy(
                handle=shell,
                )
            )

    aobasis_num_ao = ce.data_uint64_t()
    cuest_check('Query AO Basis Num AO',
        ce.cuestQuery(
            handle=cuest_handle,
            type=ce.CuestType.CUEST_AOBASIS,
            object=aobasis_handle,
            attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_AO,
            attributeValue=aobasis_num_ao,
            )
        )
    nao = aobasis_num_ao.value

    cuest_check('Destroy AO Basis Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
            parameters=aobasis_parameters,
            )
        )

    # => Build Molecular Integration Grid <= #

    atomgrid_parameters = ce.cuestAtomGridParameters()

    cuest_check('Create AtomGrid Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ATOMGRID_PARAMETERS,
            outParameters=atomgrid_parameters,
            )
        )

    atomgrids = []
    for atom,symbol in enumerate(symbols):
        # For simplicity, build an unpruned 75/302 grid using the Ahlrichs radial quadrature scheme
        ahlrichs_radius = symbol_to_ahlrichs_radius(
            symbol=symbol
            )
        radial_nodes, radial_weights = build_ahlrichs_radial_quadrature(
            npoint=75,
            R=ahlrichs_radius,
            )
        num_angular_points = [302]*len(radial_nodes)
        atomgrid_handle = ce.cuestAtomGridHandle()
        cuest_check(f'Create AtomGrid {atom}',
            ce.cuestAtomGridCreate(
                handle=cuest_handle,
                numRadialPoints=75,
                radialNodes=radial_nodes,
                radialWeights=radial_weights,
                numAngularPoints=num_angular_points,
                parameters=atomgrid_parameters,
                outAtomGrid=atomgrid_handle,
                )
            )
        atomgrids.append(atomgrid_handle)

    cuest_check('Destroy AtomGrid Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ATOMGRID_PARAMETERS,
            parameters=atomgrid_parameters,
            )
        )

    moleculargrid_parameters = ce.cuestMolecularGridParameters()

    cuest_check('Create MolecularGrid Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_MOLECULARGRID_PARAMETERS,
            outParameters=moleculargrid_parameters,
            )
        )

    cuest_check('Create MolecularGrid Workspace Query',
        ce.cuestMolecularGridCreateWorkspaceQuery(
            handle=cuest_handle,
            numAtoms=len(Zs),
            atomGrid=atomgrids,
            xyz=xyzs,
            parameters=moleculargrid_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outGrid=ce.cuestMolecularGridHandle(),
            )
        )

    grid_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    grid_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    moleculargrid=ce.cuestMolecularGridHandle()

    cuest_check('Create MolecularGrid',
        ce.cuestMolecularGridCreate(
            handle=cuest_handle,
            numAtoms=len(Zs),
            atomGrid=atomgrids,
            xyz=xyzs,
            parameters=moleculargrid_parameters,
            persistentWorkspace=grid_persistent_workspace.pointer,
            temporaryWorkspace=grid_temporary_workspace.pointer,
            outGrid=moleculargrid,
            )
        )

    del grid_temporary_workspace

    for atom,grid in enumerate(atomgrids):
        cuest_check(f'Destroy AtomGrid{atom}',
            ce.cuestAtomGridDestroy(
                atomGrid=grid,
                )
            )

    cuest_check('Destroy MolecularGrid Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_MOLECULARGRID_PARAMETERS,
            parameters=moleculargrid_parameters,
            )
        )

    # => Build XC Integral Plan <= #

    xcintplan_parameters = ce.cuestXCIntPlanParameters()
    cuest_check('Create XCIntPlan Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_XCINTPLAN_PARAMETERS,
            outParameters=xcintplan_parameters,
            )
        )

    cuest_check('Create XCIntPlan Workspace Query',
        ce.cuestXCIntPlanCreateWorkspaceQuery(
            handle=cuest_handle,
            basis=aobasis_handle,
            grid=moleculargrid,
            functional=ce.CuestXCIntPlanParametersFunctional.CUEST_XCINTPLAN_PARAMETERS_FUNCTIONAL_PBE,
            parameters=xcintplan_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outPlan=ce.cuestXCIntPlanHandle(),
            )
        )

    xcintplan = ce.cuestXCIntPlanHandle()
    xcintplan_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    xcintplan_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('Create XCIntPlan',
        ce.cuestXCIntPlanCreate(
            handle=cuest_handle,
            basis=aobasis_handle,
            grid=moleculargrid,
            functional=ce.CuestXCIntPlanParametersFunctional.CUEST_XCINTPLAN_PARAMETERS_FUNCTIONAL_PBE,
            parameters=xcintplan_parameters,
            persistentWorkspace=xcintplan_persistent_workspace.pointer,
            temporaryWorkspace=xcintplan_temporary_workspace.pointer,
            outPlan=xcintplan,
            )
        )

    del xcintplan_temporary_workspace

    cuest_check('Destroy XCIntPlan Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_XCINTPLAN_PARAMETERS,
            parameters=xcintplan_parameters,
            )
        )

    # => Configure the Nonlocal Potential <= #

    nonlocal_xc_compute_parameters = ce.cuestNonlocalXCPotentialRKSComputeParameters()

    cuest_check('NonlocalXC Parameters Create',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            outParameters=nonlocal_xc_compute_parameters,
            )
        )

    vv10_b = ce.data_double(6.0)
    cuest_check('VV10 b Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialRKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS_VV10_B,
            attributeValue=vv10_b,
            )
        )

    vv10_C = ce.data_double(0.01)
    cuest_check('VV10 C Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialRKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS_VV10_C,
            attributeValue=vv10_C,
            )
        )

    vv10_scale = ce.data_double(1.0)
    cuest_check('VV10 Scale Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialRKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS_VV10_SCALE,
            attributeValue=vv10_scale,
            )
        )

    # => RKS Potential Matrix Computation <= #

    # => Build Fake Occupied Orbitals <= #

    nocc = np.abs(int(np.sum(Zs)//2))

    Cocc_host = np.random.normal(size=(nocc, nao)).astype(np.double)
    Cocc_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nocc * nao
        )
    Cocc_device_handle = ce.Pointer(Cocc_device_pointer)
    cuda_memcpy_htod(
        device_pointer=Cocc_device_pointer,
        host_pointer=Cocc_host.ctypes.data,
        size_in_bytes=sizeof_double * nocc * nao,
        stream=stream_handle.value,
        )

    # => Call the XC Potential Compute Routines <= #

    Vxc_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nao**2
        )
    Vxc_device_handle = ce.Pointer(Vxc_device_pointer)

    Exc = ce.data_double()
    cuest_check('NonlocalXCPotentialRKSCompute Workspace Query',
        ce.cuestNonlocalXCPotentialRKSComputeWorkspaceQuery(
            handle=cuest_handle,
            plan=xcintplan,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            parameters=nonlocal_xc_compute_parameters,
            numOccupied=nocc,
            coefficientMatrix=Cocc_device_handle,
            outXCEnergy=Exc,
            outXCPotentialMatrix=Vxc_device_handle,
            )
        )

    Vxc_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('NonlocalXCPotentialRKSCompute',
        ce.cuestNonlocalXCPotentialRKSCompute(
            handle=cuest_handle,
            plan=xcintplan,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspace=Vxc_workspace.pointer,
            parameters=nonlocal_xc_compute_parameters,
            numOccupied=nocc,
            coefficientMatrix=Cocc_device_handle,
            outXCEnergy=Exc,
            outXCPotentialMatrix=Vxc_device_handle,
            )
        )

    del Vxc_workspace

    cuda_free(
        array=Cocc_device_pointer
        )
    cuest_check('Destroy NonlocalXC Parameters',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALRKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            )
        )

    Vxc_host_array = np.empty(
        (nao * nao),
        dtype=np.double
        )
    cuda_memcpy_dtoh(
        host_pointer=Vxc_host_array.ctypes.data,
        device_pointer=Vxc_device_pointer,
        size_in_bytes=sizeof_double * nao**2,
        stream=stream_handle.value,
        )

    cuda_free(
        array=Vxc_device_pointer
        )

    # => UKS Potential Matrix Computation <= #

    # => Build Fake Occupied Orbitals <= #

    # Assume a cation
    nocca = np.abs(int(np.sum(Zs)//2))
    noccb = nocca - 1

    # Alpha and Beta occupied molecular orbitals
    Cocca_host = np.random.normal(size=(nocca, nao)).astype(np.double)
    Cocca_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nocca * nao
        )
    Cocca_device_handle = ce.Pointer(Cocca_device_pointer)
    cuda_memcpy_htod(
        device_pointer=Cocca_device_pointer,
        host_pointer=Cocca_host.ctypes.data,
        size_in_bytes=sizeof_double * nocca * nao,
        stream=stream_handle.value,
        )

    Coccb_host = np.random.normal(size=(noccb, nao)).astype(np.double)
    Coccb_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * noccb * nao
        )
    Coccb_device_handle = ce.Pointer(Coccb_device_pointer)
    cuda_memcpy_htod(
        device_pointer=Coccb_device_pointer,
        host_pointer=Coccb_host.ctypes.data,
        size_in_bytes=sizeof_double * noccb * nao,
        stream=stream_handle.value,
        )

    # => Call the XC Potential Compute Routines <= #

    Vxc_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nao**2
        )
    Vxc_device_handle = ce.Pointer(Vxc_device_pointer)
    nonlocal_xc_compute_parameters = ce.cuestNonlocalXCPotentialUKSComputeParameters()

    cuest_check('NonlocalXC UKS Parameters Create',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS,
            outParameters=nonlocal_xc_compute_parameters,
            )
        )
    vv10_b = ce.data_double(6.0)
    cuest_check('VV10 b UKS Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialUKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS_VV10_B,
            attributeValue=vv10_b,
            )
        )
    vv10_C = ce.data_double(0.01)
    cuest_check('VV10 C UKS Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestParametersType.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS,
            parameters=nonlocal_xc_compute_parameters,
            attribute=ce.CuestNonlocalXCPotentialUKSComputeParametersAttributes.CUEST_NONLOCALXCPOTENTIALUKSCOMPUTE_PARAMETERS_VV10_C,
            attributeValue=vv10_C,
            )
        )
    vv10_scale = ce.data_double(1.0)
    cuest_check('VV10 Scale UKS Parameter Configure',
        ce.cuestParametersConfigure(
            parametersType=ce.CuestP
# ... truncated ...


    gbs_filename,
    ):

    # => Parse User Input <= #

    symbols, xyzs, Zs = simple_xyz_parser(
        filename=xyz_filename,
        to_bohr_scale_factor=1.0 / 0.52917720859,
        )

    shellinfo = simple_gbs_parser(
        filename=gbs_filename,
        symbols=symbols,
        )

    num_atoms = len(symbols)
    sizeof_double = np.dtype('double').itemsize

    # => Set Up cuEST <= #

    def cuest_check(
        title,
        return_code
        ):
        " A simple helper function to call cuEST functions and check the return code. "

        if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
            raise RuntimeError(f"{title} failed with code {return_code}")

    # These are user-provided functions that are responsible for allocating
    # host and device workspace arrays.
    persistent_workspace_descriptor = WorkspaceDescriptor()
    temporary_workspace_descriptor = WorkspaceDescriptor()
    # Allow 2Gb of DRAM for XC evaluation routines
    variable_buffer_size = WorkspaceDescriptor(
        host_buffer_size_in_bytes=0,
        device_buffer_size_in_bytes=2000000000,
        )

    cuest_handle_parameters = ce.cuestHandleParameters()

    cuest_check('Create Handle Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            outParameters=cuest_handle_parameters,
            )
        )


# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cuest.bindings as ce

import sys
from pathlib import Path

# Inject the directory where the example helper utilities live
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from helpers.cuda_utils import WorkspaceDescriptor, Workspace
from helpers.utilities import normalize_shell_coefficients

# This is the def2-SVP basis set definition for O and H

H_basis = {
    'n_shells' : 3,
    'shell_types' : [0, 0, 1],
    'primitive_counts' : [3, 1, 1],
    'primitive_offsets' : [0, 3, 4],
    'exponents' : [
        13.0107010, 1.9622572, 0.44453796,
        0.12194962,
        0.8000000
        ],
    'coefficients' : [
        0.019682158, 0.13796524, 0.47831935,
        1.0,
        1.0
        ],
    }

O_basis = {
    'n_shells' : 6,
    'shell_types' : [0, 0, 0, 1, 1, 2],
    'primitive_counts' : [5, 1, 1, 3, 1, 1],
    'primitive_offsets' : [0, 5, 6, 7, 10, 11],
    'exponents' : [
        2266.1767785, 340.87010191, 77.363135167, 21.479644940, 6.6589433124,
        0.80975975668,
        0.25530772234,
        17.721504317, 3.8635505440, 1.0480920883,
        0.27641544411,
        1.2
        ],
    'coefficients' : [
        -0.0053431809926, -0.039890039230, -0.17853911985, -0.46427684959, -0.44309745172,
        1.0,
        1.0,
        0.043394573193, 0.23094120765, 0.51375311064,
        1.0,
        1.0
        ],
    }

def cuest_check(
    title,
    return_code
    ):

    if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
        raise RuntimeError(f"{title} failed with code {return_code}")

cuest_handle_parameters = ce.cuestHandleParameters()
cuest_check("Parameters Create",
    ce.cuestParametersCreate(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        outParameters=cuest_handle_parameters,
        )
    )
cuest_handle = ce.cuestHandle()
cuest_check("Cuest Handle Create",
    ce.cuestCreate(
        parameters=cuest_handle_parameters,
        handle=cuest_handle,
        )
    )
cuest_check("Parameters Destroy",
    ce.cuestParametersDestroy(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        )
    )

aoshell_parameters = ce.cuestAOShellParameters()
cuest_check("AO Shell Parameters Create",
    ce.cuestParametersCreate(
        parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
        outParameters=aoshell_parameters,
        )
    )

# Make a water molecule as an example
basis_is_pure = 1
shells = []
atom_symbols = ["O", "H", "H"]
atom_bases = [O_basis, H_basis, H_basis]
for shellinfo in atom_bases:
    this_atom_shells = []

    for offset, count, L in zip(
        shellinfo['primitive_offsets'],
        shellinfo['primitive_counts'],
        shellinfo['shell_types']
        ):

        exponents = shellinfo['exponents'][offset:offset+count]
        coefficients = shellinfo['coefficients'][offset:offset+count]
        normalized_coefficients = normalize_shell_coefficients(
            coefficients=coefficients,
            exponents=exponents,
            L=L,
            normalization=1.0,
            )
        aoshell_handle = ce.cuestAOShellHandle()
        cuest_check("AO Shell Create",
            ce.cuestAOShellCreate(
                handle=cuest_handle,
                isPure=basis_is_pure,
                L=L,
                numPrimitive=len(exponents),
                exponents=exponents,
                coefficients=normalized_coefficients,
                parameters=aoshell_parameters,
                outShell=aoshell_handle
                )
            )
        this_atom_shells.append(aoshell_handle)
    shells.append(this_atom_shells)

cuest_check("AO Shell Parameters Destroy",
    ce.cuestParametersDestroy(
        parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
        parameters=aoshell_parameters,
        )
    )

# Query each shell for its metadata, and print
for atom_number, atom_shells in enumerate(shells):

    print(f"Atom {atom_number+1} ({atom_symbols[atom_number]})")

    for shell_count, atom_shell in enumerate(atom_shells):
        print(f"\tShell {shell_count+1}")
        is_pure = ce.data_int32_t()
        cuest_check("Query IS_PURE",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_IS_PURE,
                attributeValue=is_pure,
                )
            )
        print(f"\t\tIs Pure                  : {is_pure.value}")

        L = ce.data_uint64_t()
        cuest_check("Query L",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_L,
                attributeValue=L,
                )
            )
        print(f"\t\tAngular Momentum         : {L.value}")

        nprim = ce.data_uint64_t()
        cuest_check("Query Num Primitives",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_NUM_PRIMITIVE,
                attributeValue=nprim,
                )
            )
        print(f"\t\tNum. Primitives          : {nprim.value}")

        nao = ce.data_uint64_t()
        cuest_check("Query Num Basis Functions",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_NUM_AO,
                attributeValue=nao,
                )
            )
        print(f"\t\tNum. Basis Functions     : {nao.value}")

        npure = ce.data_uint64_t()
        cuest_check("Query Num Pure Functions",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_NUM_PURE,
                attributeValue=npure,
                )
            )
        print(f"\t\tNum. Pure Functions      : {npure.value}")

        ncart = ce.data_uint64_t()
        cuest_check("Query Num Cart Functions",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_NUM_CART,
                attributeValue=npure,
                )
            )
        print(f"\t\tNum. Cartesian Functions : {npure.value}")

# Destroy the shells
for atom_shells in shells:
    for atom_shell in atom_shells:
        cuest_check("Destroy AO Shell",
            ce.cuestAOShellDestroy(
                handle=atom_shell,
                )
            )

# Delete the cuEST handle
cuest_check('Destroy Cuest Handle',
    ce.cuestDestroy(
        handle=cuest_handle,
        )
    )


    'primitive_offsets' : [0, 5, 6, 7, 10, 11],
    'exponents' : [
        2266.1767785, 340.87010191, 77.363135167, 21.479644940, 6.6589433124,
        0.80975975668,
        0.25530772234,
        17.721504317, 3.8635505440, 1.0480920883,
        0.27641544411,
        1.2
        ],
    'coefficients' : [
        -0.0053431809926, -0.039890039230, -0.17853911985, -0.46427684959, -0.44309745172,
        1.0,
        1.0,
        0.043394573193, 0.23094120765, 0.51375311064,
        1.0,
        1.0
        ],
    }

def cuest_check(
    title,
    return_code
    ):

    if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
        raise RuntimeError(f"{title} failed with code {return_code}")

cuest_handle_parameters = ce.cuestHandleParameters()
cuest_check("Parameters Create",
    ce.cuestParametersCreate(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        outParameters=cuest_handle_parameters,
        )
    )
cuest_handle = ce.cuestHandle()
cuest_check("Cuest Handle Create",
    ce.cuestCreate(
        parameters=cuest_handle_parameters,
        handle=cuest_handle,
        )
    )
cuest_check("Parameters Destroy",
    ce.cuestParametersDestroy(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,


# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cuest.bindings as ce

import sys
from pathlib import Path

# Inject the directory where the example helper utilities live
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from helpers.cuda_utils import WorkspaceDescriptor, Workspace
from helpers.utilities import normalize_shell_coefficients

# This is the def2-SVP basis set definition for O and H

H_basis = {
    'n_shells' : 3,
    'shell_types' : [0, 0, 1],
    'primitive_counts' : [3, 1, 1],
    'primitive_offsets' : [0, 3, 4],
    'exponents' : [
        13.0107010, 1.9622572, 0.44453796,
        0.12194962,
        0.8000000
        ],
    'coefficients' : [
        0.019682158, 0.13796524, 0.47831935,
        1.0,
        1.0
        ],
    }

O_basis = {
    'n_shells' : 6,
    'shell_types' : [0, 0, 0, 1, 1, 2],
    'primitive_counts' : [5, 1, 1, 3, 1, 1],
    'primitive_offsets' : [0, 5, 6, 7, 10, 11],
    'exponents' : [
        2266.1767785, 340.87010191, 77.363135167, 21.479644940, 6.6589433124,
        0.80975975668,
        0.25530772234,
        17.721504317, 3.8635505440, 1.0480920883,
        0.27641544411,
        1.2
        ],
    'coefficients' : [
        -0.0053431809926, -0.039890039230, -0.17853911985, -0.46427684959, -0.44309745172,
        1.0,
        1.0,
        0.043394573193, 0.23094120765, 0.51375311064,
        1.0,
        1.0
        ],
    }

def cuest_check(
    title,
    return_code
    ):

    if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
        raise RuntimeError(f"{title} failed with code {return_code}")

cuest_handle_parameters = ce.cuestHandleParameters()
cuest_check("Parameters Create",
    ce.cuestParametersCreate(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        outParameters=cuest_handle_parameters,
        )
    )
cuest_handle = ce.cuestHandle()
cuest_check("Cuest Handle Create",
    ce.cuestCreate(
        parameters=cuest_handle_parameters,
        handle=cuest_handle,
        )
    )
cuest_check("Parameters Destroy",
    ce.cuestParametersDestroy(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        )
    )

aoshell_parameters = ce.cuestAOShellParameters()
cuest_check("AO Shell Parameters Create",
    ce.cuestParametersCreate(
        parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
        outParameters=aoshell_parameters,
        )
    )

# Make a water molecule as an example
basis_is_pure = 1
shells = []
atom_symbols = ["O", "H", "H"]
atom_bases = [O_basis, H_basis, H_basis]
for shellinfo in atom_bases:
    this_atom_shells = []

    for offset, count, L in zip(
        shellinfo['primitive_offsets'],
        shellinfo['primitive_counts'],
        shellinfo['shell_types']
        ):

        exponents = shellinfo['exponents'][offset:offset+count]
        coefficients = shellinfo['coefficients'][offset:offset+count]
        normalized_coefficients = normalize_shell_coefficients(
            coefficients=coefficients,
            exponents=exponents,
            L=L,
            normalization=1.0,
            )
        aoshell_handle = ce.cuestAOShellHandle()
        cuest_check("AO Shell Create",
            ce.cuestAOShellCreate(
                handle=cuest_handle,
                isPure=basis_is_pure,
                L=L,
                numPrimitive=len(exponents),
                exponents=exponents,
                coefficients=normalized_coefficients,
                parameters=aoshell_parameters,
                outShell=aoshell_handle
                )
            )
        this_atom_shells.append(aoshell_handle)
    shells.append(this_atom_shells)

n_shells_per_atom = [len(atom_shells) for atom_shells in shells]

cuest_check("AO Shell Parameters Destroy",
    ce.cuestParametersDestroy(
        parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
        parameters=aoshell_parameters,
        )
    )

# Query each shell for its metadata, and print
for atom_number, atom_shells in enumerate(shells):

    print(f"Atom {atom_number+1} ({atom_symbols[atom_number]})")

    for shell_count, atom_shell in enumerate(atom_shells):
        print(f"\tShell {shell_count+1}")
        is_pure = ce.data_int32_t()
        cuest_check("Query IS_PURE",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_IS_PURE,
                attributeValue=is_pure,
                )
            )
        print(f"\t\tIs Pure                  : {is_pure.value}")

        L = ce.data_uint64_t()
        cuest_check("Query L",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_L,
                attributeValue=L,
                )
            )
        print(f"\t\tAngular Momentum         : {L.value}")

        nprim = ce.data_uint64_t()
        cuest_check("Query Num Primitives",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_NUM_PRIMITIVE,
                attributeValue=nprim,
                )
            )
        print(f"\t\tNum. Primitives          : {nprim.value}")

        nao = ce.data_uint64_t()
        cuest_check("Query Num Basis Functions",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_NUM_AO,
                attributeValue=nao,
                )
            )
        print(f"\t\tNum. Basis Functions     : {nao.value}")

        npure = ce.data_uint64_t()
        cuest_check("Query Num Pure Functions",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_NUM_PURE,
                attributeValue=npure,
                )
            )
        print(f"\t\tNum. Pure Functions      : {npure.value}")

        ncart = ce.data_uint64_t()
        cuest_check("Query Num Cart Functions",
            ce.cuestQuery(
                handle=cuest_handle,
                type=ce.CuestType.CUEST_AOSHELL,
                object=atom_shell,
                attribute=ce.CuestAOShellAttributes.CUEST_AOSHELL_NUM_CART,
                attributeValue=npure,
                )
            )
        print(f"\t\tNum. Cartesian Functions : {npure.value}")

# These workspace descriptors will be used in the resource query to figure out
# how much device/host RAM will be needed in persistent/temporary workspaces to
# build the AOBasis object.
persistent_workspace_descriptor = WorkspaceDescriptor()
temporary_workspace_descriptor = WorkspaceDescriptor()

aobasis_parameters = ce.cuestAOBasisParameters()
cuest_check('Create AO Basis Params',
    ce.cuestParametersCreate(
        parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
        outParameters=aobasis_parameters,
        )
    )

# We made a list of list of shells above, to easily determine which atom each
# belongs to.  Flatten that list here, for the AOBasis creation routines.
shells = [ s for atom_shells in shells for s in atom_shells ]

aobasis_handle = ce.cuestAOBasisHandle()

# Find out the resources needed to build the AOBasis
cuest_check('Create AOBasis Workspace Query',
    ce.cuestAOBasisCreateWorkspaceQuery(
        handle=cuest_handle,
        numAtoms=len(atom_symbols),
        numShellsPerAtom=n_shells_per_atom,
        shells=shells,
        parameters=aobasis_parameters,
        persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
        temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
        outBasis=aobasis_handle,
        )
    )

print("AO Basis Persistent sizes:", persistent_workspace_descriptor)
print("AO Basis Temporary sizes:", temporary_workspace_descriptor)

# Create the workspaces.  The functions called here are just an example of how
# the host/device workspaces may be created; it is the caller's responsibility
# to allocate and free these workspaces appropriately.
aobasis_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
aobasis_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

cuest_check('Create AOBasis',
    ce.cuestAOBasisCreate(
        handle=cuest_handle,
        numAtoms=len(atom_symbols),
        numShellsPerAtom=n_shells_per_atom,
        shells=shells,
        parameters=aobasis_parameters,
        persistentWorkspace=aobasis_persistent_workspace.pointer,
        temporaryWorkspace=aobasis_temporary_workspace.pointer,
        outBasis=aobasis_handle,
        )
    )

del aobasis_temporary_workspace

for i, shell in enumerate(shells):
    cuest_check('Destroy AO Shell {i+1}',
        ce.cuestAOShellDestroy(
            handle=shell,
            )
        )

cuest_check('Destroy AO Basis Params',
    ce.cuestParametersDestroy(
        parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
        parameters=aobasis_parameters,
        )
    )

# Query the basis for metadata and print the results

aobasis_is_pure = ce.data_int32_t()
cuest_check('Query AO Basis Is Pure',
    ce.cuestQuery(
        handle=cuest_handle,
        type=ce.CuestType.CUEST_AOBASIS,
        object=aobasis_handle,
        attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_IS_PURE,
        attributeValue=aobasis_is_pure,
        )
    )
print(f"AO Basis Is Pure        : {aobasis_is_pure.value}")

aobasis_num_atom = ce.data_uint64_t()
cuest_check('Query AO Basis Num Atom',
    ce.cuestQuery(
        handle=cuest_handle,
        type=ce.CuestType.CUEST_AOBASIS,
        object=aobasis_handle,
        attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_ATOM,
        attributeValue=aobasis_num_atom,
        )
    )
print(f"AO Basis Num Atom       : {aobasis_num_atom.value}")

aobasis_num_shell = ce.data_uint64_t()
cuest_check('Query AO Basis Num Shell',
    ce.cuestQuery(
        handle=cuest_handle,
        type=ce.CuestType.CUEST_AOBASIS,
        object=aobasis_handle,
        attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_SHELL,
        attributeValue=aobasis_num_shell,
        )
    )
print(f"AO Basis Num Shell      : {aobasis_num_shell.value}")

aobasis_num_ao = ce.data_uint64_t()
cuest_check('Query AO Basis Num AO',
    ce.cuestQuery(
        handle=cuest_handle,
        type=ce.CuestType.CUEST_AOBASIS,
        object=aobasis_handle,
        attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_AO,
        attributeValue=aobasis_num_ao,
        )
    )
print(f"AO Basis Num AO         : {aobasis_num_ao.value}")

aobasis_num_cart = ce.data_uint64_t()
cuest_check('Query AO Basis Num Cart',
    ce.cuestQuery(
        handle=cuest_handle,
        type=ce.CuestType.CUEST_AOBASIS,
        object=aobasis_handle,
        attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_CART,
        attributeValue=aobasis_num_cart,
        )
    )
print(f"AO Basis Num Cart       : {aobasis_num_cart.value}")

aobasis_num_primitive = ce.data_uint64_t()
cuest_check('Query AO Basis Num Primitive',
    ce.cuestQuery(
        handle=cuest_handle,
        type=ce.CuestType.CUEST_AOBASIS,
        object=aobasis_handle,
        attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_PRIMITIVE,
        attributeValue=aobasis_num_primitive,
        )
    )
print(f"AO Basis Num Primitive  : {aobasis_num_primitive.value}")

aobasis_max_L = ce.data_uint64_t()
cuest_check('Query AO Basis Num Primitive',
    ce.cuestQuery(
        handle=cuest_handle,
        type=ce.CuestType.CUEST_AOBASIS,
        object=aobasis_handle,
        attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_MAX_L,
        attributeValue=aobasis_max_L,
        )
    )
print(f"AO Basis Max. Ang. Mom. : {aobasis_max_L.value}")

cuest_check('Destroy AOBasis',
    ce.cuestAOBasisDestroy(
        handle=aobasis_handle,
        )
    )

del aobasis_persistent_workspace

# Delete the cuEST handle
cuest_check('Destroy Cuest Handle',
    ce.cuestDestroy(
        handle=cuest_handle,
        )
    )


# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cuest.bindings as ce
import numpy as np
import argparse
from pathlib import Path
import sys

# Inject the directory where the example helper utilities live
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from helpers.cuda_utils import cuda_malloc, cuda_free, cuda_memcpy_dtoh, WorkspaceDescriptor, Workspace
from helpers.utilities import normalize_shell_coefficients
from helpers.parsers import simple_xyz_parser, simple_gbs_parser, simple_ecp_parser

def run(
    *,
    xyz_filename,
    gbs_filename,
    ecp_filename,
    ):

    # => Parse User Input <= #

    symbols, xyzs, _ = simple_xyz_parser(
        filename=xyz_filename,
        to_bohr_scale_factor=1.0 / 0.52917720859,
        )

    shellinfo = simple_gbs_parser(
        filename=gbs_filename,
        symbols=symbols,
        )

    ecp_metadata = simple_ecp_parser(
        filename=ecp_filename,
        symbols=symbols,
        )

    num_atoms = len(symbols)
    sizeof_double = np.dtype('double').itemsize

    # => Set Up cuEST <= #

    def cuest_check(
        title,
        return_code
        ):
        " A simple helper function to call cuEST functions and check the return code. "

        if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
            raise RuntimeError(f"{title} failed with code {return_code}")

    # These are user-provided functions that are responsible for allocating
    # host and device workspace arrays.
    persistent_workspace_descriptor = WorkspaceDescriptor()
    temporary_workspace_descriptor = WorkspaceDescriptor()
    # Allow 2 GB of DRAM for ECP evaluation routines
    variable_buffer_size = WorkspaceDescriptor(
        host_buffer_size_in_bytes=0,
        device_buffer_size_in_bytes=2000000000,
        )

    # => cuEST Handle Setup <= #

    cuest_handle_parameters = ce.cuestHandleParameters()

    cuest_check('Create Handle Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            outParameters=cuest_handle_parameters,
            )
        )

    cuest_handle = ce.cuestHandle()
    cuest_check('Create Cuest Handle',
        ce.cuestCreate(
            parameters=cuest_handle_parameters,
            handle=cuest_handle,
            )
        )

    cuest_check('Destroy Handle Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            parameters=cuest_handle_parameters,
            )
        )

    # => Build Gaussian Shells From Basis Parser Info <= #

    aoshell_parameters = ce.cuestAOShellParameters()
    cuest_check('Create AO Shell Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
            outParameters=aoshell_parameters,
            )
        )

    shells = []
    n_shells_per_atom = []
    shell_count = 1
    for atom_shells in shellinfo:
        n_shells_per_atom.append(len(atom_shells))
        for shell in atom_shells:
            L = shell['L']
            exponents = shell['exponents']
            raw_coefficients = shell['coefficients']
            normalized_coefficients = normalize_shell_coefficients(
                coefficients=raw_coefficients,
                exponents=exponents,
                L=L,
                normalization=shell['normalization'],
                )
            aoshell_handle = ce.cuestAOShellHandle()
            cuest_check(f'Create AO Shell {shell_count}',
                ce.cuestAOShellCreate(
                    handle=cuest_handle,
                    isPure=shell['is_pure'],
                    L=L,
                    numPrimitive=len(exponents),
                    exponents=exponents,
                    coefficients=normalized_coefficients,
                    parameters=aoshell_parameters,
                    outShell=aoshell_handle)
                )
            shells.append(aoshell_handle)
            shell_count += 1

    cuest_check('Destroy AO Shell Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
            parameters=aoshell_parameters,
            )
        )

    # => Build Basis Set From Gaussian Shells <= #

    aobasis_parameters = ce.cuestAOBasisParameters()

    cuest_check('Create AO Basis Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
            outParameters=aobasis_parameters,
            )
        )

    aobasis_handle = ce.cuestAOBasisHandle()

    cuest_check('Create AOBasis Workspace Query',
        ce.cuestAOBasisCreateWorkspaceQuery(
            handle=cuest_handle,
            numAtoms=num_atoms,
            numShellsPerAtom=n_shells_per_atom,
            shells=shells,
            parameters=aobasis_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outBasis=aobasis_handle,
            )
        )

    aobasis_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    aobasis_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('Create AOBasis',
        ce.cuestAOBasisCreate(
            handle=cuest_handle,
            numAtoms=num_atoms,
            numShellsPerAtom=n_shells_per_atom,
            shells=shells,
            parameters=aobasis_parameters,
            persistentWorkspace=aobasis_persistent_workspace.pointer,
            temporaryWorkspace=aobasis_temporary_workspace.pointer,
            outBasis=aobasis_handle,
            )
        )

    del aobasis_temporary_workspace

    for i, shell in enumerate(shells):
        cuest_check(f'Destroy AO Shell {i+1}',
            ce.cuestAOShellDestroy(
                handle=shell,
                )
            )

    aobasis_num_ao = ce.data_uint64_t()
    cuest_check('Query AO Basis Num AO',
        ce.cuestQuery(
            handle=cuest_handle,
            type=ce.CuestType.CUEST_AOBASIS,
            object=aobasis_handle,
            attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_AO,
            attributeValue=aobasis_num_ao,
            )
        )
    nao = aobasis_num_ao.value

    cuest_check('Destroy AO Basis Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
            parameters=aobasis_parameters,
            )
        )

    # => Build ECP Shells and Atoms <= #

    # Create ECP shell parameters (shared for all ECP shell creation calls)
    ecpshell_parameters = ce.cuestECPShellParameters()
    cuest_check('Create ECP Shell Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPSHELL_PARAMETERS,
            outParameters=ecpshell_parameters,
            )
        )

    # Create ECP shells once per unique element type.  When multiple atoms of the
    # same element are present, their ECP data is identical so the shell handles
    # are shared across all atoms of that element.
    unique_symbol_shells = {}  # symbol -> {'top_shell': handle, 'shells': [handles]}
    for symbol, info in zip(symbols, ecp_metadata):
        if info is None or symbol in unique_symbol_shells:
            continue

        # Create the top shell (the max-L projector)
        top_shell_info = info['top_shell']
        top_shell_handle = ce.cuestECPShellHandle()
        cuest_check(f'Create ECP Top Shell for {symbol}',
            ce.cuestECPShellCreate(
                handle=cuest_handle,
                L=top_shell_info.L,
                numPrimitive=len(top_shell_info.ns),
                radialPowers=top_shell_info.ns,
                coefficients=top_shell_info.ws,
                exponents=top_shell_info.es,
                parameters=ecpshell_parameters,
                outECPShell=top_shell_handle,
                )
            )

        # Create the angular-momentum projected shells
        shell_handles = []
        for k, shell_info in enumerate(info['shells']):
            shell_handle = ce.cuestECPShellHandle()
            cuest_check(f'Create ECP Shell {k} for {symbol}',
                ce.cuestECPShellCreate(
                    handle=cuest_handle,
                    L=shell_info.L,
                    numPrimitive=len(shell_info.ns),
                    radialPowers=shell_info.ns,
                    coefficients=shell_info.ws,
                    exponents=shell_info.es,
                    parameters=ecpshell_parameters,
                    outECPShell=shell_handle,
                    )
                )
            shell_handles.append(shell_handle)

        unique_symbol_shells[symbol] = {
            'top_shell': top_shell_handle,
            'shells': shell_handles,
            }

    cuest_check('Destroy ECP Shell Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPSHELL_PARAMETERS,
            parameters=ecpshell_parameters,
            )
        )

    # Create ECP atom parameters (shared for all ECP atom creation calls)
    ecpatom_parameters = ce.cuestECPAtomParameters()
    cuest_check('Create ECP Atom Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPATOM_PARAMETERS,
            outParameters=ecpatom_parameters,
            )
        )

    # Create one ECP atom handle per active atom (atoms that carry an ECP)
    ecp_atom_indices = []   # indices into the full atom list
    ecp_atom_handles = []   # corresponding cuestECPAtomHandle objects
    for atom_index, (symbol, info) in enumerate(zip(symbols, ecp_metadata)):
        if info is None:
            continue
        shell_data = unique_symbol_shells[symbol]
        ecpatom_handle = ce.cuestECPAtomHandle()
        cuest_check(f'Create ECP Atom {atom_index}',
            ce.cuestECPAtomCreate(
                handle=cuest_handle,
                numElectrons=info['nelectron'],
                numShells=len(shell_data['shells']),
                shells=shell_data['shells'],
                topShell=shell_data['top_shell'],
                parameters=ecpatom_parameters,
                outECPAtom=ecpatom_handle,
                )
            )
        ecp_atom_indices.append(atom_index)
        ecp_atom_handles.append(ecpatom_handle)

    # ECP shell handles are no longer needed once the atoms have been created
    for symbol, shell_data in unique_symbol_shells.items():
        cuest_check(f'Destroy ECP Top Shell for {symbol}',
            ce.cuestECPShellDestroy(
                handle=shell_data['top_shell'],
                )
            )
        for k, shell_handle in enumerate(shell_data['shells']):
            cuest_check(f'Destroy ECP Shell {k} for {symbol}',
                ce.cuestECPShellDestroy(
                    handle=shell_handle,
                    )
                )

    cuest_check('Destroy ECP Atom Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPATOM_PARAMETERS,
            parameters=ecpatom_parameters,
            )
        )

    # => Build ECP Integral Plan <= #

    ecpintplan_parameters = ce.cuestECPIntPlanParameters()
    cuest_check('Create ECPIntPlan Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPINTPLAN_PARAMETERS,
            outParameters=ecpintplan_parameters,
            )
        )

    ecpintplan_handle = ce.cuestECPIntPlanHandle()

    cuest_check('Create ECPIntPlan Workspace Query',
        ce.cuestECPIntPlanCreateWorkspaceQuery(
            handle=cuest_handle,
            basis=aobasis_handle,
            xyz=xyzs,
            numECPAtoms=len(ecp_atom_indices),
            activeIndices=ecp_atom_indices,
            activeAtoms=ecp_atom_handles,
            parameters=ecpintplan_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outPlan=ecpintplan_handle,
            )
        )

    ecpintplan_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    ecpintplan_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('Create ECPIntPlan',
        ce.cuestECPIntPlanCreate(
            handle=cuest_handle,
            basis=aobasis_handle,
            xyz=xyzs,
            numECPAtoms=len(ecp_atom_indices),
            activeIndices=ecp_atom_indices,
            activeAtoms=ecp_atom_handles,
            parameters=ecpintplan_parameters,
            persistentWorkspace=ecpintplan_persistent_workspace.pointer,
            temporaryWorkspace=ecpintplan_temporary_workspace.pointer,
            outPlan=ecpintplan_handle,
            )
        )

    del ecpintplan_temporary_workspace

    for i, ecpatom_handle in enumerate(ecp_atom_handles):
        cuest_check(f'Destroy ECP Atom {i}',
            ce.cuestECPAtomDestroy(
                handle=ecpatom_handle,
                )
            )

    cuest_check('Destroy ECPIntPlan Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPINTPLAN_PARAMETERS,
            parameters=ecpintplan_parameters,
            )
        )

    # => Compute ECP Integrals <= #

    ecpcompute_parameters = ce.cuestECPComputeParameters()
    cuest_check('Create ECPCompute Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPCOMPUTE_PARAMETERS,
            outParameters=ecpcompute_parameters,
            )
        )

    ecp_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nao**2
        )
    ecp_device_handle = ce.Pointer(ecp_device_pointer)

    # Find the temporary workspace requirements for the ECP integral evaluation
    cuest_check('ECPCompute Workspace Query',
        ce.cuestECPComputeWorkspaceQuery(
            handle=cuest_handle,
            plan=ecpintplan_handle,
            parameters=ecpcompute_parameters,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outECPMatrix=ecp_device_handle,
            )
        )

    ecp_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    # Compute the ECP matrix (nao x nao)
    cuest_check('ECPCompute',
        ce.cuestECPCompute(
            handle=cuest_handle,
            plan=ecpintplan_handle,
            parameters=ecpcompute_parameters,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspace=ecp_temporary_workspace.pointer,
            outECPMatrix=ecp_device_handle,
            )
        )

    del ecp_temporary_workspace

    cuest_check('Destroy ECPCompute Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPCOMPUTE_PARAMETERS,
            parameters=ecpcompute_parameters,
            )
        )

    ecp_host_array = np.empty(
        nao * nao,
        dtype=np.double
        )

    cuda_memcpy_dtoh(
        host_pointer=ecp_host_array.ctypes.data,
        device_pointer=ecp_device_pointer,
        size_in_bytes=sizeof_double * nao**2,
        )

    cuda_free(array=ecp_device_pointer)

    # => Cleanup <= #

    cuest_check('Destroy ECPIntPlan',
        ce.cuestECPIntPlanDestroy(
            handle=ecpintplan_handle,
            )
        )
    del ecpintplan_persistent_workspace

    cuest_check('Destroy AOBasis',
        ce.cuestAOBasisDestroy(
            handle=aobasis_handle,
            )
        )
    del aobasis_persistent_workspace

    cuest_check('Destroy Cuest Handle',
        ce.cuestDestroy(
            handle=cuest_handle,
            )
        )

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Compute ECP integrals using XYZ, GBS, and ECP files."
        )
    parser.add_argument(
        "xyz_filename",
        type=str,
        help="Molecular coordinates (Angstrom) in basic XYZ format."
        )
    parser.add_argument(
        "gbs_filename",
        type=str,
        help="Orbital basis set in G94/Psi4 GBS format."
        )
    parser.add_argument(
        "ecp_filename",
        type=str,
        help="Effective core potential basis set file."
        )

    args = parser.parse_args()

    run(
        xyz_filename=args.xyz_filename,
        gbs_filename=args.gbs_filename,
        ecp_filename=args.ecp_filename,
        )


# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cuest.bindings as ce
import numpy as np
import argparse
from pathlib import Path
import sys

# Inject the directory where the example helper utilities live
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from helpers.cuda_utils import cuda_malloc, cuda_free, cuda_memcpy_dtoh, WorkspaceDescriptor, Workspace
from helpers.utilities import normalize_shell_coefficients
from helpers.parsers import simple_xyz_parser, simple_gbs_parser, simple_ecp_parser

def run(
    *,
    xyz_filename,
    gbs_filename,
    ecp_filename,
    ):

    # => Parse User Input <= #

    symbols, xyzs, _ = simple_xyz_parser(
        filename=xyz_filename,
        to_bohr_scale_factor=1.0 / 0.52917720859,
        )

    shellinfo = simple_gbs_parser(
        filename=gbs_filename,
        symbols=symbols,
        )

    ecp_metadata = simple_ecp_parser(
        filename=ecp_filename,
        symbols=symbols,
        )

    num_atoms = len(symbols)
    sizeof_double = np.dtype('double').itemsize

    # => Set Up cuEST <= #

    def cuest_check(
        title,
        return_code
        ):
        " A simple helper function to call cuEST functions and check the return code. "

        if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
            raise RuntimeError(f"{title} failed with code {return_code}")

    # These are user-provided functions that are responsible for allocating
    # host and device workspace arrays.
    persistent_workspace_descriptor = WorkspaceDescriptor()
    temporary_workspace_descriptor = WorkspaceDescriptor()
    # Allow 2 GB of DRAM for ECP evaluation routines
    variable_buffer_size = WorkspaceDescriptor(
        host_buffer_size_in_bytes=0,
        device_buffer_size_in_bytes=2000000000,
        )

    # => cuEST Handle Setup <= #

    cuest_handle_parameters = ce.cuestHandleParameters()

    cuest_check('Create Handle Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            outParameters=cuest_handle_parameters,
            )
        )

    cuest_handle = ce.cuestHandle()
    cuest_check('Create Cuest Handle',
        ce.cuestCreate(
            parameters=cuest_handle_parameters,
            handle=cuest_handle,
            )
        )

    cuest_check('Destroy Handle Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            parameters=cuest_handle_parameters,
            )
        )

    # => Build Gaussian Shells From Basis Parser Info <= #

    aoshell_parameters = ce.cuestAOShellParameters()
    cuest_check('Create AO Shell Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
            outParameters=aoshell_parameters,
            )
        )

    shells = []
    n_shells_per_atom = []
    shell_count = 1
    for atom_shells in shellinfo:
        n_shells_per_atom.append(len(atom_shells))
        for shell in atom_shells:
            L = shell['L']
            exponents = shell['exponents']
            raw_coefficients = shell['coefficients']
            normalized_coefficients = normalize_shell_coefficients(
                coefficients=raw_coefficients,
                exponents=exponents,
                L=L,
                normalization=shell['normalization'],
                )
            aoshell_handle = ce.cuestAOShellHandle()
            cuest_check(f'Create AO Shell {shell_count}',
                ce.cuestAOShellCreate(
                    handle=cuest_handle,
                    isPure=shell['is_pure'],
                    L=L,
                    numPrimitive=len(exponents),
                    exponents=exponents,
                    coefficients=normalized_coefficients,
                    parameters=aoshell_parameters,
                    outShell=aoshell_handle)
                )
            shells.append(aoshell_handle)
            shell_count += 1

    cuest_check('Destroy AO Shell Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_AOSHELL_PARAMETERS,
            parameters=aoshell_parameters,
            )
        )

    # => Build Basis Set From Gaussian Shells <= #

    aobasis_parameters = ce.cuestAOBasisParameters()

    cuest_check('Create AO Basis Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
            outParameters=aobasis_parameters,
            )
        )

    aobasis_handle = ce.cuestAOBasisHandle()

    cuest_check('Create AOBasis Workspace Query',
        ce.cuestAOBasisCreateWorkspaceQuery(
            handle=cuest_handle,
            numAtoms=num_atoms,
            numShellsPerAtom=n_shells_per_atom,
            shells=shells,
            parameters=aobasis_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outBasis=aobasis_handle,
            )
        )

    aobasis_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    aobasis_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('Create AOBasis',
        ce.cuestAOBasisCreate(
            handle=cuest_handle,
            numAtoms=num_atoms,
            numShellsPerAtom=n_shells_per_atom,
            shells=shells,
            parameters=aobasis_parameters,
            persistentWorkspace=aobasis_persistent_workspace.pointer,
            temporaryWorkspace=aobasis_temporary_workspace.pointer,
            outBasis=aobasis_handle,
            )
        )

    del aobasis_temporary_workspace

    for i, shell in enumerate(shells):
        cuest_check(f'Destroy AO Shell {i+1}',
            ce.cuestAOShellDestroy(
                handle=shell,
                )
            )

    aobasis_num_ao = ce.data_uint64_t()
    cuest_check('Query AO Basis Num AO',
        ce.cuestQuery(
            handle=cuest_handle,
            type=ce.CuestType.CUEST_AOBASIS,
            object=aobasis_handle,
            attribute=ce.CuestAOBasisAttributes.CUEST_AOBASIS_NUM_AO,
            attributeValue=aobasis_num_ao,
            )
        )
    nao = aobasis_num_ao.value

    cuest_check('Destroy AO Basis Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_AOBASIS_PARAMETERS,
            parameters=aobasis_parameters,
            )
        )

    # => Build ECP Shells and Atoms <= #

    # Create ECP shell parameters (shared for all ECP shell creation calls)
    ecpshell_parameters = ce.cuestECPShellParameters()
    cuest_check('Create ECP Shell Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPSHELL_PARAMETERS,
            outParameters=ecpshell_parameters,
            )
        )

    # Create ECP shells once per unique element type.  When multiple atoms of the
    # same element are present, their ECP data is identical so the shell handles
    # are shared across all atoms of that element.
    unique_symbol_shells = {}  # symbol -> {'top_shell': handle, 'shells': [handles]}
    for symbol, info in zip(symbols, ecp_metadata):
        if info is None or symbol in unique_symbol_shells:
            continue

        # Create the top shell (the max-L projector)
        top_shell_info = info['top_shell']
        top_shell_handle = ce.cuestECPShellHandle()
        cuest_check(f'Create ECP Top Shell for {symbol}',
            ce.cuestECPShellCreate(
                handle=cuest_handle,
                L=top_shell_info.L,
                numPrimitive=len(top_shell_info.ns),
                radialPowers=top_shell_info.ns,
                coefficients=top_shell_info.ws,
                exponents=top_shell_info.es,
                parameters=ecpshell_parameters,
                outECPShell=top_shell_handle,
                )
            )

        # Create the angular-momentum projected shells
        shell_handles = []
        for k, shell_info in enumerate(info['shells']):
            shell_handle = ce.cuestECPShellHandle()
            cuest_check(f'Create ECP Shell {k} for {symbol}',
                ce.cuestECPShellCreate(
                    handle=cuest_handle,
                    L=shell_info.L,
                    numPrimitive=len(shell_info.ns),
                    radialPowers=shell_info.ns,
                    coefficients=shell_info.ws,
                    exponents=shell_info.es,
                    parameters=ecpshell_parameters,
                    outECPShell=shell_handle,
                    )
                )
            shell_handles.append(shell_handle)

        unique_symbol_shells[symbol] = {
            'top_shell': top_shell_handle,
            'shells': shell_handles,
            }

    cuest_check('Destroy ECP Shell Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPSHELL_PARAMETERS,
            parameters=ecpshell_parameters,
            )
        )

    # Create ECP atom parameters (shared for all ECP atom creation calls)
    ecpatom_parameters = ce.cuestECPAtomParameters()
    cuest_check('Create ECP Atom Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPATOM_PARAMETERS,
            outParameters=ecpatom_parameters,
            )
        )

    # Create one ECP atom handle per active atom (atoms that carry an ECP)
    ecp_atom_indices = []   # indices into the full atom list
    ecp_atom_handles = []   # corresponding cuestECPAtomHandle objects
    for atom_index, (symbol, info) in enumerate(zip(symbols, ecp_metadata)):
        if info is None:
            continue
        shell_data = unique_symbol_shells[symbol]
        ecpatom_handle = ce.cuestECPAtomHandle()
        cuest_check(f'Create ECP Atom {atom_index}',
            ce.cuestECPAtomCreate(
                handle=cuest_handle,
                numElectrons=info['nelectron'],
                numShells=len(shell_data['shells']),
                shells=shell_data['shells'],
                topShell=shell_data['top_shell'],
                parameters=ecpatom_parameters,
                outECPAtom=ecpatom_handle,
                )
            )
        ecp_atom_indices.append(atom_index)
        ecp_atom_handles.append(ecpatom_handle)

    # ECP shell handles are no longer needed once the atoms have been created
    for symbol, shell_data in unique_symbol_shells.items():
        cuest_check(f'Destroy ECP Top Shell for {symbol}',
            ce.cuestECPShellDestroy(
                handle=shell_data['top_shell'],
                )
            )
        for k, shell_handle in enumerate(shell_data['shells']):
            cuest_check(f'Destroy ECP Shell {k} for {symbol}',
                ce.cuestECPShellDestroy(
                    handle=shell_handle,
                    )
                )

    cuest_check('Destroy ECP Atom Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPATOM_PARAMETERS,
            parameters=ecpatom_parameters,
            )
        )

    # => Build ECP Integral Plan <= #

    ecpintplan_parameters = ce.cuestECPIntPlanParameters()
    cuest_check('Create ECPIntPlan Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPINTPLAN_PARAMETERS,
            outParameters=ecpintplan_parameters,
            )
        )

    ecpintplan_handle = ce.cuestECPIntPlanHandle()

    cuest_check('Create ECPIntPlan Workspace Query',
        ce.cuestECPIntPlanCreateWorkspaceQuery(
            handle=cuest_handle,
            basis=aobasis_handle,
            xyz=xyzs,
            numECPAtoms=len(ecp_atom_indices),
            activeIndices=ecp_atom_indices,
            activeAtoms=ecp_atom_handles,
            parameters=ecpintplan_parameters,
            persistentWorkspaceDescriptor=persistent_workspace_descriptor.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outPlan=ecpintplan_handle,
            )
        )

    ecpintplan_persistent_workspace = Workspace(workspaceDescriptor=persistent_workspace_descriptor)
    ecpintplan_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    cuest_check('Create ECPIntPlan',
        ce.cuestECPIntPlanCreate(
            handle=cuest_handle,
            basis=aobasis_handle,
            xyz=xyzs,
            numECPAtoms=len(ecp_atom_indices),
            activeIndices=ecp_atom_indices,
            activeAtoms=ecp_atom_handles,
            parameters=ecpintplan_parameters,
            persistentWorkspace=ecpintplan_persistent_workspace.pointer,
            temporaryWorkspace=ecpintplan_temporary_workspace.pointer,
            outPlan=ecpintplan_handle,
            )
        )

    del ecpintplan_temporary_workspace

    for i, ecpatom_handle in enumerate(ecp_atom_handles):
        cuest_check(f'Destroy ECP Atom {i}',
            ce.cuestECPAtomDestroy(
                handle=ecpatom_handle,
                )
            )

    cuest_check('Destroy ECPIntPlan Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPINTPLAN_PARAMETERS,
            parameters=ecpintplan_parameters,
            )
        )

    # => Compute ECP Integrals <= #

    ecpcompute_parameters = ce.cuestECPComputeParameters()
    cuest_check('Create ECPCompute Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPCOMPUTE_PARAMETERS,
            outParameters=ecpcompute_parameters,
            )
        )

    ecp_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nao**2
        )
    ecp_device_handle = ce.Pointer(ecp_device_pointer)

    # Find the temporary workspace requirements for the ECP integral evaluation
    cuest_check('ECPCompute Workspace Query',
        ce.cuestECPComputeWorkspaceQuery(
            handle=cuest_handle,
            plan=ecpintplan_handle,
            parameters=ecpcompute_parameters,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            outECPMatrix=ecp_device_handle,
            )
        )

    ecp_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    # Compute the ECP matrix (nao x nao)
    cuest_check('ECPCompute',
        ce.cuestECPCompute(
            handle=cuest_handle,
            plan=ecpintplan_handle,
            parameters=ecpcompute_parameters,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspace=ecp_temporary_workspace.pointer,
            outECPMatrix=ecp_device_handle,
            )
        )

    del ecp_temporary_workspace

    cuest_check('Destroy ECPCompute Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPCOMPUTE_PARAMETERS,
            parameters=ecpcompute_parameters,
            )
        )

    ecp_host_array = np.empty(
        nao * nao,
        dtype=np.double
        )

    cuda_memcpy_dtoh(
        host_pointer=ecp_host_array.ctypes.data,
        device_pointer=ecp_device_pointer,
        size_in_bytes=sizeof_double * nao**2,
        )

    cuda_free(array=ecp_device_pointer)

    # => Cleanup <= #

    cuest_check('Destroy ECPIntPlan',
        ce.cuestECPIntPlanDestroy(
            handle=ecpintplan_handle,
            )
        )
    del ecpintplan_persistent_workspace

    cuest_check('Destroy AOBasis',
        ce.cuestAOBasisDestroy(
            handle=aobasis_handle,
            )
        )
    del aobasis_persistent_workspace

    cuest_check('Destroy Cuest Handle',
        ce.cuestDestroy(
            handle=cuest_handle,
            )
        )

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Compute ECP integrals using XYZ, GBS, and ECP files."
        )
    parser.add_argument(
        "xyz_filename",
        type=str,
        help="Molecular coordinates (Angstrom) in basic XYZ format."
        )
    parser.add_argument(
        "gbs_filename",
        type=str,
        help="Orbital basis set in G94/Psi4 GBS format."
        )
    parser.add_argument(
        "ecp_filename",
        type=str,
        help="Effective core potential basis set file."
        )


    symbols, xyzs, _ = simple_xyz_parser(
        filename=xyz_filename,
        to_bohr_scale_factor=1.0 / 0.52917720859,
        )

    shellinfo = simple_gbs_parser(
        filename=gbs_filename,
        symbols=symbols,
        )

    ecp_metadata = simple_ecp_parser(
        filename=ecp_filename,
        symbols=symbols,
        )

    num_atoms = len(symbols)
    sizeof_double = np.dtype('double').itemsize

    # => Set Up cuEST <= #

    def cuest_check(
        title,
        return_code
        ):
        " A simple helper function to call cuEST functions and check the return code. "

        if return_code != ce.CuestStatus.CUEST_STATUS_SUCCESS:
            raise RuntimeError(f"{title} failed with code {return_code}")

    # These are user-provided functions that are responsible for allocating
    # host and device workspace arrays.
    persistent_workspace_descriptor = WorkspaceDescriptor()
    temporary_workspace_descriptor = WorkspaceDescriptor()
    # Allow 2 GB of DRAM for ECP evaluation routines
    variable_buffer_size = WorkspaceDescriptor(
        host_buffer_size_in_bytes=0,
        device_buffer_size_in_bytes=2000000000,
        )

    # => cuEST Handle Setup <= #

    cuest_handle_parameters = ce.cuestHandleParameters()

    cuest_check('Create Handle Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
            outParameters=cuest_handle_parameters,
            )