# Copyright (c) 2013 by California Institute of Technology
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 
# 3. Neither the name of the California Institute of Technology nor
#    the names of its contributors may be used to endorse or promote
#    products derived from this software without specific prior
#    written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL CALTECH
# OR THE CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
"""
Mathematical Sets and Power Sets
"""
import logging
logger = logging.getLogger(__name__)

import warnings
from itertools import chain, combinations
from collections import Iterable, Hashable, Container
from pprint import pformat
from random import randint


def compare_lists(list1, list2):
    """Compare list contents, ignoring ordering.
    
    Hashability of elements not assumed, incurring O(N**2)
    
    See Also
    ========
    L{MathSet}
    
    @type list1: list
    @type list2: list
    
    @return: True if a bijection exists between the lists.
        Note that this takes into account multiplicity of elements.
    @rtype: bool
    """
    if not isinstance(list1, list):
        raise TypeError('Not a list, instead list1:\n\t' +str(list1) )
    
    if not isinstance(list2, list):
        raise TypeError('Not a list, instead list2:\n\t' +str(list2) )
    
    dummy_list = list(list1)
    same_lists = True
    for item in list2:
        try:
            dummy_list.remove(item)
        except:
            # unique element not found to be rm'd
            same_lists = False
            break
        
    # anything remaining ?
    same_lists = same_lists and not bool(dummy_list)
    
    return same_lists

