def dict_from_hdf5(infile):
  ''' 
  Loads a data dictionary from an hdf5 file.
  
  Parameters
  ----------
  infile : str
    Name of hdf5 file
    
  Returns
  -------
  out : dict
    Data dictionary

  '''
  out = {}
  fin = h5py.File(infile,'r')
  for k in fin.keys():
    out[k] = fin[k][...]

  fin.close()
  check_data(out)
  return out   def hdf5_from_dict(outfile,data):
  ''' 
  Writes an hdf5 file from the data dictionary.
  
  Parameters
  ----------
  outfile : str
    Name of the output file
  
  data : dict
    Data dictionary      
  
  '''
  check_data(data)
  fout = h5py.File(outfile,'w') 
  for k in data.keys():
    fout[k] = data[k]
    
  fout.close()
  return @_memoize
def mjd(s,fmt):
  ''' 
  Converts a date string into Modified Julian Date (MJD)
  
  Parameters
  ----------
  s : string
    Date string

  fmt : string
    Format string indicating how to parse *s*
      
  Returns
  -------
  out : int
    Modified Julian Date

  '''
  d = datetime.strptime(s,fmt)
  out = (d - _REFERENCE_DATETIME).days
  return out @_memoize
def mjd_inv(m,fmt):
  ''' 
  Converts Modified Julian Date (MJD) to a date string 
  
  Parameters
  ----------
  m : int
    Modified Julian Date

  fmt : string
    format string indicating how to form *out*
      
  Returns
  -------
  out : str
    Date string
    
  '''
  m = int(m)
  d = _REFERENCE_DATETIME + timedelta(m)
  out = d.strftime(fmt)
  return out def make_basemap(lon,lat,resolution=None):
  ''' 
  Creates a transverse mercator projection which is centered about the 
  given positions. 

  Modifying this function should be sufficient to change the way that 
  all map projections are generated in PyGeoNS
  '''
  lon = np.asarray(lon)
  lat = np.asarray(lat)
  n = lon.shape[0]
  # the buffer should be the average shortest distance between 
  # points... but this is good enough
  lon_buff = max(0.1,lon.ptp()/np.sqrt(n))
  lat_buff = max(0.1,lat.ptp()/np.sqrt(n))
  llcrnrlon = min(lon) - lon_buff
  llcrnrlat = min(lat) - lat_buff
  urcrnrlon = max(lon) + lon_buff
  urcrnrlat = max(lat) + lat_buff
  lon_0 = (llcrnrlon + urcrnrlon)/2.0
  lat_0 = (llcrnrlat + urcrnrlat)/2.0
  return Basemap(projection='tmerc',
                 resolution=resolution,
                 lon_0 = lon_0,
                 lat_0 = lat_0,
                 llcrnrlon = llcrnrlon,
                 llcrnrlat = llcrnrlat,
                 urcrnrlon = urcrnrlon,
                 urcrnrlat = urcrnrlat) class InteractiveCleaner(InteractiveVectorViewer):
  ''' 
----------------- PyGeoNS Interactive Cleaner (PIC) -----------------

An interactive figure for removing outliers and time series jumps.

Controls :
  Left : Move back 1 time step (Ctrl-Left and Alt-Left move back 10
    and 100 respectively).
  
  Right : Move forward 1 time step (Ctrl-Right and Alt-Right move
    forward 10 and 100 respectively).
  
  Up : Move forward 1 station (Ctrl-Left and Alt-Left move back 10 and
    100 respectively).
            
  Down : Move back 1 station (Ctrl-Right and Alt-Right move forward 10
    and 100 respectively).
            
  R : Redraw figures.
  
  V : Toggle whether to hide the vertical deformation.
  
  H : Toggle whether to hide the station marker.
  
  D : Enables outlier removal mode while pressed.  Click and drag on
    the time series axes to remove outliers within a time interval.
  
  J : Enables jump removal mode while pressed. Jumps are estimated by
    taking a weighted mean of the data over a time interval before and 
    after the jump. Click on the time series axes to identify a jump and 
    drag the cursor over the desired time interval.

  Enter : Edit parameters through the command line.
      
Notes
-----
  Stations may also be selected by clicking on them.

  Exit PIC by closing the figures.

  Key bindings only work when the active window is one of the PIC
  figures

---------------------------------------------------------------------     
  '''
  def __init__(self,t,x,
               u=None,v=None,z=None,
               su=None,sv=None,sz=None,
               **kwargs):
    ''' 
    Parameters
    ----------
      t : (Nt,) array
        observation times
        
      x : (Nx,2) array
        observation positions
        
      u : (Nt,Nx) array
        east component

      v : (Nt,Nx) array
        north component

      z : (Nt,Nx) array
        vertical component

      su : (Nt,Nx) array, optional
        standard deviation of east component

      sv : (Nt,Nx) array, optional
        standard deviation of north component

      sz : (Nt,Nx) array, optional
        standard deviation of vertical component
        
    Note
    ----
      only one of u, v, and z need to be specified
    '''
    dataset_labels = kwargs.pop('dataset_labels',['edited data'])
    InteractiveVectorViewer.__init__(self,t,x,
                                     u=[u],v=[v],z=[z],
                                     su=[su],sv=[sv],sz=[sz],
                                     dataset_labels=dataset_labels,
                                     **kwargs)
    self._mode = None
    self._mouse_is_pressed = False
    # record of all removed jumps and outliers  
    self.log = []

  def get_data(self):
    u = self.data_sets[0][:,:,0]
    v = self.data_sets[0][:,:,1]
    z = self.data_sets[0][:,:,2]
    su = self.sigma_sets[0][:,:,0]
    sv = self.sigma_sets[0][:,:,1]
    sz = self.sigma_sets[0][:,:,2]
    return (u,v,z,su,sv,sz)
    
  def connect(self):
    self.ts_fig.canvas.mpl_connect('button_press_event',self.on_mouse_press)
    self.ts_fig.canvas.mpl_connect('motion_notify_event',self.on_mouse_move)
    self.ts_fig.canvas.mpl_connect('button_release_event',self.on_mouse_release)
    self.ts_fig.canvas.mpl_connect('key_release_event',self.on_key_release)
    self.map_fig.canvas.mpl_connect('key_release_event',self.on_key_release)
    InteractiveVectorViewer.connect(self)
        
  def remove_jump(self,jump_time,delta):
    ''' 
    estimates and removes a jump at time *jump_time*. *jump_time* is
    an integer and indicates the first day of the jump. The jump size
    is the difference between the mean values over an interval *delta*
    before and after the jump. If no data is available over these
    intervals then no changes will be made.

    '''
    tidx_right, = np.nonzero((self.t >= jump_time) & 
                             (self.t <= (jump_time+delta)))
    tidx_left,  = np.nonzero((self.t <  jump_time) & 
                             (self.t >= (jump_time-delta)))
    mean_right,_ = weighted_mean(self.data_sets[0][tidx_right,self.xidx,:],
                                 self.sigma_sets[0][tidx_right,self.xidx,:],
                                 axis=0)
    mean_left,_ = weighted_mean(self.data_sets[0][tidx_left,self.xidx,:],
                                self.sigma_sets[0][tidx_left,self.xidx,:],
                                axis=0)
    # jump for each component
    jump = mean_right - mean_left
    # only remove jumps when a jump can be calculated
    finite_idx, = np.isfinite(jump).nonzero()
    all_tidx_right, = np.nonzero(self.t >= jump_time)
    new_pos = self.data_sets[0][all_tidx_right,self.xidx,:]
    new_pos[:,finite_idx] = new_pos[:,finite_idx] - jump[finite_idx] 
    self.data_sets[0][all_tidx_right,self.xidx] = new_pos
    self.log.append(('jump',self.xidx,jump_time,delta))
    logger.info('removed jump for station %s at time %s' % 
                (self.xidx,jump_time))
      
  def remove_outliers(self,start_time,end_time):
    ''' 
    masks data data between *start_time* and *end_time*
    '''
    tidx, = np.nonzero((self.t >= start_time) & (self.t <= end_time))
    self.data_sets[0][tidx,self.xidx] = np.nan
    self.sigma_sets[0][tidx,self.xidx] = np.inf
    self.log.append(('outliers',self.xidx,start_time,end_time))
    logger.info('removed outliers for station %s from time %s to %s' % 
                (self.xidx,start_time,end_time))

  def on_mouse_press(self,event):
    # ignore if not the left button
    if event.button != 1: return
    # ignore if the event was not in an axis
    if event.inaxes is None: return
    # ignore if the event was not in the time series figure
    if not event.inaxes.figure is self.ts_fig: return
    self._mouse_is_pressed = True
    # use integer x data for consistency with the rest of PyGeoNS
    self._t1 = int(np.round(event.xdata))
    self.rects = []
    self.vlines = []
    for ax in self.ts_ax:
      ymin,ymax = ax.get_ylim()
      ax.set_ylim((ymin,ymax)) # prevent the ylims from changing after calls to draw
      r = Rectangle((self._t1,ymin),0.0,ymax-ymin,color='none',
                    alpha=0.5,edgecolor=None)
      self.rects += [r]
      self.vlines += [ax.vlines(self._t1,ymin,ymax,color='none')]
      ax.add_patch(r)
          
    self.ts_fig.canvas.draw()  
        
  def on_mouse_move(self,event):
    # do nothing is a mouse button is not being held down
    if not self._mouse_is_pressed: return
    # ignore if the event was not in an axis
    if event.inaxes is None: return
    # ignore if the event was not in the time series figure
    if not event.inaxes.figure is self.ts_fig: return
    # use integer x data for consistency with the rest of PyGeoNS
    self._t2 = int(np.round(event.xdata))
    for r,v in zip(self.rects,self.vlines):
      if self._mode == 'OUTLIER_REMOVAL':
        r.set_width(self._t2 - self._t1) 
        r.set_x(self._t1)
        r.set_color('r')
        v.set_color('k')
        
      elif self._mode == 'JUMP_REMOVAL':
        r.set_width(2*(self._t2 - self._t1))
        r.set_x(self._t1 - (self._t2 - self._t1))
        r.set_color('b')
        v.set_color('k')
       
      else:
        r.set_width(0.0)
        r.set_color('none')
        v.set_color('none')
            
    self.ts_fig.canvas.draw()  

  def on_mouse_release(self,event):
    # ignore if not the left button
    if event.button != 1: return
    # do nothing is a mouse button was not clicked in the axis
    if not self._mouse_is_pressed: return
    self._mouse_is_pressed = False
    # remove the rectangles no matter where the button was released
    for r,v in zip(self.rects,self.vlines):
      r.remove()
      v.remove()

    self.ts_fig.canvas.draw()  
    # only act on this event if the following conditions are met
    # ignore if the event was not in an axis
    if event.inaxes is None: return
    # ignore if the event was not in the time series figure
    if not event.inaxes.figure is self.ts_fig: return
    # _t2 needs to be set in the event that the click did not involve
    # a mouse move
    self._t2 = int(np.round(event.xdata))
    # act according to self._mode at the time of release
    if self._mode == 'OUTLIER_REMOVAL':
      mint = min(self._t1,self._t2)
      maxt = max(self._t1,self._t2)
      self.remove_outliers(mint,maxt) 
      # keep the time series axis limits fixed
      xlims = [i.get_xlim() for i in self.ts_ax]      
      ylims = [i.get_ylim() for i in self.ts_ax]      
      self.update()   
      for i,j in zip(self.ts_ax,xlims): i.set_xlim(j)
      for i,j in zip(self.ts_ax,ylims): i.set_ylim(j)
      self.ts_fig.canvas.draw()  

    elif self._mode == 'JUMP_REMOVAL':
      self.remove_jump(self._t1,abs(self._t2-self._t1))
      # keep the time series axis limits fixed
      xlims = [i.get_xlim() for i in self.ts_ax]      
      ylims = [i.get_ylim() for i in self.ts_ax]      
      self.update()   
      for i,j in zip(self.ts_ax,xlims): i.set_xlim(j)
      for i,j in zip(self.ts_ax,ylims): i.set_ylim(j)
      self.ts_fig.canvas.draw()  
    
    else: 
      return

  def _set_mode(self,name):
    if self._mode is None:
      self._mode = name
      logger.info('enabled %s mode' % name) 

  def _unset_mode(self,name):
    if self._mode is name:
      self._mode = None
      logger.info('disabled %s mode' % name) 
  
  def on_key_press(self,event):
    # disable c
    if event.key == 'c':
      return
              
    elif event.key == 'd':
      self._set_mode('OUTLIER_REMOVAL')
      self.on_mouse_move(event)

    elif event.key == 'j':
      self._set_mode('JUMP_REMOVAL')
      self.on_mouse_move(event)

    else:
      InteractiveVectorViewer.on_key_press(self,event)
        
  def on_key_release(self,event):
    if event.key == 'd':
      self._unset_mode('OUTLIER_REMOVAL')
      self.on_mouse_move(event)

    elif event.key == 'j':
      self._unset_mode('JUMP_REMOVAL')
      self.on_mouse_move(event) def unit_conversion(units,time='day',space='m'):
  ''' 
  Returns a factors that converts data with units of *units* (e.g.
  km/hr) to units that are in terms of *time* and *space*. 
  
  Example
  -------
  >>> unit_conversion('mm/yr',time='day',space='m')
      2.73785078713e-06

  '''
  # replace carets with two asterisks
  units = units.replace('^','**')
  # converts to m
  to_m = {'mm':1e-3,
          'cm':1e-2,
          'm':1e0,
          'km':1e3}
  # converts to seconds
  to_s = {'s':1.0,
          'min':60.0,
          'hr':60.0*60.0,
          'day':24.0*60.0*60.0,
          'mjd':24.0*60.0*60.0,
          'yr':365.25*24.0*60.0*60.0}
  # converts to user-specified space units and time units
  conv = dict([(k,v/to_m[space]) for k,v in to_m.iteritems()] +
              [(k,v/to_s[time]) for k,v in to_s.iteritems()])
  out = eval(units,conv)
  return out def _unit_string(space_exponent,time_exponent):
  ''' 
  returns a string indicating the units
  '''
  if space_exponent == 0:
    # if the space exponent is 0 then use units of microstrain
    space_str = '1e-6'
  elif space_exponent == 1:
    space_str = 'mm'
  else:
    space_str = 'mm^%s' % space_exponent

  if time_exponent == 0:
    time_str = ''
  elif time_exponent == -1:
    time_str = '/yr'
  else:
    time_str = '/yr^%s' % -time_exponent
  
  return space_str + time_str def _setup_map_ax(bm,ax):
  ''' 
  prepares the map axis for display
  '''
  # function which prints out the coordinates on the bottom left 
  # corner of the figure
  def coord_formatter(x,y):                         
    out = 'x : %g  y : %g  ' % (x,y)
    out += '(lon : %g E  lat : %g N)' % bm(x,y,inverse=True)
    return out

  ax.format_coord = coord_formatter
  bm.drawcountries(ax=ax,zorder=2)
  bm.drawstates(ax=ax,zorder=2) 
  bm.drawcoastlines(ax=ax,zorder=2)
  mer,par =  _get_meridians_and_parallels(bm,3)
  bm.drawmeridians(mer,
                   labels=[0,0,0,1],dashes=[2,2],
                   ax=ax,zorder=2,color=(0.3,0.3,0.3,1.0))
  bm.drawparallels(par,
                   labels=[1,0,0,0],dashes=[2,2],
                   ax=ax,zorder=2,color=(0.3,0.3,0.3,1.0))
  bm.drawmapboundary(ax=ax,fill_color=(0.9,0.9,0.9),zorder=2)
  bm.fillcontinents(ax=ax,color=(1.0,1.0,1.0),lake_color=(0.9,0.9,0.9),zorder=0)
  scale_lon,scale_lat = bm(*ax.transData.inverted().transform(ax.transAxes.transform([0.15,0.1])),
                           inverse=True)
  scale_size = one_sigfig((bm.urcrnrx - bm.llcrnrx)/5.0)/1000.0
  bm.drawmapscale(scale_lon,scale_lat,scale_lon,scale_lat,scale_size,
                  ax=ax,barstyle='fancy',fontsize=10,zorder=2)
  return def _setup_ts_ax(ax_lst,maxn=6):
  ''' 
  prepares the time series axes for display
  '''
  # display time in MJD and date on time series plot
  def coord_formatter(x,y):                         
    ''' 
    Takes coordinates, *x* and *y*, and returns their string 
    representation
    '''
    out = 'time : %g  ' % x
    out += '(date : %s)' % mjd_inv(x,'%Y-%m-%d')
    return out

  @FuncFormatter
  def xtick_formatter(x,p):
    ''' 
    Takes *x* and the number of ticks, *p*, and returns a string 
    representation of *x*
    '''
    try:
      out = mjd_inv(x,'%Y-%m-%d')
    except (ValueError,OverflowError):
      out = ''  
      
    return out
    
  for a in ax_lst: a.get_xaxis().set_major_formatter(xtick_formatter)
  for a in ax_lst: a.get_xaxis().set_major_locator(MaxNLocator(maxn))
  for a in ax_lst: a.format_coord = coord_formatter
  return
