text: #!/usr/bin/python
  ###############################################################################
  #
  # This script prepares multiple files of coordination output from MDAnalysis
  # coordination output scripts in order to facilitate numerous 1D distributions,
  # graph building, and state splitting. This script performs a variety of
  # simple processing functions including adding columns, padding the data
  # with 0's and sorting the columns by a particular column (namely z value).
  #
  # Example: For 13/26/39/...-column data with this type (describe elsewhere):
  #          1.0 -0.13 -0.193 0.522 0.0 0.0 0.0 0.0 0.0 2.0 9.0 2.0 1748.0
  #          2.0 -0.124 -0.013 0.662 0.0 0.0 1.0 0.0 0.0 2.0 8.0 2.0 1748.0
  #
  #          The following command will load two datafiles, normalize
  #          the column length to 39 and remove the first 2000 lines:
  #
  #          python Ion_Preprocessor.py -f f1.out f2.out -m 3 -c 13 -remove 2000
  #
  #
  # This script typically is the second step in a larger analysis pipeline.
  # As far as script history, it's a combination of Row_Add_Time.py,
  # Row_Pad_Columns.py, Row_Sorted_OnZKeepAllTags.py, ion_count_splitter.py.
  #
  # By Chris Ing, 2013 for Python 2.7
  #
  ###############################################################################
  from argparse import ArgumentParser
  from itertools import product
  from re import match
  from itertools import permutations
  import gzip
  
  #a great helper function to iterate over chunks of a list
  def chunker(seq, size):
      return (seq[pos:pos + size] for pos in xrange(0, len(seq), size))
  
  # A helper function to open with gzip if the file is gzipped
  def file_opener(fname):
      if fname.endswith('.gz'):
          return gzip.open(fname)
      else:
          return open(fname)
  
  # This script takes a list of lists containing floats
  # and adds time. Some datasets do not require this step.
  def add_time_column(data_floats, num_cols=13, verbose=False):
  
      data_output = []
      for line in data_floats:
          non_time_columns = line[1:]
  
          # We chunk the non-time columns into ion groupings that
          # were detected at that timestep (-1 because there is no time col)
          temp_line = []
          for ion in chunker(non_time_columns,num_cols-1):
              temp_line.append(line[0])
              temp_line.extend(ion)
              if verbose:
                  print line[0], " ".join([str(ion_col) for ion_col in ion]),
  
          # Here we're building a series of temp lines and appending
          # one for each timestep.
          data_output.append(temp_line)
          if verbose:
              print "\n",
  
      return data_output
  
  # This will produce a label for each ion at a given timestep for each binding
  # mode. If the ion is not in a particular binding mode, it will be given
  # a mode of "0". In a way, it's like a less specific version of regex_columns.
  def bindingmode_columns(data_floats, sf_col=[5,6], num_cols=13, pad_col=4):
  
      data_output = []
  
      # Compute all the binding modes and turn them into strings
      all_binding_modes = product("01",repeat=len(sf_col))
      all_binding_strs = ["".join(mode) for mode in all_binding_modes]
  
      for line in data_floats:
          temp_line = []
          for ion in chunker(line,num_cols):
              # In the case that the ion grouping has a hypen
              # we know that's a padded column and must be excluded.
              if ion[pad_col] != "-":
                  # Extract the coordination using the columns passed
                  coords = [ion[col] for col in sf_col]
                  # The same as above but in string form with booleans
                  coords_str = "".join([str(int(coord>0)) for coord in coords])
                  temp_line.append(all_binding_strs.index(coords_str))
              else:
                  temp_line.append("-")
  
          data_output.append(temp_line)
  
      return data_output
  
  # This is a script that produces a regex classification from the input data.
  # It requires sorted or sorted/padded data_floats that is returned from
  # sort_columns and pad_columns functions.
  def regex_columns(data_floats, regex_strings, num_cols=13,
                    pad_col=4, sf_col=[5,6], sort_col=3, sort_cut=0.0,
                    max_ions=3, max_coord=9, prefix=None):
  
      data_output = []
  
      # File streams for splitting purposes if the prefix flag is set
      if prefix != None:
          count_files={}
          for regex_id in range(len(regex_strings)+1):
              count_files[regex_id] = open(prefix+"_split"+str(regex_id),"w")
  
      for line in data_floats:
          temp_label = []
          num_ions = 0
          for ion in chunker(line,num_cols):
  
              # This functionality was implemented upon discovering that
              # zero coordinated ions can exist in either the central cavity
              # or the extracellular region of a channel. For this case,
              # a sort_value cutoff must be used to determine where the ion
              # actually is! Note, look at your SF bound ion distributions
              # to see if sort_cut is set correctly. 0.0 is a safe bet...
              sort_val = ion[sort_col]
              all_zeros = all([ion[col]==0.0 for col in sf_col])
  
              # I need to code a bit of logic for pre-processed data.
              if ion[pad_col] != "-":
                  if (sort_val > sort_cut) and all_zeros:
                      temp_label.extend("+"*len(sf_col))
                  else:
                      temp_label.extend([int(ion[col]) for col in sf_col])
              else:
                  temp_label.extend("-"*len(sf_col))
              num_ions += 1
  
          #print temp_label
          # Here's a fix for when coordination integer counts are too large
          # and it ruins the fixed number of digit state label paradigm that
          # is critical for regex matching.
          for digit_index, digit in enumerate(temp_label):
              if digit != "-" and digit != "+" and digit > max_coord:
                  temp_label[digit_index] = max_coord
  
          for filler in range(max_ions-num_ions):
              temp_label.extend("-"*len(sf_col))
  
          # Convert the label list to a string of length max_ions*len(sf_cols)
          temp_string = "".join([str(coord) for coord in temp_label])
          assert len(temp_string) == max_ions*len(sf_col)
  
          temp_bool = []
          for regex_id, regex in enumerate(regex_strings):
              if match(regex, temp_string) is not None:
                  temp_bool.append(True)
              else:
                  temp_bool.append(False)
  
          # Confirm there aren't multiple regex matches
          assert sum(temp_bool) < 2, \
                   "Multiple regex matches for label " +str(temp_string)
  
          # Here's the catch-all clause for when no regex matches
          if sum(temp_bool) == 0.0:
              temp_bool.append(True)
          else:
              temp_bool.append(False)
  
          # Write to filestreams if prefix is set.
          if prefix != None:
              if sum(temp_bool) == 1:
                  count_files[temp_bool.index(True)].write(
                  " ".join([str(col) for col in line]))
                  count_files[temp_bool.index(True)].write("\n")
              else:
                  count_files[len(regex_strings)].write(
                  " ".join([str(col) for col in line]))
                  count_files[len(regex_strings)].write("\n")
  
          data_output.append((temp_string, temp_bool.index(True)))
  
      # Close filestreams.
      if prefix != None:
          for key in count_files.keys():
              count_files[key].close()
  
      return data_output
  
  # This is a script that produces a state stream from the input data
  # instead of recomputing it multiple times throughout the script for things
  # like transition counting. It requires sorted or sorted/padded data
  # that are returned from sort_columns and pad_columns functions.
  def label_columns(data_floats, num_cols=13, pad_col=4, sf_col=[5,6],
                    sort_col=3, sort_cut=0.0, max_ions=3, verbose=False):
  
      data_output = []
      for line in data_floats:
          temp_label = []
          num_ions = 0
          for ion in chunker(line,num_cols):
  
              # This functionality was implemented upon discovering that
              # zero coordinated ions can exist in either the central cavity
              # or the extracellular region of a channel. For this case,
              # a sort_value cutoff must be used to determine where the ion
              # actually is! Note, look at your SF bound ion distributions
              # to see if sort_cut is set correctly.
              sort_val = ion[sort_col]
              all_zeros = all([ion[col]==0.0 for col in sf_col])
  
              # I need to code a bit of logic for pre-processed data.
              if ion[pad_col] != "-":
                  if (sort_val > sort_cut) and all_zeros:
                      temp_label.extend("+"*len(sf_col))
                  else:
                      temp_label.extend([int(ion[col]) for col in sf_col])
              else:
                  temp_label.extend("-"*len(sf_col))
              num_ions += 1
  
          for filler in range(max_ions-num_ions):
              temp_label.extend("-"*len(sf_col))
  
          # Convert the label list to a string of length max_ions*len(sf_cols)
          temp_string = "".join([str(coord) for coord in temp_label])
          data_output.append(temp_string)
  
      return data_output
  
  # This script creates a uniform number of ions in the channel by paddding
  # incidents where few ions are detected in the data_floats array.
  def pad_columns(data_floats, num_cols=13, max_ions=3,
                  time_col=0, traj_col=11, verbose=False):
  
      data_output = []
      for line in data_floats:
  
          # This is the fake line we're going to pad with.
          # We're going to assume that line[0] contains time
          # and that we have 7 additional columns other than
          # an arbitrary number of coordination count columns.
          fake_ion = [int(float(line[time_col])),0.,0.,0.] + \
                     ["-" for x in range(num_cols-7)] + \
                     [0,int(float(line[traj_col])),0]
  
          chunked_line = chunker(line,num_cols)
          num_ions = len(line)/num_cols
  
          temp_line = []
          filler_line = []
  
          # Note: if you have a timestep where num_ions exceeds your
          # max_ions variable, you are only going to select the
          # first "max_ions" ions and you may omit some!
          # For this reason it is best to sort your ions first,
          # so the deepest ion is preserved.
          # Make sure to run ion_counter() to see how much data loss
          # you are incurring by choosing a smaller max_ions value.
          for ion in list(chunked_line)[:max_ions]:
              temp_line.extend(ion)
              if verbose:
                  print " ".join([str(ion_col) for ion_col in ion]),
  
          for filler in range(max_ions-num_ions):
              filler_line.extend(fake_ion)
              if verbose:
                  print " ".join([str(ion_col) for ion_col in fake_ion]),
  
          data_output.append(temp_line+filler_line)
          if verbose:
              print "\n",
  
      return data_output
  
  # This script sorts based on a particular column value, useful for
  # ranking ions by z-value. Toggle sorting from highest to lowest
  # with the argument plus2minus but it's pretty experimental
  # when it comes to the future scripts. Namely the use of traj_col
  # and going to a padded column instead of a
  def sort_columns(data_floats, num_cols=13, sort_col=3,
                   verbose=False, plus2minus=True):
  
      data_output = []
      for line in data_floats:
          chunked_line = chunker(line,num_cols)
  
          temp_line = []
          temp_filler = []
          # Here we sort each ion grouping by the sort_col argument
          # with attention to reverse the list depending on plus2minus.
          for ion in sorted(chunked_line, key=lambda col: col[sort_col],
                            reverse=plus2minus):
              if ion[4] == "-":
                  temp_filler.extend(ion)
              else:
                  temp_line.extend(ion)
  
          if verbose:
              if plus2minus:
                  print " ".join([str(ion_col) for ion_col in temp_line]),
                  print " ".join([str(ion_col) for ion_col in temp_filler]),
              else:
                  print " ".join([str(ion_col) for ion_col in temp_filler]),
                  print " ".join([str(ion_col) for ion_col in temp_line]),
              print "\n",
  
          if plus2minus:
              data_output.append(temp_line+temp_filler)
          else:
              data_output.append(temp_filler+temp_line)
  
      return data_output
  
  # This simply writes out the data_lines passed in with a simple
  # ASCII format. If an filename for output is not specified it
  # outputs to standard output, otherwise it outputs to the file specified.
  # The write_mode argument may be used to output a series of trajectories
  # to the same output file.
  def write_columns(data_lines, outfile=None, write_mode="w"):
      if outfile is not None:
          fout = open(outfile,write_mode)
  
      for line in data_lines:
          if outfile is not None:
              fout.write(" ".join([str(col) for col in line])+"\n")
          else:
              print " ".join([str(col) for col in line])
  
      if outfile is not None:
          fout.close()
      return True
  
  # This function processes the ion species information from datafiles written
  # using "write_merged_coordination" and writes a single long list of the
  # same length as filenames with the ion species at each timestep. The
  # difference between this function and one like label_columns is that
  # it takes raw data as input. Note that this outputs a sorted list.
  def species_columns(filenames, resid_col=12, sort_col=3,
                      all_possible_states=[0,1], max_ions=3,
                      num_cols=13, remove_frames=0, sf_col=None,
                      plus2minus=True, add_time=False):
  
      data_output = []
  
      # This is a list of permutations of all species orderings where hypen
      # is the lack of an ion. Each one is given a numerical state_id.
      # Example for 2 ion occupancy: --, 00, 01, 10, 11, 0-, 1-.
      all_species_orders = ["-"*max_ions]
      for ion_occ in range(max_ions):
          temp_orders = product(all_possible_states, repeat=ion_occ+1)
          for order in temp_orders:
              order_str = "".join([str(col) for col in order])
              for filler in range(max_ions-len(order_str)):
                  order_str += "-"
              all_species_orders.append(order_str)
  
      for filename in filenames:
          with file_opener(filename) as data_file:
              data_raw = data_file.readlines()[remove_frames:]
              data_raw_split = [line.strip().split() for line in data_raw]
  
          # Newer versions of my MDAnalysis code don't require adding a time
          # column (in past version only time[0] contained the timestep)
          if add_time:
              data_raw_split = add_time_column(data_raw_split, num_cols=num_cols)
  
          for line in data_raw_split:
              chunked_line = chunker(line,num_cols)
  
              # This is created for each line because we don't know
              # if the ion is necessarily bound to the SF.
              num_ions = 0
  
              temp_line = []
              # Sort by the float of the sort_col column number and take
              # the deepest max_ions. Usually the outer ions aren't interesting.
              for ion in sorted(chunked_line,
                                key=lambda col: float(col[sort_col]),
                                reverse=plus2minus): #[:max_ions]:
  
                  # This is the secret sauce of this method,
                  # merge_coordination files have a last decimal place that
                  # encodes an ion species this is lost when the conversion
                  # to integer takes place in the normal process_input
                  # function.
                  if sf_col != None:
                      if sum([float(ion[col]) for col in sf_col]) > 0:
                          num_ions += 1
                          temp_line.append(ion[resid_col][-1])
                  else:
                      num_ions += 1
                      temp_line.append(ion[resid_col][-1])
  
                  # A little messy, but we want to keep getting ions
                  # if they exist, if we've hit a maximum, then we're done.
                  if num_ions == max_ions:
                      break
  
              # Pad the rest with with hyphens
              for filler in range(max_ions-num_ions):
                  temp_line.append("-")
  
              temp_str = "".join(temp_line)
              data_output.append((temp_str, all_species_orders.index(temp_str)))
  
      return data_output
  
  # This function preprocesses raw input and returns both the sorted
  # ion data aswell as sorted/padded list of lists. float_cols is a list
  # of column numbers that will be converted to floats, where the rest of the
  # data will be converted to integers.
  def process_input(filenames, sort_col=3, num_cols=13,
                    remove_frames=0, max_ions=3, traj_col=11,
                    add_time=False, time_increment=1,
                    padded=False, float_cols=[1,2,3], time_col=0):
  
      # There are empty lists which will contain one or multiple files
      # worth of time series data. All columns must have numeric data
      # at this stage (later you'll pad the data and add "-" strings)
      data_floats = []
  
      # Parse the input file and split and float all the columns.
      # This is the required format for all the functions in this
      # file.
      for filename in filenames:
          with file_opener(filename) as data_file:
              data_raw = data_file.readlines()[remove_frames:]
              data_raw_split = [line.strip().split() for line in data_raw]
  
          # Newer versions of my MDAnalysis code don't require adding a time
          # column (in past version only time[0] contained the timestep)
          if add_time:
              data_raw_split = add_time_column(data_raw_split, num_cols=num_cols)
  
          # This solves the problem of the first line not having a traj_num
          # by searching the input file, though it's a problem using a merged
          # trajectory.
          prev_traj = None
          for line in data_raw_split:
              if (len(line) > traj_col) and (len(line) > time_col):
                  prev_traj = int(float(line[traj_col]))
          assert prev_traj != None, \
                   "Input file " + filename + " had no traj_id column"
  
          # Loop over the data and convert everything to integer
          # except the float_cols columns.
          for line_num, line in enumerate(data_raw_split):
              temp_line = []
              # This frame number will be used when the time column
              # is not located.
              frame_num = remove_frames+line_num+1
              for colindex, colvalue in enumerate(line):
                  if (float_cols.count(colindex) > 0 or
                      float_cols.count(colindex % num_cols) > 0):
                      temp_line.append(float(colvalue))
                  else:
                      temp_line.append(int(float(colvalue)))
              # This "pads" zero ion columns with a fake ion with
              # a timestamp.
              if len(temp_line) > 1:
                  data_floats.append(temp_line)
              else:
                  # TODO: There's a bug here when prev_time is
                  # detected above as being half-way through a
                  # file.
                  data_floats.append([frame_num] +
                                     [0.,0.,0.] +
                                     ["-" for x in range(num_cols-7)] +
                                     [0,prev_traj,0])
  
      # TODO: Write something to remove duplicate lines.
  
      # Since padding normalizes your number of row entries, it's best to
      # sort your data a priori in order to capture the inner most ions
      # preferentially.
      data_floats_sorted = sort_columns(data_floats,
                                        num_cols=num_cols,
                                        sort_col=sort_col)
      if padded:
          data_floats_padded = pad_columns(data_floats_sorted,
                                       num_cols=num_cols,
                                       max_ions=max_ions,
                                       traj_col=traj_col)
          return data_floats_padded
      else:
          return data_floats_sorted
  
      return False
  
  if __name__ == '__main__':
      parser = ArgumentParser(
      description='This script parses input columnular ASCII data\
      and makes it nice and pretty for subsequent analysis.')
  
      parser.add_argument(
      '-f', dest='filenames', type=str, nargs="+", required=True,
      help='a filename of coordination data from MDAnalysis trajectory data')
      parser.add_argument(
      '-m', dest='max_ions', type=int, required=True,
      help='the maximum number of ions in the channel to consider')
      parser.add_argument(
      '-c', dest='num_cols', type=int, default=13,
      help='the number of columns per ion in the input')
      parser.add_argument(
      '-remove', dest='remove_frames', type=int, default=0,
      help='this is a number of frames to remove from the start of the data')
      parser.add_argument(
      '-s', dest='sort_col', type=int, default=3,
      help='a zero inclusive column number to sort your row on, typically x,y,z')
      parser.add_argument(
      '-t', dest='traj_col', type=int, default=11,
      help='a zero inclusive column number that contains the run number')
      parser.add_argument(
      '-o', dest='outfile', type=str, default=None,
      help='the file to output the sorted padding output of all input files')
      parser.add_argument(
      '--addtime', dest='add_time', action="store_true", default=False,
      help='an optional argument to add time columns to each ion grouping')
      # The following arguments are used for regex state stream processing
      parser.add_argument(
      '-i', dest='regex', type=str, nargs="+",
      help='a list of regex values in quotes for state stream processing')
      parser.add_argument(
      '-sc', dest='sort_cut', type=float, default=0.0,
      help='a value on the sort_col range to classify zero coordinated data')
      parser.add_argument(
      '-sf', dest='sf_col', type=int, nargs="+", default=[5,6],
      help='the coordination integer columns that define the selectivity filter')
  
      # This is only useful for process_species
      parser.add_argument(
      '-r', dest='resid_col', type=int, default=12,
      help='a zero inclusive column number that contains the ion resid')
  
      args = parser.parse_args()
  
      data_f_padded = process_input(filenames=args.filenames,
                                            num_cols=args.num_cols,
                                            max_ions=args.max_ions,
                                            remove_frames=args.remove_frames,
                                            traj_col=args.traj_col,
                                            sort_col=args.sort_col,
                                            add_time=args.add_time,
                                            padded=True)
  
      data_f_label = label_columns(data_f_padded,
                                  num_cols=args.num_cols,
                                  sort_col=args.sort_col,
                                  sort_cut=args.sort_cut,
                                  sf_col=args.sf_col,
                                  max_ions=args.max_ions)
  
      data_f_regex = regex_columns(data_f_padded, regex_strings=args.regex,
                                  num_cols=args.num_cols,
                                  sort_col=args.sort_col,
                                  sort_cut=args.sort_cut,
                                  sf_col=args.sf_col,
                                  max_ions=args.max_ions)
  
      '''
      for x,y in zip(data_f_padded, data_f_regex):
          print x,y, len(y[0])
          if len(y[0]) > args.max_ions*len(args.sf_col):
              raise ValueError("State label is too long, possible bug")
      '''
  
      #write_columns(data_f_padded, outfile=args.outfile)
  
  