#!/usr/bin/env python

# HQ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
# HQ X
# HQ X   quippy: Python interface to QUIP atomistic simulation library
# HQ X
# HQ X   Copyright James Kermode 2010
# HQ X
# HQ X   These portions of the source code are released under the GNU General
# HQ X   Public License, version 2, http://www.gnu.org/copyleft/gpl.html
# HQ X
# HQ X   If you would like to license the source code under different terms,
# HQ X   please contact James Kermode, james.kermode@gmail.com
# HQ X
# HQ X   When using this software, please cite the following reference:
# HQ X
# HQ X   http://www.jrkermode.co.uk/quippy
# HQ X
# HQ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

from qlab import *

__all__ = dir()

import optparse

p = optparse.OptionParser(usage='%prog [options] <trajectory file>...')

p.add_option('-f', '--frame', action='append', help='Initial frame to show. Can be given separately for each trajectory file. Default is last frame in file.', type=int)
p.add_option('-p', '--property', action='append', help='Property to show. Can be given separately for each trajectory file')
p.add_option('-a', '--arrows', action='append', help='Property to use to draw arrows. Can be given separately for each trajectory file')
p.add_option('-l', '--load-view', action='store', help='Load view from AtomEye command script')
p.add_option('-W', '--width', action='store', help="""Width of output movie, in pixels.""", type='int')
p.add_option('-H', '--height', action='store', help="""Height of output movie, in pixels.""", type='int')
p.add_option('-A', '--aspect', action='store', help="""Aspect ratio. Used if only one of --width or --height is given. Default 0.75.""", default=0.75, type='float')
p.add_option('-R', '--rcut', action='append', help="""Following three arguments should be SYM1 SYM2 INCREMENT, to increment cutoff distance for SYM1-SYM2 bonds.""", nargs=3)
p.add_option('-L', '--loadall', action='store_true', help="""Load all frames from all trajectories into memory on startup.""")
p.add_option('-n', '--no-interactive', action='store_true', help="""Exit when done, without starting an interactive shell.""")
p.add_option('-e', '--exec-code', action='append', help="""Python statements to execute on startup""")
p.add_option('-c', '--chain', action='store_true', help="""Chain all input files together""")
p.add_option('-t', '--time-ordered-series', action='store_true', help="""Join all input files, ordering by time and discarding duplicates""")
p.add_option('--view-args', action='store', help="""Extra arguments to be passed to view() routine.""")
p.add_option('--atom-range', action='store', help="""Range of atoms to include: should be followed by two arguments, min
and max 1-based atom indices. Using a reduced atom range dramatically
speeds up reading of large XYZ trajectories.""", nargs=2)

opt, args = p.parse_args()

if opt.chain and opt.time_ordered_series:
    p.error('-c/--chain and -t/--time-ordered-series are mutually exclusive')

view_args = {}
if opt.atom_range is not None:
    view_args['range'] = opt.atom_range
if opt.view_args is not None:
    view_args.update(eval("dict(%s)" % opt.view_args))

viewers = []
if opt.chain:
    v = view(args, loadall=opt.loadall, name='at', **view_args)
    __all__.append('at')
    viewers.append(v)
elif opt.time_ordered_series:
    try:
        pattern, = args
    except:
        p.error('when -t/--time-ordered-series is present, only one argument should be given')
    v = view(time_ordered_series(pattern), loadall=opt.loadall, name='at', **view_args)
    __all__.append('at')
    viewers.append(at)    
else:
    for filename in args:
        v = view(filename, loadall=opt.loadall, **view_args)
        __all__.append(v.name)
        viewers.append(v)

if opt.frame is None:
    opt.frame = [-1 for viewer in viewers]

show_args_list = [{} for viewer in viewers]
for arg in ['frame', 'property', 'arrows']:
    values = getattr(opt, arg)
    if values is None:
        continue
    if len(values) == 1:
        values = [values[0] for traj in viewers]

    if len(values) != len(viewers):
        p.error('Number of -%s/--%s options does not match number of trajectory files' % (arg[0], arg))

    for show_args, value in zip(show_args_list, values):
        show_args[arg] = value

for traj, show_args in zip(viewers, show_args_list):
    if opt.load_view is not None:
        print 'Loading view script %s' % opt.load_view
        traj.load_script(opt.load_view)

    if opt.rcut is not None:
        print 'Applying rcut patches %r' % opt.rcut
        for (sym1, sym2, rcut) in opt.rcut:
            traj.rcut_patch(sym1, sym2, float(rcut))

    if opt.width is not None or opt.height is not None:
        if opt.width  is None: opt.width = int(opt.height/opt.aspect)
        if opt.height is None: opt.height = int(opt.width*opt.aspect)
        print 'Setting viewer size to %d x %d pixels' % (opt.width, opt.height)
        traj.resize(opt.width, opt.height)

    if len(show_args) != 0:
        print 'Applying show_args=%r to trajectory %s' % (show_args, traj.name)
        traj.show(**show_args)

if opt.exec_code is not None:
    for line in opt.exec_code:
        exec(line)

if not opt.no_interactive:
    _interactive = True
    __all__.append('_interactive')

# clean up environment
for symbol in dir():
    if symbol not in __all__:
        del symbol

if _interactive:
    del _interactive
    from IPython import embed
    embed()

