#  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/>.

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

from sys        import stdout
from ...dl_py2f import utils

# detect haems
def _detectHaem(bioFrag, _ff, _split_frags, _segnames, _haem, _rtf_res_dicts, _topol_abspaths, _residue_aliases, _additional_prm):
    '''Detect haem and add missing RTF/PRM files if necessary'''

    from numpy    import char, isin, setdiff1d, unique
    from os.path  import dirname, join
    from ...data  import forcefield
    from ...io    import files

    ffObj = forcefield.getFFObj(_ff)

    # list of supported haem residues
    supported_haems = ffObj.haems + utils.iterutils.getFlattened([_haem])

    # loop over segments
    for iseg, seg in enumerate(_segnames):
        segFrag = _split_frags[seg][0]
        # YL 27/05/2024: now all segment names are labelled with chainID
        chainID = unique(segFrag.chainIDs)[0].decode()
        seg = seg.rstrip(f'_{chainID}')
        # YL TODO 03/09/2021: we only support single haem now
        masks_residues_in_supported_haems = isin(char.decode(char.lower(char.strip(segFrag.residues))), char.lower(supported_haems))
        if masks_residues_in_supported_haems.sum():
            haem_name = unique(segFrag.residues[masks_residues_in_supported_haems])[0]
            # detect the covalent bonded neighbour residues
            haem = bioFrag.select(segments=seg, residues=haem_name, strict=False)
            fe   = bioFrag.select(segments=seg, atoms='FE')
            haem_wo_fe = setdiff1d(haem, fe)
            id_haem = bioFrag.resIDs[fe][0]
            # find the neighbour residues and their IDs: ligands are excluded
            indices_neighbour_residues_wo_ligands = setdiff1d(bioFrag.conn[haem_wo_fe], haem.tolist()+[-1])
            neighbours_haem = char.strip(bioFrag.residues[indices_neighbour_residues_wo_ligands])
            connected_to_fe = setdiff1d(bioFrag.conn[fe], haem.tolist()+[-1])
            
            # YL 03/09/2021: changed to all residues connected to Fe
            neighbours_fe       = char.strip(bioFrag.residues[connected_to_fe])
            id_neighbours_haem  = bioFrag.resIDs[indices_neighbour_residues_wo_ligands]
            seg_neighbours_haem = char.decode(bioFrag.segments[indices_neighbour_residues_wo_ligands])
            id_neighbours_fe    = bioFrag.resIDs[connected_to_fe]
            seg_neighbours_fe   = char.decode(bioFrag.segments[connected_to_fe])

            # haem residue name
            res_haem = bioFrag.residues[fe][0].decode().strip()

            # haem C: if connected to cysteines
            if b'CYS' in unique(neighbours_haem):
                print("\n ChemShell has detected a C-type haem (residue: \"{}\")".format(res_haem))
                haem_residue_alias = 'HEMC'

            # YL 03/09/2021: haem A, B, or O
            else:
                noxygens = (bioFrag.znums[haem]==8).sum()
                if noxygens == 4:
                    print("\n ChemShell has detected a B-type haem (residue: \"{}\")".format(res_haem))
                    haem_residue_alias = 'HEMB'
                elif noxygens == 5:
                    print("\n ChemShell has detected an O-type haem (residue: \"{}\")".format(res_haem))
                    haem_residue_alias = 'HEMO'
                elif noxygens == 6:
                    print("\n ChemShell has detected an A-type haem (residue: \"{}\")".format(res_haem))
                    haem_residue_alias = 'HEMA'
                else:
                    print("\n ChemShell has detected an unknown-type haem (residue: \"{}\")".format(res_haem))
                    haem_residue_alias = 'HEM'

            # check if specified RTF files contain a HEC residue definition
            found_hec = False
            for k, v in _rtf_res_dicts.items():
                # YL NB 03/09/2021: must compare stripped strings
                if res_haem in v.keys():
                    found_hec = True
            if not found_hec:
                rtf_haemc = join(ffObj.root_dir, ffObj.haems_rtf)
                prm_haemc = join(ffObj.root_dir, ffObj.haems_prm)
                print(" Adding built-in file to the RTF list:\n {}".format(rtf_haemc))
                # update topol_abspaths, rtf_res_dicts
                _topol_abspaths += [rtf_haemc]
                _rtf_res_dicts.update(files.rtf.getResidueDicts([rtf_haemc]))
                # take patched residue names as aliases for RTF scan
                _residue_aliases.update({res_haem:f'{haem_residue_alias}'})
                # YL: update the renaming records
                bioFrag._renamed_residues.update({res_haem:f'{haem_residue_alias}'})
                # add the needed extra PRM and return to the main function
                _additional_prm += [prm_haemc]

            return res_haem, neighbours_haem, neighbours_fe, id_haem, id_neighbours_haem, id_neighbours_fe, seg_neighbours_haem, seg_neighbours_fe


