'''
Main PDB Structure Mapping (biostructmap) module. This package implements
techniques to map data onto a protein PDB structure.

Rationale:

Whilst many protein-related data (such as genetic polymorphisms) can be easily
related to a reference protein sequence, it is slightly more involved to map
this data onto a protein structure. For example, when viewing polymorphisms
that have arisen as a result of immune selection pressure, it is instructive
to be able to view the location of these polymorphisms in the context of the
protein 3D structure. Furthermore, it may also be useful to average this
data over a radius around a central residue. In our example case of protein
polymoprhisms, this allows the user to identify polymorphic hotspots on the
protein structure.

We can extend this concept further, and consider genetic tests for phenomena
which arise as a result of selection pressure at the level of protein structure.
For example, antibody-mediated immune selection pressure gives rise to balancing
selection at the level of population genetics. A number of methods exist to
identify regions of the genome under balancing selection, including Tajima's D.
Tajima's D is often calculated as a sliding window over a genome. However, when
the origins of balancing selection occur at the level of protein structure,
it can be useful to also consider protein spatial information when computing
Tajima's D using a sliding window. In essence, we could calculate Tajima's D
using a 3D sliding window over a protein structure. The biostructmap package
automates this process.

Details:

The biostructmap package makes extensive use of the Biopython Bio.PDB module for
PDB file parsing and integration with calculation of secondary structure
and relative solvent accessibility via the DSSP program. Calculation of Tajima's
D is performed using DendroPy, although several optimisations are performed
before passing a multiple sequence alignment to DendroPy in order to speed
up calculation of Tajima's D, such as removal of non-polymorphic sites and
memoization of results from previous windows when calculating a sliding window
value of Tajima's D. For a large number of sequences, we suggest interfacing
with an external command line tool to calculate Tajima's D (or other metric of
interest).

'''
from __future__ import absolute_import, division, print_function

import contextlib
from copy import deepcopy
import json
import itertools
from tempfile import NamedTemporaryFile
from Bio.PDB import DSSP, PDBIO, PDBParser, FastMMCIFParser
from Bio import AlignIO
from Bio.PDB.MMCIF2Dict import MMCIF2Dict
from . import pdbtools, gentests
from .pdbtools import match_pdb_residue_num_to_seq, SS_LOOKUP_DICT, mmcif_sequence_to_res_id
from .map_functions import (_tajimas_d, _default_mapping, _snp_mapping,
                            _map_amino_acid_scale, _count_residues,
                            _nucleotide_diversity, _wattersons_theta,
                            _shannon_entropy, _normalized_shannon_entropy)
from .seqtools import (align_protein_to_dna, _construct_sub_align, align_protein_sequences)

mapping_methods = {"default":_default_mapping,
                   "tajimasd":_tajimas_d,
                   "snps": _snp_mapping,
                   "aa_scale": _map_amino_acid_scale,
                   "count_residues": _count_residues,
                   "nucleotide_diversity": _nucleotide_diversity,
                   "wattersons_theta": _wattersons_theta,
                   "shannon_entropy": _shannon_entropy,
                   "normalized_shannon_entropy": _normalized_shannon_entropy}

@contextlib.contextmanager
def open_if_string(path_or_file, mode):
    '''Opens a file if given a string but doesn't try to re-open a file-like
    object. Used for functions that accept either a string (filename) or
    a file-like object.
    '''
    if isinstance(path_or_file, str):
        f = file_to_close = open(path_or_file, mode)
    else:
        f = path_or_file
        file_to_close = None
    try:
        yield f
    finally:
        if file_to_close:
            file_to_close.close()


