"""
Sets

AUTHORS:

- William Stein (2005) - first version

- William Stein (2006-02-16) - large number of documentation and
  examples; improved code

- Mike Hansen (2007-3-25) - added differences and symmetric
  differences; fixed operators

- Florent Hivert (2010-06-17) - Adapted to categories

- Nicolas M. Thiery (2011-03-15) - Added subset and superset methods

- Julian Rueth (2013-04-09) - Collected common code in
  :class:`Set_object_binary`, fixed ``__hash__``.

"""

#*****************************************************************************
#       Copyright (C) 2005 William Stein <wstein@gmail.com>
#                     2013 Julian Rueth <julian.rueth@fsfe.org>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code 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.
#
#  The full text of the GPL is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************
from __future__ import print_function
import six
from six import integer_types

from sage.misc.latex import latex
from sage.misc.prandom import choice

from sage.structure.category_object import CategoryObject
from sage.structure.element import Element
from sage.structure.parent import Parent, Set_generic
from sage.structure.richcmp import richcmp_method, richcmp, rich_to_bool

from sage.categories.sets_cat import Sets
from sage.categories.enumerated_sets import EnumeratedSets

import sage.rings.infinity

def has_finite_length(obj):
    """
    Return ``True`` if ``obj`` is known to have finite length.

    This is mainly meant for pure Python types, so we do not call any
    Sage-specific methods.

    EXAMPLES::

        sage: from sage.sets.set import has_finite_length
        sage: has_finite_length(tuple(range(10)))
        True
        sage: has_finite_length(list(range(10)))
        True
        sage: has_finite_length(set(range(10)))
        True
        sage: has_finite_length(iter(range(10)))
        False
        sage: has_finite_length(GF(17^127))
        True
        sage: has_finite_length(ZZ)
        False
    """
    try:
        len(obj)
    except OverflowError:
        return True
    except Exception:
        return False
    else:
        return True

def Set(X=[]):
    r"""
    Create the underlying set of ``X``.

    If ``X`` is a list, tuple, Python set, or ``X.is_finite()`` is
    ``True``, this returns a wrapper around Python's enumerated immutable
    ``frozenset`` type with extra functionality.  Otherwise it returns a
    more formal wrapper.

    If you need the functionality of mutable sets, use Python's
    builtin set type.

    EXAMPLES::

        sage: X = Set(GF(9,'a'))
        sage: X
        {0, 1, 2, a, a + 1, a + 2, 2*a, 2*a + 1, 2*a + 2}
        sage: type(X)
        <class 'sage.sets.set.Set_object_enumerated_with_category'>
        sage: Y = X.union(Set(QQ))
        sage: Y
        Set-theoretic union of {0, 1, 2, a, a + 1, a + 2, 2*a, 2*a + 1, 2*a + 2} and Set of elements of Rational Field
        sage: type(Y)
        <class 'sage.sets.set.Set_object_union_with_category'>

    Usually sets can be used as dictionary keys.

    ::

        sage: d={Set([2*I,1+I]):10}
        sage: d                  # key is randomly ordered
        {{I + 1, 2*I}: 10}
        sage: d[Set([1+I,2*I])]
        10
        sage: d[Set((1+I,2*I))]
        10

    The original object is often forgotten.

    ::

        sage: v = [1,2,3]
        sage: X = Set(v)
        sage: X
        {1, 2, 3}
        sage: v.append(5)
        sage: X
        {1, 2, 3}
        sage: 5 in X
        False

    Set also accepts iterators, but be careful to only give *finite*
    sets::

        sage: from six.moves import range
        sage: sorted(Set(range(1,6)))
        [1, 2, 3, 4, 5]
        sage: sorted(Set(list(range(1,6))))
        [1, 2, 3, 4, 5]
        sage: sorted(Set(iter(range(1,6))))
        [1, 2, 3, 4, 5]

    We can also create sets from different types::

        sage: sorted(Set([Sequence([3,1], immutable=True), 5, QQ, Partition([3,1,1])]), key=str)
        [5, Rational Field, [3, 1, 1], [3, 1]]

    Sets with unhashable objects work, but with less functionality::

        sage: A = Set([QQ, (3, 1), 5])  # hashable
        sage: sorted(A.list(), key=repr)
        [(3, 1), 5, Rational Field]
        sage: type(A)
        <class 'sage.sets.set.Set_object_enumerated_with_category'>
        sage: B = Set([QQ, [3, 1], 5])  # unhashable
        sage: sorted(B.list(), key=repr)
        Traceback (most recent call last):
        ...
        AttributeError: 'Set_object_with_category' object has no attribute 'list'
        sage: type(B)
        <class 'sage.sets.set.Set_object_with_category'>

    TESTS::

        sage: Set(Primes())
        Set of all prime numbers: 2, 3, 5, 7, ...
        sage: Set(Subsets([1,2,3])).cardinality()
        8
        sage: S = Set(iter([1,2,3])); S
        {1, 2, 3}
        sage: type(S)
        <class 'sage.sets.set.Set_object_enumerated_with_category'>
        sage: S = Set([])
        sage: TestSuite(S).run()

    Check that :trac:`16090` is fixed::

        sage: Set()
        {}
    """
    if isinstance(X, CategoryObject):
        if isinstance(X, Set_generic):
            return X
        elif X in Sets().Finite():
            return Set_object_enumerated(X)
        else:
            return Set_object(X)

    if isinstance(X, Element):
        raise TypeError("Element has no defined underlying set")

    try:
        X = frozenset(X)
    except TypeError:
        return Set_object(X)
    else:
        return Set_object_enumerated(X)

