From 372dbd55ec22ed7da5612525a2b70a28a4b7a794 Mon Sep 17 00:00:00 2001
From: Adam Clark <adam@iris.washington.edu>
Date: Fri, 22 Sep 2017 16:06:40 -0700
Subject: [PATCH] Apply "distance from event" filter to waveform selection (ie.
 don't include event/station pairs that don't pass the filter)

---
 pyweed/pyweed_core.py       | 53 ++++++++++++++++++++++++++++++++++---
 pyweed/stations_handler.py  |  3 ++-
 pyweed/waveforms_handler.py | 23 +++++++---------
 3 files changed, 62 insertions(+), 17 deletions(-)

diff --git a/pyweed/pyweed_core.py b/pyweed/pyweed_core.py
index df099ba..624a279 100755
--- a/pyweed/pyweed_core.py
+++ b/pyweed/pyweed_core.py
@@ -18,7 +18,7 @@
 
 # Pyweed UI components
 from pyweed.preferences import Preferences, user_config_path
-from pyweed.pyweed_utils import manage_cache, iter_channels, get_sncl, get_preferred_origin, DataRequest
+from pyweed.pyweed_utils import manage_cache, iter_channels, get_sncl, get_preferred_origin, DataRequest, get_distance
 from obspy.clients.fdsn import Client
 from pyweed.event_options import EventOptions
 from pyweed.station_options import StationOptions
@@ -251,12 +251,12 @@ def iter_selected_events(self):
 
     def iter_selected_event_locations(self):
         """
-        Return an iterator of (lat, lon) pairs for each event
+        Return an iterator of (id, (lat, lon)) for each event.
         """
         for event in self.iter_selected_events():
             origin = get_preferred_origin(event)
             if origin:
-                yield (origin.latitude, origin.longitude)
+                yield (event.resource_id.id, (origin.latitude, origin.longitude),)
 
     ###############
     # Stations
@@ -310,6 +310,53 @@ def iter_selected_stations(self):
                 if sncl in self.selected_station_ids:
                     yield (network, station, channel)
 
+    ###############
+    # Waveforms
+    ###############
+
+    def iter_selected_events_stations(self):
+        """
+        Iterate through the selected event/station combinations.
+
+        The main use case this method is meant to handle is where the user
+        loaded stations based on selected events.
+
+        For example, if the user selected multiple events and searched for stations
+        within 20 degrees of any event, there may be stations that are within 20 degrees
+        of one event but farther away from others -- we want to ensure that we only include
+        the event/station combinations that are within 20 degrees of each other.
+        """
+        events = list(self.iter_selected_events())
+
+        # Look for any event-based distance filter
+        distance_range = self.station_options.get_event_distances()
+
+        if distance_range:
+            # Event locations by id
+            event_locations = dict(self.iter_selected_event_locations())
+        else:
+            event_locations = {}
+
+        # Iterate through the stations
+        for (network, station, channel) in self.iter_selected_stations():
+            for event in events:
+                if distance_range:
+                    event_location = event_locations.get(event.resource_id.id)
+                    if event_location:
+                        distance = get_distance(
+                            event_location[0], event_location[1],
+                            station.latitude, station.longitude)
+                        LOGGER.debug(
+                            "Distance from (%s, %s) to (%s, %s): %s",
+                            event_location[0], event_location[1],
+                            station.latitude, station.longitude,
+                            distance
+                        )
+                        if distance < distance_range['mindistance'] or distance > distance_range['maxdistance']:
+                            continue
+                # If we reach here, include this event/station pair
+                yield (event, network, station, channel)
+
     def close(self):
         self.manage_cache(init=False)
         self.save_preferences()
diff --git a/pyweed/stations_handler.py b/pyweed/stations_handler.py
index c46d368..2d9a1f1 100644
--- a/pyweed/stations_handler.py
+++ b/pyweed/stations_handler.py
@@ -110,7 +110,8 @@ class StationsDataRequest(DataRequest):
     def __init__(self, client, base_options, distance_range, event_locations):
         super(StationsDataRequest, self).__init__(client)
         if distance_range and event_locations:
-            self.event_locations = list(event_locations)
+            # Get a list of just the (lat, lon) for each event
+            self.event_locations = list((loc[1] for loc in event_locations))
             self.distance_range = distance_range
             try:
                 combined_locations = get_combined_locations(self.event_locations, self.distance_range['maxdistance'])
diff --git a/pyweed/waveforms_handler.py b/pyweed/waveforms_handler.py
index 72fc700..fe11578 100644
--- a/pyweed/waveforms_handler.py
+++ b/pyweed/waveforms_handler.py
@@ -324,19 +324,16 @@ def create_waveforms(self, pyweed):
         """
         Create a list of waveform entries based on the current event/station selections
         """
-        self.waveforms = []
-        self.waveforms_by_id = {}
-        # Get the events as a list
-        events = list(pyweed.iter_selected_events())
-
-        # Iterate through the stations
-        for (network, station, channel) in pyweed.iter_selected_stations():
-            for event in events:
-                waveform = WaveformEntry(
-                    event, network, station, channel
-                )
-                self.waveforms.append(waveform)
-                self.waveforms_by_id[waveform.waveform_id] = waveform
+        self.waveforms = [
+            WaveformEntry(
+                event, network, station, channel
+            )
+            for (event, network, station, channel) in pyweed.iter_selected_events_stations()
+        ]
+        self.waveforms_by_id = dict(
+            (waveform.waveform_id, waveform)
+            for waveform in self.waveforms
+        )
 
     def cancel_download(self):
         """