# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_items.ipynb (unless otherwise specified).

__all__ = ['flatIndex', 'arrayIndex', 'Item', 'Collection', 'Selection', 'ItemList', 'ItemArray']

# Cell
from traitlets import HasTraits, Tuple, Unicode, Any, observe, List, Int, Float
import traitlets
import numpy as np
from traittypes import Array
from collections import namedtuple
import warnings
from IPython.display import display, Javascript

# Cell
def flatIndex(i, j, n_col):
    return i * n_col + j

def arrayIndex(index, n_col):
        row = int(index / n_col)
        col = int(index % n_col)
        return (row, col)

# Cell
class Item(HasTraits):
    '''The <code>name</code> property is used by model collections and view collections
    in order to store and retrieve models and views by name'''
    def __init__(self, name=None, **kwargs):
        super().__init__(**kwargs)
        if name:
            self.name = name
        self._index = {}

    def _repr_pretty_(self, p, cycle):
        if hasattr(self, 'name'):
            name = self.name
        else:
            name = self.__class__.__name__
        with p.group(4, str(name) +'(', ')'):
            p.breakable()
            for trait, value in self._trait_values.items():
                p.text(trait + ': ')
                p.pretty(value)
                p.text(',')
                p.breakable()

# Cell
class Collection(Item):
    '''Supports 2D and 3D arrays. 3D arrays can be accessed with tuples but are stored in
    a list.'''

    children = Tuple()

    def __init__(self, children=None, **kwargs):
        super().__init__(**kwargs)
        if children:
            self.children += tuple(children)

    def __getitem__(self, key):
        if isinstance(key, int):
            return self._children[key]
        elif isinstance(key, slice):
            return [self._children[i] for i in range(*key.indices(self._size))]
        else:
            raise KeyError('Key must be an int or slice')

    def __iter__(self):
        return self._children.__iter__()

    def get_child_index(self):
        return self._size

    @observe('children')
    def _observe_collection_children(self, change):
        self._children = ()
        self._size = 0
        for child in change['new']:
            self._add_child(child)
        return tuple(self._children)

    #def add_collection(self, collection):
    #    self._collections[id(collection)] = collection
    #    self.children += collection.children

    def _add_child(self, child):
        '''Modify child and add child to list self._children'''
        child._index[id(self)] = self.get_child_index()
        #self._indices[id(child)] = child._index[id(self)]
        self._size += 1
        self._children += (child, )

    def get_links(self, attr):
        return [(child, attr) for child in self._children]

# Cell
class Selection(Item):
    '''A collection with selectable item.'''

    children = Tuple()
    _children = Tuple()
    selected_index = Int(default_value=None, allow_none=True)
    selected_name = Unicode(default_value=None, allow_none=True)

    def __init__(self, children, **kwargs):
        super().__init__(**kwargs)
        self._size = 0
        self._children = children

    def __getitem__(self, key):
        if isinstance(key, int):
            return self._children[key]
        elif isinstance(key, slice):
            return [self._children[i] for i in range(*key._indices(self._size))]
        elif isinstance(key, str):
            if key in self._indices:
                return self._children[self._indices[key]]
            else:
                raise KeyError('Key %s not found in children or collections' % str(key))

    def __iter__(self):
        return self._children.__iter__()

    #@validate('children')
    #TODO: children can only be changed to match selected name and index

    def add_child(self, child, name=None):
        if name:
            child.name = name
        self._children += (child, )

    @observe('_children')
    def _observe_selection_children(self, change):
        self._names = []
        self._indices = {} # name: index
        ret = tuple([self._add_child(child, i) for i, child in enumerate(change['new'])])
        if ret:
            self.selected_index = 0
            self.selected_name = ret[0].name
            self.children = (ret[0], )
        self._size = len(ret)
        return ret

    def _add_child(self, child, index):
        '''Modify child and add child to list self._children'''
        child.item_index = index
        if not hasattr(child, 'name'):
            child.name = child.__class__.__name__ + str(self._size)
        self._names.append(child.name)
        self._indices[child.name] = child.item_index
        return child

    @observe('selected_index')
    def _observe_selected_index(self, change):
        selected_index = change['new']
        if self._size:
            self.children = (self._children[selected_index], )
            self.selected_name = self._names[selected_index]
        return selected_index

    @observe('selected_name')
    def _observe_selected_name(self, change):
        selected_name = change['new']
        if self._size:
            self.selected_index = self._indices[selected_name]
        return selected_name

