# -*- coding: utf-8 -*-
# ProDy: A Python Package for Protein Dynamics Analysis
# 
# Copyright (C) 2010-2012 Ahmet Bakan
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program 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 General Public License for more details.
#  
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>

"""This module defines MSA analysis functions."""

__author__ = 'Ahmet Bakan'
__copyright__ = 'Copyright (C) 2010-2012 Ahmet Bakan'

from numpy import arange

from .analysis import *
from prody import LOGGER

__all__ = ['showMSAOccupancy', 'showShannonEntropy', 'showMutinfoMatrix']


def pickSequence(msa):
    """Pick a sequence without gaps and deletions and return its residue 
    numbers and labels to be used as indices and X-axis label, or a pair
    of **None** at failure."""
    
    try:
        counts = calcMSAOccupancy(msa, 'row', count=True)
    except TypeError:
        return None, None
    else:
        length = msa.numResidues()
        split, msa.split = msa.split, True
        rows = (counts == length).nonzero()[0]
        for row in rows:
            try:
                label, indices = msa[row].getLabel(), msa[row].getResnums()
            except:
                break
            else:
                return (indices, 'Residue number ({0})'.format(label)) 
        return None, None


def showMSAOccupancy(msa, occ='res', indices=None, count=False, **kwargs):
    """Show a bar plot of occupancy calculated for :class:`.MSA` instance *msa*
    using :func:`.calcMSAOccupancy`.  *occ* may be ``'res'`` or ``'col'``, or a
    a pre-calculated occupancy array.  If x-axis *indices* are not specified,
    they will be inferred from *msa* or given *label* that may correspond to
    a sequence in the msa.
    
    Occupancy is plotted using :func:`~matplotlib.pyplot.bar` function."""
    
    try:
        numseq, lenseq = msa.numSequences(), msa.numResidues()
    except AttributeError:
        raise TypeError('msa must be an MSA instance')
    
    try:
        length = len(occ)
    except TypeError:
        raise TypeError("occ must be 'res', 'col', or an occupancy array")
    
    try:
        sw = occ.startswith
    except (TypeError, AttributeError):
        try:
            ndim = occ.ndim
        except AttributeError:
            raise TypeError("occ must be 'res', 'col', or an occupancy array")
        else:
            if ndim != 1:
                raise ValueError('occ must be a 1-dimensional array')
    else:
        occ = calcMSAOccupancy(msa, occ, count)
        length = len(occ)

    if length == numseq:
        xlabel = kwargs.pop('xlabel', None) or 'MSA sequence index'
        
    if length == lenseq:
        label = kwargs.pop('label', None)
        if label is not None:
            try:
                indices = msa[label].getResnums()
            except:
                LOGGER.info('Specified label not in msa.')
        xlabel = kwargs.pop('xlabel', None) or label or 'MSA column index'
        if xlabel is None and indices is None:
            indices, xlabel = pickSequence(msa)
    if indices is None:
        indices = arange(1, length + 1)
    
    ylabel = kwargs.pop('ylabel', 'Count' if count else 'Occupancy')  
    title = kwargs.pop('title', None)
    format = kwargs.pop('format', True)
    import matplotlib.pyplot as plt
    show = plt.bar(indices, occ, **kwargs)
    if format:
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        if title is None:
            title = 'Occupancy: ' + msa.getTitle()
        plt.title(title)
    return show


def showShannonEntropy(entropy, indices=None, **kwargs):
    """Show a bar plot of Shannon *entropy* array.  :class:`.MSA` instances 
    or Numpy character arrays storing sequence alignments are also accepted 
    as *entropy* argument, in which case :func:`.calcShannonEntropy` will 
    be used for calculations.  *indices* may be residue numbers, when not 
    specified they will be inferred from *msa* or indices starting from 1 
    will be used.
    
    Entropy is plotted using :func:`~matplotlib.pyplot.bar` function."""
    
    msa = None
    try:
        ndim = entropy.ndim
    except AttributeError:
        msa = entropy
        entropy = calcShannonEntropy(msa)
        ndim = entropy.ndim
    
    if ndim != 1:
        raise ValueError('entropy must be a 1D array')

    msa = kwargs.pop('msa', msa) 
    xlabel = kwargs.pop('xlabel', None)
    if indices is None:
        length = len(entropy)    
        if msa is not None:
            indices, xlabel = pickSequence(msa)
        if indices is None:
            indices = arange(1, length + 1)
        xlabel = xlabel or 'MSA column index'
    
    ylabel = kwargs.pop('ylabel', 'Shannon entropy')  
    title = kwargs.pop('title', None)
    format = kwargs.pop('format', True)
    import matplotlib.pyplot as plt
    show = plt.bar(indices, entropy, **kwargs)
    if format:
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        if title is None:
            if msa is None:  
                title = 'Entropy'
            else:
                title = 'Entropy: ' + str(msa)
        plt.title(title)
    return show


def showMutinfoMatrix(mutinfo, *args, **kwargs):
    """Show a heatmap of mutual information array.  :class:`.MSA` instances 
    or Numpy character arrays storing sequence alignment are also accepted 
    as *mutinfo* argument, in which case :func:`.buildMutinfoMatrix` will 
    be used for calculations.  Note that x, y axes contain indices of the
    matrix starting from 1.
    
    Mutual information is plotted using :func:`~matplotlib.pyplot.imshow`
    function. vmin and vmax values can be set by user to achieve better
    signals using this function."""
    
    msa = None
    try:
        ndim, shape = mutinfo.ndim, mutinfo.shape
    except AttributeError:
        msa = mutinfo
        mutinfo = buildMutinfoMatrix(mutinfo)
        ndim, shape = mutinfo.ndim, mutinfo.shape

    msa = kwargs.pop('msa', msa)
    if ndim != 2:
        raise ValueError('mutinfo must be a 2D matrix')
    y, x = shape
    if x != y:
        raise ValueError('mutinfo matrix must be a square matrix')
    
    kwargs.setdefault('interpolation', 'nearest') 
    kwargs.setdefault('origin', 'lower')
    
    if msa is not None:
        indices, msalabel = pickSequence(msa)
        if indices is not None:
            start = indices[0] + 0.5
            end = start + x
            extent = [start, end, start, end]
        else:
            extent = [0.5, x + 0.5, 0.5, y + 0.5]
    else:
        msalabel = None
        extent = [0.5, x + 0.5, 0.5, y + 0.5]
    
    xlabel = kwargs.pop('xlabel', None)
    if xlabel is None:
        xlabel = msalabel or 'MSA column index'
    title = kwargs.pop('title', None)
    format = kwargs.pop('format', True)
    
    import matplotlib.pyplot as plt
    show = plt.imshow(mutinfo, extent=extent, *args, **kwargs), plt.colorbar()
    
    if format:
        plt.xlabel(xlabel)
        plt.ylabel(xlabel)
        if title is None:
            if msa is None:  
                title = 'Mutual Information'
            else:
                title = 'Mutual Information: ' + str(msa)
        plt.title(title)
    return show





if __name__ == '__main__':
    from prody import *
    msa = parseMSA('piwi_seed.sth')
    print(repr(msa))
    msa = refineMSA(msa, label=msa[0][0])
    print(repr(msa))
    print(calcMSAOccupancy(msa, 'row', count=True))
    print(pickSequence(msa))
