# Path: openfermioncirq/experiments/hfvqe/gradient_hf.py
# def rhf_func_generator(rhf_objective: RestrictedHartreeFockObjective,
#                        initial_occ_vec: Optional[Union[None, np.ndarray]] = None,
#                        get_opdm_func: Optional[bool] = False):
#     """
#     Generate the energy, gradient, and unitary functions
# 
#     :param rhf_objective: objective function object
#     :param initial_occ_vec: (optional) vector for occupation numbers of the alpha-opdm
#     :return: functions for unitary, energy, gradient (in that order)
#     """
#     if initial_occ_vec is None:
#         initial_opdm = np.diag([1] * rhf_objective.nocc + [0] * rhf_objective.nvirt)
#     else:
#         initial_opdm = np.diag(initial_occ_vec)
# 
#     def energy(params):
#         u = unitary(params)
#         final_opdm_aa = u @ initial_opdm @ np.conjugate(u).T
#         tenergy = rhf_objective.energy_from_opdm(final_opdm_aa)
#         return tenergy
# 
#     def gradient(params):
#         u = unitary(params)
#         final_opdm_aa = u @ initial_opdm @ np.conjugate(u).T
#         return rhf_objective.global_gradient_opdm(params, final_opdm_aa).real
# 
#     def unitary(params):
#         kappa = rhf_params_to_matrix(params,
#                                      len(rhf_objective.occ) + len(rhf_objective.virt),
#                                      rhf_objective.occ, rhf_objective.virt)
#         return sp.linalg.expm(kappa)
# 
#     def get_opdm(params):
#         u = unitary(params)
#         return u @ initial_opdm @ np.conjugate(u).T
# 
#     if get_opdm_func:
#         return unitary, energy, gradient, get_opdm
#     return unitary, energy, gradient
# 
# def rhf_minimization(rhf_object, method='CG', initial_guess=None, verbose=True):
#     """
#     Perform Hartree-Fock energy minimization
# 
#     :param rhf_object: RestrictedHartreeFockObject
#     :param method: (optional sp opt method)
#     :return: sp result object
#     """
#     _, energy, gradient = rhf_func_generator(rhf_object)
#     if initial_guess is None:
#         init_guess = np.zeros(rhf_object.nocc * rhf_object.nvirt)
#     else:
#         init_guess = initial_guess.flatten()
# 
#     return sp.optimize.minimize(energy, init_guess, jac=gradient, method=method,
#                                    options={'disp': verbose})
# 
# Path: openfermioncirq/experiments/hfvqe/molecular_example.py
# def make_h6_1_3() -> Tuple[RestrictedHartreeFockObjective,
#                            of.MolecularData,
#                            np.ndarray,
#                            np.ndarray,
#                            np.ndarray]:
#     # load the molecule from moelcular data
#     import openfermioncirq.experiments.hfvqe as hfvqe
#     h6_1_3_path = os.path.join(
#         hfvqe.__path__[0],
#         'molecular_data/hydrogen_chains/h_6_sto-3g/bond_distance_1.3')
# 
#     molfile = os.path.join(h6_1_3_path, 'H6_sto-3g_singlet_linear_r-1.3.hdf5')
#     molecule = of.MolecularData(filename=molfile)
#     molecule.load()
# 
#     S = np.load(os.path.join(h6_1_3_path, 'overlap.npy'))
#     Hcore = np.load(os.path.join(h6_1_3_path, 'h_core.npy'))
#     TEI = np.load(os.path.join(h6_1_3_path, 'tei.npy'))
# 
#     _, X = sp.linalg.eigh(Hcore, S)
#     obi = of.general_basis_change(Hcore, X, (1, 0))
#     tbi = np.einsum('psqr', of.general_basis_change(TEI, X, (1, 0, 1, 0)))
#     molecular_hamiltonian = generate_hamiltonian(obi, tbi,
#                                                  molecule.nuclear_repulsion)
# 
#     rhf_objective = RestrictedHartreeFockObjective(molecular_hamiltonian,
#                                                    molecule.n_electrons)
# 
#     scipy_result = rhf_minimization(rhf_objective)
# 
#     return rhf_objective, molecule, scipy_result.x, obi, tbi


import numpy as np
from openfermioncirq.experiments.hfvqe.gradient_hf import (rhf_func_generator,
                               rhf_minimization)
from openfermioncirq.experiments.hfvqe.molecular_example import make_h6_1_3


def test_rhf_func_gen():
    rhf_objective, molecule, parameters, _, _ = make_h6_1_3()
    ansatz, energy, _ = rhf_func_generator(rhf_objective)
    assert np.isclose(molecule.hf_energy, energy(parameters))


