#!/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)

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:
        for kpoint_set in SET_MATCH.finditer(fhandle.read()):
            filename = "{}.set-{}.csv".format(args.bsfilename,
                                              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))