#  Copyright (C) 2021 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.ukri.org

from os import path

def patch_makefile_h(nwchem_dest):
    '''src/config/makefile.h'''

    filename = path.join(nwchem_dest, 'src', 'config', 'makefile.h')
    print("\n ChemShell: patching {}...".format(filename))
    with open(filename, 'r+') as fp:
        a0 = '        EXPLICITF = FALSE'
        a1 = '        #EXPLICITF = FALSE'
        c0 = '        FOPTIMIZE = -O2 -O scalar3,thread0,vector2,ipa2 #-rdm'
        c1 = '        FOPTIMIZE = -O2 -O scalar3,thread2,vector2,ipa3 #-rdm\n        FCONVERT = '
        makefile = fp.read().replace(a0, a1).replace(c0, c1)
        # update the file
        fp.seek(0)
        fp.write(makefile)


def patch_mm_bond_F(nwchem_dest):
    '''src/mm/mm_bond.F'''

    filename = path.join(nwchem_dest, 'src', 'mm', 'mm_bond.F')
    print("\n ChemShell: patching {}...".format(filename))
    with open(filename, 'r+') as fp:
        a0 = 'dflim = cos(phas(n))*kphi(n)*(peri(n)-gmul(inc)+gmul(inc)*cphi)'
        a1 = 'dflim = cos(phas(n))*kphi(n)*(peri(n)-gmul(int(inc))+gmul(int(inc))*cphi)'
        mm_bond = fp.read().replace(a0, a1)
        # update the file
        fp.seek(0)
        fp.write(mm_bond)


def patch_GNUmakefile(nwchem_dest):
    '''src/GNUmakefile'''

    filename = path.join(nwchem_dest, 'src', 'GNUmakefile')
    print(f'\n ChemShell: patching {filename}...')
    with open(filename, 'r+') as fp:
        GNUmakefile = fp.readlines()
        for i, l in enumerate(GNUmakefile):
            l = l.replace('\"\'', '\"\\\'')
            # make sure only replace once
            if '\\\'\"' not in l:
                GNUmakefile[i] = l.replace('\'\"', '\\\'\"')
            else:
                GNUmakefile[i] = l
        # update the file
        fp.seek(0)
        fp.writelines(GNUmakefile)


def patch_rpspsolve_c(nwchem_dest):
    '''src/nwpw/nwpwlib/pseudopotential/rpspsolve.c'''

    # patch nwpw/nwpwlib/pseudopotential/rpspsolve.c
    filename = path.join(nwchem_dest, 'src', 'nwpw', 'nwpwlib', 'pseudopotential', 'rpspsolve.c')
    print(f'\n ChemShell: patching {filename}...')
    with open(filename, 'r+') as fp:
        m_pi = '  const long double M_PI = 3.14159265358979323846264338327950288419716939937510582097494459230781640628620899L;'
        rpspsolve = fp.readlines()
        for i, l in enumerate(rpspsolve):
            # skip if already patched
            if m_pi in l:
                break
            # insert the line before using M_PI
            if 'M_PI' in l:
                rpspsolve.insert(i, m_pi+'\n')
                break
        # update the file
        fp.seek(0)
        fp.writelines(rpspsolve)


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

    from sys import argv

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

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

    if patch == 'makefile_h':
        patch_makefile_h(nwchem_dest)
    if patch == 'mm_bond_F':
        patch_mm_bond_F(nwchem_dest)
    if patch == 'GNUmakefile':
        patch_GNUmakefile(nwchem_dest)
    if patch == 'rpspsolve_c':
        patch_rpspsolve_c(nwchem_dest)
