#!/usr/bin/env python

# pmx  Copyright Notice
# ============================
#
# The pmx source code is copyrighted, but you can freely use and
# copy it as long as you don't change or remove any of the copyright
# notices.
#
# ----------------------------------------------------------------------
# pmx is Copyright (C) 2006-2013 by Daniel Seeliger
#
#                        All Rights Reserved
#
# Permission to use, copy, modify, distribute, and distribute modified
# versions of this software and its documentation for any purpose and
# without fee is hereby granted, provided that the above copyright
# notice appear in all copies and that both the copyright notice and
# this permission notice appear in supporting documentation, and that
# the name of Daniel Seeliger not be used in advertising or publicity
# pertaining to distribution of the software without specific, written
# prior permission.
#
# DANIEL SEELIGER DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS
# SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS.  IN NO EVENT SHALL DANIEL SEELIGER BE LIABLE FOR ANY
# SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER
# RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF
# CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
# CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
# ----------------------------------------------------------------------

"""Program to insert mutated residues in structure files for
free energy simulations.
"""

import sys
import argparse
from pmx import library
from pmx.model import Model
from pmx.parser import read_and_format
from pmx.utils import get_ff_path, ff_selection
from pmx.alchemy import mutate
from pmx.scripts.cli import check_unknown_cmd

# resinfo
dna_one_letter = {'A': 'adenosine',
                  'C': 'cytosine',
                  'G': 'guanine',
                  'T': 'thymine'}

rna_one_letter = {'A': 'adenosine',
                  'C': 'cytosine',
                  'G': 'guanine',
                  'U': 'uracil'}

# ================
# Helper functions
# ================
def _print_sorted_dict(d):
    for key in sorted(d.iterkeys()):
        print("{0:>5}{1:>15}".format(key, d[key]))


def _int_input():
    inp = input()
    try:
        inp = int(inp)
        return inp
    except:
        print('You entered "%s" -> Try again' % inp)
        return None


def _check_residue_name(res):
    if res.resname == 'LYS':
        if res.has_atom('HZ3'):
            res.set_resname('LYP')
    elif res.resname == 'HIS':
        if res.has_atom('HD1') and \
           res.has_atom('HE2'):
            res.set_resname('HIP')
        elif res.has_atom('HD1') and not res.has_atom('HE2'):
            res.set_resname('HID')
        elif not res.has_atom('HD1') and res.has_atom('HE2'):
            res.set_resname('HIE')
    elif res.resname == 'ASP':
        if res.has_atom('HD2'):
            res.set_resname('ASH')
    elif res.resname == 'GLU':
        if res.has_atom('HE2'):
            res.set_resname('GLH')
    elif res.resname == 'CYS':
        if not res.has_atom('HG'):
            print(' Cannot mutate SS-bonded Cys %d' % res.id, file=sys.stderr)


def _ask_next():
    sys.stdout.write('\nApply another mutation [y/n]? ')
    res = input().lower()
    if res == 'y':
        return True
    elif res == 'n':
        return False
    else:
        return _ask_next()


def _match_mutation(m, ref_m, ref_chain, ref_resid):
    """Matches chain/indices of two Models. Given the chain and resid of a
    reference Model (ref_m), return the resid of the Model (m).

    Parameters
    ----------
    m: Model
        model you want to mutate
    ref_m : Model
        reference model
    ref_chain: str
        chain of residue in reference model
    ref_resid: int
        resid of residue in reference model

    Returns
    -------
    resid: int
        residue ID of model that corresponds to the chain/resid in the
        reference.
    """

    # select non-solvent residues
    res = [r for r in m.residues if r.moltype not in ['water', 'ion']]
    ref_res = [r for r in ref_m.residues if r.moltype not in ['water', 'ion']]
    # check they have same len
    assert len(res) == len(ref_res)

    # iterate through all residue pairs
    resmap = {}
    for r, rref in zip(res, ref_res):
        # first, check that the sequence is the same
        if r.resname != rref.resname:
            raise ValueError('residue %s in the input file does not match '
                             'residue %s in the input reference file'
                             % (r.resname, rref.resname))
        # then, create a dict to map (chain_id, res_id) in the reference
        # to the residues ID in the input file
        resmap[(rref.chain_id, rref.id)] = r.id

    resid = resmap[(ref_chain, ref_resid)]
    print('log_> Residue {ref_id} (chain {ref_ch}) in file {ref} mapped to residue '
          '{m_id} in file {m} after renumbering'.format(ref_ch=ref_chain,
                                                        ref_id=ref_resid,
                                                        ref=ref_m.filename,
                                                        m_id=resid,
                                                        m=m.filename))

    return resid


