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

# you.lu@stfc.ukri.org
# Oct 2019
# revised Sept 2019

from sys import modules
try:
    from . import info
except:
    pass

from ...io    import format
from ...dl_py2f import utils

# NB
# FHEM: no ligand (His or gas) bound to Fe
# PHEM: His bound to Fe
# PLO2: His and O2 bound to Fe
# PLIG: His and CO bound to Fe
# YL 26/02/2021: do not use `from chemsh import CHEMSH_DEBUG` otherwise a different value is received
import chemsh
from ...dl_py2f import utils

class __Wrapper(object):
    '''A wrapper class that will be disguised as a module'''

    # YL 17/12/2023: this is a protection mechanism to prevent the current module (wrapped as class)
    #                being analysed by chemsh.base.getClasses(), upon ChemShell's initialisation, in
    #                which inspect.getmembers() unnecessarily parses the whole database of CHARMM FF
    @property
    def __constable(self):
        # inspect the caller
        caller = utils.miscutils.getCaller()
        if 'getmembers' in caller and '__updateClasses' in caller and 'getClasses' in caller:
            # stop the caller
            exit()


    @property
    def default_location(self):
        '''Absolute path to the library directory'''

        try:
            return info.charmm_ff_dir
        except:
            # YL 26/02/2021: CHEMSH_DEBUG can be -1 so that the warning is normally printed except for
            #                special cases such as initial scanning of the package when ChemShell is
            #                initialised
            if chemsh.CHEMSH_DEBUG > -1:
                print()
                print(format.wrap('CHARMM forcefield library is not available in the current build. To use it please rebuild with `./setup --charmm-ff`.'))
            return ''


    @property
    def default_version(self):
        '''Version phrase of the CHARMM forcefield'''

        try:
            return info.charmm_ff_version
        except:
            if chems.CHEMSH_DEBUG > -1:
                print(format.wrap('CHARMM forcefield library is not available in the current build. To use it please rebuild with `./setup --charmm-ff`.'))
            return ''


    # exceptions have to be defined otherwise our algorithm can't deal with, for example CA and FE, at the same time
    @property
    def names2znums(self):
        self._names2znums = { 
                              'CAD' :  6, # NB: in some files CAD is cadmium!
                              'NPH' :  7,
                              'NA'  :  7,
                              'NA+' : 11,
                              'SOD' : 11,
                              'POT' : 19,
                              'CLA' : 17,
                              'CLAL': 17,
                              'CL-' : 17,
                              'CAL' : 20,
                              'FE'  : 26,
                              'CU2' : 29,
                              'CU2+': 29,
                              'ZN2' : 30,
                              'ZN2+': 30,
                              'BR'  : 35,
                              'BRA' : 35,
                              'DUM' :  0, # dummy atom
                            }

        return self._names2znums

    @property
    def residue_alias(self):
        self._residue_alias = { 'HIS'  :'HSD',
                                'HEC'  :'HEMC',
                                'HOH'  :'TIP3',
                                'TIP3W':'TIP3',
                                'TP3M' :'TIP3',
                                'dA'   :'ADE',
                                'dT'   :'THY',
                                'dC'   :'CYT',
                                'dG'   :'GUA',
                                'A'    :'ADE',
                                'U'    :'URA',
                                'C'    :'CYT',
                                'G'    :'GUA',
                                'dA3'  :'ADE',
                                'dT3'  :'THY',
                                'dC3'  :'CYT',
                                'dG3'  :'GUA',
                                'A3'   :'ADE',
                                'U3'   :'URA',
                                'C3'   :'CYT',
                                'G3'   :'GUA',
                                'dA5'  :'ADE',
                                'dT5'  :'THY',
                                'dC5'  :'CYT',
                                'dG5'  :'GUA',
                                'A5'   :'ADE',
                                'U5'   :'URA',
                                'C5'   :'CYT',
                                'G5'   :'GUA',
                              }

        return self._residue_alias


    @property
    def haems_rtf(self):
        '''RTF filename for haems'''

        return 'haems.rtf'


    @property
    def haems_prm(self):
        '''PRM filename for haems'''

        return 'haems-ligands.prm'


    @property
    def haems(self):
        '''List of supported haem residues'''

        from os.path     import join
        from ...io.files import rtf

        try:
            return self._haems

        except:
            dbuff = dict()
            filename = join(self.root_dir, self.haems_rtf)
            rtf.parse(filename, receiver=dbuff, residues_only=True)

            lbuff = list(dbuff.keys())
            for key in [ 'residues', 'mass', 'decl', 'defa', 'auto' ]:
                lbuff.remove(key)

            # other common names
            lbuff += [ 'HAEM', 'HEC', 'HEME', 'HEM' ]

            self._haems = lbuff

            return self._haems


    @property
    def hemes(self):
        '''Synonym of haeme'''

        return self.haems


    @property
    def backbone_protein(self):
        '''Atom names of amino acid backbone'''

        return ['CA','C','O','N','OT1','OT2','OXT','O1','O2']


    @property
    def amino_acids(self):
        '''List of supported amino acid residues'''

        return ['ALA','ARG','ASN','ASP',
                'CYS',
                'GLN','GLU','GLY','GTT', # GTT is glutamine having a gamma peptide linkage with its side chain's carboxyl group
                'HIS','HSD','HSE','HSP',
                'ILE',
                'LEU','LYS',
                'MET',
                'PHE','PRO',
                'SER',
                'THR','TRP','TYR',
                'VAL']


    @property
    def patchable_residues(self):
        '''Names of residues to patch'''

        return { 'GTT':'GLNK',
               }


    @property
    def waters(self):
        '''List of CHARMM residue names for water (the result is a subset of charmmm.solvents)'''

        return [ 'HOH', 'TP3M', 'TIP3', 'TP3O' ]


    @property
    def solvents(self):
        '''List of CHARMM residue names for solvents'''

        # YL TODO 27/02/2021: non-water solvents
        return [ ] + self.waters


    @property
    def ions(self):
        '''List of CHARMM residue names for ions'''

        return list(self.ion_symbols2names.values())


    @property
    def ion_symbols2names(self):
        '''Dictionary of symbols-to-names mapping'''

        return {
                 'Li' :'LIT',
                 'Na' :'SOD',
                 'K'  :'POT',
                 'Mg' :'MG',
                 'Ca' :'CAL',
                 'Rb' :'RUB',
                 'Cs' :'CES',
                 'Zn' :'ZN2',
                 'Cd' :'CD2',
                 'Cl' :'CLA',
                 'Br' :'BAR',
                 'OH' :'OH',
                 'SO4':'SO4',
                }


    @property
    def nucleotide_aliases(self):
        '''Mapping between ChemShell and CHARMM names'''

        return {
                'dA':'ADE',
                'dT':'THY',
                'dG':'GUA',
                'dC':'CYT',
                'A' :'ADE',
                'G' :'GUA',
                'C' :'CYT',
                'U' :'URA',
               }


    @property
    def nucleotide_patches(self):
        '''Patches for deoxyribonucleic acid'''

        return {
                'dA' :'DEOX',
                'dT' :'DEOX',
                'dG' :'DEOX',
                'dC' :'DEOX',
                'dA3':'DEOX',
                'dT3':'DEOX',
                'dG3':'DEOX',
                'dC3':'DEOX',
                'dA5':'DEOX',
                'dT5':'DEOX',
                'dG5':'DEOX',
                'dC5':'DEOX',
               }


    @property
    def nucleotide_patches_terminal(self):
        '''Patches for 3'/5'-terminal deoxyribonucleic acid'''

        return {
                'dA3':'3TER',
                'dT3':'3TER',
                'dG3':'3TER',
                'dC3':'3TER',
                'dA5':'5TER',
                'dT5':'5TER',
                'dG5':'5TER',
                'dC5':'5TER',
                'A3' :'3TER',
                'T3' :'3TER',
                'G3' :'3TER',
                'C3' :'3TER',
                'A5' :'5TER',
                'T5' :'5TER',
                'G5' :'5TER',
                'dC5':'5TER',
               }


    @property
    def nucleotides(self):
        '''List of CHARMM residue names for nucleotides'''

        return list(self.nucleotide_aliases.keys()) + list(set(self.nucleotide_aliases.values()))


    @property
    def backbone_nucleotide(self):
        '''Atom names of nucleotide backbone'''

        return ["C1'","C2'","C3'","C4'","C5'","O3'","O4'","O5'","P"]


    @property
    def rtfs2prms(self):
        ''''''

        self._rtfs2prms = { 'toph19.inp'           :'param19.inp',
                            'top_all22_prot.rtf'   :'par_all22_prot.prm',
                            'top_all35_ethers.rtf' :'par_all35_ethers.prm',
                            'placeholder1'         :'par_all36_prot.prm',    # there used to be both par_all36_prot.prm and par_all36m_prot.prm before v2021
                            'top_all36_carb.rtf'   :'par_all36_carb.prm',
                            'top_all36_cgenff.rtf' :'par_all36_cgenff.prm',
                            'top_all36_lipid.rtf'  :'par_all36_lipid.prm',
                            'top_all36_na.rtf'     :'par_all36_na.prm',
                            'top_all36_prot.rtf'   :'par_all36m_prot.prm',   # this is the only protein parameter files in v2021
                            'placeholder2'         :'par_hbond.inp',
                            'toppar_water_ions.str':'toppar_water_ions.str',
                            'haems.rtf'            :'haems-ligands.prm',
                            'ligands.rtf'          :'haems-ligands.prm',
                            'amino_acids.rtf'      :'amino_acids.prm',
                          }

        return self._rtfs2prms


    @property
    def prms2abspaths(self):
        '''A dictionary of PRM filenames and their absolute paths'''

        from os import path

        self._prms2abspaths = { 'param19.inp'          : path.join(self.default_location, 'param19.inp'          ),
                                'par_all22_prot.prm'   : path.join(self.default_location, 'par_all22_prot.prm'   ),
                                'par_all35_ethers.prm' : path.join(self.default_location, 'par_all35_ethers.prm' ),
                                'par_all36m_prot.prm'  : path.join(self.default_location, 'par_all36m_prot.prm'  ),
                                'par_all36_carb.prm'   : path.join(self.default_location, 'par_all36_carb.prm'   ),
                                'par_all36_cgenff.prm' : path.join(self.default_location, 'par_all36_cgenff.prm' ),
                                'par_all36_lipid.prm'  : path.join(self.default_location, 'par_all36_lipid.prm'  ),
                                'par_all36_na.prm'     : path.join(self.default_location, 'par_all36_na.prm'     ),
                                'par_all36_prot.prm'   : path.join(self.default_location, 'par_all36_prot.prm'   ),
                                'par_hbond.inp'        : path.join(self.default_location, 'par_hbond.inp'        ),
                                'toppar_water_ions.str': path.join(self.default_location, 'toppar_water_ions.str'),
                                'haems-ligands.prm'    : path.join(self.root_dir        , 'haems-ligands.prm'    ),
                                'amino_acids.prm'      : path.join(self.root_dir        , 'amino_acids.prm'      ),
                              }

        return self._prms2abspaths


    @property
    def titratables(self):
        '''A dict of titratable residues'''

        # ASP  3.9
        # GLU  4.3
        # LYS 10.5
        # ARG 12.0
        # HIS  6.08
        return { 'ASP': 'ASPP',
                 'GLU': 'GLUP',
                 'LYS': 'LYSB',
                 'ARG': 'ARGB',
                 'HIS':['HSD','HSE','HSP'],
               }


    @property
    def peptide_bonds(self):
        '''A dict of connectivity relationship which forms the peptide'''

        #     |
        #  HN-N
        #     |   HB1
        #     |   |   
        #  HA-CA--CB--
        #     |   |   
        #     |   HB2
        #   O=C
        #     |

        return { 'N' : ['HN','CA','C'],
                 'CA': ['N' ,'HA','C' ,'CB'],
                 'C' : ['CA','O' ,'N'],
                 'O' : ['C'],
               }


    @property
    def protonateds(self):
        '''Flattened list of values of dict self.titratables'''

        try:
            return self._protonateds
        except:
            self._protonateds = utils.iterutils.getFlattened(list(self.titratables.values()))
            return self._protonateds


    @property
    def default_rtfs(self):
        '''A list of all built-in RTF files'''

        from os         import listdir
        from os.path    import abspath, dirname, isfile, join

        # YL 26/07/2022: check if the database has been installed but not for initialisation
        if 'getmembers()' not in utils.miscutils.getCaller():
            try:
                listdir(self.default_location)
            except FileNotFoundError:
                format.printwrap('CHARMM forcefield library is not available in the current build. To use it please rebuild with `./setup --charmm-ff`.')
                exit(999)

        _rtfs = []
        for filename in listdir(self.default_location):
            if filename.endswith('.rtf') or filename.endswith('.inp') or filename.endswith('.str'):
                # YL 21/09/2021: prevent using toph19.inp which contains definitions of amino acids with no associated FF data
                if filename == 'toph19.inp':
                    continue
                _rtfs.append(abspath(join(self.default_location, filename)))

        return _rtfs


    @property
    def builtin_rtfs(self):
        '''A list of all built-in RTF files'''

        from os      import listdir
        from os.path import abspath, dirname, isfile, join

        _rtfs = []
        for filename in listdir(dirname(__file__)):
            if filename.endswith('rtf') and isfile(join(dirname(__file__), filename)):
                _rtfs.append(abspath(join(dirname(__file__), filename)))

        return _rtfs


