#
# NOTE: THESE PATCHES ARE NEEDED ONLY FOR GULP 6.1.2
# AND SHOULD BE DELETED WHEN THE NEXT RELEASE OF GULP
# BECOMES AVAILABLE
#

#  Copyright (C) 2023 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell 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 Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

# you.lu@stfc.ac.uk

from os import path

def patch_node2var(gulp_src_dir, filename):
    '''Src/changemaxat.F90 and Src/changemaxr1at.F90'''

    # patch source file
    filename = path.join(gulp_src_dir, filename)
    print(f'\n ChemShell: Patching {filename}...')
    p1 = '  integer(i4)       :: size_node2var'
    p2 = '''  if (associated(node2var)) then
    size_node2var = size(node2var)
  else
    size_node2var = 0
  endif'''
    with open(filename, 'r+') as fp:
        lines = fp.readlines()
        skip1 = False
        skip2 = False
        for i, l in enumerate(lines):
            # skip if already patched
            if p1 in l:
                skip1 = True
            if 'associated(node2var)' in l:
                skip2 = True
            if skip1 and skip2:
                break
            # insert the line as declaration
            if 'implicit none' in l:
                # TWK need to look ahead one line here to check if already patched
                if p1 in lines[i+1]:
                    skip1 = True
                # insert after
                if not skip1:
                    lines.insert(i+1, p1+'\n')
            if 'size(node2var)' in l:
                if not skip2:
                    lines[i] = l.replace('size(node2var)', 'size_node2var')
                    # insert before
                    lines.insert(i, p2+'\n')
                    break
        # update the file
        fp.seek(0)
        fp.writelines(lines)


# TWK 17/11/2023 Patch MEAM initialisation (with fix from GULP 6.1.3)
def patch_meam(gulp_src_dir, filename):
    '''Patch MEAM initialisation in Src/initial.F90 and Src/potwordm.F90'''

    # patch source file
    filename = path.join(gulp_src_dir, filename)
    print(f'\n ChemShell: Patching MEAM initialisation in {filename}...')

    # initial.F90
    p1 = '''!
!  EAM/MEAM value initialisation
!
  ndenfncomp(1:maxeamspec) = 0'''
    s1 = "Set array of pointers from atomic number to REBO species"

    # potwordm.F90
    p2 = '''  if (meamorder.gt.1) then
    lmeamspec(neamspecloc) = .true.
  else
    lmeamspec(neamspecloc) = .false.
  endif'''
    s2 = "if (meamorder.gt.1) lmeamspec(neamspecloc) = .true."

    with open(filename, 'r+') as fp:
        lines = fp.readlines()
        skip1 = False
        skip2 = False
        for i, l in enumerate(lines):
            # skip if already patched
            if "MEAM value initialisation" in l:
                break
            # insert the new lines
            if s1 in l:
                # insert before
                lines.insert(i-1, p1+'\n')
                break
            if s2 in l:
                # delete old line
                lines[i] = l.replace(s2, '')
                # insert before
                lines.insert(i, p2+'\n')
                break
        # update the file
        fp.seek(0)
        fp.writelines(lines)


# run as a standalone script
if __name__ == '__main__':

    from sys import argv

    # get path
    try:
        gulp_src_dir = argv[1]
    except:
        pass

    # get filename
    try:
        filename = argv[2]
    except:
        pass

    if filename in ["changemaxat.F90", "changemaxr1at.F90"]:
        patch_node2var(gulp_src_dir, filename)
    elif filename in ["initial.F90", "potwordm.F90"]:
        patch_meam(gulp_src_dir, filename)


