#  Copyright (C) 2020 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/>.

__author__ = 'You Lu <you.lu@stfc.ac.uk>'

def parse(filename,
          _regex=None,
          receiver=None,
          omitted_only=False,
          ff='charmm',
          rearrange=False,
          ions_in_solvent=False,
          mute=False,
          prefix='',
          **kwargs):
    '''Parse a PQR file'''

    from numpy      import array, argwhere, count_nonzero, char, genfromtxt, loadtxt, set_printoptions, searchsorted, delete, stack, where
    from re         import compile
    from os         import path, remove
    from .          import pdb
    from ...base    import forcefield
    from ...data    import atomicdata, unitconvert, forcefield as data_ff
    from ...objects import fragment
    from sys        import stdout

    # forcefield data object
    ffObj = data_ff.getFFObj(ff)

    # YL 13/08/2022: path prefix
    filename = path.join(prefix, filename)

    version = 'pdb2pqr'
    # get the version first
    with open(filename, 'r') as fp:
        for lineno, line in enumerate(fp):
            if line.lower().startswith('remark') or line.lower().startswith('title'):
                if line.lower().strip().endswith('version=chemsh') or line.lower().strip().endswith('version=chemshell'):
                    version = 'chemsh'
                    break

    # get omitted atoms due to lack of FF definition
    if omitted_only:
        with open(filename, 'r') as fp:
            dtype = [('remark'  , 'U6'),
                     ('nremarks', 'i2'),
                     ('indices' , 'i8'),
                     ('names'   , 'U4'),
                     ('dummy'   , 'U2'),
                     ('residues', 'U5'),
                     ('resIDs'  , 'i8')]
            lines = fp.readlines()
            try:
                alpha = argwhere(char.count(lines, 'to the following atoms (omitted below):'))[0][0]+1
                omega = argwhere(char.count(lines, 'This is usually due to the fact that this residue is not'))[0][0]
            except:
                return array([], dtype=dtype)

            abuff = genfromtxt(lines[alpha:omega], dtype=dtype)

            return abuff

    tmp_filename = '_chemsh_tmp.pqr'

    # always covalent
    try:
        receiver.connect_mode = 'covalent'
    except:
        pass

    # FIXME: not in use
    # find the number of lines of comments beginning with "HEADER/REMARK/TITLE"
    skip_header = 0
    skip_footer = 0
    comments = [ 'remark', 'ter   ', 'end'
               ]
    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
        for lineno, line in enumerate(fp):
            if line.lower().startswith('atom') or line.lower().startswith('hetatm'):
                newfp.write(line)

    if version.lower() in [ 'chemshell', 'chemsh' ]:
        dtype = [('entries' , 'U6'),
                 ('indices' , 'i8'),
                 ('names'   , 'U4'),
                 ('locIDs'  , 'U1'),
                 ('residues', 'U5'),
                 ('chainIDs', 'U1'),
                 ('resIDs'  , 'i8'),
                 ('inserts' , 'U1'),
                 ('x'       , 'f8'),
                 ('y'       , 'f8'),
                 ('z'       , 'f8'),
                 ('charges' , 'f8'),
                 ('radii'   , 'f8'),
                 ('--'      , 'U6'),
                 ('segments', 'U4'),
                 ('elems'   , 'U2'),
                ]
        delimiters = [6, 6, 4, 1, 4, 1, 4, 5, 8, 8, 8, 6, 6, 5, 4, 2]
        # YL 17/07/2021: segments must be aligned to 73-76 and elements to 77-78
        # YL 02/03/2021: changed delimiter for "interts" from 5 to 4
        delimiters = [6, 6, 4, 1, 4, 1, 4, 4, 8, 8, 8, 6, 6, 6, 4, 2]
    else:
        dtype = [('entries' , 'U6'),
                 ('indices' , 'i8'),
                 ('names'   , 'U4'),
                 ('residues', 'U4'),
                 ('resIDs'  , 'i8'),
                 ('inserts' , 'U1'),
                 ('x'       , 'f8'),
                 ('y'       , 'f8'),
                 ('z'       , 'f8'),
                 ('charges' , 'f8'),
                 ('radii'   , 'f8'),
                 ('--'      , 'U5'),
                 ('segments', 'U4'),
                 ('elems'   , 'U2'),
                ]
        delimiters = [6, 6, 4, 5, 5, 4, 8, 8, 8, 8, 7, 5, 4, 2]

    #data = loadtxt(filename,
    #               dtype=dtype,
    #               delimiter=[0,6,12,16,17,21,22,26,30,38,46,54,60,72,76],
    #               comments=['TITLE ', 'REMARK', 'HEADER', 'ANISOU', 'CONECT', 'TER   '])
    abuff = genfromtxt(tmp_filename,
                       skip_header=skip_header,
                       skip_footer=skip_footer,
                       dtype=dtype,
                       #          at in na lo re ch re in x  y  z  oc te -- se el
                       delimiter=delimiters,
                       comments='ANISOU')

    # remove the unwanted lines
    for i in [ 'TITLE ', 'REMARK', 'HEADER', 'ANISOU', 'CONECT', 'TER   ' ]:
        abuff = delete(abuff, where(abuff['entries']==i)[0])

    # remove all non-alphabetics
    regex = compile('[^a-zA-Z]')
    namesbuff = char.strip(abuff['names'])
    set_printoptions(threshold=100000)
    
    # entries
    receiver._pdb_entries = abuff['entries']

    counter = 0
    receiver.resize(abuff.shape[0])


    # ChemShell version: if the PDB file comes with the elements column
    if version.lower() in [ 'chemshell', 'chemsh' ]:
        # fill in blank entries because some PDB, for example written by PSFGEN, are not complete
        blanks = char.isspace(abuff['elems']) + char.str_len(abuff['elems'])==0
        blank_names   = namesbuff[blanks]
        blank_znums   = forcefield.names2znums(blank_names, exceptions=ffObj.names2znums)
        blank_symbols = atomicdata.symbols_upper_array[blank_znums]
        abuff['elems'][blanks] = blank_symbols[:]

        # YL 24/02/2021: we first assign the znums using any available symbols in the elements column; and then
        #                correct the znums obtained from symbols using forcefield exceptions (default: CHARMM)
        try:
            # this step based on abuff['elems'] can fail due to abnormal atomic numbers
            receiver.znums = atomicdata.symbols_sort_indices[searchsorted(atomicdata.symbols_sorted_array, char.title(char.strip(abuff['elems'])))]
            znums_buff = forcefield.names2znums(namesbuff, exceptions=ffObj.names2znums)
            masks_diff = receiver.znums != znums_buff
            receiver.znums[masks_diff] = znums_buff[masks_diff]
        except:
            znums_buff = forcefield.names2znums(namesbuff, exceptions=ffObj.names2znums)
            masks_diff = receiver.znums != znums_buff
            receiver.znums[masks_diff] = znums_buff[masks_diff]

        # at last we fill in (can overwrite) the initially blank slots
        receiver.znums[blanks] = blank_znums[:]

    # PDB2PQR version
    else:
        receiver.znums = forcefield.names2znums(namesbuff, exceptions=ffObj.names2znums)

    # must assign `znums` before `names`
    receiver.names = namesbuff

    receiver.coords = stack([abuff['x'],abuff['y'],abuff['z']], axis=1)

    # Angstrom to Bohr
    receiver.coords *= unitconvert.Angstrom2Bohr

    # assign more entries
    for i in [ 'charges', 'residues', 'resIDs', 'radii', 'segments' ]:
        setattr(receiver, i, abuff[i])

    # YL 30/01/2022: disabled connectivity based on PDB2PQR radii which are too problematic!
    try:
