"""
This module provides functions to convert
NetworkX graphs to and from other formats.

The preferred way of converting data to a NetworkX graph
is through the graph constuctor.  The constructor calls
the to_networkx_graph() function which attempts to guess the
input type and convert it automatically.

Examples
--------

Create a 10 node random graph from a numpy matrix

>>> import numpy
>>> a=numpy.reshape(numpy.random.random_integers(0,1,size=100),(10,10))
>>> D=nx.DiGraph(a)

or equivalently

>>> D=nx.to_networkx_graph(a,create_using=nx.DiGraph())

Create a graph with a single edge from a dictionary of dictionaries

>>> d={0: {1: 1}} # dict-of-dicts single edge (0,1)
>>> G=nx.Graph(d)

See Also
--------
nx_pygraphviz, nx_pydot

"""
__author__ = """\n""".join(['Aric Hagberg (hagberg@lanl.gov)',
                           'Pieter Swart (swart@lanl.gov)',
                           'Dan Schult(dschult@colgate.edu)'])
#    Copyright (C) 2006-2011 by
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.

import warnings
import networkx as nx

__all__ = ['to_networkx_graph',
           'from_dict_of_dicts', 'to_dict_of_dicts',
           'from_dict_of_lists', 'to_dict_of_lists',
           'from_edgelist', 'to_edgelist',
           'from_numpy_matrix', 'to_numpy_matrix',
           'to_numpy_recarray',
           'from_scipy_sparse_matrix', 'to_scipy_sparse_matrix']

def _prep_create_using(create_using):
    """Return a graph object ready to be populated.

    If create_using is None return the default (just networkx.Graph())
    If create_using.clear() works, assume it returns a graph object.
    Otherwise raise an exception because create_using is not a networkx graph.

    """
    if create_using is None:
        G=nx.Graph()
    else:
        G=create_using
        try:
            G.clear()
        except:
            raise TypeError("Input graph is not a networkx graph type")
    return G

def to_networkx_graph(data,create_using=None,multigraph_input=False):
    """Make a NetworkX graph from a known data structure.

    The preferred way to call this is automatically
    from the class constructor

    >>> d={0: {1: {'weight':1}}} # dict-of-dicts single edge (0,1)
    >>> G=nx.Graph(d)

    instead of the equivalent

    >>> G=nx.from_dict_of_dicts(d)

    Parameters
    ----------
    data : a object to be converted
       Current known types are:
         any NetworkX graph
         dict-of-dicts
         dist-of-lists
         list of edges
         numpy matrix
         numpy ndarray
         scipy sparse matrix
         pygraphviz agraph

    create_using : NetworkX graph
       Use specified graph for result.  Otherwise a new graph is created.

    multigraph_input : bool (default False)
      If True and  data is a dict_of_dicts,
      try to create a multigraph assuming dict_of_dict_of_lists.
      If data and create_using are both multigraphs then create
      a multigraph from a multigraph.

    """
    # NX graph
    if hasattr(data,"adj"):
        try:
            result= from_dict_of_dicts(data.adj,\
                    create_using=create_using,\
                    multigraph_input=data.is_multigraph())
            if hasattr(data,'graph') and isinstance(data.graph,dict):
                result.graph=data.graph.copy()
            if hasattr(data,'node') and isinstance(data.node,dict):
                result.node=dict( (n,dd.copy()) for n,dd in data.node.items() )
            return result
        except:
            raise nx.NetworkXError("Input is not a correct NetworkX graph.")

    # pygraphviz  agraph
    if hasattr(data,"is_strict"):
        try:
            return nx.from_agraph(data,create_using=create_using)
        except:
            raise nx.NetworkXError("Input is not a correct pygraphviz graph.")

    # dict of dicts/lists
    if isinstance(data,dict):
        try:
            return from_dict_of_dicts(data,create_using=create_using,\
                    multigraph_input=multigraph_input)
        except:
            try:
                return from_dict_of_lists(data,create_using=create_using)
            except:
                raise TypeError("Input is not known type.")

    # list or generator of edges
    if (isinstance(data,list)
        or hasattr(data,'next')
        or hasattr(data, '__next__')):
        try:
            return from_edgelist(data,create_using=create_using)
        except:
            raise nx.NetworkXError("Input is not a valid edge list")

    # numpy matrix or ndarray
    try:
        import numpy
        if isinstance(data,numpy.matrix) or \
               isinstance(data,numpy.ndarray):
            try:
                return from_numpy_matrix(data,create_using=create_using)
            except:
                raise nx.NetworkXError(\
                  "Input is not a correct numpy matrix or array.")
    except ImportError:
        warnings.warn('numpy not found, skipping conversion test.',
                      ImportWarning)

    # scipy sparse matrix - any format
    try:
        import scipy
        if hasattr(data,"format"):
            try:
                return from_scipy_sparse_matrix(data,create_using=create_using)
            except:
                raise nx.NetworkXError(\
                      "Input is not a correct scipy sparse matrix type.")
    except ImportError:
        warnings.warn('scipy not found, skipping conversion test.',
                      ImportWarning)

    raise nx.NetworkXError(\
          "Input is not a known data type for conversion.")

    return

def convert_to_undirected(G):
    """Return a new undirected representation of the graph G.

    """
    return G.to_undirected()

def convert_to_directed(G):
    """Return a new directed representation of the graph G.

    """
    return G.to_directed()

def to_dict_of_lists(G,nodelist=None):
    """Return adjacency representation of graph as a dictionary of lists.

    Parameters
    ----------
    G : graph
       A NetworkX graph

    nodelist : list
       Use only nodes specified in nodelist

    Notes
    -----
    Completely ignores edge data for MultiGraph and MultiDiGraph.

    """
    if nodelist is None:
        nodelist=G

    d = {}
    for n in nodelist:
        d[n]=[nbr for nbr in G.neighbors(n) if nbr in nodelist]
    return d

def from_dict_of_lists(d,create_using=None):
    """Return a graph from a dictionary of lists.

    Parameters
    ----------
    d : dictionary of lists
      A dictionary of lists adjacency representation.

    create_using : NetworkX graph
       Use specified graph for result.  Otherwise a new graph is created.

    Examples
    --------
    >>> dol= {0:[1]} # single edge (0,1)
    >>> G=nx.from_dict_of_lists(dol)

    or
    >>> G=nx.Graph(dol) # use Graph constructor

    """
    G=_prep_create_using(create_using)
    G.add_nodes_from(d)
    if G.is_multigraph() and not G.is_directed():
        # a dict_of_lists can't show multiedges.  BUT for undirected graphs,
        # each edge shows up twice in the dict_of_lists.
        # So we need to treat this case separately.
        seen={}
        for node,nbrlist in d.items():
            for nbr in nbrlist:
                if nbr not in seen:
                    G.add_edge(node,nbr)
            seen[node]=1  # don't allow reverse edge to show up
    else:
        G.add_edges_from( ((node,nbr) for node,nbrlist in d.items()
                           for nbr in nbrlist) )
    return G

def to_dict_of_dicts(G,nodelist=None,edge_data=None):
    """Return adjacency representation of graph as a dictionary of dictionaries.

    Parameters
    ----------
    G : graph
       A NetworkX graph

    nodelist : list
       Use only nodes specified in nodelist

    edge_data : list, optional
       If provided,  the value of the dictionary will be
       set to edge_data for all edges.  This is useful to make
       an adjacency matrix type representation with 1 as the edge data.
       If edgedata is None, the edgedata in G is used to fill the values.
       If G is a multigraph, the edgedata is a dict for each pair (u,v).

    """
    dod={}
    if nodelist is None:
        if edge_data is None:
            for u,nbrdict in G.adjacency_iter():
                dod[u]=nbrdict.copy()
        else: # edge_data is not None
            for u,nbrdict in G.adjacency_iter():
                dod[u]=dod.fromkeys(nbrdict, edge_data)
    else: # nodelist is not None
        if edge_data is None:
            for u in nodelist:
                dod[u]={}
                for v,data in ((v,data) for v,data in G[u].items() if v in nodelist):
                    dod[u][v]=data
        else: # nodelist and edge_data are not None
            for u in nodelist:
                dod[u]={}
                for v in ( v for v in G[u] if v in nodelist):
                    dod[u][v]=edge_data
    return dod

def from_dict_of_dicts(d,create_using=None,multigraph_input=False):
    """Return a graph from a dictionary of dictionaries.

    Parameters
    ----------
    d : dictionary of dictionaries
      A dictionary of dictionaries adjacency representation.

    create_using : NetworkX graph
       Use specified graph for result.  Otherwise a new graph is created.

    multigraph_input : bool (default False)
       When True, the values of the inner dict are assumed
       to be containers of edge data for multiple edges.
       Otherwise this routine assumes the edge data are singletons.

    Examples
    --------
    >>> dod= {0: {1:{'weight':1}}} # single edge (0,1)
    >>> G=nx.from_dict_of_dicts(dod)

    or
    >>> G=nx.Graph(dod) # use Graph constructor

    """
    G=_prep_create_using(create_using)
    G.add_nodes_from(d)
    # is dict a MultiGraph or MultiDiGraph?
    if multigraph_input:
        # make a copy of the list of edge data (but not the edge data)
        if G.is_directed():
            if G.is_multigraph():
                G.add_edges_from( (u,v,key,data)
                                  for u,nbrs in d.items()
                                  for v,datadict in nbrs.items()
                                  for key,data in datadict.items()
                                )
            else:
                G.add_edges_from( (u,v,data)
                                  for u,nbrs in d.items()
                                  for v,datadict in nbrs.items()
                                  for key,data in datadict.items()
                                )
        else: # Undirected
            if G.is_multigraph():
                seen=set()   # don't add both directions of undirected graph
                for u,nbrs in d.items():
                    for v,datadict in nbrs.items():
                        if (u,v) not in seen:
                            G.add_edges_from( (u,v,key,data)
                                               for key,data in datadict.items()
                                              )
                            seen.add((v,u))
            else:
                seen=set()   # don't add both directions of undirected graph
                for u,nbrs in d.items():
                    for v,datadict in nbrs.items():
                        if (u,v) not in seen:
                            G.add_edges_from( (u,v,data)
                                        for key,data in datadict.items() )
                            seen.add((v,u))

    else: # not a multigraph to multigraph transfer
        if G.is_multigraph() and not G.is_directed():
            # d can have both representations u-v, v-u in dict.  Only add one.
            # We don't need this check for digraphs since we add both directions,
            # or for Graph() since it is done implicitly (parallel edges not allowed)
            seen=set()
            for u,nbrs in d.items():
                for v,data in nbrs.items():
                    if (u,v) not in seen:
                        G.add_edge(u,v,attr_dict=data)
                    seen.add((v,u))
        else:
            G.add_edges_from( ( (u,v,data)
                                for u,nbrs in d.items()
                                for v,data in nbrs.items()) )
    return G