class DataMap(dict):
    '''
    A class to hold a mapping of data to some PDB object. Extends the `dict`
    class, adding a few methods to allow correspondence between data and PDB
    chain object.

    Should be called as `DataMap(results, structure=structure, params=params)`,
    where `results` would be a dictionary of mapped data over a PDB chain for a
    particular set of parameters.
    '''
    def __init__(self, *args, **kw):
        '''Initialise a DataMap object, which stores data mapped to a PDB
        structure as a dictionary, but also provides a link to original chain object
        and parameters used to generate data.

        Args:
            *args: Standard dict args
            **kwargs: Standard dict kwargs. Kwargs `structure` and `params` are
                reserved for reference to the PDB structure object and analysis
                parameters respectively, and are required. These are removed
                from the list of kwargs before passing to dict __init__ method.
        '''
        self.structure = kw.pop('structure')
        self.params = kw.pop('params')
        super(DataMap, self).__init__(*args, **kw)


    def write_data_to_pdb_b_factor(self, default_no_value=0, fileobj=None,
                                   scale_factor=1):
        '''Write mapped data to PDB B-factor column, and save as a PDB file.

        This method allows for data visualisation over the PDB structure
        using Pymol. In particular, the `spectrum` command in Pymol can color
        residues according to B-factor values, and it is common practice to
        store data of interest in the B-factor column of PDB files.

        Args:
            default_no_value (int, optional): Default value to write to B-factor
                column if a non-numeric value exists for a particular data point
                (e.g. a None value). Default value is 0, although this could be
                altered depending the range of possible data points. For
                example, you might use a value such as 99 or -99 to indicate a
                missing data point, and then filter this out when visualising
                the data in Pymol.
            fileobj (str/file-like object, optional): Output file.
                Can be either a string representing the file path or a file-like
                object. If not provided, this defaults to a string
                describing the data being written, in the form of
                "pdbid_param1-value1_param2-value2.pdb".
            scale_factor (int/float): A factor by which to scale (multiply) all
                data values before writing to B-factor column. This is needed
                when dealing with values that are too small to be represented
                in the PDB file B-factor column.

        Returns:
            None
        '''
        # Make copy of structure object.
        # We do this so that we are not overwriting B-factor values in the
        # original PDB structure.
        # This may fail in earlier versions of Biopython if the structure
        # contains disordered residues - An appropriate fix for deepcopying
        # disordered residues has been added to Biopython (>v1.70?).
        serial_number_generator = itertools.count(start=1)
        _structure = deepcopy(self.structure)
        first_model = sorted(_structure.models)[0]
        #Set all B-factor fields to zeros/default value
        for chain in _structure[first_model].chains.values():
            for residue in chain:
                _data = self.get(residue.get_full_id()[2:4], default_no_value)
                if _data is None:
                    _data = default_no_value
                    _scale_factor = 1
                else:
                    _scale_factor = scale_factor
                for atom in residue:
                    # Currently mmcif files don't get assigned an atom serial number
                    # by Bio.PDB, which breaks trying to write to a PDB file.
                    # Simple fix is just to assign a sequential serial number.
                    # May not line up with _atom_site.id in mmcif file.
                    if atom.get_serial_number() is None:
                        atom.set_serial_number(next(serial_number_generator))
                    if atom.is_disordered():
                        for altloc in atom.disordered_get_id_list():
                            atom.disordered_select(altloc)
                            atom.set_bfactor(float(_data * _scale_factor))
                    else:
                        atom.set_bfactor(float(_data * _scale_factor))
        pdb_io = PDBIO()
        pdb_io.set_structure(_structure.structure)
        # Write structure to PDB file
        if fileobj is None:
            fileobj = _structure.pdbname + '_' + self._parameter_string() + '.pdb'
        pdb_io.save(file=fileobj, preserve_atom_numbering=True)
        return None

    def write_to_atom(self, fileobj, sep=',', model=None):
        '''Write score for each atom in a structure to a file, based on
        a data dictionary mapping output score to residue number.

        Each line of the output file contains an entry for a single atom.

        Args:
            fileobj (str/object): Output file name/path or file-like object.
            sep (str, optional): Seperator between residue and data.
                Defaults to `,`.
        Returns:
            None
        '''
        is_mmcif = self.structure._mmcif
        if is_mmcif:
            raise TypeError("Can't map to atom serial with mmcif file!")
        if model is None:
            # Use the first model. Will be the only model unless it's an NMR structure.
            model_ = self.structure[sorted(self.structure.models)[0]]
        else:
            model_ = self.structure[model]
        #For each atom in the structure, write an output score based on the data
        #given, presuming (key,value) in a dictionary with key corresponding to
        #a residue number.
        with open_if_string(fileobj, 'w') as f:
            f.write(sep.join(('atom_serial', 'score\n')))
            for res in sorted(self):
                residue = model_[res[0]][res[1]]
                for atom in residue:
                    data_pt = [str(x) for x in [atom.serial_number, self[res]]]
                    line = sep.join(data_pt) + '\n'
                    f.write(line)
        return


    def write_data_to_json(self, fileobj):
        '''Write data on contents of each 3D window and related output value.

        Data is written in JSON format. Residues are numbered according to
        PDB file numbering.

        Args:
            fileobj (str/object): Output file name/path or file-like object.
        Returns:
            None
        '''
        nearby = self.structure.nearby()
        data_to_write = []
        for res in sorted(self):
            data_to_write.append({'residue': str(res), 'score': self[res],
                                  'nearby': str(nearby[res])})
        with open_if_string(fileobj, 'w') as f:
            json.dump(data_to_write, f, indent=2)
        return

    def write_residue_data_to_csv(self, fileobj, sep=',', ref=None):
        '''Write score for each residue in a structure to a file, based on
        a data dictionary mapping output score to residue number.

        Each line of the output file contains an entry for a single residue.
        This method additionally allows for scores to be aligned to a reference
        sequence for easy comparison with other metrics or a conventional
        reference sequence.

        Args:
            fileobj (str/object): Output file name/path or file-like object.
            sep (str, optional): Seperator between residue and data.
                Defaults to `,`.
            ref (dict, optional): A dictionary of reference sequences (as strings),
                accessed by chain ID.
        Returns:
            None
        '''
        if ref is None:
            # Use existing seqs for each chain
            ref = self.structure.sequences
        data_to_write = []

        # Get data relative to PDB numbering.
        for res in sorted(self):
            data_to_write.append((res, self[res]))

        # Make lookup dict...
        # Convert PDB Residue code to something like: ('A', 123)
        pdb_to_seq = self.structure._map_pdb_numbering_to_reference(ref, map_to_dna=False)
        
        converted_data = [(*pdb_to_seq[res], score) for res, score in data_to_write if res in pdb_to_seq]

        file_contents = '\n'.join([sep.join(str(y) for y in x) for x in converted_data])
        header = sep.join(['chain', 'reference_residue', 'score\n'])
        with open_if_string(fileobj, 'w') as f:
            f.write(header)
            f.write(file_contents)
        return

    def _parameter_string(self):
        '''Create descriptive string from parameter values'''
        param_string = '_'.join(["{k}-{v}".format(k=key, v=value)
                                 for key, value in sorted(self.params.items())])
        return param_string


