#!/usr/bin/env python

import sys
import pylab as pl
from gpaw import GPAW
from optparse import OptionParser

p = OptionParser()
p.add_option('--width', default=0.1, type=float,
             help='gaussian width')
p.add_option('--angular',
             help='plot pdos for these angular momentum channels')
p.add_option('--cumul', action='store_true',
             help='plot cumulative dos')
p.add_option('--fermizero', action='store_true',
             help='shift Fermi level to zero')
#p.add_option('--allelectron', action='store_true',
#             help='plot all-electron pdos')

opts, args = p.parse_args()

npts = 1200

for i, arg in enumerate(args):
    calc = GPAW(arg, txt=None)
    #if opts.allelectron:
    #    N = len(calc.get_atoms())
    #    e, dos_e = calc.get_all_electron_ldos([N - 1], npts=npts,
    #                                          width=opts.width)
    if opts.angular:
        e, dos_e = calc.get_orbital_ldos(-1, npts=npts, width=opts.width,
                                         angular=opts.angular)
    else:
        e, dos_e = calc.get_dos(npts=4800, width=opts.width)
    epsF = calc.get_fermi_level()
    de = e[1] - e[0]
    cumul = (dos_e * de).cumsum()
    colors = 'bgrcmyk'
    color = colors[i % len(colors)]
    if opts.fermizero:
        e -= epsF
        epsF -= epsF
    if opts.cumul:
        pl.plot(e, cumul, lw=2, color=color, label=arg)
    else:
        spdos_e = dos_e[e > -10.0]
        e1 = e[e > -10.0]
        moment0 = spdos_e.sum() * de
        moment1 = (e1 * spdos_e).sum() * de
        print 'mom0', moment0
        print 'mom1', moment1
        print 'ratio', moment1/moment0
        dos_e[e > -10.0].sum() * de
        pl.plot(e, dos_e, lw=2, color=color, label=arg)
    pl.axvline(epsF, color=color, ls='--')
    pl.legend(loc='best')
pl.show()