__author__ = "Julia Goodrich"
__copyright__ = "Copyright 2007-2012, The Cogent Project"
__credits__ = ["Jens Reeder","Julia Goodrich"]
__license__ = "GPL"
__version__ = "1.5.3"
__maintainer__ = "Julia Goodrich"
__email__ = "julia.goodrich@colorado.edu"
__Status__ = "Development"


from copy import copy
from types import GeneratorType

from numpy import transpose
from numpy.random import multinomial

from cogent.util.unit_test import FakeRandom
from cogent.core.sequence import Sequence
from cogent.parse.flowgram_parser import parse_sff
from cogent.parse.flowgram import Flowgram
from cogent.core.alignment import SequenceCollection

default_floworder = "TACG"
default_keyseq = "TCAG"

def assign_sequential_names(ignored, num_seqs, base_name='seq', start_at=0):
    """Returns list of num_seqs sequential, unique names.
    
    First argument is ignored; expect this to be set as a class attribute.
    """
    return ['%s_%s' % (base_name,i) for i in range(start_at,start_at+num_seqs)]

def flows_from_array(a):
    """FlowgramCollection from array of pos x seq: names are integers.
    
    This is an InputHandler for FlowgramCollection. It converts an arbitrary
    array of numbers into Flowgram objects and leaves the flowgrams unlabeled.
    """
    return list(transpose(a)), None, None

def flows_from_flowCollection(flow):
    """FlowgramCollection from array of pos x seq: names are integers.
    
    This is an InputHandler for FlowgramCollection. It converts an arbitrary
    array of numbers into Flowgram objects and leaves the flowgrams unlabeled.
    """
    return flow.flows, flow.Names, [f.header_info for f in flow.flows]

def flows_from_kv_pairs(flows):
    """SequenceCollection from list of (key, val) pairs.

        val can be str or flowgram object
    """
    
    names, flows =  map(list, zip(*flows))
    if isinstance(flows[0], str):
        info = [None]*len(flows)
    else:
        info = [f.header_info for f in flows]
    
    return flows, names, info

def flows_from_empty(obj, *args, **kwargs):
    """SequenceCollection from empty data: raise exception."""
    raise ValueError, "Cannot create empty SequenceCollection."

def flows_from_dict(flows):
    """SequenceCollection from dict of {label:flow_as_str} or {label:flow_obj}.
    """
    names, flows = map(list, zip(*flows.items()))
    if isinstance(flows[0], str):
        info = [None]*len(flows)
    else:
        info = [f.header_info for f in flows]
    return flows, names, info

def flows_from_sff(flows):
    """lines is sff file lines.
    """
    if isinstance(flows, str):
        flows = flows.splitlines()
        
    flows, head = parse_sff(flows)
    return flows_from_generic(flows)

def flows_from_generic(flows):
    """SequenceCollection from generic seq x pos data: seq of seqs of chars.
    
    This is an InputHandler for SequenceCollection. It converts a generic list
    (each item in the list will be mapped onto an object using
    seq_constructor and assigns sequential integers (0-based) as names.
    """
    names = []
    info = []
    for f in flows:
        if hasattr(f, 'Name'):
            names.append(f.Name)
        else:
            names.append(None)
        if hasattr(f, 'header_info'):
            info.append(f.header_info)
        else:
            info.append(None)
            
    return flows, names, info