@richcmp_method
class Set_object(Set_generic):
    r"""
    A set attached to an almost arbitrary object.

    EXAMPLES::

        sage: K = GF(19)
        sage: Set(K)
        {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}
        sage: S = Set(K)

        sage: latex(S)
        \left\{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18\right\}
        sage: TestSuite(S).run()

        sage: latex(Set(ZZ))
        \Bold{Z}

    TESTS:

    See trac ticket :trac:`14486`::

        sage: 0 == Set([1]), Set([1]) == 0
        (False, False)
        sage: 1 == Set([0]), Set([0]) == 1
        (False, False)
    """
    def __init__(self, X, category=None):
        """
        Create a Set_object

        This function is called by the Set function; users
        shouldn't call this directly.

        EXAMPLES::

            sage: type(Set(QQ))
            <class 'sage.sets.set.Set_object_with_category'>
            sage: Set(QQ).category()
            Category of sets

        TESTS::

            sage: _a, _b = get_coercion_model().canonical_coercion(Set([0]), 0)
            Traceback (most recent call last):
            ...
            TypeError: no common canonical parent for objects with parents:
            '<class 'sage.sets.set.Set_object_enumerated_with_category'>'
            and 'Integer Ring'
        """
        from sage.rings.integer import is_Integer
        if isinstance(X, integer_types) or is_Integer(X):
            # The coercion model will try to call Set_object(0)
            raise ValueError('underlying object cannot be an integer')

        if category is None:
            category = Sets()
        Parent.__init__(self, category=category)
        self.__object = X

    def __hash__(self):
        """
        Return the hash value of ``self``.

        EXAMPLES::

            sage: hash(Set(QQ)) == hash(QQ)
            True
        """
        return hash(self.__object)

    def _latex_(self):
        r"""
        Return latex representation of this set.

        This is often the same as the latex representation of this
        object when the object is infinite.

        EXAMPLES::

            sage: latex(Set(QQ))
            \Bold{Q}

        When the object is finite or a special set then the latex
        representation can be more interesting.

        ::

            sage: print(latex(Primes()))
            \text{\texttt{Set{ }of{ }all{ }prime{ }numbers:{ }2,{ }3,{ }5,{ }7,{ }...}}
            sage: print(latex(Set([1,1,1,5,6])))
            \left\{1, 5, 6\right\}
        """
        return latex(self.__object)

    def _repr_(self):
        """
        Print representation of this set.

        EXAMPLES::

            sage: X = Set(ZZ)
            sage: X
            Set of elements of Integer Ring
            sage: X.rename('{ integers }')
            sage: X
            { integers }
        """
        return "Set of elements of " + repr(self.__object)

    def __iter__(self):
        """
        Iterate over the elements of this set.

        EXAMPLES::

            sage: X = Set(ZZ)
            sage: I = X.__iter__()
            sage: next(I)
            0
            sage: next(I)
            1
            sage: next(I)
            -1
            sage: next(I)
            2
        """
        return iter(self.__object)

    an_element = EnumeratedSets.ParentMethods.__dict__['_an_element_from_iterator']

    def __contains__(self, x):
        """
        Return ``True`` if `x` is in ``self``.

        EXAMPLES::

            sage: X = Set(ZZ)
            sage: 5 in X
            True
            sage: GF(7)(3) in X
            True
            sage: 2/1 in X
            True
            sage: 2/1 in ZZ
            True
            sage: 2/3 in X
            False

        Finite fields better illustrate the difference between
        ``__contains__`` for objects and their underlying sets.

            sage: X = Set(GF(7))
            sage: X
            {0, 1, 2, 3, 4, 5, 6}
            sage: 5/3 in X
            False
            sage: 5/3 in GF(7)
            False
            sage: sorted(Set(GF(7)).union(Set(GF(5))), key=int)
            [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 6]
            sage: Set(GF(7)).intersection(Set(GF(5)))
            {}
        """
        return x in self.__object

    def __richcmp__(self, right, op):
        r"""
        Compare ``self`` and ``right``.

        If ``right`` is not a :class:`Set_object`, return ``NotImplemented``.
        If ``right`` is also a :class:`Set_object`, returns comparison
        on the underlying objects.

        .. NOTE::

           If `X < Y` is true this does *not* necessarily mean
           that `X` is a subset of `Y`.  Also, any two sets can be
           compared still, but the result need not be meaningful
           if they are not equal.

        EXAMPLES::

            sage: Set(ZZ) == Set(QQ)
            False
            sage: Set(ZZ) < Set(QQ)
            True
            sage: Primes() == Set(QQ)
            False

        The following is random, illustrating that comparison of
        sets is not the subset relation, when they are not equal::

            sage: Primes() < Set(QQ)             # random  # py2
            True or False
        """
        if not isinstance(right, Set_object):
            return NotImplemented
        return richcmp(self.__object, right.__object, op)

    def union(self, X):
        """
        Return the union of ``self`` and ``X``.

        EXAMPLES::

            sage: Set(QQ).union(Set(ZZ))
            Set-theoretic union of Set of elements of Rational Field and Set of elements of Integer Ring
            sage: Set(QQ) + Set(ZZ)
            Set-theoretic union of Set of elements of Rational Field and Set of elements of Integer Ring
            sage: X = Set(QQ).union(Set(GF(3))); X
            Set-theoretic union of Set of elements of Rational Field and {0, 1, 2}
            sage: 2/3 in X
            True
            sage: GF(3)(2) in X
            True
            sage: GF(5)(2) in X
            False
            sage: sorted(Set(GF(7)) + Set(GF(3)), key=int)
            [0, 0, 1, 1, 2, 2, 3, 4, 5, 6]
        """
        if isinstance(X, Set_generic):
            if self is X:
                return self
            return Set_object_union(self, X)
        raise TypeError("X (=%s) must be a Set"%X)

    def __add__(self, X):
        """
        Return the union of ``self`` and ``X``.

        EXAMPLES::

            sage: Set(RealField()) + Set(QQ^5)
             Set-theoretic union of Set of elements of Real Field with 53 bits of precision and Set of elements of Vector space of dimension 5 over Rational Field
            sage: Set(GF(3)) + Set(GF(2))
            {0, 1, 2, 0, 1}
            sage: Set(GF(2)) + Set(GF(4,'a'))
            {0, 1, a, a + 1}
            sage: sorted(Set(GF(8,'b')) + Set(GF(4,'a')), key=str)
            [0, 0, 1, 1, a, a + 1, b, b + 1, b^2, b^2 + 1, b^2 + b, b^2 + b + 1]
        """
        return self.union(X)

    def __or__(self, X):
        """
        Return the union of ``self`` and ``X``.

        EXAMPLES::

            sage: Set([2,3]) | Set([3,4])
            {2, 3, 4}
            sage: Set(ZZ) | Set(QQ)
            Set-theoretic union of Set of elements of Integer Ring and Set of elements of Rational Field
        """

        return self.union(X)

    def intersection(self, X):
        r"""
        Return the intersection of ``self`` and ``X``.

        EXAMPLES::

            sage: X = Set(ZZ).intersection(Primes())
            sage: 4 in X
            False
            sage: 3 in X
            True

            sage: 2/1 in X
            True

            sage: X = Set(GF(9,'b')).intersection(Set(GF(27,'c')))
            sage: X
            {}

            sage: X = Set(GF(9,'b')).intersection(Set(GF(27,'b')))
            sage: X
            {}
        """
        if isinstance(X, Set_generic):
            if self is X:
                return self
            return Set_object_intersection(self, X)
        raise TypeError("X (=%s) must be a Set"%X)

    def difference(self, X):
        r"""
        Return the set difference ``self - X``.

        EXAMPLES::

            sage: X = Set(ZZ).difference(Primes())
            sage: 4 in X
            True
            sage: 3 in X
            False

            sage: 4/1 in X
            True

            sage: X = Set(GF(9,'b')).difference(Set(GF(27,'c')))
            sage: X
            {0, 1, 2, b, b + 1, b + 2, 2*b, 2*b + 1, 2*b + 2}

            sage: X = Set(GF(9,'b')).difference(Set(GF(27,'b')))
            sage: X
            {0, 1, 2, b, b + 1, b + 2, 2*b, 2*b + 1, 2*b + 2}
        """
        if isinstance(X, Set_generic):
            if self is X:
                return Set([])
            return Set_object_difference(self, X)
        raise TypeError("X (=%s) must be a Set"%X)

    def symmetric_difference(self, X):
        r"""
        Returns the symmetric difference of ``self`` and ``X``.

        EXAMPLES::

            sage: X = Set([1,2,3]).symmetric_difference(Set([3,4]))
            sage: X
            {1, 2, 4}
        """

        if isinstance(X, Set_generic):
            if self is X:
                return Set([])
            return Set_object_symmetric_difference(self, X)
        raise TypeError("X (=%s) must be a Set"%X)

    def __sub__(self, X):
        """
        Return the difference of ``self`` and ``X``.

        EXAMPLES::

            sage: X = Set(ZZ).difference(Primes())
            sage: Y = Set(ZZ) - Primes()
            sage: X == Y
            True
        """
        return self.difference(X)

    def __and__(self, X):
        """
        Returns the intersection of ``self`` and ``X``.

        EXAMPLES::

            sage: Set([2,3]) & Set([3,4])
            {3}
            sage: Set(ZZ) & Set(QQ)
            Set-theoretic intersection of Set of elements of Integer Ring and Set of elements of Rational Field
        """

        return self.intersection(X)

    def __xor__(self, X):
        """
        Returns the symmetric difference of ``self`` and ``X``.

        EXAMPLES::

            sage: X = Set([1,2,3,4])
            sage: Y = Set([1,2])
            sage: X.symmetric_difference(Y)
            {3, 4}
            sage: X.__xor__(Y)
            {3, 4}
        """
        return self.symmetric_difference(X)

    def cardinality(self):
        """
        Return the cardinality of this set, which is either an integer or
        ``Infinity``.

        EXAMPLES::

            sage: Set(ZZ).cardinality()
            +Infinity
            sage: Primes().cardinality()
            +Infinity
            sage: Set(GF(5)).cardinality()
            5
            sage: Set(GF(5^2,'a')).cardinality()
            25
        """
        if not self.is_finite():
            return sage.rings.infinity.infinity

        if self is not self.__object:
            try:
                return self.__object.cardinality()
            except (AttributeError, NotImplementedError):
                pass
            from sage.rings.integer import Integer
            try:
                return Integer(len(self.__object))
            except TypeError:
                pass

        raise NotImplementedError("computation of cardinality of %s not yet implemented"%self.__object)

    def is_empty(self):
        """
        Return boolean representing emptiness of the set.

        OUTPUT:

        True if the set is empty, False if otherwise.

        EXAMPLES::

            sage: Set([]).is_empty()
            True
            sage: Set([0]).is_empty()
            False
            sage: Set([1..100]).is_empty()
            False
            sage: Set(SymmetricGroup(2).list()).is_empty()
            False
            sage: Set(ZZ).is_empty()
            False

        TESTS::

            sage: Set([]).is_empty()
            True
            sage: Set([1,2,3]).is_empty()
            False
            sage: Set([1..100]).is_empty()
            False
            sage: Set(DihedralGroup(4).list()).is_empty()
            False
            sage: Set(QQ).is_empty()
            False
        """
        return not self

    def is_finite(self):
        """
        Return ``True`` if ``self`` is finite.

        EXAMPLES::

            sage: Set(QQ).is_finite()
            False
            sage: Set(GF(250037)).is_finite()
            True
            sage: Set(Integers(2^1000000)).is_finite()
            True
            sage: Set([1,'a',ZZ]).is_finite()
            True
        """
        obj = self.__object
        try:
            is_finite = obj.is_finite
        except AttributeError:
            return has_finite_length(obj)
        else:
            return is_finite()

    def object(self):
        """
        Return underlying object.

        EXAMPLES::

            sage: X = Set(QQ)
            sage: X.object()
            Rational Field
            sage: X = Primes()
            sage: X.object()
            Set of all prime numbers: 2, 3, 5, 7, ...
        """
        return self.__object

    def subsets(self,size=None):
        """
        Return the :class:`Subsets` object representing the subsets of a set.
        If size is specified, return the subsets of that size.

        EXAMPLES::

            sage: X = Set([1,2,3])
            sage: list(X.subsets())
            [{}, {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 3}, {1, 2, 3}]
            sage: list(X.subsets(2))
            [{1, 2}, {1, 3}, {2, 3}]

        """
        from sage.combinat.subset import Subsets
        return Subsets(self,size)

