# load python packages
import numpy as np
from pyscf.pbc import gto, dft
from pyscf.tools import trexio

# cell info
# cell
cell = gto.M(
    verbose = 4,
    a = np.eye(3)*3.5668,
    atom = '''C     0.      0.      0.    
              C     0.8917  0.8917  0.8917
              C     1.7834  1.7834  0.    
              C     2.6751  2.6751  0.8917
              C     1.7834  0.      1.7834
              C     2.6751  0.8917  2.6751
              C     0.      1.7834  1.7834
              C     0.8917  2.6751  2.6751'''
)
cell.output = pyscf_output
cell.charge = 0
cell.spin = 0
cell.symmetry = False

# basis set
cell.basis = 'ccecp-ccpvdz'
cell.exp_to_discard = 0.2

# define ecp
cell.ecp = 'ccecp'

# cell build
cell.build(cart=False)

# DFT setting
kpt = [0.0, 0.0, 0.0]
mf = dft.rks.RKS(cell, kpt=kpt)
mf.xc = "lda_x,lda_c_pz"
mf.chkfile = "pyscf_diamond.chk"

# run HF/DFT calc.
total_energy = mf.kernel()
print(f"Total HF/DFT energy = {total_energy}")

# dump to TREXIO file
trexio.to_trexio(mf, 'trexio_diamond.hdf5')