class MathSet(object):
    """Mathematical set, allows unhashable elements.
    
    Examples
    ========
    >>> s = MathSet(['a', 1, [1,2], {'a', 'b'} ] )
    
    Then print(s) shows how the elements were separately stored
    in a set and list, to optimize contains operations:
    
    >>> print(s)
    MathSet(['a', 1, [1, 2], set(['a', 'b'])])
    
    Set operations similar to the builtin type are supported:
    
    >>> p = MathSet()
    >>> p.add(1)
    >>> p |= [1, 2]
    >>> p |= {3, 4}
    >>> p |= [[1, 2], '5', {'a':1} ]
    >>> p.add_from([5, 6, '7', {8, 9} ] )
    >>> p.remove(1)
    >>> p
    MathSet([2, 3, 4, 5, 6, '5', '7', [1, 2], {'a': 1}, set([8, 9])])
    
    See Also
    ========
    L{SubSet}, L{PowerSet}, set
    """    
    def __init__(self, iterable=[]):
        """Initialize by adding elements from iterable.
        
        Example
        =======
        >>> s = MathSet([1, 2, 'a', {3, 4} ] )
        
        @param iterable: iterable from which to initialize the set S
            which underlies the PowerSet 2^S
        @type iterable: iterable, any element types allowed
        """
        self._delete_all()
        self.add_from(iterable)
    
    def __repr__(self):
        return 'MathSet(' +pformat(list(self._set) +self._list) +')'
    
    def _debug_repr(self):
        set_str = ', '.join([repr(i) for i in self._set] )
        return 'MathSet({' +set_str +'} +' +str(self._list) +')'
    
    def __or__(self, other):
        """Union with another mathematical set.
        
        See Also
        ========
        L{__ior__}
        
        @param other: any other mathematical set.
        @type other: iterable, elements not restricted to hashable
        
        @return: self | iterable
        @rtype: MathSet
        """
        s = MathSet(self)
        s.add_from(other)
        return s
    
    def cartesian(self, other):
        """Return Cartesian product with other C{MathSet},
        
        as CartesianProduct instance.
        This representation is internally more efficient,
        because it is implicit (product not explicitly stored).
        
        See Also
        ========
        L{CartesianProduct}, L{__mul__}

        UNDER DEVELOPMENT; function signature may change without
        notice.  Calling will result in NotImplementedError.
        """
        raise NotImplementedError
    
    def __mul__(self, other):
        """Return the Cartesian product with another C{MathSet}.
        
        Example
        =======
        >>> a = MathSet([1, 2] )
        >>> b = MathSet([3, 4] )
        >>> c = a *b
        >>> print(type(c) )
        
        >>> print(c)
        
        
        If we prefer a CartesianProduct returned instead:
        >>> c = a.cartesian(b)
        
        See Also
        ========
        L{cartesian}
        
        @param other: set with which to take Cartesian product
        @type other: MathSet
        
        @return: Cartesian product of C{self} with C{other}.
        @rtype: C{MathSet} (explicit construction)
        """
        cartesian = [(x, y) for x in self
                            for y in other]
        return MathSet(cartesian)
    
    def __ior__(self, iterable):
        """Union with of MathSet with iterable.
        
        Example
        =======
        >>> s = MathSet([1, 2] )
        >>> s |= [3, 4] # much cleaner & familiar
        >>> print(s)
        set([1, 2, 3, 4]) U []
        
        See Also
        ========
        L{__or__}
        
        @param iterable: any mathematical set.
        @type iterable: iterable, elements not restricted to hashable
        
        @return: self | iterable
        @rtype: MathSet
        """
        self.add_from(iterable)
        return self
    
    def __sub__(self, rm_items):
        s = MathSet(self)
        print('s = ' +str(s) )
        print(rm_items)
        for item in rm_items:
            if item in s:
                print('Removing...: ' +str(item) )
                s.remove(item)
        return s
    
    def __isub__(self, rm_items):
        """Delete multiple elements."""
        for item in rm_items:
            if item in self:
                self.remove(item)
        
        return self
    
    def __eq__(self, other):
        if not isinstance(other, MathSet):
            raise TypeError('For now comparison only to another MathSet.\n' +
                            'Got:\n\t' +str(other) +'\n of type: ' +
                            str(type(other) ) +', instead.')
        
        same_lists = compare_lists(self._list, other._list)
        
        return (self._set == other._set) and same_lists
    
    def __contains__(self, item):
        if isinstance(item, Hashable):
            try:
                return item in self._set
            except:
                logger.error('UnHashable items within Hashable.')
        
        return item in self._list
    
    def __iter__(self):
        return iter(self._list + list(self._set))
    
    def __len__(self):
        """Number of elements in set."""
        return len(self._set) +len(self._list)
    
    def _filter_hashables(self, iterable):
        return filter(lambda x: isinstance(x, Hashable), iterable)
    
    def _filter_unhashables(self, iterable):
        return list(filter(lambda x: not isinstance(x, Hashable), iterable))
    
    def _delete_all(self):
        self._set = set()
        self._list = list()
    
    def add(self, item):
        """Add element to mathematical set.
        
        Example
        =======
        >>> s = MathSet()
        >>> s.add(1)
        set([1]) U []
        
        See Also
        ========
        L{add_from}, L{__ior__}, L{remove}
        
        @param item: the new set element
        @type item: anything, if hashable it is stored in a Python set,
            otherwise stored in a list.
        """
        if isinstance(item, Hashable):
            try:
                self._set.add(item)
                return
            except TypeError:
                logger.error('UnHashable items within Hashable.')
        
        if item not in self._list:
            self._list.append(item)
        else:
            logger.warn('item already in MathSet.')
    
    def add_from(self, iterable):
        """Add multiple elements to mathematical set.
        
        Equivalent to |=
        
        Example
        =======
        >>> s = MathSet()
        >>> s.add_from([1, 2, {3} ] )
        
        is equivalent to:
        
        >>> s = MathSet()
        >>> s |= [1, 2, {3} ]
        
        See Also
        ========
        L{add}, L{__ior__}, L{remove}
        
        @param iterable: new MathSet elements
        @type iterable: iterable containing (possibly not hashable) elements
        """
        if not isinstance(iterable, Iterable):
            raise TypeError('Can only add elements to MathSet from Iterable.\n' +
                            'Got:\n\t' +str(iterable) +'\n instead.')
        
        if isinstance(iterable, MathSet):
            self._set |= set(iterable._set)
            self._list = list(unique(self._list +iterable._list) )
            return
        
        # speed up
        if isinstance(iterable, set):
            self._set |= iterable
            return
        
        # filter to optimize storage
        try:
            self._set |= set(self._filter_hashables(iterable) )
            self._list = list(unique(
                self._list + self._filter_unhashables(iterable)
            ))
            return
        except:
            # ...if contents of elements in iterable are mutable
            self._list = list(unique(self._list +list(iterable) ) )
    
    def remove(self, item):
        """Remove existing element from mathematical set.
        
        Example
        =======
        >>> p = MathSet([1, 2] )
        >>> p.remove(1)
        >>> p
        set([2]) U []
        
        See Also
        ========
        L{add}, L{add_from}, L{__or__}
        
        @param item: An item already in the set.
            For adding items, see add.
        """
        if item not in self:
            warnings.warn(
                'Set element not in set S.\n'+
                'Maybe you targeted another element for removal ?'
            )
        
        if isinstance(item, Hashable):
            try:
                self._set.remove(item)
                return
            except:
                logger.debug('item: ' +str(item) +', contains unhashables.')
            
        self._list.remove(item)
    
    def pop(self):
        """Remove and return random MathSet element.
        
        Raises KeyError if MathSet is empty.
        """
        if not self:
            raise KeyError('Nothing to pop: MathSet is empty.')
        
        if self._set and self._list:
            if randint(0, 1):
                return self._set.pop()
            else:
                return self._list.pop()
        elif self._set and not self._list:
            return self._set.pop()
        elif self._list and not self._set:
            return self._list.pop()
        else:
            raise Exception('Bug in empty MathSet: not self above' +
                            'should not reaching this point.')
    
    def intersection(self, iterable):
        """Return intersection with iterable.
        
        @param iterable: find common elements with C{self}
        @type iterable: C{Iterable}
        
        @return: intersection of C{self} with C{iterable}
        @rtype: C{MathSet}
        """
        s = MathSet()
        for item in iterable:
            print(item)
            if item in self:
                print('Adding...\n')
                s.add(item)
        return s
    
    def intersects(self, iterable):
        """Check intersection with iterable.
        
        Checks the existence of common elements with iterable.
        
        >>> s = MathSet()
        >>> s.add(1)
        >>> r = [1,2]
        >>> s.intersects(r)
        True
        
        @param iterable: with which to check intersection
        @type iterable: C{Iterable}
        
        @return: C{True} if C{self} has common element with C{iterable}.
            Otherwise C{False}.
        @rtype: C{bool}
        """
        for item in iterable:
            if item in self:
                return True
        return False

