#!/usr/bin/env python
"""Program to plot the density of states for a Dacapo calculation

Make a gnuplot:

> dos.py filename.nc

Dump the data to a file:

> dos.py filename.nc dos.dat
"""

import sys
from Dacapo import Dacapo
from ASE.Utilities.DOS import DOS
import Gnuplot as gp

# Get the filename of the Dacapo output file:
filename = sys.argv[1]

# Read in the atom and restart the calculator:
atoms = Dacapo.ReadAtoms(filename)

# Create a DOS object and extract data:
dos = DOS(atoms, width=0.2)
e = dos.GetEnergies()
dos0 = dos.GetDOS(0)
if dos.nspins == 2:
    # Spin polarized calculation:
    dos1 = dos.GetDOS(1)

if len(sys.argv) > 2:
    # Dump the numbers to a file:
    out = file(sys.argv[2], 'w')
    if dos.nspins == 1:
        for x, y in zip(e, dos0):
            print >> out, x, y
    else:
        for x, y1, y2 in zip(e, dos0, dos1):
            print >> out, x, y1, y2
else:
    # Make a gnuplot of the data:
    plot = gp.Gnuplot(persist=True)
    if dos.nspins == 1:
        plot.plot(gp.Data(e, dos0, with='lines'))
    else:
        plot.plot(gp.Data(e, dos0, with='lines'),
                  gp.Data(e, dos1, with='lines'))