#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2020 Google
End of explanation
"""


try:
    import recirq
except ImportError:
    !pip install --quiet git+https://github.com/quantumlib/ReCirq
        
try:
    import openfermionpyscf
except ImportError:
    !pip install --quiet openfermionpyscf~=0.5.0

"""
Explanation: Making Molecule Files for HFVQE
In this tutorial we describe the process of making the molecular data files necessary to run the HFVQE code.  We focus on how to use the OpenFermion plugin modules to generate molecular files with canonical Hartree-Fock and generate integrals in a given atomic orbital basis set.  We also provide helper functions to run variational Hartree-Fock simulating the experiment and generating initial parameters.
This tutorial will follow the code in recirq/hfvqe/molecular_data/ for constructing MolecularData objects and getting atomic orbital integrals.
In addition to the standard requirement of ReCirq and its dependencies, this notebook uses OpenFermion-pyscf (and pyscf) to compute some quantities. We install it below if you don't already have it.
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://quantumai.google/cirq/experiments/hfvqe/molecular_data"><img src="https://quantumai.google/site-assets/images/buttons/quantumai_logo_1x.png" />View on QuantumAI</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/quantumlib/ReCirq/blob/master/docs/hfvqe/molecular_data.ipynb"><img src="https://quantumai.google/site-assets/images/buttons/colab_logo_1x.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/quantumlib/ReCirq/blob/master/docs/hfvqe/molecular_data.ipynb"><img src="https://quantumai.google/site-assets/images/buttons/github_logo_1x.png" />View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/ReCirq/docs/hfvqe/molecular_data.ipynb"><img src="https://quantumai.google/site-assets/images/buttons/download_icon_1x.png" />Download notebook</a>
  </td>
</table>

