# 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, simple_ecp_parser

def run(
    *,
    xyz_filename,
    gbs_filename,
    ecp_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,
        )

    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 Gradients <= #

    # The ECP derivative routines contract the integral derivatives against a density
    # matrix and accumulate the result into a [num_atoms, 3] gradient array.  A real
    # calculation would supply the SCF density matrix here; we use a random symmetric
    # matrix as a stand-in.
    np.random.seed(0)
    density_host = np.random.normal(size=(nao, nao)).astype(np.double)
    density_host = 0.5 * (density_host + density_host.T)  # Symmetrize

    density_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nao**2
        )
    density_device_handle = ce.Pointer(density_device_pointer)
    cuda_memcpy_htod(
        device_pointer=density_device_pointer,
        host_pointer=density_host.ctypes.data,
        size_in_bytes=sizeof_double * nao**2,
        )

    gradient_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * num_atoms * 3
        )
    gradient_device_handle = ce.Pointer(gradient_device_pointer)

    ecpderivcompute_parameters = ce.cuestECPDerivativeComputeParameters()
    cuest_check('Create ECPDerivativeCompute Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPDERIVATIVECOMPUTE_PARAMETERS,
            outParameters=ecpderivcompute_parameters,
            )
        )

    # Find the temporary workspace requirements for the ECP gradient evaluation
    cuest_check('ECPDerivativeCompute Workspace Query',
        ce.cuestECPDerivativeComputeWorkspaceQuery(
            handle=cuest_handle,
            plan=ecpintplan_handle,
            parameters=ecpderivcompute_parameters,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            densityMatrix=density_device_handle,
            outGradient=gradient_device_handle,
            )
        )

    ecp_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    # Compute dECP/dR contracted with the density matrix; result is [num_atoms, 3]
    cuest_check('ECPDerivativeCompute',
        ce.cuestECPDerivativeCompute(
            handle=cuest_handle,
            plan=ecpintplan_handle,
            parameters=ecpderivcompute_parameters,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspace=ecp_temporary_workspace.pointer,
            densityMatrix=density_device_handle,
            outGradient=gradient_device_handle,
            )
        )

    del ecp_temporary_workspace

    cuest_check('Destroy ECPDerivativeCompute Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPDERIVATIVECOMPUTE_PARAMETERS,
            parameters=ecpderivcompute_parameters,
            )
        )

    gradient_host_array = np.empty(
        num_atoms * 3,
        dtype=np.double
        )

    cuda_memcpy_dtoh(
        host_pointer=gradient_host_array.ctypes.data,
        device_pointer=gradient_device_pointer,
        size_in_bytes=sizeof_double * num_atoms * 3,
        )

    cuda_free(array=density_device_pointer)
    cuda_free(array=gradient_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,
            )
        )

    return gradient_host_array.reshape(num_atoms, 3)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Compute ECP integral gradients contracted with a density matrix."
        )
    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,
        )


run

# 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, simple_ecp_parser

def run(
    *,
    xyz_filename,
    gbs_filename,
    ecp_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,
        )

    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 Gradients <= #

    # The ECP derivative routines contract the integral derivatives against a density
    # matrix and accumulate the result into a [num_atoms, 3] gradient array.  A real
    # calculation would supply the SCF density matrix here; we use a random symmetric
    # matrix as a stand-in.
    np.random.seed(0)
    density_host = np.random.normal(size=(nao, nao)).astype(np.double)
    density_host = 0.5 * (density_host + density_host.T)  # Symmetrize

    density_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * nao**2
        )
    density_device_handle = ce.Pointer(density_device_pointer)
    cuda_memcpy_htod(
        device_pointer=density_device_pointer,
        host_pointer=density_host.ctypes.data,
        size_in_bytes=sizeof_double * nao**2,
        )

    gradient_device_pointer = cuda_malloc(
        size_in_bytes=sizeof_double * num_atoms * 3
        )
    gradient_device_handle = ce.Pointer(gradient_device_pointer)

    ecpderivcompute_parameters = ce.cuestECPDerivativeComputeParameters()
    cuest_check('Create ECPDerivativeCompute Params',
        ce.cuestParametersCreate(
            parametersType=ce.CuestParametersType.CUEST_ECPDERIVATIVECOMPUTE_PARAMETERS,
            outParameters=ecpderivcompute_parameters,
            )
        )

    # Find the temporary workspace requirements for the ECP gradient evaluation
    cuest_check('ECPDerivativeCompute Workspace Query',
        ce.cuestECPDerivativeComputeWorkspaceQuery(
            handle=cuest_handle,
            plan=ecpintplan_handle,
            parameters=ecpderivcompute_parameters,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspaceDescriptor=temporary_workspace_descriptor.pointer,
            densityMatrix=density_device_handle,
            outGradient=gradient_device_handle,
            )
        )

    ecp_temporary_workspace = Workspace(workspaceDescriptor=temporary_workspace_descriptor)

    # Compute dECP/dR contracted with the density matrix; result is [num_atoms, 3]
    cuest_check('ECPDerivativeCompute',
        ce.cuestECPDerivativeCompute(
            handle=cuest_handle,
            plan=ecpintplan_handle,
            parameters=ecpderivcompute_parameters,
            variableBufferSize=variable_buffer_size.pointer,
            temporaryWorkspace=ecp_temporary_workspace.pointer,
            densityMatrix=density_device_handle,
            outGradient=gradient_device_handle,
            )
        )

    del ecp_temporary_workspace

    cuest_check('Destroy ECPDerivativeCompute Params',
        ce.cuestParametersDestroy(
            parametersType=ce.CuestParametersType.CUEST_ECPDERIVATIVECOMPUTE_PARAMETERS,
            parameters=ecpderivcompute_parameters,
            )
        )

    gradient_host_array = np.empty(
        num_atoms * 3,
        dtype=np.double
        )

    cuda_memcpy_dtoh(
        host_pointer=gradient_host_array.ctypes.data,
        device_pointer=gradient_device_pointer,
        size_in_bytes=sizeof_double * num_atoms * 3,
        )

    cuda_free(array=density_device_pointer)
    cuda_free(array=gradient_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,
            )
        )

    return gradient_host_array.reshape(num_atoms, 3)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Compute ECP integral gradients contracted with a density matrix."
        )
    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, Zs = 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,
            )


# 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

def cuest_check(
    title,
    return_code
    ):

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

# Create default HandleParameters object to parameterize the cuEST handle.
cuest_handle_parameters = ce.cuestHandleParameters()

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

# Make a handle to hold a C uint64_t type, and populate it with the maximum number
# of Gauss-Hermite quadrature points supported by by the cuEST handle by default.
maxgh_handle = ce.data_uint64_t()
cuest_check('Query Handle Max Num GH',
    ce.cuestParametersQuery(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_MAX_GAUSS_HERMITE,
        attributeValue=maxgh_handle,
        )
    )
print('Handle Max Num GH', maxgh_handle.value)

# Similarly, query for the maximum angular momentum supported for
# Cartesian->Spherical transforms in this cuEST handle instance.
maxl_handle = ce.data_uint64_t()
cuest_check('Query Handle Max L',
    ce.cuestParametersQuery(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_MAX_L_SOLID_HARMONIC,
        attributeValue=maxl_handle,
        )
    )
print('Handle Max L', maxl_handle.value)

# ce.cuestParametersConfigure could be called here to override the default
# values if required

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