# YL 16/02/2021: some parameter files have extension of .inp that are confused with RTF
#    @property
#    def default_prms(self):
#        '''A list of all built-in RTF files'''
#
#        from os      import listdir
#        from os.path import abspath, dirname, isfile, join
#
#        _prms = []
#        for filename in listdir(self.default_location):
#            if filename.endswith('.prm') or filename.endswith('.inp') or filename.endswith('.str'):
#                _prms.append(abspath(join(self.default_location, filename)))
#
#        return _prms


    @property
    def builtin_prms(self):
        '''A list of all built-in RTF files'''

        from os      import listdir
        from os.path import abspath, dirname, isfile, join

        _prms = []
        for filename in listdir(dirname(__file__)):
            if filename.endswith('prm') and isfile(join(dirname(__file__), filename)):
                _prms.append(abspath(join(dirname(__file__), filename)))

        return _prms


# this is a very slow method, not in use currently but may be useful in the future
#    @property
#    def residue_data_dicts(self):
#        '''A dict of dicts containing mapping from topology files to parsed residue-atoms data'''
#
#        from ...io.files import rtf
#
#        try:
#            return self._residue_data_dicts
#        except:
#            self._residue_data_dicts = rtf.getResidueDicts(self.default_rtfs+self.builtin_rtfs)
#            return self._residue_data_dicts


    @property
    def residue_names_dicts(self):
        '''A dict of dicts containing mapping from topology files to only residue names'''

        from ...io.files import rtf

        try:
            return self._residue_names_dicts
        except:
            self._residue_names_dicts = rtf.getResidueDicts(self.default_rtfs+self.builtin_rtfs, residues_only=True)
            return self._residue_names_dicts


    @property
    def root_dir(self):
        ''''''

        from os.path import abspath, dirname

        return abspath(dirname(__file__))


    def getFilesOfResidues(self, residues, rtfs=[], prms=[], return_nfounds=True):
        '''Return RTF and PRM files containing the given residues'''

        from os import path

        nfounds = 0
        # loop over missing residues
        for i, residue in enumerate(residues):
            residue_found = False
            # loop over built-in RTF files
            for rtf, data in self.residue_names_dicts.items():
                # YL 14/09/2021: because the default is prior to built-in (see `residue_names_dicts`
                #                above: self.default_rtfs+self.builtin_rtfs) we skip the rest searching
                #                when the first match is found
                #                TODO: an override_default switch?
                if residue_found:
                    break
                if residue in data:
                    residue_found = True
                    nfounds += 1
                    rtf_filename = path.basename(rtf)
                    print('\n Supplemented ChemShell/CHARMM RTF/PRM files for residue \"{}\":'.format(residue))
                    # absolute path
                    if rtf not in rtfs:
                        rtfs.append(rtf)
                        print(' {}'.format(rtf))
                    prm_filename = self.rtfs2prms[rtf_filename]
                    # absolute path
                    prm = self.prms2abspaths[prm_filename]
                    if prm not in prms:
                        prms.append(prm)
                        print(' {}'.format(prm))

        if return_nfounds:
            return rtfs, prms, nfounds
        else:
            return rtfs, prms


    def getFullPath(self, file):
        '''Return the full path to a file: return the str itself it is not found in the built-in or default libraries'''

        from os import path

        if path.isfile(file):
            file_path = file
        else:
            try:
                file_path = self.prms2abspaths[file]
            except KeyError:
                file_path = file
                print(f'\n WARNING: cannot find file {file}. Skipping...')

        return file_path


    @property
    def residues(self):
        '''Return all residue names as a list'''

        from numpy       import unique
        from ...io.files import rtf

        try:
            return self._residues
        except:
            self._residues = []
            builtin_rtf_res_dict = rtf.getResidueDicts(self.builtin_rtfs, residues_only=True)
            for f, d in builtin_rtf_res_dict.items():
                self._residues += [ r for r in list(d.keys()) if r not in ['residues','mass','decl','defa','auto'] ]
            default_rtf_res_dict = rtf.getResidueDicts(self.default_rtfs, residues_only=True)
            for f, d in default_rtf_res_dict.items():
                self._residues += [ r for r in list(d.keys()) if r not in ['residues','mass','decl','defa','auto'] ]

            self._residues = unique(self._residues).tolist()

            return self._residues



# wrap the current module around with the dummy class to use the property feature
modules[__name__] = __Wrapper()