def to_edgelist(G,nodelist=None):
    """Return a list of edges in the graph.

    Parameters
    ----------
    G : graph
       A NetworkX graph

    nodelist : list
       Use only nodes specified in nodelist

    """
    if nodelist is None:
        return G.edges(data=True)
    else:
        return G.edges(nodelist,data=True)

def from_edgelist(edgelist,create_using=None):
    """Return a graph from a list of edges.

    Parameters
    ----------
    edgelist : list or iterator
      Edge tuples

    create_using : NetworkX graph
       Use specified graph for result.  Otherwise a new graph is created.

    Examples
    --------
    >>> edgelist= [(0,1)] # single edge (0,1)
    >>> G=nx.from_edgelist(edgelist)

    or
    >>> G=nx.Graph(edgelist) # use Graph constructor

    """
    G=_prep_create_using(create_using)
    G.add_edges_from(edgelist)
    return G

def to_numpy_matrix(G, nodelist=None, dtype=None, order=None,
                    multigraph_weight=sum, weight='weight'):
    """Return the graph adjacency matrix as a NumPy matrix.

    Parameters
    ----------
    G : graph
        The NetworkX graph used to construct the NumPy matrix.

    nodelist : list, optional
       The rows and columns are ordered according to the nodes in `nodelist`.
       If `nodelist` is None, then the ordering is produced by G.nodes().

    dtype : NumPy data type, optional
        A valid single NumPy data type used to initialize the array.
        This must be a simple type such as int or numpy.float64 and
        not a compound data type (see to_numpy_recarray)
        If None, then the NumPy default is used.

    order : {'C', 'F'}, optional
        Whether to store multidimensional data in C- or Fortran-contiguous
        (row- or column-wise) order in memory. If None, then the NumPy default
        is used.

    multigraph_weight : {sum, min, max}, optional
        An operator that determines how weights in multigraphs are handled.
        The default is to sum the weights of the multiple edges.

    weight : string or None   optional (default='weight')
        The edge attribute that holds the numerical value used for
        the edge weight.  If None then all edge weights are 1.

    Returns
    -------
    M : NumPy matrix
       Graph adjacency matrix.

    See Also
    --------
    to_numpy_recarray, from_numpy_matrix

    Notes
    -----
    The matrix entries are assigned with weight edge attribute. When
    an edge does not have the weight attribute, the value of the entry is 1.
    For multiple edges, the values of the entries are the sums of the edge
    attributes for each edge.

    When `nodelist` does not contain every node in `G`, the matrix is built
    from the subgraph of `G` that is induced by the nodes in `nodelist`.

    Examples
    --------
    >>> G = nx.MultiDiGraph()
    >>> G.add_edge(0,1,weight=2)
    >>> G.add_edge(1,0)
    >>> G.add_edge(2,2,weight=3)
    >>> G.add_edge(2,2)
    >>> nx.to_numpy_matrix(G, nodelist=[0,1,2])
    matrix([[ 0.,  2.,  0.],
            [ 1.,  0.,  0.],
            [ 0.,  0.,  4.]])

    """
    try:
        import numpy as np
    except ImportError:
        raise ImportError(\
          "to_numpy_matrix() requires numpy: http://scipy.org/ ")

    if nodelist is None:
        nodelist = G.nodes()

    nodeset = set(nodelist)
    if len(nodelist) != len(nodeset):
        msg = "Ambiguous ordering: `nodelist` contained duplicates."
        raise nx.NetworkXError(msg)

    nlen=len(nodelist)
    undirected = not G.is_directed()
    index=dict(zip(nodelist,range(nlen)))

    if G.is_multigraph():
        # Handle MultiGraphs and MultiDiGraphs
        # array of nan' to start with, any leftover nans will be converted to 0
        # nans are used so we can use sum, min, max for multigraphs
        M = np.zeros((nlen,nlen), dtype=dtype, order=order)+np.nan
        # use numpy nan-aware operations
        operator={sum:np.nansum, min:np.nanmin, max:np.nanmax}
        try:
            op=operator[multigraph_weight]
        except:
            raise ValueError('multigraph_weight must be sum, min, or max')

        for u,v,attrs in G.edges_iter(data=True):
            if (u in nodeset) and (v in nodeset):
                i,j = index[u],index[v]
                e_weight = attrs.get(weight, 1)
                M[i,j] = op([e_weight,M[i,j]])
                if undirected:
                    M[j,i] = M[i,j]
        # convert any nans to zeros
        M = np.asmatrix(np.nan_to_num(M))
    else:
        # Graph or DiGraph, this is much faster than above
        M = np.zeros((nlen,nlen), dtype=dtype, order=order)
        for u,nbrdict in G.adjacency_iter():
            for v,d in nbrdict.items():
                try:
                    M[index[u],index[v]]=d.get(weight,1)
                except KeyError:
                    pass
        M = np.asmatrix(M)
    return M

def from_numpy_matrix(A,create_using=None):
    """Return a graph from numpy matrix.

    The numpy matrix is interpreted as an adjacency matrix for the graph.

    Parameters
    ----------
    A : numpy matrix
      An adjacency matrix representation of a graph

    create_using : NetworkX graph
       Use specified graph for result.  The default is Graph()

    Notes
    -----
    If the numpy matrix has a single data type for each matrix entry it
    will be converted to an appropriate Python data type.

    If the numpy matrix has a user-specified compound data type the names
    of the data fields will be used as attribute keys in the resulting
    NetworkX graph.

    See Also
    --------
    to_numpy_matrix, to_numpy_recarray

    Examples
    --------
    Simple integer weights on edges:

    >>> import numpy
    >>> A=numpy.matrix([[1,1],[2,1]])
    >>> G=nx.from_numpy_matrix(A)

    User defined compound data type on edges:

    >>> import numpy
    >>> dt=[('weight',float),('cost',int)]
    >>> A=numpy.matrix([[(1.0,2)]],dtype=dt)
    >>> G=nx.from_numpy_matrix(A)
    >>> G.edges(data=True)
    [(0, 0, {'cost': 2, 'weight': 1.0})]
    """
    kind_to_python_type={'f':float,
                         'i':int,
                         'u':int,
                         'b':bool,
                         'c':complex,
                         'S':str,
                         'V':'void'}

    try: # Python 3.x
        blurb = chr(1245) # just to trigger the exception
        kind_to_python_type['U']=str
    except ValueError: # Python 2.6+
        kind_to_python_type['U']=unicode

    # This should never fail if you have created a numpy matrix with numpy...
    try:
        import numpy as np
    except ImportError:
        raise ImportError(\
          "from_numpy_matrix() requires numpy: http://scipy.org/ ")

    G=_prep_create_using(create_using)
    n,m=A.shape
    if n!=m:
        raise nx.NetworkXError("Adjacency matrix is not square.",
                               "nx,ny=%s"%(A.shape,))
    dt=A.dtype
    try:
        python_type=kind_to_python_type[dt.kind]
    except:
        raise TypeError("Unknown numpy data type: %s"%dt)

    # make sure we get isolated nodes
    G.add_nodes_from(range(n))
    # get a list of edges
    x,y=np.asarray(A).nonzero()

    # handle numpy constructed data type
    if python_type is 'void':
        fields=sorted([(offset,dtype,name) for name,(dtype,offset) in
                       A.dtype.fields.items()])
        for (u,v) in zip(x,y):
            attr={}
            for (offset,dtype,name),val in zip(fields,A[u,v]):
                attr[name]=kind_to_python_type[dtype.kind](val)
            G.add_edge(u,v,attr)
    else: # basic data type
        G.add_edges_from( ((u,v,{'weight':python_type(A[u,v])})
                           for (u,v) in zip(x,y)) )
    return G

