import sys
from pmx import library
from pmx.utils import get_ff_path, ff_selection

# ===============================
# Class for interactive selection
# ===============================
class InteractiveSelection:
    """Class containing functions 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=None, renumbered=True):
        self.m = m
        self.ffpath = ''
        if ff==None:
            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]
            valid_ids = []
            for r in self.m.chdic[self.mut_chain].residues:
                if isinstance(r.id,str):
                    valid_ids.append(r.id.replace(" ",""))
                else:
                    valid_ids.append(r.id)
            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('%6s-%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 %s 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-%s'
              % (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

# ===============================
# Helper functions
# ===============================
def _int_input():
    inp = input()
    try:
        inp = int(inp)
        return inp
    except:
        print('You entered "%s" -> Will try to proceed assuming this is a valid ID' % inp)
        return 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 _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
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