class SubSet(MathSet):
    """Subset of selected MathSet, or other Iterable.
    
    Prior to adding new elements,
    it checks that they are in its superset.
    
    Example
    =======
    >>> superset = [1, 2]
    >>> s = SubSet(superset)
    >>> s |= [1, 2]
    >>> print(s)
    SubSet([[1, 2]])
    >>> s.add(3)
    raises exception because 3 \\notin [1,2]
    
    See Also
    ========
    L{MathSet}, L{PowerSet}
    """
    def __init__(self, superset, iterable=None):
        """Define the superset of this set.
        
        @param superset: This SubSet checked vs C{superset}
        @type superset: Iterable
        
        @param iterable: elements to add to subset
        @type iterable: Iterable
        """
        self._superset = superset
        
        super(SubSet, self).__init__([])
        
        if not isinstance(superset, Container):
            raise TypeError('superset must be Iterable,\n' +
                            'Got instead:\n\t' +str(superset) )
    
    def __repr__(self):
        return 'SubSet(' +pformat(list(self._set) +self._list) +')'
    
    def _debug_repr(self):
        set_str = ', '.join([repr(i) for i in self._set] )
        return 'SubSet({' +set_str +'} +' +str(self._list) +')'
    
    @property
    def superset(self):
        return self._superset
    
    def add(self, new_element):
        """Add state to subset.
        
        Extends MathSet.add with subset relation checking.
        
        Example
        =======
        C{new_initial_state} should already be a state.
        First use states.add to include it in set of states,
        then states.add_initial.
        
        See Also
        ========
        L{MathSet.add}
        """
        if not new_element in self._superset:
            raise Exception(
                'New element state \\notin superset.\n'
                'Add it first to states using e.g. sys.states.add()\n'
                'FYI: new element:\n\t' +str(new_element) +'\n'
                'and superset:\n\t' +str(self._superset)
            )
        
        super(SubSet, self).add(new_element)
    
    def add_from(self, new_elements):
        """Add multiple new elements to subset.
        
        Extends MathSet.add_from with subset relation checking.
        
        Note
        ====
        It would be sufficient to extend only .add provided
        MathSet.add_from called .add iteratively.
        However MathSet.add_from filters states, which is
        arguably more efficient. So both .add and .add_from
        need to be extended here.
        
        See Also
        ========
        L{add}, L{__ior__}
        """
        if not is_subset(new_elements, self._superset):
            raise Exception('All new_elements:\n\t' +str(new_elements) +
                            '\nshould already be \\in ' +
                            'self.superset = ' +str(self._superset) )
        super(SubSet, self).add_from(new_elements)

