melvin_ob/mode_control/mode/
zo_retrieval_mode.rs

1use super::{global_mode::GlobalMode, orbit_return_mode::OrbitReturnMode};
2use crate::flight_control::{FlightComputer, FlightState};
3use crate::imaging::CameraController;
4use crate::mode_control::{
5    mode_context::ModeContext,
6    signal::{ExecExitSignal, OpExitSignal, OptOpExitSignal, WaitExitSignal},
7};
8use crate::objective::KnownImgObjective;
9use crate::scheduling::task::{BaseTask, Task};
10use crate::util::Vec2D;
11use crate::{DT_0_STD, error, fatal, log, warn};
12use async_trait::async_trait;
13use chrono::{DateTime, TimeDelta, Utc};
14use fixed::types::I32F32;
15use std::{pin::Pin, sync::Arc};
16use tokio::sync::Mutex;
17use tokio_util::sync::CancellationToken;
18
19/// [`ZORetrievalMode`] is responsible for completing the final phase of a zoned objective
20/// after the spacecraft has exited its original orbit. In this mode,
21/// the spacecraft aligns, captures imagery, and uploads results.
22///
23/// The mode is considered time-sensitive, interruptible (e.g., safe mode), and does not allow
24/// velocity change tasks. It can optionally perform a secondary targeting maneuver if a
25/// secondary objective is provided.
26#[derive(Clone)]
27pub(super) struct ZORetrievalMode {
28    /// The primary zoned objective this mode attempts to complete.
29    target: KnownImgObjective,
30    /// An optional second imaging target (used for dual-image objectives).
31    add_target: Option<Vec2D<I32F32>>,
32    /// Unwrapped position of the target objective on the map (absolute), perspective from the burn exit point
33    unwrapped_pos: Arc<Mutex<Vec2D<I32F32>>>,
34}
35
36impl ZORetrievalMode {
37    /// The static name for identification/logging.
38    const MODE_NAME: &'static str = "ZORetrievalMode";
39    /// Default imaging acquisition duration for a single objective.
40    const SINGLE_TARGET_ACQ_DT: TimeDelta = TimeDelta::seconds(10);
41
42    /// Creates a new retrieval mode for the given zoned objective.
43    ///
44    /// # Arguments
45    /// * `target` – The objective to fulfill.
46    /// * `add_target` – Optional second target position for dual-acquisition.
47    /// * `unwrapped_pos` – Global position of the target on the map, perspective from the burn exit point.
48    ///
49    /// # Returns
50    /// * `ZORetrievalMode` – An initialized mode for retrieval.
51    pub(super) fn new(
52        target: KnownImgObjective,
53        add_target: Option<Vec2D<I32F32>>,
54        unwrapped_pos: Vec2D<I32F32>,
55    ) -> Self {
56        let unwrapped_lock = Arc::new(Mutex::new(unwrapped_pos));
57        Self { target, add_target, unwrapped_pos: unwrapped_lock }
58    }
59
60    /// Prepares the async future for imaging, including timing and potential
61    /// turning to a second imaging target.
62    ///
63    /// # Arguments
64    /// * `second_target` – Optional target coordinates.
65    /// * `unwrapped_pos` – Current position.
66    /// * `context` – Shared mode context.
67    ///
68    /// # Returns
69    /// * `(deadline, future)` – Deadline for task completion and associated future.
70    async fn get_img_fut(
71        second_target: Option<Vec2D<I32F32>>,
72        unwrapped_pos: Vec2D<I32F32>,
73        context: &Arc<ModeContext>,
74    ) -> (
75        DateTime<Utc>,
76        Pin<Box<dyn Future<Output = ()> + Send + Sync>>,
77    ) {
78        if let Some(add_target) = second_target {
79            let current_vel = context.k().f_cont().read().await.current_vel();
80            let to_target = {
81                let wrapped = unwrapped_pos.wrap_around_map();
82                wrapped.unwrapped_to(&add_target)
83            };
84            let target_traversal_dt =
85                TimeDelta::seconds((to_target.abs() / current_vel.abs()).to_num::<i64>());
86            let t_end = Utc::now() + Self::SINGLE_TARGET_ACQ_DT * 2 + target_traversal_dt;
87            let fut = FlightComputer::turn_for_2nd_target(context.k().f_cont(), add_target, t_end);
88            (t_end, Box::pin(fut))
89        } else {
90            let sleep_dur_std = Self::SINGLE_TARGET_ACQ_DT.to_std().unwrap_or(DT_0_STD);
91            (
92                Utc::now() + Self::SINGLE_TARGET_ACQ_DT,
93                Box::pin(tokio::time::sleep(sleep_dur_std)),
94            )
95        }
96    }
97
98    /// Executes the full retrieval task including imaging and export/upload.
99    ///
100    /// # Arguments
101    /// * `target` – The zoned objective to complete.
102    /// * `unwrapped_target` – Absolute coordinates for targeting.
103    /// * `second_target` – Optional second target for multi-point objectives.
104    /// * `context` – Shared context.
105    /// * `c_tok` – Cancellation token for task coordination.
106    async fn exec_img_task(
107        target: KnownImgObjective,
108        unwrapped_target: Vec2D<I32F32>,
109        second_target: Option<Vec2D<I32F32>>,
110        context: Arc<ModeContext>,
111        c_tok: CancellationToken,
112    ) {
113        let offset = Vec2D::new(target.zone()[0], target.zone()[1]).to_unsigned();
114        let dim = Vec2D::new(target.width(), target.height()).to_unsigned();
115
116        let c_cont = context.k().c_cont();
117        let (deadline, add_fut) =
118            Self::get_img_fut(second_target, unwrapped_target, &context).await;
119        let f_cont = context.k().f_cont();
120        let mut zoned_objective_image_buffer = None;
121        let img_fut = c_cont.execute_zo_target_cycle(
122            f_cont,
123            deadline,
124            &mut zoned_objective_image_buffer,
125            offset,
126            dim,
127        );
128        tokio::pin!(add_fut);
129        tokio::select! {
130            () = img_fut => FlightComputer::stop_ongoing_burn(context.k().f_cont()).await,
131            () = &mut add_fut => (),
132            () = c_tok.cancelled() => {
133                warn!("Zoned Objective image Task has been cancelled. Cleaning up!");
134                FlightComputer::stop_ongoing_burn(context.k().f_cont()).await;
135            }
136        }
137        let c_cont = context.k().c_cont();
138        let id = target.id();
139        let img_path = Some(CameraController::generate_zo_img_path(id));
140        c_cont
141            .export_and_upload_objective_png(
142                id,
143                offset,
144                dim,
145                img_path,
146                zoned_objective_image_buffer.as_ref(),
147            )
148            .await
149            .unwrap_or_else(|e| {
150                error!("Error exporting and uploading objective image: {e}");
151            });
152    }
153}
154
155#[async_trait]
156impl GlobalMode for ZORetrievalMode {
157    /// Returns the static name of the mode.
158    fn type_name(&self) -> &'static str { Self::MODE_NAME }
159
160    /// Initializes the mode by performing detumbling, scheduling, and target alignment.
161    ///
162    /// # Arguments
163    /// * `context` – Shared context for access to controllers and state.
164    ///
165    /// # Returns
166    /// * `OpExitSignal` – Whether to continue or reinitialize the mode.
167    async fn init_mode(&self, context: Arc<ModeContext>) -> OpExitSignal {
168        let mut unwrapped_pos = self.unwrapped_pos.lock().await;
169        let fut = FlightComputer::detumble_to(
170            context.k().f_cont(),
171            *unwrapped_pos,
172            self.target.optic_required(),
173        );
174        let safe_mon = context.super_v().safe_mon();
175        let target_t;
176        let wrapped_target;
177        let mut handle = tokio::spawn(fut);
178        tokio::select! {
179            join = &mut handle => {
180                let res = join.ok().unwrap();
181                wrapped_target =  res.1;
182                target_t = res.0;
183            },
184            () = safe_mon.notified() => {
185                handle.abort();
186                return self.safe_handler(context).await;
187            }
188        }
189        *unwrapped_pos = wrapped_target;
190        drop(unwrapped_pos);
191        let t_cont = context.k().t_cont();
192        t_cont.clear_schedule().await; // Just to be sure
193        t_cont
194            .schedule_retrieval_phase(
195                target_t,
196                wrapped_target.wrap_around_map(),
197                self.target.optic_required(),
198            )
199            .await;
200        context.k().con().send_tasklist().await;
201        OpExitSignal::Continue
202    }
203
204    /// Waits until the due time of the next task or exits early on a Safe Mode event.
205    ///
206    /// # Arguments
207    /// * `context` – Mode context.
208    /// * `due` – Scheduled execution time.
209    ///
210    /// # Returns
211    /// * `WaitExitSignal` – Indicates continuation or interruption.
212    async fn exec_task_wait(
213        &self,
214        context: Arc<ModeContext>,
215        due: DateTime<Utc>,
216    ) -> WaitExitSignal {
217        let safe_mon = context.super_v().safe_mon();
218        let dt = (due - Utc::now()).to_std().unwrap_or(DT_0_STD);
219        tokio::select! {
220            () = FlightComputer::wait_for_duration(dt, false) => {
221                WaitExitSignal::Continue
222            },
223            () = safe_mon.notified() => {
224                WaitExitSignal::SafeEvent
225            }
226        }
227    }
228
229    /// Executes a given task: imaging or state transition.
230    /// Velocity change tasks are not allowed and result in a logged error.
231    ///
232    /// # Arguments
233    /// * `context` – Shared context.
234    /// * `task` – Task to be executed.
235    ///
236    /// # Returns
237    /// * `ExecExitSignal` – Indicates result of execution.
238    async fn exec_task(&self, context: Arc<ModeContext>, task: Task) -> ExecExitSignal {
239        match task.task_type() {
240            BaseTask::TakeImage(_) => {
241                let safe_mon = context.super_v().safe_mon();
242                let c_tok = CancellationToken::new();
243                let c_tok_clone = c_tok.clone();
244                let context_clone = Arc::clone(&context);
245                let second_target = self.add_target;
246                let unwrapped_target = *self.unwrapped_pos.lock().await;
247                let target = self.target.clone();
248                let img_handle = tokio::spawn(async move {
249                    Self::exec_img_task(
250                        target,
251                        unwrapped_target,
252                        second_target,
253                        context_clone,
254                        c_tok_clone,
255                    )
256                    .await;
257                });
258                tokio::pin!(img_handle);
259                tokio::select! {
260                    _ = &mut img_handle => { },
261                    () = safe_mon.notified() => {
262                        c_tok.cancel();
263                        img_handle.await.unwrap_or_else(|e| {
264                            error!("Error joining zo image task: {e}");
265                        });
266                        return ExecExitSignal::SafeEvent;
267                    }
268                }
269            }
270            BaseTask::SwitchState(switch) => {
271                let f_cont = context.k().f_cont();
272                if matches!(
273                    switch.target_state(),
274                    FlightState::Acquisition | FlightState::Charge
275                ) {
276                    FlightComputer::set_state_wait(f_cont, switch.target_state()).await;
277                } else {
278                    fatal!("Illegal target state!");
279                }
280            }
281            BaseTask::ChangeVelocity(_) => {
282                error!("Change Velocity task is forbidden in ZORetrievalMode.");
283            }
284        }
285        ExecExitSignal::Continue
286    }
287
288    /// Handles a safe-mode event, evaluating whether the current objective is still reachable.
289    /// If reachable, reinitializes in the same mode. Otherwise, exits to `OrbitReturnMode`.
290    ///
291    /// # Arguments
292    /// * `context` – Shared context.
293    ///
294    /// # Returns
295    /// * `OpExitSignal` – ReInit or transition to fallback.
296    async fn safe_handler(&self, context: Arc<ModeContext>) -> OpExitSignal {
297        FlightComputer::escape_safe(context.k().f_cont(), false).await;
298        let (vel, pos) = {
299            let f_cont_locked = context.k().f_cont();
300            let f_cont = f_cont_locked.read().await;
301            (f_cont.current_vel(), f_cont.current_pos())
302        };
303        let to_target = pos.to(&*self.unwrapped_pos.lock().await);
304        let angle = vel.angle_to(&to_target).abs();
305        if angle < I32F32::lit("10.0") {
306            let time_cond = {
307                let state = context.k().f_cont().read().await.state();
308                if state == FlightState::Acquisition {
309                    to_target.abs() > I32F32::lit("10.0") * angle
310                } else {
311                    let transition =
312                        I32F32::from_num(state.dt_to(FlightState::Acquisition).as_secs());
313                    to_target.abs() > I32F32::lit("10.0") * angle + transition
314                }
315            };
316            if time_cond {
317                log!("Objective still reachable after safe event, staying in ZORetrievalMode");
318                FlightComputer::set_state_wait(context.k().f_cont(), FlightState::Acquisition)
319                    .await;
320                return OpExitSignal::ReInit(Box::new(self.clone()));
321            }
322        }
323        warn!("Objective not reachable after safe event, exiting ZORetrievalMode");
324        context.o_ch_lock().write().await.finish(
325            context.k().f_cont().read().await.current_pos(),
326            self.out_of_orbit_rationale(),
327        );
328        OpExitSignal::ReInit(Box::new(OrbitReturnMode::new()))
329    }
330
331    /// Not implemented – ZO handoffs do not apply during retrieval phase.
332    async fn zo_handler(&self, _: &Arc<ModeContext>, _: KnownImgObjective) -> OptOpExitSignal {
333        unimplemented!()
334    }
335
336    /// Not implemented – Beacon Objective events are ignored in this mode.
337    async fn bo_event_handler(&self, _: &Arc<ModeContext>) -> OptOpExitSignal { unimplemented!() }
338
339    /// Finalizes the retrieval mode and transitions to `OrbitReturnMode`.
340    ///
341    /// # Arguments
342    /// * `context` – Shared context.
343    ///
344    /// # Returns
345    /// * `Box<dyn GlobalMode>` – Next mode to execute.
346    async fn exit_mode(&self, context: Arc<ModeContext>) -> Box<dyn GlobalMode> {
347        context.o_ch_lock().write().await.finish(
348            context.k().f_cont().read().await.current_pos(),
349            self.tasks_done_rationale(),
350        );
351        Box::new(OrbitReturnMode::new())
352    }
353}