from IPython.display import display, Image
%matplotlib inline
import pygauss as pg
print 'pygauss version: {}'.format(pg.__version__)

"""
Explanation: Example Assessment
After installing PyGauss you should be able to open this IPython Notebook from;
https://github.com/chrisjsewell/PyGauss/blob/master/Example_Assessment.ipynb, and run the following...
End of explanation
"""


folder = pg.get_test_folder()
len(folder.list_files())

"""
Explanation: The test folder has a number of example Gaussian outputs to play around with.
End of explanation
"""


mol = pg.molecule.Molecule(folder_obj=folder,
                init_fname='CJS1_emim-cl_B_init.com', 
                opt_fname=['CJS1_emim-cl_B_6-311+g-d-p-_gd3bj_opt-modredundant_difrz.log',
                           'CJS1_emim-cl_B_6-311+g-d-p-_gd3bj_opt-modredundant_difrz_err.log',
                           'CJS1_emim-cl_B_6-311+g-d-p-_gd3bj_opt-modredundant_unfrz.log'],
                freq_fname='CJS1_emim-cl_B_6-311+g-d-p-_gd3bj_freq_unfrz.log',
                nbo_fname='CJS1_emim-cl_B_6-311+g-d-p-_gd3bj_pop-nbo-full-_unfrz.log', 
                atom_groups={'emim':range(20), 'cl':[20]},
                alignto=[3,2,1])

"""
Explanation: Note: the folder object will act identical whether using a local path or one on a server over ssh (using paramiko):
folder = pg.Folder('/path/to/folder', 
                ssh_server='login.server.com',
                ssh_username='username')

Single Molecule Analysis
A molecule can be created containg data about the inital geometry, optimisation process and analysis of the final configuration.
End of explanation
"""


#mol.show_initial(active=True)
vdw = mol.show_initial(represent='vdw', rotations=[[0,0,90], [-90, 90, 0]])
ball_stick = mol.show_optimisation(represent='ball_stick', rotations=[[0,0,90], [-90, 90, 0]])
display(vdw, ball_stick)

print 'Cl optimised polar coords from aromatic ring : ({0}, {1},{2})'.format(
    *[round(i, 2) for i in mol.calc_polar_coords_from_plane(20,3,2,1)])
ax = mol.plot_opt_trajectory(20, [3,2,1])
ax.set_title('Cl optimisation path')
ax.get_figure().set_size_inches(4, 3)

"""
Explanation: Geometric Analysis
Molecules can be viewed statically or interactively.
End of explanation
"""


print('Optimised? {0}, Conformer? {1}, Energy = {2} a.u.'.format(
    mol.is_optimised(), mol.is_conformer(), 
    round(mol.get_opt_energy(units='hartree'),3)))
ax = mol.plot_opt_energy(units='hartree')
ax.get_figure().set_size_inches(3, 2)
ax = mol.plot_freq_analysis()
ax.get_figure().set_size_inches(4, 2)

"""
Explanation: Energetics and Frequency Analysis
End of explanation
"""


mol2 = pg.molecule.Molecule(folder_obj=folder, alignto=[3,2,1],
            pes_fname=['CJS_emim_6311_plus_d3_scan.log', 
                       'CJS_emim_6311_plus_d3_scan_bck.log'])   
ax, data = mol2.plot_pes_scans([1,4,9,10], rotation=[0,0,90], img_pos='local_maxs', zoom=0.5)
ax.set_title('Ethyl chain rotational conformer analysis')
ax.get_figure().set_size_inches(7, 3)

"""
Explanation: Potential Energy Scan analysis of geometric conformers...
End of explanation
"""


print '+ve charge centre polar coords from aromatic ring: ({0} {1},{2})'.format(
    *[round(i, 2) for i in mol.calc_nbo_charge_center(3, 2, 1)])
display(mol.show_nbo_charges(represent='ball_stick', axis_length=0.4, 
                              rotations=[[0,0,90], [-90, 90, 0]]))

"""
Explanation: Partial Charge Analysis
using Natural Bond Orbital (NBO) analysis
End of explanation
"""


print 'Number of Orbitals: {}'.format(mol.get_orbital_count())
homo, lumo = mol.get_orbital_homo_lumo()
homoe, lumoe = mol.get_orbital_energies([homo, lumo])
print 'HOMO at {} eV'.format(homoe)
print 'LUMO at {} eV'.format(lumoe)

ax = mol.plot_dos(per_energy=1, lbound=-20, ubound=10, legend_size=12)

"""
Explanation: Density of States Analysis
End of explanation
"""


print 'H inter-bond energy = {} kJmol-1'.format(
        mol.calc_hbond_energy(eunits='kJmol-1', atom_groups=['emim', 'cl']))
print 'Other inter-bond energy = {} kJmol-1'.format(
    mol.calc_sopt_energy(eunits='kJmol-1', no_hbonds=True, atom_groups=['emim', 'cl']))
display(mol.show_sopt_bonds(min_energy=1, eunits='kJmol-1',
                            atom_groups=['emim', 'cl'],
                            no_hbonds=True,
                            rotations=[[0, 0, 90]]))
display(mol.show_hbond_analysis(cutoff_energy=5.,alpha=0.6, 
                                atom_groups=['emim', 'cl'],
                                rotations=[[0, 0, 90], [90, 0, 0]]))

"""
Explanation: Bonding Analysis
Using Second Order Perturbation Theory.
End of explanation
"""