class CartesianProduct(object):
    """List of MathSets, with Cartesian semantics.
    """
    def __init__(self):
        self.mathsets = []
    
    def __contains__(self, element):
        #TODO check ordered
        if not isinstance(element, Iterable):
            raise TypeError(
                'Argument element must be Iterable, otherwise cannot ' +
                'recover which item in it belongs to which set in the ' +
                'Cartesian product.'
            )
        
        for idx, item in enumerate(element):
            if item not in self.mathsets[idx]:
                return False
        return True
    
    def __mul__(self, mathsets):
        """Multiply Cartesian products."""
        if not isinstance(mathsets, list):
            raise TypeError('mathsets given must be a list of MathSet.')
    
    def add(self, mathset):
        self.mathsets += [mathset]
    
    def add_from(self, mathsets):
        self.mathsets += mathsets
    
    def remove(self, mathset):
        self.mathsets.remove(mathset)
    
    def remove_from(self, mathsets):
        for mathset in mathsets:
            self.remove(mathset)

def unique(iterable):
    """Return unique elements.
    
    Note
    ====
    Always returning a list for consistency was tempting,
    however this defeats the purpose of creating this function
    to achieve brevity elsewhere in the code.
    
    @return: iterable with duplicates removed, as C{set} if possible.
    @rtype:
        - If all items in C{iterable} are hashable,
            then returns C{set}.
        - If iterable contains unhashable item,
            then returns C{list} of unique elements.
    """
    # hashable items ?
    try:
        unique_items = set(iterable)
    except:
        unique_items = []
        for item in iterable:
            if item not in unique_items:
                unique_items.append(item)
    
    return unique_items

def contains_multiple(iterable):
    """Does iterable contain any item multiple times ?"""    
    return len(iterable) != len(unique(iterable) )

def is_subset(small_iterable, big_iterable):
    """Comparison for handling list <= set, and lists with unhashable items.
    """
    # asserts removed when compiling with optimization on...
    # it would have been elegant to use instead:
    #   assert(isinstance(big_iterable, Iterable))
    # since the error msg is succintly stated by the assert itself
    if not isinstance(big_iterable, (Iterable, Container) ):
        raise TypeError('big_iterable must be either Iterable or Container, '
                        'otherwise subset relation undefined.\n'
                        'Got:\n\t' +str(big_iterable) +'\ninstead.')
    
    if not isinstance(small_iterable, Iterable):
        raise TypeError('small_iterable must be Iterable, '
                        'otherwise subset relation undefined.\n'
                        'Got:\n\t' +str(small_iterable) +'\ninstead.')
    
    # nxor
    if isinstance(small_iterable, str) != isinstance(big_iterable, str):
        raise TypeError('Either both or none of small_iterable, '
                        'big_iterable should be strings.\n'
                        'Otherwise subset relation between string '
                        'and non-string may introduce bugs.\nGot:\n\t' +
                        str(small_iterable) +',\t' +str(big_iterable) +
                        '\ninstead.')
    
    try:
        # first, avoid object duplication
        if not isinstance(small_iterable, set):
            small_iterable = set(small_iterable)
        
        if not isinstance(big_iterable, set):
            big_iterable = set(big_iterable)
        
        return small_iterable <= big_iterable
    except TypeError:
        # not all items hashable...
        
        try:
            # list to avoid: unhashable \in set ? => error
            if not isinstance(big_iterable, list):
                # avoid object duplication
                big_iterable = list(big_iterable)
        except:
            logger.error('Could not convert big_iterable to list.')
        
        for item in small_iterable:
            if item not in big_iterable:
                return False
        return True
    except:
        raise Exception('Failed to compare iterables.')

def powerset(iterable):
        """powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)
        
        From http://docs.python.org/2/library/itertools.html,
        also in https://pypi.python.org/pypi/more-itertools
        """
        s = list(iterable)
        return chain.from_iterable(combinations(s, r) for r in range(len(s)+1) )