class Set_object_enumerated(Set_object):
    """
    A finite enumerated set.
    """
    def __init__(self, X):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(GF(19)); S
            {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}
            sage: S.category()
            Category of finite sets
            sage: print(latex(S))
            \left\{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18\right\}
            sage: TestSuite(S).run()
        """
        Set_object.__init__(self, X, category=Sets().Finite())

    def random_element(self):
        r"""
        Return a random element in this set.

        EXAMPLES::

            sage: Set([1,2,3]).random_element() # random
            2
        """
        try:
            return self.object().random_element()
        except AttributeError:
            # TODO: this very slow!
            return choice(self.list())

    def is_finite(self):
        r"""
        Return ``True`` as this is a finite set.

        EXAMPLES::

            sage: Set(GF(19)).is_finite()
            True
        """
        return True

    def cardinality(self):
        """
        Return the cardinality of ``self``.

        EXAMPLES::

            sage: Set([1,1]).cardinality()
            1
        """
        from sage.rings.integer import Integer
        return Integer(len(self.set()))

    def __len__(self):
        """
        EXAMPLES::

            sage: len(Set([1,1]))
            1
        """
        return len(self.set())

    def __iter__(self):
        r"""
        Iterating through the elements of ``self``.

        EXAMPLES::

            sage: S = Set(GF(19))
            sage: I = iter(S)
            sage: next(I)
            0
            sage: next(I)
            1
            sage: next(I)
            2
            sage: next(I)
            3
        """
        return iter(self.set())

    def _latex_(self):
        r"""
        Return the LaTeX representation of ``self``.

        EXAMPLES::

            sage: S = Set(GF(2))
            sage: latex(S)
            \left\{0, 1\right\}
        """
        return '\\left\\{' + ', '.join([latex(x) for x in self.set()])  + '\\right\\}'

    def _repr_(self):
        r"""
        Return the string representation of ``self``.

        EXAMPLES::

            sage: S = Set(GF(2))
            sage: S
            {0, 1}

        TESTS::

            sage: Set()
            {}
        """
        py_set = self.set()
        if six.PY3:
            if not py_set:
                return "{}"
            return repr(py_set)
        else:
            return "{" + repr(py_set)[5:-2] + "}"

    def list(self):
        """
        Return the elements of ``self``, as a list.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: X
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: X.list()
            [0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1]
            sage: type(X.list())
            <... 'list'>

        .. TODO::

            FIXME: What should be the order of the result?
            That of ``self.object()``? Or the order given by
            ``set(self.object())``? Note that :meth:`__getitem__` is
            currently implemented in term of this list method, which
            is really inefficient ...
        """
        return list(set(self.object()))

    def set(self):
        """
        Return the Python set object associated to this set.

        Python has a notion of finite set, and often Sage sets
        have an associated Python set.  This function returns
        that set.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: X
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: X.set()
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: type(X.set())
            <... 'set'>
            sage: type(X)
            <class 'sage.sets.set.Set_object_enumerated_with_category'>
        """
        return set(self.object())

    def frozenset(self):
        """
        Return the Python frozenset object associated to this set,
        which is an immutable set (hence hashable).

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: X
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: s = X.set(); s
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: hash(s)
            Traceback (most recent call last):
            ...
            TypeError: unhashable type: 'set'
            sage: s = X.frozenset(); s
            frozenset({0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1})

            sage: hash(s) != hash(tuple(X.set()))
            True

            sage: type(s)
            <... 'frozenset'>
        """
        return frozenset(self.object())

    def __hash__(self):
        """
        Return the hash of ``self`` (as a ``frozenset``).

        EXAMPLES::

            sage: s = Set(GF(8,'c'))
            sage: hash(s) == hash(s)
            True
        """
        return hash(self.frozenset())

    def __richcmp__(self, other, op):
        """
        Compare the sets ``self`` and ``other``.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: X == Set(GF(8,'c'))
            True
            sage: X == Set(GF(4,'a'))
            False
            sage: Set(QQ) == Set(ZZ)
            False
            sage: Set([1]) == set([1])
            True
        """
        if not isinstance(other, Set_object_enumerated):
            if isinstance(other, (set, frozenset)):
                return self.set() == other
            return NotImplemented
        if self.set() == other.set():
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def issubset(self, other):
        r"""
        Return whether ``self`` is a subset of ``other``.

        INPUT:

         - ``other`` -- a finite Set

        EXAMPLES::

            sage: X = Set([1,3,5])
            sage: Y = Set([0,1,2,3,5,7])
            sage: X.issubset(Y)
            True
            sage: Y.issubset(X)
            False
            sage: X.issubset(X)
            True

        TESTS::

            sage: len([Z for Z in Y.subsets() if Z.issubset(X)])
            8
        """
        if not isinstance(other, Set_object_enumerated):
            raise NotImplementedError
        return self.set().issubset(other.set())

    def issuperset(self, other):
        r"""
        Return whether ``self`` is a superset of ``other``.

        INPUT:

         - ``other`` -- a finite Set

        EXAMPLES::

            sage: X = Set([1,3,5])
            sage: Y = Set([0,1,2,3,5])
            sage: X.issuperset(Y)
            False
            sage: Y.issuperset(X)
            True
            sage: X.issuperset(X)
            True

        TESTS::

            sage: len([Z for Z in Y.subsets() if Z.issuperset(X)])
            4
        """
        if not isinstance(other, Set_object_enumerated):
            raise NotImplementedError
        return self.set().issuperset(other.set())

    def union(self, other):
        """
        Return the union of ``self`` and ``other``.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: Y = Set([GF(8,'c').0, 1, 2, 3])
            sage: X
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: sorted(Y)
            [1, 2, 3, c]
            sage: sorted(X.union(Y), key=str)
            [0, 1, 2, 3, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1]
        """
        if not isinstance(other, Set_object_enumerated):
            return Set_object.union(self, other)
        return Set_object_enumerated(self.set().union(other.set()))

    def intersection(self, other):
        """
        Return the intersection of ``self`` and ``other``.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: Y = Set([GF(8,'c').0, 1, 2, 3])
            sage: X.intersection(Y)
            {1, c}
        """
        if not isinstance(other, Set_object_enumerated):
            return Set_object.intersection(self, other)
        return Set_object_enumerated(self.set().intersection(other.set()))

    def difference(self, other):
        """
        Return the set difference ``self - other``.

        EXAMPLES::

            sage: X = Set([1,2,3,4])
            sage: Y = Set([1,2])
            sage: X.difference(Y)
            {3, 4}
            sage: Z = Set(ZZ)
            sage: W = Set([2.5, 4, 5, 6])
            sage: W.difference(Z)
            {2.50000000000000}
        """
        if not isinstance(other, Set_object_enumerated):
            return Set([x for x in self if x not in other])
        return Set_object_enumerated(self.set().difference(other.set()))

    def symmetric_difference(self, other):
        """
        Return the symmetric difference of ``self`` and ``other``.

        EXAMPLES::

            sage: X = Set([1,2,3,4])
            sage: Y = Set([1,2])
            sage: X.symmetric_difference(Y)
            {3, 4}
            sage: Z = Set(ZZ)
            sage: W = Set([2.5, 4, 5, 6])
            sage: U = W.symmetric_difference(Z)
            sage: 2.5 in U
            True
            sage: 4 in U
            False
            sage: V = Z.symmetric_difference(W)
            sage: V == U
            True
            sage: 2.5 in V
            True
            sage: 6 in V
            False
        """
        if not isinstance(other, Set_object_enumerated):
            return Set_object.symmetric_difference(self, other)
        return Set_object_enumerated(self.set().symmetric_difference(other.set()))

class Set_object_binary(Set_object):
    r"""
    An abstract common base class for sets defined by a binary operation (ex.
    :class:`Set_object_union`, :class:`Set_object_intersection`,
    :class:`Set_object_difference`, and
    :class:`Set_object_symmetric_difference`).

    INPUT:

    - ``X``, ``Y`` -- sets, the operands to ``op``

    - ``op`` -- a string describing the binary operation

    - ``latex_op`` -- a string used for rendering this object in LaTeX

    EXAMPLES::

        sage: X = Set(QQ^2)
        sage: Y = Set(ZZ)
        sage: from sage.sets.set import Set_object_binary
        sage: S = Set_object_binary(X, Y, "union", "\\cup"); S
        Set-theoretic union of Set of elements of Vector space of dimension 2
         over Rational Field and Set of elements of Integer Ring
    """
    def __init__(self, X, Y, op, latex_op):
        r"""
        Initialization.

        TESTS::

            sage: from sage.sets.set import Set_object_binary
            sage: X = Set(QQ^2)
            sage: Y = Set(ZZ)
            sage: S = Set_object_binary(X, Y, "union", "\\cup")
            sage: type(S)
            <class 'sage.sets.set.Set_object_binary_with_category'>
        """
        self._X = X
        self._Y = Y
        self._op = op
        self._latex_op = latex_op
        Set_object.__init__(self, self)

    def _repr_(self):
        r"""
        Return a string representation of this set.

        EXAMPLES::

            sage: Set(ZZ).union(Set(GF(5)))
            Set-theoretic union of Set of elements of Integer Ring and {0, 1, 2, 3, 4}
        """
        return "Set-theoretic {} of {} and {}".format(self._op, self._X, self._Y)

    def _latex_(self):
        r"""
        Return a latex representation of this set.

        EXAMPLES::

            sage: latex(Set(ZZ).union(Set(GF(5))))
            \Bold{Z} \cup \left\{0, 1, 2, 3, 4\right\}
        """
        return latex(self._X) + self._latex_op + latex(self._Y)

    def __hash__(self):
        """
        The hash value of this set.

        EXAMPLES:

        The hash values of equal sets are in general not equal since it is not
        decidable whether two sets are equal::

            sage: X = Set(GF(13)).intersection(Set(ZZ))
            sage: Y = Set(ZZ).intersection(Set(GF(13)))
            sage: hash(X) == hash(Y)
            False

        TESTS:

        Test that :trac:`14432` has been resolved::

            sage: S = Set(ZZ).union(Set([infinity]))
            sage: T = Set(ZZ).union(Set([infinity]))
            sage: hash(S) == hash(T)
            True
        """
        return hash((self._X, self._Y, self._op))

class Set_object_union(Set_object_binary):
    """
    A formal union of two sets.
    """
    def __init__(self, X, Y):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(QQ^2)
            sage: T = Set(ZZ)
            sage: X = S.union(T); X
            Set-theoretic union of Set of elements of Vector space of dimension 2 over Rational Field and Set of elements of Integer Ring

            sage: latex(X)
            \Bold{Q}^{2} \cup \Bold{Z}

            sage: TestSuite(X).run()
        """
        Set_object_binary.__init__(self, X, Y, "union", "\\cup")

    def is_finite(self):
        r"""
        Return whether this set is finite.

        EXAMPLES::

            sage: X = Set(range(10))
            sage: Y = Set(range(-10,0))
            sage: Z = Set(Primes())
            sage: X.union(Y).is_finite()
            True
            sage: X.union(Z).is_finite()
            False
        """
        return self._X.is_finite() and self._Y.is_finite()

    def __richcmp__(self, right, op):
        r"""
        Try to compare ``self`` and ``right``.

        .. NOTE::

           Comparison is basically not implemented, or rather it could
           say sets are not equal even though they are.  I don't know
           how one could implement this for a generic union of sets in
           a meaningful manner.  So be careful when using this.

        EXAMPLES::

            sage: Y = Set(ZZ^2).union(Set(ZZ^3))
            sage: X = Set(ZZ^3).union(Set(ZZ^2))
            sage: X == Y
            True
            sage: Y == X
            True

        This illustrates that equality testing for formal unions
        can be misleading in general.

        ::

            sage: Set(ZZ).union(Set(QQ)) == Set(QQ)
            False
        """
        if not isinstance(right, Set_generic):
            return rich_to_bool(op, -1)
        if not isinstance(right, Set_object_union):
            return rich_to_bool(op, -1)
        if self._X == right._X and self._Y == right._Y or \
           self._X == right._Y and self._Y == right._X:
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def __iter__(self):
        """
        Return iterator over the elements of ``self``.

        EXAMPLES::

            sage: [x for x in Set(GF(3)).union(Set(GF(2)))]
            [0, 1, 2, 0, 1]
        """
        for x in self._X:
            yield x
        for y in self._Y:
            yield y

    def __contains__(self, x):
        """
        Return ``True`` if ``x`` is an element of ``self``.

        EXAMPLES::

            sage: X = Set(GF(3)).union(Set(GF(2)))
            sage: GF(5)(1) in X
            False
            sage: GF(3)(2) in X
            True
            sage: GF(2)(0) in X
            True
            sage: GF(5)(0) in X
            False
        """
        return x in self._X or x in self._Y

    def cardinality(self):
        """
        Return the cardinality of this set.

        EXAMPLES::

            sage: X = Set(GF(3)).union(Set(GF(2)))
            sage: X
            {0, 1, 2, 0, 1}
            sage: X.cardinality()
            5

            sage: X = Set(GF(3)).union(Set(ZZ))
            sage: X.cardinality()
            +Infinity
        """
        return self._X.cardinality() + self._Y.cardinality()

