content: """ Support for LMA data display in brawl4d.
      
      These are meant to be lightweight wrappers to coordinate data formats 
      understood by the lmatools package.
  
  """
  import numpy as np
  
  from lmatools.flashsort.autosort.LMAarrayFile import LMAdataFile
  
  from stormdrain.bounds import Bounds, BoundsFilter
  from stormdrain.data import NamedArrayDataset, indexed
  from stormdrain.pipeline import Branchpoint, coroutine, ItemModifier
  from stormdrain.support.matplotlib.artistupdaters import PanelsScatterController
  from stormdrain.support.matplotlib.poly_lasso import LassoPayloadController
  
  class LMAAnimator(object):
      
      
      def __init__(self, duration, variable='time'):
          self.tstart = time.time()
          self.duration = duration
          
      def draw_frame(self, animator, time_fraction):
          pass
          
      
      def init_draw(self, animator):
          pass
  
  
  class LMAController(object):
      """ Manages bounds object with LMA-specific criteria. Convenience functions for loading LMA data.
      """
      
      z_alt_mapping = {'z':('alt', (lambda v: (v[0]*1.0e3 - 1.0e3, v[1]*1.0e3 + 1.0e3)) ) }
      
      def __init__(self, *args, **kwargs):
          super(LMAController, self).__init__(*args, **kwargs)
          self.bounds = Bounds(chi2=(0.0, 1.0), stations=(6, 99))
          self.default_color_bounds = Bounds(parent=self.bounds, charge=(-1,1))
          self.datasets = set()
          self.flash_datasets = set()
      
      def pipeline_for_dataset(self, d, panels, 
                               names4d=('lon', 'lat', 'alt', 'time'),
                               transform_mapping=None,
                               scatter_kwargs = {}
                               ):
          """ Set 4d_names to the spatial coordinate names in d that provide 
              longitude, latitude, altitude, and time. Default of 
              lon, lat, alt, and time which are assumed to be in deg, deg, meters, seconds
          
              entries in the scatter_kwargs dictionary are passed as kwargs to the matplotlib
              scatter call.
          """
          # Set up dataset -> time-height bound filter -> brancher
          branch = Branchpoint([])
          brancher = branch.broadcast()
          
          # strictly speaking, z in the map projection and MSL alt aren't the same - z is somewhat distorted by the projection.
          # therefore, add some padding. filtered again later after projection.
          
          quality_filter = BoundsFilter(target=brancher, bounds=self.bounds).filter()
          if transform_mapping is None:
              transform_mapping = self.z_alt_mapping
          # Use 'time', which is the name in panels.bounds, and not names4d[3], which should
          # is linked to 'time' by transform_mapping if necessary
          bound_filter = BoundsFilter(target=quality_filter, bounds=panels.bounds, 
                                      restrict_to=('time'), transform_mapping=transform_mapping)
          filterer = bound_filter.filter()
          d.target = filterer
          
          # Set up brancher -> coordinate transform -> final_filter -> mutli-axis scatter updater
          scatter_ctrl = PanelsScatterController(
                              panels=panels, 
                              color_field=names4d[3], 
                              default_color_bounds=self.default_color_bounds,
                              **scatter_kwargs)
          scatter_outlet_broadcaster = scatter_ctrl.branchpoint
          scatter_updater = scatter_outlet_broadcaster.broadcast() 
          final_bound_filter = BoundsFilter(target=scatter_updater, bounds=panels.bounds)
          final_filterer = final_bound_filter.filter()
          cs_transformer = panels.cs.project_points(
                              target=final_filterer, 
                              x_coord='x', y_coord='y', z_coord='z', 
                              lat_coord=names4d[1], lon_coord=names4d[0], alt_coord=names4d[2],
                              distance_scale_factor=1.0e-3)
          branch.targets.add(cs_transformer)
          
          # return each broadcaster so that other things can tap into results of transformation of this dataset
          return branch, scatter_ctrl
      
      @coroutine
      def flash_stat_printer(self, min_points=10):
          while True:
              ev, fl = (yield)
              template = "{0} of {1} flashes have > {3} points. Their average area = {2:5.1f} km^2"
              N = len(fl)
              good = (fl['n_points'] >= min_points)
              N_good = len(fl[good])
              area = np.mean(fl['area'][good])
              print template.format(N_good, N, area, min_points)
          
      def flash_stats_for_dataset(self, d, selection_broadcaster):
          
          flash_stat_branchpoint = Branchpoint([self.flash_stat_printer()])
          flash_stat_brancher = flash_stat_branchpoint.broadcast()
          
          @coroutine
          def flash_data_for_selection(target, flash_id_key = 'flash_id'):
              """ Accepts an array of event data from the pipeline, and sends 
                  event and flash data.
              """
              while True:
                  ev = (yield) # array of event data
                  fl_dat = d.flash_data
                  
                  flash_ids = set(ev[flash_id_key])
                  flashes = np.fromiter(
                              (fl for fl in fl_dat if fl[flash_id_key] in flash_ids), 
                              dtype=fl_dat.dtype)
                  target.send((ev, flashes))
                  
          selection_broadcaster.targets.add(flash_data_for_selection(flash_stat_brancher))
          return flash_stat_branchpoint
          
          
          
          
      @indexed()
      def read_dat(self, *args, **kwargs):
          """ All args and kwargs are passed to the LMAdataFile object from lmatools"""
          lma = LMAdataFile(*args, **kwargs)
          stn = lma.stations # adds stations to lma.data as a side-effect
          d = NamedArrayDataset(lma.data)
          self.datasets.add(d)
          return d
          
      def load_dat_to_panels(self, panels, *args, **kwargs):
          """ All args and kwargs are passed to the LMAdataFile object from lmatools"""
          d = self.read_dat(*args, **kwargs)
          post_filter_brancher,  scatter_ctrl = self.pipeline_for_dataset(d, panels)
          branch_to_scatter_artists = scatter_ctrl.branchpoint
          # ask for a copy of the array from each selection operation, so that
          # it's saved and ready for any lasso operations
          
          charge_lasso = LassoChargeController(
                              target=ItemModifier(
                              target=d.update(field_names=['charge']), 
                                              item_name='charge').modify())
          branch_to_scatter_artists.targets.add(charge_lasso.cache_segment.cache_segment())
          
          return d, post_filter_brancher, scatter_ctrl, charge_lasso
          
      @indexed(index_name='hdf_row_idx')     
      def read_hdf5(self, LMAfileHDF):
          try:
              import tables
          except ImportError:
              print "couldn't import pytables"
              return None
          from hdf5_lma import HDF5Dataset
          
          # get the HDF5 table name
          LMAh5 = tables.openFile(LMAfileHDF, 'r')
          table_names = LMAh5.root.events._v_children.keys()
          table_path = '/events/' + table_names[0]
          LMAh5.close()
          d = HDF5Dataset(LMAfileHDF, table_path=table_path, mode='a')
          self.datasets.add(d)
          
          if d.flash_table is not None:
              print "found flash data"
          
          return d
          
      def load_hdf5_to_panels(self, panels, LMAfileHDF, scatter_kwargs={}):
          d = self.read_hdf5(LMAfileHDF)
          post_filter_brancher, scatter_ctrl = self.pipeline_for_dataset(d, panels, scatter_kwargs=scatter_kwargs)
          branch_to_scatter_artists = scatter_ctrl.branchpoint
          charge_lasso = LassoChargeController(
                              target=ItemModifier(
                              target=d.update(index_name='hdf_row_idx',
                                              field_names=['charge']), 
                                              item_name='charge').modify())
          branch_to_scatter_artists.targets.add(charge_lasso.cache_segment.cache_segment())            
          
          return d, post_filter_brancher, scatter_ctrl, charge_lasso
          
      def load_hdf5_flashes_to_panels(self, panels, hdf5dataset, min_points=10):
          """ Set up a flash dataset display. The sole argument is usually the HDF5 
              LMA dataset returned by a call to self.load_hdf5_to_panels """
          from hdf5_lma import HDF5FlashDataset
          if hdf5dataset.flash_table is not None:
              point_count_dtype = hdf5dataset.flash_data['n_points'].dtype
              self.bounds.n_points = (min_points, np.iinfo(point_count_dtype))
              flash_d = HDF5FlashDataset(hdf5dataset)
              transform_mapping = {}
              transform_mapping['time'] = ('start', (lambda v: (v[0], v[1])) )
              transform_mapping['lat'] = ('init_lat', (lambda v: (v[0], v[1])) )
              transform_mapping['lon'] = ('init_lon', (lambda v: (v[0], v[1])) )
              transform_mapping['z'] = ('init_alt',  (lambda v: (v[0]*1.0e3 - 1.0e3, v[1]*1.0e3 + 1.0e3)) )
              flash_post_filter_brancher, flash_scatter_ctrl = self.pipeline_for_dataset(flash_d, panels, 
                      transform_mapping=transform_mapping, 
                      names4d=('init_lon', 'init_lat', 'init_alt', 'start') )
              for art in flash_scatter_ctrl.artist_outlet_controllers:
                  # there is no time variable, but the artist updater is set to expect
                  # time. Patch that up.
                  if art.coords == ('time', 'z'):
                      art.coords = ('start', 'z')
                  # Draw flash markers in a different style
                  art.artist.set_edgecolor('k')
          self.flash_datasets.add(flash_d)
          return flash_d, flash_post_filter_brancher, flash_scatter_ctrl
  
  class LassoChargeController(LassoPayloadController):
      """ The "charge" attribute is one of {-1, 0, 1} to set 
          negative, unclassified, or positive charge, or None
          to do nothing.
      """
      charge = LassoPayloadController.Payload()        