#!/usr/bin/env python
# vim: set fileencoding=utf-8 ts=8 sw=4 tw=0 :

"""
Convert the discrete CP2K PDOS points to a smoothed curve using convoluted gaussians.

Also shifts the energies by the Fermi energy (so the Fermi energy will afterwards be at 0),
and normalizes by the number of atoms of this kind.
"""

# Copyright (c) 2017 Tiziano Müller <tiziano.mueller@chem.uzh.ch>,
# based on a Fortran tool written by Marcella Iannuzzi <marcella.iannuzzi@chem.uzh.ch>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.


from __future__ import print_function

import sys
import re
import argparse
import contextlib

import numpy as np


HEADER_MATCH = re.compile(
    r'\# Projected DOS for atomic kind (?P<element>\w+) at iteration step i = \d+, E\(Fermi\) = [ \t]* (?P<Efermi>[^\t ]+) a\.u\.')

# Column indexes, starting from 0
EIGENVALUE_COLUMN = 1
DENSITY_COLUMN = 3


# from https://stackoverflow.com/a/17603000/1400465
@contextlib.contextmanager
def smart_open(filename=None):
    if filename and filename != '-':
        fhandle = open(filename, 'w')
    else:
        fhandle = sys.stdout

    try:
        yield fhandle
    finally:
        if fhandle is not sys.stdout:
            fhandle.close()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('pdosfilenames', metavar='<PDOS-file>', type=str, nargs='+',
                        help="the PDOS file generated by CP2K, specify two (alpha/beta) files for a common grid")
    parser.add_argument('--sigma', '-s', type=float, default=0.02,
                        help="sigma for the gaussian distribution (default: 0.02)")
    parser.add_argument('--de', '-d', type=float, default=0.001,
                        help="integration step size (default: 0.001)")
    parser.add_argument('--scale', '-c', type=float, default=1,
                        help="scale the density by this factor (default: 1)")
    parser.add_argument('--total-sum', action='store_true',
                        help="calculate and print the total sum for each orbital (default: no)")
    parser.add_argument('--no-header', action='store_true', default=False,
                        help="do not print a header (default: print header)")
    parser.add_argument('--output', '-o', type=str, default=None,
                        help="write output to specified file (default: write to standard output)")
    args = parser.parse_args()

    alldata = []
    orb_headers = []

    for pdosfilename in args.pdosfilenames:
        with open(pdosfilename, 'r') as fhandle:
            match = HEADER_MATCH.match(fhandle.readline().rstrip())
            if not match:
                print(("The file '{}' does not look like a CP2K PDOS output.\n"
                       "If it is indeed a correct output file, please report an issue at\n"
                       "    https://github.com/dev-zero/cp2k-tools/issues").format(pdosfilename))
                sys.exit(1)

            efermi = float(match.group('Efermi'))
            # header is originally: ['#', 'MO', 'Eigenvalue', '[a.u.]', 'Occupation', 's', 'py', ...]
            header = fhandle.readline().rstrip().split()[1:]  # remove the comment directly
            header[1:3] = [' '.join(header[1:3])]  # rejoin "Eigenvalue" and its unit
            data = np.loadtxt(fhandle)  # load the rest directly with numpy

        alldata.append(data)

        orb_headers += header[DENSITY_COLUMN:]

    # take the boundaries over all energy eigenvalues (not guaranteed to be the same)
    # add a margin to not cut-off Gaussians at the borders
    margin = 10 * args.sigma
    emin = min(np.min(data[:, EIGENVALUE_COLUMN]) for data in alldata) - margin
    emax = max(np.max(data[:, EIGENVALUE_COLUMN]) for data in alldata) + margin
    ncols = sum(data.shape[1] - DENSITY_COLUMN for data in alldata)
    nmesh = int((emax-emin)/args.de)+1  # calculate manually instead of using np.arange to ensure emax inside the mesh
    xmesh = np.linspace(emin, emax, nmesh)
    ymesh = np.zeros((nmesh, ncols))

    # printing to stderr makes it possible to simply redirect the stdout to a file
    print("Integration step:  {:14.5f}".format(args.de), file=sys.stderr)
    print("Sigma:             {:14.5f}".format(args.sigma), file=sys.stderr)
    print("Minimum energy:    {:14.5f} - {:.5f}".format(emin+margin, margin), file=sys.stderr)
    print("Maximum energy:    {:14.5f} + {:.5f}".format(emax-margin, margin), file=sys.stderr)
    print("Nr of mesh points: {:14d}".format(nmesh), file=sys.stderr)

    fact = args.de/(args.sigma*np.sqrt(2.0*np.pi))

    coloffset = 0
    for fname, data in zip(args.pdosfilenames, alldata):
        print("Nr of lines:       {:14d} in {}".format(data.shape[0], fname), file=sys.stderr)
        ncol = data.shape[1] - DENSITY_COLUMN

        for idx in range(nmesh):
            func = np.exp(-(xmesh[idx]-data[:, EIGENVALUE_COLUMN])**2/(2.0*args.sigma**2))*fact
            ymesh[idx, coloffset:(coloffset+ncol)] = func.dot(data[:, DENSITY_COLUMN:])

        coloffset += ncol

    if args.total_sum:
        finalsum = np.sum(ymesh, 0)*args.de
        print("Sum over all meshpoints, per orbital:", file=sys.stderr)
        print(("{:16.8f}"*ncols).format(*finalsum), file=sys.stderr)

    xmesh -= efermi  # put the Fermi energy at 0
    xmesh *= 27.211384  # convert to eV
    ymesh *= args.scale  # scale

    with smart_open(args.output) as fhandle:
        if not args.no_header:
            print(("{:>16}" + " {:>16}"*ncols).format("Energy_[eV]", *orb_headers), file=fhandle)
        for idx in range(nmesh):
            print(("{:16.8f}" + " {:16.8f}"*ncols).format(xmesh[idx], *ymesh[idx, :]), file=fhandle)


if __name__ == '__main__':
    main()

#  vim: set ts=4 sw=4 tw=0 :