A combined logging and debugging function.

  Arguments:
    x: The thing to log. x is also the return value. Defaults to '', although
       it's somewhat odd to call this function without passing x.
    s: Optional string to prefix the log message with. Defaults to '', which
       results in the function prefixing the log with the source code at the
       call site, or the type of x if it can't extract the source code.
    l: Optional lambda to be run after logging x. Allows useful things like
       inspecting other variables in the current context when x gets logged. Its
       return value is logged immediately after x is logged. Defaults to None.
    d: Optional bool to drop into the debugger if x is logged. Defaults to
       False.
    p: Optional bool to log the public properties of x, including basic call
       signatures of functions, if x is logged.
    t: Optional bool to wrap x in a tensorflow.Print call. If x is a Tensor,
       its value will be printed at graph execution time. The tensorflow.Print
       calls will be prefixed with the value in s, or qj's default logging
       prefix, just like any other qj log message. It also sets the
       tensorflow.Print call to only print qj.MAX_FRAME_LOGS times, so that
       it isn't too spammy when running the graph. The call attempts to validate
       that tensorflow is available and that x can be passed to tensorflow.Print
       before calling tensorflow.Print.
    n: Optional bool to log the shape, min, mean, and max values of x if numpy
       is available.
    r: Optional alternate return value to use instead of x if x is logged. Any
       value passed to r will be returned (even None). Only the private value
       _QJ_R_MAGIC is ignored.
    z: Optional bool to zero out the frame count in the current frame, so that
       logging continues. Most useful when using qj in colabs at the top
       (module) level, since that stack frame effectively lives forever.
       Defaults to False.
    b: Optional bool to enable or disable the logging of x. Defaults to True (so
       that x is logged, assuming other conditions don't prevent logging).
    pad: Optional bool to add padding blank lines before and after the logs.
       Useful for visually extracting particular logs.
    tfc: Optional bool to wrap x in a tensorflow.check_numerics call if x is a
       Tensor.
    tic: Optional bool to begin recording a duration.
    toc: Optional bool to end recording a duration started with a previous
       `tic`. Logs the corresponding duration if there was a previous `tic`.
       `tic` and `toc` can be set in the same call -- `toc` is handled first,
       which allows you to measure the body of a loop or comprehension with a
       single call to `qj(tic=1, toc=1)`.
    time: Optional bool to turn on timing of a function call. Can be used as a
       decorator.  E.g., `@qj(time=100) def foo()...` will print timing stats
       every 100 calls to foo.
    catch: Optional bool to decorate a function with exception catching that
       drops into the debugger.
    log_all_calls: Optional bool to wrap x in a new object that logs every call
       to x.  Experimental.
    _depth: Private parameter used to specify which stack frame should be used
            for both logging and debugging operations. If you're not wrapping
            qj or adding features to qj, you should leave this at it's default.

  Returns:
    x, which allows you to insert a call to qj just about anywhere.

def qj(x='',  # pylint: disable=invalid-name
       s='',
       l=None,
       d=False,
       p=False,
       t=False,
       n=False,
       r=_QJ_R_MAGIC,
       z=False,
       b=True,
       pad=False,
       tfc=False,
       tic=False,
       toc=False,
       time=False,
       catch=False,
       log_all_calls=False,
       _depth=1):
  """A combined logging and debugging function.

  Arguments:
    x: The thing to log. x is also the return value. Defaults to '', although
       it's somewhat odd to call this function without passing x.
    s: Optional string to prefix the log message with. Defaults to '', which
       results in the function prefixing the log with the source code at the
       call site, or the type of x if it can't extract the source code.
    l: Optional lambda to be run after logging x. Allows useful things like
       inspecting other variables in the current context when x gets logged. Its
       return value is logged immediately after x is logged. Defaults to None.
    d: Optional bool to drop into the debugger if x is logged. Defaults to
       False.
    p: Optional bool to log the public properties of x, including basic call
       signatures of functions, if x is logged.
    t: Optional bool to wrap x in a tensorflow.Print call. If x is a Tensor,
       its value will be printed at graph execution time. The tensorflow.Print
       calls will be prefixed with the value in s, or qj's default logging
       prefix, just like any other qj log message. It also sets the
       tensorflow.Print call to only print qj.MAX_FRAME_LOGS times, so that
       it isn't too spammy when running the graph. The call attempts to validate
       that tensorflow is available and that x can be passed to tensorflow.Print
       before calling tensorflow.Print.
    n: Optional bool to log the shape, min, mean, and max values of x if numpy
       is available.
    r: Optional alternate return value to use instead of x if x is logged. Any
       value passed to r will be returned (even None). Only the private value
       _QJ_R_MAGIC is ignored.
    z: Optional bool to zero out the frame count in the current frame, so that
       logging continues. Most useful when using qj in colabs at the top
       (module) level, since that stack frame effectively lives forever.
       Defaults to False.
    b: Optional bool to enable or disable the logging of x. Defaults to True (so
       that x is logged, assuming other conditions don't prevent logging).
    pad: Optional bool to add padding blank lines before and after the logs.
       Useful for visually extracting particular logs.
    tfc: Optional bool to wrap x in a tensorflow.check_numerics call if x is a
       Tensor.
    tic: Optional bool to begin recording a duration.
    toc: Optional bool to end recording a duration started with a previous
       `tic`. Logs the corresponding duration if there was a previous `tic`.
       `tic` and `toc` can be set in the same call -- `toc` is handled first,
       which allows you to measure the body of a loop or comprehension with a
       single call to `qj(tic=1, toc=1)`.
    time: Optional bool to turn on timing of a function call. Can be used as a
       decorator.  E.g., `@qj(time=100) def foo()...` will print timing stats
       every 100 calls to foo.
    catch: Optional bool to decorate a function with exception catching that
       drops into the debugger.
    log_all_calls: Optional bool to wrap x in a new object that logs every call
       to x.  Experimental.
    _depth: Private parameter used to specify which stack frame should be used
            for both logging and debugging operations. If you're not wrapping
            qj or adding features to qj, you should leave this at it's default.

  Returns:
    x, which allows you to insert a call to qj just about anywhere.
  """
  if qj.LOG and b:
    try:
      # Compute and collect values needed for logging.
      # We need the caller's stack frame both for logging the function name and
      # line number qj was called from, and to store some state that makes the
      # more magical features work.
      f = inspect.currentframe()
      for _ in range(_depth):
        f = f.f_back

      # This is the magic dictionary where we write state that gives log output
      # that can represent the underlying function's code structure, as well as
      # tracking how many times we logged from the stack frame, which allows us
      # to minimize log spam from logs in loops and comprehensions.
      qj_dict = f.f_locals.get('__qj_magic_wocha_doin__', {})
      qj_dict = {} if z else qj_dict
      log_count_key = 'frame_log_count_%d' % f.f_lasti
      qj_dict[log_count_key] = qj_dict.get(log_count_key, 0) + 1

      if qj_dict[log_count_key] > qj.MAX_FRAME_LOGS:
        return x

      # We're going to log things, so go ahead and collect information about the
      # caller's stack frame.
      func_name = qj_dict.get('func_name')
      if func_name is None:
        func_name = inspect.getframeinfo(f).function
        if func_name == '<dictcomp>':
          func_name = inspect.getframeinfo(f.f_back).function
        if func_name == '<genexpr>':
          func_name = inspect.getframeinfo(f.f_back).function
        if func_name == '<listcomp>':
          func_name = inspect.getframeinfo(f.f_back).function
        elif func_name == '<setcomp>':
          func_name = inspect.getframeinfo(f.f_back).function
        elif func_name == '<lambda>':
          func_name = inspect.getframeinfo(f.f_back).function + '.lambda'
        if func_name.startswith('<module>'):
          func_name = func_name.replace('<module>', 'module_level_code')

        filename = os.path.basename(f.f_code.co_filename)
        # Don't include the filename when logging in ipython contexts.
        if filename[0] != '<':
          filename = filename.replace('.py', '')
          func_name = '<{}> {}'.format(filename, func_name)
        qj_dict['func_name'] = func_name

      # If we are dealing with module-level code, don't limit logging, since
      # large amounts of module-level logs generally means we're running in a
      # colab, and it's annoying to have your logs suddenly stop after k runs.
      if 'module_level_code' in func_name:
        qj_dict[log_count_key] = 1

      # This is the magic that allows us to indent the logs in a sensible
      # manner. f_lasti is the last instruction index executed in the frame
      # (i.e., the instruction that executed the call to qj). We add each
      # instruction index into the dictionary, setting the value to the length
      # of the dictionary after that addition, so the first instruction we see
      # gets a value of 1, the second a value of 2, etc.
      qj_instructions_dict = qj_dict.get('instructions', {})
      qj_dict['instructions'] = qj_instructions_dict
      qj_instructions_dict[f.f_lasti] = qj_instructions_dict.get(
          f.f_lasti, len(qj_instructions_dict) + 1)
      # Here, we use that value to determine how many spaces we need after the
      # log prefix.
      spaces = ' ' * qj_instructions_dict[f.f_lasti]
      # And we store the dictionary back in the caller's frame.
      f.f_locals['__qj_magic_wocha_doin__'] = qj_dict

      # Try to extract the source code of this call if a string wasn't specified.
      if not s:
        try:
          code_key = '%s:%r:%s' % (f.f_code.co_filename, f.f_code.co_firstlineno, f.f_code.co_code)
          fn_calls = qj._FN_MAPS.get(code_key, {})
          if f.f_lasti not in fn_calls:
            qj._DEBUG_QJ and qj._DISASSEMBLE_FN(f.f_code, f.f_lasti)
            fn_calls[f.f_lasti] = _find_current_fn_call(f.f_code, f.f_lasti)
            qj._FN_MAPS[f.f_code.co_code] = fn_calls
          s = fn_calls.setdefault(f.f_lasti, '').strip()
        except IOError:
          # Couldn't get the source code, fall back to showing the type.
          s = ''

      # Now that we've computed the call count and the indentation, we can log.
      prefix = '%s:%s%s <%d>:' % (func_name, spaces, s or type(x), f.f_lineno)
      log = ''

      # First handle parameters that might change how x is logged.
      if n and 'numpy' in sys.modules:
        try:
          np = sys.modules['numpy']
          np_x = np.array(x)
          log = str((np_x.shape, (np.min(np_x),
                                  (np.mean(np_x), np.std(np_x)),
                                  np.max(np_x)),
                     np.histogram(np_x,
                                  bins=max(int(n),
                                           min(np.prod(np_x.shape), 5)))[0]
                    ))
          s = s or str(type(x))
          s += ' (shape (min (mean std) max) hist)'
          prefix = '%s:%s%s <%d>:' % (func_name, spaces, s, f.f_lineno)
        except:  # pylint: disable=bare-except
          pass

      if tic and x is '':
        log = 'Adding tic.'

      # toc needs to be processed after tic here so that the log messages make sense
      # when using tic/toc in a single call in a loop.
      if toc and x is '':
        if len(qj._tics):  # pylint: disable=g-explicit-length-test
          log = 'Computing toc.'
        else:
          log = 'Unable to compute toc -- no unmatched tic.'
          toc = False

      if time and x is '':
        log = 'Preparing decorator to measure timing...' + ('\n%s' % log if log else '')

      if catch and x is '':
        log = 'Preparing decorator to catch exceptions...' + ('\n%s' % log if log else '')

      # Now, either we have set the log message, or we are ready to build it directly from x.
      if not log:
        log = qj.STR_FN(x)
      log = '(multiline log follows)\n%s' % log if '\n' in log else log

      padding_string = ''
      if pad:
        if isinstance(pad, str):
          # Turn pad into a character string with no newlines as long as the
          # log string.
          log_len = (len(qj.PREFIX) + len(prefix.split('\n')[-1]) +
                     len(log.split('\n')[0]) + 1)
          padding_string = (pad.replace('\n', ' ') * log_len)[:log_len]
        else:
          try:
            padding_string = '\n' * (int(pad) - 1) + ' '
          except ValueError:
            padding_string = '\n'

      if padding_string:
        qj.LOG_FN(padding_string)

      # Log the primary log message.
      qj.LOG_FN('%s%s %s%s' % (qj.PREFIX, prefix, qj._COLOR_LOG(), log))

      # If there's a lambda, run it and log it.
      if l:
        log = qj.STR_FN(l(x))
        log = '(multiline log follows)\n%s' % log if '\n' in log else log
        qj.LOG_FN('%s%s %s%s' % (qj.PREFIX, ' ' * len(prefix), qj._COLOR_LOG(),
                                 log))

      # If we requested x's properties, compute them and log them.
      if p:
        try:
          if hasattr(inspect, 'signature'):
            argspec_func = lambda f: str(inspect.signature(f))
          else:
            argspec_func = lambda f: inspect.formatargspec(*inspect.getargspec(f))
          docs = [
              '%s%s' % (n,
                        argspec_func(v)
                        if inspect.isroutine(v) and not inspect.isbuiltin(v)
                        else '')
              for n, v in inspect.getmembers(x)
              if n == '__init__' or not n.startswith('_')
          ]
        except:  # pylint: disable=bare-except
          docs = [n for n in dir(x) if not n.startswith('_')]
        prefix_spaces = ' ' * len(prefix)
        qj.LOG_FN('%s%s %sPublic properties:\n    %s' %
                  (qj.PREFIX, prefix_spaces, qj._COLOR_LOG(), '\n    '.join(docs)))

      # If we requested tensorflow printing, wrap x in a tf.Print.
      if t:
        if (hasattr(x, '__module__') and
            'tensorflow' in x.__module__ and
            'tensorflow' in sys.modules):
          tf = sys.modules['tensorflow']
          prefix_spaces = ' ' * len(prefix)
          if 'session' in x.__module__:
            try:
              # pylint: disable=g-import-not-at-top
              try:
                from tensorflow.python import debug as tf_debug
              except ImportError:
                from google3.third_party.tensorflow.python import debug as tf_debug
              # pylint: enable=g-import-not-at-top
              x = tf_debug.LocalCLIDebugWrapperSession(x)
              x.add_tensor_filter('has_inf_or_nan', tf_debug.has_inf_or_nan)

              qj.LOG_FN('%s%s %sWrapping tf session in tfdbg session.' %
                        (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces))
            except:  # pylint: disable=bare-except
              qj.LOG_FN('%s%s %sUnable to wrap tf session in tfdbg session. '
                        'Make sure your BUILD rule depends on '
                        '//tensorflow/python/debug:debug_py.' %
                        (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces))
          else:
            qj.LOG_FN('%s%s %sWrapping return value in tf.Print operation.' %
                      (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces))
            x = tf.Print(
                x, [tf.shape(x), x],
                message='%s%s%s%s' % (qj._COLOR_PREFIX(), qj.PREFIX, prefix[:-1],
                                      qj._COLOR_END()),
                first_n=qj.MAX_FRAME_LOGS if t is True else int(t),
                summarize=qj.MAX_FRAME_LOGS,
                name='qj_print_%s_%d_%s' % (
                    func_name.split(' ')[-1], f.f_lineno,
                    x.name.replace('/', '_').replace(':', '_') if hasattr(x, 'name') else ''))

      if tfc:
        if (hasattr(x, '__module__') and
            'tensorflow' in x.__module__ and
            'tensorflow' in sys.modules):
          tf = sys.modules['tensorflow']
          prefix_spaces = ' ' * len(prefix)
          qj.LOG_FN('%s%s %sWrapping return value in tf.check_numerics.' %
                    (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces))
          x = tf.check_numerics(
              x,
              message='%s%s%s%s' % (
                  qj._COLOR_PREFIX(), qj.PREFIX, prefix[:-1], qj._COLOR_END()),
              name='qj_check_numerics_%s_%d_%s' % (
                  func_name.split(' ')[-1], f.f_lineno,
                  x.name.replace('/', '_').replace(':', '_') if hasattr(x, 'name') else ''))

      # toc needs to be processed before tic, so that single call tic/toc works in loops.
      if toc:
        if len(qj._tics):  # pylint: disable=g-explicit-length-test
          prefix_spaces = ' ' * len(prefix)
          toc = int(toc)
          if toc < 0:
            toc = len(qj._tics)
          toc = min(toc, len(qj._tics))
          toc_time = _time.time()
          for _ in range(toc):
            tic_ = qj._tics.pop()
            qj.LOG_FN('%s%s %s%.4f seconds since %s.' %
                      (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces, toc_time - tic_[1], tic_[0]))

      if tic:
        tic_ = (s, _time.time())
        qj._tics.append(tic_)
        if x is not '':
          prefix_spaces = ' ' * len(prefix)
          qj.LOG_FN('%s%s %sAdded tic.' %
                    (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces))

      if time:
        prefix_spaces = ' ' * len(prefix)
        if isinstance(x, types.FunctionType):
          qj.LOG_FN('%s%s %sWrapping return value in timing function.' %
                    (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces))
          # pylint: disable=no-value-for-parameter
          x = _timing(logs_every=int(time))(x)
          # pylint: enable=no-value-for-parameter
        elif x is '':
          # x is '', so we'll assume it's the default value and we're decorating
          # a function
          x = lambda f: (
              (qj.LOG_FN('%s%s %sDecorating %s with timing function.' %
                         (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces, str(f)))
               and False)
              # pylint: disable=no-value-for-parameter
              or _timing(logs_every=int(time))(f))
          # pylint: enable=no-value-for-parameter

      if catch:
        prefix_spaces = ' ' * len(prefix)
        if isinstance(x, types.FunctionType):
          qj.LOG_FN('%s%s %sWrapping return value in exception function.' %
                    (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces))
          # pylint: disable=no-value-for-parameter
          x = _catch(exception_type=catch)(x)
          # pylint: enable=no-value-for-parameter
        elif x is '':
          # x is '', so we'll assume it's the default value and we're decorating
          # a function
          x = lambda f: (
              (qj.LOG_FN('%s%s %sDecorating %s with exception function.' %
                         (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces, str(f)))
               and False)
              # pylint: disable=no-value-for-parameter
              or _catch(exception_type=catch)(f))
          # pylint: enable=no-value-for-parameter

      if log_all_calls:
        prefix_spaces = ' ' * len(prefix)
        qj.LOG_FN('%s%s %sWrapping all public method calls for object.' %
                  (qj.PREFIX, qj._COLOR_LOG(), prefix_spaces))

        def wrap(member_name, member_fn):
          """Wrap member_fn in a lambda that logs."""
          wrapped = (lambda *a, **kw:
                     qj('%s(%s)' % (member_name,
                                    ', '.join(['%r' % a_ for a_ in a]
                                              + ['%s=%r' % (k, v) for k, v in kw.items()])),
                        'calling %s' % member_name, _depth=2)
                     and qj(member_fn(*a, **kw),
                            'returning from %s' % member_name, _depth=2))
          if hasattr(member_fn, '__doc__'):
            wrapped.__doc__ = member_fn.__doc__
          return wrapped

        class Wrapper(type(x)):

          def __init__(self, x):
            method_types = (
                types.BuiltinFunctionType, types.BuiltinMethodType,
                types.FunctionType, types.LambdaType, types.MethodType
            )
            for m in inspect.getmembers(x):
              name = m[0]
              if not name.startswith('_'):
                member = m[1]
                if isinstance(member, method_types):
                  wrapped_fn = wrap(name, member)
                  setattr(self, name, wrapped_fn)
                else:
                  # pylint: disable=line-too-long
                  # TODO(iansf): This may be wrong. See
                  #              https://stackoverflow.com/questions/1325673/how-to-add-property-to-a-class-dynamically
                  # pylint: enable=line-too-long
                  setattr(self.__class__, name, member)

        x = Wrapper(x)

      # If we requested an alternative return value, log it.
      if r != _QJ_R_MAGIC:
        prefix = '%s:%s%s <%d>:' % (func_name, spaces, s or type(r), f.f_lineno)
        prefix_spaces = ' ' * len(prefix)
        log = qj.STR_FN(r)
        log = '(multiline log follows)\n%s' % log if '\n' in log else log
        qj.LOG_FN('%s%s %sOverridden return value: %s' % (qj.PREFIX, prefix_spaces,
                                                          qj._COLOR_LOG(), log))

      if padding_string:
        qj.LOG_FN(padding_string)

      # vvvvvvvv NO LOGS PERMITTED AFTER THIS BLOCK vvvvvvvv
      if qj_dict[log_count_key] == qj.MAX_FRAME_LOGS:
        qj.LOG_FN('%s%s:%s%sMaximum per-frame logging hit (%d). '
                  'No more logs will print at this call within this stack frame. '
                  'Set qj.MAX_FRAME_LOGS to change the limit or pass z=1 to this qj call '
                  'to zero out the frame log count.' %
                  (qj.PREFIX, func_name, spaces, qj._COLOR_LOG(),
                   qj.MAX_FRAME_LOGS))
      # ^^^^^^^^ NO LOGS PERMITTED AFTER THIS BLOCK ^^^^^^^^

      # If we requested debugging, drop into the debugger.
      if d:
        if not qj.DEBUG_FN:
          try:
            from colabtools import _debugger  # pylint: disable=g-import-not-at-top
            qj.DEBUG_FN = lambda frame: _debugger.ColabPdb().set_trace(frame=frame)
          except ImportError:
            try:
              # Import ipdb here because importing it at the top slows down execution.
              import ipdb  # pylint: disable=g-import-not-at-top
              qj.DEBUG_FN = ipdb.set_trace
            except ImportError:
              import pdb  # pylint: disable=g-import-not-at-top
              qj.DEBUG_FN = lambda frame: pdb.Pdb().set_trace(frame=frame)
        qj.DEBUG_FN(frame=f)

      # If we requested an alternative return value, return it now that
      # everything else is done.
      if r != _QJ_R_MAGIC:
        return r

    finally:
      # Delete the stack frame to ensure there are no memory leaks, as suggested
      # by https://docs.python.org/2/library/inspect.html#the-interpreter-stack
      del f

  # After everything else is done, return x.
  return x