class FlowgramCollection(object):
    """stores Flowgrams.
    
    - InputHandlers: support for different data types
    - flows: behaves like list of Flowgram objects
    - Names: behaves like list of names for the Flowgram objects
    - NamedFlows: behaves like dict of {name:flow}
    """    
    InputHandlers = {   'array': flows_from_array,
                        'dict': flows_from_dict,
                        'sff' :flows_from_sff,
                        'empty': flows_from_empty,
                        'kv_pairs':flows_from_kv_pairs,
                        'flowcoll':flows_from_flowCollection,
                        'generic':flows_from_generic
                    }
    HeaderInfo = ['Magic Number', 'Version', 'Index Offset', 'Index Length',
                  '# of Reads', 'Header Length', 'Key Length', '# of Flows',
                  'Flowgram Code', 'Flow Chars', 'Key Sequence']

    DefaultNameFunction = assign_sequential_names
    
    def __init__(self, data, Name = None, Names = None,
                 header_info = None, conversion_f = None,
                 name_conversion_f = None, remove_duplicate_names = False):
        """Initialize self with data and optionally Info.
        

        Parameters:
        
        data:           Data to convert into a FlowgramCollection
        
        Name:           Name of the FlowgramCollection.
                
        conversion_f:   Function to convert data into Flowgram.
                        
        name_conversion_f: if present, converts name into f(name).

        header_info: contains info to be printed in the common header of an sff
                    file, it is a dictionary. ex: Key Sequence:"ATCG"
        """

        #read all the data in if we were passed a generator
        if isinstance(data, GeneratorType):
            data = list(data)

        #set the Name
        self.Name = Name

        if header_info is not None:
            self._check_header_info(header_info)
        
            for i in header_info:
                setattr(self,i,header_info[i])
            if 'Key Sequence' in header_info:
                keyseq = header_info['Key Sequence']
            else:
                keyseq = None
            if 'Flow Chars' in header_info:
                floworder = header_info['Flow Chars']
            else:
                floworder = default_floworder
        else:
            keyseq = None
            floworder = default_floworder


        self.header_info = header_info

        per_flow_names, flows, name_order, info = \
                self._names_flows_order(conversion_f, data, Names, \
                name_conversion_f, remove_duplicate_names)
        self.Names = name_order

        #will take only the flows and names that are in name_order
        if per_flow_names != name_order:
            good_indices = []
            for n in name_order:
                good_indices.append(per_flow_names.index(n))
            flows = [flows[i] for i in good_indices]
            info = [info[i] for i in good_indices]
            per_flow_names = name_order

        self.flow_str= flows  
        self.flows = [Flowgram(f,n,keyseq,floworder, i)\
                      for f,n, i in zip(flows,self.Names,info)]

        #create NamedFlows dict for fast lookups
        self.NamedFlows = self._make_named_flows(self.Names, self.flows)
                

    def _strip_duplicates(self, names, flows, info):
        """Internal function to strip duplicates from list of names"""
        if len(set(names)) == len(names):
            return set(), names, flows, info
        #if we got here, there are duplicates
        unique_names = {}
        duplicates = {}
        fixed_names = []
        fixed_flows = []
        fixed_info = []
        for n, f, i in zip(names, flows,info):
            if n in unique_names:
                duplicates[n] = 1
            else:
                unique_names[n] = 1
                fixed_names.append(n)
                fixed_flows.append(f)
                fixed_info.append(i)
                
        return duplicates, fixed_names, fixed_flows, fixed_info


    def _names_flows_order(self, conversion_f, data, Names, \
            name_conversion_f, remove_duplicate_names):
        """Internal function to figure out names, flows, and name_order."""
        #figure out conversion function and whether it's an array
        if not conversion_f:
            input_type = self._guess_input_type(data)
            conversion_f = self.InputHandlers[input_type]
        #set seqs, names, and handler_info as properties
        
        flows, names, info = conversion_f(data)
        if names and name_conversion_f:
            names = map(name_conversion_f, names)

        #if no names were passed in as Names, if we obtained them from
        #the seqs we should use them, but otherwise we should use the
        #default names
        if Names is None:
            if (names is None) or (None in names):
                per_flow_names = name_order = \
                    self.DefaultNameFunction(len(flows))
            else:   #got names from seqs
                per_flow_names = name_order = names
        else:
            #otherwise, names were passed in as Names: use this as the order
            #if we got names from the sequences, but otherwise assign the
            #names to successive sequences in order
            if (names is None) or (None in names):
                per_flow_names = name_order = Names
            else: #got names from seqs, so assume name_order is in Names
                per_flow_names = names
                name_order = Names
        #check for duplicate names
        duplicates, fixed_names, fixed_flows, fixed_info = \
            self._strip_duplicates(per_flow_names, flows, info)
        if duplicates:
            if remove_duplicate_names:
                per_flow_names, flows, info =fixed_names,fixed_flows,fixed_info
                #if name_order doesn't have the same names as per_seq_names,
                #replace it with per_seq_names
                if (set(name_order) != set(per_flow_names)) or\
                    (len(name_order) != len(per_flow_names)):
                    name_order = per_flow_names
            else:
                raise ValueError, \
                "Some names were not unique. Duplicates are:\n" + \
                str(sorted(duplicates.keys()))
        return per_flow_names, flows, name_order, info


    def _check_header_info(self, info):
        for h in info:
            if h not in self.HeaderInfo:
                raise ValueError, "invalid key in header_info"
        
    def _make_named_flows(self, names, flows):
        """Returns NamedFlows: dict of name:flow."""
        name_flow_tuples = zip(names, flows)
        for n, f in name_flow_tuples:
            f.Name = n
        return dict(name_flow_tuples)
    
    def _guess_input_type(self, data):
        """Guesses input type of data; returns result as key of InputHandlers.
               
        Returns 'empty' if check fails, i.e. if it can't recognize the data
        as a specific type. Note that bad data is not guaranteed to
        return 'empty', and may be recognized as another type incorrectly.
        """
        
        if isinstance(data, dict):
            return 'dict'
        if isinstance(data, str):
            return 'sff'
        if isinstance(data, FlowgramCollection):
            return 'flowcoll'
        
        first = None
        try:
            first = data[0]
        except (IndexError, TypeError):
            pass
        try:
            first = iter(data).next()
        except (IndexError, TypeError, StopIteration):
            pass
        if first is None:
            return 'empty'
        try:
            if isinstance(first, Flowgram):     #model sequence base type
                return 'generic'
            if hasattr(first, 'dtype'):    #array object
                return 'array'
            elif isinstance(first, str) and first.startswith('Common'):
                return 'sff'
            else:
                try:
                    dict(data)
                    return 'kv_pairs'
                except (TypeError, ValueError):
                    pass
            return "generic"
        except (IndexError, TypeError), e:
            return 'empty'

    def __len__(self):
        """returns the number of flowgrams in the collection"""
        return len(self.flows)

    def __iter__(self):
        """iterates over the flows in the collection"""
        for f in self.flows:
            yield f

    def __str__(self):
        """returns string like sff file given the flowgrams and header_info"""
        lines = self.createCommonHeader()
        lines.append('')
        lines.extend([f.createFlowHeader() for f in self.flows])
        return '\n'.join(lines)

    def __cmp__(self, other):
        """cmp first tests as dict, then as str."""
        c = cmp(self.NamedFlows, other)
        if not c:
            return 0
        else:
            return cmp(str(self), str(other))

    def keys(self):
        """keys uses self.Names
        
        Note: returns copy, not original.
        """
        return self.Names[:]
    
    def values(self):
        """values returns values corresponding to self.Names."""
        return [self.NamedFlows[n] for n in self.Names]

    def items(self):
        """items returns (name, value) pairs."""
        return [(n, self.NamedFlows[n]) for n in self.Names]
    
    def writeToFile(self, filename=None, **kwargs):
        """Write the flowgrams to a file.
        
        Arguments:
        - filename: name of the output file
        
        """
        
        if filename is None:
            raise DataError('no filename specified')
        
        f = open(filename, 'w')
        f.write(str(self))
        f.close()

    def createCommonHeader(self):
        """header_info dict turned into flowgram common header"""
        lines = ["Common Header:"]
        
        if self.header_info is not None:
            lines.extend(["  %s:\t%s" % (param,self.header_info[param]) \
                      for param in self.header_info])

        return lines

    def toFasta(self, exclude_ambiguous = False,Bases=False,
                make_seqlabel = None):
        """Return flowgram collection in Fasta format
        
        Arguments:
            - make_seqlabel: callback function that takes the seq object and
              returns a label str

        if Bases is True then a fasta string will be made using
            self.Bases instead of translating the flowgram
              
        """
        if exclude_ambiguous:
            flows = flows.omitAmbiguousFlows()
        else:
            flows = self

        seqs = flows.toSequenceCollection(Bases)
            
        return seqs.toFasta(make_seqlabel = make_seqlabel)

    def toPhylip(self, exclude_ambiguous = False,Bases=False,generic_label=True,
                 make_seqlabel=None):
        """
        Return alignment in PHYLIP format and mapping to sequence ids
        
        raises exception if invalid alignment
        
        Arguments:
            - make_seqlabel: callback function that takes the seq object and
              returns a label str

        if Bases is True then a fasta string will be made using
            self.Bases instead of translating the flowgram 
        """
        if exclude_ambiguous:
            flows = flows.omitAmbiguousFlows()
        else:
            flows = self

        seqs = flows.toSequenceCollection( Bases)
        return seqs.toPhylip(make_seqlabel = make_seqlabel,
                             generic_label= generic_label)
    
    def toNexus(self,seq_type, exclude_ambiguous = False, Bases = False,
                interleave_len=50):
        """
        Return alignment in NEXUS format and mapping to sequence ids
        
        **NOTE** Not that every sequence in the alignment MUST come from
            a different species!! (You can concatenate multiple sequences from
            same species together before building tree)
        
        seq_type: dna, rna, or protein
        
        Raises exception if invalid alignment
        if Bases is True then a fasta string will be made using
        self.Bases instead of translating the flowgram
        """

        if exclude_ambiguous:
            flows = flows.omitAmbiguousFlows()
        else:
            flows = self

        seqs = flows.toSequenceCollection(Bases)
        return seqs.toNexus(seq_type,interleave_len = interleave_len)
    
    def toSequenceCollection(self, Bases = False):
        names = self.Names
        flow_dict = self.NamedFlows
        flows = [flow_dict[f].toSeq(Bases = Bases) for f in names]
        return SequenceCollection(flows)

    def addFlows(self, other):
        """Adds flowgrams from other to self. Returns a new object.
        
        other must be of same class as self or coerceable to that class..
        """
        assert not isinstance(other, str), "Must provide a series of flows "+\
                                            "or an flowgramCollection"
        self_flow_class = self.flows[0].__class__
        try:
            combined = self.flows + other.flows
        except AttributeError:
            combined = self.flows + list(other)
        try:
            combined_info =copy(self.header_info)
            combined_info.update(other.header_info)
        except AttributeError:
            combined_info =self.header_info
        
        
        for flow in combined:
            assert flow.__class__ == self_flow_class,\
                "classes different: Expected %s, Got %s" % \
                    (flow.__class__, self_flow_class)
        return self.__class__(data=combined,header_info=combined_info)

    def iterFlows(self, flow_order=None):
        """Iterates over values (sequences) in the alignment, in order.
        
        seq_order: list of keys giving the order in which seqs will be returned.
        Defaults to self.Names. Note that only these sequences will be
        returned, and that KeyError will be raised if there are sequences
        in order that have been deleted from the Alignment. If self.Names
        is None, returns the sequences in the same order as
        self.NamedSeqs.values().
        
        Use map(f, self.seqs()) to apply the constructor f to each seq. f must
        accept a single list as an argument.
        
        Always returns references to the same objects that are values of the
        alignment.
        """
        ns = self.NamedFlows
        get = ns.__getitem__
        for key  in flow_order or self.Names:
            yield get(key)


    def iterItems(self, flow_order=None, pos_order=None):
        """Iterates over elements in the flowgram collection.
        
        seq_order (names) can be used to select a subset of seqs.
        pos_order (positions) can be used to select a subset of positions.
        
        Always iterates along a seq first, then down a position (transposes
        normal order of a[i][j]; possibly, this should change)..
        
        WARNING: FlowgramCollection.iterItems() is not the same as fc.iteritems()
        (which is the built-in dict iteritems that iterates over key-value
        pairs).
        """
        if pos_order:
            for row in self.iterFlows(flow_order):
                for i in pos_order:
                    yield row[i]
        else:
            for row in self.iterFlows(flow_order):
                for i in row:
                    yield i

    Items = property(iterItems)
                    
    def copy(self):
        """Returns deep copy of self."""
        result = self.__class__(self)

    def _take_flows(self): return list(self.iterFlows())
    
    Flows = property(_take_flows)   #access as attribute if using default order.
    
    def takeFlows(self, flows, negate=False, **kwargs):
        """Returns new FlowgramCollection containing only specified flows.
        
        Note that the flows in the new collection will be references to the
        same objects as the seqs in the old collection.
        """
        get = self.NamedFlows.__getitem__
        result = {}
        if negate:
            #copy everything except the specified seqs
            negated_names = []
            row_lookup = dict.fromkeys(flows)
            for r, row in self.NamedFlows.items():
                if r not in row_lookup:
                    result[r] = row
                    negated_names.append(r)
            flows = negated_names #remember to invert the list of names
        
        else:
            #copy only the specified seqs
            for r in flows:
                result[r] = get(r)
        if result:
            return self.__class__(result, Names=flows, **kwargs)
        else:
            return {}   #safe value; can't construct empty collection

    def getFlowIndices(self, f, negate=False, Bases = False):
        """Returns list of keys of flows where f(Seq) is True.
        
        List will be in the same order as self.Names, if present.
        If Bases is True it will use the flowgram Bases attribute otherwise
        it uses the translation to a sequence
        """
        get = self.NamedFlows.__getitem__
        #negate function if necessary
        if negate:
            new_f = lambda x: not f(x)
        else:
            new_f = f
        #get all the seqs where the function is True
        return [key for key in self.Names \
               if new_f(get(key).toSeq(Bases = Bases))]
    
    def takeFlowsIf(self, f, negate=False, Bases = False, **kwargs):
        """Returns new Alignment containing seqs where f(row) is True.
        
        Note that the seqs in the new Alignment are the same objects as the
        seqs in the old Alignment, not copies.

        If Bases is True it will use the flowgram Bases attribute otherwise
        it uses the translation to a sequence
        """
        #pass negate to get SeqIndices
        return self.takeFlows(self.getFlowIndices(f, negate, Bases= Bases),
                              **kwargs)

    def getFlow(self, flowname):
        """Return a flowgram object for the specified seqname.
        """
        return self.NamedFlows[flowname]

    def getFlowNames(self):
        """Return a list of Flowgram names."""
        return self.Names[:]

    def getIntMap(self,prefix='seq_'):
        """Returns a dict with names mapped to enumerates integer names.
            
            - prefix: prefix for sequence label. Default = 'seq_'
            - int_keys is a dict mapping int names to sorted original names.
        """
        get = self.NamedFlows.__getitem__
        int_keys = dict([(prefix+str(i),k) for i,k in \
                enumerate(sorted(self.NamedFlows.keys()))])
        int_map = dict([(k, copy(get(v))) for k,v in int_keys.items()])
        return int_map, int_keys

    def toDict(self):
        """Returns the collection as dict of names -> strings.

        Note: returns strings, NOT Flowgram objects.
        """
        collection_dict = {}
        
        for flow_name in self.Names:
            collection_dict[flow_name] = self.NamedFlows[flow_name]._flowgram
        
        return collection_dict
    
    def omitAmbiguousFlows(self, Bases = False):
        """Returns an object containing only the sequences without N's"""
        is_ambiguous = lambda x: 'N' not in x
        
        return self.takeFlowsIf(is_ambiguous, Bases = Bases)

    def setBases(self):
        """Sets the Bases property for each flowgram using toSeq"""
        for f in self.values():
            f.Bases = f.toSeq()
    

