#  Copyright (C) 2019 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>'

# PDB v3.3
# COLUMNS      DATA TYPE        FIELD      DEFINITION
# ------------------------------------------------------
#  1 -  6      Record name      "ATOM    "
#  7 - 11      Integer          serial     Atom serial number
# 12 - 12
# 13 - 16      Atom             name       Atom name
# 17           Character        altLoc     Alternate location indicator
# 18 - 20      Residue name     resName    Residue name
# 21 - 21
# 22           Character        chainID    Chain identifier
# 23 - 26      Integer          resSeq     Residue sequence number
# 27           AChar            iCode      Code for insertion of residues
# 28 - 30
# 31 - 38      Real(8.3)        x          Orthogonal coordinates for X in
#                                          Angstroms
# 39 - 46      Real(8.3)        y          Orthogonal coordinates for Y in
#                                          Angstroms
# 47 - 54      Real(8.3)        z          Orthogonal coordinates for Z in
#                                          Angstroms
# 55 - 60      Real(6.2)        occupancy  Occupancy
# 61 - 66      Real(6.2)        tempFactor Temperature factor
# 77 - 78      LString(2)       element    Element symbol, right-justified
# 79 - 80      LString(2)       charge     Charge on the atom
# see: http://www.wwpdb.org/documentation/file-format-content/format23/sect9.html

def parse(filename,
          receiver=None,
          threshold_occ=0.6,
          choose_locations=True,
          ff='charmm',
          rearrange=False,
          ions_in_solvent=False,
          force_reassign_segments=False,
          attach_to_protein=True,
          mute=False,
          prefix='',
          **kwargs):
    '''Parse a PDB file'''

    from numpy      import all, arange, count_nonzero, char, genfromtxt, loadtxt, mean, set_printoptions, searchsorted, delete, stack, where, unique
    from re         import compile
    from io         import StringIO
    from os         import path, remove
    from sys        import stdout
    from ...base    import forcefield
    from ...data    import atomicdata, unitconvert, forcefield as data_ff
    from ...objects import fragment

    tmp_filename = '_chemsh_tmp.pdb'

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

    # 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 = [ 'title ', 'remark', 'header', 'anisou', 'AUTHOR',
                 'conect', 'ter   ', 'end'   , 'master', 'REVDAT',
                 'COMPND', 'SOURCE', 'KEYWDS', 'EXPDTA', 'JRNL'  ,
                 'DBREF' , 'SEQADV', 'SEQRES', 'HET'   , 'HETNAM',
                 'FORMUL', 'HELIX' , 'SHEET' , 'SSBOND', 'SITE'  ,
                 'CRYST1', 'ORIGX1', 'ORIGX2', 'ORIGX3',
                 'SCALE1', 'SCALE2', 'SCALE3'
               ]

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

    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)
                    
    # in the standard PDB format residues are only 3-character long,
    # but in practice they usually are 5-long (CHARMM and PSFGEN)
    dtype = [('entries' , 'U6'),
             ('indices' , 'i8'),
             ('names'   , 'U4'),
             ('locIDs'  , 'U1'),
             ('residues', 'U5'),
             ('chainIDs', 'U1'),
             ('resIDs'  , 'i8'),
             ('inserts' , 'U1'),
             ('x'       , 'f8'),
             ('y'       , 'f8'),
             ('z'       , 'f8'),
             ('occ'     , 'f8'),
             ('bfactors', 'f8'),
             ('--'      , 'U6'),
             ('segments', 'U4'),
             ('elems'   , 'U2'),
            ]

    delimiters = [6, 6, 4, 1, 5, 0, 4, 5, 8, 8, 8, 6, 6, 5, 4, 2]
    # YL 02/03/2021: changed delimiter for "interts" from 5 to 4
    # YL 17/07/2021: segments must be aligned to 73-76 and elements to 77-78
    #    columns: 6 12 16 17 21 22 26 30 38 46 54 60 66 72 76 78
    delimiters = [6, 6, 4, 1, 4, 1, 4, 4, 8, 8, 8, 6, 6, 6, 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])

    # assign znums
