import psi4
import forte
import forte.utils

"""
Explanation: Forte Tutorial 2.00: Visualizing orbitals in Jupyter notebooks

In this tutorial we are going to explore how to generate orbitals and visualize them in Jupyter notebooks using the Python API.
Let's import the psi4 and forte modules, including the forte.utils submodule
End of explanation
"""


xyz = """
0 1
C      0.722300    0.722300    0.000000
C      0.722300   -0.722300    0.000000
C     -0.722300    0.722300    0.000000
C     -0.722300   -0.722300    0.000000
H      1.483571    1.483571    0.000000
H      1.483571   -1.483571    0.000000
H     -1.483571    1.483571    0.000000
H     -1.483571   -1.483571    0.000000
symmetry c1"""

E_scf, wfn = forte.utils.psi4_scf(xyz, 'sto-3g', 'rhf', functional = 'hf')

"""
Explanation: Note: in this tutorial we do not call forte.startup()/cleanup() because we are not going to do any computation using forte
We will start by generating SCF orbitals via psi4 using the forte.util
End of explanation
"""


forte.utils.psi4_cubeprop(wfn,path='cubes',nocc=2,nvir=2)

"""
Explanation: Generating the cube files using the forte.utils.psi4_cubeprop function
Next we use a convenience function that generates cube files. Here is some information on it:
```python
def psi4_cubeprop(wfn, path = '.', orbs = [], nocc = 3, nvir = 3, load = False):
    """
    Run a psi4 cubeprop computation to generate cube files from a given Wavefunction object
    By default this function plots from the HOMO -2 to the LUMO + 2
Parameters
----------
wfn : psi4Wavefunction
    A psi4 Wavefunction object
path : str
    The path of the directory that will contain the cube files
orbs : list
    The list of orbitals to convert to cube files (one based).
nocc : int
    The number of occupied orbitals
nvir : int
    The number of virtual orbitals

```
We want to plot two occupied and two virtual orbitals so we pass the arguments nocc=2 and nvir=2.
End of explanation
"""


import os
os.listdir('./cubes')

"""
Explanation: After we execute this command, four cube files are generated in the folder ./cubes. We can list them with the os.listdir function:
End of explanation
"""


cube = forte.CubeFile('cubes/Psi_a_15_15-A.cube')

"""
Explanation: Reading cube files
To load these cube files we will use the CubeFile class. This class reads a cube file and lets us manipulate them. Let's load the HOMO, which is orbital 15 (the cube files use 1-based indexing)
End of explanation
"""


# number of atoms
print(f'cube.natoms() -> {cube.natoms()}')

# number of grid points along each direction
print(f'cube.num() -> {cube.num()}')

"""
Explanation: A CubeFile object cannot be directly visualized. However, we can access all the information about the cube file.
End of explanation
"""


levels85 = cube.compute_levels('mo',0.85)
levels85

"""
Explanation: The CubeFile class provides a convenient method (compute_levels) to determine the isocontour levels that contain a given fraction of the electron density. For example, to compute the isovalues that encompass 85% of the density we call compute_levels('mo',0.85), where the argument 'mo' indicates that the cube file stores the values of a molecular orbital (this is necessary, because the levels are computed differently for the density)
End of explanation
"""


with open('cubes/Psi_a_15_15-A.cube') as f:
    print("".join(f.readlines()[:6]))

"""
Explanation: This information is also generated and stored by psi4 in the cube file (this is not standard, only psi4 implements this feature)
End of explanation
"""


dens = forte.CubeFile(cube)
dens.pointwise_product(dens)

"""
Explanation: Manipulating cube files
The CubeFile class supports three type of operations:
- scale(double factor): scale all the values on the grid by factor:
$$ \phi(\mathbf{r}_i) \leftarrow \mathtt{factor} * \phi(\mathbf{r}_i) $$
- add(CubeFile cube): add to this cube file the grid values stored in cube:
$$ \phi(\mathbf{r}_i) \leftarrow \phi(\mathbf{r}_i) + \psi(\mathbf{r}_i) $$
- pointwise_product(CubeFile cube): multiply each point of this cube file with the values stored in cube:
$$ \phi(\mathbf{r}_i) \leftarrow \phi(\mathbf{r}_i) * \psi(\mathbf{r}_i) $$
To evaluate the density corresponing to the HOMO we can do just compute the pointwise product with itself
End of explanation
"""


dens.save('cubes/dens.cube')

"""
Explanation: and then save a cubefile to disk using the save function
End of explanation
"""