# Cell
class ItemList(Collection):
    '''A WidgetCollectin in which every child is the same type of Widget.
    Any child trait name listed in `trait_names` will be turned into an
    observable list, for example if 'value' is listed, a List trait 'value_list'
    will give the user access to a list of child values so that when the list
    is updated, so is the value of the child'''

    def __init__(self, children, traits=('value', ), **kwargs):
        super().__init__(**kwargs)
        self._traits = traits
        self.children = tuple(children)

    @observe('children')
    def _observe_collection_children(self, change):
        super()._observe_collection_children(change)
        if self._children and self._traits:
            for trait in self._traits:
                # initialize and observe child handler
                for child in self._children:
                    child._child_handlers[trait] = traitlets.observe(trait, type='change')
                    child._child_handlers[trait]._init_call(self._child_handler_factory(trait))
                    child.observe(child._child_handlers[trait], trait)

    def _add_child(self, child):
        child._child_handlers = {}
        super()._add_child(child)

    def set_list(self, values, trait='value'):
        if isinstance(values, np.ndarray):
            values = values.tolist()
        for i, value in enumerate(values):
            setattr(self._children[i], trait, value)

    def get_list(self, values, trait='value'):
        ret = []
        for i, child in enumerate(self._children):
            ret.append(getattr(child, trait))
        return ret

    def apply_all(self, method_name, *args, **kwargs):
        for child in self._children:
            method = getattr(child, method_name)
            method(*args, **kwargs)

    def observe_all(self, handler, names=None, type='change'):
        for child in self._children:
            if names:
                child.observe(handler, names, type)
            else:
                child.observe(handler)

    def unobserve_all(self, handler, names=None, type='change'):
        for child in self._children:
            if names:
                child.unobserve(handler, names, type)
            else:
                child.unobserve(handler)

    def observe_children(self, trait):
        for child in self._children:
            child.observe(child._child_handlers[trait], trait)

    def unobserve_children(self, trait):
        for child in self._children:
            child.unobserve(child._child_handlers[trait], trait)

    def _child_handler_factory(self, name):
        def wrapper(change):
            return self._child_handler(name, change)
        return wrapper

    def _child_handler(self, name, change):
        # add to change dict and notify change
        change['index'] = change['owner']._index[id(self)]
        change['type'] = 'child_change'
        change['list'] = self.get_list(name)
        self.notify_change(change)
        return change['new']

# Cell
class ItemArray(ItemList):
    '''A CollectionList with two dimensions. Children can be accessed
    with tuple notation similar to the numpy library. Any trait_name listed
    will cause the creation of an array of child traits, for example
    if 'value' is listed, the trait 'value_array' will become available'''
    def __init__(self, n_rows, n_columns, children=None, traits=('value',), **kwargs):
        self.n_rows = n_rows
        self.n_columns = n_columns
        super().__init__(children, traits, **kwargs)

    def get_child_index(self):
        return self._2D_index(self._size)

    def set_array(self, values, trait='value'):
        #self.unobserve_children(trait)
        if isinstance(values, np.ndarray):
            values = values.tolist()
        for i in range(self.n_rows):
            for j in range(self.n_columns):
                setattr(self._children[self._flat_index(i, j)], trait, values[i][j])
        #self.observe_children(trait)

    def get_array(self, trait='value'):
        ret = []
        for i in range(self.n_rows):
            ret.append([])
            for j in range(self.n_columns):
                ret[i].append(getattr(self._children[self._flat_index(i, j)], trait))
        return ret

    def __getitem__(self, index):
        if not isinstance(index, tuple):
            raise IndexError('Array index must be a tuple')
        x, y = index
        if isinstance(x, int) and isinstance(y, int):
            assert x <= self.n_rows and y <= self.n_columns, "index (%i, %i) not in range" % (i, j)
            return self._children[self._flat_index(x, y)]
        else:
            ret = []
            for x in self._to_range(x, self.n_rows):
                for y in self._to_range(y, self.n_columns):
                    ret.append(self._children[self._flat_index(x, y)])
            return ret

    def _to_range(self, s, end):
        if isinstance(s, int):
            return range(s, s+1)
        else:
            if hasattr(s, 'stop') and s.stop:
                assert s.stop <= end, "slice not in range"
                end = s.stop
            return range(*s.indices(end))

    def _flat_index(self, i, j):
        return i * self.n_columns + j

    def _2D_index(self, index):
        row = int(index / self.n_columns)
        col = int(index % self.n_columns)
        return (row, col)

    def _child_handler(self, name, change):
        # add to change dict and notify change
        change['index'] = change['owner']._index[id(self)]
        change['type'] = 'child_change'
        change['array'] = self.get_array(name)
        self.notify_change(change)
        return change['new']