class Structure(object):
    '''A class to hold a PDB structure object.

    Attributes:
        structure: Underlying Bio.PDB.Structure.Structure object
        sequences (dict): A dictionary of all protein sequences within
            structure, accessed by chain id.
        models (dict): Dictionary of all models within structure, accessed by
            model id.
        pdbname (str): A descriptive name for the PDB file.
    '''
    def __init__(self, pdbfile, pdbname='pdb_file', mmcif=False):
        '''Initialise PDB Structure object.

        Args:
            pdbfile (str/file-object): A filename string or a file-like object
                that contains a PDB file.
            pdbname (str, optional): A descriptive name for the PDB file. This is used
                as part of the default naming options when writing output files.
            mmcif (bool, optional): Set to true if reading a PDBx/mmCIF file,
                otherwise this defaults to reading a PDB file.
        '''
        # Create PDB parser
        if mmcif:
            parser = FastMMCIFParser()
        else:
            parser = PDBParser()
        self._mmcif_dict = None
        self._mmcif = mmcif
        #Get Bio.PDB structure
        self._pdbfile = pdbfile
        self.structure = parser.get_structure(pdbname, self.pdb_file())
        #Get PDB sequences
        if mmcif:
            self.sequences = pdbtools.get_mmcif_canonical_seq(self.mmcif_dict())
        else:
            self.sequences = pdbtools.get_pdb_seq(self.pdb_file())
        self.models = {model.get_id():Model(self, model) for
                       model in self.structure}
        self.pdbname = pdbname
        self._nearby = {}

    def __iter__(self):
        '''Iterate over all models within structure'''
        for key in sorted(self.models):
            yield self.models[key]

    def __getitem__(self, key):
        '''Get a model from structure'''
        return self.models[key]

    def pdb_file(self):
        '''Return the PDB file object, which can either be a string, or a
        file-like object.

        Notes:
            If it is a file-like object, sets read point to start of the
            file or stream. This enables us to perform analysis on PDB files
            downloaded from RCSB and stored in memory (as a StringIO object).

        Returns:
            str/file object: Either a string representing a filename or a file
                object with read point set to the start of the file or stream.
        '''
        if not isinstance(self._pdbfile, str):
            self._pdbfile.seek(0)
        return self._pdbfile

    def mmcif_dict(self):
        '''Return the mmcif dictionary.

        Only applicable if using an mmCIF file.

        Returns:
            dict: A dictionary containing mmCIF data.
        '''
        if self._mmcif and self._mmcif_dict is None:
            self._mmcif_dict = MMCIF2Dict(self.pdb_file())
        elif not self._mmcif:
            raise TypeError("Not an mmCIF file!")
        return self._mmcif_dict

    def nearby(self, radius=15, atom='all'):
        '''Take a Bio.PDB Structure object, and find all residues within a
        radius of a given residue.

        Note that this method uses the first model within the structure for
        all distance calculations.

        Args:
            radius (int/float): Radius within which to find nearby residues for
                each residue in the structure.
            atom (str): The atom with which to compute distances. By default
                this is 'all', which gets all non-heterologous atoms. Other
                potential options include 'CA', 'CB' etc. If an atom is not
                found within a residue object, then method reverts to using
                'CA'.

        Returns:
            dict: A dictionary containing a list of nearby residues for each
                residue in the structure.
        '''
        parameter_key = (radius, atom)
        # Run on first model in structure
        first_model = sorted(self.models)[0]
        #Calculate distance matrix and store it for retrieval in future queries.
        if parameter_key not in self._nearby:
            dist_map = pdbtools.nearby(self.structure[first_model], radius, atom)
            self._nearby[parameter_key] = dist_map
        return self._nearby[parameter_key]

    def map(self, data, method='default', ref=None, radius=15, selector='all',
            rsa_range=None, map_to_dna=False, method_params=None):
        '''Perform a mapping of some parameter or function to a pdb structure,
        with the ability to apply the function over a '3D sliding window'.

        The residues within a radius of a central residue are passed to the
        mapping function, which computes an output value for the central
        residue. This is performed for each residue in the structure.
        For calculation of Tajima's D, reference sequence is a genomic sequence.

        Args:
            data (dict/object): Generally a dictionary containing data to
                to be mapped over structure. The exact form for this data
                depends on the function being used to map data - each mapping
                function is passed this data object.
                Often the data will be in the form of a dictionary mapping
                multiple chains to a single shared data object (for example,
                when performing a calculation with multiple chains that all
                align to the one reference sequence).
                e.g.:
                    {('A', 'B'): data_object}
            method (str): A string representing a method for mapping data.
                Internally, these strings are keys for a dictionary of
                functions, and these can be extended with custom user-provided
                functions if desired.
            ref (dict): A reference protein sequence for each chain.
                Each chain to be mapped must have a reference sequence supplied,
                with one dictionary key per chain:
                    {'A': 'a sequence string...',
                     'B': 'another sequence string...',
                     ...}
                For calculation of Tajima's D, this is a reference genomic
                sequence. Data to be mapped should be aligned/referenced
                according to this reference sequence.
            radius (int/float, optional): The radius (Angstrom) over which to
                select nearby residues for inclusion within each 3D window.
                This defaults to 15 Angstrom, which is the typical maximum
                dimension for an antibody epitope.
            selector (str, optional): A string indicating the atom with which
                to compute distances between residues. By default this is 'all',
                which gets all non-heterologous atoms. Other potential options
                include 'CA', 'CB' etc. If an atom is not found within a residue
                object, then the selection method reverts to using 'CA'.
            rsa_range (tuple, optional): A tuple giving (minimum, maximum)
                values of relative solvent accessibility with which to filter
                all residues on (ie. map method will ignore all residues
                outside this range). This is useful when wanting to examine only
                surface exposed residues.
            map_to_dna (bool, optional): Set True if the mapping method involves
                aligning to a DNA sequence. Defaults to False.
            method_params (dict): Additional parameters to pass to a data
                aggregation method. Dictionary keys/values are function keyword
                arguments and associated values.

        Returns:
            biostructmap.DataMap: A dictionary-like object which contains mapped
                values for each residue (key). This object extends the standard
                dict type, adding methods to allow writing of data to PDB
                B-factor columns for easy viewing using Pymol or other similar
                programs.
        '''
        #Note: This method attempts to deal with 3 different ways of identifying
        #residue position: i) Within a PDB file, residues are labelled with a
        #number, which should increment according to position in sequence, but
        #is not necessarily gapless, and may include negative numbering in
        #order to preserve alignment with similar sequences. ii) Residue
        #numbering according to the position of the residue within the
        #protein sequence extracted from the PDB file. That is, the first
        #residue in the sequence is numbering '1', and incrementing from there.
        #iii) Residue numbering according to a reference sequence provided, or
        #according to a dna sequence (which could include introns).
        if method in mapping_methods:
            method = mapping_methods[method]

        if method_params is None:
            method_params = {}

        if map_to_dna and ref is None:
            raise ValueError("Must provide a reference DNA sequence if you "\
                             "are mapping to DNA.")
        elif ref is None:
            ref = self.sequences

        # Generate a map of nearby residues for each residue in pdb file.
        # Numbering is according to pdb residue numbering from file
        residue_map = self.nearby(radius=radius, atom=selector)

        # Map pdb numbering by file to the reference sequence
        # (dna or protein) provided, as long as the residues exists within the PDB
        # structure (ie has coordinates)
        pdbnum_to_ref = self._map_pdb_numbering_to_reference(ref, map_to_dna)
       
        results = {}

        #For each residue within the sequence, apply a function and return result.
        for residue in residue_map:
            if rsa_range:
                residues = self._filter_rsa(residue_map[residue], rsa_range)
                if residue not in residues:
                    results[residue] = None
                    continue
            else:
                residues = residue_map[residue]
            results[residue] = method(self, data, residues, pdbnum_to_ref, **method_params)
        params = {'radius':radius, 'selector': selector}
        return DataMap(results, structure=self, params=params)

    def _map_pdb_numbering_to_reference(self, ref, map_to_dna=False):
        '''Create a lookup dictionary mapping PDB numbering as given by Biopython to
        the indices of a reference sequence.

        Args:
            ref (dict): A dictionary of reference sequences accessed by chain ID.
            map_to_dna (bool, optional): Set to true if reference sequences are DNA sequences.
                Will align DNA to protein sequences using Exonerate (if available),
                otherwise will perform a simple translation (first reading frame) and align
                translated protein to PDB sequence.
        Returns:
            dict: A map of PDB numbering (key) to reference sequence index (value)
        '''
        # Use the first model. Will be the only model unless it's an NMR structure.
        model = self[sorted(self.models)[0]]
        # Create a map of pdb sequence index (1-indexed) to pdb residue
        # numbering from file
        is_mmcif = self._mmcif
        if is_mmcif:
            mmcif_dict = self.mmcif_dict()
            _, seq_index_to_pdb_numb = mmcif_sequence_to_res_id(mmcif_dict)
        else:
            seq_index_to_pdb_numb = match_pdb_residue_num_to_seq(model,
                                                                 self.sequences)

        pdb_index_to_ref = {}
        # For each protein chain, map provided reference sequence to PDB residue identifier.
        # Output should be in the form {('A', 17): ('A', (' ', 24, ' ')), ...}
        for chain_id, ref_seq in ref.items():
            if map_to_dna:
                chain_pdbindex_to_ref = align_protein_to_dna(self.sequences[chain_id], ref_seq)
            #Generate mapping of pdb sequence index to reference sequence (also indexed by position)
            else:
                chain_pdbindex_to_ref, _ = align_protein_sequences(self.sequences[chain_id],
                                                                   ref_seq)

            pdb_index_to_ref.update({(chain_id, key): (chain_id, value) for
                                     key, value in chain_pdbindex_to_ref.items()})

        #Finally, map pdb numbering by file to the reference sequence
        #(dna or protein) provided, as long as the residues exists within the PDB
        #structure (ie has coordinates)
        pdbnum_to_ref = {seq_index_to_pdb_numb[x]:pdb_index_to_ref[x] for x in
                         pdb_index_to_ref if x in seq_index_to_pdb_numb}
        return pdbnum_to_ref

    def _filter_rsa(self, residues, rsa_range):
        '''
        Function to remove residues with relative solvent accessibility
        values outside the requested range.

        If a residue solvent accessibility cannot be calculated by DSSP,
        then that residue is ignored.

        Args:
            residues (list): List of residues to filter.
            rsa_range (list/tuple): A tuple/list of the form (min, max),
                that gives the minimum and maximum relative solvent
                accessibility values with which to filter residues. Note that
                these values should be between 0 and 1.

        Returns:
            list: A list of residues (int) that have a relative solvent
                accessibility within the required range (inclusive).
        '''
        first_model = sorted(self.models)[0]
        filtered_residues = []
        for residue in residues:
            chain_id = residue[0]
            rsa = self[first_model][chain_id].rel_solvent_access()
            if (residue[1] in rsa and rsa[residue[1]] is not None and
                    rsa_range[0] <= rsa[residue[1]] <= rsa_range[1]):
                filtered_residues.append(residue)
        return filtered_residues


