#!/usr/bin/env python
"""
Convert the CP2K band structure output to CSV files
"""

import re
import argparse

SET_MATCH = re.compile(r'''
[ ]*
  SET: [ ]* (?P<setnr>\d+) [ ]*
  TOTAL [ ] POINTS: [ ]* (?P<totalpoints>\d+) [ ]*
  \n
(?P<content>
  [\s\S]*?(?=\n.*?[ ] SET|$)  # match everything until next 'SET' or EOL
)
''', re.VERBOSE)

SPOINTS_MATCH = re.compile(r'''
[ ]*
  POINT [ ]+ (?P<pointnr>\d+) [ ]+ (?P<a>\S+) [ ]+ (?P<b>\S+) [ ]+ (?P<c>\S+)
''', re.VERBOSE)

POINTS_MATCH = re.compile(r'''
[ ]*
  Nr\. [ ]+ (?P<nr>\d+) [ ]+
  Spin [ ]+ (?P<spin>\d+) [ ]+
  K-Point [ ]+ (?P<a>\S+) [ ]+ (?P<b>\S+) [ ]+ (?P<c>\S+) [ ]*
  \n
[ ]* (?P<npoints>\d+) [ ]* \n
(?P<values>
  [\s\S]*?(?=\n.*?[ ] Nr|$)  # match everything until next 'Nr.' or EOL
)
''', re.VERBOSE)


def convert_bs2csv(ifile, csvpattern="bs.set-{setnr}.csv"):
    """
    Convert the input from the given input file handle and write
    CSV output files based on the given pattern.
    """
    for kpoint_set in SET_MATCH.finditer(ifile.read()):
        filename = csvpattern.format(setnr=kpoint_set.group('setnr'))
        set_content = kpoint_set.group('content')

        with open(filename, 'w') as csvout:
            print(("writing point set {}"
                   " (total number of k-points: {totalpoints})"
                   .format(filename, **kpoint_set.groupdict())))

            print("  with the following special points:")
            for point in SPOINTS_MATCH.finditer(set_content):
                print("  {pointnr}: {a}/{b}/{c}".format(
                    **point.groupdict()))

            for point in POINTS_MATCH.finditer(set_content):
                results = point.groupdict()
                results['values'] = " ".join(results['values'].split())
                csvout.write("{a} {b} {c} {values}\n".format(**results))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('bsfilename', metavar='bandstructure-file', type=str,
                        help="the band structure file generated by CP2K")

    args = parser.parse_args()

    with open(args.bsfilename, 'r') as fhandle:
        convert_bs2csv(fhandle, "{}.set-{{setnr}}.csv".format(args.bsfilename))