#!/usr/bin/env python
###############################################################################
#
# Convolves band.out-like files of DFTB+ into DOS/PDOS via Gaussian broadening
#
###############################################################################
import optparse
import numpy as np
from dptools.bandout import BandOut

USAGE = """usage: %prog [options] INPUT OUTPUT

Reads the band structure information stored in a file INPUT created by DFTB+
(usually band.out for DOS and region*.out for PDOS) and convolves the
eigenlevels with gaussians to produce nice DOS/PDOS curves. The result is stored
in the file OUTPUT in NXY-format.

IMPORTANT: If you create PDOS, you have to weight the levels by the
'occupation' numbers (using option '-w'), otherwise you will obtain the total
DOS.

For spin unpolarized calculations, the output contains one Y-column only. For
spin polarized calculations the first column contains the total value, while
the further Y-columns contain the values for each spin-channel separately."""

# Default values:
DEFSIGMA = 0.1                          # broadening width
DEFGRIDSEP = 0.01                       # grid separation
DEFSIGMARANGE = 3.5                     # gaussian range to consider


def main():
    parser = optparse.OptionParser(usage=USAGE)
    parser.add_option("-b", "--broaden", type="float", dest="width",
                      default=DEFSIGMA,
                      help="broadening width sigma (default: {0:.2f})".format(
                          DEFSIGMA))
    parser.add_option("-g", "--gridres", type="float", dest="gridres",
                      default=DEFGRIDSEP,
                      help="grid separation (default: {0:.2f})".format(
                          DEFGRIDSEP))
    parser.add_option("-w", "--weight-occ", action="store_true",
                      dest="occweight", default=False,
                      help="create pdos or occupation weighted dos")
    parser.add_option("-s", "--sigma-range", type="float", dest="sigmarange",
                      default=DEFSIGMARANGE,
                      help="number of sigmas after which the Gaussian is "
                      "considered to be zero (default: {0:.2f})".format(
                          DEFSIGMARANGE))
    options, args = parser.parse_args()

    if len(args) != 2:
        parser.error("You must specify two arguments (input and the output "
                     "file name)")
    infile = args[0]
    outfile = args[1]

    bandout = BandOut.fromfile(infile)
    eigvals = bandout.eigvalarray
    if not options.occweight:
        eigvals[:,:,:,1] = 1.0

    sigma = options.width
    aa = 1.0 / (sigma * np.sqrt(2.0 * np.pi))
    bb = -0.5 / sigma**2
    sigmarange = options.sigmarange
    dsigma = sigmarange * sigma
    gridres = options.gridres

    minval = np.min(eigvals[:,:,:,0])
    maxval = np.max(eigvals[:,:,:,0])

    # First and last grid points, x-grid
    xmin = np.floor((minval - dsigma) / gridres) * gridres
    xmax = np.ceil((maxval + dsigma) / gridres) * gridres
    xvals = np.arange(xmin, xmax + gridres, gridres, dtype=float)
    nval = len(xvals)

    # Empty container for y-values
    yvals = np.zeros((bandout.nspin, nval), dtype=float)

    # Calculate the Gauss broadened values around each state on the grid.
    for ispin in range(bandout.nspin):
        for ik in range(bandout.nkpt):
            prefac = aa * bandout.kweights[ispin, ik]
            for eigval, occ in eigvals[ispin, ik]:
                # Grid points for the current curve (first, last)
                ilower = int(np.floor((eigval - dsigma - xmin) / gridres))
                iupper = int(np.ceil((eigval + dsigma - xmin) / gridres))
                dx = eigval - xvals[ilower:iupper+1]
                yvals[ispin, ilower:iupper+1] += ((prefac * occ)
                                                  * np.exp(bb * dx * dx))

    # Write resulting DOS
    fp = open(outfile, "w")
    if bandout.nspin == 1:
        for xx, yy in zip(xvals, yvals[0]):
            fp.write("{0:18.10E} {1:18.10E}\n".format(xx, yy))
    else:
        ytotal = np.sum(yvals, axis=0)
        formstr0 = "{0:18.10E} "
        tmp = [ "{" + "{0:d}".format(ii) + ":18.10E}" for ii in
                range(1,  bandout.nspin + 2) ]
        formstr = formstr0 + " ".join(tmp) + "\n"
        for ii in range(len(xvals)):
            fp.write(formstr.format(xvals[ii], ytotal[ii], *yvals[:,ii]))
    fp.close()


if __name__ == "__main__":
    main()
