text: import logging
  import sys
  
  import numpy as np
  
  from westpa.tools import WESTTool, WESTDataReader, BinMappingComponent
  import westpa
  
  from westpa.tools.binning import write_bin_info
  
  
  log = logging.getLogger('w_bins')
  
  
  class WBinTool(WESTTool):
      prog = 'w_bins'
      description = '''\
  Display information and statistics about binning in a WEST simulation, or
  modify the binning for the current iteration of a WEST simulation.
  -------------------------------------------------------------------------------
  '''
  
      def __init__(self):
          super().__init__()
          self.subcommand = None
          self.data_reader = WESTDataReader()
          self.binning = BinMappingComponent()
          self.args = None
          self.n_iter = None
  
      # Interface for command-line tools
      def add_args(self, parser):
          self.data_reader.add_args(parser)
  
          subparsers = parser.add_subparsers(help='available commands')
  
          info_parser = subparsers.add_parser('info', help='Display information about binning.')
          info_parser.add_argument(
              '-n', '--n-iter', type=int, help='''Consider initial points of segment N_ITER (default: current iteration).'''
          )
          info_parser.add_argument(
              '--detail',
              action='store_true',
              help='''Display detailed per-bin information in addition to summary
                                   information.''',
          )
          self.binning.add_args(info_parser)
          info_parser.set_defaults(func=self.cmd_info)
  
          rebin_parser = subparsers.add_parser('rebin', help='Rebuild current iteration with new binning.')
          rebin_parser.add_argument(
              '--confirm',
              action='store_true',
              help='''Commit the revised iteration to HDF5; without this option, the effects of the
                                    new binning are only calculated and printed.''',
          )
          rebin_parser.add_argument(
              '--detail',
              action='store_true',
              help='''Display detailed per-bin information in addition to summary
                                       information.''',
          )
          rebin_parser.add_argument(
              '-n', '--n-iter', type=int, help='''Consider initial points of segment N_ITER (default: current iteration).'''
          )
          self.binning.add_args(rebin_parser, suppress=['--bins-from-file'])
          self.binning.add_target_count_args(rebin_parser)
          rebin_parser.set_defaults(func=self.cmd_rebin)
  
      def process_args(self, args):
  
          self.data_reader.process_args(args)
          self.data_reader.open(mode='r+')
          self.n_iter = getattr(args, 'n_iter', None) or self.data_reader.current_iteration
  
          # we cannot read bin information during rebins
          # interesting note: '==' is required here; 'is' fails
          if args.func == self.cmd_rebin:
              self.binning.target_counts_required = True
          else:
              self.binning.set_we_h5file_info(self.n_iter, self.data_reader)
  
          self.binning.process_args(args)
  
          self.args = args
          self.subcommand = args.func
  
      def go(self):
          self.subcommand()
  
      def cmd_info(self):
          mapper = self.binning.mapper
  
          # Get target states and their assignments
          target_states = self.data_reader.get_target_states(self.n_iter)
          n_target_states = len(target_states)
  
          iter_group = self.data_reader.get_iter_group(self.n_iter)
  
          # bin initial pcoords for iteration n_iter
          initial_pcoords = iter_group['pcoord'][:, 0, :]
          assignments = mapper.assign(initial_pcoords)
          del initial_pcoords
  
          print('Bin information for iteration {:d}'.format(self.n_iter))
  
          # Get bin counts and weights
          weights = iter_group['seg_index']['weight']
  
          write_bin_info(mapper, assignments, weights, n_target_states, detailed=self.args.detail)
  
      def cmd_rebin(self):
          mapper = self.binning.mapper
          assert mapper is not None
          if self.n_iter == 1:
              sys.stderr.write('rebin is not supported for the first iteration; reinitialize with w_init instead\n')
              sys.exit(1)
          n_target_states = len(self.data_reader.get_target_states(self.n_iter))
          we_driver = westpa.rc.get_we_driver()
          data_manager = self.data_reader.data_manager
  
          segments = data_manager.get_segments(self.n_iter, load_pcoords=True)
          last_iter_segments = data_manager.get_segments(self.n_iter - 1, load_pcoords=False)
  
          # Bin on this iteration's initial points
          # We don't have to worry about recycling because we are binning on
          # initial points rather than final points, so recycling has already
          # occurred for this iteration.
          # We do need initial states, in case we merge a newly-created walker out of existence
          # avail_initial_states = {state.state_id: state
          #                        for state in data_manager.get_unused_initial_states(n_iter = self.n_iter)}
          avail_initial_states = data_manager.get_unused_initial_states(n_iter=self.n_iter)
          used_initial_states = data_manager.get_segment_initial_states(segments)
          we_driver.new_iteration(
              initial_states=avail_initial_states, bin_mapper=mapper, bin_target_counts=self.binning.bin_target_counts
          )
          we_driver.used_initial_states = {state.state_id: state for state in used_initial_states}
          we_driver.assign(segments, initializing=True)
          we_driver.rebin_current(parent_segments=last_iter_segments)
  
          weights = np.array([segment.weight for segment in we_driver.next_iter_segments])
          assignments = np.fromiter(we_driver.next_iter_assignments, dtype=int, count=len(weights))
          write_bin_info(mapper, assignments, weights, n_target_states, detailed=self.args.detail)
  
          if self.args.confirm:
              data_manager.prepare_iteration(self.n_iter, list(we_driver.next_iter_segments))
  
              # manually update endpoint statuses only
              endpoint_types = sorted([(segment.seg_id, segment.endpoint_type) for segment in last_iter_segments])
              last_iter_group = data_manager.get_iter_group(self.n_iter - 1)
              last_iter_index = last_iter_group['seg_index'][...]
              last_iter_index['endpoint_type'] = [pair[1] for pair in endpoint_types]
              last_iter_group['seg_index'][...] = last_iter_index
  
              data_manager.save_iter_binning(
                  self.n_iter, self.binning.mapper_hash, self.binning.mapper_pickle, we_driver.bin_target_counts
              )
              data_manager.update_initial_states(we_driver.all_initial_states)
              data_manager.flush_backing()
  
  
  def entry_point():
      WBinTool().main()
  
  
  if __name__ == '__main__':
      entry_point()
  