class Model(object):
    '''A class to hold a PDB model object.

    Attributes:
        model: Underlying Bio.PDB.Model.Model object
        dssp (dict): DSSP results for the model, ONLY IF model is the first in
            PDB file. Otherwise set to an empty dict. This is because the DSSP
            program will only read the first model in a PDB file.
        chains (dict): Dictionary of all chains within model, accessed by chain
            id.
    '''
    def __init__(self, structure, model):
        '''Initialise A PDB Model object.

        Args:
            structure (biostructmap.Structure): Parent structure object.
            model (Bio.PDB.Model.Model): Bio.PDB Model object.
        '''
        self._id = model.get_id()
        self.model = model
        self._parent = structure
        self.chains = {chain.get_id():Chain(self, chain) for
                       chain in self.model}
        self._dssp = None # May not need this, so don't compute unless needed.

    def __iter__(self):
        '''Iterate over all chains in model'''
        for key in sorted(self.chains):
            yield self.chains[key]

    def __getitem__(self, key):
        '''Get selected chain'''
        return self.chains[key]

    def dssp(self):
        '''Get DSSP object'''
        if self._dssp is not None:
            return self._dssp
        #DSSP only works on the first model in the PDB file
        if isinstance(self._parent.pdb_file(), str) and self._id == 0:
            try:
                self._dssp = DSSP(self.model, self._parent.pdb_file())
            except OSError:
                self._dssp = DSSP(self.model, self._parent.pdb_file(),
                                  dssp="mkdssp")
        elif self._id == 0:
            if self._parent._mmcif:
                suffix = '.cif'
            else:
                suffix = '.pdb'
            with NamedTemporaryFile(mode='w', suffix=suffix) as temp_pdb_file:
                temp_pdb_file.write(self._parent.pdb_file().read())
                temp_pdb_file.flush()
                try:
                    self._dssp = DSSP(self.model, temp_pdb_file.name)
                except OSError:
                    self._dssp = DSSP(self.model, temp_pdb_file.name,
                                      dssp="mkdssp")
        else:
            self._dssp = {}
        return self._dssp

    def parent(self):
        '''Get parent Structure object'''
        return self._parent

    def get_id(self):
        '''Get model ID'''
        return self._id