# ===============================
# Class for interactive selection
# ===============================
class InteractiveSelection:
    """Class containing fucntions related to the interactive selection of
    residues to be mutated.

    Parameters
    ----------
    m : Model object
        instance of pmx.model.Model
    ffpath : str
        path to forcefield files

    Attributes
    ----------
    mut_resid : int
        index of residue to be mutated
    mut_resname : str
        one-letter code of target residue

    """

    def __init__(self, m, ff, renumbered=True):
        self.m = m
        self.ffpath = get_ff_path(ff)

        # get selection
        if renumbered is True:
            self.mut_chain = None
        elif renumbered is False:
            self.mut_chain = self.select_chain()

        self.mut_resid = self.select_residue()
        self.mut_resname = self.select_mutation()

    def select_chain(self):
        """Ask for the chain id to mutate.
        """
        # show selection
        valid_ids = [c.id for c in self.m.chains]
        print('\nSelect a chain:')
        for c in self.m.chains:
            print('{0:>6}'.format(c.id))

        # select id
        selected_chain_id = None
        while selected_chain_id is None:
            sys.stdout.write('Enter chain ID: ')
            selected_chain_id = input()
            if selected_chain_id is not None and selected_chain_id not in valid_ids:
                print('Chain id %s not among valid IDs -> Try again' % selected_chain_id)
                selected_chain_id = None
        return selected_chain_id

    def select_residue(self):
        """Ask for the residue id to mutate.
        """
        # show selection if we do not need chain ID
        if self.mut_chain is None:
            valid_ids = [r.id for r in self.m.residues]
            print('\nSelect residue to mutate:')
            for i, r in enumerate(self.m.residues):
                if r.moltype not in ['water', 'ion']:
                    sys.stdout.write('%6d-%s-%s' % (r.id, r.resname, r.chain_id))
                    if (i+1) % 6 == 0:
                        print("")
        elif self.mut_chain is not None:
            valid_ids = [r.id for r in self.m.chdic[self.mut_chain].residues]
            print('\nSelect residue to mutate:')
            for i, r in enumerate(self.m.chdic[self.mut_chain].residues):
                if r.moltype not in ['water', 'ion']:
                    sys.stdout.write('%6d-%s-%s' % (r.id, r.resname, r.chain_id))
                    if (i+1) % 6 == 0:
                        print("")
        print("")

        # select id
        selected_residue_id = None
        while not selected_residue_id:
            sys.stdout.write('Enter residue number: ')
            selected_residue_id = _int_input()
            if selected_residue_id is not None and selected_residue_id not in valid_ids:
                print('Residue id %d not among valid IDs -> Try again' % selected_residue_id)
                selected_residue_id = None
        return selected_residue_id

    def select_mutation(self):
        """Ask which residue to mutate to.
        """

        residue = self.m.fetch_residue(idx=self.mut_resid, chain=self.mut_chain)
        if residue.moltype == 'protein':
            aa = self.select_aa_mutation(residue)
        elif residue.moltype in ['dna', 'rna']:
            aa = self.select_nuc_mutation(residue)
        return aa

    def select_aa_mutation(self, residue):
        """Selection for protein residues.
        """

        _check_residue_name(residue)
        print('\nSelect new amino acid for %s-%s: ' % (residue.id, residue.resname))
        sys.stdout.write('Three- or one-letter code (or four-letter for ff specific residues): ')
        if residue.resname in ['HIE', 'HISE', 'HSE']:
            rol = 'X'
        elif residue.resname in ['HIP', 'HISH', 'HSP']:
            rol = 'Z'
        elif residue.resname in ['GLH', 'GLUH', 'GLUP']:
            rol = 'J'
        elif residue.resname in ['ASH', 'ASPH', 'ASPP']:
            rol = 'B'
        elif residue.resname in ['LYN', 'LYS', 'LSN']:
            rol = 'O'
        else:
            rol = library._one_letter[residue.resname]
        aa = None
        ol = list(library._aacids_dic.keys())
        tl = list(library._aacids_dic.values())
        ffpathlower = self.ffpath.lower()
        if 'amber' in ffpathlower:
                ol = list(library._aacids_ext_amber.keys())
                tl = list(library._aacids_ext_amber.values())
        if 'opls' in ffpathlower:
                ol = list(library._aacids_ext_oplsaa.keys())
                tl = list(library._aacids_ext_oplsaa.values()) + ['ASPP', 'GLUP', 'LSN']
        if 'charmm' in ffpathlower:
                ol = list(library._aacids_ext_charmm.keys())
                tl = list(library._aacids_ext_charmm.values())

        while aa is None:
            aa = input().upper()
            # some special residues:
            #   CM - deprotonated cysteine
            #   YM - deprotonated tyrosine
            if aa == 'CM':
                sys.stdout.write('Special case for deprotonated residue')
            elif len(aa) != 1 and len(aa) != 3 and len(aa) != 4:
                sys.stdout.write('Nope!\nThree- or one-letter code (or four-letter for ff specific residues): ')
                aa = None
            elif (len(aa) == 1 and aa not in ol+['B', 'J', 'O', 'X', 'Z']) or \
                 (len(aa) == 3 and aa not in tl) or \
                 (len(aa) == 4 and aa not in tl):
                sys.stdout.write('Unknown aa "%s"!\nThree- or one-letter code (or four-letter for ff specific residues): ' % aa)
                aa = None
            if aa and (len(aa) == 3 or len(aa) == 4):
                aa = library._ext_one_letter[aa]
        print('Will apply mutation %s->%s on residue %s-%d'
              % (rol, aa, residue.resname, residue.id))
        return aa

    def select_nuc_mutation(self, residue):
        """Selection for nucleic acids.
        """
        aa = None
        print('\nSelect new base for %s-%s: ' % (residue.id, residue.resname))
        sys.stdout.write('One-letter code: ')
        while aa is None:
            aa = input().upper()
            if residue.moltype == 'dna' and aa not in ['A', 'C', 'G', 'T']:
                sys.stdout.write('Unknown DNA residue "%s"!\nOne-letter code: ' % aa)
                aa = None
            elif residue.moltype == 'rna' and aa not in ['A', 'C', 'G', 'U']:
                sys.stdout.write('Unknown RNA residue "%s"!\nOne-letter code: ' % aa)
                aa = None
            if aa:
                print('Will apply mutation %s->%s on residue %s-%d'
                      % (residue.resname[1], aa, residue.resname, residue.id))
            return aa


