circuit_exploration_with_h

"""
pip install openfermion pyscf openfermionpyscf fqe
"""
import numpy as np
import cirq
import openfermion as of
from openfermionpyscf import run_pyscf
import fqe
from fqe.openfermion_utils import integrals_to_fqe_restricted


def main():
    num_atoms = 4
    atom_type = 'H'
    bond_distance = 1.2
    molecule = of.chem.chemical_series.make_atomic_lattice(nx_atoms=2,
                                                           ny_atoms=2,
                                                           nz_atoms=1,
                                                           basis='sto-3g',
                                                           atom_type=atom_type,
                                                           spacing=bond_distance
                                                           )
    print(molecule)
    print(molecule.geometry)
    print(molecule.basis)
    molecule = run_pyscf(molecule, run_scf=True, run_fci=True)
    print(molecule.hf_energy)
    print(molecule.canonical_orbitals)

    n_electrons = molecule.n_electrons
    sz = 0
    norbs = molecule.n_orbitals

    oei, tei = molecule.get_integrals()
    print(oei)
    fermion_hamiltonian = of.get_fermion_operator(molecule.get_molecular_hamiltonian())
    fqe_ham = integrals_to_fqe_restricted(oei, tei)

    print(of.jordan_wigner(fermion_hamiltonian))

    # etc, etc, etc,

    hf_wf = fqe.Wavefunction([[n_electrons, sz, norbs]])
    hf_wf.set_wfn(strategy='hartree-fock')
    hf_wf.print_wfn()
    print(hf_wf.expectationValue(fqe_ham).real + molecule.nuclear_repulsion,
          molecule.hf_energy)  # these should be the same

    # gs_e, gs_wf = of.get_ground_state(of.get_sparse_operator(fermion_hamiltonian).real)
    # print(gs_e, molecule.fci_energy)
    # fqe_wf = fqe.from_cirq(gs_wf.flatten(), 1.0E-12)
    # fqe_wf.print_wfn()
    return of.get_sparse_operator(fermion_hamiltonian).real, of.jordan_wigner(fermion_hamiltonian)


def qcircuit_expectation(sparse_ham):
    qubits = cirq.LineQubit.range(8)
    circuit = cirq.Circuit([cirq.X.on(xx) for xx in qubits[:4]])
    circuit += cirq.Circuit([cirq.I.on(xx) for xx in qubits[4:]])

    for _ in range(0):
        rint = np.random.randint(0, 3)
        if rint == 0:
            rqubit_idx = np.random.randint(0, 8)
            circuit.append(cirq.H.on(qubits[rqubit_idx]))
        elif rint == 1:
            rqubit_idx = np.random.randint(0, 8)
            circuit.append(cirq.S.on(qubits[rqubit_idx]))
        else:
            rqubit_idx_0 = np.random.randint(0, 8)
            circuit.append(cirq.CNOT.on(qubits[rqubit_idx_0], qubits[(rqubit_idx_0 + 1) % 8]))

    print(circuit.to_text_diagram(transpose=True))
    final_wf = circuit.final_wavefunction()
    print(final_wf)


    expetation_value_ham = final_wf.reshape((-1, 1)).conj().T @ sparse_ham @ final_wf.reshape((-1, 1))
    print(expetation_value_ham)


def cirq_method(qubit_ham):
    pass


if __name__ == "__main__":
    sham, qubit_ham = main()
    # qcircuit_expectation(sham)
    cirq_method(qubit_ham)