#  Copyright (C) 2018 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@stfc.ac.uk'
# Oct 2018

import sys
import chemsh
from ...dl_py2f       import utils
from ...dl_py2f.utils import miscutils

__all__ = 'pdb2pqr'


# CHARMM aliases for nucleotides
charmm2pqr = {
               'dA' :'DA',
               'dT' :'DT',
               'dC' :'DC',
               'dG' :'DG',
               'A'  :'RA',
               'C'  :'RC',
               'G'  :'RG',
               'U'  :'RU',
               'dA3':'DA3',
               'dT3':'DT3',
               'dC3':'DC3',
               'dG3':'DG3',
               'A3' :'RA3',
               'C3' :'RC3',
               'G3' :'RG3',
               'U3' :'RU3',
               'dA5':'DA5',
               'dT5':'DT5',
               'dC5':'DC5',
               'dG5':'DG5',
               'A5' :'RA5',
               'C5' :'RC5',
               'G5' :'RG5',
               'U5' :'RU5',
             }

# map to ChemShell notations
pqr2chemsh = {
               'dA' :'dA',
               'dT' :'dT',
               'dC' :'dC',
               'dG' :'dG',
               'rA' :'A',
               'rC' :'C',
               'rG' :'G',
               'rU' :'U',
               'dA3':'dA3',
               'dT3':'dT3',
               'dC3':'dC3',
               'dG3':'dG3',
               'rA3':'A3',
               'rC3':'C3',
               'rG3':'G3',
               'rU3':'U3',
               'dA5':'dA5',
               'dT5':'dT5',
               'dC5':'dC5',
               'dG5':'dG5',
               'rA5':'A5',
               'rC5':'C5',
               'rG5':'G5',
               'rU5':'U5',
             }

def runExec(path,
            input,
            ff='charmm',
            ph_method='propka',
            output='',
            output_dir='_pdb2pqr',
            verbose=True,
            summary=True):
    '''Run PDB2PQR system command (only for PDB2PQR v2)'''

    from os.path  import abspath, dirname, join
    from ...base  import run
    from ...data  import references
    from ...io    import banner, format

    output_log = f'{output_dir}.out'
    output_err = f'{output_dir}.err'
    propka_log = f'{output_dir}.propka'

    utils.fileutils.archiveDir(output_dir)

    # output filename
    if not output:
        output = utils.fileutils.getBaseName(input) + '.pqr'

    # output level
    if chemsh.CHEMSH_DEBUG:
        verbose, summary = '--verbose', ''
    else:
        verbose, summary = '', ''
    verbose, summary = '--verbose', ''
        
    # force field format
    ff = ff.lower()

    # for example: pdb2pqr --ff=charmm --ffout=charmm --verbose --ph-calc-method=propka 2yli_atom_no_pca.pdb tmp.pqr
    # path to exec can be overridden by user-specified path
    binary = path or join(abspath(dirname(__file__)), 'src', 'pdb2pqr')
    syscmd = '{} --ff={} --ffout={} --ph-calc-method={} --include-header {} {} {}'.format(binary, ff, ff, ph_method, verbose, input, output)

    if chemsh.CHEMSH_DEBUG:
        print("\n PDB2PQR: system command:\n {}".format(syscmd))

    utils.miscutils.runSysCmd(syscmd, stdout=output_log, stderr=output_err)

    postprocessV2(output_dir, input, output)


# YL 19/03/2021: run of PDB2PQR v3 directly as a Python module
def runLib(input,
           ff='charmm',
           ph=7.0,
           output='',
           output_dir='_pdb2pqr',
           path='',
           verbose=True,
           summary=True):
    '''Run PDB2PQR as a Python module'''

    import sys
    from warnings import catch_warnings, resetwarnings, simplefilter
    from os.path  import abspath, dirname, join
    from logging  import basicConfig, captureWarnings, DEBUG, getLogger
    from sys      import stdout, stderr
    from numpy    import char
    from ...data  import references
    from ...io    import banner, format

    banner.printModuleBanner('ChemShell/PDB2PQR Interface')