# Once the handle itself has been created, its parameters can be destroyed
cuest_check('Destroy Handle Params',
    ce.cuestParametersDestroy(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        )
    )

# The handle is ready to be used for calculations

# Destroying the handle will also destoy CUDA stream, cuBLAS, and cuSolver
# handles held by 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

def cuest_check(
    title,
    return_code
    ):

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

# Create default HandleParameters object to parameterize the cuEST handle.
cuest_handle_parameters = ce.cuestHandleParameters()

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

# Make a handle to hold a C uint64_t type, and populate it with the maximum number
# of Gauss-Hermite quadrature points supported by by the cuEST handle by default.
maxgh_handle = ce.data_uint64_t()
cuest_check('Query Handle Max Num GH',
    ce.cuestParametersQuery(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_MAX_GAUSS_HERMITE,


# 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 make_stream,\
                               make_cublas_handle, free_cublas_handle,\
                               make_cusolver_handle, free_cusolver_handle,\
                               set_cublas_stream, set_cusolver_stream

def cuest_check(
    title,
    return_code
    ):

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

stream = make_stream()
cublas_c_handle = make_cublas_handle()
cusolver_c_handle = make_cusolver_handle()

set_cublas_stream(
    handle=cublas_c_handle,
    stream=stream.__int__(),
    )

set_cusolver_stream(
    handle=cusolver_c_handle,
    stream=stream.__int__(),
    )

# Create default HandleParameters object to parameterize the cuEST handle.
cuest_handle_parameters = ce.cuestHandleParameters()

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

# => Configure the cuEST handle paramters <= #

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

cublas_handle = ce.data_cublasHandle_t(cublas_c_handle.value)
cuest_check('Configure Handle cuBLAS',
    ce.cuestParametersConfigure(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_CUBLAS,
        attributeValue=cublas_handle,
        )
    )

cusolver_handle = ce.data_cusolverDnHandle_t(cusolver_c_handle.value)
cuest_check('Configure Handle cuSolver',
    ce.cuestParametersConfigure(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_CUSOLVER,
        attributeValue=cusolver_handle,
        )
    )

# => Make cuEST Handle <= #

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

# Once the handle itself has been created, its parameters can be destroyed
cuest_check('Destroy Handle Params',
    ce.cuestParametersDestroy(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        parameters=cuest_handle_parameters,
        )
    )

# The handle is ready to be used for calculations

# Destroying the handle will also destoy CUDA stream, cuBLAS, and cuSolver
# handles held by the cuEST handle.
cuest_check('Destroy Cuest Handle',
    ce.cuestDestroy(
        handle=cuest_handle,
        )
    )

free_cublas_handle(
    handle=cublas_c_handle,
    )
free_cusolver_handle(
    handle=cusolver_c_handle,
    )
# A key difference between using cuda.bindings to wrap CUDA (as the stream is
# in this example) vs. using ctypes (used for cuBLAS and cuSolver) is that the
# former takes care of memory management automatically, so no call to free the
# stream is needed.  We can optionally free it as follows:
del stream


#
# 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 make_stream,\
                               make_cublas_handle, free_cublas_handle,\
                               make_cusolver_handle, free_cusolver_handle,\
                               set_cublas_stream, set_cusolver_stream

def cuest_check(
    title,
    return_code
    ):

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

stream = make_stream()
cublas_c_handle = make_cublas_handle()
cusolver_c_handle = make_cusolver_handle()

set_cublas_stream(
    handle=cublas_c_handle,
    stream=stream.__int__(),
    )

set_cusolver_stream(
    handle=cusolver_c_handle,
    stream=stream.__int__(),
    )

# Create default HandleParameters object to parameterize the cuEST handle.
cuest_handle_parameters = ce.cuestHandleParameters()

cuest_check('Create Handle Params',


# 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

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('Create Handle Params',
    ce.cuestParametersCreate(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        outParameters=cuest_handle_parameters,
        )
    )

# Create a list of 2 cuEST handles.  Each one will hold its own unique
# resources Stream, cuBLAS, cuSolver.
cuest_handles = []
for _ in range(2):
    cuest_handle = ce.cuestHandle()
    cuest_check('Create Cuest Handle',
        ce.cuestCreate(
            parameters=cuest_handle_parameters,
            handle=cuest_handle,
            )
        )
    cuest_handles.append(cuest_handle)

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

# The cuEST handles are ready to be used for calculations

# Free both of the handles when finished, to free up resources
for cuest_handle in cuest_handles:
    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

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('Create Handle Params',
    ce.cuestParametersCreate(
        parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
        outParameters=cuest_handle_parameters,
        )
    )

# Create a list of 2 cuEST handles.  Each one will hold its own unique
# resources Stream, cuBLAS, cuSolver.
cuest_handles = []
for _ in range(2):
    cuest_handle = ce.cuestHandle()
    cuest_check('Create Cuest Handle',
        ce.cuestCreate(
            parameters=cuest_handle_parameters,
            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 re
from collections import namedtuple

def simple_xyz_parser(
    *,
    filename,
    to_bohr_scale_factor,
    ):

    elements = [
        "X",  "H",  "HE", "LI", "BE", "B",  "C",  "N",  "O",  "F",  "NE",
        "NA", "MG", "AL", "SI", "P",  "S",  "CL", "AR", "K",  "CA",
        "SC", "TI", "V",  "CR", "MN", "FE", "CO", "NI", "CU", "ZN",
        "GA", "GE", "AS", "SE", "BR", "KR", "RB", "SR", "Y",  "ZR",
        "NB", "MO", "TC", "RU", "RH", "PD", "AG", "CD", "IN", "SN",
        "SB", "TE", "I",  "XE", "CS", "BA", "LA", "CE", "PR", "ND",
        "PM", "SM", "EU", "GD", "TB", "DY", "HO", "ER", "TM", "YB",
        "LU", "HF", "TA", "W",  "RE", "OS", "IR", "PT", "AU", "HG",
        "TL", "PB", "BI", "PO", "AT", "RN", "FR", "RA", "AC", "TH",
        "PA", "U",  "NP", "PU", "AM", "CM", "BK", "CF", "ES", "FM",
        "MD", "NO", "LR", "RF", "DB", "SG", "BH", "HS", "MT", "DS",
        "RG", "CN", "NH", "FL", "MC", "LV", "TS", "OG"
        ]
    symbols = []
    xyzs = []
    Zs = []
    with open(filename) as fp:
        atom_count = int(fp.readline())
        comment = fp.readline()
        for _ in range(atom_count):
            parts = fp.readline().split()
            symbol = parts[0].upper()
            # The potential integrals routine compute electron-electron interaction
            # integrals; to make this an electron-nucleus interaction we scale Z here
            Z = -1.0 * elements.index(symbol)
            x = float(parts[1]) * to_bohr_scale_factor
            y = float(parts[2]) * to_bohr_scale_factor
            z = float(parts[3]) * to_bohr_scale_factor
            xyzs.extend([x, y, z])
            symbols.append(symbol)
            Zs.append(Z)
    return symbols, xyzs, Zs

ECPShellInfo = namedtuple('ECPShellInfo', ['L', 'ns', 'ws', 'es'])

def _ecp_shell_info_from_ecp_lines(
    *,
    lines,
    ):
    _shell_data_ = {
        'S': 0,
        'P': 1,
        'D': 2,
        'F': 3,
        'G': 4,
        'H': 5,
        'I': 6,
        'K': 7,
        'L': 8,
        'M': 9,
        'N': 10,
        'O': 11,
        'Q': 12,
        'R': 13,
        'T': 14,
        'U': 15,
        'V': 16,
        'W': 17,
        'X': 18,
        'Y': 19,
        'Z': 20,
        }
    mobj = re.match(r'^\s*(\S+)-ECP\s+(\d+)\s+(\d+)\s*$', lines[0], re.IGNORECASE)
    if mobj is None:
        raise RuntimeError('Where is the "ATOM-ECP max_L nelectron" line?')
    max_L = int(mobj.group(2))
    nelectron = int(mobj.group(3))

    lines2 = lines[1:]

    potential_indices = [
        index
        for index, line in enumerate(lines2)
        if re.match(r'^\s*(\S+)\s+potential\s*$', line, re.IGNORECASE)
    ]

    if len(potential_indices) < 1:
        raise RuntimeError('len(potential_indices) < 1')

    potential_indices.append(len(lines2))

    # Top shell

    index = 0

    potential_start = potential_indices[index]
    potential_stop = potential_indices[index + 1]

    mobj = re.match(r'^\s*(\S)\s+potential\s*$', lines2[potential_start], re.IGNORECASE)
    if mobj is None:
        raise RuntimeError('Where is the "max_L potential" line?')

    shell_char = mobj.group(1).upper()
    if shell_char not in _shell_data_:
        raise RuntimeError(f'Unknown shell type: {shell_char}')

    if _shell_data_[shell_char] != max_L:
        raise RuntimeError('Invalid max_L')

    mobj = re.match(r'^\s*(\d+)\s*$', lines2[potential_start + 1])
    if mobj is None:
        raise RuntimeError('Where is the "nprimitive" line?')

    nprimitive = int(mobj.group(1))

    lines3 = lines2[potential_start + 2 : potential_stop]

    if len(lines3) != nprimitive:
        raise RuntimeError('nprimitive is incorrect')

    ns = []
    ws = []
    es = []
    for line in lines3:
        mobj = re.match(r'^\s*(\d+)\s+(\S+)\s+(\S+)\s*$', line)
        if mobj is None:
            raise RuntimeError('where is the "n e w" line?')
        ns.append(int(mobj.group(1)))
        es.append(float(mobj.group(2)))
        ws.append(float(mobj.group(3)))

    top_shell = ECPShellInfo(L=max_L, ns=ns, ws=ws, es=es)

    # Other shells

    shells = []
    for index in range(1, len(potential_indices) - 1):
        potential_start = potential_indices[index]
        potential_stop = potential_indices[index + 1]

        mobj = re.match(r'^\s*(\S)-(\S)\s+potential\s*$', lines2[potential_start], re.IGNORECASE)
        if mobj is None:
            raise RuntimeError('Where is the "L-max_L potential" line?')

        if _shell_data_[mobj.group(2).upper()] != max_L:
            raise RuntimeError('Invalid max_L')

        L = _shell_data_[mobj.group(1).upper()]

        mobj = re.match(r'^\s*(\d+)\s*$', lines2[potential_start + 1])
        if mobj is None:
            raise RuntimeError('Where is the "nprimitive" line?')

        nprimitive = int(mobj.group(1))

        lines3 = lines2[potential_start + 2 : potential_stop]

        if len(lines3) != nprimitive:
            raise RuntimeError('nprimitive is incorrect')

        ns = []
        ws = []
        es = []
        for line in lines3:
            mobj = re.match(r'^\s*(\d+)\s+(\S+)\s+(\S+)\s*$', line)
            if mobj is None:
                raise RuntimeError('where is the "n e w" line?')
            ns.append(int(mobj.group(1)))
            es.append(float(mobj.group(2)))
            ws.append(float(mobj.group(3)))

        shells.append(ECPShellInfo(L=L, ns=ns, ws=ws, es=es))

    return {
        'nelectron': nelectron,
        'shells': shells,
        'top_shell': top_shell,
        }

def simple_ecp_parser(
    *,
    filename,
    symbols,
    ):

    with open(filename) as fp:
        lines = fp.readlines()

    # => Cleaning <= #

    # Strip blank lines out
    lines = [_ for _ in lines if len(_.strip())]
    # Strip comment lines out
    re_comment = re.compile(r'\s*!')
    lines = [_ for _ in lines if not re.match(re_comment, _)]
    # Lines must be nonzero
    if len(lines) == 0: raise RuntimeError('GBS lines are blank')

    # Find ECP entries matching this kind of pattern
    # SI     0
    # SI-ECP     2     10
    # But joined onto a single line
    ecp_re = re.compile(r'^\s*(\S+)\s+(\d+)\s*(\S+)-ECP\s+(\d+)\s+(\d+)\s*$')
    ecp_inds = []
    for ind in range(len(lines) - 1):
        line = lines[ind] + lines[ind + 1]  # Join into a single line
        if ecp_re.match(line):
            ecp_inds.append(ind)
    ecp_inds.append(len(lines))

    # Extract the lines pertaining to each atom symbol
    symbol_ecp_info = {}
    for k in range(len(ecp_inds) - 1):
        ind1 = ecp_inds[k]
        ind2 = ecp_inds[k + 1]
        if (ind2 - ind1) <= 0:  # Guard against empty entries
            continue
        mobj = re.match(
            r'^\s*(\S+)\s+(\d+)\s*$', lines[ind1]
        )  # Check if the line is a regular basis entry or an ECP entry
        if mobj == None:
            raise Exception("Where is the ID V line?")
        if mobj.group(2) != '0':
            continue  # Regular basis entry, not an ECP entry
        # Try to match the next line to something of the form
        # SR-ECP     3     28
        mobj = re.match(r'^\s*(\S+)-ECP\s+(\d+)\s+(\d+)\s*$', lines[ind1 + 1])
        if mobj is None:
            # This is a regular atom entry, not an ECP entry
            continue
        symbol_ecp_info[mobj.group(1).upper()] = _ecp_shell_info_from_ecp_lines(lines=lines[ind1 + 1 : ind2])

    ecp_metadata = []
    for symbol in symbols:
        ecp_metadata.append(symbol_ecp_info.get(symbol.upper(), None))

    return ecp_metadata

def simple_gbs_parser(
    *,
    filename,
    symbols,
    ):

    shell_types = [
        'S', 'P', 'D', 'F', 'G', 'H', 'I',
        'K', 'L', 'M', 'N', 'O', 'Q', 'R',
        'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
        ]
    with open(filename) as fp:
        lines = fp.readlines()

    # => Cleaning <= #

    # Strip blank lines out
    lines = [_ for _ in lines if len(_.strip())]
    # Strip comment lines out
    re_comment = re.compile(r'\s*!')
    lines = [_ for _ in lines if not re.match(re_comment, _)]
    # Lines must be nonzero
    if len(lines) == 0: raise RuntimeError('GBS lines are blank')

    # => Spherical / Cartesian Tag Line <= #

    mobj = re.match(r'^\s*(spherical|cartesian)\s*$', lines[0])
    if mobj is None:
        raise RuntimeError("First line of GBS file must be 'spherical' or 'cartesian', instead is: %s" % lines[0])

    if mobj.group(1) == 'spherical':
        is_pure = True
    elif mobj.group(1) == 'cartesian':
        is_pure = False
    else:
        raise RuntimeError('Invalid cartesian/spherical label: %s' % mobj.group(1))

    lines = lines[1:]

    # => Atom Block Location <= #

    re_separator = re.compile(r'\s*\*\*\*\*\s*$')
    separator_indices = [k for k, line in enumerate(lines) if re.match(re_separator, line)]
    if len(separator_indices) == 0:
        raise RuntimeError('No **** separators present')
    if separator_indices[-1] + 1 != len(lines):
        raise RuntimeError('Last line must be ****, instead is: %s' % lines[-1])
    lines = lines[:-1]
    separator_indices = separator_indices[:-1]

    if len(lines) == 0: raise RuntimeError('GBS lines are blank')
    if len(separator_indices) == 0: raise RuntimeError('No **** separators present')

    # => Atom IDs and corresponding block index <= #

    re_atom_id = re.compile(r'^\s*(\S+)\s+(\d+)\s*$')
    atom_id_lines = [lines[_ + 1] for _ in separator_indices]
    symbol_index_map = {}
    for index, atom_id_line in enumerate(atom_id_lines):
        mobj = re.match(re_atom_id, atom_id_line)
        if mobj is None:
            raise RuntimeError('Malformed atom ID line: %s' % (atom_id_line))
        symbol = mobj.group(1)
        atom_index = int(mobj.group(2))
        if atom_index != 0:
            raise RuntimeError('"Symbol 0" is only allowed atom line - multiple basis sets per atom type in GBS files is not supported by this library')
        symbol_upper = symbol.upper()
        if symbol in symbol_index_map:
            raise RuntimeError(f'Duplicate atomic symbol in GBS file: {symbol}')
        symbol_index_map[symbol_upper] = index

    re_shell_type = re.compile(r'^\s*(\S+)\s+(\d+)\s+(\S+)\s*$')
    re_primitive = re.compile(r'^\s*(\S+)\s+(\S+)\s*$')

    shellinfo = []
    for symbol in symbols:
        symbol = symbol.upper()
        if symbol not in symbol_index_map:
            raise RuntimeError(f'{symbol} not defined in GBS file')
        block_index = symbol_index_map[symbol]
        index1 = separator_indices[block_index + 0]
        index2 = separator_indices[block_index + 1] if block_index + 1 < len(separator_indices) else len(lines)

        block_lines = lines[index1+2:index2]
        shell_type_indices = [k for k, line in enumerate(block_lines) if re.match(re_shell_type, line)]

        atom_shells = []
        for index in shell_type_indices:
            shell_type_line = block_lines[index]
            mobj = re.match(re_shell_type, shell_type_line)
            am_symbol_upper = mobj.group(1).upper()
            if am_symbol_upper not in shell_types:
                raise RuntimeError(f'Unknown angular momentum symbol: {am_symbol_upper}')
            L = shell_types.index(am_symbol_upper)
            nprimitive = int(mobj.group(2))
            normalization = float(mobj.group(3))
            exponents = []
            coefficients_raw = []
            for K in range(nprimitive):
                primitive_line = block_lines[index + 1 + K]
                # Replace D with E for FORTRAN notation
                primitive_line = primitive_line.replace('D', 'E')
                primitive_line = primitive_line.replace('d', 'e')
                mobj = re.match(re_primitive, primitive_line)
                exponents.append(float(mobj.group(1)))
                coefficients_raw.append(float(mobj.group(2)))
            atom_shells.append({
                'is_pure' : is_pure,
                'L' : L,
                'exponents' : exponents,
                'coefficients' : coefficients_raw,
                'normalization' : normalization,
                })
        shellinfo.append(atom_shells)
    return shellinfo


# 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 re
from collections import namedtuple

def simple_xyz_parser(
    *,
    filename,
    to_bohr_scale_factor,
    ):

    elements = [
        "X",  "H",  "HE", "LI", "BE", "B",  "C",  "N",  "O",  "F",  "NE",
        "NA", "MG", "AL", "SI", "P",  "S",  "CL", "AR", "K",  "CA",
        "SC", "TI", "V",  "CR", "MN", "FE", "CO", "NI", "CU", "ZN",
        "GA", "GE", "AS", "SE", "BR", "KR", "RB", "SR", "Y",  "ZR",
        "NB", "MO", "TC", "RU", "RH", "PD", "AG", "CD", "IN", "SN",
        "SB", "TE", "I",  "XE", "CS", "BA", "LA", "CE", "PR", "ND",
        "PM", "SM", "EU", "GD", "TB", "DY", "HO", "ER", "TM", "YB",
        "LU", "HF", "TA", "W",  "RE", "OS", "IR", "PT", "AU", "HG",
        "TL", "PB", "BI", "PO", "AT", "RN", "FR", "RA", "AC", "TH",
        "PA", "U",  "NP", "PU", "AM", "CM", "BK", "CF", "ES", "FM",
        "MD", "NO", "LR", "RF", "DB", "SG", "BH", "HS", "MT", "DS",
        "RG", "CN", "NH", "FL", "MC", "LV", "TS", "OG"
        ]
    symbols = []
    xyzs = []
    Zs = []
    with open(filename) as fp:
        atom_count = int(fp.readline())
        comment = fp.readline()
        for _ in range(atom_count):
            parts = fp.readline().split()
            symbol = parts[0].upper()
            # The potential integrals routine compute electron-electron interaction
            # integrals; to make this an electron-nucleus interaction we scale Z here
            Z = -1.0 * elements.index(symbol)
            x = float(parts[1]) * to_bohr_scale_factor
            y = float(parts[2]) * to_bohr_scale_factor
            z = float(parts[3]) * to_bohr_scale_factor
            xyzs.extend([x, y, z])
            symbols.append(symbol)
            Zs.append(Z)
    return symbols, xyzs, Zs

ECPShellInfo = namedtuple('ECPShellInfo', ['L', 'ns', 'ws', 'es'])

def _ecp_shell_info_from_ecp_lines(
    *,
    lines,
    ):
    _shell_data_ = {
        'S': 0,
        'P': 1,
        'D': 2,
        'F': 3,
        'G': 4,
        'H': 5,
        'I': 6,
        'K': 7,
        'L': 8,
        'M': 9,
        'N': 10,


    with open(filename) as fp:
        atom_count = int(fp.readline())
        comment = fp.readline()
        for _ in range(atom_count):
            parts = fp.readline().split()
            symbol = parts[0].upper()
            # The potential integrals routine compute electron-electron interaction
            # integrals; to make this an electron-nucleus interaction we scale Z here
            Z = -1.0 * elements.index(symbol)
            x = float(parts[1]) * to_bohr_scale_factor
            y = float(parts[2]) * to_bohr_scale_factor
            z = float(parts[3]) * to_bohr_scale_factor
            xyzs.extend([x, y, z])
            symbols.append(symbol)
            Zs.append(Z)
    return symbols, xyzs, Zs

ECPShellInfo = namedtuple('ECPShellInfo', ['L', 'ns', 'ws', 'es'])

def _ecp_shell_info_from_ecp_lines(
    *,
    lines,
    ):
    _shell_data_ = {
        'S': 0,
        'P': 1,
        'D': 2,
        'F': 3,
        'G': 4,
        'H': 5,
        'I': 6,
        'K': 7,
        'L': 8,
        'M': 9,
        'N': 10,
        'O': 11,
        'Q': 12,
        'R': 13,
        'T': 14,
        'U': 15,
        'V': 16,
        'W': 17,
        'X': 18,
        'Y': 19,
        'Z': 20,
        }
    mobj = re.match(r'^\s*(\S+)-ECP\s+(\d+)\s+(\d+)\s*$', lines[0], re.IGNORECASE)
    if mobj is None:
        raise RuntimeError('Where is the "ATOM-ECP max_L nelectron" line?')
    max_L = int(mobj.group(2))
    nelectron = int(mobj.group(3))

    lines2 = lines[1:]

    potential_indices = [
        index
        for index, line in enumerate(lines2)
        if re.match(r'^\s*(\S+)\s+potential\s*$', line, re.IGNORECASE)
    ]

    if len(potential_indices) < 1:
        raise RuntimeError('len(potential_indices) < 1')

    potential_indices.append(len(lines2))

    # Top shell

    index = 0

    potential_start = potential_indices[index]
    potential_stop = potential_indices[index + 1]

    mobj = re.match(r'^\s*(\S)\s+potential\s*$', lines2[potential_start], re.IGNORECASE)
    if mobj is None:
        raise RuntimeError('Where is the "max_L potential" line?')

    shell_char = mobj.group(1).upper()
    if shell_char not in _shell_data_:
        raise RuntimeError(f'Unknown shell type: {shell_char}')

    if _shell_data_[shell_char] != max_L:
        raise RuntimeError('Invalid max_L')

    mobj = re.match(r'^\s*(\d+)\s*$', lines2[potential_start + 1])
    if mobj is None:
        raise RuntimeError('Where is the "nprimitive" line?')

    nprimitive = int(mobj.group(1))

    lines3 = lines2[potential_start + 2 : potential_stop]

    if len(lines3) != nprimitive:
        raise RuntimeError('nprimitive is incorrect')

    ns = []
    ws = []
    es = []
    for line in lines3:
        mobj = re.match(r'^\s*(\d+)\s+(\S+)\s+(\S+)\s*$', line)
        if mobj is None:
            raise RuntimeError('where is the "n e w" line?')
        ns.append(int(mobj.group(1)))
        es.append(float(mobj.group(2)))
        ws.append(float(mobj.group(3)))

    top_shell = ECPShellInfo(L=max_L, ns=ns, ws=ws, es=es)

    # Other shells

    shells = []
    for index in range(1, len(potential_indices) - 1):
        potential_start = potential_indices[index]
        potential_stop = potential_indices[index + 1]

        mobj = re.match(r'^\s*(\S)-(\S)\s+potential\s*$', lines2[potential_start], re.IGNORECASE)
        if mobj is None:
            raise RuntimeError('Where is the "L-max_L potential" line?')

        if _shell_data_[mobj.group(2).upper()] != max_L:
            raise RuntimeError('Invalid max_L')

        L = _shell_data_[mobj.group(1).upper()]

        mobj = re.match(r'^\s*(\d+)\s*$', lines2[potential_start + 1])
        if mobj is None:
            raise RuntimeError('Where is the "nprimitive" line?')

        nprimitive = int(mobj.group(1))

        lines3 = lines2[potential_start + 2 : potential_stop]

        if len(lines3) != nprimitive:
            raise RuntimeError('nprimitive is incorrect')

        ns = []
        ws = []
        es = []
        for line in lines3:
            mobj = re.match(r'^\s*(\d+)\s+(\S+)\s+(\S+)\s*$', line)
            if mobj is None:
                raise RuntimeError('where is the "n e w" line?')
            ns.append(int(mobj.group(1)))
            es.append(float(mobj.group(2)))
            ws.append(float(mobj.group(3)))

        shells.append(ECPShellInfo(L=L, ns=ns, ws=ws, es=es))

    return {
        'nelectron': nelectron,
        'shells': shells,
        'top_shell': top_shell,
        }

def simple_ecp_parser(
    *,
    filename,
    symbols,
    ):

    with open(filename) as fp:
        lines = fp.readlines()

    # => Cleaning <= #

    # Strip blank lines out
    lines = [_ for _ in lines if len(_.strip())]
    # Strip comment lines out
    re_comment = re.compile(r'\s*!')
    lines = [_ for _ in lines if not re.match(re_comment, _)]
    # Lines must be nonzero
    if len(lines) == 0: raise RuntimeError('GBS lines are blank')


        ns = []
        ws = []
        es = []
        for line in lines3:
            mobj = re.match(r'^\s*(\d+)\s+(\S+)\s+(\S+)\s*$', line)
            if mobj is None:
                raise RuntimeError('where is the "n e w" line?')
            ns.append(int(mobj.group(1)))
            es.append(float(mobj.group(2)))
            ws.append(float(mobj.group(3)))

        shells.append(ECPShellInfo(L=L, ns=ns, ws=ws, es=es))

    return {
        'nelectron': nelectron,
        'shells': shells,
        'top_shell': top_shell,
        }

def simple_ecp_parser(
    *,
    filename,
    symbols,
    ):

    with open(filename) as fp:
        lines = fp.readlines()

    # => Cleaning <= #

    # Strip blank lines out
    lines = [_ for _ in lines if len(_.strip())]
    # Strip comment lines out
    re_comment = re.compile(r'\s*!')
    lines = [_ for _ in lines if not re.match(re_comment, _)]
    # Lines must be nonzero
    if len(lines) == 0: raise RuntimeError('GBS lines are blank')

    # Find ECP entries matching this kind of pattern
    # SI     0
    # SI-ECP     2     10
    # But joined onto a single line
    ecp_re = re.compile(r'^\s*(\S+)\s+(\d+)\s*(\S+)-ECP\s+(\d+)\s+(\d+)\s*$')
    ecp_inds = []
    for ind in range(len(lines) - 1):
        line = lines[ind] + lines[ind + 1]  # Join into a single line
        if ecp_re.match(line):
            ecp_inds.append(ind)
    ecp_inds.append(len(lines))

    # Extract the lines pertaining to each atom symbol
    symbol_ecp_info = {}
    for k in range(len(ecp_inds) - 1):
        ind1 = ecp_inds[k]
        ind2 = ecp_inds[k + 1]
        if (ind2 - ind1) <= 0:  # Guard against empty entries
            continue
        mobj = re.match(
            r'^\s*(\S+)\s+(\d+)\s*$', lines[ind1]
        )  # Check if the line is a regular basis entry or an ECP entry
        if mobj == None:
            raise Exception("Where is the ID V line?")
        if mobj.group(2) != '0':
            continue  # Regular basis entry, not an ECP entry
        # Try to match the next line to something of the form
        # SR-ECP     3     28
        mobj = re.match(r'^\s*(\S+)-ECP\s+(\d+)\s+(\d+)\s*$', lines[ind1 + 1])
        if mobj is None:
            # This is a regular atom entry, not an ECP entry
            continue
        symbol_ecp_info[mobj.group(1).upper()] = _ecp_shell_info_from_ecp_lines(lines=lines[ind1 + 1 : ind2])

    ecp_metadata = []
    for symbol in symbols:
        ecp_metadata.append(symbol_ecp_info.get(symbol.upper(), None))

    return ecp_metadata

def simple_gbs_parser(
    *,
    filename,
    symbols,
    ):

    shell_types = [
        'S', 'P', 'D', 'F', 'G', 'H', 'I',
        'K', 'L', 'M', 'N', 'O', 'Q', 'R',
        'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
        ]
    with open(filename) as fp:
        lines = fp.readlines()

    # => Cleaning <= #

    # Strip blank lines out
    lines = [_ for _ in lines if len(_.strip())]


        )  # Check if the line is a regular basis entry or an ECP entry
        if mobj == None:
            raise Exception("Where is the ID V line?")
        if mobj.group(2) != '0':
            continue  # Regular basis entry, not an ECP entry
        # Try to match the next line to something of the form
        # SR-ECP     3     28
        mobj = re.match(r'^\s*(\S+)-ECP\s+(\d+)\s+(\d+)\s*$', lines[ind1 + 1])
        if mobj is None:
            # This is a regular atom entry, not an ECP entry
            continue
        symbol_ecp_info[mobj.group(1).upper()] = _ecp_shell_info_from_ecp_lines(lines=lines[ind1 + 1 : ind2])

    ecp_metadata = []
    for symbol in symbols:
        ecp_metadata.append(symbol_ecp_info.get(symbol.upper(), None))

    return ecp_metadata

def simple_gbs_parser(
    *,
    filename,
    symbols,
    ):

    shell_types = [
        'S', 'P', 'D', 'F', 'G', 'H', 'I',
        'K', 'L', 'M', 'N', 'O', 'Q', 'R',
        'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
        ]
    with open(filename) as fp:
        lines = fp.readlines()

    # => Cleaning <= #

    # Strip blank lines out
    lines = [_ for _ in lines if len(_.strip())]
    # Strip comment lines out
    re_comment = re.compile(r'\s*!')
    lines = [_ for _ in lines if not re.match(re_comment, _)]
    # Lines must be nonzero
    if len(lines) == 0: raise RuntimeError('GBS lines are blank')

    # => Spherical / Cartesian Tag Line <= #

    mobj = re.match(r'^\s*(spherical|cartesian)\s*$', lines[0])
    if mobj is None:
        raise RuntimeError("First line of GBS file must be 'spherical' or 'cartesian', instead is: %s" % lines[0])

    if mobj.group(1) == 'spherical':
        is_pure = True
    elif mobj.group(1) == 'cartesian':
        is_pure = False
    else:
        raise RuntimeError('Invalid cartesian/spherical label: %s' % mobj.group(1))

    lines = lines[1:]

    # => Atom Block Location <= #

    re_separator = re.compile(r'\s*\*\*\*\*\s*$')
    separator_indices = [k for k, line in enumerate(lines) if re.match(re_separator, line)]
    if len(separator_indices) == 0:
        raise RuntimeError('No **** separators present')
    if separator_indices[-1] + 1 != len(lines):
        raise RuntimeError('Last line must be ****, instead is: %s' % lines[-1])
    lines = lines[:-1]
    separator_indices = separator_indices[:-1]

    if len(lines) == 0: raise RuntimeError('GBS lines are blank')
    if len(separator_indices) == 0: raise RuntimeError('No **** separators present')

    # => Atom IDs and corresponding block index <= #

    re_atom_id = re.compile(r'^\s*(\S+)\s+(\d+)\s*$')
    atom_id_lines = [lines[_ + 1] for _ in separator_indices]
    symbol_index_map = {}
    for index, atom_id_line in enumerate(atom_id_lines):
        mobj = re.match(re_atom_id, atom_id_line)
        if mobj is None:
            raise RuntimeError('Malformed atom ID line: %s' % (atom_id_line))
        symbol = mobj.group(1)
        atom_index = int(mobj.group(2))
        if atom_index != 0:
            raise RuntimeError('"Symbol 0" is only allowed atom line - multiple basis sets per atom type in GBS files is not supported by this library')
        symbol_upper = symbol.upper()
        if symbol in symbol_index_map:
            raise RuntimeError(f'Duplicate atomic symbol in GBS file: {symbol}')
        symbol_index_map[symbol_upper] = index

    re_shell_type = re.compile(r'^\s*(\S+)\s+(\d+)\s+(\S+)\s*$')
    re_primitive = re.compile(r'^\s*(\S+)\s+(\S+)\s*$')

    shellinfo = []
    for symbol in symbols:
        symbol = symbol.upper()
        if symbol not in symbol_index_map:
            raise RuntimeError(f'{symbol} not defined in GBS file')
        block_index = symbol_index_map[symbol]
        index1 = separator_indices[block_index + 0]
        index2 = separator_indices[block_index + 1] if block_index + 1 < len(separator_indices) else len(lines)

        block_lines = lines[index1+2:index2]
        shell_type_indices = [k for k, line in enumerate(block_lines) if re.match(re_shell_type, line)]

        atom_shells = []
        for index in shell_type_indices:
            shell_type_line = block_lines[index]
            mobj = re.match(re_shell_type, shell_type_line)
            am_symbol_upper = mobj.group(1).upper()
            if am_symbol_upper not in shell_types:
                raise RuntimeError(f'Unknown angular momentum symbol: {am_symbol_upper}')
            L = shell_types.index(am_symbol_upper)
            nprimitive = int(mobj.group(2))
            normalization = float(mobj.group(3))
            exponents = []
            coefficients_raw = []
            for K in range(nprimitive):
                primitive_line = block_lines[index + 1 + K]
                # Replace D with E for FORTRAN notation
                primitive_line = primitive_line.replace('D', 'E')
                primitive_line = primitive_line.replace('d', 'e')
                mobj = re.match(re_primitive, primitive_line)
                exponents.append(float(mobj.group(1)))
                coefficients_raw.append(float(mobj.group(2)))
            atom_shells.append({
                'is_pure' : is_pure,
                'L' : L,
                'exponents' : exponents,
                'coefficients' : coefficients_raw,
                'normalization' : normalization,
                })
        shellinfo.append(atom_shells)
    return shellinfo


# 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.

try:
    import numpy as np
except ImportError:
    raise RuntimeError("numpy could not be imported. It is available via\n\tpip install numpy")

def normalize_shell_coefficients(
    coefficients,
    exponents,
    L,
    normalization,
    ):
    """
    Normalizes a Gaussian shell of primitive functions, so that their self overlap is the
    value provided by the "normalization" argument.
    """

    exponents = np.array(exponents)
    coefficients = np.array(coefficients)
    assert len(exponents) == len(coefficients)

    dfact = 1
    for l in range(1, L+1):
        dfact *= 2 * l - 1

    # Primitive normalization
    output = np.sqrt( 2**L * (2 * exponents)**(L+1.5) / (np.pi**1.5 * dfact) ) * coefficients

    # Contraction normalization
    Q = 0.0
    for c1, e1 in zip(coefficients, exponents):
        for c2, e2 in zip(coefficients, exponents):
            Q += (np.sqrt(4 * e1 * e2) / (e1 + e2)) ** (L + 1.5) * c1 * c2
    Q = Q**(-0.5)
    Q *= np.sqrt(normalization)

    return output * Q


# 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 numpy as np

def symbol_to_ahlrichs_radius(
    *,
    symbol,
    ):

    symbol = symbol.upper()

    radii = {
         "X" : 1.00,
         "H" : 0.80,
        "HE" : 0.90,
        "LI" : 1.80,
        "BE" : 1.40,
         "B" : 1.30,
         "C" : 1.10,
         "N" : 0.90,
         "O" : 0.90,
         "F" : 0.90,
        "NE" : 0.90,
        "NA" : 1.40,
        "MG" : 1.30,
        "AL" : 1.30,
        "SI" : 1.20,
         "P" : 1.10,
         "S" : 1.00,
        "CL" : 1.00,
        "AR" : 1.00,
         "K" : 1.50,
        "CA" : 1.40,
        "SC" : 1.30,
        "TI" : 1.20,
         "V" : 1.20,
        "CR" : 1.20,
        "MN" : 1.20,
        "FE" : 1.20,
        "CO" : 1.20,
        "NI" : 1.10,
        "CU" : 1.10,
        "ZN" : 1.10,
        "GA" : 1.10,
        "GE" : 1.00,
        "AS" : 0.90,
        "SE" : 0.90,
        "BR" : 0.90,
        "KR" : 0.90,
    }

    return radii.get(symbol, 1.0)

def build_ahlrichs_radial_quadrature(
    *,
    npoint,
    R,
    ):

    alpha = 0.6

    n = np.arange(1.0, npoint+1.0)
    z = n * np.pi / (npoint + 1.0)
    x = np.cos(z)
    y = np.sin(z)
    u = np.log((1.0 - x) / 2.0)
    v = ((1.0 + x)**alpha) / np.log(2.0)
    radial_nodes = - R * v * u
    radial_weights = np.pi / (npoint + 1.0) * y * R  * v * (-alpha * u / (1.0 + x) + 1.0 / (1.0 - x)) * radial_nodes**2

    return np.flip(radial_nodes), np.flip(radial_weights)


# 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.

# There are numerous way to access GPU resources from Python.  High level
# packages such as CuPy, TensorFlow and Pytorch can simplify allocations and
# transfers between host and device.  Here we demonstrate two lower level
# approaches to accessing CUDA functions to provide these utilities: 1) using
# the cuda-bindings package to access native CUDA in Pythonic way, and 2) using
# the ctypes module to directly access these functions from the CUDA runtime
# library.  A real implementation should ideally use a single unified
# mechanism, but a mixture is used here for demonstration purposes.

import ctypes
try:
    import numpy as np
except ImportError:
    raise RuntimeError("numpy could not be imported. It is available via\n\tpip install numpy")

try:
    import cuda.bindings.runtime as cuda
except ImportError:
    raise RuntimeError("cuda.bindings could not be imported. It is available via\n\tpip install cuda-bindings")

# These will be used by the ctypes-powered utility functions below.
try:
    _cublas = ctypes.cdll.LoadLibrary('libcublas.so')
except OSError as err:
    raise RuntimeError("Could not load cuBLAS library.") from err

try:
    _cusolver = ctypes.cdll.LoadLibrary('libcusolver.so')
except OSError as err:
    raise RuntimeError("Could not load cuSolver library.") from err

def make_stream():
    """
    This function should return a new cuestStream_t (or Python
    equivalent) representing a CUDA stream.
    """

    ret, stream = cuda.cudaStreamCreate()

    if ret != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Unable to make CUDA stream:", ret)

    return stream

def cuda_malloc(
    *,
    size_in_bytes,
    ):
    """
    This function should allocate device memory, and return an integer
    representation of the address of the resulting block.
    """
    ret, pointer = cuda.cudaMalloc(size_in_bytes)

    if ret != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Unable to allocate CUDA array:", ret)

    return pointer

def cuda_memcpy_dtoh(
    *,
    host_pointer,
    device_pointer,
    size_in_bytes,
    stream=0,
    ):
    """
    This function should take integer representations of pointers to host
    and device memory allocations, and copy size_in_bytes bytes from the
    device memory to the host memory.  The stream to execute this can be
    optionally provided.
    """

    ret = cuda.cudaMemcpyAsync(
        dst=host_pointer,
        src=device_pointer,
        count=size_in_bytes,
        kind=cuda.cudaMemcpyKind.cudaMemcpyDeviceToHost,
        stream=stream)

    if ret[0] != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Copy from device to host failed:", ret[0])

    cuda.cudaStreamSynchronize(stream)

def cuda_memcpy_htod(
    *,
    device_pointer,
    host_pointer,
    size_in_bytes,
    stream=0,
    ):
    """
    This function should take integer representations of pointers to host
    and device memory allocations, and copy size_in_bytes bytes from the
    host memory to the device memory.  The stream to execute this can be
    optionally provided.
    """

    ret = cuda.cudaMemcpyAsync(
        dst=device_pointer,
        src=host_pointer,
        count=size_in_bytes,
        kind=cuda.cudaMemcpyKind.cudaMemcpyHostToDevice,
        stream=stream)

    if ret[0] != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Copy from host to device failed:", ret[0])

    cuda.cudaStreamSynchronize(stream)

def cuda_free(
    array
    ):
    """
    This function should free memory allocated on the device, passed in
    as an integer representation of its device pointer
    """

    ret = cuda.cudaFree(array)

    if ret[0] != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Unable to free CUDA array:", ret[0])

def make_cublas_handle():
    """
    Helper function demonstrating how to use ctypes to access the C
    cuBLAS API to make a new cuBLAS handle.
    """
    handle = ctypes.c_void_p()

    CUBLAS_STATUS_SUCCESS = 0

    _cublas.cublasCreate_v2.restype = ctypes.c_int
    _cublas.cublasCreate_v2.argtypes = [ctypes.POINTER(ctypes.c_void_p)]

    status = _cublas.cublasCreate_v2(ctypes.byref(handle))
    if status != CUBLAS_STATUS_SUCCESS:
        raise RuntimeError(f"cublasCreate_v2 failed with status {status}")

    return handle

def free_cublas_handle(
    *,
    handle
    ):

    """
    Helper function to free cuBLAS resource handle.
    """

    CUBLAS_STATUS_SUCCESS = 0

    _cublas.cublasDestroy_v2.restype = ctypes.c_int
    _cublas.cublasDestroy_v2.argtypes = [ctypes.c_void_p]

    status = _cublas.cublasDestroy_v2(handle)
    if status != CUBLAS_STATUS_SUCCESS:
        raise RuntimeError(f"cublasDestroy_v2 failed with status {status}")

def make_cusolver_handle():
    """
    Make a cuSolver handle via the CUDA C API.
    """
    handle = ctypes.c_void_p()

    CUSOLVER_STATUS_SUCCESS = 0

    _cusolver.cusolverDnCreate.restype = ctypes.c_int
    _cusolver.cusolverDnCreate.argtypes = [ctypes.POINTER(ctypes.c_void_p)]

    status = _cusolver.cusolverDnCreate(ctypes.byref(handle))
    if status != CUSOLVER_STATUS_SUCCESS:
        raise RuntimeError(f"cusolverDnCreate failed with status {status}")

    return handle

def free_cusolver_handle(
    handle
    ):
    """
    Free resources associated with a cuSolver handle.
    """

    CUSOLVER_STATUS_SUCCESS = 0

    _cusolver.cusolverDnDestroy.restype = ctypes.c_int
    _cusolver.cusolverDnDestroy.argtypes = [ctypes.c_void_p]

    status = _cusolver.cusolverDnDestroy(handle)
    if status != CUSOLVER_STATUS_SUCCESS:
        raise RuntimeError(f"cusolverDnDestroy failed with status {status}")

def set_cublas_stream(
    *,
    handle,
    stream,
    ):
    """
    Set the active stream for a given cuBLAS handle
    """

    CUBLAS_STATUS_SUCCESS = 0

    _cublas.cublasSetStream_v2.restype = ctypes.c_int
    _cublas.cublasSetStream_v2.argtypes = [ctypes.c_void_p, ctypes.c_void_p]

    status = _cublas.cublasSetStream_v2(handle, stream)
    if status != CUBLAS_STATUS_SUCCESS:
        raise RuntimeError(f"cublasSetStream failed with status {status}")

def set_cusolver_stream(
    *,
    handle,
    stream,
    ):
    """
    Set the stream for a given cuSolver handle.
    """

    CUSOLVER_STATUS_SUCCESS = 0

    _cusolver.cusolverDnSetStream.restype = ctypes.c_int
    _cusolver.cusolverDnSetStream.argtypes = [ctypes.c_void_p, ctypes.c_void_p]

    status = _cusolver.cusolverDnSetStream(handle, stream)
    if status != CUSOLVER_STATUS_SUCCESS:
        raise RuntimeError(f"cusolverDnSetStream failed with status {status}")

class WorkspaceDescriptor:
    """
    This helper function needs to implement the following C structure,
    and define a pointer method that returns a pointer to it.

    typedef struct {
        size_t hostBufferSizeInBytes;      ///< Required size of host workspace buffer in bytes
        size_t deviceBufferSizeInBytes;    ///< Required size of device workspace buffer in bytes
    } cuestWorkspaceDescriptor_t;

    This implementation uses Numpy, as the ctypes.data member provides a
    convenient way to access a pointer to the object
    """

    def __init__(
        self,
        *,
        host_buffer_size_in_bytes = 0,
        device_buffer_size_in_bytes = 0,
        ):

        _workspace_descriptor_dtype = np.dtype([
            ("hostBufferSizeInBytes", np.uint64, ),
            ("deviceBufferSizeInBytes", np.uint64, ),
            ], align=True
            )

        self.struct = np.empty(1, dtype=_workspace_descriptor_dtype)
        self.struct['deviceBufferSizeInBytes'] = device_buffer_size_in_bytes
        self.struct['hostBufferSizeInBytes'] = host_buffer_size_in_bytes

    def __str__(
        self,
        ):

        host_size = self.struct['hostBufferSizeInBytes'].item()
        device_size = self.struct['deviceBufferSizeInBytes'].item()
        return f'host buffer size = {host_size} bytes, device buffer size = {device_size} bytes'

    @property
    def pointer(self):
        return self.struct.ctypes.data

class Workspace:
    """
    This helper function needs to implement the following C structure,
    and define a pointer method that returns a pointer to it.

    typedef struct {
        uintptr_t hostBuffer;              ///< Opaque pointer to host-side workspace buffer
        size_t hostBufferSizeInBytes;      ///< Size of host workspace in bytes
        uintptr_t deviceBuffer;            ///< Opaque pointer to device-side (GPU) workspace buffer
        size_t deviceBufferSizeInBytes;    ///< Size of device workspace in bytes
    } cuestWorkspace_t;

    This implementation uses Numpy, as the ctypes.data member provides a
    convenient way to access a pointer to the object
    """
    def __init__(
        self,
        *,
        workspaceDescriptor,
        ):
        _workspace_dtype = np.dtype([
            ("hostBuffer", np.uintp, ),
            ("hostBufferSizeInBytes", np.uint64, ),
            ("deviceBuffer", np.uintp, ),
            ("deviceBufferSizeInBytes", np.uint64, ),
            ], align=True
            )

        host_buffer_size_in_bytes = workspaceDescriptor.struct['hostBufferSizeInBytes'].item()
        device_buffer_size_in_bytes = workspaceDescriptor.struct['deviceBufferSizeInBytes'].item()

        self.struct = np.empty(1, dtype=_workspace_dtype)
        self.struct['deviceBufferSizeInBytes'] = device_buffer_size_in_bytes
        self.struct['hostBufferSizeInBytes'] = host_buffer_size_in_bytes

        if device_buffer_size_in_bytes:
            ret, pointer = cuda.cudaMalloc(device_buffer_size_in_bytes)
            if ret != cuda.cudaError_t.cudaSuccess:
                raise RuntimeError("Failed to allocate workspace", ret)
            self.struct['deviceBuffer'] = pointer
        else:
            self.struct['deviceBuffer'] = 0

        if host_buffer_size_in_bytes:
            self.cpu_memory = np.zeros(host_buffer_size_in_bytes, dtype=np.int8)
            self.struct['hostBuffer'] = self.cpu_memory.ctypes.data
        else:
            self.struct['hostBuffer'] = 0
            self.cpu_memory = None

    def __del__(
        self,
        ):

        if self.struct['deviceBuffer']:
            ret = cuda.cudaFree(self.struct['deviceBuffer'].item())
            if ret[0] != cuda.cudaError_t.cudaSuccess:
                raise RuntimeError("Unable to deallocate workspace:", ret[0])

    @property
    def pointer(self):
        return self.struct.ctypes.data


try:
    import cuda.bindings.runtime as cuda
except ImportError:
    raise RuntimeError("cuda.bindings could not be imported. It is available via\n\tpip install cuda-bindings")

# These will be used by the ctypes-powered utility functions below.
try:
    _cublas = ctypes.cdll.LoadLibrary('libcublas.so')
except OSError as err:
    raise RuntimeError("Could not load cuBLAS library.") from err

try:
    _cusolver = ctypes.cdll.LoadLibrary('libcusolver.so')
except OSError as err:
    raise RuntimeError("Could not load cuSolver library.") from err

def make_stream():
    """
    This function should return a new cuestStream_t (or Python
    equivalent) representing a CUDA stream.
    """

    ret, stream = cuda.cudaStreamCreate()

    if ret != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Unable to make CUDA stream:", ret)

    return stream

def cuda_malloc(
    *,
    size_in_bytes,
    ):
    """
    This function should allocate device memory, and return an integer
    representation of the address of the resulting block.
    """
    ret, pointer = cuda.cudaMalloc(size_in_bytes)

    if ret != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Unable to allocate CUDA array:", ret)

    return pointer

def cuda_memcpy_dtoh(
    *,


    _cusolver = ctypes.cdll.LoadLibrary('libcusolver.so')
except OSError as err:
    raise RuntimeError("Could not load cuSolver library.") from err

def make_stream():
    """
    This function should return a new cuestStream_t (or Python
    equivalent) representing a CUDA stream.
    """

    ret, stream = cuda.cudaStreamCreate()

    if ret != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Unable to make CUDA stream:", ret)

    return stream

def cuda_malloc(
    *,
    size_in_bytes,
    ):
    """
    This function should allocate device memory, and return an integer
    representation of the address of the resulting block.
    """
    ret, pointer = cuda.cudaMalloc(size_in_bytes)

    if ret != cuda.cudaError_t.cudaSuccess:
        raise RuntimeError("Unable to allocate CUDA array:", ret)

    return pointer

def cuda_memcpy_dtoh(
    *,
    host_pointer,
    device_pointer,
    size_in_bytes,
    stream=0,
    ):
    """
    This function should take integer representations of pointers to host
    and device memory allocations, and copy size_in_bytes bytes from the
    device memory to the host memory.  The stream to execute this can be
    optionally provided.
    """

    ret = cuda.cudaMemcpyAsync(
        dst=host_pointer,
        src=device_pointer,
        count=size_in_bytes,