class Set_object_intersection(Set_object_binary):
    """
    Formal intersection of two sets.
    """
    def __init__(self, X, Y):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(QQ^2)
            sage: T = Set(ZZ)
            sage: X = S.intersection(T); X
            Set-theoretic intersection of Set of elements of Vector space of dimension 2 over Rational Field and Set of elements of Integer Ring
            sage: latex(X)
            \Bold{Q}^{2} \cap \Bold{Z}

            sage: X = Set(IntegerRange(100)).intersection(Primes())
            sage: X.is_finite()
            True
            sage: TestSuite(X).run()
        """
        Set_object_binary.__init__(self, X, Y, "intersection", "\\cap")

    def is_finite(self):
        r"""
        Return whether this set is finite.

        EXAMPLES::

            sage: X = Set(IntegerRange(100))
            sage: Y = Set(ZZ)
            sage: X.intersection(Y).is_finite()
            True
            sage: Y.intersection(X).is_finite()
            True
            sage: Y.intersection(Set(QQ)).is_finite()
            Traceback (most recent call last):
            ...
            NotImplementedError
        """
        if self._X.is_finite():
            return True
        elif self._Y.is_finite():
            return True
        raise NotImplementedError

    def __richcmp__(self, right, op):
        r"""
        Try to compare ``self`` and ``right``.

        .. NOTE::

           Comparison is basically not implemented, or rather it could
           say sets are not equal even though they are.  I don't know
           how one could implement this for a generic intersection of
           sets in a meaningful manner.  So be careful when using this.

        EXAMPLES::

            sage: Y = Set(ZZ).intersection(Set(QQ))
            sage: X = Set(QQ).intersection(Set(ZZ))
            sage: X == Y
            True
            sage: Y == X
            True

        This illustrates that equality testing for formal unions
        can be misleading in general.

        ::

            sage: Set(ZZ).intersection(Set(QQ)) == Set(QQ)
            False
        """
        if not isinstance(right, Set_generic):
            return rich_to_bool(op, -1)
        if not isinstance(right, Set_object_intersection):
            return rich_to_bool(op, -1)
        if self._X == right._X and self._Y == right._Y or \
           self._X == right._Y and self._Y == right._X:
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def __iter__(self):
        """
        Return iterator through elements of ``self``.

        ``self`` is a formal intersection of `X` and `Y` and this function is
        implemented by iterating through the elements of `X` and for
        each checking if it is in `Y`, and if yielding it.

        EXAMPLES::

            sage: X = Set(ZZ).intersection(Primes())
            sage: I = X.__iter__()
            sage: next(I)
            2

        Check that known finite intersections have finite iterators (see
        :trac:`18159`)::

            sage: P = Set(ZZ).intersection(Set(range(10,20)))
            sage: list(P)
            [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
        """
        X = self._X
        Y = self._Y
        if not self._X.is_finite() and self._Y.is_finite():
            X,Y = Y,X
        for x in X:
            if x in Y:
                yield x

    def __contains__(self, x):
        """
        Return ``True`` if ``self`` contains ``x``.

        Since ``self`` is a formal intersection of `X` and `Y` this function
        returns ``True`` if both `X` and `Y` contains ``x``.

        EXAMPLES::

            sage: X = Set(QQ).intersection(Set(RR))
            sage: 5 in X
            True
            sage: ComplexField().0 in X
            False

        Any specific floating-point number in Sage is to finite precision,
        hence it is rational::

            sage: RR(sqrt(2)) in X
            True

        Real constants are not rational::

            sage: pi in X
            False
        """
        return x in self._X and x in self._Y

class Set_object_difference(Set_object_binary):
    """
    Formal difference of two sets.
    """
    def __init__(self, X, Y):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(QQ)
            sage: T = Set(ZZ)
            sage: X = S.difference(T); X
            Set-theoretic difference of Set of elements of Rational Field and Set of elements of Integer Ring
            sage: latex(X)
            \Bold{Q} - \Bold{Z}

            sage: TestSuite(X).run()
        """
        Set_object_binary.__init__(self, X, Y, "difference", "-")

    def is_finite(self):
        r"""
        Return whether this set is finite.

        EXAMPLES::

            sage: X = Set(range(10))
            sage: Y = Set(range(-10,5))
            sage: Z = Set(QQ)
            sage: X.difference(Y).is_finite()
            True
            sage: X.difference(Z).is_finite()
            True
            sage: Z.difference(X).is_finite()
            False
            sage: Z.difference(Set(ZZ)).is_finite()
            Traceback (most recent call last):
            ...
            NotImplementedError
        """
        if self._X.is_finite():
            return True
        elif self._Y.is_finite():
            return False
        raise NotImplementedError

    def __richcmp__(self, right, op):
        r"""
        Try to compare ``self`` and ``right``.

        .. NOTE::

           Comparison is basically not implemented, or rather it could
           say sets are not equal even though they are.  I don't know
           how one could implement this for a generic intersection of
           sets in a meaningful manner.  So be careful when using
           this.

        EXAMPLES::

            sage: Y = Set(ZZ).difference(Set(QQ))
            sage: Y == Set([])
            False
            sage: X = Set(QQ).difference(Set(ZZ))
            sage: Y == X
            False
            sage: Z = X.difference(Set(ZZ))
            sage: Z == X
            False

        This illustrates that equality testing for formal unions
        can be misleading in general.

        ::

            sage: X == Set(QQ).difference(Set(ZZ))
            True
        """
        if not isinstance(right, Set_generic):
            return rich_to_bool(op, -1)
        if not isinstance(right, Set_object_difference):
            return rich_to_bool(op, -1)
        if self._X == right._X and self._Y == right._Y:
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def __iter__(self):
        """
        Return iterator through elements of ``self``.

        ``self`` is a formal difference of `X` and `Y` and this function
        is implemented by iterating through the elements of `X` and for
        each checking if it is not in `Y`, and if yielding it.

        EXAMPLES::

            sage: X = Set(ZZ).difference(Primes())
            sage: I = X.__iter__()
            sage: next(I)
            0
            sage: next(I)
            1
            sage: next(I)
            -1
            sage: next(I)
            -2
            sage: next(I)
            -3
        """
        for x in self._X:
            if x not in self._Y:
                yield x

    def __contains__(self, x):
        """
        Return ``True`` if ``self`` contains ``x``.

        Since ``self`` is a formal intersection of `X` and `Y` this function
        returns ``True`` if both `X` and `Y` contains ``x``.

        EXAMPLES::

            sage: X = Set(QQ).difference(Set(ZZ))
            sage: 5 in X
            False
            sage: ComplexField().0 in X
            False
            sage: sqrt(2) in X     # since sqrt(2) is not a numerical approx
            False
            sage: sqrt(RR(2)) in X # since sqrt(RR(2)) is a numerical approx
            True
            sage: 5/2 in X
            True
        """
        return x in self._X and x not in self._Y

class Set_object_symmetric_difference(Set_object_binary):
    """
    Formal symmetric difference of two sets.
    """
    def __init__(self, X, Y):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(QQ)
            sage: T = Set(ZZ)
            sage: X = S.symmetric_difference(T); X
            Set-theoretic symmetric difference of Set of elements of Rational Field and Set of elements of Integer Ring
            sage: latex(X)
            \Bold{Q} \bigtriangleup \Bold{Z}

            sage: TestSuite(X).run()
        """
        Set_object_binary.__init__(self, X, Y, "symmetric difference", "\\bigtriangleup")

    def is_finite(self):
        r"""
        Return whether this set is finite.

        EXAMPLES::

            sage: X = Set(range(10))
            sage: Y = Set(range(-10,5))
            sage: Z = Set(QQ)
            sage: X.symmetric_difference(Y).is_finite()
            True
            sage: X.symmetric_difference(Z).is_finite()
            False
            sage: Z.symmetric_difference(X).is_finite()
            False
            sage: Z.symmetric_difference(Set(ZZ)).is_finite()
            Traceback (most recent call last):
            ...
            NotImplementedError
        """
        if self._X.is_finite():
            return self._Y.is_finite()
        elif self._Y.is_finite():
            return False
        raise NotImplementedError

    def __richcmp__(self, right, op):
        r"""
        Try to compare ``self`` and ``right``.

        .. NOTE::

           Comparison is basically not implemented, or rather it could
           say sets are not equal even though they are.  I don't know
           how one could implement this for a generic symmetric
           difference of sets in a meaningful manner.  So be careful
           when using this.

        EXAMPLES::

            sage: Y = Set(ZZ).symmetric_difference(Set(QQ))
            sage: X = Set(QQ).symmetric_difference(Set(ZZ))
            sage: X == Y
            True
            sage: Y == X
            True

        """
        if not isinstance(right, Set_generic):
            return rich_to_bool(op, -1)
        if not isinstance(right, Set_object_symmetric_difference):
            return rich_to_bool(op, -1)
        if self._X == right._X and self._Y == right._Y or \
           self._X == right._Y and self._Y == right._X:
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def __iter__(self):
        """
        Return iterator through elements of ``self``.

        This function is implemented by first iterating through the elements
        of `X` and  yielding it if it is not in `Y`.
        Then it will iterate throw all the elements of `Y` and yielding it if
        it is not in `X`.

        EXAMPLES::

            sage: X = Set(ZZ).symmetric_difference(Primes())
            sage: I = X.__iter__()
            sage: next(I)
            0
            sage: next(I)
            1
            sage: next(I)
            -1
            sage: next(I)
            -2
            sage: next(I)
            -3
        """
        for x in self._X:
            if x not in self._Y:
                yield x

        for y in self._Y:
            if y not in self._X:
                yield y

    def __contains__(self, x):
        """
        Return ``True`` if ``self`` contains ``x``.

        Since ``self`` is the formal symmetric difference of `X` and `Y`
        this function returns ``True`` if either `X` or `Y` (but not both)
        contains ``x``.

        EXAMPLES::

            sage: X = Set(QQ).symmetric_difference(Primes())
            sage: 4 in X
            True
            sage: ComplexField().0 in X
            False
            sage: sqrt(2) in X      # since sqrt(2) is currently symbolic
            False
            sage: sqrt(RR(2)) in X # since sqrt(RR(2)) is currently approximated
            True
            sage: pi in X
            False
            sage: 5/2 in X
            True
            sage: 3 in X
            False
        """
        return (x in self._X and x not in self._Y) \
               or (x in self._Y and x not in self._X)

def is_Set(x):
    """
    Deprecated. Use ``isinstance(x, Set_generic)`` instead.

    TESTS::

        sage: from sage.sets.set import is_Set
        sage: is_Set(Primes())
        doctest:...: DeprecationWarning: Please use isinstance(x, Set_generic)
        See http://trac.sagemath.org/24443 for details.
        True
    """
    from sage.misc.superseded import deprecation
    deprecation(24443, "Please use isinstance(x, Set_generic)")
    return isinstance(x, Set_generic)


# -*- coding: utf-8 -*-

"""
@author: krakowiakpawel9@gmail.com
@site: e-smartdata.org
"""

import numpy as np
import pandas as pd

s = pd.Series(['Apple', '   Microsoft', np.nan, '  Google  ', 'Anazon'])

# %%
s = s.str.strip()

# %%
lower = s.str.lower()

# %%
upper = s.str.upper()

# %%
length = s.str.len()

# %% case
df = pd.DataFrame(np.random.rand(10, 2),
                  columns=['          ID value ', '  Price'])

# %%
df.columns = df.columns.str.strip()

# %%
df.columns = df.columns.str.lower()

# %%
df.columns = df.columns.str.replace(' ', '_')


#!/usr/bin/env python
u"""
grace_months_html.py
Written by Tyler Sutterley (10/2020)

Creates a html file with the start and end days for each dataset
Shows the range of each month for CSR/GFZ/JPL (RL06) and GSFC (rl06v1.0)
Shows which months are missing for each dataset as **missing**

Similar to ftp://podaac.jpl.nasa.gov/allData/tellus/L3/Doc/GraceMonths.html
    ftp://podaac.jpl.nasa.gov/allData/tellus/L3/Doc/gracemonths_20160112.html

INPUTS:
    base_dir: Working data directory for GRACE/GRACE-FO data

OPTIONS:
    DREL: GRACE/GRACE-FO data release (RL06,rl06v1.0)

OUTPUTS:
    GRACE_months.txt
    Column 1: GRACE Month Number
    Column 2: Calendar Date
    Column 3: CSR RL06 Date Range
    Column 4: GFZ RL06 Date Range
    Column 5: GSFC rl06v1.0 Date Range
    Column 6: JPL RL06 Date Range

COMMAND LINE OPTIONS:
    --help: list the command line options
    -D X, --directory=X: Working GRACE/GRACE-FO data directory
    -R X, --release=X: GRACE/GRACE-FO data releases to run (RL06,rl06v1.0)

PYTHON DEPENDENCIES:
    numpy: Scientific Computing Tools For Python (https://numpy.org)

UPDATE HISTORY:
    Updated 03/2021: added options for GSFC Release-6 Version 1.0
    Updated 10/2020: use argparse to set command line parameters
    Updated 09/2020: add link to plain text table
    Updated 08/2020: using git lfs for image storage
    Updated 06/2020: use full calendar years to not require local dependencies
    Updated 03/2020: local import of required dependencies
    Updated 02/2020: add favicon to html header
    Updated 10/2019: no longer show Release-5 data by default
    Updated 06/2019: added notes for GRACE-FO data
    Updated 04/2019: set default releases for each data center
    Updated 07/2018: link images if hovering over a GRACE month for a center
        add navigation side bar and symbols to footer. include Wahr et al. 2015
        added column for GSFC mascon solutions between GFZ and JPL harmonics
    Updated 05/2018: added options for release 6
    Updated 09/2017: added more metadata to output html file
    Updated 05-06/2016: using __future__ print function. format month lines
        Highlight table rows on mouse hover
    Forked 04/2016: forked for HTML table creation
    Updated 03/2016: using getopt to set RL04 parameter, added new help module
        forked for markdown table creation
    Updated 10/2015: cleaned up and added a few comments
    Updated 11/2014: minor updates to code. added main definition
    Updated 10/2014: updated comments, current Sean Geocenter file
    Updated 05/2014: added OPTION to not run RL04
    Updated 07/2013: minor update: new Sean geocenter file
        moved geocenter files to grace.dir/geocenter.dir/
    Updated 05/2013: converted to Python and added years to month label
    Updated 03/2013: changed degree 1 to show both RL04 and RL05
    Updated 02/2013: new degree 1 file from Sean Swenson
        Changed to read from ascii files created from grace_date.pro
    Updated 11/2012: added DEG1 and SLR outputs
    Written 07/2012
"""
from __future__ import print_function

import sys
import os
import inspect
import argparse
import numpy as np
import calendar,time

#-- PURPOSE: create HTML file of GRACE "nominal" months
def grace_months(base_dir, DREL=['RL06','rl06v1.0']):

    #-- Opening output GRACE months HTML file
    filename = inspect.getframeinfo(inspect.currentframe()).filename
    filepath = os.path.dirname(os.path.abspath(filename))
    fid = open(os.path.join(filepath,'GRACE-Months.html'), 'w')

    #-- Initial parameters
    #-- processing centers
    PROC = ['CSR', 'GFZ', 'JPL', 'GSFC']
    #-- read from GSM datasets
    DSET = 'GSM'
    #-- maximum month of the datasets
    #-- checks for the maximum month between processing centers
    max_mon = 0
    #-- contain the information for each dataset
    var_info = {}

    #-- Looping through data releases first (all RL04 then all RL05)
    #-- for each considered data release (RL04,RL05)
    for rl in DREL:
        #-- for each processing centers (CSR, GFZ, JPL)
        for pr in PROC:
            #-- Setting the data directory for processing center and release
            grace_dir = os.path.join(base_dir,pr,rl,DSET)
            #-- read GRACE date ascii file
            #-- file created in read_grace.py or grace_dates.py
            grace_date_file = '{0}_{1}_DATES.txt'.format(pr,rl)
            if os.access(os.path.join(grace_dir,grace_date_file), os.F_OK):
                #-- skip the header line
                date_input = np.loadtxt(os.path.join(grace_dir,grace_date_file),
                    skiprows=1)
                #-- number of months
                nmon = np.shape(date_input)[0]

                #-- Setting the dictionary key e.g. 'CSR RL04'
                var_name = '{0} {1}'.format(pr,rl)

                #-- Creating a python dictionary for each dataset with parameters:
                #-- month #, start year, start day, end year, end day
                #-- Purpose is to get all of the dates loaded for each dataset
                #-- Adding data to dictionary for data processing and release
                var_info[var_name] = {}
                #-- allocate for output variables
                var_info[var_name]['mon'] = np.zeros((nmon),dtype=np.int)
                var_info[var_name]['styr'] = np.zeros((nmon),dtype=np.int)
                var_info[var_name]['stday'] = np.zeros((nmon),dtype=np.int)
                var_info[var_name]['endyr'] = np.zeros((nmon),dtype=np.int)
                var_info[var_name]['endday'] = np.zeros((nmon),dtype=np.int)
                #-- place output variables in dictionary
                for i,key in enumerate(['mon','styr','stday','endyr','endday']):
                    #-- first column is date in decimal form (start at 1 not 0)
                    var_info[var_name][key] = date_input[:,i+1].astype(np.int)
                #-- Finding the maximum month measured
                if (var_info[var_name]['mon'].max() > max_mon):
                    #-- if the maximum month in this dataset is greater
                    #-- than the previously read datasets
                    max_mon = np.int(var_info[var_name]['mon'].max())

    #-- print HTML headers
    print('<!DOCTYPE html>', file=fid)
    print('', file=fid)
    print('\t', file=fid)
    print('\t<meta charset="utf-8">', file=fid)
    print('\t<meta name="author" content="Tyler Sutterley">', file=fid)
    print('\t<meta name="viewport" content="width=device-width">', file=fid)
    print('\t<title>GRACE/GRACE-FO Months</title>', file=fid)
    print('\t<link rel="icon" href="../assets/img/favicon.ico" type="image/x-icon"/>', file=fid)
    print('\t<link rel="stylesheet" href="../assets/css/styles.css">', file=fid)
    print('\t<link rel="stylesheet" href="../assets/css/font-awesome.min.css">', file=fid)
    print('\t<link rel="stylesheet" href="../assets/css/academicons.min.css">', file=fid)
    print('\t', file=fid)
    print('\t', file=fid)
    print('\t', file=fid)
    print('\t\t
', file=fid)
    print('\t\t\t×', file=fid)
    print('\t\t\tHome', file=fid)
    print('\t\t\tPublications', file=fid)
    print('\t\t\tPresentations', file=fid)
    print('\t\t\tDatasets', file=fid)
    print('\t\t\tDocumentation', file=fid)
    print('\t\t\tCurriculum Vitae', file=fid)
    print('\t\t\tNews', file=fid)
    print('\t\t\tResources', file=fid)
    print('\t\t\tGRACE Greenland Animation', file=fid)
    print('\t\t\tGRACE Antarctic Animation', file=fid)
    print(('\t\t
\n\t\t☰'), file=fid)
    print('\t\t
', file=fid)
    #-- print table header
    print('\t\t
', file=fid)
    print('\t\t
', file=fid)
    print('\t\t\tMonth	', file=fid)
    print('\t\t\tDate	', file=fid)
    #-- sort datasets alphanumerically
    var_name = sorted(var_info.keys())
    for v in var_name:
        print('\t\t\t{0}	'.format(v),file=fid)
    print('\t\t
', file=fid)
    print('\t\t
', file=fid)
    #-- print table body
    print('\t\t
', file=fid)
    #-- for each possible month
    #-- GRACE starts at month 004 (April 2002)
    #-- max_mon+1 to include max_mon
    for m in range(4, max_mon+1):
        #-- finding the month name e.g. Apr
        calendar_year = 2002 + (m-1)//12
        calendar_month = (m-1) % 12 + 1
        month_string = calendar.month_abbr[calendar_month]
        #-- printing table lines to file
        print('\t\t
',file=fid)
        print('\t\t\t{0:03d}	'.format(m),
            file=fid)
        print('\t\t\t{0}{1:4d}	'.format(
            month_string,calendar_year), file=fid)
        #-- for each processing center and data release
        for var in var_name:
            #-- split var name for data processing center and release
            PROC,DREL = var.split()
            #-- find if the month of data exists
            #-- exists will be greater than 0 if there is a match
            exists = np.count_nonzero(var_info[var]['mon'] == m)
            if (exists != 0):
                #-- if there is a matching month
                #-- indice of matching month
                ind, = np.nonzero(var_info[var]['mon'] == m)
                #-- start date
                st_yr, = var_info[var]['styr'][ind]
                st_day, = var_info[var]['stday'][ind]
                #-- end date
                end_yr, = var_info[var]['endyr'][ind]
                end_day, = var_info[var]['endday'][ind]
                #-- output table element is the date range
                #-- string format: 2002_102--2002_120
                args = (st_yr, st_day, end_yr, end_day)
                print(('\t\t\t'
                    '{0:4d}_{1:03d}–{2:4d}_{3:03d}').format(*args),file=fid)
                print('\t\t\t\t',file=fid)
                src = '{0}-{1}-{2:03d}.jpg'.format(PROC,DREL,m)
                print('\t\t\t\t\t'.format(src),file=fid)
                print('\t\t\t\t',file=fid)
                print('\t\t\t	', file=fid)
            else:
                #-- if there is no matching month: missing or not yet processed
                print(('\t\t\t'
                    '**missing**	'), file=fid)
        #-- end of table row
        print('\t\t
', file=fid)
    #-- print table body footer text
    print('\t\t
', file=fid)
    print('\t\t
', file=fid)

    #-- print references
    print('\n\t\t
', file=fid)
    print('\t\t
GRACE/GRACE-FO anomalies for harmonic solutions are calculated in reference to the 2003'
        '–2010 mean\n\t\tand are smoothed using a 350km radius Gaussian filter', file=fid)
    print(('''\t\t'''
        '(Wahr et al., 1998)'), file=fid)
    print('\t\tafter destriping with a decorrelation algorithm ', file=fid)
    print(('''\t\t'''
        '(Swenson and Wahr, 2006).'), file=fid)
    #-- pole tide drift if showing Release-5 products
    if ('RL05' in DREL):
        print(('\t\tGRACE Release-5 data products are corrected for pole tide '
            'drift following '), file=fid)
        print(('''\t\t'''
            'Wahr et al. (2015).'), file=fid)
    print(('\t\tGSFC GRACE/GRACE-FO mascon data products are calculated as described in '), file=fid)
    print(('''\t\t'''
        'Loomis et al. (2019).'), file=fid)
    print('\t\tGRACE/GRACE-FO fields have been corrected for Glacial Isostatic '
        'Adjustment (GIA) using coefficients from ICE6G Version-D', file=fid)
    print(('''\t\t'''
        '(Peltier et al. , 2018).'), file=fid)

    print('\t\t
', file=fid)
    print('\t\t\t
', file=fid)
    print('\t\t\t\t	', file=fid)
    print('\t\t\t\t', file=fid)
    print('\t\t\t\tS. Swenson and J. Wahr.', file=fid)
    print('\t\t\t\tPost-processing removal of correlated errors in GRACE data.', file=fid)
    print('\t\t\t\tGeophysical Research Letters, 33(8), 2006.', file=fid)
    print('\t\t\t\t[ bib |', file=fid)
    print('\t\t\t\thttp ]', file=fid)
    print('\t\t\t\t	', file=fid)
    print('\t\t\t
', file=fid)

    print('\t\t\t
', file=fid)
    print('\t\t\t\t	', file=fid)
    print('\t\t\t\t', file=fid)
    print('\t\t\t\tJ. Wahr, M. Molenaar and F. Bryan.', file=fid)
    print("\t\t\t\tTime variability of the Earth's gravity field: Hydrological and", file=fid)
    print('\t\t\t\toceanic effects and their possible detection using GRACE.', file=fid)
    print('\t\t\t\tJournal of Geophysical Research: Solid Earth,', file=fid)
    print('\t\t\t\t103(B12):30205–30229, 1998.', file=fid)
    print('\t\t\t\t[ bib |', file=fid)
    print('\t\t\t\thttp ]', file=fid)
    print('\t\t\t\t	', file=fid)
    print('\t\t\t
', file=fid)

    #-- pole tide drift if showing Release-5 products
    if ('RL05' in DREL):
        print('\t\t\t
', file=fid)
        print('\t\t\t\t	', file=fid)
        print('\t\t\t\t', file=fid)
        print('\t\t\t\tJ. Wahr, R. S. Nerem and S. V. Bettadpur.', file=fid)
        print('\t\t\t\tThe pole tide and its effect on GRACE time-variable gravity measurements:', file=fid)
        print('\t\t\t\tImplications for estimates of surface mass variations.', file=fid)
        print('\t\t\t\tJournal of Geophysical Research: Solid Earth,', file=fid)
        print('\t\t\t\t120(6):4597–4615, 2015.', file=fid)
        print('\t\t\t\t[ bib |', file=fid)
        print('\t\t\t\thttp ]', file=fid)
        print('\t\t\t\t	', file=fid)
        print('\t\t\t
', file=fid)

    print('\t\t\t
', file=fid)
    print('\t\t\t\t	', file=fid)
    print('\t\t\t\t', file=fid)
    print('\t\t\t\tB. D. Loomis, S. B. Luthcke, T. J. Sabaka.', file=fid)
    print('\t\t\t\tRegularization and error characterization of GRACE mascons.', file=fid)
    print('\t\t\t\tJournal of Geodesy,', file=fid)
    print('\t\t\t\t93(9):1381–1398, 2019.', file=fid)
    print('\t\t\t\t[ bib |', file=fid)
    print('\t\t\t\thttp ]', file=fid)

    print('\t\t\t
', file=fid)
    print('\t\t\t\t	', file=fid)
    print('\t\t\t\t', file=fid)
    print('\t\t\t\tW. R. Peltier, D. F. Argus, R. Drummond.', file=fid)
    print('\t\t\t\tComment on "An Assessment of the ICE-6G_C (VM5a) Glacial ', file=fid)
    print('\t\t\t\tIsostatic Adjustment Model" by Purcell et al.', file=fid)
    print('\t\t\t\tJournal of Geophysical Research: Solid Earth,', file=fid)
    print('\t\t\t\t123(2):2019–2028, 2018.', file=fid)
    print('\t\t\t\t[ bib |', file=fid)
    print('\t\t\t\thttp ]', file=fid)
    print('\t\t\t\t	', file=fid)
    print('\t\t\t
\n\t\t
\n\t\t
\n\t\t
', file=fid)

    #-- print footer text
    args = (time.strftime('%Y-%m-%d',time.localtime()), os.path.basename(sys.argv[0]))
    print(('\n\t\t
Table generated on {0} with \n'
        '\t\t\t{1}
').format(*args), file=fid)
    print('\t\tTable as plain text
', file=fid)
    #-- print navigation symbols
    print(('\t\t
\n\t\t\t'
        '\n\t\t\t'
        ''
        ''
        '\n\t\t
'), file=fid)

    #-- print javascript commands
    print(('\t\t'), file=fid)
    print(('\t\t'), file=fid)
    print(('\t\t'), file=fid)
    print(('\t\t'), file=fid)
    #-- print HTML footers
    print('\t\n', file=fid)
    #-- close output HTML file
    fid.close()

#-- PURPOSE: functional call to grace_months() if running as program
def main():
    #-- Read the system arguments listed after the program
    parser = argparse.ArgumentParser(
        description="""SCreates a html file with the start and end days for
            each dataset
            """
    )
    #-- command line parameters
    #-- working data directory
    parser.add_argument('--directory','-D',
        type=lambda p: os.path.abspath(os.path.expanduser(p)),
        default=os.getcwd(),
        help='Working data directory')
    #-- GRACE/GRACE-FO data release
    parser.add_argument('--release','-r',
        metavar='DREL', type=str, nargs='+',
        default=['RL06','rl06v1.0'],
        help='GRACE/GRACE-FO data release')
    args = parser.parse_args()

    #-- run GRACE/GRACE-FO months program
    grace_months(args.directory, DREL=args.release)

#-- run main program
if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
"""
Example script to serve as starting point for display evaluation of reconstructions
using the Parallel Level Sets (PLS) prior.

Prerequisite:
You should have executed the following on your command prompt
    ./run_simulation_brain.sh
    ./run_reconstruction_brain.sh
    ./run_reconstruction_brain_PSF.sh
    ./run_reconstruction_brain_PLS.sh

This will use the existing data in the folder. The exercise gets more interesting
if you added noise to the data.

Author: Kris Thielemans
"""
#%% Initial imports
import numpy
import matplotlib.pyplot as plt
import stir
from stirextra import *
import os
#%% go to directory with input files
# adapt this path to your situation (or start everything in the exercises directory)
os.chdir(os.getenv('STIR_exercises_PATH'))
#%% change directory to where the output files are
os.chdir('working_folder/brain')
#%% Read origin images
groundtruth=to_numpy(stir.FloatVoxelsOnCartesianGrid.read_from_file('ground_truth.hv'));
anatomical=to_numpy(stir.FloatVoxelsOnCartesianGrid.read_from_file('anatomical_image.hv'));

#%% variables for future display
maxforplot=groundtruth.max();
# pick central slice
slice=numpy.int(groundtruth.shape[0]/2)
#%% Display
plt.figure();

ax=plt.subplot(1,2,1);
plt.imshow(groundtruth[slice,:,:,]);
plt.clim(0,maxforplot)
plt.colorbar();
plt.axis('off');
ax.set_title('ground truth');

ax=plt.subplot(1,2,2);
plt.imshow(anatomical[slice,:,:,]);
plt.colorbar();
plt.axis('off');
ax.set_title('anatomical image');

#%% Read in images
# All reconstruction were run with 240 subiterations using OSL
# OSL_PLS_240.hv is the output with the PLS prior
# OSLPSF_PLS_240.hv is the output with the PLS prior when using resolution modelling
# We also read the OSEM imges with resolution modelling for comparison
OSEMPSF240=to_numpy(stir.FloatVoxelsOnCartesianGrid.read_from_file('OSEMPSF_240.hv'));
OSLPSF240=to_numpy(stir.FloatVoxelsOnCartesianGrid.read_from_file('OSL_PLS_240.hv'));
OSLPLSPSF240=to_numpy(stir.FloatVoxelsOnCartesianGrid.read_from_file('OSLPSF_PLS_240.hv'));
#%% bitmap display of images
plt.figure();

ax=plt.subplot(2,2,1);
plt.imshow(groundtruth[slice,:,:,]);
plt.clim(0,maxforplot)
plt.colorbar();
plt.axis('off');
ax.set_title('ground truth');

ax=plt.subplot(2,2,2);
plt.imshow(OSEMPSF240[slice,:,:,]);
plt.clim(0,maxforplot)
plt.colorbar();
plt.axis('off');
ax.set_title('OSEM (with PSF)');

ax=plt.subplot(2,2,3);
plt.imshow(OSLPSF240[slice,:,:,]);
plt.clim(0,maxforplot)
plt.colorbar();
plt.axis('off');
ax.set_title('PLS (no PSF)');

ax=plt.subplot(2,2,4);
plt.imshow(OSLPLSPSF240[slice,:,:,]);
plt.clim(0,maxforplot)
plt.colorbar();
plt.axis('off');
ax.set_title('PLS (with PSF)');

#%% What now?
# You can change the parameters of the PLS prior (edit OSMAPOSLPSF_PLS.par)


"""

"""

# IMPORT modules. Must have unittest, and probably coast.
import coast
import unittest
import numpy as np
import os.path as path
import xarray as xr
import unit_test_files as files

class test_gridded_initialisation(unittest.TestCase):
    def test_gridded_load_of_data_and_domain(self):
        # Check successfully load example data and domain
        sci = coast.Gridded(files.fn_nemo_dat, files.fn_nemo_dom, config=files.fn_config_t_grid)
        sci_attrs_ref = dict(
            [
                ("name", "AMM7_1d_20070101_20070131_25hourm_grid_T"),
                ("description", "ocean T grid variables, 25h meaned"),
                ("title", "ocean T grid variables, 25h meaned"),
                ("Conventions", "CF-1.6"),
                ("timeStamp", "2019-Dec-26 04:35:28 GMT"),
                ("uuid", "96cae459-d3a1-4f4f-b82b-9259179f95f7"),
            ]
        )

        # checking is LHS is a subset of RHS
        check1 = sci_attrs_ref.items() <= sci.dataset.attrs.items()

        self.assertTrue(check1, msg="Check1")

    def test_gridded_load_of_data_only(self):
        # Check load only data
        ds = xr.open_dataset(files.fn_nemo_dat)
        sci_load_ds = coast.Gridded(config=files.fn_config_t_grid)
        sci_load_ds.load_dataset(ds)
        sci_load_file = coast.Gridded(config=files.fn_config_t_grid)
        sci_load_file.load(files.fn_nemo_dat)
        check1 = sci_load_ds.dataset.identical(sci_load_file.dataset)
        self.assertTrue(check1, msg="check1")

    def test_gridded_load_variables_correctly_renamed(self):
        # Check temperature is correctly renamed
        sci = coast.Gridded(files.fn_nemo_dat, files.fn_nemo_dom, config=files.fn_config_t_grid)
        check1 = "temperature" in sci.dataset
        self.assertTrue(check1, msg="check1")

    def test_gridded_load_dimensions_correctly_renamed(self):
        # Check gridded dimensions are correctly renamed
        sci = coast.Gridded(files.fn_nemo_dat, files.fn_nemo_dom, config=files.fn_config_t_grid)
        check1 = sci.dataset.temperature.dims == ("t_dim", "z_dim", "y_dim", "x_dim")
        self.assertTrue(check1, msg="check1")

    def test_gridded_load_domain_only(self):
        # Check gridded load domain only
        nemo_f = coast.Gridded(fn_domain=files.fn_nemo_dom, config=files.fn_config_f_grid)

        check1 = False
        if nemo_f.dataset._coord_names == {"depth_0", "latitude", "longitude"}:
            var_name_list = []
            for var_name in nemo_f.dataset.data_vars:
                var_name_list.append(var_name)
            if var_name_list == ["bathymetry", "e1", "e2", "e3_0"]:
                check1 = True
        self.assertTrue(check1, msg="check1")

    def test_gridded_calculate_depth0_for_tuvwf(self):
        nemo_t = coast.Gridded(
            fn_data=files.fn_nemo_grid_t_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_t_grid
        )
        if not np.isclose(np.nansum(nemo_t.dataset.depth_0.values), 1705804300.0):
            raise ValueError(" X - Gridded depth_0 failed on t-grid failed")
        nemo_u = coast.Gridded(
            fn_data=files.fn_nemo_grid_u_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_u_grid
        )
        if not np.isclose(np.nansum(nemo_u.dataset.depth_0.values), 1705317600.0):
            raise ValueError(" X - Gridded depth_0 failed on u-grid failed")
        nemo_v = coast.Gridded(
            fn_data=files.fn_nemo_grid_v_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_v_grid
        )
        if not np.isclose(np.nansum(nemo_v.dataset.depth_0.values), 1705419100.0):
            raise ValueError(" X - Gridded depth_0 failed on v-grid failed")
        nemo_f = coast.Gridded(fn_domain=files.fn_nemo_dom, config=files.fn_config_f_grid)
        if not np.isclose(np.nansum(nemo_f.dataset.depth_0.values), 1704932600.0):
            raise ValueError(" X - Gridded depth_0 failed on f-grid failed")

    def test_gridded_load_subregion_with_domain(self):
        amm7 = coast.Gridded(files.fn_nemo_dat_subset, files.fn_nemo_dom, config=files.fn_config_t_grid)

        # checking all the coordinates mapped correctly to the dataset object
        check1 = amm7.dataset._coord_names == {"depth_0", "latitude", "longitude", "time"}
        self.assertTrue(check1, msg="check1")

    def test_gridded_load_multiple(self):
        amm7 = coast.Gridded(files.file_names_amm7, files.fn_nemo_dom, config=files.fn_config_t_grid, multiple=True)

        # checking all the coordinates mapped correctly to the dataset object
        check1 = amm7.dataset.time.size == 14
        self.assertTrue(check1, msg="check1")

    def test_gridded_compute_e3_from_ssh(self):
        nemo_t = coast.Gridded(
            fn_data=files.fn_nemo_grid_t_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_t_grid
        )

        e3t, e3u, e3v, e3f, e3w = coast.Gridded.get_e3_from_ssh(nemo_t, True, True, True, True, True)
        cksum = np.array([e3t.sum(), e3u.sum(), e3v.sum(), e3f.sum(), e3w.sum()])
        # these references are based on the example file's ssh field
        reference = np.array([8.337016e08, 8.333972e08, 8.344886e08, 8.330722e08, 8.265948e08])
        check1 = np.allclose(cksum, reference)
        self.assertTrue(check1, msg="check1")


#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Note: See python_module_template.py first

###
# Name: YOUR_FULL_NAME_HERE
# Student ID: ID_HERE
# Email: CHAPMAN_EMAIL_HERE
# Course: PHYS220/MATH220 Fall 2019
# Assignment: HOMEWORK_OR_CLASSWORK_NUMBER
###

# Many useful test functions are available in these modules
import nose.tools
import numpy.testing
# Look for the functions starting with "assert_" in each module
# Remember that you can use a python interpreter to see what they do
# For example:
#   >>> help(nose.tools.assert_almost_equal)

"""Test Module Description (Replace this docstring with your own documentation)

This docstring should contain an overview of the tests contained in the file.

All test modules should start with the prefix "test_" so that the "nose"
framework can locate which tests to run.
"""

# Test functions follow below
# For each logically independent test, write one test function.
# Each logical test may test multiple cases using multiple assert expressions.

def test_function_1():
    """Test function for nose

    Any function name starting with prefix "test_" will be automatically run
    by nose. Include docstrings for your tests so that nose can report
    information about the tests. Use function names that are themselves
    descriptive.

    In a test function, use an assert command to test a Boolean statement
    about the execution of your code.  If the assert fails, it throws
    an exception, which is caught by nose and reported as a failure.
    Anything that is printed to the screen during this function is
    suppressed unless there is a failure, where it can be used for
    debugging.
    """
    # Any print statements inside your test function will only be printed
    # in the event of a failure, and can be very useful for debugging

    print("Helpful test message")

    # This is an "assert" statement that uses the keyword "assert" like an
    # "if" test to check that an expression correctly evaluates to True.
    # If the if check fails, the "assert" expression throws an exception
    # and (optionally) prints the provided string for context.

    assert 2 == 1+1, "2 == 1+1"

    # You might consider using nose.tools.assert_equals(left_side, right_side)
    # instead of the raw assert command, since it gives more information
    # in the event of a failure

    nose.tools.assert_equals(2, 1+1)

    # When testing floating point numbers (with a decimal), always use
    # nose.tools.assert_almost_equals(left_side, right_side) instead to avoid
    # rounding errors (or alternatives in nose.tools or numpy.testing)

def test_function_2():
    """Second test for nose

    Each function is run as a separate logical test.
    """
    assert True


import pytest
import sys
import os
import numpy as np

sys.path.append(os.path.dirname(os.path.realpath(__file__)) + "/../src")

try:
    from base import switch, clean_contents, read_file, clean_contents_control
except ImportError as mod: # If the user didn't install the required modules beore trying to run SEED 2.0
    print("Install the required modules before starting:\n" + str(mod))
except Exception as err: # Any other exception that should occur (nothing else should happen, hence generalising all other exceptions)
    print("Error while importing:\n" + str(err))

class TestSwitch(object):
    def test_switch_with_correct_input(self):
        actual = [switch("finite_difference"), switch("savitzky_golay"), switch("spectral"), switch("spline"), switch("trend_filtered")]
        expected = [0, 1, 2, 3, 4]

        assert actual == expected

    def test_switch_with_bad_input(self):
        assert switch(200) == 0

class TestReadData(object):
    def test_default_lorenz(self):
        ts, dt, cont, var = clean_contents(read_file("data_Lorenz3d.csv", ""))
        ts = ts[0:4]
        cont = cont[0:4]

        expected_ts = [0., 0.002, 0.004, 0.006]
        expected_dt = 0.002
        expected_cont = [   [-8,8,27],
                            [-7.683508382,7.966250523,26.73151873],
                            [-7.373984577,7.92929174,26.46998112],
                            [-7.071350036,7.889537629,26.21523996]
                            ]
        expected_vars = ["x", "y", "z"]

        assert np.all(ts == expected_ts)
        assert np.all(dt == expected_dt)
        assert np.all(cont == expected_cont)
        assert np.all(var == expected_vars)

    def test_random_5d(self):
        ts, dt, cont, var = clean_contents(read_file("random_5d.csv", ""))
        ts = ts[0:4]
        cont = cont[0:4]

        expected_ts = [0., 0.001, 0.002, 0.003]
        expected_dt = 0.001
        expected_cont = [   [3.216510998,9.914181513,3.866592635,0.746334938,4.606170334],
                            [0.650685986,4.486606167,5.84805047,1.991102499,5.211598163],
                            [1.318906913,8.869120427,2.594125945,9.226401504,2.357597024],
                            [8.26282217,2.762088665,1.941167642,5.280758711,4.70049113]
                            ]
        expected_vars = ["a", "b", "c", "d", "e"]

        assert np.all(ts == expected_ts)
        assert np.all(dt == expected_dt)
        assert np.all(cont == expected_cont)
        assert np.all(var == expected_vars)

    def test_clean_contents_control(self):
        ts, dt, cont, u, var = clean_contents_control(read_file("predatorpreydata.csv", ""))
        ts = ts[0:4]
        cont = cont[0:4]
        u = u[0:4]

        expected_ts = [0.1, 0.2, 0.3, 0.4]
        expected_dt = 0.1
        expected_cont = [
            [1.0,1.0],
            [0.93815867,1.0],
            [0.894741036,0.993815867],
            [0.878213584,0.983355064]
        ]
        expected_u = np.array(['0.2196665', '0.437335995', '0.651031414', '0.858815353'])
        expected_vars = ["x", "y", "u"]

        assert np.all(ts == expected_ts)
        assert np.all(dt == expected_dt)
        assert np.all(cont == expected_cont)
        assert np.all(u == expected_u)
        assert np.all(var == expected_vars)


# Copyright 2018-2020 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Unit tests for the draw transform.
"""
import functools
import pytest

import pennylane as qml
from pennylane import numpy as np

def test_drawing():
    """Test circuit drawing"""

    x = np.array(0.1, requires_grad=True)
    y = np.array([0.2, 0.3], requires_grad=True)
    z = np.array(0.4, requires_grad=True)

    dev = qml.device("default.qubit", wires=2)

    @qml.beta.qnode(dev, interface="autograd")
    def circuit(p1, p2=y, **kwargs):
        qml.RX(p1, wires=0)
        qml.RY(p2[0] * p2[1], wires=1)
        qml.RX(kwargs["p3"], wires=0)
        qml.CNOT(wires=[0, 1])
        return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

    result = qml.draw(circuit)(p1=x, p3=z)
    expected = """\
 0: ──RX(0.1)───RX(0.4)──╭C──╭┤ ⟨Z ⊗ X⟩
 1: ──RY(0.06)───────────╰X──╰┤ ⟨Z ⊗ X⟩
"""

    assert result == expected

def test_drawing_tf():
    """Test circuit drawing when using TensorFlow"""
    tf = pytest.importorskip("tensorflow")

    x = tf.constant(0.1)
    y = tf.constant([0.2, 0.3])
    z = tf.Variable(0.4)

    dev = qml.device("default.qubit", wires=2)

    @qml.beta.qnode(dev, interface="tf")
    def circuit(p1, p2=y, **kwargs):
        qml.RX(p1, wires=0)
        qml.RY(p2[0] * p2[1], wires=1)
        qml.RX(kwargs["p3"], wires=0)
        qml.CNOT(wires=[0, 1])
        return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

    result = qml.draw(circuit)(p1=x, p3=z)
    expected = """\
 0: ──RX(0.1)───RX(0.4)──╭C──╭┤ ⟨Z ⊗ X⟩
 1: ──RY(0.06)───────────╰X──╰┤ ⟨Z ⊗ X⟩
"""

    assert result == expected

def test_drawing_torch():
    """Test circuit drawing when using Torch"""
    torch = pytest.importorskip("torch")

    x = torch.tensor(0.1, requires_grad=True)
    y = torch.tensor([0.2, 0.3], requires_grad=True)
    z = torch.tensor(0.4, requires_grad=True)

    dev = qml.device("default.qubit", wires=2)

    @qml.beta.qnode(dev, interface="torch")
    def circuit(p1, p2=y, **kwargs):
        qml.RX(p1, wires=0)
        qml.RY(p2[0] * p2[1], wires=1)
        qml.RX(kwargs["p3"], wires=0)
        qml.CNOT(wires=[0, 1])
        return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

    result = qml.draw(circuit)(p1=x, p3=z)
    expected = """\
 0: ──RX(0.1)───RX(0.4)──╭C──╭┤ ⟨Z ⊗ X⟩
 1: ──RY(0.06)───────────╰X──╰┤ ⟨Z ⊗ X⟩
"""

    assert result == expected

def test_drawing_jax():
    """Test circuit drawing when using JAX"""
    jax = pytest.importorskip("jax")
    jnp = jax.numpy

    x = jnp.array(0.1)
    y = jnp.array([0.2, 0.3])
    z = jnp.array(0.4)

    dev = qml.device("default.qubit", wires=2)

    @qml.beta.qnode(dev, interface="jax")
    def circuit(p1, p2=y, **kwargs):
        qml.RX(p1, wires=0)
        qml.RY(p2[0] * p2[1], wires=1)
        qml.RX(kwargs["p3"], wires=0)
        qml.CNOT(wires=[0, 1])
        return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

    result = qml.draw(circuit)(p1=x, p3=z)
    expected = """\
 0: ──RX(0.1)───RX(0.4)──╭C──╭┤ ⟨Z ⊗ X⟩
 1: ──RY(0.06)───────────╰X──╰┤ ⟨Z ⊗ X⟩
"""

    assert result == expected

def test_drawing_ascii():
    """Test circuit drawing when using ASCII characters"""
    from pennylane import numpy as np

    x = np.array(0.1, requires_grad=True)
    y = np.array([0.2, 0.3], requires_grad=True)
    z = np.array(0.4, requires_grad=True)

    dev = qml.device("default.qubit", wires=2)

    @qml.beta.qnode(dev, interface="autograd")
    def circuit(p1, p2=y, **kwargs):
        qml.RX(p1, wires=0)
        qml.RY(p2[0] * p2[1], wires=1)
        qml.RX(kwargs["p3"], wires=0)
        qml.CNOT(wires=[0, 1])
        return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

    result = qml.draw(circuit, charset="ascii")(p1=x, p3=z)
    expected = """\
 0: --RX(0.1)---RX(0.4)--+C--+| <Z @ X>
 1: --RY(0.06)-----------+X--+| <Z @ X>
"""

    assert result == expected

def test_show_all_wires_error():
    """Test that show_all_wires will raise an error if the provided wire
    order does not contain all wires on the device"""

    dev = qml.device("default.qubit", wires=[-1, "a", "q2", 0])

    @qml.beta.qnode(dev)
    def circuit():
        qml.Hadamard(wires=-1)
        qml.CNOT(wires=[-1, "q2"])
        return qml.expval(qml.PauliX(wires="q2"))

    with pytest.raises(ValueError, match="must contain all wires"):
        qml.draw(circuit, show_all_wires=True, wire_order=[-1, "a"])()

def test_missing_wire():
    """Test that wires not specifically mentioned in the wire
    reordering are appended at the bottom of the circuit drawing"""

    dev = qml.device("default.qubit", wires=["a", -1, "q2"])

    @qml.beta.qnode(dev)
    def circuit():
        qml.Hadamard(wires=-1)
        qml.CNOT(wires=["a", "q2"])
        qml.RX(0.2, wires="a")
        return qml.expval(qml.PauliX(wires="q2"))

    # test one missing wire
    res = qml.draw(circuit, wire_order=["q2", "a"])()
    expected = [
        " q2: ──╭X───────────┤ ⟨X⟩ ",
        "  a: ──╰C──RX(0.2)──┤     ",
        " -1: ───H───────────┤     \n",
    ]

    assert res == "\n".join(expected)

    # test one missing wire
    res = qml.draw(circuit, wire_order=["q2", -1])()
    expected = [
        " q2: ─────╭X───────────┤ ⟨X⟩ ",
        " -1: ──H──│────────────┤     ",
        "  a: ─────╰C──RX(0.2)──┤     \n",
    ]

    assert res == "\n".join(expected)

    # test multiple missing wires
    res = qml.draw(circuit, wire_order=["q2"])()
    expected = [
        " q2: ─────╭X───────────┤ ⟨X⟩ ",
        " -1: ──H──│────────────┤     ",
        "  a: ─────╰C──RX(0.2)──┤     \n",
    ]

    assert res == "\n".join(expected)

def test_invalid_wires():
    """Test that an exception is raised if a wire in the wire
    ordering does not exist on the device"""
    dev = qml.device("default.qubit", wires=["a", -1, "q2"])

    @qml.beta.qnode(dev)
    def circuit():
        qml.Hadamard(wires=-1)
        qml.CNOT(wires=["a", "q2"])
        qml.RX(0.2, wires="a")
        return qml.expval(qml.PauliX(wires="q2"))

    with pytest.raises(ValueError, match="contains wires not contained on the device"):
        qml.draw(circuit, wire_order=["q2", 5])()

@pytest.mark.parametrize(
    "transform",
    [qml.gradients.param_shift(shift=0.2), functools.partial(qml.gradients.param_shift, shift=0.2)],
)
def test_draw_batch_transform(transform):
    """Test that drawing a batch transform works correctly"""
    dev = qml.device("default.qubit", wires=1)

    @transform
    @qml.beta.qnode(dev)
    def circuit(x):
        qml.Hadamard(wires=0)
        qml.RX(x, wires=0)
        return qml.expval(qml.PauliZ(wires=0))

    # the parameter-shift transform will create two circuits; one with x+0.2
    # and one with x-0.2.
    res = qml.draw(circuit)(0.6)
    expected = [" 0: ──H──RX(0.8)──┤ ⟨Z⟩ ", "", " 0: ──H──RX(0.4)──┤ ⟨Z⟩ ", ""]
    assert res == "\n".join(expected)

def test_direct_qnode_integration():
    """Test that a QNode renders correctly."""
    dev = qml.device("default.qubit", wires=2)

    @qml.beta.qnode(dev)
    def qfunc(a, w):
        qml.Hadamard(0)
        qml.CRX(a, wires=[0, 1])
        qml.Rot(w[0], w[1], w[2], wires=[1])
        qml.CRX(-a, wires=[0, 1])

        return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

    a, w = 2.3, [1.2, 3.2, 0.7]

    assert qml.draw(qfunc)(a, w) == (
        " 0: ──H──╭C────────────────────────────╭C─────────╭┤ ⟨Z ⊗ Z⟩ \n"
        + " 1: ─────╰RX(2.3)──Rot(1.2, 3.2, 0.7)──╰RX(-2.3)──╰┤ ⟨Z ⊗ Z⟩ \n"
    )

    assert qml.draw(qfunc, charset="ascii")(a, w) == (
        " 0: --H--+C----------------------------+C---------+| <Z @ Z> \n"
        + " 1: -----+RX(2.3)--Rot(1.2, 3.2, 0.7)--+RX(-2.3)--+| <Z @ Z> \n"
    )

def test_same_wire_multiple_measurements():
    """Test that drawing a QNode with multiple measurements on certain wires works correctly."""
    dev = qml.device("default.qubit", wires=4)

    @qml.beta.qnode(dev)
    def qnode(x, y):
        qml.RY(x, wires=0)
        qml.Hadamard(0)
        qml.RZ(y, wires=0)
        return [
            qml.expval(qml.PauliX(wires=[0]) @ qml.PauliX(wires=[1]) @ qml.PauliX(wires=[2])),
            qml.expval(qml.PauliX(wires=[0]) @ qml.PauliX(wires=[3])),
        ]

    expected = (
        " 0: ──RY(1)──H──RZ(2)──╭┤ ⟨X ⊗ X ⊗ X⟩ ╭┤ ⟨X ⊗ X⟩ \n"
        + " 1: ───────────────────├┤ ⟨X ⊗ X ⊗ X⟩ │┤         \n"
        + " 2: ───────────────────╰┤ ⟨X ⊗ X ⊗ X⟩ │┤         \n"
        + " 3: ────────────────────┤             ╰┤ ⟨X ⊗ X⟩ \n"
    )
    assert qml.draw(qnode)(1.0, 2.0) == expected

def test_same_wire_multiple_measurements_many_obs():
    """Test that drawing a QNode with multiple measurements on certain
    wires works correctly when there are more observables than the number of
    observables for any wire.
    """
    dev = qml.device("default.qubit", wires=4)

    @qml.beta.qnode(dev)
    def qnode(x, y):
        qml.RY(x, wires=0)
        qml.Hadamard(0)
        qml.RZ(y, wires=0)
        return [
            qml.expval(qml.PauliZ(0)),
            qml.expval(qml.PauliZ(1)),
            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
        ]

    expected = (
        " 0: ──RY(0.3)──H──RZ(0.2)──┤ ⟨Z⟩ ┤     ╭┤ ⟨Z ⊗ Z⟩ \n"
        + " 1: ───────────────────────┤     ┤ ⟨Z⟩ ╰┤ ⟨Z ⊗ Z⟩ \n"
    )
    assert qml.draw(qnode)(0.3, 0.2) == expected

class TestWireOrdering:
    """Tests for wire ordering functionality"""

    def test_default_ordering(self):
        """Test that the default wire ordering matches the device"""

        dev = qml.device("default.qubit", wires=["a", -1, "q2"])

        @qml.beta.qnode(dev)
        def circuit():
            qml.Hadamard(wires=-1)
            qml.CNOT(wires=["a", "q2"])
            qml.RX(0.2, wires="a")
            return qml.expval(qml.PauliX(wires="q2"))

        res = qml.draw(circuit)()
        expected = [
            "  a: ─────╭C──RX(0.2)──┤     ",
            " -1: ──H──│────────────┤     ",
            " q2: ─────╰X───────────┤ ⟨X⟩ \n",
        ]

        assert res == "\n".join(expected)

    def test_wire_reordering(self):
        """Test that wires are correctly reordered"""

        dev = qml.device("default.qubit", wires=["a", -1, "q2"])

        @qml.beta.qnode(dev)
        def circuit():
            qml.Hadamard(wires=-1)
            qml.CNOT(wires=["a", "q2"])
            qml.RX(0.2, wires="a")
            return qml.expval(qml.PauliX(wires="q2"))

        res = qml.draw(circuit, wire_order=["q2", "a", -1])()
        expected = [
            " q2: ──╭X───────────┤ ⟨X⟩ ",
            "  a: ──╰C──RX(0.2)──┤     ",
            " -1: ───H───────────┤     \n",
        ]

        assert res == "\n".join(expected)

    def test_include_empty_wires(self):
        """Test that empty wires are correctly included"""

        dev = qml.device("default.qubit", wires=[-1, "a", "q2", 0])

        @qml.beta.qnode(dev)
        def circuit():
            qml.Hadamard(wires=-1)
            qml.CNOT(wires=[-1, "q2"])
            return qml.expval(qml.PauliX(wires="q2"))

        res = qml.draw(circuit, show_all_wires=True)()
        expected = [
            " -1: ──H──╭C──┤     ",
            "  a: ─────│───┤     ",
            " q2: ─────╰X──┤ ⟨X⟩ ",
            "  0: ─────────┤     \n",
        ]

        assert res == "\n".join(expected)

    def test_show_all_wires_error(self):
        """Test that show_all_wires will raise an error if the provided wire
        order does not contain all wires on the device"""

        dev = qml.device("default.qubit", wires=[-1, "a", "q2", 0])

        @qml.beta.qnode(dev)
        def circuit():
            qml.Hadamard(wires=-1)
            qml.CNOT(wires=[-1, "q2"])
            return qml.expval(qml.PauliX(wires="q2"))

        with pytest.raises(ValueError, match="must contain all wires"):
            qml.draw(circuit, show_all_wires=True, wire_order=[-1, "a"])()

    def test_missing_wire(self):
        """Test that wires not specifically mentioned in the wire
        reordering are appended at the bottom of the circuit drawing"""

        dev = qml.device("default.qubit", wires=["a", -1, "q2"])

        @qml.beta.qnode(dev)
        def circuit():
            qml.Hadamard(wires=-1)
            qml.CNOT(wires=["a", "q2"])
            qml.RX(0.2, wires="a")
            return qml.expval(qml.PauliX(wires="q2"))

        # test one missing wire
        res = qml.draw(circuit, wire_order=["q2", "a"])()
        expected = [
            " q2: ──╭X───────────┤ ⟨X⟩ ",
            "  a: ──╰C──RX(0.2)──┤     ",
            " -1: ───H───────────┤     \n",
        ]

        assert res == "\n".join(expected)

        # test one missing wire
        res = qml.draw(circuit, wire_order=["q2", -1])()
        expected = [
            " q2: ─────╭X───────────┤ ⟨X⟩ ",
            " -1: ──H──│────────────┤     ",
            "  a: ─────╰C──RX(0.2)──┤     \n",
        ]

        assert res == "\n".join(expected)

        # test multiple missing wires
        res = qml.draw(circuit, wire_order=["q2"])()
        expected = [
            " q2: ─────╭X───────────┤ ⟨X⟩ ",
            " -1: ──H──│────────────┤     ",
            "  a: ─────╰C──RX(0.2)──┤     \n",
        ]

        assert res == "\n".join(expected)

    def test_invalid_wires(self):
        """Test that an exception is raised if a wire in the wire
        ordering does not exist on the device"""
        dev = qml.device("default.qubit", wires=["a", -1, "q2"])

        @qml.beta.qnode(dev)
        def circuit():
            qml.Hadamard(wires=-1)
            qml.CNOT(wires=["a", "q2"])
            qml.RX(0.2, wires="a")
            return qml.expval(qml.PauliX(wires="q2"))

        with pytest.raises(ValueError, match="contains wires not contained on the device"):
            res = qml.draw(circuit, wire_order=["q2", 5])()

    def test_no_ops_draws(self):
        """Test that a QNode with no operations still draws correctly"""
        dev = qml.device("default.qubit", wires=3)

        @qml.beta.qnode(dev)
        def qnode():
            return qml.expval(qml.PauliX(wires=[0]) @ qml.PauliX(wires=[1]) @ qml.PauliX(wires=[2]))

        res = qml.draw(qnode)()
        expected = [
            " 0: ──╭┤ ⟨X ⊗ X ⊗ X⟩ \n",
            " 1: ──├┤ ⟨X ⊗ X ⊗ X⟩ \n",
            " 2: ──╰┤ ⟨X ⊗ X ⊗ X⟩ \n",
        ]

        assert res == "".join(expected)


"""Demonstration of some basic structural features."""

# various methods for importing modules from other packages
from math import factorial
import datetime
import numpy as np
import sys

print("\nbasic structure allows use of functions and statements...")
print("time is ", datetime.datetime.now())
# we can use factorial rather than math.factorial()
print("5! = ", factorial(5))
array = np.array([1, 2, 3])     # use simpler/shorter module name
print(array)

# conditional statements
if np.size(array) > 0:
    print("array is not empty!")

print("\ncase sensitive so my_var and MY_VAR are different variables...")
my_var = 3
MY_VAR = 4
print("my_var at", id(my_var), "with value", my_var)
print("MY_VAR at", id(MY_VAR), "with value", MY_VAR)

def my_function(bob, susan):
    # indentation is key in Python
    try:
        if bob is None:
            if susan is None:
                print("anybody there?")
                return True    # the return is also conditional on bob and susan both None
            return False    # bob is None but susan isn't
        return False    # bob is not None, who cares about susan
    except ValueError:
        print("whoops")
        pass                # does nothing
    return False

# use line continuation after the operator
if None is not None and\
   None is None:
    print("how odd")

# you can use line continuations like this
my_variable = 1 + 4 +\
    6 + 10
print(my_variable)

# line continuation not required for example in lists or dictionaries
my_colors = ['red',
             'green']
print(my_colors)

my_dictionary = {
    'a': 10,
    'b': 23
}
print(my_dictionary)

# multiple assignment
print("\nyou can use multiple assignment syntax such as a, b,c = 1, 2, 3...")
a, b, c = 1, 2, 3
print(a, b, c)
d = e = f = 4
print(d)

print("\nuse help() to get help on functions and modules etc...")
help(dir)

# note __name__ is typically pronounced 'dunder name' - for double underscore

if __name__ == '__main__':      # typical entry point for python program from interpreter, not another module
    print("in main!")           # you would expect to call some function here and pass the return value to
    sys.exit()                  # ..sys.exit. Note a None value is equivalent to passing 0 (ok)


import random
import numpy as np
import torch

# This module is created to make our programs reproducible.
# To do this we control the seeds of all random sources.
# When will this break?
#     If the data provided to the program is not in the same order or the data file has been changed,
#     we cannot reproduce the results!!

# This function needs to be called before any function that uses randomness.
def reproduce(seed = None):
    if seed == None:
        seed = random.randint(0,2**32)

    print("Execution ID =", seed)
    print("Use the above Execution ID to reproduce the results.")

    random.seed(seed)
    np.random.seed(seed + 1)
    torch.manual_seed(seed + 2)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed + 3)

    return seed