Setup
End of explanation
"""


import os
import numpy as np
import scipy

from recirq.hfvqe.molecular_data.molecular_data_construction import (
    h6_linear_molecule, h8_linear_molecule, 
    h10_linear_molecule, h12_linear_molecule, 
    get_ao_integrals)
from recirq.hfvqe.gradient_hf import rhf_minimization, rhf_func_generator
from recirq.hfvqe.objective import RestrictedHartreeFockObjective, generate_hamiltonian

def make_rhf_objective(molecule):
    S, Hcore, TEI = get_ao_integrals(molecule)
    _, X = scipy.linalg.eigh(Hcore, S)

    molecular_hamiltonian = generate_hamiltonian(
        general_basis_change(Hcore, X, (1, 0)),
        numpy.einsum('psqr', general_basis_change(TEI, X, (1, 0, 1, 0)),
                     molecule.nuclear_repulsion))

    rhf_objective = RestrictedHartreeFockObjective(molecular_hamiltonian,
                                                   molecule.n_electrons)
    return rhf_objective, S, Hcore, TEI

"""
Explanation: Now we can import the packages required for this notebook.
End of explanation
"""


from openfermion import MolecularData, general_basis_change
from openfermionpyscf import run_pyscf

n_hydrogens = 6
bond_distance = 1.3 # in Angstroms
molecule = MolecularData(
    geometry=[('H', (0, 0, i * bond_distance)) for i in range(n_hydrogens)],
    charge=0,
    basis='6-31g',
    multiplicity=1,
    description=f"linear_r-{bond_distance}")

"""
Explanation: Hydrogen Chain MolecularData
We provide helper functions in the hfvqe module to generate the Hydrogen chain data.  Each chain is constructed using OpenFermion and Psi4 via the OpenFermion-Psi4 plugin.  We will use creating H6 with equal spacing between Hydrogen atoms as an example.
End of explanation
"""


molecule = run_pyscf(molecule, run_scf=True, run_cisd=True)
print("Hartree-Fock energy:", molecule.hf_energy, 
      "\nCISD energy:", molecule.cisd_energy)

"""
Explanation: The previous lines set up the MolecularData file.  We can now use pyscf to either run a full self-consistent-field Hartree-Fock calculation or get atomic integrals.  Via Openfermion-Pyscf we provide an interface to running Hartree-Fock, coupled-cluster, second order perturbation theory, configuration-interaction singles-doubles, and full configuration interaction.  Many of these methods depend on parameters such as convergence criteria or initial vectors in the subspace expansion.  run_pyscf assumes common defaults which are appropriate for most systems. Below we will run default Hartree-Fock and CISD.
End of explanation
"""


print("Overlap Ints")
print(molecule.overlap_integrals)

print()
print("One-electron integrals")
print(molecule.one_body_integrals)

"""
Explanation: The MolecularData file holds almost all information that is required for post-Hartree-Fock correlated calculations.  For example, we provide access to integrals as attributes of MolecularData.
End of explanation
"""


oei_mo,  tei_mo = molecule.one_body_integrals, molecule.two_body_integrals
C = molecule.canonical_orbitals
S = molecule.overlap_integrals
oei_ao = general_basis_change(oei_mo, C.conj().T @ S, key=(1, 0))

print(oei_ao)

print(oei_mo)

# Use pyscf to get atomic integrals and compare to transformed integrals from above
pyscf_mol = molecule._pyscf_data['mol']
t = pyscf_mol.intor('int1e_kin')
v = pyscf_mol.intor('int1e_nuc')
hcore = t + v
assert np.allclose(hcore, oei_ao)

"""
Explanation: For the Hartree-Fock experiment we will need to get the atomic basis integrals from the molecular integrals.  We can use the identity $C^{\dagger}SC = I$ to reverse the transformation on the one and two electron integrals.
End of explanation
"""


# Transform the two-electron integrals
tei_ao = general_basis_change(tei_mo, C.conj().T @ S, key=(1, 1, 0, 0))  

# re-ordering for chem->physics storage of the integrals
eri_ao = np.einsum('ijkl->iklj', pyscf_mol.intor('int2e', aosym='s1'))  

assert np.allclose(tei_ao, eri_ao)
assert not np.allclose(tei_ao, tei_mo)

"""
Explanation: The two-electron integrals can also be acquired from the pyscf_mol object or we can use the general_basis_change to transform the two-electron integrals back into the AO basis.
End of explanation
"""


# diagonalize the AO-core Hamiltonian
_, X = scipy.linalg.eigh(oei_ao, S)  

obi = general_basis_change(oei_ao, X, (1, 0))
tbi = np.einsum('psqr', general_basis_change(pyscf_mol.intor('int2e', aosym='s1'), X, (1, 0, 1, 0)))
molecular_hamiltonian = generate_hamiltonian(obi, tbi,
                                             molecule.nuclear_repulsion)

rhf_objective = RestrictedHartreeFockObjective(molecular_hamiltonian,
                                               molecule.n_electrons)

"""
Explanation: We also provide a function in recirq.hfvqe.molecular_data.molecular_data_construction that uses psi4 to generate atomic integrals for the HF-VQE study.  Once the atomic orbital integrals are obtained we can perform the first step in setting up the HF-VQE experiment.  This involes transforming the integrals to the core-orbital basis and building an RestrictedHartreeFockObjective.
End of explanation
"""


from recirq.hfvqe.gradient_hf import rhf_minimization, rhf_func_generator
# Uses conjugate gradient to solve non-linear hartree-fock functional:
scipy_result = rhf_minimization(rhf_objective, verbose=True)  

print(molecule.hf_energy, scipy_result.fun)

print("Initial Parameters for HF-VQE Study ", scipy_result.x)

# Explicity build kappa matrix
from recirq.hfvqe.circuits import rhf_params_to_matrix
import matplotlib.pyplot as plt
kappa = rhf_params_to_matrix(scipy_result.x, len(rhf_objective.occ) + len(rhf_objective.virt), rhf_objective.occ,
            rhf_objective.virt)
plt.imshow(kappa)
plt.colorbar()

"""
Explanation: To get initial parameters we can simulate the Hartree-Fock experiment which is performing variational-Hartree-Fock theory.
End of explanation
"""