analysis = pg.Analysis(folder_obj=folder)
errors = analysis.add_runs(headers=['Cation', 'Anion', 'Initial'], 
                               values=[['emim'], ['cl'],
                                       ['B', 'BE', 'BM', 'F', 'FE']],
            init_pattern='*{0}-{1}_{2}_init.com',
            opt_pattern='*{0}-{1}_{2}_6-311+g-d-p-_gd3bj_opt*unfrz.log',
            freq_pattern='*{0}-{1}_{2}_6-311+g-d-p-_gd3bj_freq*.log',
            nbo_pattern='*{0}-{1}_{2}_6-311+g-d-p-_gd3bj_pop-nbo-full-*.log',
            alignto=[3,2,1], atom_groups={'emim':range(1,20), 'cl':[20]}, 
            ipython_print=True)

"""
Explanation: Multiple Computations Analysis
Multiple computations, for instance of different starting conformations, can be grouped into an Analysis class and anlaysed collectively.
End of explanation
"""


fig, caption = analysis.plot_mol_images(mtype='optimised', max_cols=3,
                        info_columns=['Cation', 'Anion', 'Initial'],
                        rotations=[[0,0,90]])
print caption

"""
Explanation: Molecular Comparison
End of explanation
"""


fig, caption = analysis.plot_mol_graphs(gtype='dos', max_cols=3, 
                                        lbound=-20, ubound=10, legend_size=0,
                                        band_gap_value=False,
                                   info_columns=['Cation', 'Anion', 'Initial'])
print caption

"""
Explanation: Data Comparison
End of explanation
"""


analysis.add_mol_property_subset('Opt', 'is_optimised', rows=[2,3])
analysis.add_mol_property('Energy (au)', 'get_opt_energy', units='hartree')
analysis.add_mol_property('Cation chain, $\\psi$', 'calc_dihedral_angle', [1, 4, 9, 10])
analysis.add_mol_property('Cation Charge', 'calc_nbo_charge', 'emim')
analysis.add_mol_property('Anion Charge', 'calc_nbo_charge', 'cl')
analysis.add_mol_property(['Anion-Cation, $r$', 'Anion-Cation, $\\theta$', 'Anion-Cation, $\\phi$'], 
                               'calc_polar_coords_from_plane', 3, 2, 1, 20)
analysis.add_mol_property('Anion-Cation h-bond', 'calc_hbond_energy', 
                          eunits='kJmol-1', atom_groups=['emim', 'cl'])
analysis.get_table(row_index=['Anion', 'Cation', 'Initial'], 
                   column_index=['Cation', 'Anion', 'Anion-Cation'])

"""
Explanation: The methods mentioned for indivdiual molecules can be applied to all or a subset of these computations.
End of explanation
"""


analysis.get_table(row_index=['Anion', 'Cation', 'Initial'],
                   column_index=['Cation', 'Anion', 'Anion-Cation'],
                   as_image=True, font_size=12)

"""
Explanation: There is also an option (requiring pdflatex and ghostscript+imagemagik) to output the tables as a latex formatted image.
End of explanation
"""


ax = analysis.plot_radviz_comparison('Anion', columns=range(4, 10))

"""
Explanation: Multi-Variate Analysis
RadViz is a way of visualizing multi-variate data.
End of explanation
"""


pg.utils.imgplot_kmean_groups(
    analysis, 'Anion', 'cl', 4, range(4, 10), 
    output=['Initial'], mtype='optimised', 
    rotations=[[0, 0, 90], [-90, 90, 0]],
    axis_length=0.3)

"""
Explanation: The KMeans algorithm clusters data by trying to separate samples into n groups of equal variance.
End of explanation
"""


file_path = folder.save_ipyimg(vdw, 'image_of_molecule')
Image(file_path)

"""
Explanation: Documentation (MS Word)
After analysing the computations, it would be reasonable to want to document some of our findings. This can be achieved by outputting individual figure or table images via the folder object.
End of explanation
"""


import matplotlib.pyplot as plt
d = pg.MSDocument()
d.add_heading('A Pygauss Example Assessment', level=0)

d.add_docstring("""
# Introduction

We have looked at the following aspects
of [EMIM]^{+}[Cl]^{-} (C_{6}H_{11}ClN_{2});

- Geometric conformers
- Electronic structure

# Geometric Conformers
""")

fig, caption = analysis.plot_mol_images(max_cols=2, 
                rotations=[[90,0,0], [0,0,90]], 
                info_columns=['Anion', 'Cation', 'Initial'])
d.add_mpl(fig, dpi=96, height=9, caption=caption)
plt.close()
d.add_paragraph()
df = analysis.get_table(
        columns=['Anion Charge', 'Cation Charge'],
        row_index=['Anion', 'Cation', 'Initial'])
d.add_dataframe(df, incl_indx=True, style='Medium Shading 1 Accent 1',
                caption='Analysis of Conformer Charge')

d.add_docstring("""
# Molecular Orbital Analysis
## Density of States

It is **important** to *emphasise* that the
computations have only been run in the gas phase.
""")
fig, caption = analysis.plot_mol_graphs(gtype='dos', max_cols=3, 
                                    lbound=-20, ubound=10, legend_size=0,
                                    band_gap_value=False,
                               info_columns=['Cation', 'Anion', 'Initial'])
d.add_mpl(fig, dpi=96, height=9, caption=caption)
plt.close()

d.save('exmpl_assess.docx')

"""
Explanation: But you may also want to produce a more full record of your analysis, and this is where python-docx steps in. Building on this package the pygauss MSDocument class can produce a full document of your analysis.
End of explanation
"""