def to_numpy_recarray(G,nodelist=None,
                      dtype=[('weight',float)],
                      order=None):
    """Return the graph adjacency matrix as a NumPy recarray.

    Parameters
    ----------
    G : graph
        The NetworkX graph used to construct the NumPy matrix.

    nodelist : list, optional
       The rows and columns are ordered according to the nodes in `nodelist`.
       If `nodelist` is None, then the ordering is produced by G.nodes().

    dtype : NumPy data-type, optional
        A valid NumPy named dtype used to initialize the NumPy recarray.
        The data type names are assumed to be keys in the graph edge attribute
        dictionary.

    order : {'C', 'F'}, optional
        Whether to store multidimensional data in C- or Fortran-contiguous
        (row- or column-wise) order in memory. If None, then the NumPy default
        is used.

    Returns
    -------
    M : NumPy recarray
       The graph with specified edge data as a Numpy recarray

    Notes
    -----
    When `nodelist` does not contain every node in `G`, the matrix is built
    from the subgraph of `G` that is induced by the nodes in `nodelist`.

    Examples
    --------
    >>> G = nx.Graph()
    >>> G.add_edge(1,2,weight=7.0,cost=5)
    >>> A=nx.to_numpy_recarray(G,dtype=[('weight',float),('cost',int)])
    >>> print(A.weight)
    [[ 0.  7.]
     [ 7.  0.]]
    >>> print(A.cost)
    [[0 5]
     [5 0]]
    """
    try:
        import numpy as np
    except ImportError:
        raise ImportError(\
          "to_numpy_matrix() requires numpy: http://scipy.org/ ")

    if G.is_multigraph():
        raise nx.NetworkXError("Not implemented for multigraphs.")

    if nodelist is None:
        nodelist = G.nodes()

    nodeset = set(nodelist)
    if len(nodelist) != len(nodeset):
        msg = "Ambiguous ordering: `nodelist` contained duplicates."
        raise nx.NetworkXError(msg)

    nlen=len(nodelist)
    undirected = not G.is_directed()
    index=dict(zip(nodelist,range(nlen)))
    M = np.zeros((nlen,nlen), dtype=dtype, order=order)

    names=M.dtype.names
    for u,v,attrs in G.edges_iter(data=True):
        if (u in nodeset) and (v in nodeset):
            i,j = index[u],index[v]
            values=tuple([attrs[n] for n in names])
            M[i,j] = values
            if undirected:
                M[j,i] = M[i,j]

    return M.view(np.recarray)

def to_scipy_sparse_matrix(G, nodelist=None, dtype=None,
                           weight='weight', format='csr'):
    """Return the graph adjacency matrix as a SciPy sparse matrix.

    Parameters
    ----------
    G : graph
        The NetworkX graph used to construct the NumPy matrix.

    nodelist : list, optional
       The rows and columns are ordered according to the nodes in `nodelist`.
       If `nodelist` is None, then the ordering is produced by G.nodes().

    dtype : NumPy data-type, optional
        A valid NumPy dtype used to initialize the array. If None, then the
        NumPy default is used.

    weight : string or None   optional (default='weight')
        The edge attribute that holds the numerical value used for
        the edge weight.  If None then all edge weights are 1.

    format : str in {'bsr', 'csr', 'csc', 'coo', 'lil', 'dia', 'dok'}
        The type of the matrix to be returned (default 'csr').  For
        some algorithms different implementations of sparse matrices
        can perform better.  See [1]_ for details.

    Returns
    -------
    M : SciPy sparse matrix
       Graph adjacency matrix.

    Notes
    -----
    The matrix entries are populated using the edge attribute held in
    parameter weight. When an edge does not have that attribute, the
    value of the entry is 1.

    For multiple edges the matrix values are the sums of the edge weights.

    When `nodelist` does not contain every node in `G`, the matrix is built
    from the subgraph of `G` that is induced by the nodes in `nodelist`.

    Uses lil_matrix format. To convert to other formats specify the
    format= keyword.

    Examples
    --------
    >>> G = nx.MultiDiGraph()
    >>> G.add_edge(0,1,weight=2)
    >>> G.add_edge(1,0)
    >>> G.add_edge(2,2,weight=3)
    >>> G.add_edge(2,2)
    >>> S = nx.to_scipy_sparse_matrix(G, nodelist=[0,1,2])
    >>> S.todense()
    matrix([[ 0.,  2.,  0.],
            [ 1.,  0.,  0.],
            [ 0.,  0.,  4.]])

    References
    ----------
    .. [1] Scipy Dev. References,
       "Sparse Matrices"
       http://docs.scipy.org/doc/scipy/reference/sparse.html
    """
    try:
        from scipy import sparse
    except ImportError:
        raise ImportError(\
          "to_scipy_sparse_matrix() requires scipy: http://scipy.org/ ")

    if nodelist is None:
        nodelist = G.nodes()

    nodeset = set(nodelist)
    if len(nodelist) != len(nodeset):
        msg = "Ambiguous ordering: `nodelist` contained duplicates."
        raise nx.NetworkXError(msg)

    nlen=len(nodelist)
    undirected = not G.is_directed()
    index=dict(zip(nodelist,range(nlen)))
    M = sparse.lil_matrix((nlen,nlen), dtype=dtype)

    for u,v,attrs in G.edges_iter(data=True):
        if (u in nodeset) and (v in nodeset):
            i,j = index[u],index[v]
            M[i,j] += attrs.get(weight, 1)
            if undirected:
                M[j,i] = M[i,j]
    try:
        return M.asformat(format)
    except AttributeError:
        raise nx.NetworkXError("Unknown sparse matrix format: %s"%format)

def from_scipy_sparse_matrix(A,create_using=None):
    """Return a graph from scipy sparse matrix adjacency list.

    Parameters
    ----------
    A : scipy sparse matrix
      An adjacency matrix representation of a graph

    create_using : NetworkX graph
       Use specified graph for result.  The default is Graph()

    Examples
    --------
    >>> import scipy.sparse
    >>> A=scipy.sparse.eye(2,2,1)
    >>> G=nx.from_scipy_sparse_matrix(A)

    """
    G=_prep_create_using(create_using)

    # convert all formats to lil - not the most efficient way
    AA=A.tolil()
    n,m=AA.shape

    if n!=m:
        raise nx.NetworkXError(\
              "Adjacency matrix is not square. nx,ny=%s"%(A.shape,))
    G.add_nodes_from(range(n)) # make sure we get isolated nodes

    for i,row in enumerate(AA.rows):
        for pos,j in enumerate(row):
            G.add_edge(i,j,**{'weight':AA.data[i][pos]})
    return G

# fixture for nose tests
def setup_module(module):
    from nose import SkipTest
    try:
        import numpy
    except:
        raise SkipTest("NumPy not available")
    try:
        import scipy
    except:
        raise SkipTest("SciPy not available")

import numpy,from numpy,import scipy,from scipy,import networkx


import numpy as np
import pytest

import pandas.util._test_decorators as td

import pandas as pd
import pandas._testing as tm

# TODO(ArrayManager) concat with reindexing
pytestmark = td.skip_array_manager_not_yet_implemented

def test_error():
    df = pd.DataFrame(
        {"A": pd.Series([[0, 1, 2], np.nan, [], (3, 4)], index=list("abcd")), "B": 1}
    )
    with pytest.raises(ValueError, match="column must be a scalar"):
        df.explode(list("AA"))

    df.columns = list("AA")
    with pytest.raises(ValueError, match="columns must be unique"):
        df.explode("A")

def test_basic():
    df = pd.DataFrame(
        {"A": pd.Series([[0, 1, 2], np.nan, [], (3, 4)], index=list("abcd")), "B": 1}
    )
    result = df.explode("A")
    expected = pd.DataFrame(
        {
            "A": pd.Series(
                [0, 1, 2, np.nan, np.nan, 3, 4], index=list("aaabcdd"), dtype=object
            ),
            "B": 1,
        }
    )
    tm.assert_frame_equal(result, expected)

def test_multi_index_rows():
    df = pd.DataFrame(
        {"A": np.array([[0, 1, 2], np.nan, [], (3, 4)], dtype=object), "B": 1},
        index=pd.MultiIndex.from_tuples([("a", 1), ("a", 2), ("b", 1), ("b", 2)]),
    )

    result = df.explode("A")
    expected = pd.DataFrame(
        {
            "A": pd.Series(
                [0, 1, 2, np.nan, np.nan, 3, 4],
                index=pd.MultiIndex.from_tuples(
                    [
                        ("a", 1),
                        ("a", 1),
                        ("a", 1),
                        ("a", 2),
                        ("b", 1),
                        ("b", 2),
                        ("b", 2),
                    ]
                ),
                dtype=object,
            ),
            "B": 1,
        }
    )
    tm.assert_frame_equal(result, expected)

def test_multi_index_columns():
    df = pd.DataFrame(
        {("A", 1): np.array([[0, 1, 2], np.nan, [], (3, 4)], dtype=object), ("A", 2): 1}
    )

    result = df.explode(("A", 1))
    expected = pd.DataFrame(
        {
            ("A", 1): pd.Series(
                [0, 1, 2, np.nan, np.nan, 3, 4],
                index=pd.Index([0, 0, 0, 1, 2, 3, 3]),
                dtype=object,
            ),
            ("A", 2): 1,
        }
    )
    tm.assert_frame_equal(result, expected)

def test_usecase():
    # explode a single column
    # gh-10511
    df = pd.DataFrame(
        [[11, range(5), 10], [22, range(3), 20]], columns=list("ABC")
    ).set_index("C")
    result = df.explode("B")

    expected = pd.DataFrame(
        {
            "A": [11, 11, 11, 11, 11, 22, 22, 22],
            "B": np.array([0, 1, 2, 3, 4, 0, 1, 2], dtype=object),
            "C": [10, 10, 10, 10, 10, 20, 20, 20],
        },
        columns=list("ABC"),
    ).set_index("C")

    tm.assert_frame_equal(result, expected)

    # gh-8517
    df = pd.DataFrame(
        [["2014-01-01", "Alice", "A B"], ["2014-01-02", "Bob", "C D"]],
        columns=["dt", "name", "text"],
    )
    result = df.assign(text=df.text.str.split(" ")).explode("text")
    expected = pd.DataFrame(
        [
            ["2014-01-01", "Alice", "A"],
            ["2014-01-01", "Alice", "B"],
            ["2014-01-02", "Bob", "C"],
            ["2014-01-02", "Bob", "D"],
        ],
        columns=["dt", "name", "text"],
        index=[0, 0, 1, 1],
    )
    tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize(
    "input_dict, input_index, expected_dict, expected_index",
    [
        (
            {"col1": [[1, 2], [3, 4]], "col2": ["foo", "bar"]},
            [0, 0],
            {"col1": [1, 2, 3, 4], "col2": ["foo", "foo", "bar", "bar"]},
            [0, 0, 0, 0],
        ),
        (
            {"col1": [[1, 2], [3, 4]], "col2": ["foo", "bar"]},
            pd.Index([0, 0], name="my_index"),
            {"col1": [1, 2, 3, 4], "col2": ["foo", "foo", "bar", "bar"]},
            pd.Index([0, 0, 0, 0], name="my_index"),
        ),
        (
            {"col1": [[1, 2], [3, 4]], "col2": ["foo", "bar"]},
            pd.MultiIndex.from_arrays(
                [[0, 0], [1, 1]], names=["my_first_index", "my_second_index"]
            ),
            {"col1": [1, 2, 3, 4], "col2": ["foo", "foo", "bar", "bar"]},
            pd.MultiIndex.from_arrays(
                [[0, 0, 0, 0], [1, 1, 1, 1]],
                names=["my_first_index", "my_second_index"],
            ),
        ),
        (
            {"col1": [[1, 2], [3, 4]], "col2": ["foo", "bar"]},
            pd.MultiIndex.from_arrays([[0, 0], [1, 1]], names=["my_index", None]),
            {"col1": [1, 2, 3, 4], "col2": ["foo", "foo", "bar", "bar"]},
            pd.MultiIndex.from_arrays(
                [[0, 0, 0, 0], [1, 1, 1, 1]], names=["my_index", None]
            ),
        ),
    ],
)
def test_duplicate_index(input_dict, input_index, expected_dict, expected_index):
    # GH 28005
    df = pd.DataFrame(input_dict, index=input_index)
    result = df.explode("col1")
    expected = pd.DataFrame(expected_dict, index=expected_index, dtype=object)
    tm.assert_frame_equal(result, expected)

