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

__all__ = ['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

# 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='values', **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'):
                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']