#!/usr/bin/env python
from sys import exit
import string
from optparse import OptionParser
from math import log10
from Numeric import sign
 
from Dacapo import Dacapo
from ASE.IO.Cube import WriteCube
 
def getNumberOfSpins(calc):
  if (calc.GetSpinPolarized()):
    return 2
  else:
    return 1
 
def index_error(idx, calc):
  (nband, nkpt, nspin) = idx
  NumberOfBands = calc.GetNumberOfBands()
  NumberOfKpoints = len(calc.GetIBZKPoints())
  isSpinPolarized = calc.GetSpinPolarized()
  error = False
  if not (nband < NumberOfBands):
    print "error: nband must be less than %d" % NumberOfBands
    error = True
  if not (nkpt < NumberOfKpoints):
    print "error: nkpt must be less than %d" % NumberOfKpoints 
    error = True
  if isSpinPolarized:
    if not ((nspin == 0) or (nspin == 1)):
      print "error: nspin must be 0 or 1"
      error = True
  else:
    if not (nspin == 0):
      print "error: nspin must be 0"
      error = True
  return error
 
def getRealWavefunction(idx, calc):
  (nb, nk, ns) = idx
  wf = calc.GetWaveFunctionArray(band = nb, kpt = nk, spin = ns)
  density = sign(wf.real)*abs(wf)
  return density
 
def getHighestOccupiedState(calc):
  NumberOfBands = calc.GetNumberOfBands()
  NumberOfKpoints = len(calc.GetIBZKPoints())
  NumberOfSpins = getNumberOfSpins(calculator)
  HighestOccupiedEnergy = calc.GetPotentialEnergy()
  idx = (0, 0, 0)
  for nspin in range(NumberOfSpins):
    for nkpt in range(NumberOfKpoints):
      energies = calc.GetEigenvalues(kpt = nkpt, spin = nspin) 
      for nband in range(NumberOfBands):
        if energies[nband] < 0.0:
          if energies[nband] > HighestOccupiedEnergy:
            HighestOccupiedEnergy = energies[nband]
            idx = (nband, nkpt, nspin)
  return idx
 
def indexed_file(given_filename, idx, calc, default_extension):
  (nband, nkpt, nspin) = idx
  NumberOfBands = calculator.GetNumberOfBands()
  NumberOfKpoints = len(calc.GetIBZKPoints())
  NumberOfSpins = getNumberOfSpins(calculator)
  filename = string.split(given_filename, '.')
  basename = filename[0]
  if len(filename) > 1:
    extension = filename[-1]
  else:
    extension = default_extension
  retrunname = '%s' % basename
  fmts = '%%0%dd' % (log10(NumberOfBands) + 1)
  retrunname += fmts % nband
  if (NumberOfKpoints > 1):
    fmts = 'k%%0%dd' % (log10(NumberOfKpoints) + 1)
    retrunname += fmts % nkpt
  if (NumberOfSpins > 1):
    retrunname += 's%1d' % nspin
  retrunname += '.%s' % extension
  return retrunname
 
cmd = OptionParser(usage = '%prog [-i|-a] input_nc_file output_cube_file')
cmd.add_option('-a', '--all', action = "store_true", default = False,
               help = 'create cube files for all wavefunctions')
cmd.add_option('-i', '--indecies', type = 'int', nargs = 3,
               help = 'density of N-th Band, K-th Kpoint, and S-th Spin state',
               metavar = 'N K S')
 
 
(opt, argv) = cmd.parse_args()
 
if len(argv) != 2:
    cmd.print_help()
    raise SystemExit
 
if (opt.indecies is not None) and (opt.all is True):
  print "error: do not specify both options -i and -a"
  error = True
  exit()
 
ncfile = argv[0]
model = Dacapo.ReadAtoms(ncfile)
calculator = model.GetCalculator()
 
if opt.all is True:
  NumberOfBands = calculator.GetNumberOfBands()
  NumberOfKpoints = len(calculator.GetIBZKPoints())
  NumberOfSpins = getNumberOfSpins(calculator)
  for nband in range(NumberOfBands):
    for nkpt in range(NumberOfKpoints):
      for nspin in range(NumberOfSpins):
        idx = (nband, nkpt, nspin)
        density = getRealWavefunction(idx, calculator)  
        cubefile = indexed_file(argv[1], idx, calculator, 'cube')
        WriteCube(model, density, cubefile)
else:
  if opt.indecies is not None:
    idx = opt.indecies
    if index_error(idx, calculator):
      exit()
  else:
    idx = getHighestOccupiedState(calculator)
  density = getRealWavefunction(idx, calculator)  
  cubefile = argv[1]
  WriteCube(model, density, cubefile)

usage: getwfs [-i|-a] input_nc_file output_cube_file
 
options:
  -h, --help            show this help message and exit
  -a, --all             create cube files for all wavefunctions
  -iN K S, --indecies=N K S
                        density of N-th Band, K-th Kpoint, and S-th Spin state