class Chain(object):
    '''A class to hold a PDB chain object.

    Attributes:
        chain: Underlying Bio.PDB.Chain.Chain object
        dssp (dict): DSSP results for the parent model, ONLY IF model is the
            first in PDB file. Otherwise set to an empty dict. This is because
            the DSSP program will only read the first model in a PDB file.
        sequence (str): Chain protein sequence. Extracted from PDB file header
            if present, otherwise constructed from PDB atom records for that
            chain.
    '''
    def __init__(self, model, chain):
        '''Initialise A PDB Chain object.

        Args:
            model (biostructmap.Model): Parent Model object.
            chain (Bio.PDB.Chain.Chain): Bio.PDB Chain object.
        '''
        self._id = chain.get_id()
        self.chain = chain
        self._parent = model
        #Refer to model for DSSP data
        self.dssp = self._parent.dssp
        #Some PDB files do not contain sequences in the header, and
        #hence we need to parse the atom records for each chain
        try:
            self.sequence = model.parent().sequences[self.get_id()]
        except KeyError:
            if self._parent._parent._mmcif:
                raise KeyError("mmCIF file doesn't contain sequences " \
                              "for chain {chain_id}".format(chain_id=self._id))
            self.sequence = pdbtools.get_pdb_seq_from_atom(chain)
            self._parent.parent().sequences[self.get_id()] = self.sequence
        self._rsa = {}

    def __iter__(self):
        '''Iterate over all residues in the chain'''
        for residue in self.chain:
            yield residue

    def __getitem__(self, key):
        '''Get residue within chain'''
        return self.chain[key]

    def parent(self):
        '''Get parent Model object'''
        return self._parent

    def get_id(self):
        '''Get chain ID'''
        return self._id

    def rel_solvent_access(self):
        '''Use Bio.PDB DSSP tools to calculate relative solvent accessibility.

        Returns:
            dict: A dictionary with RSA values for each residue (key).
        '''
        if not self._rsa:
            rsa = {}
            for residue in self.chain:
                key = (self.get_id(), residue.get_id())
                if key in self.dssp():
                    try:
                        rsa[key[1]] = float(self.dssp()[key][3])
                    except ValueError:
                        rsa[key[1]] = None
            self._rsa = rsa
        return self._rsa

    def secondary_structure(self, numeric_ss_code=False):
        '''Use DSSP to calculate secondary structure elements.

        Notes:
            Secondary structure assignment notation follows that of the DSSP
            program.

            The DSSP codes for secondary structure used here are:
                =====     ====
                Code      Structure
                =====     ====
                 H        Alpha helix (4-12)
                 B        Isolated beta-bridge residue
                 E        Strand
                 G        3-10 helix
                 I        Pi helix
                 T        Turn
                 S        Bend
                 -       None
                =====     ====
            If kwarg `numeric_ss_code` == True, then secondary structure values
            will be returned as integers according to the following lookup
            table:
            {
            'H': 0,
            'B': 1,
            'E': 2,
            'G': 3,
            'I': 4,
            'T': 5,
            'S': 6,
            '-', 7
            }
            These are provided in pdbtools.SS_LOOKUP_DICT, as well as the
            reverse lookups (i.e. {0: 'H', ...}).

        Args:
            numeric_ss_code (bool): Return secondary structure elements as a
                numeric lookup code if True. Set to False by default.

        Returns:
            dict: A dictionary with secondary structure assignment (value)
                for each residue (key) within the chain.
        '''
        keys = [x for x in self.dssp().keys() if x[0] == self._id]
        ss_dict = {x[1]: self.dssp().property_dict[x][2] for x in keys}
        if numeric_ss_code:
            return {key:SS_LOOKUP_DICT[item] for key, item in ss_dict.items()}
        else:
            return ss_dict

    def residue_to_atom_map(self):
        '''Return a map of residue number (per numbering in PDB file) to atom
        number for this chain.

        Returns:
            dict: A dictionary mapping residue number (key) to a tuple (value)
                containing all atom serial numbers for that residue.
        '''
        is_mmcif = self._parent._parent._mmcif
        if not is_mmcif:
            mapping = {residue.id[1]:tuple(atom.serial_number for atom in residue)
                       for residue in self.chain}
        else:
            raise TypeError("Can't map to atom serial with mmcif file!")
        return mapping

    def _filter_rsa(self, residues, rsa_range):
        '''
        Function to remove residues with relative solvent accessibility
        values outside the requested range.

        If a residue solvent accessibility cannot be calculated by DSSP,
        then that residue is ignored.

        Args:
            residues (list): List of residues to filter.
            rsa_range (list/tuple): A tuple/list of the form (min, max),
                that gives the minimum and maximum relative solvent
                accessibility values with which to filter residues. Note that
                these values should be between 0 and 1.

        Returns:
            list: A list of residues (int) that have a relative solvent
                accessibility within the required range (inclusive).
        '''
        rsa = self.rel_solvent_access()
        filtered_residues = [x for x in residues if rsa.get(x, None) and
                             rsa_range[0] <= rsa[x] <= rsa_range[1]]
        return filtered_residues


