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

If you have separate alpha/beta-spin pdos files (spin-unrestricted calculation),
pass both files as arguments to get one common grid for both of them.
"""

# Copyright (c) 2020 Tiziano Müller <tiziano.mueller@chem.uzh.ch>,
# based on a Fortran tool written by Marcella Iannuzzi <marcella.iannuzzi@chem.uzh.ch>


import argparse
import contextlib
import re
import sys

import numpy as np

HEADER_MATCH = re.compile(
    r"\# Projected DOS for (?:atomic kind \w+|list \d+ of \d+ atoms,)"
    r" 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 cp2k_pdos():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "pdosfilenames",
        metavar="<PDOS-file>",
        type=str,
        nargs="+",
        help="the PDOS file generated by CP2K, specify two files (alpha/beta) in a spin-unrestricted case",
    )
    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(
                    f"The file '{pdosfilename}' 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/cp2k/cp2k-output-tools/issues"
                )
                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(
        f"Integration step:  {args.de:14.5f}\n"
        f"Sigma:             {args.sigma:14.5f}\n"
        f"Minimum energy:    {emin+margin:14.5f} - {margin:.5f}\n"
        f"Maximum energy:    {emax-margin:14.5f} + {margin:.5f}\n"
        f"Nr of mesh points: {nmesh:14d}",
        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(f"Nr of lines:       {data.shape[0]:14d} in {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)


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