text: # Licensed under a 3-clause BSD style license - see LICENSE.rst
  # -*- coding: utf-8 -*-
  """
  Base class for handling bit masks.
  
  *License*:
      Copyright (c) 2015, SDSS-IV/MaNGA Pipeline Group
          Licensed under BSD 3-clause license - see LICENSE.rst
  
  *Class usage examples*:
      TODO
  
  *Revision history*:
      | **01 Jun 2015**: Original implementation by K. Westfall (KBW)
      | **06 Sep 2018**: (KBW) Added to PypeIt repo, removed
          functionality of instantiating a bitmask from a file, code
          update, and slight doc changes.
  """
  import numpy
  import os
  import textwrap
  
  class BitMask:
      r"""
      Generic class to handle and manipulate bitmasks.  The input list of
      bit names (keys) must be unique, except that values of 'NULL' are
      ignored.  The index in the input keys determines the bit value;
      'NULL' keys are included in the count.  For example::
  
          >>> from mangadap.util.bitmask import BitMask
          >>> keys = [ 'key1', 'key2', 'NULL', 'NULL', 'key3' ]
          >>> bm = BitMask(keys)
          >>> bm.info()
                   Bit: key1 = 0
  
                   Bit: key2 = 1
  
                   Bit: key3 = 4
  
      .. todo::
          - Have the class keep the mask values internally instead of
            having it only operate on the mask array...
  
      Args:
          keys (:obj:`str`, :obj:`list`):
              List of keys (or single key) to use as the bit name.  Each
              key is given a bit number ranging from 0..N-1.
  
          descr (:obj:`str`, :obj:`list`, optional):
              List of descriptions (or single discription) provided by
              :func:`info` for each bit.  No descriptions by default.
  
      Raises:
          ValueError:
              Raised if more than 64 bits are provided.
          TypeError:
              Raised if the provided `keys` do not have the correct type.
  
      Attributes:
          nbits (int):
              Number of bits
          bits (dict):
              A dictionary with the bit name and value
          descr (numpy.ndarray):
              List of bit descriptions
          max_value (int):
              The maximum valid bitmask value given the number of bits.
      """
      prefix = 'BIT'
      def __init__(self, keys, descr=None):
  
          _keys = keys if hasattr(keys, '__iter__') else [keys]
          _keys = numpy.atleast_1d(_keys).ravel()
          _descr = None if descr is None else numpy.atleast_1d(descr).ravel()
  
  #        if not numpy.all([isinstance(k, str) for k in _keys]):
  #            raise TypeError('Input keys must have string type.')
          if _descr is not None:
              if not all([isinstance(d, str) for d in _descr]):
                  raise TypeError('Input descriptions must have string type.')
              if len(_descr) != len(_keys):
                  raise ValueError('Number of listed descriptions not the same as number of keys.')
  
          # Do not allow for more that 64 bits
          if len(_keys) > 64:
              raise ValueError('Can only define up to 64 bits!')
  
          # Allow for multiple NULL keys; but check the rest for
          # uniqueness
          diff = set(_keys) - set(['NULL'])
          if len(diff) != numpy.unique(_keys[[k != 'NULL' for k in _keys]]).size:
              raise ValueError('All input keys must be unique.')
  
          # Initialize the attributes
          self.nbits = len(_keys)
          self.bits = { k:i for i,k in enumerate(_keys) }
          self.max_value = (1 << self.nbits)-1
          self.descr = _descr
          
      def _prep_flags(self, flag):
          """Prep the flags for use."""
          # Flags must be a numpy array
          _flag = numpy.array(self.keys()) if flag is None else numpy.atleast_1d(flag).ravel()
          # NULL flags not allowed
          if numpy.any([f == 'NULL' for f in _flag]):
              raise ValueError('Flag name NULL is not allowed.')
          # Flags should be among the bitmask keys
          if numpy.any([f not in self.keys() for f in _flag]):
              raise ValueError('Some bit names not recognized.')
  #        # Flags should be strings
  #        if numpy.any([ not isinstance(f, str) for f in _flag ]):
  #            raise TypeError('Provided bit names must be strings!')
          return _flag
  
      @staticmethod
      def _fill_sequence(keys, vals, descr=None):
          r"""
          Fill bit sequence with NULL keys if bit values are not
          sequential.
  
          The instantiation of :class:`BitMask` does not include the
          value of the bit, it just assumes that the bits are in
          sequence such that the first key has a value of 0, and the
          last key has a value of N-1. This is a convenience function
          that fills the list of keys with 'NULL' for bit values that
          are non-sequential. This is used primarily for instantiation
          the BitMask from bits written to a file where the NULL bits
          have been skipped.
          Args:
              keys (:obj:`list`, :obj:`str`):
                  Bit names
              vals (:obj:`list`, :obj:`int`):
                  Bit values
              descr (:obj:`list`, :obj:`str`, optional):
                  The description of each bit. If None, no bit
                  descriptions are defined.
  
          Returns:
              `numpy.ndarray`_: Three 1D arrays with the filled keys,
              values, and descriptions.
  
          Raises:
              ValueError: Raised if a bit value is less than 0.
          """
          _keys = numpy.atleast_1d(keys).ravel()
          _vals = numpy.atleast_1d(vals).ravel()
          _descr = None if descr is None else numpy.atleast_1d(descr).ravel()
  
          if numpy.amin(_vals) < 0:
              raise ValueError('No bit cannot be less than 0!')
          minv = numpy.amin(_vals)
          maxv = numpy.amax(_vals)
  
          if minv != 0 or maxv != len(_vals)-1:
              diff = list(set(numpy.arange(maxv)) - set(_vals))
              _vals = numpy.append(_vals, diff)
              _keys = numpy.append(_keys, numpy.array(['NULL']*len(diff)))
              if _descr is not None:
                  _descr = numpy.append(_descr, numpy.array(['']*len(diff)))
  
          return _keys, _vals, _descr
  
      def keys(self):
          """
          Return a list of the bits; 'NULL' keywords are ignored.
          
          Returns:
              list: List of bit keywords.
          """
          return list(set(self.bits.keys())-set(['NULL']))
  
      def info(self):
          """
          Print the list of bits and, if available, their descriptions.
          """
          try:
              tr, tcols = numpy.array(os.popen('stty size', 'r').read().split()).astype(int)
              tcols -= int(tcols*0.1)
          except:
              tr = None
              tcols = None
  
          for k,v in sorted(self.bits.items(), key=lambda x:(x[1],x[0])):
              if k == 'NULL':
                  continue
              print('         Bit: {0} = {1}'.format(k,v))
              if self.descr is not None:
                  if tcols is not None:
                      print(textwrap.fill(' Description: {0}'.format(self.descr[v]), tcols))
                  else:
                      print(' Description: {0}'.format(self.descr[v]))
              print(' ')
  
      def minimum_dtype(self, asuint=False):
          """
          Return the smallest int datatype that is needed to contain all
          the bits in the mask.  Output as an unsigned int if requested.
  
          Args:
              asuint (:obj:`bool`, optional):
                  Return an unsigned integer type.  Signed types are
                  returned by default.
  
          .. warning::
              uses int16 if the number of bits is less than 8 and
              asuint=False because of issue astropy.io.fits has writing
              int8 values.
          """
          if self.nbits < 8:
              return numpy.uint8 if asuint else numpy.int16
          if self.nbits < 16:
              return numpy.uint16 if asuint else numpy.int16
          if self.nbits < 32:
              return numpy.uint32 if asuint else numpy.int32
          return numpy.uint64 if asuint else numpy.int64
  
      def flagged(self, value, flag=None):
          """
          Determine if a bit is on in the provided bitmask value.  The
          function can be used to determine if any individual bit is on or
          any one of many bits is on.
  
          Args:
              value (int, array-like):
                  Bitmask value.  It should be less than or equal to
                  :attr:`max_value`; however, that is not checked.
              flag (str, array-like, optional):
                  One or more bit names to check.  If None, then it checks
                  if *any* bit is on.
          
          Returns:
              bool: Boolean flags that the provided flags (or any flag) is
              on for the provided bitmask value.  Shape is the same as
              `value`.
  
          Raises:
              KeyError: Raised by the dict data type if the input *flag*
                  is not one of the valid :attr:`flags`.
              TypeError: Raised if the provided *flag* does not contain
                  one or more strings.
          """
          _flag = self._prep_flags(flag)
  
          out = value & (1 << self.bits[_flag[0]]) != 0
          if len(_flag) == 1:
              return out
  
          nn = len(_flag)
          for i in range(1,nn):
              out |= (value & (1 << self.bits[_flag[i]]) != 0)
          return out
  
      def flagged_bits(self, value):
          """
          Return the list of flagged bit names for a single bit value.
  
          Args:
              value (int):
                  Bitmask value.  It should be less than or equal to
                  :attr:`max_value`; however, that is not checked.
          
          Returns:
              list: List of flagged bit value keywords.
  
          Raises:
              KeyError:
                  Raised by the dict data type if the input *flag* is not
                  one of the valid :attr:`flags`.
              TypeError:
                  Raised if the provided *flag* does not contain one or
                  more strings.
          """
          if not numpy.issubdtype(type(value), numpy.integer):
              raise TypeError('Input must be a single integer.')
          if value <= 0:
              return []
          keys = numpy.array(self.keys())
          indx = numpy.array([1<<self.bits[k] & value != 0 for k in keys])
          return list(keys[indx])
  
      def toggle(self, value, flag):
          """
          Toggle a bit in the provided bitmask value.
  
          Args:
              value (int, array-like):
                  Bitmask value.  It should be less than or equal to
                  :attr:`max_value`; however, that is not checked.
              flag (str, array-like):
                  Bit name(s) to toggle.
  
          Returns:
              array-like: New bitmask value after toggling the selected
              bit.
  
          Raises:
              ValueError:
                  Raised if the provided flag is None.
          """ 
          if flag is None:
              raise ValueError('Provided bit name cannot be None.')
  
          _flag = self._prep_flags(flag)
  
          out = value ^ (1 << self.bits[_flag[0]])
          if len(_flag) == 1:
              return out
  
          nn = len(_flag)
          for i in range(1,nn):
              out ^= (1 << self.bits[_flag[i]])
          return out
  
      def turn_on(self, value, flag):
          """
          Ensure that a bit is turned on in the provided bitmask value.
  
          Args:
              value (uint or array): Bitmask value.  It should be less
                  than or equal to :attr:`max_value`; however, that is not
                  checked.
              flag (list, numpy.ndarray, or str): Bit name(s) to turn on.
          
          Returns:
              uint: New bitmask value after turning on the selected bit.
  
          Raises:
              KeyError: Raised by the dict data type if the input *flag*
                  is not one of the valid :attr:`flags`.
              Exception: Raised if the provided *flag* is not a string.
          """
          if flag is None:
              raise ValueError('Provided bit name cannot be None.')
  
          _flag = self._prep_flags(flag)
  
          out = value | (1 << self.bits[_flag[0]])
          if len(_flag) == 1:
              return out
  
          nn = len(_flag)
          for i in range(1,nn):
              out |= (1 << self.bits[_flag[i]])
          return out
  
      def turn_off(self, value, flag):
          """
          Ensure that a bit is turned off in the provided bitmask value.
  
          Args:
              value (uint or array): Bitmask value.  It should be less
                  than or equal to :attr:`max_value`; however, that is not
                  checked.
              flag (list, numpy.ndarray, or str): Bit name(s) to turn off.
          
          Returns:
              uint: New bitmask value after turning off the selected bit.
  
          Raises:
              KeyError: Raised by the dict data type if the input *flag*
                  is not one of the valid :attr:`flags`.
              Exception: Raised if the provided *flag* is not a string.
          """
          if flag is None:
              raise ValueError('Provided bit name cannot be None.')
  
          _flag = self._prep_flags(flag)
  
          out = value & ~(1 << self.bits[_flag[0]])
          if len(_flag) == 1:
              return out
  
          nn = len(_flag)
          for i in range(1,nn):
              out &= ~(1 << self.bits[_flag[i]])
          return out
  
      def consolidate(self, value, flag_set, consolidated_flag):
          """
          Consolidate a set of flags into a single flag.
          """
          indx = self.flagged(value, flag=flag_set)
          value[indx] = self.turn_on(value[indx], consolidated_flag)
          return value
  
      def unpack(self, value, flag=None):
          """
          Construct boolean arrays with the selected bits flagged.
  
          Args:
              value (numpy.ndarray):
                  The bitmask values to unpack.
              flag (:obj:`str`, :obj:`list`, optional):
                  The specific bits to unpack.  If None, all values are
                  unpacked.
          Returns:
              tuple: A tuple of boolean numpy.ndarrays flagged according
              to each bit.
          """
          _flag = self._prep_flags(flag)
          return tuple([self.flagged(value, flag=f) for f in _flag])
  
      def to_header(self, hdr, prefix=None, quiet=False):
          """
          Write the bits to a fits header.
  
          .. todo::
              - This is very similar to the function in ParSet.  Abstract
                to a general routine?
              - The comment might have a limited length and be truncated.
  
          Args:
              hdr (`astropy.io.fits.Header`):
                  Header object for the parameters. Modified in-place.
              prefix (:obj:`str`, optional):
                  Prefix to use for the header keywords, which
                  overwrites the string defined for the class. If None,
                  uses the default for the class.
              quiet (:obj:`bool`, optional):
                  Suppress print statements.
          """
          if prefix is None:
              prefix = self.prefix
          maxbit = max(list(self.bits.values()))
          ndig = int(numpy.log10(maxbit))+1 
          for key, value in sorted(self.bits.items(), key=lambda x:(x[1],x[0])):
              if key == 'NULL':
                  continue
              hdr['{0}{1}'.format(prefix, str(value).zfill(ndig))] = (key, self.descr[value])
  
      @classmethod
      def from_header(cls, hdr, prefix=None):
          """
          Instantiate the BitMask using data parsed from a fits header.
  
          .. todo::
              - This is very similar to the function in ParSet.  Abstract
                to a general routine?
              - If comments are truncated by the comment line length,
                they'll be different than a direct instantiation.
  
          Args:
              hdr (`astropy.io.fits.Header`):
                  Header object with the bits.
              prefix (:obj:`str`, optional):
                  Prefix of the relevant header keywords, which
                  overwrites the string defined for the class. If None,
                  uses the default for the class.
          """
          if prefix is None:
              prefix = cls.prefix
          # Parse the bits from the header
          keys, values, descr = cls.parse_bits_from_hdr(hdr, prefix)
          # Fill in any missing bits
          keys, values, descr = cls._fill_sequence(keys, values, descr=descr)
          # Make sure the bits are sorted
          srt = numpy.argsort(values)
          # Instantiate the BitMask
          return cls(keys[srt], descr=descr[srt])
  
      @staticmethod
      def parse_bits_from_hdr(hdr, prefix):
          """
          Parse bit names, values, and descriptions from a fits header.
  
          .. todo::
              - This is very similar to the function in ParSet.  Abstract
                to a general routine?
  
          Args:
              hdr (`astropy.io.fits.Header`):
                  Header object with the bits.
              prefix (:obj:`str`):
                  The prefix used for the header keywords.
          
          Returns:
              Three lists are returned providing the bit names, values,
              and descriptions.
          """
          keys = []
          values = []
          descr = []
          for k, v in hdr.items():
              # Check if this header keyword starts with the required
              # prefix
              if k[:len(prefix)] == prefix:
                  try:
                      # Try to convert the keyword without the prefix
                      # into an integer. Bits are 0 indexed and written
                      # to the header that way.
                      i = int(k[len(prefix):])
                  except ValueError:
                      # Assume the value is some other random keyword that
                      # starts with the prefix but isn't a parameter
                      continue
  
                  # Assume we've found a bit entry. Parse the bit name
                  # and description and add to the compiled list
                  keys += [v]
                  values += [i]
                  descr += [hdr.comments[k]]
          return keys, values, descr
  
  
  