# =============
# Input Options
# =============
def parse_options():
    parser = argparse.ArgumentParser(description='''
This script applies mutations of residues in a structure file for subsequent
free energy calculations. It supports mutations to protein, DNA, and RNA
molecules.

The mutation information and dummy placements are taken from the hybrid residue
database "mutres.mtp". The best way to use this script is to take a pdb/gro file
that has been written with pdb2gmx with all hydrogen atoms present.

By default, all residues are renumbered starting from 1, so to have unique
residue IDs. If you want to keep the original residue IDs, you can use the flag
--keep_resid. In this case, you will also need to provide chain information
in order to be able to mutate the desired residue. Alternatively, if you would
like to use the original residue IDs but these have been changed, e.g. by gromacs,
you can provide a reference PDB file (with chain information too) using the --ref
flag. The input structure will be mutated according to the IDs chosen for the
reference structure after having mapped the two residue indices.

The program can either be executed interactively or via script. The script file
simply has to consist of "residue_id target_residue_name" pairs (just with some
space between the id and the name), or "chain_id residue_id target_residue_name"
if you are keeping the original residue IDs or providing a reference structure.

The script uses an extended one-letter code for amino acids to account for
different protonation states. Use the --resinfo flag to print the dictionary.

''', formatter_class=argparse.RawTextHelpFormatter)

    exclus = parser.add_mutually_exclusive_group()

    parser.add_argument('-f',
                        metavar='infile',
                        dest='infile',
                        type=str,
                        help='Input structure file in PDB or GRO format. '
                        'Default is "protein.pdb"',
                        default='protein.pdb')
    parser.add_argument('-fB',
                        metavar='infileB',
                        dest='infileB',
                        type=str,
                        help='Input structure file of the B state in PDB '
                        'or GRO format (optional).',
                        default=None)
    parser.add_argument('-o',
                        metavar='outfile',
                        dest='outfile',
                        type=str,
                        help='Output structure file in PDB or GRO format. '
                        'Default is "mutant.pdb"',
                        default='mutant.pdb')
    parser.add_argument('-ff',
                        metavar='ff',
                        dest='ff',
                        type=str.lower,
                        help='Force field to use. If none is provided, \n'
                        'a list of available ff will be shown.',
                        default=None)
    parser.add_argument('--script',
                        metavar='script',
                        dest='script',
                        type=str,
                        help='Text file with list of mutations (optional).',
                        default=None)
    exclus.add_argument('--keep_resid',
                        dest='renumber',
                        help='Whether to renumber all residues or to keep the\n'
                        'original residue IDs. By default, all residues are\n'
                        'renumbered so to have unique IDs. With this flags set,\n'
                        'the original IDs are kept. Because the IDs might not\n'
                        'be unique anymore, you will also be asked to choose\n'
                        'the chain ID where the residue you want to mutate is.',
                        default=True,
                        action='store_false')
    exclus.add_argument('--ref',
                        metavar='',
                        dest='ref_infile',
                        help='Provide a reference PDB structure from which to map\n'
                        'the chain and residue IDs onto the file to be mutated (-f).\n'
                        'This can be useful when wanting to mutate a file that\n'
                        'has had its residues renumbered or the chain information\n'
                        'removed (e.g. after gmx grompp). As in the --keep_resid\n'
                        'option, if --ref is chosen, you will need to provide chain\n'
                        'information either interactively or via the --script flag.',
                        default=None)
    parser.add_argument('--resinfo',
                        dest='resinfo',
                        help='Show the list of 3-letter -> 1-letter residues',
                        default=False,
                        action='store_true')

    args, unknown = parser.parse_known_args()
    check_unknown_cmd(unknown)

    # ------------------
    # residue dictionary
    # ------------------
    if args.resinfo is True:
        moltype = Model(args.infile).moltype
        if moltype == 'protein':
            print('\n ---------------------------')
            print(' Protein residues dictionary')
            print(' ---------------------------')
            _print_sorted_dict(library._ext_one_letter)
            print(' ---------------------------\n')
        elif moltype == 'dna':
            print('\n -----------------------')
            print(' DNA residues dictionary')
            print(' -----------------------')
            _print_sorted_dict(dna_one_letter)
            print(' -----------------------\n')
        elif moltype == 'rna':
            print('\n -----------------------')
            print(' RNA residues dictionary')
            print(' -----------------------')
            _print_sorted_dict(rna_one_letter)
            print(' -----------------------\n')
        else:
            print('\n ---------------------------')
            print(' Protein residues dictionary')
            print(' ---------------------------')
            _print_sorted_dict(library._ext_one_letter)
            print(' ---------------------------')
            print(' DNA residues dictionary')
            print(' ---------------------------')
            _print_sorted_dict(dna_one_letter)
            print(' ---------------------------')
            print(' RNA residues dictionary')
            print(' ---------------------------')
            _print_sorted_dict(rna_one_letter)
            print(' ---------------------------\n')
        exit()

    # ------------
    # ff selection
    # ------------
    if args.ff is None:
        args.ff = ff_selection()

    return args


