#!/usr/bin/python

import sys
from datetime import datetime
import getopt

usage = """
This is a simple free program to process the output of the
Mighty Ohm geiger counter connected through an FTDI
usb/serial adapter.

See http://www-cs-faculty.stanford.edu/~nick/mrgeiger/

This is free software released under the Apache v2 license
Version 1.0
Nick Parlante
nick.parlante@cs.stanford.edu

Usage:

-i <input> -- *required*
 Using the very common FTDI serial chip, the input filenames look like:
  /dev/ttyUSB0 on Linux
  /dev/cu.usbserial-XXXXXXXX on the Mac
  Something like "com1" or "com2" on Windows
  
-o <output-file> -- write to named output file
-a -- auto output filename, e.g. 2012.04.01-23.14.12.csv (implies -o)
-e  -- echo the raw lines from the geiger counter
-g N -- group into samples of N seconds
        Output will be CSV format: seconds, count, cpm, uSv/hr
-t N -- timed stop after N seconds
-u X -- specify X as cpm to uSv/hr conversion factor

Use Ctrl-c to stop if not using -t N to stop after a number
of seconds.

Example:
Group data by 20 seconds, ending at 120 seconds,
write the csv data to afile.csv as well as stdout:

mrgeiger.py -i /dev/cu.usbserial-A9007CPG -g 20 -t 120 -o afile.csv

See the README.txt for more details.
"""


class MrGeiger:
  def __init__(self, argv):
    if len(argv) == 1:
      print usage
      sys.exit(1)
  
    optlist, args = getopt.getopt(argv[1:], 'i:o:g:t:u:ae')  
    opts = dict(optlist)

    # Here get string value or None
    self.in_fname = opts.get('-i')
    if not self.in_fname:
      sys.stderr.write('-i <input> required, typically:\n' +
        ' /dev/cu.usbserial-XXXXXXXX on the Mac\n' + 
        ' /dev/ttyUSB0 on Linux\n' +
        ' com1 or com2 ... on Windows\n')
      sys.exit(1)
  
    # writing to output, and what is the filename
    self.out_on = '-a' in opts
    self.out_fname = opts.get('-o')
    if self.out_fname:
      self.out_on = 1
    self.outf = None
    self.echo = '-e' in opts

    # ints: group end group_mode
    self.group = int(opts.get('-g', 0))
    self.group_mode = self.group > 0
    self.end = int(opts.get('-t', 0))
    
    # conversion of cpm to uSv/hr
    # 0.0057 is the factor for a SBM-20 tube
    self.usv_factor = float(opts.get('-u', 0.0057))

  header_csv = 'seconds, count, cpm, uSv/hr\n'

  def open_output(self):
    "Opens up self.outf as appropriate"
    if self.out_on:
      suffix = '.geiger'
      if self.group_mode:
        suffix = '.csv'
      
      if self.out_fname:
        outfilename = self.out_fname
      else:
        outfilename = datetime.now().strftime('%Y.%m.%d-%H.%M.%S') + suffix
      self.outf = open(outfilename, 'w')
  
  def print_both_data(self, seconds, count):
    "Prints one csv data line to output and stdout as appropriate."
    cpm = count * 60.0 / self.group
    s = '%d, %d, %g, %g\n' % (seconds, count, cpm, self.usv_factor * cpm)
    self.print_both(s)

  def print_both(self, s):
    "Prints to file output and stdout as appropriate."
    if self.outf:
      self.outf.write(s)
    sys.stdout.write(s)

  def readloop(self):
    "Read from device, gathering stats, writing output to file and stdout."
    count_inner = 0
    seconds_inner = 0
    count_total = 0
    seconds_total = 0

    self.open_output()  # sets self.outf for writing
    inf = open(self.in_fname, 'rU')
    
    if self.group_mode:
      self.print_both(self.header_csv)
        
    while True:
      try:
        line = inf.readline()
        if line == '': break
        if line == '\n': continue

        # like this:
        # CPS, 1, CPM, 22, uSv/hr, 0.12, SLOW
        parts = line.split(',')
        
        # This catches obvious serial garbled data, and if it's not in SLOW mode.
        # Skip the garbled line.
        # The second line read is all garbled for some reason on the mac.
        if len(parts) != 7 or not 'SLOW' in parts[6]:
          sys.stderr.write('ERROR line ignored:' + line + '\n')
          continue
        
        if self.group_mode:
          if self.echo:
            sys.stdout.write(line)
          # accumulate to the "inner" vars
          count_inner += int(parts[1])
          seconds_inner += 1
          if seconds_inner == self.group:
            count_total += count_inner
            seconds_total += seconds_inner
            self.print_both_data(seconds_total, count_inner)
            count_inner = 0
            seconds_inner = 0
        else:
          count_total += int(parts[1])
          seconds_total += 1
          self.print_both(line)

        if self.end and seconds_total >= self.end:
          break
      except KeyboardInterrupt:
        print
        break
    # print summary to stdout, but not to the output file.
    cpm = count_total * 60.0 / seconds_total
    print 'total seconds:%d count:%d cpm:%g uSv/hr:%g' % (seconds_total, count_total,
      cpm, self.usv_factor * cpm)


def main():
  geiger = MrGeiger(sys.argv)
  geiger.readloop()


if __name__ == '__main__':
  main()