import numpy as np
import logging
import matplotlib.pyplot as plt
from pygeons.io.convert import dict_from_hdf5,hdf5_from_dict
from pygeons.mjd import mjd,mjd_inv
from pygeons.basemap import make_basemap
from pygeons.clean.iclean import InteractiveCleaner
from pygeons.units import unit_conversion
from pygeons.plot.plot import (_unit_string,
                               _setup_map_ax,
                               _setup_ts_ax)                               
Defines cleaning functions which are called by the PyGeoNS executable.
'''
from __future__ import division
logger = logging.getLogger(__name__)


def _remove_extension(f):
  '''remove file extension if one exists'''
  if '.' not in f:
    return f
  else:
    return '.'.join(f.split('.')[:-1])


def pygeons_clean(input_file,resolution='i',
                  input_edits_file=None,
                  break_lons=None,break_lats=None,
                  break_conn=None,no_display=False,
                  output_stem=None,**kwargs):
  ''' 
  runs the PyGeoNS Interactive Cleaner
  
  Parameters
  ----------
    data : dict
      data dictionary

    resolution : str
      basemap resolution    
    
    input_edits_file : str
      Name of the file containing edits which will automatically be 
      applied before opening up the interactive viewer.
    
    output_edits_file : str
      Name of the file where all edits will be recorded.   
      
    **kwargs : 
      gets passed to pygeons.clean.clean
         
  Returns
  -------
    out : dict
      output data dictionary 
    
  '''
  logger.info('Running pygeons clean ...')
  data = dict_from_hdf5(input_file)
  out = dict((k,np.copy(v)) for k,v in data.iteritems())

  ts_fig,ts_ax = plt.subplots(3,1,sharex=True,num='Time Series View',facecolor='white')
  _setup_ts_ax(ts_ax)
  map_fig,map_ax = plt.subplots(num='Map View',facecolor='white')
  bm = make_basemap(data['longitude'],data['latitude'],resolution=resolution)
  _setup_map_ax(bm,map_ax)
  x,y = bm(data['longitude'],data['latitude'])
  pos = np.array([x,y]).T
  t = data['time']
  dates = [mjd_inv(ti,'%Y-%m-%d') for ti in t]
  units = _unit_string(data['space_exponent'],data['time_exponent'])
  conv = 1.0/unit_conversion(units,time='day',space='m')
  u = conv*data['east']
  v = conv*data['north']
  z = conv*data['vertical']
  su = conv*data['east_std_dev']
  sv = conv*data['north_std_dev']
  sz = conv*data['vertical_std_dev']
  ic = InteractiveCleaner(
         t,pos,u=u,v=v,z=z,su=su,sv=sv,sz=sz,
         map_ax=map_ax,ts_ax=ts_ax,
         time_labels=dates,
         units=units,
         station_labels=data['id'],
         **kwargs)

  # make edits to the data set prior to displaying it
  if input_edits_file is not None:
    with open(input_edits_file,'r') as fin:
      for line in fin: 
        # ignore blank lines
        if line.isspace():
          continue
          
        type,sta,a,b = line.strip().split()
        # set the current station in *ic* to the station for this edit
        xidx, = (data['id'] == sta).nonzero()
        if len(xidx) == 0:
          # continue because the station does not exist in this 
          # dataset
          continue
          
        ic.xidx = xidx[0]
        if type == 'outliers':
          start_time = mjd(a,'%Y-%m-%d')
          stop_time = mjd(b,'%Y-%m-%d')
          ic.remove_outliers(start_time,stop_time)
        elif type == 'jump':
          jump_time = mjd(a,'%Y-%m-%d')
          delta = int(b)
          ic.remove_jump(jump_time,delta)
        else:
          raise ValueError('edit type must be either "outliers" or "jump"')

  if not no_display:
    ic.update()
    ic.connect()
    
  # set output file name
  if output_stem is None:
    output_stem = _remove_extension(input_file) + '.clean'

  output_file = output_stem + '.h5'
  output_edits_file = output_stem + '.txt'
  
  with open(output_edits_file,'w') as fout:
    for i in ic.log:
      type,xidx,a,b = i
      if type == 'outliers':
        station = data['id'][xidx]
        start_date = mjd_inv(a,'%Y-%m-%d')

        stop_date = mjd_inv(b,'%Y-%m-%d')