import numpy as np
from pyscf import gto, scf, mcscf, mcpdft
import unittest
#from mrh.my_pyscf.fci import csf_solver
from pyscf.fci.addons import fix_spin_

geom_h2o='''
O  0.00000000   0.08111156   0.00000000
H  0.78620605   0.66349738   0.00000000
H -0.78620605   0.66349738   0.00000000
'''
geom_furan= '''
C        0.000000000     -0.965551055     -2.020010585
C        0.000000000     -1.993824223     -1.018526668
C        0.000000000     -1.352073201      0.181141565
O        0.000000000      0.000000000      0.000000000
C        0.000000000      0.216762264     -1.346821565
H        0.000000000     -1.094564216     -3.092622941
H        0.000000000     -3.062658055     -1.175803180
H        0.000000000     -1.688293885      1.206105691
H        0.000000000      1.250242874     -1.655874372
'''
mol_h2o = gto.M(atom = geom_h2o, basis = 'aug-cc-pVDZ', symmetry='c2v', output='/dev/null', verbose=0)
mol_furan_cation = gto.M(atom = geom_furan, basis = 'sto-3g', charge=1, spin=1, symmetry=False, output='/dev/null', verbose=0)

# Three singlets all of A1 symmetry
def get_h2o(mol,iroots=3): 
    weights = [1/iroots]*iroots
    mf = scf.RHF(mol).run()
    mc = mcpdft.CASSCF(mf,'tPBE', 4, 4, grids_level=1)
    #mc.fcisolver = csf_solver(mol, smult=1, symm='A1')
    fix_spin_(mc.fcisolver, ss=0)
    mc.fcisolver.wfnsym = 'A1'
    mc = mc.multi_state(weights)
    mo = mcscf.sort_mo(mc, mf.mo_coeff, [4,5,8,9])
    mc.conv_tol = 1e-11
    mc.kernel(mo)
    return mc

# Three doublets of A2, B2, and A2 symmetries
def get_furan_cation(mol,iroots=3):
    weights = [1/iroots]*iroots
    mf = scf.RHF(mol).run()
    mc = mcpdft.CASSCF(mf,'tBLYP', 5, 5, grids_level=1)
    #mc.fcisolver = csf_solver(mol, smult=2)
    fix_spin_(mc.fcisolver, ss=0.75)
    mc = mc.multi_state(weights, 'cms')
    mc.max_cycle_macro = 200
    mo = mcscf.sort_mo(mc, mf.mo_coeff, [12,17,18,19,20])
    mc.conv_tol = 1e-11
    mc.kernel(mo)
    return mc

def tearDownModule():
    global mol_h2o, mol_furan_cation
    mol_h2o.stdout.close ()
    mol_furan_cation.stdout.close ()
    del mol_h2o, mol_furan_cation

class KnownValues(unittest.TestCase):
    '''
    The reference values were obtained by numeric differentiation of the energy 
    with respect to the electric field strength extrapolated to zero step size. 
    The fields were applied along each direction in the XYZ frame, and derivatives 
    were evaluated using 2-point central difference formula.   
    '''
    def test_h2o_cms3_tpbe_631g(self):
        dm_ref = np.array(\
            [[0.0000, 1.5991, 0.0000],
            [ 0.0000,-0.2234, 0.0000],
            [ 0.0000, 3.7302, 0.0000]])
        delta = 0.001
        message = "Dipoles are not equal within {} D".format(delta)

        iroots=3
        mc = get_h2o(mol_h2o, iroots)
        for i in range(3):
            with self.subTest (i=i):
                dm_test = mc.dip_moment(unit='Debye', origin="Coord_center",state=i)
                for dmt,dmr in zip(dm_test,dm_ref[i]):
                    self.assertAlmostEqual(dmt, dmr, None, message, delta)
        mc.stdout.close()
        
    def test_furan_cation_cms3_tblyp_sto3g(self):
        # Numerical ref from this software using 2-point central difference formula
        dm_ref = np.array(\
            [[0.0000, -1.0146, -1.0418],
            [ 0.0000, -1.1001, -1.1296],
            [ 0.0000, -0.3954, -0.4061]])
        delta = 0.001
        message = "Dipoles are not equal within {} D".format(delta)

        iroots=3
        mc = get_furan_cation(mol_furan_cation, iroots)
        for i in range(iroots):
            with self.subTest (i=i):
                dm_test = mc.dip_moment(unit='Debye', origin="charge_center",state=i)
                for dmt,dmr in zip(dm_test,dm_ref[i]):
                    self.assertAlmostEqual(dmt, dmr, None, message, delta)

if __name__ == "__main__":
    print("Test for CMS-PDFT permanent dipole moments")
    unittest.main()
