from pyscf import gto, scf, cc, lib, lo
import numpy as np
import time

mol_coords = """C           0.71745       -0.01896       -0.00104
O          -0.69434        0.09886       -0.00123
H           1.04711       -0.55920        0.89128
H           1.02646       -0.56030       -0.90233
H           1.15836        0.97587        0.00620
H          -1.04673       -0.79917        0.00371"""

time_1 = time.time()

mol = gto.M(atom=mol_coords, basis="cc-pVQZ", verbose=1)
mf = scf.RHF(mol).run()


mo_coeff_can = mf.mo_coeff


mo_occ = mo_coeff_can[:, mf.mo_occ > 0]
mo_virt = mo_coeff_can[:, mf.mo_occ == 0]

time_2 = time.time()

loc_occ = lo.PM(mf.mol, mo_occ)
loc_occ.kernel()

loc_virt = lo.PM(mf.mol, mo_virt)
loc_virt.kernel()

time_3 = time.time()

print("Time for SCF: ", time_2 - time_1)
print("Time for Localization: ", time_3 - time_2)