def pick_from_prob_density(pvals,bin_size):
    l = multinomial(1,pvals)
    return (l.nonzero()[0][0]) * bin_size

def seqs_to_flows(seqs, keyseq = default_keyseq, floworder = default_floworder,
                  numflows = None, probs = None, bin_size = 0.01,
                  header_info = {}):
    """ Transfrom a sequence into an ideal flow
        seqs: a list of name sequence object tuples (name,tuple)
        keyseq: the flowgram key Sequence
        floworder: The chars needed to convert seq to flow
        numflows: number of total flows in each flowgram, if it is specified
            the flowgram will be padded to that number
        probs: dictionary defining the probability distribution for each
            homopolymer

        WARNING:each distributions probabilities must add to 1.0
    """
    flows = []
    homopolymer_counter = 1.0
    if probs:
        for p in probs:
            if round(sum(probs[p]),1) != 1.0:
                raise ValueError, 'probs[%s] does not add to 1.0' % p

    for name,seq in seqs:
        flow_seq = FakeRandom(floworder,True)
        flow = []
        seq_len = len(seq)
        for i, nuc in enumerate(seq):
            if i < seq_len-1 and seq[i+1] == nuc:
                homopolymer_counter += 1.0
            else:
                while flow_seq() != nuc:
                    if probs is None:
                        val = 0.0
                    else:
                        val = pick_from_prob_density(probs[0],bin_size)
                    flow.append(val)
                if (probs is None) or (homopolymer_counter > 9):
                    val = homopolymer_counter
                else:
                    val = pick_from_prob_density(probs[int(homopolymer_counter)],bin_size)
                 
                flow.append(val)
                homopolymer_counter = 1.0

        len_flow = len(flow)
        len_order = len(floworder)

        if numflows is not None and numflows % len_order != 0:
            raise ValueError, "numflows must be divisable by the length of floworder"
        if (len_flow % len_order != 0):
            right_missing = len_order - (len_flow % len_order)
            if numflows != (len_flow + right_missing) and numflows is not None:
                right_missing += (numflows - (len_flow+right_missing))
            if probs is None:
                flow.extend([0.0]*right_missing)
            else:
                for i in range(0,right_missing):
                    flow.append(pick_from_prob_density(probs[0],bin_size))

        flows.append((name, Flowgram(flow, id, keyseq, floworder)))
    if keyseq is not None:
        keylen = len(keyseq)
    else:
        keylen = None
    header_info.update({'Key Sequence':keyseq,'Flow Chars':floworder,
                        'Key Length':keylen})
    return FlowgramCollection(flows, header_info = header_info)