#!/usr/bin/env python3
from optparse import OptionParser
from ase.visualize import view
from ase.calculators.jacapo import Jacapo
from Scientific.IO.NetCDF import NetCDFFile as netCDF
import os

os.environ['DACAPO_READONLY'] = 'On'

parser = OptionParser(usage='ncsum',
                      version='0.1')
parser.add_option('-e',
                  nargs=0,
                  help = 'print only the energy')

parser.add_option('-p',
                  nargs=0,
                  help='view atoms')

parser.add_option('-r',
                  nargs=1,
                  help='repeat plotted atoms in n1,n2,n3 times')

options,args = parser.parse_args()

for arg in args:

    calc = Jacapo(arg)
    atoms = calc.get_atoms()

    if options.e is not None:
        nc = netCDF(arg,'r')
        energy = nc.variables.get('TotalEnergy',None)
        if energy is not None:
            print(energy[:][-1])
        else:
            print(None)
        nc.close()
    else:
        print(calc)

    if options.p is not None:
        if options.r is not None:
            n1,n2,n3 = [int(x) for x in options.r.split(',')]
        else:
            n1,n2,n3 = (1,1,1)

        view(atoms.repeat([n1,n2,n3]))

