import pandas as pd
import numpy as np
import exatomic

"""
Explanation: Welcome to exatomic! Let's get started
End of explanation
"""


u = exatomic.Universe()
u

"""
Explanation: Here are some test demo containers to play around with
End of explanation
"""


#myxyz = exatomic.XYZ('../data/examples/porphyrin.xyz')
myxyz = exatomic.XYZ('porphyrin.xyz')

myxyz.head()

"""
Explanation: exatomic universes in principle contain a QM/MD calculation or set of calculations
The following dataframes are currently supported as properties of the universe with their associated required columns
<li> Frame -- ['atom_count']
<li> Atom -- ['symbol', 'x', 'y', 'z', 'frame']
<li> Two -- ['distance', 'atom0', 'atom1', 'frame']

This constitutes all the required information to visualize an MD trajectory (or geometry optimization, etc.). However, there are more dataframes that allow for increased functionality.
<li> AtomicField
<li> Molecule
<li> Overlap
<li> BasisSet

An exhaustive list can be found in the documentation or on readthedocs.org

### There are convenience methods for immediate access to your data
exatomic.XYZ('/path/to/xyz/or/trajectory')

exatomic.Cube('/path/to/cube/file')

### Try it out!
End of explanation
"""


myxyz.atom.head()    # Atomic units are used throughout the exatomic package

myuni = myxyz.to_universe()
myuni.two.head()

myuni

"""
Explanation: Just a textfile...?
End of explanation
"""


from exa.relational import Isotope
import random

nat = 10**4    # Be careful changing this value...
x = nat**0.5 * np.random.rand(nat)
y = nat**0.5 * np.random.rand(nat)
z = nat**0.5 * np.random.rand(nat)
symbols = Isotope.to_frame().drop_duplicates('symbol')['symbol'].tolist()
symbol = [random.choice(symbols) for i in range(nat)]
atom = pd.DataFrame.from_dict({'x': x, 'y': y, 'z': z, 'symbol': symbol})
atom['frame'] = 0
scuni = exatomic.Universe(atom=atom)

scuni

"""
Explanation: There we go. Our porphyrin looks pretty good. Check out the GUI controls in the animation
So what happened above?
exatomic.XYZ is a wrapper around exatomic.Editor, the base class for dealing with file I/O in exatomic. The base class has a to_universe method which converts an exatomic.Editor to an exatomic.Universe, which ships our data to javscript to be visualized right in a widget in the notebook.
So... Avogadro in the notebook? Surely it won't scale...
End of explanation
"""


