text: #!/usr/bin/env python
  
  # Copyright (c) 2019, IRIS-HEP
  # All rights reserved.
  # 
  # Redistribution and use in source and binary forms, with or without
  # modification, are permitted provided that the following conditions are met:
  # 
  # * Redistributions of source code must retain the above copyright notice, this
  #   list of conditions and the following disclaimer.
  # 
  # * Redistributions in binary form must reproduce the above copyright notice,
  #   this list of conditions and the following disclaimer in the documentation
  #   and/or other materials provided with the distribution.
  # 
  # * Neither the name of the copyright holder nor the names of its
  #   contributors may be used to endorse or promote products derived from
  #   this software without specific prior written permission.
  # 
  # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  import collections
  import numbers
  
  import awkward.array.base
  import awkward.type
  import awkward.util
  
  class MaskedArray(awkward.array.base.AwkwardArrayWithContent):
      """
      MaskedArray
      """
  
      masked = None
  
      def __init__(self, mask, content, maskedwhen=True):
          self.mask = mask
          self.content = content
          self.maskedwhen = maskedwhen
  
      def copy(self, mask=None, content=None, maskedwhen=None):
          out = self.__class__.__new__(self.__class__)
          out._mask = self._mask
          out._content = self._content
          out._maskedwhen = self._maskedwhen
          out._isvalid = self._isvalid
          if mask is not None:
              out.mask = mask
          if content is not None:
              out.content = content
          if maskedwhen is not None:
              out.maskedwhen = maskedwhen
          return out
  
      def deepcopy(self, mask=None, content=None):
          out = self.copy(mask=mask, content=content)
          out._mask = self._util_deepcopy(out._mask)
          out._content = self._util_deepcopy(out._content)
          return out
  
      def _mine(self, overrides):
          mine = {}
          mine["maskedwhen"] = overrides.pop("maskedwhen", self._maskedwhen)
          return mine
  
      def empty_like(self, **overrides):
          mine = self._mine(overrides)
          if isinstance(self._content, self.numpy.ndarray):
              return self.copy(content=self.numpy.empty_like(self._content), **mine)
          else:
              return self.copy(content=self._content.empty_like(**overrides), **mine)
  
      def zeros_like(self, **overrides):
          mine = self._mine(overrides)
          if isinstance(self._content, self.numpy.ndarray):
              return self.copy(content=self.numpy.zeros_like(self._content), **mine)
          else:
              return self.copy(content=self._content.zeros_like(**overrides), **mine)
  
      def ones_like(self, **overrides):
          mine = self._mine(overrides)
          if isinstance(self._content, self.numpy.ndarray):
              return self.copy(content=self.numpy.ones_like(self._content), **mine)
          else:
              return self.copy(content=self._content.ones_like(**overrides), **mine)
  
      def __awkward_persist__(self, ident, fill, prefix, suffix, schemasuffix, storage, compression, **kwargs):
          self._valid()
          return {"id": ident,
                  "call": ["awkward", "MaskedArray"],
                  "args": [fill(self._mask, "MaskedArray.mask", prefix, suffix, schemasuffix, storage, compression, **kwargs),
                           fill(self._content, "MaskedArray.content", prefix, suffix, schemasuffix, storage, compression, **kwargs),
                           {"json": bool(self._maskedwhen)}]}
  
      @property
      def mask(self):
          return self._mask
  
      @mask.setter
      def mask(self, value):
          value = self._util_toarray(value, self.MASKTYPE, self.numpy.ndarray)
          if self.check_prop_valid:
              if len(value.shape) != 1:
                  raise ValueError("mask must have 1-dimensional shape")
          if not issubclass(value.dtype.type, (self.numpy.bool_, self.numpy.bool)):
              value = (value != 0)
          self._mask = value
          self._isvalid = False
  
      def boolmask(self, maskedwhen=None):
          if maskedwhen is None:
              maskedwhen = self._maskedwhen
          if maskedwhen == self._maskedwhen:
              return self._mask
          else:
              return self.numpy.logical_not(self._mask)
  
      @property
      def ismasked(self):
          return self.boolmask(maskedwhen=True)
  
      @property
      def isunmasked(self):
          return self.boolmask(maskedwhen=False)
  
      @property
      def content(self):
          return self._content
  
      @content.setter
      def content(self, value):
          self._content = self._util_toarray(value, self.DEFAULTTYPE)
          self._isvalid = False
  
      @property
      def maskedwhen(self):
          return self._maskedwhen
  
      @maskedwhen.setter
      def maskedwhen(self, value):
          self._maskedwhen = bool(value)
  
      def _getnbytes(self, seen):
          if id(self) in seen:
              return 0
          else:
              seen.add(id(self))
              return self._mask.nbytes + (self._content.nbytes if isinstance(self._content, self.numpy.ndarray) else self._content._getnbytes(seen))
  
      def __len__(self):
          return len(self._mask)
  
      def _gettype(self, seen):
          return awkward.type.OptionType(awkward.type._fromarray(self._content, seen))
  
      def _valid(self):
          if self.check_whole_valid:
              if not self._isvalid:
                  if len(self._mask) > len(self._content):
                      raise ValueError("mask length ({0}) must be the same as (or shorter than) the content length ({1})".format(len(self._mask), len(self._content)))
  
                  self._isvalid = True
  
      def __iter__(self, checkiter=True):
          if checkiter:
              self._checkiter()
          self._valid()
  
          mask = self._mask
          lenmask = len(mask)
          content = self._content
          maskedwhen = self._maskedwhen
          masked = self.masked
  
          i = 0
          while i < lenmask:
              if mask[i] == maskedwhen:
                  yield masked
              else:
                  yield content[i]
              i += 1
  
      def __getitem__(self, where):
          self._valid()
  
          if self._util_isstringslice(where):
              content = self._content[where]
              cls = awkward.array.objects.Methods.maybemixin(type(content), self.MaskedArray)
              out = cls.__new__(cls)
              out.__dict__.update(self.__dict__)
              out._content = content
              return out
  
          if isinstance(where, tuple) and len(where) == 0:
              return self
          if not isinstance(where, tuple):
              where = (where,)
          head, tail = where[0], where[1:]
  
          if self._util_isinteger(head):
              if self._mask[head] == self._maskedwhen:
                  if tail != ():
                      raise ValueError("masked element ({0}) is not subscriptable".format(self.masked))
                  return self.masked
              else:
                  return self._content[:len(self._mask)][(head,) + tail]
  
          else:
              mask = self._mask[head]
              if tail != () and ((self._maskedwhen and mask.any()) or (not self._maskedwhen and not mask.all())):
                  raise ValueError("masked element ({0}) is not subscriptable".format(self.masked))
              else:
                  return self.copy(mask=mask, content=self._content[:len(self._mask)][(head,) + tail])
          
      def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
          if "out" in kwargs:
              raise NotImplementedError("in-place operations not supported")
  
          if method != "__call__":
              return NotImplemented
  
          tokeep = None
          for x in inputs:
              if isinstance(x, MaskedArray):
                  x._valid()
                  if tokeep is None:
                      tokeep = x.boolmask(maskedwhen=False)
                  else:
                      tokeep = tokeep & x.boolmask(maskedwhen=False)
  
          assert tokeep is not None
  
          inputs = list(inputs)
          for i in range(len(inputs)):
              if isinstance(inputs[i], IndexedMaskedArray):
                  inputs[i] = inputs[i]._content[inputs[i]._mask[tokeep]]
              elif isinstance(inputs[i], MaskedArray):
                  inputs[i] = inputs[i]._content[tokeep]
              elif isinstance(inputs[i], (self.numpy.ndarray, awkward.array.base.AwkwardArray)):
                  inputs[i] = inputs[i][tokeep]
              else:
                  try:
                      for first in inputs[i]:
                          break
                  except TypeError:
                      pass
                  else:
                      inputs[i] = self.numpy.array(inputs[i], copy=False)[tokeep]
  
          # compute only the non-masked elements
          result = getattr(ufunc, method)(*inputs, **kwargs)
  
          # put the masked out values back
          index = self.numpy.full(len(tokeep), -1, dtype=self.INDEXTYPE)
          index[tokeep] = self.numpy.arange(self.numpy.count_nonzero(tokeep))
  
          if isinstance(result, tuple):
              return tuple(self.Methods.maybemixin(type(x), IndexedMaskedArray)(index, x, maskedwhen=-1) if isinstance(x, (self.numpy.ndarray, awkward.array.base.AwkwardBase)) else x for x in result)
          elif method == "at":
              return None
          else:
              return self.Methods.maybemixin(type(result), IndexedMaskedArray)(index, result, maskedwhen=-1)
  
      def indexed(self):
          maskindex = self.numpy.arange(len(self), dtype=self.INDEXTYPE)
          maskindex[self.boolmask(maskedwhen=True)] = -1
          return IndexedMaskedArray(maskindex, self._content, maskedwhen=-1)
  
      def _reduce(self, ufunc, identity, dtype, regularaxis):
          if self._util_hasjagged(self._content):
              return self.copy(content=self._content._reduce(ufunc, identity, dtype, regularaxis))
  
          elif isinstance(self._content, awkward.array.table.Table):
              out = awkward.array.table.Table()
              for n, x in self._content._contents.items():
                  out[n] = self.copy(content=x)
              return out._reduce(ufunc, identity, dtype, regularaxis)
  
          else:
              return ufunc.reduce(self._prepare(identity, dtype))
  
      def _prepare(self, identity, dtype):
          if isinstance(self._content, self.numpy.ndarray):
              if dtype is None and issubclass(self._content.dtype.type, (self.numpy.bool_, self.numpy.bool)):
                  dtype = self.numpy.dtype(type(identity))
              if dtype is None:
                  content = self._content
              else:
                  content = self._content.astype(dtype)
          else:
              content = self._content._prepare(identity, dtype)
  
          if content is self._content or not content.flags.owndata:
              content = content.copy()
  
          content[self.ismasked] = identity
          return content
  
  class BitMaskedArray(MaskedArray):
      """
      BitMaskedArray
      """
  
      # TODO for 1.0: need a maskshape parameter to apply length and multidimensional shape to the output
  
      def __init__(self, mask, content, maskedwhen=True, lsborder=False):
          super(BitMaskedArray, self).__init__(mask, content, maskedwhen=maskedwhen)
          self.lsborder = lsborder
  
      @classmethod
      def fromboolmask(cls, mask, content, maskedwhen=True, lsborder=False):
          return cls(cls.bool2bit(mask, lsborder=lsborder), content, maskedwhen=maskedwhen, lsborder=lsborder)
          
      def copy(self, mask=None, content=None, maskedwhen=None, lsborder=None):
          out = super(BitMaskedArray, self).copy(mask=mask, content=content, maskedwhen=maskedwhen)
          out._lsborder = self._lsborder
          if lsborder is not None:
              out._lsborder = lsborder
          return out
  
      def _mine(self, overrides):
          mine = {}
          mine["maskedwhen"] = overrides.pop("maskedwhen", self._maskedwhen)
          mine["lsborder"] = overrides.pop("lsborder", self._lsborder)
          return mine
  
      def __awkward_persist__(self, ident, fill, prefix, suffix, schemasuffix, storage, compression, **kwargs):
          self._valid()
          return {"id": ident,
                  "call": ["awkward", "BitMaskedArray"],
                  "args": [fill(self._mask, "BitMaskedArray.mask", prefix, suffix, schemasuffix, storage, compression, **kwargs),
                           fill(self._content, "BitMaskedArray.content", prefix, suffix, schemasuffix, storage, compression, **kwargs),
                           {"json": bool(self._maskedwhen)},
                           {"json": bool(self._lsborder)}]}
  
      @property
      def mask(self):
          return self._mask
  
      @mask.setter
      def mask(self, value):
          value = self._util_toarray(value, self.BITMASKTYPE, self.numpy.ndarray)
          if self.check_prop_valid:
              if len(value.shape) != 1:
                  raise ValueError("mask must have 1-dimensional shape")
          self._mask = value.view(self.BITMASKTYPE)
          self._isvalid = False
  
      def __len__(self):
          return len(self._content)
  
      @staticmethod
      def _ceildiv8(x):
          return -(-x >> 3)   # this is int(math.ceil(x / 8))
  
      @classmethod
      def bit2bool(cls, bitmask, lsborder=False):
          out = cls.numpy.unpackbits(bitmask)
          if lsborder:
              out = out.reshape(-1, 8)[:,::-1].reshape(-1)
          return out.view(cls.MASKTYPE)
          
      @classmethod
      def bool2bit(cls, boolmask, lsborder=False):
          boolmask = cls._util_toarray(boolmask, cls.MaskedArray.fget(None).MASKTYPE, cls.numpy.ndarray)
          if len(boolmask.shape) != 1:
              raise ValueError("boolmask must have 1-dimensional shape")
          if not issubclass(boolmask.dtype.type, (cls.numpy.bool_, cls.numpy.bool)):
              boolmask = (boolmask != 0)
  
          if lsborder:
              # maybe pad the length for reshape
              length = cls._ceildiv8(len(boolmask)) * 8
              if length != len(boolmask):
                  out = cls.numpy.empty(length, dtype=boolmask.dtype)
                  out[:len(boolmask)] = boolmask
              else:
                  out = boolmask
  
              # reverse the order in groups of 8
              out = out.reshape(-1, 8)[:,::-1].reshape(-1)
  
          else:
              # numpy.packbits encodes as msb (most significant bit); already in the right order
              out = boolmask
  
          return cls.numpy.packbits(out)
  
      def boolmask(self, maskedwhen=None):
          if maskedwhen is None:
              maskedwhen = self._maskedwhen
          if maskedwhen == self._maskedwhen:
              bitmask = self._mask
          else:
              bitmask = self.numpy.bitwise_not(self._mask)
          return self.bit2bool(bitmask, lsborder=self._lsborder)[:len(self._content)]
  
      @property
      def lsborder(self):
          return self._lsborder
  
      @lsborder.setter
      def lsborder(self, value):
          self._lsborder = bool(value)
  
      def _valid(self):
          if self.check_whole_valid:
              if not self._isvalid:
                  if len(self._mask) != self._ceildiv8(len(self._content)):
                      raise ValueError("mask length ({0}) must be equal to ceil(content length / 8) ({1})".format(len(self._mask), self._ceildiv8(len(self._content))))
  
                  self._isvalid = True
  
      def __iter__(self, checkiter=True):
          if checkiter:
              self._checkiter()
          self._valid()
  
          one = self.numpy.uint8(1)
          zero = self.numpy.uint8(0)
          mask = self._mask
          content = self._content
          lencontent = len(content)
          maskedwhen = self._maskedwhen
          masked = self.masked
  
          if self._lsborder:
              byte = i = 0
              bit = start = self.numpy.uint8(1)
              while i < lencontent:
                  if ((mask[byte] & bit) != 0) == self._maskedwhen:
                      yield masked
                  else:
                      yield content[i]
                  bit <<= one
                  if bit == zero:
                      bit = start
                      byte += 1
                  i += 1
  
          else:
              byte = i = 0
              bit = start = self.numpy.uint8(128)
              while i < lencontent:
                  if ((mask[byte] & bit) != 0) == self._maskedwhen:
                      yield masked
                  else:
                      yield content[i]
                  bit >>= one
                  if bit == zero:
                      bit = start
                      byte += 1
                  i += 1
  
      def _maskat(self, where):
          bytepos = self.numpy.right_shift(where, 3)    # where // 8
          bitpos  = where - 8*bytepos                           # where % 8
  
          if self._lsborder:
              bitmask = self.numpy.left_shift(1, bitpos)
          else:
              bitmask = self.numpy.right_shift(128, bitpos)
  
          if isinstance(bitmask, self.numpy.ndarray):
              bitmask = bitmask.astype(self.BITMASKTYPE)
          else:
              bitmask = self.BITMASKTYPE.type(bitmask)
  
          return bytepos, bitmask
  
      def _maskwhere(self, where):
          if self._util_isinteger(where):
              bytepos, bitmask = self._maskat(where)
              return self.numpy.bitwise_and(self._mask[bytepos], bitmask) != 0
  
          elif isinstance(where, slice):
              # assumes a small slice; for a big slice, it could be faster to unpack the whole mask
              return self._maskwhere(self.numpy.arange(*where.indices(len(self._content))))
  
          else:
              where = self.numpy.array(where, copy=False)
              if len(where.shape) == 1 and self._util_isintegertype(where.dtype.type):
                  byteposes, bitmasks = self._maskat(where)
                  self.numpy.bitwise_and(bitmasks, self._mask[byteposes], bitmasks)
                  return bitmasks.astype(self.numpy.bool_)
          
              elif len(where.shape) == 1 and issubclass(where.dtype.type, (self.numpy.bool, self.numpy.bool_)):
                  # scales with the size of the mask anyway, so go ahead and unpack the whole mask
                  unpacked = self.numpy.unpackbits(self._mask).view(self.MASKTYPE)
  
                  if self._lsborder:
                      unpacked = unpacked.reshape(-1, 8)[:,::-1].reshape(-1)[:len(where)]
                  else:
                      unpacked = unpacked[:len(where)]
  
                  return unpacked[where]
  
              else:
                  raise TypeError("cannot interpret shape {0}, dtype {1} as a fancy index or mask".format(where.shape, where.dtype))
  
      def __getitem__(self, where):
          self._valid()
  
          if self._util_isstringslice(where):
              content = self._content[where]
              cls = awkward.array.objects.Methods.maybemixin(type(content), self.BitMaskedArray)
              out = cls.__new__(cls)
              out.__dict__.update(self.__dict__)
              out._content = content
              return out
  
          if isinstance(where, tuple) and len(where) == 0:
              return self
          if not isinstance(where, tuple):
              where = (where,)
          head, tail = where[0], where[1:]
  
          if self._util_isinteger(head):
              if self._maskwhere(head) == self._maskedwhen:
                  if tail != ():
                      raise ValueError("masked element ({0}) is not subscriptable".format(self.masked))
                  return self.masked
              else:
                  return self._content[(head,) + tail]
  
          else:
              mask = self._maskwhere(head)
              if tail != () and ((self._maskedwhen and mask.any()) or (not self._maskedwhen and not mask.all())):
                  raise ValueError("masked element ({0}) is not subscriptable".format(self.masked))
              else:
                  return self.copy(mask=self.bool2bit(mask, lsborder=self._lsborder), content=self._content[(head,) + tail], lsborder=self._lsborder)
  
  class IndexedMaskedArray(MaskedArray):
      """
      IndexedMaskedArray
      """
  
      # TODO for 1.0: remove maskedwhen and instead check for any negative value (can't be allowed to inherit methods that assume self.maskedwhen!)
  
      def __init__(self, mask, content, maskedwhen=-1):
          super(IndexedMaskedArray, self).__init__(mask, content, maskedwhen=maskedwhen)
          self._isvalid = False
  
      def copy(self, mask=None, content=None, maskedwhen=None):
          out = self.__class__.__new__(self.__class__)
          out._mask = self._mask
          out._content = self._content
          out._maskedwhen = self._maskedwhen
          out._isvalid = self._isvalid
          if mask is not None:
              out._mask = mask
          if content is not None:
              out._content = content
          if maskedwhen is not None:
              out._maskedwhen = maskedwhen
          return out
  
      def __awkward_persist__(self, ident, fill, prefix, suffix, schemasuffix, storage, compression, **kwargs):
          self._valid()
          return {"id": ident,
                  "call": ["awkward", "IndexedMaskedArray"],
                  "args": [fill(self._mask, "IndexedMaskedArray.mask", prefix, suffix, schemasuffix, storage, compression, **kwargs),
                           fill(self._content, "IndexedMaskedArray.content", prefix, suffix, schemasuffix, storage, compression, **kwargs),
                           {"json": int(self._maskedwhen)}]}
  
      @property
      def mask(self):
          return self._mask
  
      @mask.setter
      def mask(self, value):
          value = self._util_toarray(value, self.INDEXTYPE, self.numpy.ndarray)
          if self.check_prop_valid:
              if not self._util_isintegertype(value.dtype.type):
                  raise TypeError("starts must have integer dtype")
              if len(value.shape) != 1:
                  raise ValueError("mask must have 1-dimensional shape")
          self._mask = value
          self._isvalid = False
  
      @property
      def content(self):
          return self._content
  
      @content.setter
      def content(self, value):
          self._content = self._util_toarray(value, self.DEFAULTTYPE)
          self._isvalid = False
  
      @property
      def maskedwhen(self):
          return self._maskedwhen
  
      @maskedwhen.setter
      def maskedwhen(self, value):
          if self.check_prop_valid:
              if not self._util_isinteger(value):
                  raise TypeError("maskedwhen must be an integer for IndexedMaskedArray")
          self._maskedwhen = value
  
      def boolmask(self, maskedwhen=True):
          if maskedwhen is None:
              raise TypeError("maskedwhen must be True or False")
          if maskedwhen:
              return self._mask == self._maskedwhen
          else:
              return self._mask != self._maskedwhen
  
      def _valid(self):
          if self.check_whole_valid:
              if not self._isvalid:
                  if len(self._mask) != 0:
                      if self._mask.max() > len(self._content):
                          raise ValueError("maximum mask-index ({0}) is beyond the length of the content ({1})".format(self._mask.max(), len(self._content)))
                      if (self._mask[self._mask != self._maskedwhen] < 0).any():
                          raise ValueError("mask-index has negative values (other than maskedwhen)")
  
                  self._isvalid = True
  
      def __iter__(self, checkiter=True):
          if checkiter:
              self._checkiter()
          self._valid()
  
          mask = self._mask
          lenmask = len(mask)
          content = self._content
          maskedwhen = self._maskedwhen
          masked = self.masked
  
          i = 0
          while i < lenmask:
              maskindex = mask[i]
              if maskindex == maskedwhen:
                  yield masked
              else:
                  yield content[maskindex]
              i += 1
  
      def __getitem__(self, where):
          self._valid()
  
          if self._util_isstringslice(where):
              content = self._content[where]
              cls = awkward.array.objects.Methods.maybemixin(type(content), self.IndexedMaskedArray)
              out = cls.__new__(cls)
              out.__dict__.update(self.__dict__)
              out._content = content
              return out
  
          if isinstance(where, tuple) and len(where) == 0:
              return self
          if not isinstance(where, tuple):
              where = (where,)
          head, tail = where[0], where[1:]
  
          if self._util_isinteger(head):
              maskindex = self._mask[head]
              if maskindex == self._maskedwhen:
                  if tail != ():
                      raise ValueError("masked element ({0}) is not subscriptable".format(self.masked))
                  return self.masked
              else:
                  return self._content[(maskindex,) + tail]
  
          else:
              maskindex = self._mask[head]
              if tail != () and (maskindex == self._maskedwhen).any():
                  raise ValueError("masked element ({0}) is not subscriptable".format(self.masked))
              else:
                  return self.copy(mask=maskindex)
  
      def indexed(self):
          return self
  
      def _prepare(self, identity, dtype):
          if isinstance(self._content, self.numpy.ndarray):
              if dtype is None and issubclass(self._content.dtype.type, (self.numpy.bool_, self.numpy.bool)):
                  dtype = self.numpy.dtype(type(identity))
              if dtype is None:
                  content = self._content
              else:
                  content = self._content.astype(dtype)
          else:
              content = self._content._prepare(identity, dtype)
  
          out = self.numpy.full(self._mask.shape + content.shape[1:], identity, dtype=content.dtype)
          out[self.isunmasked] = content[self.mask[self.mask >= 0]]
          return out
  