from pyscf import gto, scf
import asf
import numpy as np
from pyscf.mcscf import avas

ASF = asf.asf()

mol = gto.Mole()
mol.atom = """
C          0.00000        0.00000        0.00000
C          0.00000        0.00000        1.20000
"""
mol.basis = 'def2-svp'
mol.charge = 0
mol.spin = 0
mol.build()
# UHF for UNOs
mf = scf.RHF(mol).run(max_cycle=100)
mo_new = mf.stability()[0]
while mo_new is not mf.mo_coeff:
    mf.kernel(dm0=mf.make_rdm1(mo_coeff=mo_new))
    mo_new = mf.stability()[0]

# AVAS initial guess.
ao_labels = ['C 2s', 'C 2p']
norb, ne_act, orbs = avas.avas(mf, ao_labels, canonicalize=False)

# Plot AVAS selected orbitals.
# orbital list
fao = asf.act_fao(mf.mo_occ, ne_act)
orblist = fao + np.array(list(range(norb)))
asf.visualize_mos(mol, orbs, orblist)

# Select an active space using entropies.
ele, mos = ASF.find_active_space(mol, ne_act, norb, orbs, plot=True)
