"""
This script tests the scf module.
"""


import os
import numpy as np

import pytest

from frankenstein.be import BE


@pytest.mark.parametrize(
    "geom, basis, fsites, frank, incore, nibath, solver, jac, B0, bad_con, "
        "good_con, e_ps_ref",
    [
    # Frankenstein BE, FCI solver
    ("geom/h10.zmat", "sto-3g", [0,1,2], True, True, False, "FCI", None, None,
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # frank/FCI/QN/eye
    ("geom/h10.zmat", "sto-3g", [0,1,2], True, True, False, "FCI", None, "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # frank/FCI/QN/scf
    ("geom/h10.zmat", "sto-3g", [0,1,2], True, True, False, "FCI", 2, None,
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # frank/FCI/NR/eye
    ("geom/h10.zmat", "sto-3g", [0,1,2], True, True, False, "FCI", 2, "mp2",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # frank/FCI/NR/mp2
    # Frankenstein BE, FCI solver, nibath
    ("geom/h10.zmat", "sto-3g", [0,1,2], True, True, True, "FCI", None,
        "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8207132177, -2.8129821608, -2.8207131811]),
    # PySCF BE, FCI solver
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, True, False, "FCI", None, None,
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # pyscf/FCI/QN/eye
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, True, False, "FCI", None, "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # pyscf/FCI/QN/scf
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, True, False, "FCI", 2, None,
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # pyscf/FCI/NR/eye
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, True, False, "FCI", 2, "mp2",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # pyscf/FCI/NR/mp2
    # PySCF BE, different solvers
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, True, False, "mp2", None, "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8123768, -2.8133190, -2.8123768]),  # pyscf/mp2/QN/scf
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, True, False, "ccsd", None,
        "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127582, -2.8149148, -2.8127582]),  # pyscf/ccsd/QN/scf
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, True, False, "cisd", None,
        "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127194, -2.8148940, -2.8127194]),  # pyscf/cisd/QN/scf
    # PySCF BE, different fragment size
    ("geom/h10.zmat", "sto-3g", [0,1,2,3], False, True, False, "fci", None,
        "scf",
        [[[0,0],[3,3]],[[1,1],[2,2]], [[0,1],[2,3]]], [0.5, 0.5, [1,2]],
        [-2.8137710, -2.8142425, -2.8142424, -2.8137710]),  # pyscf/FCI/QN/scf
    # PySCF BE, not incore
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, False, False, "FCI", None,
        "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127889, -2.8149386, -2.8127889]),  # pyscf/FCI/QN/scf
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, False, False, "mp2", None,
        "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8123768, -2.8133190, -2.8123768]),  # pyscf/mp2/QN/scf
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, False, False, "ccsd", None,
        "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127582, -2.8149148, -2.8127582]),  # pyscf/ccsd/QN/scf
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, False, False, "cisd", None,
        "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8127194, -2.8148940, -2.8127194]),  # pyscf/cisd/QN/scf
    # PySCF BE, nibath
    ("geom/h10.zmat", "sto-3g", [0,1,2], False, True, True, "FCI", None,
        "scf",
        [[[0,0],[2,2]],[[1,1]]], [0.5, 0.5],
        [-2.8207132177, -2.8129821608, -2.8207131811]),
    ])
def test_be(geom, basis, fsites, frank, incore, nibath, solver, jac, B0,
    bad_con, good_con, e_ps_ref):
    verbose = 2
    if frank:
        from frankenstein import molecule, scf
        from frankenstein.be import SD
        mol = molecule.MOL(geom, basis, ao_direct=False)
        mf = scf.RHF(mol, opt="diis")
        mf.kernel()
        msd = SD(mf, fsites, nibath=nibath, verbose=verbose)
    else:
        from frankenstein.tools.pyscf_utils import get_pymol
        from pyscf import scf as pyscf_scf
        from frankenstein.pyscf_be import pySD
        pymol = get_pymol(geom, basis)
        pymf = pyscf_scf.RHF(pymol)
        pymf.kernel()
        msd = pySD(pymf, fsites, incore=incore, nibath=nibath, verbose=verbose)

    mb = BE(msd, imp_sol=solver, jac=jac, B0=B0,
        bad_con=bad_con, good_con=good_con)
    mb.kernel()
    print(np.diag(mb.rdm1s))

    mb.msd.delete_eri()
    mb.msd.delete_erifile()

    assert(np.allclose(mb.e_persite, e_ps_ref))


@pytest.mark.parametrize(
    "geom, basis, natm, incore, nibath, frzcmo, fast_ao2mo, auxbasis, mtl,"
        "imp_sol, B0, ec_ref",
    [
    ("geom/c6h6.xyz", "sto-3g", 2, True, False, False, False, None,
        [["intra", "1epop"]], "mp2", "scf", -0.3766784783), # incore, int_bath
    ("geom/c6h6.xyz", "sto-3g", 2, False, False, False, False, None,
        [["intra", "1epop"]], "mp2", "scf", -0.3766784783), # outcore, int_bath
    ("geom/c6h6.xyz", "sto-3g", 2, True, False, True, True, None,
        [["intra", "1epop"]], "mp2", "scf", -0.3769771553), # incore, int_bath
    ("geom/c6h6.xyz", "sto-3g", 2, True, False, True, True, "def2-svp-c",
        [["intra", "1epop"]], "mp2", "scf", -0.3769766986), # incore, int_bath
    ("geom/si10h16.xyz", "sto-3g", 1, True, False, True, True, "def2-svp-c",
        [["intra", "1epop"]], "mp2", "scf", -0.4488356791), # incore, int_bath
    ])
def test_beatm(geom, basis, natm, incore, nibath, frzcmo, fast_ao2mo, auxbasis,
    mtl, imp_sol, B0, ec_ref):
    from frankenstein.tools.pyscf_utils import get_pymol
    from pyscf import scf as pyscf_scf
    from frankenstein.pyscf_be.pysd_atom import pySDATOM
    from frankenstein.be.bemol import BEatom
    pymol = get_pymol(geom, basis)
    pymf = pyscf_scf.RHF(pymol)
    pymf.kernel()
    msdatm = pySDATOM(pymf, natm, incore=incore, nibath=nibath,
        frzcore_mo=frzcmo, matchtypelist=mtl, auxbasis=auxbasis,
        fast_ao2mo=fast_ao2mo, verbose=2)
    mbatm = BEatom(msdatm, imp_sol=imp_sol, B0=B0)
    mbatm.kernel()

    assert(np.abs(mbatm.e_corr-ec_ref) < 1.E-4)


@pytest.mark.parametrize(
    "geom, basis, natm, rdm_type, mtl, ec_ref",
    [("geom/c6h6.xyz", "sto-3g", 1, "relaxed", [["intra", "1epop"]],
        -0.5042406923),
    ("geom/c6h6.xyz", "sto-3g", 1, "unrelaxed1", [["intra", "1epop"]],
        -0.5068129669),
    ("geom/c6h6.xyz", "sto-3g", 1, "unrelaxed2", [["intra", "1epop"]],
        -0.4264953732)]
    )
def test_ccsd_rdm(geom, basis, natm, rdm_type, mtl, ec_ref):
    from frankenstein.tools.pyscf_utils import get_pymol
    from pyscf import scf as pyscf_scf
    from frankenstein.pyscf_be.pysd_atom import pySDATOM
    from frankenstein.be.bemol import BEatom
    pymol = get_pymol(geom, basis)
    pymf = pyscf_scf.RHF(pymol)
    pymf.kernel()
    msdatm = pySDATOM(pymf, natm, fast_ao2mo=True,
        frzcore=False, frzcore_mo=True, matchtypelist=mtl, verbose=2)
    mbatm = BEatom(msdatm, B0="scf", imp_sol="ccsd",
        sol_params={"rdm":rdm_type},)
    mbatm.kernel()

    assert(np.abs(mbatm.e_corr-ec_ref) < 1.E-4)