#    print(' Please cite:')
#    format.printwrap(references.references['pdb2pqr'], prepend_blanc_line=False)

    # to suppress DeprecationWarning
    with catch_warnings() as w:
        simplefilter('ignore')
# YL 11/04/2022: pandas is no longer a dependency since v3.2.0, reported by HMS
        # check dependency pandas
#        try:
#            import pandas
#        except:
#            format.printwrap(">>> ChemShell/PDB2PQR ERROR: PDB2PQR dependency pandas not available. Exiting...\n")
#            raise
#            exit(112)
        # check dependency mmcif_pdbx
        try:
            import pdbx
        except:
            format.printwrap(">>> ChemShell/PDB2PQR ERROR: PDB2PQR dependency mmcif_pdbx not available. Exiting...\n")
            raise
#            exit(112)
        try:
            from .src.build.lib import pdb2pqr
            # PDB2PQR will search in sys.path for the dat directory
            sys.path.append(join(abspath(dirname(__file__)), 'src', 'build', 'lib'))
        except ModuleNotFoundError as e:
            format.printwrap(">>> ChemShell/PDB2PQR ERROR: Missing dependencies")
            print(f' {e}')
        except:
            format.printwrap(">>> ChemShell/PDB2PQR ERROR: PDB2PQR library not available in the current installation. Rebuild ChemShell with: ./setup --pdb2pqr\n")
#            exit(112)
    resetwarnings()

    basicConfig(filename=f'{output_dir}.err')
#    captureWarnings(True)

    stdout.flush()
    stderr.flush()

    # instantiate the args object with the two positional arguments
    # YL 02/07/2021: there is a bug in Python for embedded code resulting in
    #                AttributeError: 'module' object has no attribute 'argv'
    #                it's not fixed until Python 3.8
    #                see: https://bugs.python.org/issue32573
    #                we create a `sys.argv` if it doesn't exist
    if not hasattr(sys, 'argv'):
        sys.argv  = ['']

    # YL 06/01/2021: the code structure has been changed in the latest versions of PDB2PQR where pdb2pqr.main is overloaded and cannot be directly imported
    #                here __package__ must be the long name otherwise in biomolecule.py the class name becomes `pdb2pqr.pdb.ATOM` and cannot be matched
    package = 'chemsh.addons.pdb2pqr.src.build.lib.pdb2pqr'
    pdb2pqr.main = utils.modutils.importModule('main', join(abspath(dirname(__file__)), 'src', 'build', 'lib', 'pdb2pqr', 'main.py'), __package__=package)

    # argument parser
    pdb2pqrParser = pdb2pqr.main.build_main_parser()
    pdb2pqrArgs   = pdb2pqrParser.parse_args([input, output])

    # debug
    version = pdb2pqr._version.__version__
    if chemsh.CHEMSH_DEBUG > 4:
        print(f'\n Version: {version}')
        pdb2pqrParser.print_usage()

    output_log = f'{output_dir}.out'
    output_err = f'{output_dir}.err'

    utils.fileutils.archiveDir(output_dir)

    # output filename
    if not output:
        output = utils.fileutils.getBaseName(input) + '.pqr'

    # pass in other options
    pdb2pqrArgs.ff             = ff.upper()
    pdb2pqrArgs.ffout          = ff.upper()
    pdb2pqrArgs.include_header = True
    pdb2pqrArgs.verbose        = True
    # YL 19/08/2021: suppress warnings which are printed in _pdb2pqr.err anyway when no debug level
    #                PDB2PQR options: "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"
    if chemsh.CHEMSH_DEBUG >= 4:
        pdb2pqrArgs.log_level      ='DEBUG'
    elif chemsh.CHEMSH_DEBUG == 3:
        pdb2pqrArgs.log_level      ='INFO'
    elif chemsh.CHEMSH_DEBUG == 2:
        pdb2pqrArgs.log_level      ='WARNING'
    elif chemsh.CHEMSH_DEBUG == 1:
        pdb2pqrArgs.log_level      ='CRITICAL'
    elif chemsh.CHEMSH_DEBUG == 0:
        pdb2pqrArgs.log_level      ='ERROR'
    pdb2pqrArgs.pdb_output     = utils.fileutils.getBaseName(input) + '.out.pdb'
    pdb2pqrArgs.output_pqr     = output
    if ph != 7.0:
        pdb2pqrArgs.ph = ph

    # run
    with catch_warnings() as w:
        simplefilter('ignore')

        # we have to adapt and override the original main_driver function because
        # the current version v3.1.0 (as of 19/03/2021) strangely doesn't write the
        # lines of header and omitted atoms