class PowerSet(object):
    """Efficiently store power set of a mathematical set.
    
    Set here isn't necessarily a Python set,
    i.e., it may comprise of unhashable elements.
      
    Example
    =======
    Specify the mathematical set S underlying the PowerSet.

    >>> S = [[1, 2], '3', {'a':1}, 1]
    >>> p = PowerSet(S)
    
    >>> q = PowerSet()
    >>> q.math_set = S
    
    Add new element to underlying set S.

    >>> p.math_set.add({3: 'a'} )
    
    Add multiple new elements to underlying set S.

    >>> p.math_set.add_from({3, 'a'} )
    >>> p.math_set |= [1,2]
    
    Remove existing element from set S.

    >>> p.remove(1)
    
    See Also
    ========
    L{MathSet}, L{SubSet}, L{is_subset}
    
    @param iterable: mathematical set S of elements, on which this 2^S defined.
    @type iterable: iterable container
    """    
    def __init__(self, iterable=None):
        """Create new PowerSet over elements contained in S = C{iterable}.
        
        This powerset is 2^S.
        
        @param iterable: contains elements of set S underlying the PowerSet.
        @type iterable: iterable of elements which can be hashable or not.
        """
        if iterable is None:
            iterable = []
        self.math_set = MathSet(iterable)
    
    def __get__(self, instance, value):
        return self()
    
    def __repr__(self):
        return 'PowerSet(' +str(self.math_set) +' )'
    
    def __contains__(self, item):
        """Is item \\in 2^iterable = this powerset(iterable)."""
        if not isinstance(item, Iterable):
            raise Exception('Not iterable:\n\t' +str(item) +',\n'
                            'this is a powerset, so it contains (math) sets.')
        
        return is_subset(item, self.math_set)
    
    def __iter__(self):
        return powerset(self.math_set)
    
    def __len__(self):
        return 2**len(self.math_set)
    
    def __add__(self, other):
        if not isinstance(other, PowerSet):
            raise TypeError('Addition defined only between PowerSets.\n' +
                            'Got instead:\n\t other = ' +str(other) )
        
        list1 = self.math_set
        list2 = other.math_set
        union = list1 | list2
        return PowerSet(union)
    
    def __eq__(self, other):
        if not isinstance(other, PowerSet):
            raise TypeError('Can only compare to another PowerSet.')
        
        return other.math_set == self.math_set
    
    def __setattr__(self, name, value):
        if name is 'math_set' and not isinstance(value, MathSet):
            msg = 'PowerSet.math_set must be of class MathSet.\n'
            msg += 'Got instead:\n\t' +str(value)
            msg += '\nof class:\nt\t' +str(type(value) )
            raise Exception(msg)
        
        object.__setattr__(self, name, value)

class TypedDict(dict):
    """dict subclass where values can be constrained by key.
    
    For each key, a domain can optionally be defined,
    which restricts the admissible values that can be
    paired with that key.
    """
    # credits for debugging this go here:
    #   http://stackoverflow.com/questions/2060972/
    def __init__(self, *args, **kwargs):
        self.update(*args, **kwargs)
    
    def __setitem__(self, i, y):
        """Raise ValueError if value y not allowed for key i.
        """
        valid_y = True
        try:
            if i in self.allowed_values:
                if y not in self.allowed_values[i]:
                    valid_y = False
        except:
            pass
        
        if not valid_y:
            msg = 'key: ' + str(i) + ', cannot be'
            msg += ' assigned value: ' + str(y) + '\n'
            msg += 'Admissible values are:\n\t'
            msg += str(self.allowed_values[i])
            raise ValueError(msg)
        
        super(TypedDict, self).__setitem__(i, y)
    
    def __str__(self):
        return 'TypedDict(' + dict.__str__(self) + ')'
    
    def update(self, *args, **kwargs):
        if args:
            if len(args) > 1:
                raise TypeError("update expected at most 1 arguments, "
                                "got %d" % len(args))
            other = dict(args[0])
            for key in other:
                self[key] = other[key]
        for key in kwargs:
            self[key] = kwargs[key]

    def setdefault(self, key, value=None):
        if key not in self:
            self[key] = value
        return self[key]
    
    def set_types(self, allowed_values):
        """Restrict values the key can be paired with.
        
        @param allowed_values: dict of the form:
            
                {key : values}
            
                C{values} must implement C{__contains__}
                to enable checking validity of values.
        """
        self.allowed_values = allowed_values