# Path: openfermioncirq/experiments/hfvqe/gradient_hf.py
# def rhf_func_generator(rhf_objective: RestrictedHartreeFockObjective,
#                        initial_occ_vec: Optional[Union[None, np.ndarray]] = None,
#                        get_opdm_func: Optional[bool] = False):
#     """
#     Generate the energy, gradient, and unitary functions
# 
#     :param rhf_objective: objective function object
#     :param initial_occ_vec: (optional) vector for occupation numbers of the alpha-opdm
#     :return: functions for unitary, energy, gradient (in that order)
#     """
#     if initial_occ_vec is None:
#         initial_opdm = np.diag([1] * rhf_objective.nocc + [0] * rhf_objective.nvirt)
#     else:
#         initial_opdm = np.diag(initial_occ_vec)
# 
#     def energy(params):
#         u = unitary(params)
#         final_opdm_aa = u @ initial_opdm @ np.conjugate(u).T
#         tenergy = rhf_objective.energy_from_opdm(final_opdm_aa)
#         return tenergy
# 
#     def gradient(params):
#         u = unitary(params)
#         final_opdm_aa = u @ initial_opdm @ np.conjugate(u).T
#         return rhf_objective.global_gradient_opdm(params, final_opdm_aa).real
# 
#     def unitary(params):
#         kappa = rhf_params_to_matrix(params,
#                                      len(rhf_objective.occ) + len(rhf_objective.virt),
#                                      rhf_objective.occ, rhf_objective.virt)
#         return sp.linalg.expm(kappa)
# 
#     def get_opdm(params):
#         u = unitary(params)
#         return u @ initial_opdm @ np.conjugate(u).T
# 
#     if get_opdm_func:
#         return unitary, energy, gradient, get_opdm
#     return unitary, energy, gradient
# 
# def rhf_minimization(rhf_object, method='CG', initial_guess=None, verbose=True):
#     """
#     Perform Hartree-Fock energy minimization
# 
#     :param rhf_object: RestrictedHartreeFockObject
#     :param method: (optional sp opt method)
#     :return: sp result object
#     """
#     _, energy, gradient = rhf_func_generator(rhf_object)
#     if initial_guess is None:
#         init_guess = np.zeros(rhf_object.nocc * rhf_object.nvirt)
#     else:
#         init_guess = initial_guess.flatten()
# 
#     return sp.optimize.minimize(energy, init_guess, jac=gradient, method=method,
#                                    options={'disp': verbose})
# 
# Path: openfermioncirq/experiments/hfvqe/molecular_example.py
# def make_h6_1_3() -> Tuple[RestrictedHartreeFockObjective,
#                            of.MolecularData,
#                            np.ndarray,
#                            np.ndarray,
#                            np.ndarray]:
#     # load the molecule from moelcular data
#     import openfermioncirq.experiments.hfvqe as hfvqe
#     h6_1_3_path = os.path.join(
#         hfvqe.__path__[0],
#         'molecular_data/hydrogen_chains/h_6_sto-3g/bond_distance_1.3')
# 
#     molfile = os.path.join(h6_1_3_path, 'H6_sto-3g_singlet_linear_r-1.3.hdf5')
#     molecule = of.MolecularData(filename=molfile)
#     molecule.load()
# 
#     S = np.load(os.path.join(h6_1_3_path, 'overlap.npy'))
#     Hcore = np.load(os.path.join(h6_1_3_path, 'h_core.npy'))
#     TEI = np.load(os.path.join(h6_1_3_path, 'tei.npy'))
# 
#     _, X = sp.linalg.eigh(Hcore, S)
#     obi = of.general_basis_change(Hcore, X, (1, 0))
#     tbi = np.einsum('psqr', of.general_basis_change(TEI, X, (1, 0, 1, 0)))
#     molecular_hamiltonian = generate_hamiltonian(obi, tbi,
#                                                  molecule.nuclear_repulsion)
# 
#     rhf_objective = RestrictedHartreeFockObjective(molecular_hamiltonian,
#                                                    molecule.n_electrons)
# 
#     scipy_result = rhf_minimization(rhf_objective)
# 
#     return rhf_objective, molecule, scipy_result.x, obi, tbi


# Path: openfermioncirq/experiments/hfvqe/gradient_hf_test.py
import numpy as np
from openfermioncirq.experiments.hfvqe.gradient_hf import (rhf_func_generator,
                               rhf_minimization)
from openfermioncirq.experiments.hfvqe.molecular_example import make_h6_1_3



def test_rhf_func_gen():
    rhf_objective, molecule, parameters, _, _ = make_h6_1_3()
    ansatz, energy, _ = rhf_func_generator(rhf_objective)
    assert np.isclose(molecule.hf_energy, energy(parameters))


    ansatz, energy, _, opdm_func = rhf_func_generator(