#        pdb2pqr.main.main_driver(pdb2pqrArgs)

        # logger
        pdb2pqr.io.setup_logger(pdb2pqrArgs.output_pqr, pdb2pqrArgs.log_level)
        pdb2pqr.main.print_splash_screen(pdb2pqrArgs)

        # transform args
        args = pdb2pqr.main.transform_arguments(pdb2pqrArgs)

        # sanity checks
        pdb2pqr.main.check_files(args)
        pdb2pqr.main.check_options(args)

        # parsing
        definition = pdb2pqr.io.get_definitions()
        pdblist, is_cif = pdb2pqr.io.get_molecule(args.input_path)

        # create molecule objects
        biomolecule, definition, ligand = pdb2pqr.main.setup_molecule(pdblist, definition, args.ligand)
        biomolecule.set_termini(args.neutraln, args.neutralc)
        biomolecule.update_bonds()

        # run
        try:
            results = pdb2pqr.main.non_trivial(args=args,
                                               biomolecule=biomolecule,
                                               ligand=ligand,
                                               definition=definition,
                                               is_cif=is_cif)
        except ValueError as err:
            format.printwrap('{}'.format(err))
            print('\n ChemShell/PDB2PQR ERROR: Failed running PDB2PQR. Exiting...\n')
            # YL 12/05/2021: do not issue `exit` because the error may be fixed
            return err

        # convert the header (<class str>) to list of str
        header = char.add(results['header'].splitlines(), '\n').tolist()

        # YL TODO 19/03/2021: missed is a list of PDB2PQR <class Atom> objects and we can use them directly rather than parsing from the PQR header (see io.files.pqr.parse())
        missed = results['missed_residues']
        if len(missed):
            print(f'\n ChemShell/PDB2PQR WARNING: {len(missed)} atoms cannot be processed, see {output_err}')
        print(f'\n PDB2PQR v{version} successfully ran.')

        # write out the PQR: concatenate the header and content lines because v3.1.0 doesn't print the header but only throws out a warning instead
        pdb2pqr.main.print_pqr(args=args,
                               pqr_lines=header+results['lines'],
                               header_lines='',
                               missing_lines='',
                               is_cif=is_cif)

        # write out a PDB: concatenate the header and content lines because v3.1.0 doesn't print the header but only throws out a warning instead
        pdb_lines = pdb2pqr.io.print_biomolecule_atoms(biomolecule.atoms, chainflag=args.keep_chain, pdbfile=True)
        pdb2pqr.main.print_pdb(args=args,
                               pdb_lines=header+pdb_lines,
                               header_lines='',
                               missing_lines='',
                               is_cif=is_cif)

    resetwarnings()

    postprocessV3(output_dir, input, output)
    postprocessNucleotides(output_dir, input, output)


@miscutils.time
def run(*args, path='', **kwargs):
    '''Run PDB2PQR'''

    preprocess(*args, **kwargs)

    if path:
        return runExec(path, *args, **kwargs)
    else:
        return runLib(*args, **kwargs)


def preprocess(pdb2pqr_input, *args, ff='charmm', **kwargs):
    '''Preprocess the biofragment to become readable for PDB2PQR'''

    from ...objects import fragbio
    from ...io      import file

    # nucleotides must be renamed

    # load the biofragment
    bioFrag = fragbio.FragmentBio()
    file.parse(pdb2pqr_input, 'pdb', receiver=bioFrag, mute=True)

    bioFrag.detectNucleotides(aliases=charmm2pqr, delete_terminal_hydrogen=True, delete_terminal_phosphate=True)
    bioFrag.save(pdb2pqr_input, version='pdb2pqr')