#    abuff = genfromtxt(filename,
#                       usecols=2,
#                       skip_header=skip_header,
#                       skip_footer=1,
#                       comments='REMARK',
#                       dtype='U8')
    # remove all non-alphabetics
    regex = compile('[^a-zA-Z]')

    # YL 03/03/2021: we allow the names to be overwirtten by another source when creating
    #                the fragment instance, for example enzyme = Fragment(coords='enz.pdb', names=['A','B','C',...])
    if all(~(receiver.names==b'X')):
        print("\n ChemShell PDB parser: PDB atom names are prevailed by predefined values")
        namesbuff = char.strip(char.decode(receiver.names))
    else:
        namesbuff = char.strip(abuff['names'])

    set_printoptions(threshold=100000)
    
    # entries
    receiver._pdb_entries = abuff['entries']

    counter = 0
   
    # first assign names directly found in the list of conventions; then remove the trailing numbers to match the name conventions
#    receiver.znums = [ forcefield.names2znums[symbol] if symbol in forcefield.names2znums else
#                       forcefield.names2znums[regex.sub('', symbol)]
#                       for symbol in namesbuff ]

    # must assign `znums` before `names`
    # if the PDB file comes with the elements column
    # 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[:]
    receiver.names = namesbuff


#    # assign names
#    abuff = genfromtxt(filename,
#                       usecols=2,
#                       skip_header=skip_header,
#                       skip_footer=1,
#                       comments='REMARK',
#                       dtype='U8')
#    receiver.names = abuff
#
#    # assign coords
#    abuff = genfromtxt(filename,
#                       usecols=[5,6,7],
#                       skip_header=skip_header,
#                       skip_footer=1,
#                       comments='REMARK')

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

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

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

    # segments are not always defined in PDB
    renameResidues(receiver, ff=ff, mute=mute)
    assignSegments(receiver,
                   ff=ff,
                   rearrange=rearrange,
                   mute=mute,
                   ions_in_solvent=ions_in_solvent,
                   force_reassign=force_reassign_segments,
                   attach_to_protein=attach_to_protein)

    # remove candidate atom positions with less occupancy
    if choose_locations:
        indices_removed_uncertainty = chooseCandidateStructures(receiver, abuff, threshold_occ=threshold_occ)
    else:
        indices_removed_uncertainty = ...

    # we can't use the getSelected() function
    receiver._master = receiver._master[indices_removed_uncertainty]
    # set all occupancies to 1.0 since we have made selections
    receiver._master['occ'] = 1.0
    receiver._pdb_entries = receiver._pdb_entries[indices_removed_uncertainty]

    # get the terminal
    with open(filename, 'r') as fp:
        dtype = [('entry'  , 'U6'),
                 ('index'  , 'i8'),
                 ('name'   , 'U4'),
                 ('locID'  , 'U1'),
                 ('residue', 'U5'),
                 ('chainID', 'U1'),
                 ('resID'  , 'i8')]
        delimiters = [6, 6, 4, 1, 4, 1, 4]
        for lineno, line in enumerate(fp):
            if line.lower().startswith('ter'):
                sbuff = StringIO(line)
                abuff = genfromtxt(sbuff,
                                   dtype=dtype,
                                   delimiter=delimiters)
                receiver.terminal.residue = abuff['residue']
                receiver.terminal.resID   = abuff['resID']
# YL TODO 27/02/2021: it's difficult to set the terminal reasonably becasue there might
#                     be multiple terminals
#                receiver.terminal.segment = abuff['resID']

    # delete the tmp PDB
    remove(tmp_filename)

    # assign title
#    strbuff = genfromtxt(filename, skip_header=1, max_rows=1, dtype='U255')
#    receiver.title = str(strbuff)

    if not mute:
        printSummary(filename, receiver)


