#! /usr/bin/env python3
#
# Copyright 2009-2019 The VOTCA Development Team (http://www.votca.org)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

VERSION = '@PROJECT_VERSION@ #CSG_GIT_ID#'
import sys
import argparse
import re
import lxml.etree as lxml


PROGTITLE = 'THE VOTCA::XTP MAPPING FILE UPDATER'
PROGDESCR = 'Updates the CTP mapping file to VOTCA XTP'
VOTCAHEADER = '''\
==================================================
========   VOTCA (http://www.votca.org)   ========
==================================================

{progtitle}

please submit bugs to bugs@votca.org 
xtp_update_mapfile, version {version}

'''.format(version=VERSION, progtitle=PROGTITLE)


def okquit(what=''):
    if what != '':
        print(what)
    sys.exit(0)
# =============================================================================
# PROGRAM OPTIONS
# =============================================================================


class XtpHelpFormatter(argparse.HelpFormatter):
    def _format_usage(self, usage, action, group, prefix):
        return VOTCAHEADER


progargs = argparse.ArgumentParser(prog='xtp_update_mapfile',
                                   formatter_class=lambda prog: XtpHelpFormatter(
                                       prog, max_help_position=70),
                                   description=PROGDESCR)

progargs.add_argument('-sin', '--S_input',
                      dest='S_input',
                      action='store',
                      required=True,
                      type=argparse.FileType('r'),
                      default='',
                      help='Mapping file to update.')

progargs.add_argument('-sout', '--S_output',
                      dest='S_output',
                      action='store',
                      required=True,
                      type=argparse.FileType('w'),
                      default='',
                      help='Filename of new Mapping file.')


OPTIONS = progargs.parse_args()


# =============================================================================
# LXML
# =============================================================================
def RepresentsInt(s):
    try:
        int(s)
        return True
    except ValueError:
        return False


def changeframe(localframetext):
    ids = localframetext.split()
    newids = []
    for i in ids:
        if RepresentsInt(i):
            integer = int(i)
        else:
            print("Cannot convert {} to integer".format(i))
        newinteger = integer-1
        newids.append("{:d}".format(newinteger))
    return ' '.join(newids)


def changeids(atomidtext):
    re_s = re.compile(r'(\S+)')
    atoms_and_whitespaces = re_s.split(atomidtext)
    newstring = []
    for atom in atoms_and_whitespaces:
        if not atom or atom.isspace():
            newstring.append(atom)
        else:
            array = atom.split(":")
            atomtype = array[1]
            if RepresentsInt(array[0]):
                integer = int(array[0])
            else:
                print("Cannot convert {} to integer".format(array[0]))
            newinteger = integer-1
            newstring.append("{:d}:{}".format(newinteger, atomtype))
    return ''.join(newstring)


print("Reading in {}".format(OPTIONS.S_input.name))
Tree = lxml.parse(OPTIONS.S_input.name)
Root = Tree.getroot()
Molecules = Root.find("molecules")
for molecule in Molecules.iter('molecule'):
    segments = molecule.find("segments")
    for segment in segments.iter('segment'):
        fragments = segment.find("fragments")
        for fragment in fragments.iter('fragment'):
            qmatoms = fragment.find("qmatoms")
            qmatoms.text = changeids(qmatoms.text)
            mpoles = fragment.find("mpoles")
            mpoles.text = changeids(mpoles.text)
            localframe = fragment.find("localframe")
            localframe.text = changeframe(localframe.text)
print("Writing to {}".format(OPTIONS.S_output.name))
Tree.write(OPTIONS.S_output.name, pretty_print=True)
sys.exit(0)