import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence

class EncoderCNN(nn.Module):
    def __init__(self, embed_size):
        super(EncoderCNN, self).__init__()
        self.embed = nn.Linear(2048, embed_size)

    def forward(self, features):
        features = features.mean(1).unsqueeze(1)
        features = self.embed(features)
        return features

class DecoderRNN(nn.Module):
    __units = {"elman": nn.RNN, "lstm": nn.LSTM, "gru": nn.GRU}

    def __init__(self, embed_size, hidden_size, vocab_size, num_layers, unit='gru'):
        """
        Args:
            embed_size: embedding size of each word
            hidden_size: size of hidden weight in RNN
            vocab_size: size of vocab
            num_layers: layers of RNN
            unit: type of RNN used
        """
        super(DecoderRNN, self).__init__()
        self.embed_size = embed_size
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size
        self.num_layers = num_layers
        self.unit_name = unit
        # Model
        self.unit = DecoderRNN.__units[unit](embed_size, hidden_size, num_layers, batch_first=True)
        # word embedding
        self.embedding = nn.Embedding(num_embeddings=self.vocab_size, embedding_dim=embed_size, padding_idx=0)
        self.relu = nn.ReLU(inplace=True)
        self.linear = nn.Linear(in_features=self.hidden_size, out_features=self.vocab_size)

    def forward(self, features, captions, lengths):
        """
        Shapes:
            features: (B, Embed_size)
            captions: (B, L)
            pattern: Connect Method
        Return: (B, L, Vocab_size)
        """
        # features: (B, Embed_size) -> (B, 1, Num_features), caption_embed: (B, L, Embed_size)
        caption_embed = self.embedding(captions)
        # image_caption_embed: (B, L + 1, Embed_size)
        image_caption_embed = torch.cat((features, caption_embed), dim=1)
        # lengths: list of sequences lengths of each batch element
        inputs_packed = pack_padded_sequence(image_caption_embed, lengths, batch_first=True)
        outputs, _ = self.unit(inputs_packed)  # outputs[0]: (*, Hidden_size)
        outputs = self.relu(self.linear(outputs[0]))  # outputs[0]: (*, Vocab_size)
        return outputs

    def sample(self, features, max_len=25, search_mode="beam", beam_size=10):
        """
        Sample from RNN using greedy search or beam search
        Args:
            features: features from CNN feature extractor
            max_len: max length of sentences
            search_mode: sample mode, greedy search or beam search
            beam_size: beam search size
        Return:
            list of predicted image captions
        """
        output_ids = []
        states = None
        inputs = features
        if search_mode == "greedy":
            for i in range(max_len):
                # pass data through recurrent network
                outputs, states = self.unit(inputs, states)
                outputs = self.relu(self.linear(outputs.squeeze(1)))  # outputs: (1, vocab_size)
                # find maximal predictions
                predicted = outputs.max(1)[1]  # get maximal indices i.e. max()[1], specific [*]
                # append results from given step to global results
                output_ids.append(predicted)
                # prepare chosen words for next decoding step
                inputs = self.embedding(predicted)
                inputs = inputs.unsqueeze(1)  # inputs:(1, 1, embed_size)
            output_ids = torch.stack(output_ids, 1)  # output_ids: (1, max_len)
            return output_ids.squeeze().data.cpu().numpy()  # squeeze dim 0
        if search_mode == "beam":
            feeds = torch.ones((1, 1), device="cuda:0", dtype=torch.long)  # initial feed, feeds: (1, 1)
            hidden = None
            for t in range(max_len):
                if t == 0:  # t == 0 is specified for image input
                    outputs, hidden = self.unit(inputs)
                    outputs = self.relu(self.linear(outputs.squeeze(1)))

                    scores = F.softmax(outputs.squeeze(), -1)  # scores: (vocab_size)
                    scores = scores.log()
                    scores = scores.unsqueeze(0)  # scores: (1, vocab_size)
                    sorted_scores, idxes = scores.topk(beam_size, dim=1)  # sorted_scores, idxes: (1, beam_size)

                    candidates = idxes.unsqueeze(-1)  # candidates: (1, beam_size, 1)
                    feeds = torch.reshape(candidates, (-1, 1))  # feeds: (10, 1)

                    V = scores.size(-1)  # V: vocab_size
                    if self.unit_name == "gru":
                        # hidden: (1, beam_size, hidden_size)
                        hidden = torch.reshape(torch.cat([hidden.squeeze(0)] * beam_size, -1), (hidden.size(0), -1, hidden.size(-1)))
                    else:
                        # LSTM has h0 and c0
                        h, c = hidden
                        h = torch.reshape(torch.cat([h] * beam_size, -1), (h.size(0), -1, h.size(-1)))  # (L, N*B, H)
                        c = torch.reshape(torch.cat([c] * beam_size, -1), (c.size(0), -1, c.size(-1)))
                        hidden = (h, c)
                else:
                    # feeds: (10, 1) or (beam_size, 1) means 10 sentences which each one includes one word
                    caption_embed = self.embedding(feeds)  # caption_embed: (beam_size, 1, embed_size)
                    scores, hidden = self.unit(caption_embed, hidden)  # scores: (beam_size, 1, embed_size)
                    scores = self.relu(self.linear(scores))
                    scores = F.softmax(scores.squeeze(), -1)  # scores: (beam_size, vocab_size)
                    # scores: (beam_size, vocab_size), sorted_scores: (beam_size, 1) and broadcasting
                    scores = scores.log() + sorted_scores.view(-1, 1)  # scores: (beam_size, vocab_size)
                    scores = torch.reshape(scores, (1, -1))  # flatten scores to get index by // and %, scores: (1, beam_size * vocab_size)
                    sorted_scores, idxes = scores.topk(beam_size, dim=1)  # sorted_scores, idxes: (1, 10)
                    # Ref: https://cloud.tencent.com/developer/article/1675791
                    prior_candidate = candidates[np.arange(1)[:, None], idxes // V]  # prior_candidate: (1, beam_size, t)
                    current_candidate = (idxes % V).unsqueeze(-1)  # current_candidate: (1, beam_size, 1)
                    candidates = torch.cat([prior_candidate, current_candidate], -1)  # candidates: (1, beam_size, t + 1)
                    feeds = torch.reshape(candidates[:, :, -1], (-1, 1))  # feeds are for the next loop. feeds: (beam_size, 1)

            captions = candidates[:, 0, :].squeeze().data.cpu().numpy()  # pick best candidates
            return captions

class NIC(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers, unit='gru'):
        super(NIC, self).__init__()
        self.encoder = EncoderCNN(embed_size)
        self.decoder = DecoderRNN(embed_size, hidden_size, vocab_size, num_layers, unit)

    def forward(self, features, captions, lengths):
        features = self.encoder(features)
        outputs = self.decoder(features, captions, lengths)
        return outputs


from pynwb import register_class
from pynwb.file import NWBFile
from pynwb.icephys import IntracellularElectrode, PatchClampSeries
from pynwb.base import TimeSeries
import numpy as np
try:
    from pynwb.core import DynamicTable, DynamicTableRegion, VectorIndex, VectorData   # pragma: no cover
except ImportError:                                                        # pragma: no cover
    from hdmf.common import DynamicTable, DynamicTableRegion, VectorIndex  # pragma: no cover
from hdmf.utils import docval, popargs, getargs, call_docval_func, get_docval, fmt_docval_args
import warnings
import pandas as pd
from collections import OrderedDict
from copy import copy

namespace = 'ndx-icephys-meta'

class HierarchicalDynamicTableMixin:
    """
    Mixin class for defining specialized functionality for hierarchical dynamic tables.

    Assumptions:

    1) The current implementation assumes that there is only one DynamicTableRegion column
    that needs to be expanded as part of the hierarchy.  Allowing multiple hierarchical
    columns in a single table get tricky, because it is unclear how those rows should
    be joined. To clarify, allowing multiple DynamicTableRegion should be fine, as long
    as only one of them should be expanded as part of the hierarchy.

    2) The default implementation of the get_hierarchy_column_name function assumes that
    the first DynamicTableRegion that references a DynamicTable that inherits from
    HierarchicalDynamicTableMixin is the one that should be expanded as part of the
    hierarchy of tables. If there is no such column, then the default implementation
    assumes that the first DynamicTableRegion column is the one that needs to be expanded.
    These assumption of get_hierarchy_column_name can be easily fixed by overwriting
    the function in the subclass to return the name of the approbritate column.
    """

    def get_hierarchy_column_name(self):
        """
        Get the name of column that references another DynamicTable that
        is itself a HierarchicalDynamicTableMixin table.

        :returns: String with the column name or None
        """
        first_col = None
        for col_index, col in enumerate(self.columns):
            if isinstance(col, DynamicTableRegion):
                first_col = col.name
                if isinstance(col.table, HierarchicalDynamicTableMixin):
                    return col.name
        return first_col

    def get_referencing_column_names(self):
        """
        Determine the names of all columns that reference another table, i.e.,
        find all DynamicTableRegion type columns

        Returns: List of strings with the column names
        """
        col_names = []
        for col_index, col in enumerate(self.columns):
            if isinstance(col, DynamicTableRegion):
                col_names.append(col.name)
        return col_names

    def get_targets(self, include_self=False):
        """
        Get a list of the full table hierarchy, i.e., recursively list all
        tables referenced in the hierarchy.

        Returns: List of DynamicTable objects

        """
        hcol_name = self.get_hierarchy_column_name()
        hcol = self[hcol_name]
        hcol_target = hcol.table if isinstance(hcol, DynamicTableRegion) else hcol.target.table
        if isinstance(hcol_target, HierarchicalDynamicTableMixin):
            re = [self, ] if include_self else []
            re += [hcol_target, ]
            re += hcol_target.get_targets()
            return re
        else:
            return [hcol_target, ]

    def to_denormalized_dataframe(self, flat_column_index=False):
        """
        Shorthand for 'self.to_hierarchical_dataframe().reset_index()'

        The function denormalizes the hierarchical table and represents all data as
        columns in the resulting dataframe.
        """
        hier_df = self.to_hierarchical_dataframe(flat_column_index=True)
        flat_df = hier_df.reset_index()
        if not flat_column_index:
            # cn[0] is the level, cn[1:] is the label. If cn has only 2 elements than use cn[1] instead to
            # avoid creating column labels that are tuples with just one element
            mi_tuples = [(cn[0], cn[1:] if len(cn) > 2 else cn[1])
                         for cn in flat_df.columns]
            flat_df.columns = pd.MultiIndex.from_tuples(mi_tuples, names=('source_table', 'label'))

        return flat_df

    def to_hierarchical_dataframe(self, flat_column_index=False):
        """
        Create a Pandas dataframe with a hierarchical MultiIndex index that represents the
        hierarchical dynamic table.
        """
        # Get the references column
        hcol_name = self.get_hierarchy_column_name()
        hcol = self[hcol_name]
        hcol_target = hcol.table if isinstance(hcol, DynamicTableRegion) else hcol.target.table

        # Create the data variables we need to collect the data for our output dataframe and associated index
        index = []
        data = []
        columns = None
        index_names = None

        # If we have indexed columns (other than our hierarchical column) then our index data for our
        # MultiIndex will contain lists as elements (which are not hashable) and as such create an error.
        # As such we need to check if we have any affected columns so we can  fix our data
        indexed_column_indicies = np.where([isinstance(self[colname], VectorIndex)
                                            for colname in self.colnames if colname != hcol_name])[0]
        indexed_column_indicies += 1  # Need to increment by 1 since we add the row id in our iteration below

        # Case 1:  Our DynamicTableRegion column points to a regular DynamicTable
        #          If this is the case than we need to de-normalize the data and flatten the hierarchy
        if not isinstance(hcol_target, HierarchicalDynamicTableMixin):
            # 1) Iterate over all rows in our hierarchical columns (i.e,. the DynamicTableRegion column)
            for row_index, row_df in enumerate(hcol[:]):  # need hcol[:] here in case this is an h5py.Dataset
                # 1.1): Since hcol is a DynamicTableRegion, each row returns another DynamicTable so we
                #       next need to iterate over all rows in that table to denormalize our data
                for row in row_df.itertuples(index=True):
                    # 1.1.1) Determine the column data for our row. Each selected row from our target table
                    #        becomes a row in our flattened table
                    data.append(row)
                    # 1.1.2) Determine the multi-index tuple for our row, consisting of: i) id of the row in this
                    #        table, ii) all columns (except the hierarchical column we are flattening), and
                    #        iii) the index (i.e., id) from our target row
                    index_data = ([self.id[row_index], ] +
                                  [self[row_index, colname] for colname in self.colnames if colname != hcol_name])
                    for i in indexed_column_indicies:  # Fix data from indexed columns
                        index_data[i] = tuple(index_data[i])  # Convert from list to tuple (which is hashable)
                    index.append(tuple(index_data))
                    # Determine the names for our index and columns of our output table if this is the first row.
                    # These are constant for all rows so we only need to do this onle once for the first row.
                    if row_index == 0:
                        index_names = ([(self.name, 'id')] +
                                       [(self.name, colname)
                                        for colname in self.colnames if colname != hcol_name])
                        if flat_column_index:
                            columns = [(hcol_target.name, 'id'), ] + list(row_df.columns)
                        else:
                            columns = pd.MultiIndex.from_tuples([(hcol_target.name, 'id'), ] +
                                                                [(hcol_target.name, c) for c in row_df.columns],
                                                                names=('source_table', 'label'))
            #  if we had an empty data table then at least define the columns
            if index_names is None:
                index_names = ([(self.name, 'id')] +
                               [(self.name, colname)
                                for colname in self.colnames if colname != hcol_name])
                if flat_column_index:
                    columns = [(hcol_target.name, 'id'), ] + list(hcol_target.colnames)
                else:
                    columns = pd.MultiIndex.from_tuples([(hcol_target.name, 'id'), ] +
                                                        [(hcol_target.name, c) for c in hcol_target.colnames],
                                                        names=('source_table', 'label'))

        # Case 2:  Our DynamicTableRegion columns points to another HierarchicalDynamicTable.
        else:
            # 1) First we need to recursively flatten the hierarchy by calling 'to_hierarchical_dataframe()'
            #    (i.e., this function) on the target of our hierarchical column
            hcol_hdf = hcol_target.to_hierarchical_dataframe(flat_column_index=flat_column_index)
            # 2) Iterate over all rows in our hierarchcial columns (i.e,. the DynamicTableRegion column)
            for row_index, row_df_level1 in enumerate(hcol[:]):   # need hcol[:] here  in case this is an h5py.Dataset
                # 1.1): Since hcol is a DynamicTableRegion, each row returns another DynamicTable so we
                #       next need to iterate over all rows in that table to denormalize our data
                for row_df_level2 in row_df_level1.itertuples(index=True):
                    # 1.1.2) Since our target is itself a HierarchicalDynamicTable each target row itself
                    #        may expand into multiple rows in flattened hcol_hdf. So we now need to look
                    #        up the rows in hcol_hdf that correspond to the rows in row_df_level2.
                    #        NOTE: In this look-up we assume that the ids (and hence the index) of
                    #              each row in the table are in fact unique.
                    for row_tuple_level3 in hcol_hdf.loc[[row_df_level2[0]]].itertuples(index=True):
                        # 1.1.2.1) Determine the column data for our row.
                        data.append(row_tuple_level3[1:])
                        # 1.1.2.2) Determine the multi-index tuple for our row,
                        index_data = ([self.id[row_index], ] +
                                      [self[row_index, colname] for colname in self.colnames if colname != hcol_name] +
                                      list(row_tuple_level3[0]))
                        for i in indexed_column_indicies:  # Fix data from indexed columns
                            index_data[i] = tuple(index_data[i])  # Convert from list to tuple (which is hashable)
                        index.append(tuple(index_data))
                        # Determine the names for our index and columns of our output table if this is the first row
                        if row_index == 0:
                            index_names = ([(self.name, "id")] +
                                           [(self.name, colname)
                                            for colname in self.colnames if colname != hcol_name] +
                                           hcol_hdf.index.names)
                            columns = hcol_hdf.columns
            # if we had an empty table, then at least define the columns
            if index_names is None:
                index_names = ([(self.name, "id")] +
                               [(self.name, colname)
                                for colname in self.colnames if colname != hcol_name] +
                               hcol_hdf.index.names)
                columns = hcol_hdf.columns

        # Construct the pandas dataframe with the hierarchical multi-index
        multi_index = pd.MultiIndex.from_tuples(index, names=index_names)
        out_df = pd.DataFrame(data=data, index=multi_index, columns=columns)
        return out_df

@register_class('AlignedDynamicTable', namespace)
class AlignedDynamicTable(DynamicTable):
    """
    DynamicTable container that subports storing a collection of subtables. Each sub-table is a
    DynamicTable itself that is aligned with the main table by row index. I.e., all
    DynamicTables stored in this group MUST have the same number of rows. This type effectively
    defines a 2-level table in which the main data is stored in the main table implementd by this type
    and additional columns of the table are grouped into categories, with each category being'
    represented by a separate DynamicTable stored within the group.
    """
    __fields__ = ({'name': 'category_tables', 'child': True}, )

    @docval(*get_docval(DynamicTable.__init__),
            {'name': 'category_tables', 'type': list,
             'doc': 'List of DynamicTables to be added to the container', 'default': None},
            {'name': 'categories', 'type': 'array_data',
             'doc': 'List of names with the ordering of category tables', 'default': None})
    def __init__(self, **kwargs):
        in_category_tables = popargs('category_tables', kwargs)
        in_categories = popargs('categories', kwargs)
        if in_categories is None and in_category_tables is not None:
            in_categories = [tab.name for tab in in_category_tables]
        if in_categories is not None and in_category_tables is None:
            raise ValueError("Categories provided but no category_tables given")
        # at this point both in_categories and in_category_tables should either both be None or both be a list
        if in_categories is not None:
            if len(in_categories) != len(in_category_tables):
                raise ValueError("%s category_tables given but %s categories specified" %
                                 (len(in_category_tables), len(in_categories)))
        # Initialize the main dynamic table
        call_docval_func(super().__init__, kwargs)
        # Create and set all sub-categories
        dts = OrderedDict()
        # Add the custom categories given as inputs
        if in_category_tables is not None:
            # We may need to resize our main table when adding categories as the user may not have set ids
            if len(in_category_tables) > 0:
                # We have categories to process
                if len(self.id) == 0:
                    # The user did not initialize our main table id's nor set columns for our main table
                    for i in range(len(in_category_tables[0])):
                        self.id.append(i)
            # Add the user-provided categories in the correct order as described by the categories
            # This is necessary, because we do not store the categories explicitly but we maintain them
            # as the order of our self.category_tables. In this makes sure look-ups are consistent.
            lookup_index = OrderedDict([(k, -1) for k in in_categories])
            for i, v in enumerate(in_category_tables):
                # Error check that the name of the table is in our categories list
                if v.name not in lookup_index:
                    raise ValueError("DynamicTable %s does not appear in categories %s" % (v.name, str(in_categories)))
                # Error check to make sure no two tables with the same name are given
                if lookup_index[v.name] >= 0:
                    raise ValueError("Duplicate table name %s found in input dynamic_tables" % v.name)
                lookup_index[v.name] = i
            for table_name, tabel_index in lookup_index.items():
                # This error case should not be able to occur since the length of the in_categories and
                # in_category_tables must match and we made sure that each DynamicTable we added had its
                # name in the in_categories list. We, therefore, exclude this check from coverage testing
                # but we leave it in just as a backup trigger in case something unexpected happens
                if tabel_index < 0:  # pragma: no cover
                    raise ValueError("DynamicTable %s listed in categories but does not appear in category_tables" %
                                     table_name)  # pragma: no cover
                # Test that all category tables have the correct number of rows
                category = in_category_tables[tabel_index]
                if len(category) != len(self):
                    raise ValueError('Category DynamicTable %s does not align, it has %i rows expected %i' %
                                     (category.name, len(category), len(self)))
                # Add the category table to our category_tables.
                dts[category.name] = category
        # Set the self.category_tables attribute, which will set the parent/child relationships for the category_tables
        self.category_tables = dts

    def __contains__(self, val):
        """
        Check if the given value (i.e., column) exists in this table

        :param val: If val is a string then check if the given category exists. If val is a tuple
        of two strings (category, colname) then check for the given category if the given colname exists.
        """
        if isinstance(val, str):
            return val in self.category_tables or val in self.colnames
        elif isinstance(val, tuple):
            if len(val) != 2:
                raise ValueError("Expected tuple of strings of length 2 got tuple of length %i" % len(val))
            return val[1] in self.get_category(val[0])
        else:
            return False

    @property
    def categories(self):
        """
        Get the list of names the categories

        Short-hand for list(self.category_tables.keys())

        :raises: KeyError if the given name is not in self.category_tables
        """
        return list(self.category_tables.keys())

    @docval({'name': 'category', 'type': DynamicTable, 'doc': 'Add a new DynamicTable category'},)
    def add_category(self, **kwargs):
        """
        Add a new DynamicTable to the AlignedDynamicTable to create a new category in the table.

        NOTE: The table must align with (i.e, have the same number of rows as) the main data table (and
        other category tables). I.e., if the AlignedDynamicTable is already populated with data
        then we have to populate the new category with the corresponding data before adding it.

        :raises: ValueError is raised if the input table does not have the same number of rows as the main table
        """
        category = getargs('category', kwargs)
        if len(category) != len(self):
            raise ValueError('New category DynamicTable does not align, it has %i rows expected %i' %
                             (len(category), len(self)))
        if category.name in self.category_tables:
            raise ValueError("Category %s already in the table" % category.name)
        self.category_tables[category.name] = category
        category.parent = self

    @docval({'name': 'name', 'type': str, 'doc': 'Name of the category we want to retrieve', 'default': None})
    def get_category(self, **kwargs):
        name = popargs('name', kwargs)
        if name is None or (name not in self.category_tables and name == self.name):
            return self
        else:
            return self.category_tables[name]

    @docval(*get_docval(DynamicTable.add_column),
            {'name': 'category', 'type': str, 'doc': 'The category the column should be added to',
             'default': None})
    def add_column(self, **kwargs):
        """
        Add a column to the table

        :raises: KeyError if the category does not exist

        """
        category_name = popargs('category', kwargs)
        if category_name is None:
            # Add the column to our main table
            call_docval_func(super().add_column, kwargs)
        else:
            # Add the column to a sub-category table
            try:
                category = self.get_category(category_name)
            except KeyError:
                raise KeyError("Category %s not in table" % category_name)
            category.add_column(**kwargs)

    @docval({'name': 'data', 'type': dict, 'doc': 'the data to put in this row', 'default': None},
            {'name': 'id', 'type': int, 'doc': 'the ID for the row', 'default': None},
            {'name': 'enforce_unique_id', 'type': bool, 'doc': 'enforce that the id in the table must be unique',
             'default': False},
            allow_extra=True)
    def add_row(self, **kwargs):
        """
        We can either provide the row data as a single dict or by specifying a dict for each category
        """
        data, row_id, enforce_unique_id = popargs('data', 'id', 'enforce_unique_id', kwargs)
        data = data if data is not None else kwargs

        # extract the category data
        category_data = {k: data.pop(k) for k in self.categories if k in data}

        # Check that we have the approbriate categories provided
        missing_categories = set(self.categories) - set(list(category_data.keys()))
        if missing_categories:
            raise KeyError(
                '\n'.join([
                    'row data keys don\'t match available categories',
                    'missing {} category keys: {}'.format(len(missing_categories), missing_categories)
                ])
            )
        # Add the data to our main dynamic table
        data['id'] = row_id
        data['enforce_unique_id'] = enforce_unique_id
        call_docval_func(super().add_row, data)

        # Add the data to all out dynamic table categories
        for category, values in category_data.items():
            self.category_tables[category].add_row(**values)

    @docval({'name': 'ignore_category_ids', 'type': bool,
             'doc': "Ignore id columns of sub-category tables", 'default': False})
    def to_dataframe(self, **kwargs):
        """Convert the collection of tables to a single pandas DataFrame"""
        dfs = [super().to_dataframe().reset_index(), ]
        if getargs('ignore_category_ids', kwargs):
            dfs += [category.to_dataframe() for category in self.category_tables.values()]
        else:
            dfs += [category.to_dataframe().reset_index() for category in self.category_tables.values()]
        names = [self.name, ] + list(self.category_tables.keys())
        res = pd.concat(dfs, axis=1, keys=names)
        res.set_index((self.name, 'id'), drop=True, inplace=True)
        return res

    def __getitem__(self, item):
        """
        If item is:
        * int : Return a single row of the table
        * string : Return a single category of the table
        * tuple: Get a column, row, or cell from a particular category. The tuple is expected to consist
                 of (category, selection) where category may be a string with the name of the sub-category
                 or None (or the name of this AlignedDynamicTable) if we want to slice into the main table.

        :returns: DataFrame when retrieving a row or category. Returns scalar when selecting a cell.
                 Returns a VectorData/VectorIndex when retrieving a single column.
        """
        if isinstance(item, (int, list, np.ndarray, slice)):
            # get a single full row from all tables
            dfs = ([super().__getitem__(item).reset_index(), ] +
                   [category[item].reset_index() for category in self.category_tables.values()])
            names = [self.name, ] + list(self.category_tables.keys())
            res = pd.concat(dfs, axis=1, keys=names)
            res.set_index((self.name, 'id'), drop=True, inplace=True)
            return res
        elif isinstance(item, str) or item is None:
            if item in self.colnames:
                # get a specfic column
                return super().__getitem__(item)
            else:
                # get a single category
                return self.get_category(item).to_dataframe()
        elif isinstance(item, tuple):
            if len(item) == 2:
                return self.get_category(item[0])[item[1]]
            elif len(item) == 3:
                return self.get_category(item[0])[item[1]][item[2]]
            else:
                raise ValueError("Expected tuple of length 2 or 3 with (category, column, row) as value.")

@register_class('TimeSeriesReferenceVectorData', namespace)
class TimeSeriesReferenceVectorData(VectorData):
    """
    Column storing references to a TimeSeries (rows). For each TimeSeries this VectorData
    column stores the start_index and count to indicate the range in time to be selected
    as well as an object reference to the TimeSeries.
    """

    @docval(*get_docval(VectorData.__init__))
    def __init__(self, **kwargs):
        call_docval_func(super().__init__, kwargs)

@register_class('IntracellularElectrodesTable', namespace)
class IntracellularElectrodesTable(DynamicTable):
    """
    Table for storing intracellular electrode related metadata'
    """
    __columns__ = (
        {'name': 'electrode',
         'description': 'Column for storing the reference to the intracellular electrode',
         'required': True,
         'index': False,
         'table': False},
    )

    @docval(*get_docval(DynamicTable.__init__, 'id', 'columns', 'colnames'))
    def __init__(self, **kwargs):
        # Define defaultb name and description settings
        kwargs['name'] = 'electrodes'
        kwargs['description'] = ('Table for storing intracellular electrode related metadata')
        # Initialize the DynamicTable
        call_docval_func(super().__init__, kwargs)

@register_class('IntracellularStimuliTable', namespace)
class IntracellularStimuliTable(DynamicTable):
    """
    Table for storing intracellular electrode related metadata'
    """
    __columns__ = (
        {'name': 'stimulus',
         'description': 'Column storing the reference to the recorded stimulus for the recording (rows)',
         'required': True,
         'index': False,
         'table': False,
         'class': TimeSeriesReferenceVectorData},
    )

    @docval(*get_docval(DynamicTable.__init__, 'id', 'columns', 'colnames'))
    def __init__(self, **kwargs):
        # Define defaultb name and description settings
        kwargs['name'] = 'stimuli'
        kwargs['description'] = ('Table for storing intracellular stimulus related metadata')
        # Initialize the DynamicTable
        call_docval_func(super().__init__, kwargs)

@register_class('IntracellularResponsesTable', namespace)
class IntracellularResponsesTable(DynamicTable):
    """
    Table for storing intracellular electrode related metadata'
    """
    __columns__ = (
        {'name': 'response',
         'description': 'Column storing the reference to the recorded response for the recording (rows)',
         'required': True,
         'index': False,
         'table': False,
         'class': TimeSeriesReferenceVectorData},
    )

    @docval(*get_docval(DynamicTable.__init__, 'id', 'columns', 'colnames'))
    def __init__(self, **kwargs):
        # Define defaultb name and description settings
        kwargs['name'] = 'responses'
        kwargs['description'] = ('Table for storing intracellular response related metadata')
        # Initialize the DynamicTable
        call_docval_func(super().__init__, kwargs)

@register_class('IntracellularRecordingsTable', namespace)
class IntracellularRecordingsTable(AlignedDynamicTable):
    """
    A table to group together a stimulus and response from a single electrode and
    a single simultaneous_recording. Each row in the table represents a single recording consisting
    typically of a stimulus and a corresponding response.
    """
    @docval(*get_docval(AlignedDynamicTable.__init__, 'id', 'columns', 'colnames', 'category_tables', 'categories'))
    def __init__(self, **kwargs):
        kwargs['name'] = 'intracellular_recordings'
        kwargs['description'] = ('A table to group together a stimulus and response from a single electrode '
                                 'and a single simultaneous recording and for storing metadata about the '
                                 'intracellular recording.')
        in_category_tables = getargs('category_tables', kwargs)
        if in_category_tables is None or len(in_category_tables) == 0:
            kwargs['category_tables'] = [IntracellularElectrodesTable(),
                                         IntracellularStimuliTable(),
                                         IntracellularResponsesTable()]
            kwargs['categories'] = None
        else:
            # Check if our required data tables are supplied, otherwise add them to the list
            required_dynamic_table_given = [-1 for i in range(3)]  # The first three are our required tables
            for i, tab in enumerate(in_category_tables):
                if isinstance(tab, IntracellularElectrodesTable):
                    required_dynamic_table_given[0] = i
                elif isinstance(tab, IntracellularStimuliTable):
                    required_dynamic_table_given[1] = i
                elif isinstance(tab, IntracellularResponsesTable):
                    required_dynamic_table_given[2] = i
            # Check if the supplied tables contain data but not all required tables have been supplied
            required_dynamic_table_missing = np.any(np.array(required_dynamic_table_given[0:3]) < 0)
            if len(in_category_tables[0]) != 0 and required_dynamic_table_missing:
                raise ValueError("IntracellularElectrodeTable, IntracellularStimuliTable, and "
                                 "IntracellularResponsesTable are required when adding custom, non-empty "
                                 "tables to IntracellularRecordingsTable as the missing data for the required "
                                 "tables cannot be determined automatically")
            # Compile the complete list of tables
            dynamic_table_arg = copy(in_category_tables)
            categories_arg = [] if getargs('categories', kwargs) is None else copy(getargs('categories', kwargs))
            if required_dynamic_table_missing:
                if required_dynamic_table_given[2] < 0:
                    dynamic_table_arg.append(IntracellularResponsesTable)
                    if not dynamic_table_arg[-1].name in categories_arg:
                        categories_arg.insert(0, dynamic_table_arg[-1].name)
                if required_dynamic_table_given[1] < 0:
                    dynamic_table_arg.append(IntracellularStimuliTable())
                    if not dynamic_table_arg[-1].name in categories_arg:
                        categories_arg.insert(0, dynamic_table_arg[-1].name)
                if required_dynamic_table_given[0] < 0:
                    dynamic_table_arg.append(IntracellularElectrodesTable())
                    if not dynamic_table_arg[-1].name in categories_arg:
                        categories_arg.insert(0, dynamic_table_arg[-1].name)
            kwargs['category_tables'] = dynamic_table_arg
            kwargs['categories'] = categories_arg

        call_docval_func(super().__init__, kwargs)

    @docval({'name': 'electrode', 'type': IntracellularElectrode, 'doc': 'The intracellular electrode used'},
            {'name': 'stimulus_start_index', 'type': 'int', 'doc': 'Start index of the stimulus', 'default': -1},
            {'name': 'stimulus_index_count', 'type': 'int', 'doc': 'Stop index of the stimulus', 'default': -1},
            {'name': 'stimulus', 'type': TimeSeries,
             'doc': 'The TimeSeries (usually a PatchClampSeries) with the stimulus',
             'default': None},
            {'name': 'response_start_index', 'type': 'int', 'doc': 'Start index of the response', 'default': -1},
            {'name': 'response_index_count', 'type': 'int', 'doc': 'Stop index of the response', 'default': -1},
            {'name': 'response', 'type': TimeSeries,
             'doc': 'The TimeSeries (usually a PatchClampSeries) with the response',
             'default': None},
            {'name': 'electrode_metadata', 'type': dict,
             'doc': 'Additional electrode metadata to be stored in the electrodes table', 'default': None},
            {'name': 'stimulus_metadata', 'type': dict,
             'doc': 'Additional stimulus metadata to be stored in the stimuli table', 'default': None},
            {'name': 'response_metadata', 'type': dict,
             'doc': 'Additional resposnse metadata to be stored in the responses table', 'default': None},
            returns='Integer index of the row that was added to this table',
            rtype=int,
            allow_extra=True)
    def add_recording(self, **kwargs):
        """
        Add a single recording to the IntracellularRecordingsTable table.

        Typically, both stimulus and response are expected. However, in some cases only a stimulus
        or a resposne may be recodred as part of a recording. In this case, None, may be given
        for either stimulus or response, but not both. Internally, this results in both stimulus
        and response pointing to the same timeseries, while the start_index and index_count for
        the invalid series will both be set to -1.
        """
        # Get the input data
        stimulus_start_index, stimulus_index_count, stimulus = popargs('stimulus_start_index',
                                                                       'stimulus_index_count',
                                                                       'stimulus',
                                                                       kwargs)
        response_start_index, response_index_count, response = popargs('response_start_index',
                                                                       'response_index_count',
                                                                       'response',
                                                                       kwargs)
        electrode = popargs('electrode', kwargs)
        # Confirm that we have at least a valid stimulus or response
        if stimulus is None and response is None:
            raise ValueError("stimulus and response cannot both be None.")

        # Compute the start and stop index if necessary
        if stimulus is not None:
            stimulus_start_index, stimulus_index_count = self.__compute_index(stimulus_start_index,
                                                                              stimulus_index_count,
                                                                              stimulus, 'stimulus')
        if response is not None:
            response_start_index, response_index_count = self.__compute_index(response_start_index,
                                                                              response_index_count,
                                                                              response, 'response')

        # If either stimulus or response are None, then set them to the same TimeSeries to keep the I/O happy
        response = response if response is not None else stimulus
        stimulus = stimulus if stimulus is not None else response

        # Make sure the types are compatible
        if ((response.neurodata_type.startswith("CurrentClamp") and
                stimulus.neurodata_type.startswith("VoltageClamp")) or
                (response.neurodata_type.startswith("VoltageClamp") and
                 stimulus.neurodata_type.startswith("CurrentClamp"))):
            raise ValueError("Incompatible types given for 'stimulus' and 'response' parameters. "
                             "'stimulus' is of type %s and 'response' is of type %s." %
                             (stimulus.neurodata_type, response.neurodata_type))
        if response.neurodata_type == 'IZeroClampSeries':
            if stimulus is not None:
                raise ValueError("stimulus should usually be None for IZeroClampSeries response")
        if isinstance(response, PatchClampSeries) and isinstance(stimulus, PatchClampSeries):
            # # We could also check sweep_number, but since it is mostly relevant to the deprecated SweepTable
            # # we don't really need to enforce it here
            # if response.sweep_number != stimulus.sweep_number:
            #     warnings.warn("sweep_number are usually expected to be the same for PatchClampSeries type "
            #                   "stimulus and response pairs in an intracellular recording.")
            if response.electrode != stimulus.electrode:
                raise ValueError("electrodes are usually expected to be the same for PatchClampSeries type "
                                 "stimulus and response pairs in an intracellular recording.")

        # Compile the electrodes table data
        electrodes = copy(popargs('electrode_metadata', kwargs))
        if electrodes is None:
            electrodes = {}
        electrodes['electrode'] = electrode

        # Compile the stimuli table data
        stimuli = copy(popargs('stimulus_metadata', kwargs))
        if stimuli is None:
            stimuli = {}
        stimuli['stimulus'] = (stimulus_start_index, stimulus_index_count, stimulus)

        # Compile the reponses table data
        responses = copy(popargs('response_metadata', kwargs))
        if responses is None:
            responses = {}
        responses['response'] = (response_start_index, response_index_count, response)

        _ = super().add_row(enforce_unique_id=True,
                            electrodes=electrodes,
                            responses=responses,
                            stimuli=stimuli,
                            **kwargs)
        return len(self) - 1

    @staticmethod
    def __compute_index(start_index, index_count, time_series, name):
        start_index = start_index if start_index >= 0 else 0
        num_samples = time_series.num_samples
        index_count = (index_count
                       if index_count >= 0
                       else ((num_samples - start_index)
                             if num_samples is not None
                             else None))
        if index_count is None:
            raise IndexError("Invalid %s_index_count cannot be determined from %s data." % (name, name))
        if num_samples is not None:
            if start_index >= num_samples:
                raise IndexError("%s_start_index out of range" % name)
            if (start_index + index_count) > num_samples:
                raise IndexError("%s_start_index + %s_index_count out of range" % (name, name))
        return start_index, index_count

    @docval(*get_docval(AlignedDynamicTable.to_dataframe, 'ignore_category_ids'),
            {'name': 'electrode_refs_as_objectids', 'type': bool,
             'doc': 'replace object references in the electrode column with object_ids',
             'default': False},
            {'name': 'stimulus_refs_as_objectids', 'type': bool,
             'doc': 'replace object references in the stimulus column with object_ids',
             'default': False},
            {'name': 'response_refs_as_objectids', 'type': bool,
             'doc': 'replace object references in the response column with object_ids',
             'default': False}
            )
    def to_dataframe(self, **kwargs):
        """Convert the collection of tables to a single pandas DataFrame"""
        res = super().to_dataframe(ignore_category_ids=getargs('ignore_category_ids', kwargs))
        if getargs('electrode_refs_as_objectids', kwargs):
            res[('electrodes', 'electrode')] = [e.object_id for e in res[('electrodes', 'electrode')]]
        if getargs('stimulus_refs_as_objectids', kwargs):
            res[('stimuli', 'stimulus')] = [(e[0], e[1],  e[2].object_id) for e in res[('stimuli', 'stimulus')]]
        if getargs('response_refs_as_objectids', kwargs):
            res[('responses', 'response')] = [(e[0], e[1],  e[2].object_id) for e in res[('responses', 'response')]]
        return res

@register_class('SimultaneousRecordingsTable', namespace)
class SimultaneousRecordingsTable(HierarchicalDynamicTableMixin, DynamicTable):
    """
    A table for grouping different intracellular recordings from the
    IntracellularRecordingsTable table together that were recorded simultaneously
    from different electrodes.
    """

    __columns__ = (
        {'name': 'recordings',
         'description': 'Column with a references to one or more rows in the IntracellularRecordingsTable table',
         'required': True,
         'index': True,
         'table': True},
    )

    @docval({'name': 'intracellular_recordings_table',
             'type': IntracellularRecordingsTable,
             'doc': 'the IntracellularRecordingsTable table that the recordings column indexes. May be None when '
                    'reading the Container from file as the table attribute is already populated in this case '
                    'but otherwise this is required.',
             'default': None},
            *get_docval(DynamicTable.__init__, 'id', 'columns', 'colnames'))
    def __init__(self, **kwargs):
        intracellular_recordings_table = popargs('intracellular_recordings_table', kwargs)
        # Define default name and description settings
        kwargs['name'] = 'simultaneous_recordings'
        kwargs['description'] = ('A table for grouping different intracellular recordings from the'
                                 'IntracellularRecordingsTable table together that were recorded simultaneously '
                                 'from different electrodes.')
        # Initialize the DynamicTable
        call_docval_func(super().__init__, kwargs)
        if self['recordings'].target.table is None:
            if intracellular_recordings_table is not None:
                self['recordings'].target.table = intracellular_recordings_table
            else:
                raise ValueError("intracellular_recordings constructor argument required")

    @docval({'name': 'recordings',
             'type': 'array_data',
             'doc': 'the indices of the recordings belonging to this simultaneous recording'},
            returns='Integer index of the row that was added to this table',
            rtype=int,
            allow_extra=True)
    def add_simultaneous_recording(self, **kwargs):
        """
        Add a single Sweep consisting of one-or-more recordings and associated custom
        SimultaneousRecordingsTable metadata to the table.
        """
        _ = super().add_row(enforce_unique_id=True, **kwargs)
        return len(self.id) - 1

@register_class('SequentialRecordingsTable', namespace)
class SequentialRecordingsTable(HierarchicalDynamicTableMixin, DynamicTable):
    """
    A table for grouping different intracellular recording simultaneous_recordings from the
    SimultaneousRecordingsTable table together. This is typically used to group together simultaneous_recordings
    where the a sequence of stimuli of the same type with varying parameters
    have been presented in a sequence.
    """

    __columns__ = (
        {'name': 'simultaneous_recordings',
         'description': 'Column with a references to one or more rows in the SimultaneousRecordingsTable table',
         'required': True,
         'index': True,
         'table': True},
        {'name': 'stimulus_type',
         'description': 'Column storing the type of stimulus used for the sequential recording',
         'required': True,
         'index': False,
         'table': False}
    )

    @docval({'name': 'simultaneous_recordings_table',
             'type': SimultaneousRecordingsTable,
             'doc': 'the SimultaneousRecordingsTable table that the simultaneous_recordings '
                    'column indexes. May be None when reading the Container from file as the '
                    'table attribute is already populated in this case but otherwise this is required.',
             'default': None},
            *get_docval(DynamicTable.__init__, 'id', 'columns', 'colnames'))
    def __init__(self, **kwargs):
        simultaneous_recordings_table = popargs('simultaneous_recordings_table', kwargs)
        # Define defaultb name and description settings
        kwargs['name'] = 'sequential_recordings'
        kwargs['description'] = ('A table for grouping different intracellular recording simultaneous_recordings '
                                 'from the SimultaneousRecordingsTable table together. This is typically used to '
                                 'group together simultaneous_recordings where the a sequence of stimuli of the '
                                 'same type with varying parameters have been presented in a sequence.')
        # Initialize the DynamicTable
        call_docval_func(super().__init__, kwargs)
        if self['simultaneous_recordings'].target.table is None:
            if simultaneous_recordings_table is not None:
                self['simultaneous_recordings'].target.table = simultaneous_recordings_table
            else:
                raise ValueError('simultaneous_recordings_table constructor argument required')

    @docval({'name': 'stimulus_type',
             'type': str,
             'doc': 'the type of stimulus used for the sequential recording'},
            {'name': 'simultaneous_recordings',
             'type': 'array_data',
             'doc': 'the indices of the simultaneous_recordings belonging to this sequential recording'},
            returns='Integer index of the row that was added to this table',
            rtype=int,
            allow_extra=True)
    def add_sequential_recording(self, **kwargs):
        """
        Add a sequential recording (i.e., one row)  consisting of one-or-more recording simultaneous_recordings
        and associated custom sequential recording  metadata to the table.
        """
        _ = super().add_row(enforce_unique_id=True, **kwargs)
        return len(self.id) - 1

@register_class('RepetitionsTable', namespace)
class RepetitionsTable(HierarchicalDynamicTableMixin, DynamicTable):
    """
    A table for grouping different intracellular recording sequential recordings together.
    With each SweepSequence typically representing a particular type of stimulus, the
    RepetitionsTable table is typically used to group sets of stimuli applied in sequence.
    """

    __columns__ = (
        {'name': 'sequential_recordings',
         'description': 'Column with a references to one or more rows in the SequentialRecordingsTable table',
         'required': True,
         'index': True,
         'table': True},
    )

    @docval({'name': 'sequential_recordings_table',
             'type': SequentialRecordingsTable,
             'doc': 'the SequentialRecordingsTable table that the sequential_recordings column indexes. May '
                    'be None when reading the Container from file as the table attribute is already populated '
                    'in this case but otherwise this is required.',
             'default': None},
            *get_docval(DynamicTable.__init__, 'id', 'columns', 'colnames'))
    def __init__(self, **kwargs):
        sequential_recordings_table = popargs('sequential_recordings_table', kwargs)
        # Define default name and description settings
        kwargs['name'] = 'repetitions'
        kwargs['description'] = ('A table for grouping different intracellular recording sequential recordings '
                                 'together. With each SimultaneousRecording typically representing a particular type '
                                 'of stimulus, the RepetitionsTable table is typically used to group sets '
                                 'of stimuli applied in sequence.')
        # Initialize the DynamicTable
        call_docval_func(super().__init__, kwargs)
        if self['sequential_recordings'].target.table is None:
            if sequential_recordings_table is not None:
                self['sequential_recordings'].target.table = sequential_recordings_table
            else:
                raise ValueError('sequential_recordings_table constructor argument required')

    @docval({'name': 'sequential_recordings',
             'type': 'array_data',
             'doc': 'the indices of the sequential recordings belonging to this repetition',
             'default': None},
            returns='Integer index of the row that was added to this table',
            rtype=int,
            allow_extra=True)
    def add_repetition(self, **kwargs):
        """
        Add a repetition (i.e., one row)  consisting of one-or-more recording sequential recordings
        and associated custom repetition  metadata to the table.
        """
        _ = super().add_row(enforce_unique_id=True, **kwargs)
        return len(self.id) - 1

@register_class('ExperimentalConditionsTable', namespace)
class ExperimentalConditionsTable(HierarchicalDynamicTableMixin, DynamicTable):
    """
    A table for grouping different intracellular recording repetitions together that
    belong to the same experimental conditions.
    """

    __columns__ = (
        {'name': 'repetitions',
         'description': 'Column with a references to one or more rows in the RepetitionsTable table',
         'required': True,
         'index': True,
         'table': True},
    )

    @docval({'name': 'repetitions_table',
             'type': RepetitionsTable,
             'doc': 'the RepetitionsTable table that the repetitions column indexes',
             'default': None},
            *get_docval(DynamicTable.__init__, 'id', 'columns', 'colnames'))
    def __init__(self, **kwargs):
        repetitions_table = popargs('repetitions_table', kwargs)
        # Define default name and description settings
        kwargs['name'] = 'experimental_conditions'
        kwargs['description'] = ('A table for grouping different intracellular recording repetitions together that '
                                 'belong to the same experimental experimental_conditions.')
        # Initialize the DynamicTable
        call_docval_func(super().__init__, kwargs)
        if self['repetitions'].target.table is None:
            if repetitions_table is not None:
                self['repetitions'].target.table = repetitions_table
            else:
                raise ValueError('repetitions_table constructor argument required')

    @docval({'name': 'repetitions',
             'type': 'array_data',
             'doc': 'the indices of the repetitions  belonging to this condition',
             'default': None},
            returns='Integer index of the row that was added to this table',
            rtype=int,
            allow_extra=True)
    def add_experimental_condition(self, **kwargs):
        """
        Add a condition (i.e., one row)  consisting of one-or-more recording repetitions of sequential recordings
        and associated custom experimental_conditions  metadata to the table.
        """
        _ = super().add_row(enforce_unique_id=True, **kwargs)
        return len(self.id) - 1

@register_class('ICEphysFile', namespace)
class ICEphysFile(NWBFile):
    """
    Extension of the NWBFile class to allow placing the new icephys
    metadata types in /general/intracellular_ephys in the NWBFile
    NOTE: If this proposal for extension to NWB gets merged with
    the core schema, then this type would be removed and the
    NWBFile specification updated instead
    """

    __nwbfields__ = ({'name': 'intracellular_recordings',
                      'child': True,
                      'required_name': 'intracellular_recordings',
                      'doc': 'IntracellularRecordingsTable table to group together a stimulus and response '
                             'from a single intracellular electrode and a single simultaneous recording.'},
                     {'name': 'icephys_simultaneous_recordings',
                      'child': True,
                      'required_name': 'simultaneous_recordings',
                      'doc': 'SimultaneousRecordingsTable table for grouping different intracellular recordings from'
                             'the IntracellularRecordingsTable table together that were recorded simultaneously '
                             'from different electrodes'},
                     {'name': 'icephys_sequential_recordings',
                      'child': True,
                      'required_name': 'sequential_recordings',
                      'doc': 'A table for grouping different simultaneous intracellular recording from the '
                             'SimultaneousRecordingsTable table together. This is typically used to group '
                             'together simultaneous recordings where the a sequence of stimuli of the same '
                             'type with varying parameters have been presented in a sequence.'},
                     {'name': 'icephys_repetitions',
                      'child': True,
                      'required_name': 'repetitions',
                      'doc': 'A table for grouping different intracellular recording sequential recordings together.'
                             'With each SweepSequence typically representing a particular type of stimulus, the '
                             'RepetitionsTable table is typically used to group sets of stimuli applied in sequence.'},
                     {'name': 'icephys_experimental_conditions',
                      'child': True,
                      'required_name': 'experimental_conditions',
                      'doc': 'A table for grouping different intracellular recording repetitions together that '
                             'belong to the same experimental experimental_conditions.'},
                     )

    @docval(*get_docval(NWBFile.__init__),
            {'name': 'intracellular_recordings', 'type': IntracellularRecordingsTable, 'default': None,
             'doc': 'the IntracellularRecordingsTable table that belongs to this NWBFile'},
            {'name': 'icephys_simultaneous_recordings', 'type': SimultaneousRecordingsTable, 'default': None,
             'doc': 'the SimultaneousRecordingsTable table that belongs to this NWBFile'},
            {'name': 'icephys_sequential_recordings', 'type': SequentialRecordingsTable, 'default': None,
             'doc': 'the SequentialRecordingsTable table that belongs to this NWBFile'},
            {'name': 'icephys_repetitions', 'type': RepetitionsTable, 'default': None,
             'doc': 'the RepetitionsTable table that belongs to this NWBFile'},
            {'name': 'icephys_experimental_conditions', 'type': ExperimentalConditionsTable, 'default': None,
             'doc': 'the ExperimentalConditionsTable table that belongs to this NWBFile'},
            {'name': 'ic_filtering', 'type': str, 'default': None,
             'doc': '[DEPRECATED] Use IntracellularElectrode.filtering instead. Description of filtering used.'})
    def __init__(self, **kwargs):
        # Get the arguments to pass to NWBFile and remove arguments custum to this class
        intracellular_recordings = kwargs.pop('intracellular_recordings', None)
        icephys_simultaneous_recordings = kwargs.pop('icephys_simultaneous_recordings', None)
        icephys_sequential_recordings = kwargs.pop('icephys_sequential_recordings', None)
        icephys_repetitions = kwargs.pop('icephys_repetitions', None)
        icephys_experimental_conditions = kwargs.pop('icephys_experimental_conditions', None)
        if kwargs.get('sweep_table') is not None:
            warnings.warn("Use of SweepTable is deprecated. Use the intracellular_recordings, "
                          "simultaneous_recordings, sequential_recordings, repetitions and/or "
                          "experimental_conditions table(s) instead.", DeprecationWarning)
        # Initialize the NWBFile parent class
        pargs, pkwargs = fmt_docval_args(super().__init__, kwargs)
        super().__init__(*pargs, **pkwargs)
        # Set ic filtering if requested
        self.ic_filtering = kwargs.get('ic_filtering')
        # Set the intracellular_recordings if available
        setattr(self, 'intracellular_recordings', intracellular_recordings)
        setattr(self, 'icephys_simultaneous_recordings', icephys_simultaneous_recordings)
        setattr(self, 'icephys_sequential_recordings', icephys_sequential_recordings)
        setattr(self, 'icephys_repetitions', icephys_repetitions)
        setattr(self, 'icephys_experimental_conditions', icephys_experimental_conditions)

    @property
    def ic_filtering(self):
        return self.fields.get('ic_filtering')

    @ic_filtering.setter
    def ic_filtering(self, val):
        if val is not None:
            warnings.warn("Use of ic_filtering is deprecated. Use the IntracellularElectrode.filtering"
                          "field instead", DeprecationWarning)
            self.fields['ic_filtering'] = val

    @docval(*get_docval(NWBFile.add_stimulus),
            {'name': 'use_sweep_table', 'type': bool, 'default': False, 'doc': 'Use the deprecated SweepTable'})
    def add_stimulus(self, **kwargs):
        """
        Overwrite behavior from NWBFile to avoid use of the deprecated SweepTable
        """
        timeseries = popargs('timeseries', kwargs)
        self._add_stimulus_internal(timeseries)
        use_sweep_table = popargs('use_sweep_table', kwargs)
        if use_sweep_table:
            if self.sweep_table is None:
                warnings.warn("Use of SweepTable is deprecated. Use the IntracellularRecordingsTable, "
                              "SimultaneousRecordingsTable tables instead. See the add_intracellular_recordings, "
                              "add_icephsy_simultaneous_recording, add_icephys_sequential_recording, "
                              "add_icephys_repetition, add_icephys_condition functions.",
                              DeprecationWarning)
            self._update_sweep_table(timeseries)

    @docval(*get_docval(NWBFile.add_stimulus),
            {'name': 'use_sweep_table', 'type': bool, 'default': False, 'doc': 'Use the deprecated SweepTable'})
    def add_stimulus_template(self, **kwargs):
        """
        Overwrite behavior from NWBFile to avoid use of the deprecated SweepTable
        """
        timeseries = popargs('timeseries', kwargs)
        self._add_stimulus_template_internal(timeseries)
        use_sweep_table = popargs('use_sweep_table', kwargs)
        if use_sweep_table:
            if self.sweep_table is None:
                warnings.warn("Use of SweepTable is deprecated. Use the IntracellularRecordingsTable, "
                              "SimultaneousRecordingsTable tables instead. See the add_intracellular_recordings, "
                              "add_icephsy_simultaneous_recording, add_icephys_sequential_recording, "
                              "add_icephys_repetition, add_icephys_condition functions.",
                              DeprecationWarning)
            self._update_sweep_table(timeseries)

    @docval(*get_docval(NWBFile.add_acquisition),
            {'name': 'use_sweep_table', 'type': bool, 'default': False, 'doc': 'Use the deprecated SweepTable'})
    def add_acquisition(self, **kwargs):
        """
        Overwrite behavior from NWBFile to avoid use of the deprecated SweepTable
        """
        nwbdata = popargs('nwbdata', kwargs)
        self._add_acquisition_internal(nwbdata)
        use_sweep_table = popargs('use_sweep_table', kwargs)
        if use_sweep_table:
            if self.sweep_table is None:
                warnings.warn("Use of SweepTable is deprecated. Use the IntracellularRecordingsTable, "
                              "SimultaneousRecordingsTable tables instead. See the add_intracellular_recordings, "
                              "add_icephsy_simultaneous_recording, add_icephys_sequential_recording, "
                              "add_icephys_repetition, add_icephys_condition functions.",
                              DeprecationWarning)
            self._update_sweep_table(nwbdata)

    @docval(returns='The NWBFile.intracellular_recordings table', rtype=IntracellularRecordingsTable)
    def get_intracellular_recordings(self):
        """
        Get the NWBFile.intracellular_recordings table.

        In contrast to NWBFile.intracellular_recordings, this function will create the
        IntracellularRecordingsTable table if not yet done, whereas NWBFile.intracellular_recordings
        will return None if the table is currently not being used.
        """
        if self.intracellular_recordings is None:
            self.intracellular_recordings = IntracellularRecordingsTable()
        return self.intracellular_recordings

    @docval(*get_docval(IntracellularRecordingsTable.add_recording),
            returns='Integer index of the row that was added to IntracellularRecordingsTable',
            rtype=int,
            allow_extra=True)
    def add_intracellular_recording(self, **kwargs):
        """
        Add a intracellular recording to the intracellular_recordings table. If the
        electrode, stimulus, and/or response do not exsist yet in the NWBFile, then
        they will be added to this NWBFile before adding them to the table.
        """
        # Add the stimulus, response, and electrode to the file if they don't exist yet
        stimulus, response, electrode = getargs('stimulus', 'response', 'electrode', kwargs)
        if (stimulus is not None and
                (stimulus.name not in self.stimulus and
                 stimulus.name not in self.stimulus_template)):
            self.add_stimulus(stimulus, use_sweep_table=False)
        if response is not None and response.name not in self.acquisition:
            self.add_acquisition(response, use_sweep_table=False)
        if electrode is not None and electrode.name not in self.icephys_electrodes:
            self.add_icephys_electrode(electrode)
        # make sure the intracellular recordings table exists and if not create it using get_intracellular_recordings
        # Add the recoding to the intracellular_recordings table
        return call_docval_func(self.get_intracellular_recordings().add_recording, kwargs)

    @docval(returns='The NWBFile.icephys_simultaneous_recordings table', rtype=SimultaneousRecordingsTable)
    def get_icephys_simultaneous_recordings(self):
        """
        Get the NWBFile.icephys_simultaneous_recordings table.

        In contrast to NWBFile.icephys_simultaneous_recordings, this function will create the
        SimultaneousRecordingsTable table if not yet done, whereas NWBFile.icephys_simultaneous_recordings
        will return None if the table is currently not being used.
        """
        if self.icephys_simultaneous_recordings is None:
            self.icephys_simultaneous_recordings = SimultaneousRecordingsTable(self.get_intracellular_recordings())
        return self.icephys_simultaneous_recordings

    @docval(*get_docval(SimultaneousRecordingsTable.add_simultaneous_recording),
            returns='Integer index of the row that was added to SimultaneousRecordingsTable',
            rtype=int,
            allow_extra=True)
    def add_icephys_simultaneous_recording(self, **kwargs):
        """
        Add a new simultaneous recording to the icephys_simultaneous_recordings table
        """
        return call_docval_func(self.get_icephys_simultaneous_recordings().add_simultaneous_recording, kwargs)

    @docval(returns='The NWBFile.icephys_sequential_recordings table', rtype=SequentialRecordingsTable)
    def get_icephys_sequential_recordings(self):
        """
        Get the NWBFile.icephys_sequential_recordings table.

        In contrast to NWBFile.icephys_sequential_recordings, this function will create the
        IntracellularRecordingsTable table if not yet done, whereas NWBFile.icephys_sequential_recordings
        will return None if the table is currently not being used.
        """
        if self.icephys_sequential_recordings is None:
            self.icephys_sequential_recordings = SequentialRecordingsTable(self.get_icephys_simultaneous_recordings())
        return self.icephys_sequential_recordings

    @docval(*get_docval(SequentialRecordingsTable.add_sequential_recording),
            returns='Integer index of the row that was added to SequentialRecordingsTable',
            rtype=int,
            allow_extra=True)
    def add_icephys_sequential_recording(self, **kwargs):
        """
        Add a new sequential recording to the icephys_sequential_recordings table
        """
        self.get_icephys_sequential_recordings()
        return call_docval_func(self.icephys_sequential_recordings.add_sequential_recording, kwargs)

    @docval(returns='The NWBFile.icephys_repetitions table', rtype=RepetitionsTable)
    def get_icephys_repetitions(self):
        """
        Get the NWBFile.icephys_repetitions table.

        In contrast to NWBFile.icephys_repetitions, this function will create the
        RepetitionsTable table if not yet done, whereas NWBFile.icephys_repetitions
        will return None if the table is currently not being used.
        """
        if self.icephys_repetitions is None:
            self.icephys_repetitions = RepetitionsTable(self.get_icephys_sequential_recordings())
        return self.icephys_repetitions

    @docval(*get_docval(RepetitionsTable.add_repetition),
            returns='Integer index of the row that was added to RepetitionsTable',
            rtype=int,
            allow_extra=True)
    def add_icephys_repetition(self, **kwargs):
        """
        Add a new repetition to the RepetitionsTable table
        """
        return call_docval_func(self.get_icephys_repetitions().add_repetition, kwargs)

    @docval(returns='The NWBFile.icephys_experimental_conditions table', rtype=ExperimentalConditionsTable)
    def get_icephys_experimental_conditions(self):
        """
        Get the NWBFile.icephys_experimental_conditions table.

        In contrast to NWBFile.icephys_experimental_conditions, this function will create the
        RepetitionsTable table if not yet done, whereas NWBFile.icephys_experimental_conditions
        will return None if the table is currently not being used.
        """
        if self.icephys_experimental_conditions is None:
            self.icephys_experimental_conditions = ExperimentalConditionsTable(self.get_icephys_repetitions())
        return self.icephys_experimental_conditions

    @docval(*get_docval(ExperimentalConditionsTable.add_experimental_condition),
            returns='Integer index of the row that was added to ExperimentalConditionsTable',
            rtype=int,
            allow_extra=True)
    def add_icephys_experimental_condition(self, **kwargs):
        """
        Add a new condition to the ExperimentalConditionsTable table
        """
        return call_docval_func(self.get_icephys_experimental_conditions().add_experimental_condition, kwargs)

    def get_icephys_meta_parent_table(self):
        """
        Get the top-most table in the intracellular ephys metadata table hierarchy that exists in this NWBFile.

        The intracellular ephys metadata consists of a hierarchy of DynamicTables, i.e.,
        experimental_conditions --> repetitions --> sequential_recordings -->
        simultaneous_recordings --> intracellular_recordings etc.
        In a given NWBFile not all tables may exist. This convenience functions returns the top-most
        table that exists in this file. E.g., if the file contains only the simultaneous_recordings
        and intracellular_recordings tables then the function would return the simultaneous_recordings table.
        Similarly, if the file contains all tables then it will return the experimental_conditions table.

        :returns: DynamicTable object or None
        """
        if self.icephys_experimental_conditions is not None:
            return self.icephys_experimental_conditions
        elif self.icephys_repetitions is not None:
            return self.icephys_repetitions
        elif self.icephys_sequential_recordings is not None:
            return self.icephys_sequential_recordings
        elif self.icephys_simultaneous_recordings is not None:
            return self.icephys_simultaneous_recordings
        elif self.intracellular_recordings is not None:
            return self.intracellular_recordings
        else:
            return None


# ------------------------------------------
# DESCRIPTION OF MODULE
# ------------------------------------------

"""
CamPong Score class

It stores the score of a bat and may be used to display the
corresponding text on the screen
"""

# ------------------------------------------
# IMPORTS
# ------------------------------------------

try:
    import os
    os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
    import pygame
    import math
    from Constants import *
except ImportError as err:
    print ("Error: couldn't load module" + str(err) + ". Exiting...")
    exit()

# ------------------------------------------
# CLASSES DEFINITIONS
# ------------------------------------------

class Score(pygame.sprite.Sprite):
    """
    Attributes:
        - font:Font
        - text:String
        - image:Surface
        - rect:Rect. Determines position and hitbox
        - position:(Int,Int). (X,Y)
        - area:Rect. Screen the object is moving across
        - sprite:Sprite
    Methods: __init__, update, calcnewpos
    """

    TEXT_SIZE = 20

    def __init__(self, position_init):
        pygame.sprite.Sprite.__init__(self)
        self.font = pygame.font.Font(None, 36)
        self.text = "0 - 0"
        # image (appearance)
        self.image = self.font.render(self.text, True, COLOUR_WHITE)
        # rect (hitbox)
        self.rect = self.image.get_rect()
        self.position = position_init
        self.rect.center = self.position
        pygame.draw.rect(self.image, COLOUR_WHITE, self.rect)
        # area that contains the text (i.e. the board itself)
        screen = pygame.display.get_surface()
        self.area = screen.get_rect()
        # sprite
        self.sprite = pygame.sprite.RenderPlain(self)

    def update(self, score1, score2, screen, background):
        self.text = str(score1) + " - " + str(score2)
        self.image = self.font.render(self.text, True, COLOUR_WHITE)
        self.rect = self.image.get_rect()
        self.rect.center = self.position
        screen.blit(background, self.rect, self.rect)
        self.sprite.draw(screen)

    def calcnewpos(self,rect,vector):
        (angle,z) = vector
        (dx,dy) = (z*math.cos(angle),z*math.sin(angle))
        return rect.move(dx,dy)


#!/usr/bin/env python

from __future__ import print_function
import yaml
import copy
import os
import json

def add_to_yaml(a, b, ignore_keys=set(("name", "description", "PIPELINE"))):
    '''
    If an object is in b but not a, add it to a.

    Fail if confilicting values are found
    '''
    if type(a) != type(b):
        raise ValueError("Incomptable yaml entries found")
    if type(a) is dict:
        for k in b.keys():
            if k in ignore_keys:
                continue
            if k in a:
                add_to_yaml(a[k], b[k])
            else:
                a[k] = b[k]
    elif type(a) is list:  # In this case, order is unimportant
        try:  # Only applicable with zones.
            a_set, b_set = set(a), set(b)
            intersection = list(a_set.intersection(b_set))
            for i in range(len(a) - 1, -1, -1):
                if a[i] not in intersection:
                    a.pop(i)
            if len(a) == 0:
                raise ValueError("No values left after intersection")
            return
        except TypeError:
            pass
        # item["name"] essentially acts as a key
        a_dict = dict([(x["name"], x) for x in a])
        b_dict = dict([(x["name"], x) for x in b])
        for k, v in b_dict.items():
            if k in ignore_keys:
                continue
            if k in a_dict:
                if v != a_dict[k]:
                    print("'{}' does not equal '{}'".format(v, a_dict[k]))
                    raise ValueError("Paramters are unequal")
            else:
                a.append(v)
    else:
        if a != b:
            print("'{}' does not equal '{}'".format(a, b))
            raise ValueError("Values are unequal")

script_dir = os.path.dirname(os.path.realpath(__file__))
germline = yaml.load(open(script_dir + "/germline.yaml"))
tn = None
try:
    tn = yaml.load(open(script_dir + "/somatic.yaml"))
except IOError:
    pass
output = script_dir + "/runner_default.json"

additional_input_params = {
        "ZONES": None,
        "DISK_SIZE": 300,
        "MIN_CPU": 64,
        "MIN_RAM_GB": 56,
        "PIPELINE": "DNA",
        "PROJECT_ID": None,
        "DOCKER_IMAGE": "sentieon/sentieon-google-cloud:0.2.0",
        "PREEMPTIBLE_TRIES": 0,
        "NONPREEMPTIBLE_TRY": True
}

out_yaml = copy.deepcopy(germline)
if tn:
    add_to_yaml(out_yaml, tn)
out_json = {}
for param in out_yaml["inputParameters"]:
    k = param["name"]
    if "defaultValue" not in param:
        v = None
    elif param["defaultValue"] == "None":
        v = None
    else:
        v = param["defaultValue"]
    out_json[k] = v
out_json.update(additional_input_params)

with open(output, 'w') as f:
    json.dump(out_json, f, indent=2)


import tqdm
import numpy as np
import csv
import six
from collections import OrderedDict
from collections import Iterable

import keras
from keras.callbacks import ModelCheckpoint
from keras.callbacks import TensorBoard
from keras.callbacks import CSVLogger
from keras.callbacks import LearningRateScheduler

class Evaluate(keras.callbacks.Callback):
    def __init__(
        self,
        generator,
        eval_steps,
        is_patch,
        verbose=1):

        super(Evaluate, self).__init__()
        self.generator = generator
        self.eval_steps = eval_steps
        self.is_patch = is_patch
        self.verbose = verbose

        self.tot_dice = 0.
        self.pneu_dice = 0.

    def on_epoch_end(self, epoch, logs=None):
        if epoch % 5 == 0:
            logs = logs or {}
            self.tot_dice, self.pneu_dice = self.evaluate()

            if self.verbose == 1:
                print('Epoch {:4d} - Dice score : {:.4f} | Pneu score : {:.4f}'.format(epoch+1, self.tot_dice, self.pneu_dice))

        logs['eval_dice'] = self.tot_dice
        logs['eval_dice_pneu'] = self.pneu_dice

    def dice(self, y_true, y_pred, classes, smooth=1.):
        loss = 0.
        y_pred = np.clip(y_pred, 1e-7, 1-1e-7)
        if classes > 1:
            for num_label in range(classes):
                y_true_f = y_true[...,num_label].flatten()
                y_pred_f = y_pred[...,num_label].flatten()
                intersection = np.sum(y_true_f * y_pred_f)
                loss += (2. * intersection + smooth) / (np.sum(y_true_f) + np.sum(y_pred_f) + smooth)

        else:
            y_true_f = y_true.flatten()
            y_pred_f = y_pred.flatten()
            intersection = np.sum(y_true_f * y_pred_f)
            loss += (2. * intersection + smooth) / (np.sum(y_true_f) + np.sum(y_pred_f) + smooth)

        return loss / classes

    def onehot(self, x, classes):
        result = np.zeros(x.shape+(classes,))
        for i in range(classes):
            result[...,i][np.where(x == i)] = 1.
        return result

    def evaluate(self, classes=2):
        tot_dice = []
        pneu_dice = []
        for i in tqdm.trange(self.eval_steps):
            y_true = np.zeros((1024, 1024, 2))
            y_pred = np.zeros((1024, 1024, 2))
            loc_cnt = np.zeros((1024, 1024, 2))

            g = next(self.generator)
            result = self.model.predict_on_batch(g[0])

            if self.is_patch:
                ratio = 1024//result.shape[1]
                for l in range(result.shape[0]):
                    # result = model.predict_on_batch(g[0][l][np.newaxis,...])
                    y_true[(l//(2*ratio-1))*1024//(ratio*2):(l//(2*ratio-1)+2)*1024//(ratio*2),
                           (l%(2*ratio-1))*1024//(ratio*2):(l%(2*ratio-1)+2)*1024//(ratio*2)] += g[1][l]

                    y_pred[(l//(2*ratio-1))*1024//(ratio*2):(l//(2*ratio-1)+2)*1024//(ratio*2),
                           (l%(2*ratio-1))*1024//(ratio*2):(l%(2*ratio-1)+2)*1024//(ratio*2)] += result[l]

                    loc_cnt[(l//(2*ratio-1))*1024//(ratio*2):(l//(2*ratio-1)+2)*1024//(ratio*2),
                            (l%(2*ratio-1))*1024//(ratio*2):(l%(2*ratio-1)+2)*1024//(ratio*2)] += 1

                y_true /= loc_cnt
                y_pred /= loc_cnt

            else:
                y_pred = result

            y_pred = self.onehot(np.argmax(y_pred, axis=-1), classes)
            tot_dice.append(self.dice(y_true, y_pred, classes))
            pneu_dice.append(self.dice(y_true[...,1], y_pred[...,1], 1))

        return sum(tot_dice)/eval_steps, sum(pneu_dice)/eval_steps

class CustomBatchLogger(CSVLogger):
    def __init__(self, filename, separator=',', append=False):
        super().__init__(filename, separator, append)

    def on_train_begin(self, logs=None):
        super().on_train_begin(logs)

    def on_batch_end(self, batch, logs=None):
        logs = logs or {}

        def handle_value(k):
            is_zero_dim_ndarray = isinstance(k, np.ndarray) and k.ndim == 0
            if isinstance(k, six.string_types):
                return k
            elif isinstance(k, Iterable) and not is_zero_dim_ndarray:
                return '"[%s]"' % (', '.join(map(str, k)))
            else:
                return k

        if self.keys is None:
            self.keys = sorted(logs.keys())

        if self.model.stop_training:
            # We set NA so that csv parsers do not fail for this last batch.
            logs = dict([(k, logs[k] if k in logs else 'NA') for k in self.keys])

        if not self.writer:
            class CustomDialect(csv.excel):
                delimiter = self.sep
            fieldnames = ['batch'] + self.keys
            if six.PY2:
                fieldnames = [unicode(x) for x in fieldnames]
            self.writer = csv.DictWriter(self.csv_file,
                                         fieldnames=fieldnames,
                                         dialect=CustomDialect)
            if self.append_header:
                self.writer.writeheader()

        row_dict = OrderedDict({'batch': batch})
        row_dict.update((key, handle_value(logs[key])) for key in self.keys)
        self.writer.writerow(row_dict)
        self.csv_file.flush()

    def on_epoch_end(self, epoch, logs=None):
        pass

    def on_train_end(self, logs=None):
        super().on_train_end(logs)

def callback_checkpoint(filepath, monitor, verbose, mode, save_best_only, save_weights_only):
    return ModelCheckpoint(filepath=filepath,
                           monitor=monitor,
                           verbose=verbose,
                           mode=mode,
                           save_best_only=save_best_only,
                           save_weights_only=save_weights_only)

def callback_epochlogger(filename, separator, append):
    return CSVLogger(filename=filename,
                     separator=separator,
                     append=append)

def callback_batchlogger(filename, separator, append):
    return CustomBatchLogger(filename=filename,
                             separator=separator,
                             append=append)

def callback_tensorboard(log_dir, batch_size):
    return TensorBoard(log_dir=log_dir, batch_size=batch_size)

def main_schedule(initlr, mode, value, duration, total_epoch):
    def _cosine_anneal_schedule(epoch, total_epoch):
        total_epoch = np.maximum(total_epoch, 50)
        cos_inner = np.pi * (epoch % total_epoch)
        cos_inner /= total_epoch
        cos_out = np.cos(cos_inner) + 1
        return float(initlr / 2 * cos_out)

    def _exponential(epoch, value, duration):
        update = epoch // duration
        return initlr * (value ** update)

    def _schedule(epoch, lr):
        if mode == 'exponential':
            return _exponential(epoch, value, duration)
        elif mode == 'cosine':
            return _cosine_anneal_schedule(epoch, total_epoch)

    return _schedule

def callback_learningrate(initlr, mode, value, duration, total_epoch):
    return LearningRateScheduler(schedule=main_schedule(initlr=initlr,
                                                        mode=mode,
                                                        value=value,
                                                        duration=duration,
                                                        total_epoch=total_epoch), verbose=1)

def callback_evaluate(generator, eval_steps, is_patch):
    return Evaluate(generator=generator, eval_steps=eval_steps, is_patch=is_patch, verbose=1)


submit = True
from collections import Counter
from aocd import data
import u

img = [int(c) for c in data.strip()]

def a():
    w = 25
    h = 6
    best_zero = None
    res = None
    for layer in u.chunks(img, w * h):
        counts = Counter(layer)
        if best_zero is None or counts[0] < best_zero:
            res = counts[1] * counts[2]
            best_zero = counts[0]
    return res

def b():
    w = 25
    h = 6
    out = [None] * (w * h)
    layers = list(u.chunks(img, w * h))
    for i in range(w * h):
        for layer in layers:
            if layer[i] != 2:
                out[i] = layer[i]
                break
    m = {
        (x, y): row[x] for y, row in enumerate(list(u.chunks(out, w))) for x in range(w)
    }
    u.print_matrix(m, " #")

u.main(a, b, submit=globals().get("submit", False))