def printSummary(filename, fragObj):
    '''Print a summary'''

    from sys   import stdout
    from numpy import char, unique
    from ...io import format

    segments_unique = fragObj.all_segments

    print(' '+'-'*70)
    print(f' Summary of loaded fragment')
    print(' '+'-'*70)
    print(f' Filename: {filename}')
    print(f' Number of atoms: {fragObj.natoms}')
    print(f' Segments and residue names contained:')
    for segment in segments_unique:
        residues_unique = unique(fragObj.residues[fragObj.segments==segment])
        print(' {:4}: '.format(segment.decode().strip()), end='')
        if not len(residues_unique):
            continue
        resnames = format.wrap(' '.join(char.strip(char.decode(residues_unique))),
                               indent='',
                               initial_indent='',
                               width=70-6,
                               subsequent_indent=' '*7)
        print(resnames)
    print(' '+'-'*70)
    stdout.flush()


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

    from numpy      import all, arange, array, char, hstack, savetxt, unique
    from time       import asctime, localtime, time
    from io         import BytesIO
    from ...data    import atomicdata, unitconvert
    from ...objects import fragment, fragbio


    # YL 27/04/2022: due to a bug in PDB2PQR (as of 3.5) the residue name is assigned wrong if the segment contains only one residue, for example in
    #                atom   :  1  2  3  4  5  6  7  8  9
    #                resID  :  1  1  2  2  1  1  1  1  2
    #                residue:  k  k  l  l  m  m  n  n  o
    #                segment:  A  A  A  A  B  B  C  C  C
    #                the residue m of segment B will be wrongly assigned as "n" because PDB2PQR only compares their resIDs
    if version == 'pdb2pqr':
        for segment in fragObj.all_segments:
            masks_segment = char.lower(char.strip(fragObj.segments)) == segment.strip().lower()
            resIDs_segment = unique(fragObj.resIDs[masks_segment])
            if len(resIDs_segment) == 1:
                fragObj.resIDs[masks_segment] = 9999

    now = asctime(localtime(time()))
    strbuff = ''
    # YL 07/09/2020: PDB2PQR can't process REMARK???
    if fragObj.title:
        strbuff += 'TITLE  {:3} {}\n'.format(1, fragObj.title)
        strbuff += 'REMARK   1 Generated by Py-ChemShell according to PDB v3.3\n'
        strbuff += 'REMARK   2 {}\n'.format(now)
        strbuff += 'REMARK   3  Column      Abbreviation     Entry\n'
        strbuff += 'REMARK   3  1 -  6                       Record name\n'
        strbuff += 'REMARK   3  7 - 11      I                Atom serial number\n'
        strbuff += 'REMARK   3 12 - 12\n'
        strbuff += 'REMARK   3 13 - 16      N                Atom name\n'
        strbuff += 'REMARK   3 17 - 17      L                Location indicator\n'
        strbuff += 'REMARK   3 18 - 20      RS               Residue name\n'
        strbuff += 'REMARK   3 21 - 21\n'
        strbuff += 'REMARK   3 22 - 22      C                Chain ID\n'
        strbuff += 'REMARK   3 23 - 26      CH               Residue sequence number\n'
        strbuff += 'REMARK   3 27 - 27      IN               Insertion\n'
        strbuff += 'REMARK   3 28 - 30\n'
        strbuff += 'REMARK   3 31 - 38      X                X coorindate (Angstrom)\n'
        strbuff += 'REMARK   3 39 - 46      Y                Y coorindate (Angstrom)\n'
        strbuff += 'REMARK   3 47 - 54      Z                Z coorindate (Angstrom)\n'
        strbuff += 'REMARK   3 55 - 60      O                Occupancy\n'
        strbuff += 'REMARK   3 61 - 66      B                Temperature factor\n'
        strbuff += 'REMARK   3 77 - 78      E                Element symbol\n'
        strbuff += 'REMARK   3I    N   LRS  CCH  IN  X       Y       Z       O     B               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 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(char.strip(fragObj.residues), 4)
        chainIDs = char.ljust(fragObj.chainIDs, 1)
        fmt_rs   = '%-4s'
        fmt_id   = '%1s'

    # pretty print
    segments = char.ljust(char.strip(fragObj.segments), 4)

    # YL 03/05/2024: when the indices exceed 99999 we have to use hex
    indices = arange(1,fragObj.natoms+1)
    bint = BytesIO()
    savetxt(bint, indices[:99999], fmt='%5d', newline=' ')
    sint = bint.getvalue().decode()
    aint = array(sint.split(), dtype=str)
    bhex = BytesIO()
    savetxt(bhex, indices[99999:], fmt='%5x', newline=' ')
    shex = bhex.getvalue().decode()
    ahex = array(shex.split(), dtype=str)
    a_indices = hstack((aint,ahex))

    strbuff += fragObj.coords2str(coords='%8.3f',                            # 31-54 (cols 9-11)
                                  inserts=[( 0,  entries         , '%-6s' ), #  1- 6
                                           ( 1,  a_indices       , '%5s'  ), #  7-11 (YL 03/05/2024: changed as above to prevent exceeding 5 columns)
                                           ( 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.occ     , '%6.2f'), # 55-60
                                           (13,  fragObj.bfactors, '%6.2f'), # 61-66
                                           (14, ' '*6            , '%6s'  ), #       TODO
                                           (15,  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:
        strbuff += '{:<6s}{:>5d}      {:<4s} {:>4d}\n'.format('TER', fragObj.natoms+1, fragObj.terminal.residue, fragObj.terminal.resID)
    except:
        pass

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


def renameResidues(fragObj, ff='charmm', mute=False):
    '''Rename residue names to be recognised by forcefiled data'''

    strbuff = '\n ChemShell PDB parser: automatically renaming `residues`...'
    to_print = False

    # 'K'->'POT'
    selected = fragObj.select(residues='K', strict=True)
    if len(selected):
        to_print = True
        fragObj.residues[selected] = 'POT'
        strbuff += '\n residue "K" -> residue "POT" (potassium)'
        fragObj._renamed_residues.update({'K':'POT'})

    # ligands
    selected = fragObj.select(residues='OXY', strict=True)
    if len(selected):
        to_print = True
        fragObj.residues[selected] = 'O2'
        strbuff += '\n residue "OXY" -> residue "O2" (oxygen)'
        fragObj._renamed_residues.update({'OXY':'O2'})

    if to_print and not mute:
        print(strbuff)


def assignSegments(fragObj, ff='charmm', rearrange=False, mute=False, ions_in_solvent=False, force_reassign=False, attach_to_protein=True):
    '''Fill in the segment entries'''

    from sys     import stdout
    from numpy   import all, any, arange, argmin, array_equal, char, count_nonzero, diff, hstack, isin, sort, unique
    # not to be confused with base.forcefield!
    from ...data import forcefield as ff
    from ...io   import format

    # forcefield object
    ffObj = ff.getFFObj(ff)

    # segments are not always defined in PDB, especially in PQR where we will
    # have no idea what the residues are if their IDs are not sequential
    if not count_nonzero(fragObj.segments) or force_reassign:

        if not mute:
            print("\n ChemShell PDB parser: automatically assigning `segments`...")

        # do some automatical guess
        # water types:
        # HOH  - crystal PDB
        # TP3M - PDB2PQR
        selected = fragObj.select(residues=ffObj.waters)
        if len(selected):
            if not mute:
                print(' Residue "HOH" -> segment "CW" (crystal water)')
            fragObj.segments[selected] = 'CW'

        # water
        selected = fragObj.select(residues='TP3M')
        if len(selected):
            if not mute:
                print(' Residue "TP3M" -> segment "CW" (crystal water)')
            fragObj.segments[selected] = 'CW'

        # ligands
        selected = fragObj.select(residues='O2')
        if len(selected):
            if not mute:
                print(' Residue "O2" -> segment "LIG" (ligand)')
            fragObj.segments[selected] = 'LIG'

        # ions
        selected = fragObj.select(residues=['CLA','SO4','SOD','POT','CES','MG','CAL','ZN'])
        if len(selected):
            if not mute:
                print(' Ions ("CLA", "SO4", "SOD", etc.) -> segment "CW" (crystal water)')
            fragObj.segments[selected] = 'CW'
        
        # amino acids
        selected_protein = fragObj.select(residues=ffObj.amino_acids)
        if len(selected_protein):
            if not mute:
                print(' Amino acid residues -> segment "PROT" (protein)')
            fragObj.segments[selected_protein] = 'PROT'

        # haems
        for haem in ffObj.haems:
            selected = fragObj.select(residues=haem)
            if len(selected):
                if not mute:
                    print(f' Residue "{haem}" -> segment "HAEM" (haem)')
                fragObj.segments[selected] = 'HAEM'

        # unknown
        fragObj.segments[char.isspace(fragObj.segments)] = 'MISC'

    if not mute:
        print("\n Checking segments connectivity and contiguousness...\n")

    # YL 12/02/2021: scan all protein segments
    from time import time
    if attach_to_protein:
        segments_protein = unique(fragObj.segments[fragObj.selectByProtein()])

        t0 = time()
        # YL 14/03/2024: we check if the protein segments are considered connected by some residues
        masks_protein = fragObj.selectBio(segments=segments_protein, return_masks=True)
        indices_protein = fragObj._indices_all[masks_protein]
        indices_non_protein = fragObj._indices_all[~masks_protein]
        fragIDs_protein = unique(fragObj.fragIDs[indices_protein])
        masks_in_protein_fragsIDs = isin(fragObj.fragIDs, fragIDs_protein)*(~masks_protein)
        segments_in_protein_fragsIDs = unique(fragObj.segments[masks_in_protein_fragsIDs])
        strbuff = ''
        for segment_in_protein_fragsIDs in segments_in_protein_fragsIDs:
            masks_segment_in_protein_fragsIDs = segment_in_protein_fragsIDs == fragObj.segments
            # using indices instead of masks is much faster, especially in the loop below over resIDs_in_current_segment
            indices_segment_in_protein_fragsIDs = fragObj._indices_all[masks_segment_in_protein_fragsIDs]
            resIDs_in_current_segment = unique(fragObj.resIDs[indices_segment_in_protein_fragsIDs])
            resIDs_all_in_current_segment = fragObj.resIDs[indices_segment_in_protein_fragsIDs]
            # should always loop over resIDs in a segment
            for resID_in_current_segment in resIDs_in_current_segment:
                indices_current_resID = indices_segment_in_protein_fragsIDs[resID_in_current_segment==resIDs_all_in_current_segment]
                indices_connected = unique(fragObj.conn[indices_current_resID])
                indices_connected = list(set(indices_connected[indices_connected != -1]) - set(indices_current_resID))
                resIDs_connected = fragObj.resIDs[indices_connected]
                segments_connected = fragObj.segments[indices_connected]
                # iff the current residue has only one fragID (single fragment)
                if unique(fragObj.fragIDs[indices_current_resID]).size == 1 and isin(segments_protein, segments_connected).sum() >= 2:
                    tbuff = tuple(zip(segments_connected, resIDs_connected))
                    residue_in_current_segment = fragObj.resID2residue(resID_in_current_segment, segment_in_protein_fragsIDs)
                    strbuff += f'{segment_in_protein_fragsIDs.decode().strip()}:{residue_in_current_segment.strip()}{resID_in_current_segment} => '
                    lbuff = [ f'{k.decode().strip()}:{fragObj.resID2residue(v,k).strip()}{v}' for k,v in tbuff ]
                    strbuff += ' '.join(unique(lbuff))
        if strbuff:
            plural = {1:' is'}.get(len(strbuff.splitlines()), 's are')
            format.printwrap(f'NB: The following residue{plural} connected to at least two protein segments (seg:resID):', subsequent_indent=' '*4)
            format.printwrap(strbuff, prepend_blanc_line=False, indent=' '*5, subsequent_indent=' '*12)

        format.printwrap(f'>>> Time used for calculating connection between protein segments: {time()-t0:.3f} s', subsequent_indent=' '*4)
        print()

        # connect all bond-connected residues to protein
        # YL 20/09/2021: this is required by DL_FIELD but should not be used for PSFGEN!
        for protein_segname in segments_protein:
            # YL NB: this functions contains an operation of fragObj.optimiseResIDs()
            fragObj.mergeConnectedResiduesToSegment(protein_segname, ff=ff, ions_in_solvent=ions_in_solvent)

    # YL 28/01/2021: a DL_FIELD/DL_POLY segment ("group") must be a contiguous sequence in PDB
    # YL 24/02/2021: realign all segments
    # YL 02/06/2021: we can't always rearrange the sequence otherwise there will be great inconvenience
    #                and chaos, for example when AMBER prmtop is used and no DL_FIELD is needed
    if rearrange:
        for segname in char.decode(unique(fragObj.segments)):
            fragObj.keepSegmentContiguous(segname)

    # YL 04/08/2021: but we should always make sure the resIDs are correct!
    for segname in char.decode(unique(fragObj.segments)):
        fragObj.optimiseResIDs(segname)


# 04/02/2021: filter the residues of multiple candidate positions according to both occupancy and b-factor
def chooseCandidateStructures(receiver, abuff, threshold_occ=0.6):
    '''Return an array of indices containing only residue structure of greater occupancy or less b-factors'''

    from numpy import append, arange, array, char, copy, delete, mean, sort, where, unique, vstack

    abuff['segments'] = getattr(receiver, 'segments')

    indices_all = arange(abuff.shape[0])
    masks_all_w_candidate = ~char.isspace(abuff['locIDs'])
    indices_no_candidate = indices_all[~masks_all_w_candidate]

    abuff_all_w_candidate = abuff[masks_all_w_candidate]

    # loop over segments containing candidate structures
    candidate_segments = unique(abuff_all_w_candidate['segments'])
    if candidate_segments.size:
        sbuff = '\n ChemShell PDB parser: candidate residue structures* kept in raw PDB:\n'
        sbuff += ' '+70*'-'+'\n'
        sbuff += ' segment  resID  residue  locID  mean occ.  % min. b-factors  removed\n'
        sbuff += ' '+70*'-'+'\n'

    for candidate_segment in candidate_segments:
        masks_current_segment_w_candidate = abuff_all_w_candidate['segments'] == candidate_segment

        # loop over residues containing candidate structures
        candidate_resIDs_current_segment = unique(abuff_all_w_candidate[masks_current_segment_w_candidate]['resIDs'])
        for candidate_resID_current_segment in candidate_resIDs_current_segment:
            masks_current_resID_segment_w_candidate = abuff_all_w_candidate[masks_current_segment_w_candidate]['resIDs'] == candidate_resID_current_segment
            indices_current_resID_segment_w_candidate = indices_all[masks_all_w_candidate][masks_current_segment_w_candidate][masks_current_resID_segment_w_candidate]
            abuff_current_resID_segment_w_candidate = abuff_all_w_candidate[masks_current_segment_w_candidate][masks_current_resID_segment_w_candidate]
            candidate_residue_current_segment = unique(abuff_current_resID_segment_w_candidate['residues'])[0].strip()

            # all location IDs of the current residue
            locIDs_current_resID_segment = unique(abuff_current_resID_segment_w_candidate['locIDs'])

            # all occupancies of the current residue
            occs_current_resID_segment = abuff_current_resID_segment_w_candidate['occ']

            # initialise
            kept_ID = locIDs_current_resID_segment[0]
            masks_determined_locID_current_resID_segment = abuff_current_resID_segment_w_candidate['locIDs']==kept_ID

            mean_occ0 = mean(occs_current_resID_segment[masks_determined_locID_current_resID_segment])

            nentries = [ masks_determined_locID_current_resID_segment.sum() ]

            # loop over candidate structure IDs
            for locID in locIDs_current_resID_segment[1:]:
                masks_locID_current_resID_segment = abuff_current_resID_segment_w_candidate['locIDs'] == locID
                nentries += [ masks_locID_current_resID_segment.sum() ]
                mean_occ = mean(occs_current_resID_segment[masks_locID_current_resID_segment])
                if mean_occ0 <= mean_occ:
                    mean_occ0 = mean_occ
                    kept_ID = locID
                    # determined masks
                    masks_determined_locID_current_resID_segment = masks_locID_current_resID_segment

            ratio_current_bfactors = '--'

            # if mean occupancy lower than the threshold (default: 0.6)
            if mean_occ0 < threshold_occ:
                # all b-factors of the current residue
                bfactors_current_resID_segment = abuff_current_resID_segment_w_candidate['bfactors']

                # initialise
                kept_ID = locIDs_current_resID_segment[0]
                masks_determined_locID_current_resID_segment = abuff_current_resID_segment_w_candidate['locIDs']==kept_ID

                # construct an array of bfactors by looping over locIDs again
                bfactors = array([]).reshape(0,max(nentries))
                for ilocID, locID in enumerate(locIDs_current_resID_segment):
                    # NB: the lengths may be uneven!
                    if nentries[ilocID] < max(nentries):
                        continue
                    masks_locID_current_resID_segment = abuff_current_resID_segment_w_candidate['locIDs'] == locID
                    bfactors = vstack((bfactors, bfactors_current_resID_segment[masks_locID_current_resID_segment]))
                    nbfactors = bfactors.shape

                # find which ID features the greatest number of minimum b-factors
                iIDs_min_bfactors = bfactors.argmin(axis=0)
                unique_iIDs_min_bfactors, counts_min_bfactors = unique(iIDs_min_bfactors, return_counts=True)
                iID_min_bfactors = unique_iIDs_min_bfactors[counts_min_bfactors.argmax()]
                kept_ID = locIDs_current_resID_segment[iID_min_bfactors]
                # determined masks
                masks_determined_locID_current_resID_segment = abuff_current_resID_segment_w_candidate['locIDs']==kept_ID
                ratio_current_bfactors = str(int(round(counts_min_bfactors.max()/max(nentries)*100, 0)))+'%'

                # update the mean occupancy 
                mean_occ0 = mean(occs_current_resID_segment[masks_determined_locID_current_resID_segment])

            removed_IDs = delete(locIDs_current_resID_segment, where(locIDs_current_resID_segment==kept_ID))

            sbuff += "  {}     {:>4d}   {}      {}      {:4.2f}           {:>4s}          {}\n".format(candidate_segment,
                                                                                                       candidate_resID_current_segment,
                                                                                                       candidate_residue_current_segment,
                                                                                                       kept_ID,
                                                                                                       mean_occ0,
                                                                                                       ratio_current_bfactors,
                                                                                                       ' '.join(removed_IDs))
            # indices of the residue atoms of the determined structure
            determined_global_indices = indices_current_resID_segment_w_candidate[masks_determined_locID_current_resID_segment]

            # concat the indices
            indices_no_candidate = append(indices_no_candidate, determined_global_indices)

        sbuff += ' '+70*'-'+'\n'
        sbuff += ' *: Residues with greater mean occupancy or more number of mininum\n'
        sbuff += '    b-factors when the mean occupancy is less than {}\n'.format(threshold_occ)

    if indices_no_candidate.size and candidate_segments.size:
        print(sbuff)

    return sort(indices_no_candidate)