#        receiver.radii *= unitconvert.Angstrom2Bohr
        receiver.radii = atomicdata.covalent_radii[receiver.znums]
    except:
        pass

    pdb.assignSegments(receiver, ff=ff, rearrange=rearrange, mute=mute, ions_in_solvent=ions_in_solvent)

    # delete the tmp PDB
    remove(tmp_filename)

    # print a summary
    if not mute:
        pdb.printSummary(filename, receiver)


def write(filename, fragObj, version='chemshell', uppercase=False, **kwargs):
    '''Write out a PQR file'''

    from sys        import stdout
    from numpy      import arange, argwhere, char
    from time       import asctime, localtime, time
    from ...data    import atomicdata, unitconvert
    from ...objects import fragment, fragbio

    # YL 07/09/2020: PDB2PQR can't process REMARK???

    now = asctime(localtime(time()))
    strbuff = ''
    if version.strip().lower().startswith('chemsh'):
        strbuff += 'TITLE  {:3} {}\n'.format(1, 'VERSION=CHEMSH')
        strbuff += 'TITLE  {:3} {}\n'.format(2,  fragObj.title)
#        strbuff += 'REMARK {}\n'.format(now)
#        strbuff += 'REMARK {:3} {}\n'.format(1, 'ChemShell PQR format is same as PDB except that')
#        strbuff += 'REMARK {:3} {}\n'.format(1, 'Columns 55-60 for occupancies replaced with charges')
#        strbuff += 'REMARK {:3} {}\n'.format(1, 'Columns 61-66 for b-factors replaced with covalent radii')
#        strbuff += 'REMARK  Column      Abbreviation     Entry\n'
#        strbuff += 'REMARK  1 -  6                       Record name\n'
#        strbuff += 'REMARK  7 - 11      I                Atom serial number\n'
#        strbuff += 'REMARK 12 - 12\n'
#        strbuff += 'REMARK 13 - 16      N                Atom name\n'
#        strbuff += 'REMARK 17\n'
#        strbuff += 'REMARK 18 - 20      RS               Residue name\n'
#        strbuff += 'REMARK 21 - 21\n'
#        strbuff += 'REMARK 22           C                Chain ID\n'
#        strbuff += 'REMARK 23 - 26      CH               Residue sequence number\n'
#        strbuff += 'REMARK 27\n'
#        strbuff += 'REMARK 28 - 30\n'
#        strbuff += 'REMARK 31 - 38      X                X coorindate (Angstrom)\n'
#        strbuff += 'REMARK 39 - 46      Y                Y coorindate (Angstrom)\n'
#        strbuff += 'REMARK 47 - 54      Z                Z coorindate (Angstrom)\n'
#        strbuff += 'REMARK 55 - 60      Q                Charge\n'
#        strbuff += 'REMARK 61 - 66      R                Radius\n'
#        strbuff += 'REMARK 77 - 78      E                Element symbol\n'
#        strbuff += 'REMARK I    N    RS  CCH      X       Y       Z       Q     R               E  0\n'

    if not isinstance(fragObj, fragbio.FragmentBio):
        fragObj = fragObj.castTo('bio')

    # in PDB resIDs greater than 9999 cause serious problems
    if fragObj.resIDs.max() > 9999:
        print("\n >>> WARNING: PDB resID greater than 9999 can cause serious problems\n")

    try:
        entries = fragObj._pdb_entries
    except:
        entries = 'ATOM'

    if uppercase:
        elements = char.rjust(atomicdata.symbols_upper_array[fragObj.znums], 2)
    else:
        elements = char.rjust(atomicdata.symbols_array[fragObj.znums], 2)

    # YL 28/01/2022: the radii array must be in Bohr!
    masks_zero_radii = fragObj.radii == 0
    fragObj.radii[masks_zero_radii] = atomicdata.covalent_radii[fragObj.znums[masks_zero_radii]]
    # YL 28/01/2022: in the .pqr file it should be written in Angstrom
    radii = fragObj.radii*unitconvert.Bohr2Angstrom

    # YL 25/01/2022: make sure residues and chainIDs together occupy at most 5 columns
    if all(char.strip(fragObj.chainIDs)==b''):
        residues = fragObj.residues
        chainIDs = ''
        fmt_rs   = '%-5s'
        fmt_id   = '%0s'
    else:
        residues = char.ljust(fragObj.residues, 4)
        chainIDs = char.ljust(fragObj.chainIDs, 1)
        fmt_rs   = '%-4s'
        fmt_id   = '%1s'

    strbuff += fragObj.coords2str(coords='%8.3f',                                      # 31-54 (cols 9-11)
                                  inserts=[( 0,  entries                   , '%-6s' ), #  1- 6
                                           ( 1,  arange(1,fragObj.natoms+1), '%5d'  ), #  7-11
                                           ( 2, ' '*1                      , '%1s'  ), #
                                           ( 3,  fragObj.names             , '%-4s' ), # 13-16
                                           ( 4, ' '*1                      , '%1s'  ), #    17 TODO
                                           # non-standard PDB with residue of more than 3 characters (especially needed for TIP3 or TIP3P)
                                           ( 5,  residues                  ,  fmt_rs), # 18-21
                                           ( 6,  chainIDs                  ,  fmt_id), #    22
                                           ( 7,  fragObj.resIDs            , '%4d'  ), # 23-26
                                           ( 8, ' '*4                      , '%4s'  ), #       TODO
                                           (12,  fragObj.charges           , '%6.3f'), # 55-60
                                           (13,  radii                     , '%6.3f'), # 61-66
                                           (14, ' '*6                      , '%6s'  ), #       TODO
                                           (15,  fragObj.segments          , '%4s'  ), # 73-76 not defined in v3.3
                                           (16,  elements                  , '%-2s' ), # 77-78 elements
                                           (17,     0                      , '%2d'  ), # 79-80 charges
                                          ],
                                  atoms=None,
                                  charges=None,
                                  order='ac',
                                  delimiter='',
                                  unit='angstrom')

    try:
        with open(filename, 'w') as fp:
            fp.write(strbuff)
    # YL 08/01/2020: allow file pointer
    except:
        filename.write(strbuff)