def test_ignore_index():
    # GH 34932
    df = pd.DataFrame({"id": range(0, 20, 10), "values": [list("ab"), list("cd")]})
    result = df.explode("values", ignore_index=True)
    expected = pd.DataFrame(
        {"id": [0, 0, 10, 10], "values": list("abcd")}, index=[0, 1, 2, 3]
    )
    tm.assert_frame_equal(result, expected)

def test_explode_sets():
    # https://github.com/pandas-dev/pandas/issues/35614
    df = pd.DataFrame({"a": [{"x", "y"}], "b": [1]}, index=[1])
    result = df.explode(column="a").sort_values(by="a")
    expected = pd.DataFrame({"a": ["x", "y"], "b": [1, 1]}, index=[1, 1])
    tm.assert_frame_equal(result, expected)


import numpy as np #the Numpy library
import matplotlib.pyplot as plt #Matplotlib's pyplot

import sys #gives access to C-like sys library
import os #gives accesss to operating system

print(sys.argv) #prints any command line arguments, incl. program name
print(os.getcwd())


from jupyter.data import get_fremont_bicycle_data
import pandas as pd
import numpy as np
##mm error
## data = get_fremont_bicycle_data()
def test_fremont_data():
    URL = 'https://data.seattle.gov/api/views/65db-xm6k/rows.csv?accessType=DOWNLOAD'
    FILENAME = 'fremont_bicycle.csv'
    data = get_fremont_bicycle_data(FILENAME, URL)
    assert all(data.columns == ['Total', 'East', 'West'])
    assert isinstance(data.index, pd.DatetimeIndex)
    assert len(np.unique(data.index.time)) == 24


import numpy as np
arr1 = np.full (2, True)
print("Full 1D Array")
print(arr1)
print(arr1.ndim)

#[True True]
# arr1.ndim prints dimensions of array i.e. 1


# -*- coding: utf-8 -*-
########################################################################
#
# multidimension tests for blaze.carray.
# based on carray tests, adapted to blaze and nosetest
#
########################################################################

import sys
import struct

import numpy as np
from numpy.testing import assert_array_equal, assert_array_almost_equal
import blaze.carray as ca
from blaze.carray.carrayExtension import chunk
from blaze.carray.tests import common
from common import MayBeDiskTest
import unittest