def main(args):

    # input variables
    infile = args.infile
    infileB = args.infileB
    outfile = args.outfile
    ff = args.ff
    script = args.script
    renumber = args.renumber
    ref_infile = args.ref_infile

    # initialise Model
    m = Model(infile, renumber_residues=renumber, bPDBTER=True,
              rename_atoms=True, scale_coords='A')

    # if reference structure provided, initialise that Model too
    if ref_infile is not None:
        ref_m = Model(ref_infile, renumber_residues=False,
                      bPDBTER=True, rename_atoms=True, scale_coords='A')

    # if script is provided, do the mutations in that file
    # ----------------------------------------------------
    if script is not None:
        # 1) defualt: renumbering and not providing a ref struct
        if renumber is True and ref_infile is None:
            mutations_to_make = read_and_format(script, "is")
            # modify mut lists in mutations_to_make so that they have same
            # len of 3, both in the case where renumber is True or False
            mutations_to_make = [[None]+x for x in mutations_to_make]
        # 2) not renumbering and not providing a ref struct
        elif renumber is False and ref_infile is None:
            mutations_to_make = read_and_format(script, "sis")
        # 3) renumbering and providing a ref struct
        elif renumber is True and ref_infile is not None:
            # we read mutations according to numbering in the reference
            ref_mutations_to_make = read_and_format(script, "sis")
            mutations_to_make = []
            # we map these onto the file to be mutated
            for ref_mut in ref_mutations_to_make:
                m_resid = _match_mutation(m=m, ref_m=ref_m,
                                          ref_chain=ref_mut[0],
                                          ref_resid=ref_mut[1])
                # mut = [No Chain, resid in m, target residue]
                mut = [None, m_resid, ref_mut[2]]
                mutations_to_make.append(mut)

        # 4) NOT ALLOWED: providing a ref struct while not renumbering. This is
        # not allowed because we want to make sure there are unique indices for
        # each residue

        for mut in mutations_to_make:
            _check_residue_name(m.fetch_residue(idx=mut[1], chain=mut[0]))
            mutate(m=m,
                   mut_chain=mut[0],
                   mut_resid=mut[1],
                   mut_resname=mut[2],
                   ff=ff,
                   refB=infileB,
                   inplace=True,
                   verbose=True)

    # if script not provided, interactive selection
    # ---------------------------------------------
    else:
        do_more = True
        while do_more:
            # if no reference provided, choose from infile (Model m)
            if ref_infile is None:
                sele = InteractiveSelection(m=m, ff=ff, renumbered=renumber)
            # if reference IS provided, then choose from reference, then map
            # mutation onto infile (Model m)
            elif ref_infile is not None:
                sele = InteractiveSelection(m=ref_m, ff=ff, renumbered=False)
                sele.mut_resid = _match_mutation(m=m, ref_m=ref_m,
                                                 ref_chain=sele.mut_chain,
                                                 ref_resid=sele.mut_resid)
                sele.mut_chain = None

            # we have the needed info ==> carry out the mutation
            mutate(m=m,
                   mut_chain=sele.mut_chain,
                   mut_resid=sele.mut_resid,
                   mut_resname=sele.mut_resname,
                   ff=ff,
                   refB=infileB,
                   inplace=True,
                   verbose=True)

            # ask whether to do more mutations or stop
            if not _ask_next():
                do_more = False

    m.write(outfile)
    print('')
    print('mutations done...........')
    print('')


def entry_point():
    args = parse_options()
    main(args)


if __name__ == '__main__':
    entry_point()