def postprocessV2(prefix, input, output):
    '''Check if the run was successful for PDB2PQR v2'''

    from os.path  import isfile, join
    from ...io    import format

    # archive files
    utils.fileutils.archiveFile(input              , dir=prefix, suffix='', copy=True) # retain the input file
    utils.fileutils.archiveFile(output             , dir=prefix, suffix='')
    utils.fileutils.archiveFile(f'{prefix}.out.pdb', dir=prefix, suffix='')
    utils.fileutils.archiveFile(f'{prefix}.log'    , dir=prefix, suffix='')
    utils.fileutils.archiveFile(f'{prefix}.err'    , dir=prefix, suffix='')
    utils.fileutils.archiveFile(utils.fileutils.getBaseName(input)+'.out.pdb', dir=prefix, suffix='')

    # check if the run was successful
    output_pqr = join(prefix, output)
    output_out = join(prefix, f'{prefix}.log')

    # failed
    if not isfile(output_pqr):
        with open(output_out) as fp:
            for lineno, line in enumerate(fp):
                if line.strip().lower().startswith('error encountered:'):
                    print("\n >>> However PDB2PQR has failed:")
                    format.printwrap(line.strip())
                    print()
                    return 112

    # succeeded
    else:
        print(f'\n Output PQR written to:\n {output}')
        return 0


def postprocessV3(prefix, input, output):
    '''Check if the run was successful for PDB2PQR v3'''

    from os.path  import isfile, join
    from ...io    import format

    # archive files
    utils.fileutils.archiveFile(input              , dir=prefix, suffix='', copy=True) # retain the input file
    utils.fileutils.archiveFile(output             , dir=prefix, suffix='')
    utils.fileutils.archiveFile(f'{prefix}.out.pdb', dir=prefix, suffix='')
    utils.fileutils.archiveFile(f'{prefix}.log'    , dir=prefix, suffix='')
    utils.fileutils.archiveFile(f'{prefix}.err'    , dir=prefix, suffix='')
    utils.fileutils.archiveFile(utils.fileutils.getBaseName(input)+'.out.pdb', dir=prefix, suffix='')

    # check if the run was successful
    output_pqr = join(prefix, output)
    output_out = join(prefix, f'{prefix}.log')

    # failed
    if not isfile(output_pqr):
        with open(output_out) as fp:
            for lineno, line in enumerate(fp):
                if line.strip().lower().startswith('error encountered:'):
                    print("\n >>> However PDB2PQR has failed:")
                    format.printwrap(line.strip())
                    print()
                    return 112

    # succeeded
    else:
        print(f'\n Output PQR written to:\n {output_pqr}')
        return 0


def postprocessNucleotides(prefix, input, output):
    '''Revert the nucleotide residue names'''

    from os.path    import join
    from numpy      import copy
    from ...objects import fragbio
    from ...io      import file

    output_pdb  = join(prefix, f'{prefix}.out.pdb')
    output_pqr = join(prefix, output)

    # load the biofragment
    bioFrag = fragbio.FragmentBio()
    file.parse(output_pdb, 'pdb', receiver=bioFrag, choose_locations=False, mute=True)
    # must force to connect otherwise it will connect later with PQR radii
    bioFrag.connect()

    has_nucleotide = bioFrag.detectNucleotides(aliases=pqr2chemsh, delete_terminal_hydrogen=False, delete_terminal_phosphate=False)

    # do nothing if there's no nucleotide
    # TODO: protein-DNA complex
    if not has_nucleotide:
        return

    # to avoid reconnect because the PQR radii are strange and will cause C2 and C4 to connect
    tmpFrag = fragbio.FragmentBio()
    file.parse(output_pqr, 'pqr', receiver=tmpFrag, mute=True)
    bioFrag._master.charges = tmpFrag.charges
    bioFrag._master.radii   = tmpFrag.radii

    bioFrag.save(output_pqr)