class SequenceAlignment(object):
    '''A class to hold a multiple sequence alignment object.

    This class is a basic wrapper for a Bio.Align.MultipleSequenceAlignment
    object, and allows calculation of Tajima's D over selected codons, as well
    as methods for slightly more efficient selection of sub-alignments based on
    selected codons, which may or may not be continuous.

    Attributes:
        alignment (Bio.Align.MultipleSequenceAlignment): Multiple Sequence
            Alignment object from Bio.Align.MultipleSequenceAlignment.
    '''
    def __init__(self, alignfile, file_format='fasta'):
        '''Initialises a SequenceAlignment object with a multiple sequence
        alignment file or file-like object.

        Args:
            alignfile (str/file handle): A multiple sequence alignment
                file path (str) or file-like object. This is the same as the
                Bio.AlignIO.read(...) inputs.
            file_format (str, optional): File format that the sequence
                alignment is stored as. This is the same as the arguments for
                Bio.AlignIO.read(...), and defaults to 'fasta'.
        '''
        self.alignment = AlignIO.read(alignfile, file_format)
        self._alignment_position_dict = None
        self._isolate_ids = None

    def __getitem__(self, key):
        return self.alignment[key]

    def get_alignment_position_dict(self):
        '''Returns a dictionary with single base pair alignments for each
        position in the multiple sequence alignment.

        This method utilises memoization to speed up multiple slicing operations
        on MultipleSeqAlignment objects, which can be expensive. This is
        patricularly useful when constructing a sub-alignment representing
        selected residues in a protein structure, as we would typically have
        to make many such sub-alignments when iterating over all residues within
        a structure.

        Returns:
            dict: A dictionary containing Bio.Align.MultipleSeqAlignment objects
                with a single base pair represented for each indexed position.
                Dictionary is of the form {int: MultipleSeqAlignment, ...}.
        '''
        if self._alignment_position_dict is None:
            self._alignment_position_dict = {}
            for i in range(len(self.alignment[0])):
                self._alignment_position_dict[i] = self.alignment[:, i]

        return self._alignment_position_dict

    def get_isolate_ids(self):
        '''Returns the isolate IDs for each sequence in the multiple sequence
        alignment.

        Returns:
            list: List of isolates/strains as strings.
        '''
        if self._isolate_ids is None:
            self._isolate_ids = [seq.id for seq in self.alignment]
        return self._isolate_ids


    def tajimas_d(self, window=None, step=3, protein_ref=None, genome_ref=None,
                  output_protein_num=False):
        '''Calculate Tajima's D on a SequenceAlignment object.

        If no window parameter is passed to the function, then the function
        calculates Tajima's D over the whole sequence alignment and
        returns a single numerical result.

        If a window size is given, then the function returns a dictionary
        of Tajima's D values over a sliding window.

        Note:

        Args:
            window (int, optional): The size of the sliding window over which
                Tajima's D is calculated. Default is None, in which case a
                single Tajima's D value is calculated for the multiple sequence
                alignment.
            step (int, optional): Step size for sliding window calculation.
                Default step size of 3 (ie. one codon).
            protein_ref (str, optional): A reference protein sequence which is
                used to optionally determine the protein coding region of the
                genome, and only calculate Tajima's D over these regions.
            genome_ref (str/Bio.Seq.Seq, optional): A genome reference sequence
                that is aligned to the multiple sequence alignment. Required if
                providing a protein reference sequence.
            output_protein_num (bool, optional): If set to True, will also
                return Tajima's D values for each position in the protein
                reference sequence.
        Returns:
            float/dict: If window parameter is None, returns a single value for
                Tajima's D. Otherwise a dict of genome window midpoint to
                calculated Tajima's D values is returned
            list (optional): Map of protein residue number to calculated
                Tajima's D values as tuple pairs of the form (residue, TajimasD)
                This is not returned as a dictionary as some residues may have
                multiple mapped Tajima's D value depending on step size.
        '''
        # If given a protein reference sequence, align genome to reference
        # sequence, and only perform a Tajima's D analysis over the protein
        # coding region. Otherwise perform a Tajima's D test over the full
        # genomic sequence.
        if genome_ref is not None and protein_ref is not None:
            #Align genome to reference
            prot_to_genome = align_protein_to_dna(protein_ref, ''.join([x for x in genome_ref]))
            #Get sorted list of codons
            codons = [prot_to_genome[x] for x in sorted(prot_to_genome)]
            #Construct a sub-alignment
            alignment = _construct_sub_align(self, codons, fasta=True)
        elif genome_ref is None and protein_ref is None:
            alignment = self.alignment
        elif protein_ref is None:
            raise TypeError("Missing protein_ref assignment")
        elif genome_ref is None:
            raise TypeError("Missing genome_ref assignment")

        #Perform Tajima's D calculation
        try:
            tajd = gentests.tajimas_d(alignment, window, step)
        except TypeError:
            print("Error calculating Tajima's D. Please check inputs to " +
                  "Tajima's D function.")
            raise
        if window is not None and protein_ref is not None:
            #Unpack list of codons
            codon_list = [x for sublist in codons for x in sublist]
            #Reference to genome numbering
            tajd_ref_to_genome = {codon_list[int(x) - 1]: tajd[x] for x in tajd}
            #Dictionary comprehension to reverse and unpack map of protein
            #position to genome codons. Initial dictionary is in the form
            # {1:(1,2,3), 2:(4,5,6,),...}, hence the need for the multiple
            #layers within the dictionary comprehension (needed to unpack each
            #codon tuple).
            genome_to_prot = {i:x for x in prot_to_genome for i in prot_to_genome[x]}
            tajd_ref_to_protein = [(genome_to_prot[x], tajd_ref_to_genome[x])
                                   for x in tajd_ref_to_genome]
            if output_protein_num:
                return tajd_ref_to_genome, tajd_ref_to_protein
            else:
                return tajd_ref_to_genome
        else:
            return tajd