# YL 27/05/2024: added chainID to differentiate segments
def _patchHaem(haem, bioFrag, ffObj, haem_results, seg, chainID, residues_seg, residue_patches, patch_disulphides=True, mute=True):
    '''Patch the haem residue for PSFGEN'''

    from numpy    import char, isin

    # result string
    strbuff = ''
    patch = ''
    seg = seg.strip()

    # possible residue names
    histidines = ['his','hsd','hse','hsp']

    # list of supported haem residues
    supported_haems = ffObj.haems + utils.iterutils.getFlattened([haem])
    if isin(char.lower(residues_seg), char.lower(supported_haems)).sum() and haem_results is not None:

        strbuff += '\n# automatic residue patch for haem generated by ChemShell\n'
        # returned res_haem, neighbours_haem, neighbours_fe, id_haem, id_neighbours_haem, id_neighbours_fe
        res_haem, neighbours_haem, neighbours_fe, id_haem, id_neighbours_haem, id_neighbours_fe, seg_neighbours_haem, seg_neighbours_fe = haem_results[:]

        # YL 24/11/2021: patch liganding amino acids
        masks_amino_acids_in_neighbours_fe = isin(char.decode(neighbours_fe), ffObj.amino_acids)
        for i, amino_acid_in_neighbours_fe in enumerate(char.decode(neighbours_fe)[masks_amino_acids_in_neighbours_fe]):
            # patch deprotonated cysteine
            if amino_acid_in_neighbours_fe == 'CYS':
                strbuff += '\n# automatic patch to liganding cysteine generated by ChemShell\n'
                strbuff += f'patch CYSD {seg_neighbours_fe[i]}_{chainID}:{id_neighbours_fe[i]}\n'
                # take patched residue names as aliases for RTF scan
                residue_patches.update({'CYS':'CYSD'})

        # PHC1 and PHC2: patch for HAEM C - CYS (C-S connection)
        if len(neighbours_haem):
            print("\n Residues connected to the haem (residue: \"{}\"):".format(res_haem))
            stdout.flush()
        if not patch_disulphides and res_haem.lower().strip() in ['hec','hemc']:
            for i, n in enumerate(char.decode(neighbours_haem)):
                print(" {}:{}".format(id_neighbours_haem[i],n))
                stdout.flush()
                if char.lower(char.strip(char.decode(neighbours_haem)))[i] in ['cys']:
                    strbuff += 'patch PHC{} {}:{} {}:{}\n'.format(i+1,
                                                                  seg_neighbours_haem[i]+f'_{chainID}',
                                                                  id_neighbours_haem[i],
                                                                  seg+f'_{chainID}',
                                                                  id_haem)

        strbuff += '\nregenerate angles dihedrals\n\n'

        # H4__: patch for free haem (Fe is not coordinated to a ligand)
        # previously named FHEM
        if not len(neighbours_fe):
            patch = 'H4__'
            if not mute:
                print(f" Applying {patch} patching for non-liganded haem...")
            strbuff += '# patch non-coordinated haem\n'
            strbuff += 'patch {} {}:{}\n'.format(patch,
                                                 seg+f'_{chainID}',
                                                 id_haem)

        else:
            print("\n Residues connected to the haem's Fe: [segment:resID residue]")
            ns = char.decode(neighbours_fe)
            for i, n in enumerate(ns):
                print(" {:<4s}:{:>4d} {:<s}".format(seg_neighbours_fe[i],id_neighbours_fe[i],n))

            # 5-coordinated haems
            if len(neighbours_fe) == 1:
                n = char.decode(neighbours_fe)[0].lower().strip()
                # H5H_: patch for HIS - HAEM
                # previously named PHEM
                if n in histidines:
                    patch = 'H5H_'
                    if not mute:
                        print(f" Applying {patch} patching for haem 5-coordinated to histidine...")
                    strbuff += '# patch haem 5-coordinated to histidine\n'
                    strbuff += 'patch {} {}:{} {}:{}\n'.format(patch,
                                                               seg_neighbours_fe[0]+f'_{chainID}',
                                                               id_neighbours_fe[0],
                                                               seg+f'_{chainID}',
                                                               id_haem)

                # H5M_: patch for MET - HAEM
                # previously named PMET
                if n in ['met']:
                    patch = 'H5M_'
                    if not mute:
                        print(f" Applying {patch} patching for haem 5-coordinated to methionine...")
                    strbuff += '# patch haem 5-coordinated to methionine\n'
                    strbuff += 'patch {} {}:{} {}:{}\n'.format(patch,
                                                               seg_neighbours_fe[0]+f'_{chainID}',
                                                               id_neighbours_fe[0],
                                                               seg+f'_{chainID}',
                                                               id_haem)

                # H5C_: patch for CYS - HAEM
                # previously named PCYS
                if n in ['cys','cysd']:
                    patch = 'H5C_'
                    if not mute:
                        print(f" Applying {patch} patching for haem 5-coordinated to cysteine...")
                    strbuff += '# patch haem 5-coordinated to cysteine\n'
                    strbuff += 'patch {} {}:{} {}:{}\n'.format(patch,
                                                               seg_neighbours_fe[0]+f'_{chainID}',
                                                               id_neighbours_fe[0],
                                                               seg+f'_{chainID}',
                                                               id_haem)

                # H5M1: patch for CH3S(-) - HAEM
                # previously named PSUL
                if n in ['mes1']:
                    patch = 'H5M1'
                    if not mute:
                        print(f" Applying {patch} patching for haem 5-coordinated to CH3S(-) (methylthiolate anion)...")
                    strbuff += '# patch haem 5-coordinated to CH3S(-) (methylthiolate anion)\n'
                    strbuff += 'patch {} {}:{} {}:{}\n'.format(patch,
                                                               seg_neighbours_fe[0]+f'_{chainID}',
                                                               id_neighbours_fe[0],
                                                               seg+f'_{chainID}',
                                                               id_haem)

                # C15C: patch for CYS - CPM1
                # previously named PCC1
                if n in ['cys','cysd'] and res_haem.lower().strip() in ['cpd1']:
                    patch = 'C15C'
                    if not mute:
                        print(f" Applying {patch} patching for Compound I (type-B haem) 5-coordinated to cysteine...")
                    strbuff += '# patch Compound I (type-B haem) 5-coordinated to cysteine\n'
                    strbuff += 'patch {} {}:{} {}:{}\n'.format(patch,
                                                               seg_neighbours_fe[0]+f'_{chainID}',
                                                               id_neighbours_fe[0],
                                                               seg+f'_{chainID}',
                                                               id_haem)

            # 6-coordinated haems
            elif len(neighbours_fe) == 2:
                n0,n1 = char.lower(char.strip(char.decode(neighbours_fe)))[:]
                i0,i1 = id_neighbours_fe[:]
                s0,s1 = seg_neighbours_fe[:]
                # YL 13/01/2022: there shouldn't be any space
                s0,s1 = s0.strip(), s1.strip()

                # H6TH: patch for TYR - HAEM - HIS
                # previously named PLOH
                if (n0 in ['tyr'] and n1 in histidines) or (n1 in ['tyr'] and n0 in histidines):
                    # YL 16/12/2021: must respect the TYR - HAEM - HIS order
                    if n1 in ['tyr'] and n0 in histidines:
                        s0, s1 = s1, s0
                        i0, i1 = i1, i0
                    patch = 'H6TH'
                    if not mute:
                        print(f" Applying {patch} patching for haem 6-coordinated to tyrosine and histidine...")
                    strbuff += '# patch haem 6-coordinated to tyrosine and histidine\n'
                    strbuff += 'patch {} {}:{} {}:{} {}:{}\n'.format(patch, s0+f'_{chainID}', i0, seg+f'_{chainID}', id_haem, s1+f'_{chainID}', i1)

                # H6HO: patch for O2 - HAEM - HIS
                # previously named PLO2
                if (n0 in ['o2','oxy','oo'] and n1 in histidines) or (n1 in ['o2','oxy','oo'] and n0 in histidines):
                    # YL 16/12/2021: must respect the O2 - HAEM - HIS order
                    if n1 in ['o2','oxy','oo'] and n0 in histidines:
                        s0, s1 = s1, s0
                        i0, i1 = i1, i0
                    patch = 'H6H0'
                    if not mute:
                        print(f" Applying {patch} patching for haem 6-coordinated to O2 and histidine...")
                    strbuff += '# patch haem 6-coordinated to O2 and histidine\n'
                    strbuff += 'patch {} {}:{} {}:{} {}:{}\n'.format(patch, s0+f'_{chainID}', i0, seg+f'_{chainID}', id_haem, s1+f'_{chainID}', i1)

                # H6CO: patch for O2 - HAEM - CYS
                # previously named CLO2
                if (n0 in ['o2','oxy','oo'] and n1 in ['cys','cysd']) or (n1 in ['o2','oxy','oo'] and n0 in ['cys','cysd']):
                    # YL 16/12/2021: must respect the O2 - HAEM - CYS order
                    if n1 in ['o2','oxy','oo'] and n0 in ['cys','cysd']:
                        s0, s1 = s1, s0
                        i0, i1 = i1, i0
                    patch = 'H6CO'
                    if not mute:
                        print(f" Applying {patch} patching for haem 6-coordinated to O2 and cysteine...")
                    strbuff += '# patch haem 6-coordinated to O2 and cysteine\n'
                    strbuff += 'patch {} {}:{} {}:{} {}:{}\n'.format(patch, s0+f'_{chainID}', i0, seg+f'_{chainID}', id_haem, s1+f'_{chainID}', i1)

                # H6HC: patch for CO - HAEM - HIS
                # previously named PLIG
                if (n0 in ['co'] and n1 in histidines) or (n1 in ['co'] and n0 in histidines):
                    # YL 16/12/2021: must respect the CO - HAEM - HIS order
                    if n1 in ['co'] and n0 in histidines:
                        s0, s1 = s1, s0
                        i0, i1 = i1, i0
                    patch = 'H6HC'
                    if not mute:
                        print(f" Applying {patch} patching for haem 6-coordinated to O2 and histidine...")
                    strbuff += '# patch haem 6-coordinated to CO and cysteine\n'
                    strbuff += 'patch {} {}:{} {}:{} {}:{}\n'.format(patch, s0+f'_{chainID}', i0, seg+f'_{chainID}', id_haem, s1+f'_{chainID}', i1)

            else:
                print(f' WARNING: Unknown circumstance that haem is {len(neighbours_fe)+4}-coordinated!')

            strbuff += '\n# do NOT use regenerate again after the above patch!\n'

    if patch:
        residue_patches.update({haem:patch})

    stdout.flush()
    return strbuff
