From 686fd831e51c175a37a47ccf51be23272e546acd Mon Sep 17 00:00:00 2001
From: Bryan Miller <millerwbryan@gmail.com>
Date: Mon, 21 Aug 2023 11:41:55 -0400
Subject: [PATCH] Fixed bugs related to atoms and observation splitting

---
 scheduler/core/components/collector/__init__.py  | 4 ++--
 scheduler/core/components/optimizer/greedymax.py | 2 +-
 scheduler/core/output/__init__.py                | 3 ++-
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/scheduler/core/components/collector/__init__.py b/scheduler/core/components/collector/__init__.py
index c152321c..5b857450 100644
--- a/scheduler/core/components/collector/__init__.py
+++ b/scheduler/core/components/collector/__init__.py
@@ -508,14 +508,14 @@ def time_accounting(self, site_plans: Plans) -> None:
                 observation = self.get_observation(v.obs_id)
                 obs_seq = observation.sequence
                 # check that Observation is Observed
-                if v.atom_end_idx == len(obs_seq)-1:
+                if v.atom_end_idx == len(obs_seq) - 1:
                     logger.warning(f'Marking observation complete: {observation.id.id}')
                     observation.status = ObservationStatus.OBSERVED
                 else:
                     observation.status = ObservationStatus.ONGOING
                       
                 # Update by atom in the sequence
-                for atom_idx in range(v.atom_start_idx, v.atom_end_idx):
+                for atom_idx in range(v.atom_start_idx, v.atom_end_idx + 1):
                     obs_seq[atom_idx].program_used = obs_seq[atom_idx].prog_time
                     obs_seq[atom_idx].partner_used = obs_seq[atom_idx].part_time
 
diff --git a/scheduler/core/components/optimizer/greedymax.py b/scheduler/core/components/optimizer/greedymax.py
index 3fc76da2..a094dfe6 100644
--- a/scheduler/core/components/optimizer/greedymax.py
+++ b/scheduler/core/components/optimizer/greedymax.py
@@ -112,7 +112,7 @@ def cumulative_seq_exec_times(sequence: Sequence) -> List[timedelta]:
         for atom in sequence:
             if not atom.observed:
                 total_exec += atom.exec_time
-                cumul_seq.append(total_exec)
+            cumul_seq.append(total_exec)
         if len(cumul_seq) == 0:
             cumul_seq.append(total_exec)
 
diff --git a/scheduler/core/output/__init__.py b/scheduler/core/output/__init__.py
index 2ae2f7dd..bf18cd7f 100644
--- a/scheduler/core/output/__init__.py
+++ b/scheduler/core/output/__init__.py
@@ -130,7 +130,8 @@ def print_plans(all_plans: List[Plans]) -> None:
         for plan in plans:
             print(f'Plan for site: {plan.site.name}')
             for visit in plan.visits:
-                print(f'\t{visit.start_time}   {visit.obs_id.id:20} {visit.score:8.2f}')
+                print(f'\t{visit.start_time}   {visit.obs_id.id:20} {visit.score:8.2f} {visit.atom_start_idx:4d} '
+                      f'{visit.atom_end_idx:4d}')
 
 
 def plans_table(all_plans: List[Plans]) -> List[Dict[Site, pd.DataFrame]]: