import exatomic
from exatomic.base import resource         # Easy access to static files
from exatomic import UniverseWidget as UW  # The visualization system

"""
Explanation: Visualize Orbitals
End of explanation
"""


from exatomic import gaussian

uni = gaussian.Output(resource('g09-ch3nh2-631g.out')).to_universe()
fni = gaussian.Fchk(resource('g09-ch3nh2-631g.fchk')).to_universe()

uni.atom # Inspect the geometry

uni.basis_set_order.head() # And the first few basis functions

# Let's add the first 10 molecular orbitals 
uni.add_molecular_orbitals(vector=range(10))
fni.add_molecular_orbitals(vector=range(10))

"""
Explanation: Quantum chemistry codes don't always make it easy to get all the necessary information but provided the outputs have the data and the parser is implemented to handle the print-out from a specific code, this is what the API looks like for orbital viewing.
- First you parse the output 
     - must contain full basis set specification and a matrix of basis function coefficients
ed = exatomic.myqmcode.Output('/path/to/my/output')

Inspect the parsed data. Does it look correct?

ed.basis_set         # houses primitive exponents, coefficients
ed.momatrix          # contains the basis function coefficients
ed.basis_set_order   # indices specifying full basis set order
- ed is an Editor, an object that makes parsing files fun!
    - ed must be converted to a universe for the magic to happen
uni = ed.to_universe()
- A Universe has an add_molecular_orbitals method
uni.add_molecular_orbitals?   # to see all keyword arguments
uni.add_molecular_orbitals()  # tries to guess smart defaults

Don't forget python uses 0-based indexing, so check your vector values!

Gaussian API
End of explanation
"""


UW(uni, fni)

"""
Explanation: The UniverseWidget accepts up to 9 universes

Play around with the buttons:
Clear removes contents from the scene. (Press Fill to get the geometry back)
Active Scenes controls which scenes are controlled by the buttons (default all)
Image allows for saving PNGs of the scenes (can specify a directory and file names)
Camera allows to link the camera between scenes (load and save are buggy)
Fill switches between WebGL fragment shaders and fancy three.js Spheres
Axis adds a unit axis (often hidden within an atom if centered at the origin)
Animate will play dynamics (if a universe has multiple frames)
Fields expands to show:
Isosurfaces tab shows the molecular orbitals
Contours tab shows contour lines in x, y, or z axes




Be sure to use the close button to clean up the javascript side of things
End of explanation
"""


uo2 = gaussian.Output(resource('g09-uo2.out')).to_universe()
uo2.add_molecular_orbitals()
UW(uo2)

"""
Explanation: Recursive solid harmonics means no hard-coded maximum angular momentum
End of explanation
"""


from exatomic import molcas

mol = molcas.Output(resource('mol-uo2-anomb.out'))
orb = molcas.Orb(resource('mol-uo2-anomb.scforb'))
# Just attach it to the universe
# mol.momatrix = orb.momatrix
# mol.orbital = orb.orbital

# Alternatively there's a convenience method on molcas.Output
mol.add_orb(resource('mol-uo2-anomb.scforb'))  # adds momatrix and orbital
uni = mol.to_universe()

uni.add_molecular_orbitals(vector=range(40, 60))

UW(uni)

"""
Explanation: Molcas API

You must include the BSSHOW keyword in the SEWARD module or GATEWAY with a higher than normal print threshold to print the basis set specification
Molcas prints out Orb files which contain basis function coefficients so we use 2 Editors
End of explanation
"""


from exatomic import nwchem

nw = nwchem.Output(resource('nw-ch3nh2-631g.out')).to_universe()

nw.add_molecular_orbitals(vector=range(20))
UW(nw)

"""
Explanation: NWChem API
End of explanation
"""