class constructorTest(MayBeDiskTest, unittest.TestCase):

    open = False

    def test00a(self):
        """Testing `carray` reshape"""
        a = np.arange(16).reshape((2,2,4))
        b = ca.arange(16, rootdir=self.rootdir).reshape((2,2,4))
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test00b(self):
        """Testing `carray` reshape (large shape)"""
        a = np.arange(16000).reshape((20,20,40))
        b = ca.arange(16000, rootdir=self.rootdir).reshape((20,20,40))
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test01a(self):
        """Testing `zeros` constructor (I)"""
        a = np.zeros((2,2,4), dtype='i4')
        b = ca.zeros((2,2,4), dtype='i4', rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test01b(self):
        """Testing `zeros` constructor (II)"""
        a = np.zeros(2, dtype='(2,4)i4')
        b = ca.zeros(2, dtype='(2,4)i4', rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test01c(self):
        """Testing `zeros` constructor (III)"""
        a = np.zeros((2,2), dtype='(4,)i4')
        b = ca.zeros((2,2), dtype='(4,)i4', rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test02(self):
        """Testing `ones` constructor"""
        a = np.ones((2,2), dtype='(4,)i4')
        b = ca.ones((2,2), dtype='(4,)i4', rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test03a(self):
        """Testing `fill` constructor (scalar default)"""
        a = np.ones((2,200), dtype='(4,)i4')*3
        b = ca.fill((2,200), 3, dtype='(4,)i4', rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test03b(self):
        """Testing `fill` constructor (array default)"""
        a = np.ones((2,2), dtype='(4,)i4')*3
        b = ca.fill((2,2), [3,3,3,3], dtype='(4,)i4', rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test04(self):
        """Testing `fill` constructor with open and resize (array default)"""
        a = np.ones((3,200), dtype='(4,)i4')*3
        b = ca.fill((2,200), [3,3,3,3], dtype='(4,)i4', rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        c = np.ones((1,200), dtype='(4,)i4')*3
        b.append(c)
        #print "b->", `b`, len(b), b[1]
        assert_array_equal(a, b, "Arrays are not equal")

    def test05(self):
        """Testing `fill` constructor with open and resize (nchunks>1)"""
        a = np.ones((3,2000), dtype='(4,)i4')*3
        b = ca.fill((2,2000), [3,3,3,3], dtype='(4,)i4', rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        c = np.ones((1,2000), dtype='(4,)i4')*3
        b.append(c)
        #print "b->", `b`
        # We need to use the b[:] here to overcome a problem with the
        # assert_array_equal() function
        assert_array_equal(a, b[:], "Arrays are not equal")

class constructorDiskTest(constructorTest, unittest.TestCase):
    disk = True
    open = False

class constructorOpenTest(constructorTest, unittest.TestCase):
    disk = True
    open = True

class getitemTest(MayBeDiskTest, unittest.TestCase):

    open = False

    def test00a(self):
        """Testing `__getitem()__` method with only a start (scalar)"""
        a = np.ones((2,3), dtype="i4")*3
        b = ca.fill((2,3), 3, dtype="i4", rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = 1
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test00b(self):
        """Testing `__getitem()__` method with only a start (slice)"""
        a = np.ones((27,2700), dtype="i4")*3
        b = ca.fill((27,2700), 3, dtype="i4", rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = slice(1)
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test01(self):
        """Testing `__getitem()__` method with a start and a stop"""
        a = np.ones((5,2), dtype="i4")*3
        b = ca.fill((5,2), 3, dtype="i4", rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = slice(1,4)
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test02(self):
        """Testing `__getitem()__` method with a start, stop, step"""
        a = np.ones((10,2), dtype="i4")*3
        b = ca.fill((10,2), 3, dtype="i4", rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = slice(1,9,2)
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test03a(self):
        """Testing `__getitem()__` method with several slices (I)"""
        a = np.arange(12).reshape((4,3))
        b = ca.carray(a, rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = (slice(1,3,1), slice(1,4,2))
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test03b(self):
        """Testing `__getitem()__` method with several slices (II)"""
        a = np.arange(24*1000).reshape((4*1000,3,2))
        b = ca.carray(a, rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = (slice(1,3,2), slice(1,4,2), slice(None))
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test03c(self):
        """Testing `__getitem()__` method with several slices (III)"""
        a = np.arange(120*1000).reshape((5*1000,4,3,2))
        b = ca.carray(a, rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = (slice(None,None,3), slice(1,3,2), slice(1,4,2))
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test04a(self):
        """Testing `__getitem()__` method with shape reduction (I)"""
        a = np.arange(12000).reshape((40,300))
        b = ca.carray(a, rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = (1,1)
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test04b(self):
        """Testing `__getitem()__` method with shape reduction (II)"""
        a = np.arange(12000).reshape((400,30))
        b = ca.carray(a, rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = (1,slice(1,4,2))
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test04c(self):
        """Testing `__getitem()__` method with shape reduction (III)"""
        a = np.arange(6000).reshape((50,40,3))
        b = ca.carray(a, rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)
        sl = (1,slice(1,4,2),2)
        #print "b[sl]->", `b[sl]`
        self.assert_(a[sl].shape == b[sl].shape, "Shape is not equal")
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

class getitemDiskTest(getitemTest, unittest.TestCase):
    disk = True
    open = False

class getitemOpenTest(getitemTest, unittest.TestCase):
    disk = True
    open = True

class setitemTest(MayBeDiskTest, unittest.TestCase):

    open = False

    def test00a(self):
        """Testing `__setitem()__` method with only a start (scalar)"""
        a = np.ones((2,3), dtype="i4")*3
        b = ca.fill((2,3), 3, dtype="i4", rootdir=self.rootdir)
        sl = slice(1)
        a[sl,:] = 0
        b[sl] = 0
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "b[sl]->", `b[sl]`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test00b(self):
        """Testing `__setitem()__` method with only a start (vector)"""
        a = np.ones((200,300), dtype="i4")*3
        b = ca.fill((200,300), 3, dtype="i4", rootdir=self.rootdir)
        sl = slice(1)
        a[sl,:] = range(300)
        b[sl] = range(300)
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "b[sl]->", `b[sl]`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test01a(self):
        """Testing `__setitem()__` method with start,stop (scalar)"""
        a = np.ones((500,200), dtype="i4")*3
        b = ca.fill((500,200), 3, dtype="i4", rootdir=self.rootdir,
                    cparams=ca.cparams())
        sl = slice(100,400)
        a[sl,:] = 0
        b[sl] = 0
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "b[sl]->", `b[sl]`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")
        #assert_array_equal(a[:], b[:], "Arrays are not equal")

    def test01b(self):
        """Testing `__setitem()__` method with start,stop (vector)"""
        a = np.ones((5,2), dtype="i4")*3
        b = ca.fill((5,2), 3, dtype="i4", rootdir=self.rootdir)
        sl = slice(1,4)
        a[sl,:] = range(2)
        b[sl] = range(2)
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "b[sl]->", `b[sl]`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test02a(self):
        """Testing `__setitem()__` method with start,stop,step (scalar)"""
        a = np.ones((1000,200), dtype="i4")*3
        b = ca.fill((1000,200), 3, dtype="i4", rootdir=self.rootdir)
        sl = slice(100,800,3)
        a[sl,:] = 0
        b[sl] = 0
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "b[sl]->", `b[sl]`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test02b(self):
        """Testing `__setitem()__` method with start,stop,step (scalar)"""
        a = np.ones((10,2), dtype="i4")*3
        b = ca.fill((10,2), 3, dtype="i4", rootdir=self.rootdir)
        sl = slice(1,8,3)
        a[sl,:] = range(2)
        b[sl] = range(2)
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "b[sl]->", `b[sl]`, `b`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test03a(self):
        """Testing `__setitem()__` method with several slices (I)"""
        a = np.arange(12000).reshape((400,30))
        b = ca.carray(a, rootdir=self.rootdir)
        sl = (slice(1,3,1), slice(1,None,2))
        #print "before->", `b[sl]`
        a[sl] = [[1],[2]]
        b[sl] = [[1],[2]]
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "after->", `b[sl]`
        assert_array_equal(a[:], b[:], "Arrays are not equal")

    def test03b(self):
        """Testing `__setitem()__` method with several slices (II)"""
        a = np.arange(24000).reshape((400,3,20))
        b = ca.carray(a, rootdir=self.rootdir)
        sl = (slice(1,3,1), slice(1,None,2), slice(1))
        #print "before->", `b[sl]`
        a[sl] = [[[1]],[[2]]]
        b[sl] = [[[1]],[[2]]]
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "after->", `b[sl]`
        assert_array_equal(a[:], b[:], "Arrays are not equal")

    def test03c(self):
        """Testing `__setitem()__` method with several slices (III)"""
        a = np.arange(120).reshape((5,4,3,2))
        b = ca.carray(a, rootdir=self.rootdir)
        sl = (slice(1,3), slice(1,3,1), slice(1,None,2), slice(1))
        #print "before->", `b[sl]`
        a[sl] = [[[[1]],[[2]]]]*2
        b[sl] = [[[[1]],[[2]]]]*2
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "after->", `b[sl]`
        assert_array_equal(a[:], b[:], "Arrays are not equal")

    def test03d(self):
        """Testing `__setitem()__` method with several slices (IV)"""
        a = np.arange(120).reshape((5,4,3,2))
        b = ca.carray(a, rootdir=self.rootdir)
        sl = (slice(1,3), slice(1,3,1), slice(1,None,2), slice(1))
        #print "before->", `b[sl]`
        a[sl] = 2
        b[sl] = 2
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "after->", `b[sl]`
        assert_array_equal(a[:], b[:], "Arrays are not equal")

    def test04a(self):
        """Testing `__setitem()__` method with shape reduction (I)"""
        a = np.arange(12).reshape((4,3))
        b = ca.carray(a, rootdir=self.rootdir)
        sl = (1,1)
        #print "before->", `b[sl]`
        a[sl] = 2
        b[sl] = 2
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "after->", `b[sl]`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test04b(self):
        """Testing `__setitem()__` method with shape reduction (II)"""
        a = np.arange(12).reshape((4,3))
        b = ca.carray(a, rootdir=self.rootdir)
        sl = (1,slice(1,4,2))
        #print "before->", `b[sl]`
        a[sl] = 2
        b[sl] = 2
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "after->", `b[sl]`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

    def test04c(self):
        """Testing `__setitem()__` method with shape reduction (III)"""
        a = np.arange(24).reshape((4,3,2))
        b = ca.carray(a, rootdir=self.rootdir)
        sl = (1,2,slice(None,None,None))
        #print "before->", `b[sl]`
        a[sl] = 2
        b[sl] = 2
        if self.open:
            b.flush()
            b = ca.open(rootdir=self.rootdir)
        #print "after->", `b[sl]`
        assert_array_equal(a[sl], b[sl], "Arrays are not equal")

class setitemDiskTest(setitemTest, unittest.TestCase):
    disk = True

class setitemOpenTest(setitemTest, unittest.TestCase):
    disk = True
    open = True

class appendTest(MayBeDiskTest, unittest.TestCase):

    def test00a(self):
        """Testing `append()` method (correct shape)"""
        a = np.ones((2,300), dtype="i4")*3
        b = ca.fill((1,300), 3, dtype="i4", rootdir=self.rootdir)
        b.append([(3,)*300])
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test00b(self):
        """Testing `append()` method (correct shape, single row)"""
        a = np.ones((2,300), dtype="i4")*3
        b = ca.fill((1,300), 3, dtype="i4", rootdir=self.rootdir)
        b.append((3,)*300)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test01(self):
        """Testing `append()` method (incorrect shape)"""
        a = np.ones((2,3), dtype="i4")*3
        b = ca.fill((1,3), 3, dtype="i4", rootdir=self.rootdir)
        self.assertRaises(ValueError, b.append, [(3,3)])

    def test02(self):
        """Testing `append()` method (several rows)"""
        a = np.ones((4,3), dtype="i4")*3
        b = ca.fill((1,3), 3, dtype="i4", rootdir=self.rootdir)
        b.append([(3,3,3)]*3)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

class appendDiskTest(appendTest, unittest.TestCase):
    disk = True

class resizeTest(MayBeDiskTest, unittest.TestCase):

    def test00a(self):
        """Testing `resize()` (trim)"""
        a = np.ones((2,3), dtype="i4")
        b = ca.ones((3,3), dtype="i4", rootdir=self.rootdir)
        b.resize(2)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test00b(self):
        """Testing `resize()` (trim to zero)"""
        a = np.ones((0,3), dtype="i4")
        b = ca.ones((3,3), dtype="i4", rootdir=self.rootdir)
        b.resize(0)
        #print "b->", `b`
        # The next does not work well for carrays with shape (0,)
        #assert_array_equal(a, b, "Arrays are not equal")
        self.assert_("a.dtype.base == b.dtype.base")
        self.assert_("a.shape == b.shape+b.dtype.shape")

    def test01(self):
        """Testing `resize()` (enlarge)"""
        a = np.ones((4,3), dtype="i4")
        b = ca.ones((3,3), dtype="i4", rootdir=self.rootdir)
        b.resize(4)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

class resizeDiskTest(resizeTest, unittest.TestCase):
    disk = True

class iterTest(unittest.TestCase):

    def test00(self):
        """Testing `iter()` (no start, stop, step)"""
        a = np.ones((3,), dtype="i4")
        b = ca.ones((1000,3), dtype="i4")
        #print "b->", `b`
        for r in b.iter():
            assert_array_equal(a, r, "Arrays are not equal")

    def test01(self):
        """Testing `iter()` (w/ start, stop)"""
        a = np.ones((3,), dtype="i4")
        b = ca.ones((1000,3), dtype="i4")
        #print "b->", `b`
        for r in b.iter(start=10):
            assert_array_equal(a, r, "Arrays are not equal")

    def test02(self):
        """Testing `iter()` (w/ start, stop, step)"""
        a = np.ones((3,), dtype="i4")
        b = ca.ones((1000,3), dtype="i4")
        #print "b->", `b`
        for r in b.iter(15, 100, 3):
            assert_array_equal(a, r, "Arrays are not equal")

class reshapeTest(unittest.TestCase):

    def test00a(self):
        """Testing `reshape()` (unidim -> ndim)"""
        a = np.ones((3,4), dtype="i4")
        b = ca.ones(12, dtype="i4").reshape((3,4))
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test00b(self):
        """Testing `reshape()` (unidim -> ndim, -1 in newshape (I))"""
        a = np.ones((3,4), dtype="i4")
        b = ca.ones(12, dtype="i4").reshape((-1,4))
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test00c(self):
        """Testing `reshape()` (unidim -> ndim, -1 in newshape (II))"""
        a = np.ones((3,4), dtype="i4")
        b = ca.ones(12, dtype="i4").reshape((3,-1))
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test01(self):
        """Testing `reshape()` (ndim -> unidim)"""
        a = np.ones(12, dtype="i4")
        c = ca.ones(12, dtype="i4").reshape((3,4))
        b = c.reshape(12)
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test02a(self):
        """Testing `reshape()` (ndim -> ndim, I)"""
        a = np.arange(12, dtype="i4").reshape((3,4))
        c = ca.arange(12, dtype="i4").reshape((4,3))
        b = c.reshape((3,4))
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test02b(self):
        """Testing `reshape()` (ndim -> ndim, II)"""
        a = np.arange(24, dtype="i4").reshape((2,3,4))
        c = ca.arange(24, dtype="i4").reshape((4,3,2))
        b = c.reshape((2,3,4))
        #print "b->", `b`
        assert_array_equal(a, b, "Arrays are not equal")

    def test03(self):
        """Testing `reshape()` (0-dim)"""
        a = np.ones((0,4), dtype="i4")
        b = ca.ones(0, dtype="i4").reshape((0,4))
        #print "b->", `b`
        # The next does not work well for carrays with shape (0,)
        #assert_array_equal(a, b, "Arrays are not equal")
        self.assert_(a.dtype.base == b.dtype.base)
        self.assert_(a.shape == b.shape+b.dtype.shape)

class compoundTest:

    def test00(self):
        """Testing compound types (creation)"""
        a = np.ones((300,4), dtype=self.dtype)
        b = ca.ones((300,4), dtype=self.dtype)
        #print "b.dtype-->", b.dtype
        #print "b->", `b`
        self.assert_(a.dtype == b.dtype.base)
        assert_array_equal(a, b[:], "Arrays are not equal")

    def test01(self):
        """Testing compound types (append)"""
        a = np.ones((300,4), dtype=self.dtype)
        b = ca.carray([], dtype=self.dtype).reshape((0,4))
        b.append(a)
        #print "b.dtype-->", b.dtype
        #print "b->", `b`
        self.assert_(a.dtype == b.dtype.base)
        assert_array_equal(a, b[:], "Arrays are not equal")

    def test02(self):
        """Testing compound types (iter)"""
        a = np.ones((3,), dtype=self.dtype)
        b = ca.ones((1000,3), dtype=self.dtype)
        #print "b->", `b`
        for r in b.iter():
            #print "r-->", r
            assert_array_equal(a, r, "Arrays are not equal")

class plainCompoundTest(compoundTest, unittest.TestCase):
    dtype = np.dtype("i4,i8")

class nestedCompoundTest(compoundTest, unittest.TestCase):
    dtype = np.dtype([('f1', [('f1', 'i2'), ('f2', 'i4')])])

class stringTest(unittest.TestCase):

    def test00(self):
        """Testing string types (creation)"""
        a = np.array([["ale", "ene"], ["aco", "ieie"]], dtype="S4")
        b = ca.carray(a)
        #print "b.dtype-->", b.dtype
        #print "b->", `b`
        self.assert_(a.dtype == b.dtype.base)
        assert_array_equal(a, b[:], "Arrays are not equal")

    def test01(self):
        """Testing string types (append)"""
        a = np.ones((300,4), dtype="S4")
        b = ca.carray([], dtype="S4").reshape((0,4))
        b.append(a)
        #print "b.dtype-->", b.dtype
        #print "b->", `b`
        self.assert_(a.dtype == b.dtype.base)
        assert_array_equal(a, b[:], "Arrays are not equal")

    def test02(self):
        """Testing string types (iter)"""
        a = np.ones((3,), dtype="S40")
        b = ca.ones((1000,3), dtype="S40")
        #print "b->", `b`
        for r in b.iter():
            #print "r-->", r
            assert_array_equal(a, r, "Arrays are not equal")

class unicodeTest(unittest.TestCase):

    def test00(self):
        """Testing unicode types (creation)"""
        a = np.array([[u"aŀle", u"eñe"], [u"açò", u"áèâë"]], dtype="U4")
        b = ca.carray(a)
        #print "b.dtype-->", b.dtype
        #print "b->", `b`
        self.assert_(a.dtype == b.dtype.base)
        assert_array_equal(a, b[:], "Arrays are not equal")

    def test01(self):
        """Testing unicode types (append)"""
        a = np.ones((300,4), dtype="U4")
        b = ca.carray([], dtype="U4").reshape((0,4))
        b.append(a)
        #print "b.dtype-->", b.dtype
        #print "b->", `b`
        self.assert_(a.dtype == b.dtype.base)
        assert_array_equal(a, b[:], "Arrays are not equal")

    def test02(self):
        """Testing unicode types (iter)"""
        a = np.ones((3,), dtype="U40")
        b = ca.ones((1000,3), dtype="U40")
        #print "b->", `b`
        for r in b.iter():
            #print "r-->", r
            assert_array_equal(a, r, "Arrays are not equal")

class computeMethodsTest(unittest.TestCase):

    def test00(self):
        """Testing sum()."""
        a = np.arange(1e5).reshape(10, 1e4)
        sa = a.sum()
        ac = ca.carray(a)
        sac = ac.sum()
        #print "numpy sum-->", sa
        #print "carray sum-->", sac
        self.assert_(sa.dtype == sac.dtype, "sum() is not working correctly.")
        self.assert_(sa == sac, "sum() is not working correctly.")

class carrayConstructorDimensionTest(MayBeDiskTest, unittest.TestCase):
    """
    This test is related to issue #14 in blaze github repo.
    Check that when using a carray constructor the dimensionality of the array
    is not lost. Neither in the implicit dtype or explicit dtype cases.
    """
    open = False

    def testImplicitDtype(self):
        """Testing carray construction keeping dimensions (implicit dtype)"""
        a = np.eye(6) # 2d
        b = ca.carray(a, rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)

        # array equality implies having the same shape
        assert_array_equal(a, b, "Arrays are not equal")

    def testExplicitDtype(self):
        """Testing carray construction keeping dimensions (explicit dtype)"""
        dtype = np.dtype(np.float64)
        a = np.eye(6, dtype=dtype)
        b = ca.carray(a, dtype=dtype, rootdir=self.rootdir)
        if self.open:
            b = ca.open(rootdir=self.rootdir)

        # array equality implies having the same shape
        assert_array_equal(a, b, "Arrays are not equal")

class carrayConstructorDimensionDiskTest(carrayConstructorDimensionTest,
                                         unittest.TestCase):
    disk = True
    open = False

class carrayConstructorDimensionOpenTest(carrayConstructorDimensionTest,
                                         unittest.TestCase):
    disk = True
    open = True

## Local Variables:
## mode: python
## py-indent-offset: 4
## tab-width: 4
## fill-column: 72
## End:


from os.path import split
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelWithLMHead
import torch
import numpy as np
import pandas as pd
from tqdm import tqdm
import os
import pickle
import utils

SAVE_PATH = 'output/gpt2/imdb'
if not os.path.exists(SAVE_PATH):
    os.makedirs(SAVE_PATH)

LR = '5e-5'
EPOCHS = 1

def compute_all(model, all_encodings, fname, n=None):
    print("Finding perplexities")
    perplexities, lls = [], []
    # all_encodings = all_encodings[:n]
    pbar = tqdm(total=len(all_encodings))
    for idx, encodings in enumerate(all_encodings):
        try:
            pp, ll = compute_perplexity(model, encodings, device=device)
            perplexities.append(pp)
            lls.append(ll)
        except Exception as e:
            print("Exception at idx", idx)
            print(e)
            continue
        finally:
            pbar.update(1)

    pbar.close()

    perplexities = np.array(perplexities)
    np.save(f'{SAVE_PATH}/{fname}_{LR}_pps.npy', perplexities)
    print(f"\nMean: {perplexities.mean()}, Std: {perplexities.std()}")

    with open(f'{SAVE_PATH}/{fname}_{LR}_lls.pkl', 'wb') as fw:
        pickle.dump(lls, fw)

    return perplexities

def compute_perplexity(model, encodings, stride=None, device='cuda'):
    max_length = model.config.n_positions
    lls = []
    if stride is None:
        stride = 1

    for i in range(1, encodings.input_ids.size(1), stride):
        begin_loc = max(i + stride - max_length, 0)
        end_loc = i + stride
        input_ids = encodings.input_ids[:,begin_loc:end_loc].to(device)
        target_ids = input_ids.clone()
        target_ids[:,:-stride] = -100

        with torch.no_grad():
            outputs = model(input_ids, labels=target_ids)
            log_likelihood = outputs[0] * stride

        lls.append(log_likelihood)

    ppl = torch.exp(torch.stack(lls).sum() / i)
    return ppl.item(), torch.stack(lls).detach().cpu().numpy()

def setup(path='lvwerra/gpt2-imdb'):
    if torch.cuda.is_available():
        device = 'cuda'
    else:
        device = 'cpu'
    n = None
    model = AutoModelWithLMHead.from_pretrained(path).to(device)
    tokenizer = AutoTokenizer.from_pretrained(path)
    return model, tokenizer, device, n

def process_dataset(dataset_name, model, tokenizer, device, n=None, key='text', configs=None, fname=None):
    print(f"\n------Processing perplexity for dataset: {dataset_name}-------")

    if configs is None:
        dataset = load_dataset(dataset_name, split='test')
    else:
        dataset = [load_dataset(dataset_name, config, split='test') for config in configs]

    print("Tokenizing data")
    if configs is None:
        all_encodings = [tokenizer(text + ' <|endoftext|>', return_tensors='pt') for text in dataset[key][:n]]
    else:
        all_encodings = [tokenizer(text + ' <|endoftext|>', return_tensors='pt') for _dataset in dataset for text in _dataset[key][:n]]

    if fname is None:
        fname = dataset_name
    return compute_all(model, all_encodings, fname, n)

def process_entailment(dataset_name, model, tokenizer, device, n=None, dataset_subname=None, fname=None, key1='premise', key2='hypothesis'):
    print(f"\n------Processing perplexity for dataset: {dataset_name}_{dataset_subname}-------")

    if dataset_subname is None:
        dataset = load_dataset(dataset_name, split='validation')
    else:
        dataset = load_dataset(dataset_name, dataset_subname, split='validation')
    dataset_texts = []
    for ex in dataset:
        dataset_texts.append(ex[key1] + ' ' + ex[key2])

    print("Tokenizing data")
    all_encodings = [tokenizer(text + ' <|endoftext|>', return_tensors='pt') for text in dataset_texts[:n]]

    if fname is None:
        fname = dataset_name
    return compute_all(model, all_encodings, fname, n)

def process_counterfactual(dataset_name, model, tokenizer, device, n=None, fname=None, key='Text'):
    print('Loading data...')
    train_df = pd.read_table(os.path.join(os.getcwd(), 'data', dataset_name, (utils.train_split_keys[dataset_name] + '.tsv')))
    eval_df = pd.read_table(os.path.join(os.getcwd(), 'data', dataset_name, (utils.eval_split_keys[dataset_name] + '.tsv')))
    test_df = pd.read_table(os.path.join(os.getcwd(), 'data', dataset_name, (utils.test_split_keys[dataset_name] + '.tsv')))

    num_labels = len(np.unique(pd.Categorical(train_df[utils.label_keys[dataset_name]], ordered=True)))
    dataset = pd.concat([train_df, eval_df, test_df])
    print(dataset)

    print("Tokenizing data")
    all_encodings = [tokenizer(text + ' <|endoftext|>', return_tensors='pt') for text in dataset[key][:n]]

    if fname is None:
        fname = dataset_name
    return compute_all(model, all_encodings, fname, n)

if __name__ == '__main__':
    print("Loading model...")
    path = f'/scratch/ua388/nlp/ckpts/gpt2-imdb-{EPOCHS}_epochs-{LR}_lr'
    # path = f'/scratch/ua388/nlp/ckpts/gpt2-glue_sst2-{EPOCHS}_epochs-{LR}_lr'
    print("Loading model...", path)
    model, tokenizer, device, n = setup(path)
    # process_dataset('imdb', model, tokenizer, device, n=3000, key='text')
    # process_dataset('yelp_polarity', model, tokenizer, device, n=3000, key='text')

    process_counterfactual('counterfactual-imdb', model, tokenizer, device, n=2)

    # process_dataset('sentiment140', model, tokenizer, device, n=2, key='text')

    # process_dataset('glue', model, tokenizer, device, configs=['sst2'], key='sentence', fname='sst2')

    # process_entailment('glue', model, tokenizer, device, dataset_subname='rte', fname='rte', key1='sentence1', key2='sentence2')
    # process_entailment('snli', model, tokenizer, device, key1='premise', key2='hypothesis')

    print("\n\n--------DONE--------")


"""
Autor: Blanca Alonso

Chuleta y ayudas para comandos en Python

"""

# =============================================================================
# 0.- ESTABLECER DIRECTORIO
# =============================================================================

import os
os.chdir("path")

# ¡IMPORTANTE! Las barras del directorio deben ser dobles \\

# =============================================================================
# 1.- CARGA DE PAQUETES
# =============================================================================

# Librerías numpy: biblioteca de álgebra lineal más conocida

import numpy as np

# Librería pandas: tratamiento de datos estructurados, semi y no estructurados
import pandas as pd

# =============================================================================
# 2.- LECTURA Y ESCRITURA
# =============================================================================

# CSV
df = pd.read_csv("Input.csv", sep = ";", decimal = ",")  # lectura
df.to_csv("Output.csv", sep = ";", index = False) # escritura

# EXCEL
df = pd.read_excel("Input.xlsx")
df.to_excel("Output.xlsx", index = False)

# =============================================================================
# 3.- TRATAMNIENTO BBDD
# =============================================================================

# Limpiar datos
df['campo'].str.replace(r" \(.*\)","")

# Tipos de datos de cada columna del dataframe
df.dtypes

# Cambiar el tipo de dato
df['campo'].astype(str)

# Nombres de tablas a mayusculas (dos formas distintas)
df.columns = df.columns.str.upper()
df.columns = map(lambda x: str(x).upper(), df.columns)

# Convertir todos los valores de un campo a mayúsculas
df['campo'] = df['campo'].str.upper()

# Lista con x nombres de las columnas definido con un rango
colnames_df = list(df.columns[2:4])

# Cambiar los nombres de todos los campos de la tabla segun el orden en el que están
df.columns = ['campo_nuevo1', 'campo_nuevo2', 'campo_nuevo3']
df = df.rename(columns = {"campo_now1": "campo_new1", "campo_now2": "campo_new2"})

# Quitar espacios en blanco en los nombres de las columnas o en un campo
df.columns = df.columns.str.strip()                             # trim
df.columns = df.columns.str.strip().str.replace(' ', '')        # replace
df['campo'] = df.campo.str.strip().str.replace(' ', 'vacío')
df['campo'] = df.campo.str.strip().replace('...', 'vacío')

# Reemplazar valores según sea igual o no a un valor
df['campo'] = np.where(df.TOTAL_PAQUETE.isin(['condicion1', 'condicion2']),
              df['campo'] , 'Otros')

# Select
df.columna  # Una columna
df = df[["columna1", "columna2", "columna3"]]   # Varias columnas

# Seleccionar x columnas según un filtro
df.loc[df['columna1' > 10], ['columna1', 'columna2']]

# Eliminar una columna
df = df.drop(columns = 'Columna')
df = df.drop(['Columna'], axis = 1)

# NaN
df.isnull()                     # Para comprobar si hay NaN
df.notnull()                     # Para comprobar si no hay NaN
df.dropna()                     # Elimina cualquier fila con NaN
df.fillna(value = 0)            # Reemplaza los NaN por un valor
df = df[df['campo'].notna()]    # Elimina los NaN de una columna

list(df['campo'].isnull()).count(True) # Comprobar NA's en una columna

# Imputation Na
from sklearn.impute import SimpleImputer
my_imputer = SimpleImputer(strategy = 'XX') #a que se quieren imputar: median, constant (númerico), most_frequent
imputed_X_train = pd.DataFrame(my_imputer.fit_transform(X_train))

# Quitar duplicados
df = df.drop_duplicates()

# Valores distintos de una columna
df['columna'].nunique() # devuelve un int de cuántos valores != hay
df['columna'].unique()  # devuelve un array con los != valores

# Apendizar dos dataframes por filas
df1 = df1.append(df2)
df = pd.concat([df1, df2])

# Apendizar dos dataframes por filas
df = pd.concat([df1, df2])

# Calcular el máximo según un campo
df[df['CAMPO'] == 'FILTRO'].CAMPO_NUMERICO.max()

# Recuento de la frecuencia de los valores de una columna
df['CAMPO'].value_counts()

# Dividir una columna en dos por un delimitador
df[['CAMPO1','CAMPO2']] = df['CAMPO1'].str.split(',',expand=True)

# Borrar NA´s
df.dropna(axis=0)

# Ordenar
DF.sort_values (by )
DF.sort_index()

# 3.1.- FECHAS
# --------------------------------------------------------------------------#

import datetime as dt

# Transformar un campo a fecha
df['FECHA'] = pd.to_datetime(df['FECHA'])

# Extraer el año, mes, día
df['ANIO'] = df['FECHA'].dt.year

# =============================================================================
# 4.-  OPERACIONES CON LAS BBDD
# =============================================================================

# Resumen de los datos con la media, desviación tipica, cuartiles
DF.describe()

# Primeras filas
DF.head()

# Aparecen las columnas
DF.columns

# Cuenta las veces que aparece un elemento en una lista
lista.count(valor)

# En un dataframe
list(df['CAMPO']).count('valor')

# Crear columnas condicionales
obs_pres_med = [x for x in obs if "Prescripción Médica" in x]

df['Prescr_Medica'] = df.Observaciones.apply(lambda x: 'SI' if x in obs_pres_med else 'NO')

# Transformar un dataframe a diccionario
df_dict = df.set_index('key').T.to_dict('list')

# =============================================================================
# 5.- CRUCES
# =============================================================================

# ¡OJO! Asegurarse que df2 no tenga duplicados en la clave

# Inner Join: valores coincidentes
df1.merge(df2, how = "inner", on = "CLAVE")

# Left y right join
df1.merge(df2, how = "left", on = "CLAVE")
df1.merge(df2, how = "right", on = "CLAVE")

# Full join: unimos las dos tablas crucen o no
df1.merge(df2, how = "outer", on = "CLAVE")

# Si los campos no tienen el mismo nombre
df1.merge(df2, how = "left", left_on = "CLAVE_df_uno", right_on = "CLAVE_df2")

# Si queremos seleccionar sólo algunos campos del dataframe 2
df1.merge(df2[["campos"]], how = "left", left_on = "CLAVE_df_uno",
          right_on = "CLAVE_df2")

# Cruzar por varios campos
df1.merge(df2, how = "left", on = ["CAMPO1", "CAMPO2"])

# =============================================================================
# 6.- BUCLES FOR & IF
# =============================================================================

# Para obtener los diferentes valores de error si se cambian parametros en cada modelo

from sklearn.metrics import mean_absolute_error
from sklearn.tree import DecisionTreeRegressor

def get_mae(max_leaf_nodes, train_X, val_X, train_y, val_y):
    model = DecisionTreeRegressor(max_leaf_nodes=max_leaf_nodes, random_state=0)
    model.fit(train_X, train_y)
    preds_val = model.predict(val_X)
    mae = mean_absolute_error(val_y, preds_val)
    return(mae)

 for max_leaf_nodes in [5, 50, 500, 5000]:
    my_mae = get_mae(max_leaf_nodes, train_X, val_X, train_y, val_y)
    print("Max leaf nodes: %d  \t\t Mean Absolute Error:  %d" %(max_leaf_nodes, my_mae))

# =============================================================================
# 7.- TOTALIZACIONES
# =============================================================================

# Totalizar
df = df.groupby("CAMPO").sum()                  # Por una variable
df = df.groupby(["CAMPO1", "CAMPO2"]).sum()     # Por varias variables

# Totalizar realizando distintas operaciones segun el campo
df = df.groupby("CAMPO").agg({'CAMPO_num1': sum, 'CAMPO_num2' : max})

# Totalizar definiendo un campo nuevo en base a otro
df = df.set_index('CAMPO').groupby(level = 0)['CAMPO_num1'].agg({'CAMPO_new' : np.sum})

# =============================================================================
# 8.- FUNCIONES LAMBDA
# =============================================================================

# =============================================================================
# 9.- MACHINE LEARNING
# =============================================================================

# Definir una semilla
import random
random.seed(1234)

# Dividimos en train y test
from sklearn.model_selection import train_test_split

datos_spliteados = train_test_split(df,
                                    train_size = 0.8,     # 80% training
                                    test_size = 0.2       # 20% testing
                                   )

df_training_set = datos_spliteados[0]
df_test_set = datos_spliteados[1]

# 9.1.- Regresión Logístico ----------------------------------------------
# --------------------------------------------------------------------------#

from sklearn.linear_model import LogisticRegression

# Creamos una instancia del modelo, con todos los hiperparámetros y argumentos
# por defecto:

lr = LogisticRegression()

# Ahora vamos a entrenarla con el conjunto de training:
lr.fit(X = dataset_clasificacion_training[["hormona_a", "hormona_b"]], # features
       y = dataset_clasificacion_training["label"])                    # labels

# 9.2.- Arboles----------------------------------------------
# --------------------------------------------------------------------------#
from sklearn.tree import DecisionTreeRegressor

# Define model. Specify a number for random_state to ensure same results each run
TR_model = DecisionTreeRegressor(random_state=1) # = 1 para que siempre se cree el mismo modelo

# Ahora vamos a entrenarla con el conjunto de training:
DF.fit(X_train, y_train)

# 9.3.- Random Forest----------------------------------------------
# --------------------------------------------------------------------------#
from sklearn.ensemble import RandomForestRegressor

# Define model. Specify a number for random_state to ensure same results each run
forest_model = RandomForestRegressor(random_state=1) # = 1 para que siempre se cree el mismo modelo

# Ahora vamos a entrenarla con el conjunto de training:
DF.fit(X_train, y_train)

# 9.4.- XGBoost---------------------------------------------
# --------------------------------------------------------------------------#
from xgboost import XGBRegressor

my_model = XGBRegressor(n_estimators=1000, learning_rate=0.05)
my_model.fit(X_train, y_train,  early_stopping_rounds=5,
             eval_set=[(X_valid, y_valid)],
             verbose=False)

# 9.5.- Cross Validation---------------------------------------------
# --------------------------------------------------------------------------#
from sklearn.ensemble import RandomForestRegressor
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer

my_pipeline = Pipeline(steps=[('preprocessor', SimpleImputer()), # imputar los missings a lo que se decida
                              ('model', RandomForestRegressor(n_estimators=50,
                                                              random_state=0))
                             ])

from sklearn.model_selection import cross_val_score

# Multiply by -1 since sklearn calculates *negative* MAE
scores = -1 * cross_val_score(my_pipeline, X, y,
                              cv=5,
                              scoring='neg_mean_absolute_error')

print("MAE scores:\n", scores)

# 9.6.- Predicción y error----------------------------------------------
# --------------------------------------------------------------------------#
from sklearn.metrics import mean_absolute_error

prediccion = Df.predict(X)
mean_absolute_error(y, prediccion)

# 9.7.- Output---------------------------------------------
# --------------------------------------------------------------------------#
output = pd.DataFrame({'Id': X_test.index,
                       'Target': preds_test})
output.to_csv('nombre_archivo.csv', index=False)

# =============================================================================
# 10.- OTRAS COSAS
# =============================================================================

from  pandas.io.json  import  json_normalize

# Normalizar los datos
df  =  json_normalize(creds)


from logging import WARNING
from warnings import WarningMessage, warn_explicit
import warnings
import numpy as np
from pytest import raises
import pytest
from strapvizpy.bootstrap import bootstrap_distribution, calculate_boot_stats

def test_bootstrap_distribution():
    """
    Tests the functionality of `bootstrap_distribution()`, which includes
    checking the type of object returned, its shape and its values.

    5 tests in total.
    """

    dist = bootstrap_distribution([1, 2, 3], 3, 3, random_seed=1)

    # checks the return type
    assert isinstance(dist, np.ndarray)

    # checks the shape of returned object
    assert dist.shape == (3,)

    # checks the values of returned object
    assert np.array_equal(dist.tolist(),
                          [1.3333333333333333,
                           1.6666666666666667,
                           1.3333333333333333])

    dist_2 = bootstrap_distribution([1, 2, 3],
                                  3,
                                  estimator="var",
                                  random_seed=1)

    # checks with different estimator
    assert np.array_equal(dist_2.tolist(),
                          [0.2222222222222222,
                           0.2222222222222222,
                           0.2222222222222222])

    dist_3 = bootstrap_distribution([1, 2, 3],
                                  3,
                                  estimator="var",
                                  random_seed=1)

    # checks if random_seed works
    assert np.array_equal(dist_2, dist_3)

def test_bootstrap_distribution_errors():
    """
    Tests error cases and messages thrown by `bootstrap_distribution()`.

    10 tests in total.
    """

    # tests with invalid input type of sample
    with raises(TypeError) as e:
        bootstrap_distribution({1, 2, 3}, 3, 3)
    assert str(e.value) == (
        "sample should be one of the types"
        "[list, numpy.ndarray, pandas.core.series.Series]"
    )

    # tests with invalid input type of rep
    with raises(TypeError) as e:
        bootstrap_distribution([1, 2, 3], 3.3, 3)
    assert str(e.value) == "rep should be of type 'int'"

    # tests with invalid input value of rep
    with raises(ValueError) as e:
        bootstrap_distribution([1, 2, 3], -1, 3)
    assert str(e.value) == "Invalid value for rep"

    # tests with invalid input type of n
    with raises(TypeError) as e:
        bootstrap_distribution([1, 2, 3], 3, 3.3)
    assert str(e.value) == "n should be of type 'str' or 'int'"

    # tests with invalid input value of n
    with raises(ValueError) as e:
        bootstrap_distribution([1, 2, 3], 3, "catch me")
    assert str(e.value) == "Invalid value for n. Did you intend n='auto'?"

    # tests with invalid input value of n
    with raises(ValueError) as e:
        bootstrap_distribution([1, 2, 3], 3, -3)
    assert str(e.value) == "Invalid value for n"

    # tests with invalid input type of estimator
    with raises(TypeError) as e:
        bootstrap_distribution([1, 2, 3], 3, 3, estimator=9)
    assert str(e.value) == "estimator should be of type 'str'"

    # tests with invalid input value of estimator
    with raises(ValueError) as e:
        bootstrap_distribution([1, 2, 3], 3, 3, estimator="outlier")
    assert str(e.value) == "Supported estimators are mean, median, var, sd"

    # tests with invalid input type of random_seed
    with raises(TypeError) as e:
        bootstrap_distribution([1, 2, 3], 3, 3, random_seed="I'm not a seed")
    assert str(e.value) == "random_seed should be None or of type 'int'"

    # tests with invalid input value of random_seed
    with raises(ValueError) as e:
        bootstrap_distribution([1, 2, 3], 3, 3, random_seed=-3)
    assert str(e.value) == "Invalid value for random_seed"

def test_calculate_boot_stats():
    """
    Tests functionality of calculate_boot_stats()

    16 tests in total.
    """
    # test integration with bootstrap dist function
    test_dist = bootstrap_distribution(sample=[1, 2, 3],
                                       rep=100)

    assert isinstance(test_dist, np.ndarray)

    # set up test dicts
    test_dict = calculate_boot_stats(
        [1, 2, 3, 4],
        1000,
        random_seed=123)

    test_dict_2 = calculate_boot_stats(
        [1000, 2000, 3000, 4000],
        n=4,
        rep=1000,
        level=0.9,
        random_seed=1234,
        estimator='var')

    test_dict_3, test_sample_dist = calculate_boot_stats(
        [1000, 2000, 3000, 4000],
        n=5,
        rep=1000,
        level=0.9,
        random_seed=1234,
        estimator='var',
        pass_dist=True)

    # checks the return type
    assert isinstance(test_dict, dict)

    assert isinstance(test_sample_dist, np.ndarray)

    # check properties and values of dictionary output
    assert len(test_dict) == 9

    assert test_dict["lower"] == 1.5

    assert test_dict["upper"] == 3.5

    assert test_dict["std_err"] == 0.5414773771820943

    assert test_dict["sample_size"] == 4

    assert test_dict["n"] == 'auto'

    assert test_dict['rep'] == 1000

    assert test_dict['level'] == 0.95

    assert test_dict['estimator'] == 'mean'

    # increasing sampling number should decrease the standard error
    assert test_dict_2['std_err'] > test_dict_3['std_err']

    # changing level parameter should change its value
    assert test_dict_2['level'] != test_dict['level']

    # changing estimate parameter should change its value
    assert test_dict_3['estimator'] == 'var'

def test_calculate_boot_stats_errors():
    """
    Tests error cases and messages thrown by `calculate_boot_stats()`.

    4 tests in total.
    """

    # tests with invalid input type of sample
    with raises(TypeError) as e:
        calculate_boot_stats([1, 2, 3, 4],
        1000,
        level='ninety-five',
        estimator="mean",
        random_seed=123)
    assert str(e.value) == ("level should be of type 'float'")

    with raises(TypeError) as e:
        calculate_boot_stats([1, 2, 3, 4],
        1000,
        level=0.95,
        estimator="mean",
        random_seed=123,
        pass_dist='True')
    assert str(e.value) == ("pass_dist should be of type 'bool'")

    with raises(ValueError) as e:
        calculate_boot_stats([1, 2, 3, 4],
        1000,
        level=1.0,
        estimator="mean",
        random_seed=123)
    assert str(e.value) == ("level should be between 0 and 1")

    with raises(ValueError) as e:
        calculate_boot_stats([1, 2, 3, 4],
        1000,
        level=0.0,
        estimator="mean",
        random_seed=123)
    assert str(e.value) == ("level should be between 0 and 1")

    with pytest.warns(UserWarning) as w:
        calculate_boot_stats([1, 2, 3, 4],
                            1000,
                            level=0.05)
    assert w