from __future__ import absolute_import, division, print_function

import tempfile
import os
import inspect

from itertools import islice
from contextlib import contextmanager
from collections import Iterator

import numpy as np

# Imports that replace older utils.
from cytoolz import count, unique, partition_all, nth, groupby, reduceby
from blaze.compatibility import map, zip


def nth_list(n, seq):
    """

    >>> tuple(nth_list([0, 1, 4], 'Hello'))
    ('H', 'e', 'o')
    >>> tuple(nth_list([4, 1, 0], 'Hello'))
    ('o', 'e', 'H')
    >>> tuple(nth_list([0, 0, 0], 'Hello'))
    ('H', 'H', 'H')
    """
    seq = iter(seq)
    sn = sorted(n)

    result = []
    old = 0
    item = next(seq)
    for index in sorted(n):
        for i in range(index - old):
            item = next(seq)
        result.append(item)
        old = index

    order = [x[1] for x in sorted(zip(n, range(len(n))))]
    return (result[i] for i in order)


def get(ind, coll, lazy=False):
    """

    >>> get(0, 'Hello')
    'H'

    >>> get([1, 0], 'Hello')
    ('e', 'H')

    >>> get(slice(1, 4), 'Hello')
    ('e', 'l', 'l')

    >>> get(slice(1, 4), 'Hello', lazy=True)  # doctest: +SKIP
    <itertools.islice object at 0x25ac470>
    """
    if isinstance(ind, list):
        result = nth_list(ind, coll)
    elif isinstance(ind, slice):
        result = islice(coll, ind.start, ind.stop, ind.step)
    else:
        if isinstance(coll, Iterator):
            result = nth(ind, coll)
        else:
            result = coll[ind]
    if lazy==False and isinstance(result, Iterator):
        result = tuple(result)
    return result


def ndget(ind, data):
    """
    Get from N-Dimensional getable

    Can index with elements, lists, or slices.  Mimic's numpy fancy indexing on
    generic indexibles.

    >>> data = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
    >>> ndget(0, data)
    [[1, 2], [3, 4]]
    >>> ndget((0, 1), data)
    [3, 4]
    >>> ndget((0, 0, 0), data)
    1
    >>> ndget((slice(0, 2), [0, 1], 0), data)
    ((1, 3), (5, 7))
    """
    if isinstance(ind, tuple) and len(ind) == 1:
        ind = ind[0]
    if not isinstance(ind, tuple):
        return get(ind, data)
    result = get(ind[0], data)
    if isinstance(ind[0], (list, slice)):
        return type(result)(ndget(ind[1:], row) for row in result)
    else:
        return ndget(ind[1:], result)


@contextmanager
def filetext(text, extension='', open=open):
    with tmpfile(extension=extension) as filename:
        f = open(filename, "wt")
        f.write(text)
        try:
            f.close()
        except AttributeError:
            pass

        yield filename


@contextmanager
def filetexts(d, open=open):
    """ Dumps a number of textfiles to disk

    d - dict
        a mapping from filename to text like {'a.csv': '1,1\n2,2'}
    """
    for filename, text in d.items():
        f = open(filename, 'wt')
        f.write(text)
        try:
            f.close()
        except AttributeError:
            pass

    yield list(d)

    for filename in d:
        if os.path.exists(filename):
            os.remove(filename)


@contextmanager
def tmpfile(extension=''):
    extension = '.' + extension.lstrip('.')
    handle, filename = tempfile.mkstemp(extension)

    try:
        yield filename
    finally:
        try:
            if os.path.exists(filename):
                os.remove(filename)
        except OSError:  # Sometimes Windows can't close files
            if os.name == 'nt':
                os.close(handle)
                os.remove(filename)


def raises(err, lamda):
    try:
        lamda()
        return False
    except err:
        return True


def keywords(func):
    """ Get the argument names of a function

    >>> def f(x, y=2):
    ...     pass

    >>> keywords(f)
    ['x', 'y']
    """
    if isinstance(func, type):
        return keywords(func.__init__)
    return inspect.getargspec(func).args


def assert_allclose(lhs, rhs):
    for tb in map(zip, lhs, rhs):
        for left, right in tb:
            if isinstance(left, (np.floating, float)):
                # account for nans
                assert np.all(np.isclose(left, right, equal_nan=True))
            else:
                assert left == right


def example(filename, datapath=os.path.join('examples', 'data')):
    import blaze
    return os.path.join(os.path.dirname(blaze.__file__), os.pardir, datapath,
                        filename)

@@ -2506,16 +2506,33 @@
 , %22wt%22)%0A
+        try:%0A    
         
@@ -2537,32 +2537,53 @@
   f.write(text)%0A
+        finally:%0A    
         try:%0A   
@@ -2583,32 +2583,36 @@
 ry:%0A            
+    
 f.close()%0A      
@@ -2597,32 +2597,36 @@
       f.close()%0A
+    
         except A
@@ -2632,32 +2632,36 @@
 AttributeError:%0A
+    
             pass
@@ -2935,16 +2935,33 @@
 , 'wt')%0A
+        try:%0A    
         
@@ -2966,32 +2966,53 @@
   f.write(text)%0A
+        finally:%0A    
         try:%0A   
@@ -3012,32 +3012,36 @@
 ry:%0A            
+    
 f.close()%0A      
@@ -3026,32 +3026,36 @@
       f.close()%0A
+    
         except A
@@ -3061,32 +3061,36 @@
 AttributeError:%0A
+    
             pass

Make sure we close files properly