From 2e110779c49af693c042d80a3c7485847a874511 Mon Sep 17 00:00:00 2001
From: SGCMarkus <markusornik@gmail.com>
Date: Mon, 18 Mar 2024 13:44:48 +0100
Subject: [PATCH] astrodatadownloader: show later/better matching stage for
 fetching which files to download

---
 .../astrodatadownloader.py                    | 40 +++++++++++--------
 main/astrodatagui/ui/NewStarDialog.py         | 19 +++++----
 2 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/main/astrodatadownloader/astrodatadownloader.py b/main/astrodatadownloader/astrodatadownloader.py
index b15e357..eddde02 100644
--- a/main/astrodatadownloader/astrodatadownloader.py
+++ b/main/astrodatadownloader/astrodatadownloader.py
@@ -57,7 +57,19 @@ def parse_manifest(manifest):
 
     return results
 
-def getStarObservations(starName: str, sources: list[ObservationSource], sequences: list[list[int]] = []) -> str:
+def fitsFilenameFilterFunc(table, key_colnames):
+    if(str(table["productFilename"]).lower().endswith(".fits") or
+       str(table["productFilename"]).lower().endswith(".fit")):
+        return True
+    return False
+
+def getStarObservations(starName: str, keplerCadences, k2Cadences,
+                        sources: list[ObservationSource], sequences: list[list[int]] = []) -> str:
+    if(len(keplerCadences) != 2):
+        raise Exception("keplerCadences needs to have length 2")
+    if(len(k2Cadences) != 2):
+        raise Exception("k2Cadences need to have length 2")
+
     obs = Observations.query_object(objectname=starName, radius="0 deg")
 
     obsWantedFilter = obs["dataproduct_type"] == "timeseries"
@@ -79,21 +91,7 @@ def getStarObservations(starName: str, sources: list[ObservationSource], sequenc
 
     obsWantedFilter &= obsSourceFilter
 
-    return obs[obsWantedFilter]
-
-def fitsFilenameFilterFunc(table, key_colnames):
-    if(str(table["productFilename"]).lower().endswith(".fits") or
-       str(table["productFilename"]).lower().endswith(".fit")):
-        return True
-    return False
-
-def downloadStarProducts(obs_wanted, keplerCadences, k2Cadences):
-    if(len(keplerCadences) != 2):
-        raise Exception("keplerCadences needs to have length 2")
-    if(len(k2Cadences) != 2):
-        raise Exception("k2Cadences need to have length 2")
-
-    dataProducts = Observations.get_product_list(obs_wanted)
+    dataProducts = Observations.get_product_list(obs[obsWantedFilter])
 
     # LC:  TESS lightcurve
     # SLC: Kepler short cadence lightcurve
@@ -123,5 +121,15 @@ def downloadStarProducts(obs_wanted, keplerCadences, k2Cadences):
     productsWanted = productsWanted.group_by("productFilename")
     productsWanted = productsWanted.groups.filter(fitsFilenameFilterFunc)
 
+    newCol = []
+    for pwCol in productsWanted:
+        for owCol in obs[obsWantedFilter]:
+            if(pwCol["obsID"] == owCol["obsid"]):
+                newCol.append(owCol["sequence_number"])
+
+    productsWanted.add_column(newCol, name="sequence_number")
+    return productsWanted
+
+def downloadStarProducts(productsWanted):
     filenames = Observations.download_products(productsWanted)
     return parse_manifest(filenames)
diff --git a/main/astrodatagui/ui/NewStarDialog.py b/main/astrodatagui/ui/NewStarDialog.py
index 2ac6e9b..c180f3b 100644
--- a/main/astrodatagui/ui/NewStarDialog.py
+++ b/main/astrodatagui/ui/NewStarDialog.py
@@ -44,7 +44,12 @@ class NewStarDialog(QDialog):
                                   "Please select atleast one source before proceeding")
             return
 
-        obs = getStarObservations(star, sources)
+        keplerKadences = [self.cbKeplerShortCadence.isChecked(),
+                          self.cbKeplerLongCadence.isChecked()]
+        k2Kadences = [self.cbK2ShortCadence.isChecked(),
+                          self.cbK2LongCadence.isChecked()]
+
+        obs = getStarObservations(star, keplerKadences, k2Kadences, sources)
         if(len(obs) == 0):
             self.showErrorMessage("No observations found",
                                   "No observationnal data has been found with the current filters")
@@ -52,7 +57,8 @@ class NewStarDialog(QDialog):
 
         self.listPreview.clear()
         for o in obs:
-            self.listPreview.addItem(QListWidgetItem(f"{o['target_name']} - {o['obs_collection']} - {o['sequence_number']}"))
+            print(o)
+            self.listPreview.addItem(QListWidgetItem(f"{star} - {o['obs_collection']} - {o['sequence_number']}"))
 
         self.currentObservations = obs
         self.starIdentifier = star
@@ -62,14 +68,7 @@ class NewStarDialog(QDialog):
             if(self.leStarIdentifier.text() != self.starIdentifier):
                 self.showErrorMessage("Star Identifier", "Star identifier of fetched data does not match current name")
                 return
-
-            keplerKadences = [self.cbKeplerShortCadence.isChecked(),
-                              self.cbKeplerLongCadence.isChecked()]
-            k2Kadences = [self.cbK2ShortCadence.isChecked(),
-                              self.cbK2LongCadence.isChecked()]
-
-            self.productManifest = downloadStarProducts(self.currentObservations,
-                                                        keplerKadences, k2Kadences)
+            self.productManifest = downloadStarProducts(self.currentObservations)
             self.accept()
         else:
             self.showErrorMessage("No data fetched", "You have to fetch data first before you can proceed!")