rscm_core/python/
state.rs

1//! Python bindings for typed state access
2//!
3//! This module provides Python wrappers for TimeseriesWindow, GridTimeseriesWindow,
4//! and typed output slices (FourBoxSlice, HemisphericSlice).
5
6use crate::spatial::{FourBoxRegion, HemisphericRegion};
7use crate::state::{FourBoxSlice, HemisphericSlice, StateValue};
8use crate::timeseries::FloatValue;
9use numpy::{PyArray1, ToPyArray};
10use pyo3::exceptions::PyValueError;
11use pyo3::prelude::*;
12use pyo3::{pymodule, Bound, PyResult};
13
14// =============================================================================
15// Python Typed Output Slices
16// =============================================================================
17
18/// Python wrapper for FourBoxSlice
19///
20/// Provides type-safe access to regional values with named keyword arguments.
21///
22/// Example:
23///     slice = FourBoxSlice(
24///         northern_ocean=15.0,
25///         northern_land=14.0,
26///         southern_ocean=10.0,
27///         southern_land=9.0
28///     )
29#[pyclass]
30#[pyo3(name = "FourBoxSlice")]
31#[derive(Debug, Clone)]
32pub struct PyFourBoxSlice(pub FourBoxSlice);
33
34#[pymethods]
35impl PyFourBoxSlice {
36    #[new]
37    #[pyo3(signature = (northern_ocean=f64::NAN, northern_land=f64::NAN, southern_ocean=f64::NAN, southern_land=f64::NAN))]
38    fn new(
39        northern_ocean: FloatValue,
40        northern_land: FloatValue,
41        southern_ocean: FloatValue,
42        southern_land: FloatValue,
43    ) -> Self {
44        Self(FourBoxSlice::from_array([
45            northern_ocean,
46            northern_land,
47            southern_ocean,
48            southern_land,
49        ]))
50    }
51
52    /// Create a slice with all regions set to the same value
53    #[staticmethod]
54    fn uniform(value: FloatValue) -> Self {
55        Self(FourBoxSlice::uniform(value))
56    }
57
58    /// Create a slice from an array [northern_ocean, northern_land, southern_ocean, southern_land]
59    #[staticmethod]
60    fn from_array(values: [FloatValue; 4]) -> Self {
61        Self(FourBoxSlice::from_array(values))
62    }
63
64    /// Get the northern ocean value
65    #[getter]
66    fn northern_ocean(&self) -> FloatValue {
67        self.0.get(FourBoxRegion::NorthernOcean)
68    }
69
70    /// Set the northern ocean value
71    #[setter]
72    fn set_northern_ocean(&mut self, value: FloatValue) {
73        self.0.set(FourBoxRegion::NorthernOcean, value);
74    }
75
76    /// Get the northern land value
77    #[getter]
78    fn northern_land(&self) -> FloatValue {
79        self.0.get(FourBoxRegion::NorthernLand)
80    }
81
82    /// Set the northern land value
83    #[setter]
84    fn set_northern_land(&mut self, value: FloatValue) {
85        self.0.set(FourBoxRegion::NorthernLand, value);
86    }
87
88    /// Get the southern ocean value
89    #[getter]
90    fn southern_ocean(&self) -> FloatValue {
91        self.0.get(FourBoxRegion::SouthernOcean)
92    }
93
94    /// Set the southern ocean value
95    #[setter]
96    fn set_southern_ocean(&mut self, value: FloatValue) {
97        self.0.set(FourBoxRegion::SouthernOcean, value);
98    }
99
100    /// Get the southern land value
101    #[getter]
102    fn southern_land(&self) -> FloatValue {
103        self.0.get(FourBoxRegion::SouthernLand)
104    }
105
106    /// Set the southern land value
107    #[setter]
108    fn set_southern_land(&mut self, value: FloatValue) {
109        self.0.set(FourBoxRegion::SouthernLand, value);
110    }
111
112    /// Get value by region index
113    fn get(&self, region: usize) -> PyResult<FloatValue> {
114        match region {
115            0 => Ok(self.0.get(FourBoxRegion::NorthernOcean)),
116            1 => Ok(self.0.get(FourBoxRegion::NorthernLand)),
117            2 => Ok(self.0.get(FourBoxRegion::SouthernOcean)),
118            3 => Ok(self.0.get(FourBoxRegion::SouthernLand)),
119            _ => Err(PyValueError::new_err(format!(
120                "Invalid region index: {}. Must be 0-3.",
121                region
122            ))),
123        }
124    }
125
126    /// Set value by region index
127    fn set(&mut self, region: usize, value: FloatValue) -> PyResult<()> {
128        match region {
129            0 => self.0.set(FourBoxRegion::NorthernOcean, value),
130            1 => self.0.set(FourBoxRegion::NorthernLand, value),
131            2 => self.0.set(FourBoxRegion::SouthernOcean, value),
132            3 => self.0.set(FourBoxRegion::SouthernLand, value),
133            _ => {
134                return Err(PyValueError::new_err(format!(
135                    "Invalid region index: {}. Must be 0-3.",
136                    region
137                )))
138            }
139        }
140        Ok(())
141    }
142
143    /// Convert to numpy array
144    fn to_array<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<FloatValue>> {
145        self.0.as_array().to_pyarray(py)
146    }
147
148    /// Convert to list
149    fn to_list(&self) -> Vec<FloatValue> {
150        self.0.to_vec()
151    }
152
153    /// Convert to dict with region names as keys
154    fn to_dict(&self) -> std::collections::HashMap<String, FloatValue> {
155        let mut map = std::collections::HashMap::new();
156        map.insert(
157            "northern_ocean".to_string(),
158            self.0.get(FourBoxRegion::NorthernOcean),
159        );
160        map.insert(
161            "northern_land".to_string(),
162            self.0.get(FourBoxRegion::NorthernLand),
163        );
164        map.insert(
165            "southern_ocean".to_string(),
166            self.0.get(FourBoxRegion::SouthernOcean),
167        );
168        map.insert(
169            "southern_land".to_string(),
170            self.0.get(FourBoxRegion::SouthernLand),
171        );
172        map
173    }
174
175    fn __repr__(&self) -> String {
176        format!(
177            "FourBoxSlice(northern_ocean={}, northern_land={}, southern_ocean={}, southern_land={})",
178            self.0.get(FourBoxRegion::NorthernOcean),
179            self.0.get(FourBoxRegion::NorthernLand),
180            self.0.get(FourBoxRegion::SouthernOcean),
181            self.0.get(FourBoxRegion::SouthernLand)
182        )
183    }
184
185    fn __getitem__(&self, index: usize) -> PyResult<FloatValue> {
186        self.get(index)
187    }
188
189    fn __setitem__(&mut self, index: usize, value: FloatValue) -> PyResult<()> {
190        self.set(index, value)
191    }
192
193    fn __len__(&self) -> usize {
194        4
195    }
196}
197
198/// Python wrapper for HemisphericSlice
199///
200/// Provides type-safe access to hemispheric values with named keyword arguments.
201///
202/// Example:
203///     slice = HemisphericSlice(northern=15.0, southern=10.0)
204#[pyclass]
205#[pyo3(name = "HemisphericSlice")]
206#[derive(Debug, Clone)]
207pub struct PyHemisphericSlice(pub HemisphericSlice);
208
209#[pymethods]
210impl PyHemisphericSlice {
211    #[new]
212    #[pyo3(signature = (northern=f64::NAN, southern=f64::NAN))]
213    fn new(northern: FloatValue, southern: FloatValue) -> Self {
214        Self(HemisphericSlice::from_array([northern, southern]))
215    }
216
217    /// Create a slice with both hemispheres set to the same value
218    #[staticmethod]
219    fn uniform(value: FloatValue) -> Self {
220        Self(HemisphericSlice::uniform(value))
221    }
222
223    /// Create a slice from an array [northern, southern]
224    #[staticmethod]
225    fn from_array(values: [FloatValue; 2]) -> Self {
226        Self(HemisphericSlice::from_array(values))
227    }
228
229    /// Get the northern hemisphere value
230    #[getter]
231    fn northern(&self) -> FloatValue {
232        self.0.get(HemisphericRegion::Northern)
233    }
234
235    /// Set the northern hemisphere value
236    #[setter]
237    fn set_northern(&mut self, value: FloatValue) {
238        self.0.set(HemisphericRegion::Northern, value);
239    }
240
241    /// Get the southern hemisphere value
242    #[getter]
243    fn southern(&self) -> FloatValue {
244        self.0.get(HemisphericRegion::Southern)
245    }
246
247    /// Set the southern hemisphere value
248    #[setter]
249    fn set_southern(&mut self, value: FloatValue) {
250        self.0.set(HemisphericRegion::Southern, value);
251    }
252
253    /// Get value by region index
254    fn get(&self, region: usize) -> PyResult<FloatValue> {
255        match region {
256            0 => Ok(self.0.get(HemisphericRegion::Northern)),
257            1 => Ok(self.0.get(HemisphericRegion::Southern)),
258            _ => Err(PyValueError::new_err(format!(
259                "Invalid region index: {}. Must be 0-1.",
260                region
261            ))),
262        }
263    }
264
265    /// Set value by region index
266    fn set(&mut self, region: usize, value: FloatValue) -> PyResult<()> {
267        match region {
268            0 => self.0.set(HemisphericRegion::Northern, value),
269            1 => self.0.set(HemisphericRegion::Southern, value),
270            _ => {
271                return Err(PyValueError::new_err(format!(
272                    "Invalid region index: {}. Must be 0-1.",
273                    region
274                )))
275            }
276        }
277        Ok(())
278    }
279
280    /// Convert to numpy array
281    fn to_array<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<FloatValue>> {
282        self.0.as_array().to_pyarray(py)
283    }
284
285    /// Convert to list
286    fn to_list(&self) -> Vec<FloatValue> {
287        self.0.to_vec()
288    }
289
290    /// Convert to dict with region names as keys
291    fn to_dict(&self) -> std::collections::HashMap<String, FloatValue> {
292        let mut map = std::collections::HashMap::new();
293        map.insert(
294            "northern".to_string(),
295            self.0.get(HemisphericRegion::Northern),
296        );
297        map.insert(
298            "southern".to_string(),
299            self.0.get(HemisphericRegion::Southern),
300        );
301        map
302    }
303
304    fn __repr__(&self) -> String {
305        format!(
306            "HemisphericSlice(northern={}, southern={})",
307            self.0.get(HemisphericRegion::Northern),
308            self.0.get(HemisphericRegion::Southern)
309        )
310    }
311
312    fn __getitem__(&self, index: usize) -> PyResult<FloatValue> {
313        self.get(index)
314    }
315
316    fn __setitem__(&mut self, index: usize, value: FloatValue) -> PyResult<()> {
317        self.set(index, value)
318    }
319
320    fn __len__(&self) -> usize {
321        2
322    }
323}
324
325// =============================================================================
326// Python StateValue
327// =============================================================================
328
329/// Python wrapper for StateValue enum
330///
331/// Represents a value that can be scalar or spatially-resolved (FourBox/Hemispheric).
332///
333/// Example:
334///     scalar = StateValue.scalar(15.0)
335///     four_box = StateValue.four_box(FourBoxSlice.uniform(10.0))
336///     hemispheric = StateValue.hemispheric(HemisphericSlice(northern=15.0, southern=10.0))
337#[pyclass]
338#[pyo3(name = "StateValue")]
339#[derive(Debug, Clone)]
340pub struct PyStateValue(pub StateValue);
341
342#[pymethods]
343impl PyStateValue {
344    /// Create a scalar StateValue
345    #[staticmethod]
346    fn scalar(value: FloatValue) -> Self {
347        Self(StateValue::Scalar(value))
348    }
349
350    /// Create a FourBox StateValue from a FourBoxSlice
351    #[staticmethod]
352    fn four_box(slice: PyFourBoxSlice) -> Self {
353        Self(StateValue::FourBox(slice.0))
354    }
355
356    /// Create a Hemispheric StateValue from a HemisphericSlice
357    #[staticmethod]
358    fn hemispheric(slice: PyHemisphericSlice) -> Self {
359        Self(StateValue::Hemispheric(slice.0))
360    }
361
362    /// Check if this is a scalar value
363    fn is_scalar(&self) -> bool {
364        self.0.is_scalar()
365    }
366
367    /// Check if this is a FourBox grid value
368    fn is_four_box(&self) -> bool {
369        self.0.is_four_box()
370    }
371
372    /// Check if this is a Hemispheric grid value
373    fn is_hemispheric(&self) -> bool {
374        self.0.is_hemispheric()
375    }
376
377    /// Get the scalar value if this is a Scalar variant, otherwise None
378    fn as_scalar(&self) -> Option<FloatValue> {
379        self.0.as_scalar()
380    }
381
382    /// Get the FourBoxSlice if this is a FourBox variant, otherwise None
383    fn as_four_box(&self) -> Option<PyFourBoxSlice> {
384        self.0.as_four_box().map(|s| PyFourBoxSlice(*s))
385    }
386
387    /// Get the HemisphericSlice if this is a Hemispheric variant, otherwise None
388    fn as_hemispheric(&self) -> Option<PyHemisphericSlice> {
389        self.0.as_hemispheric().map(|s| PyHemisphericSlice(*s))
390    }
391
392    /// Convert to a scalar value, aggregating grid values if necessary
393    ///
394    /// For Scalar: returns the value directly
395    /// For FourBox: returns the arithmetic mean of all 4 regional values
396    ///              (sum of all regions divided by 4, giving equal weight to each region)
397    /// For Hemispheric: returns the arithmetic mean of both hemispheres
398    ///                  (sum of northern and southern divided by 2, giving equal weight to each)
399    fn to_scalar(&self) -> FloatValue {
400        self.0.to_scalar()
401    }
402
403    fn __repr__(&self) -> String {
404        match &self.0 {
405            StateValue::Scalar(v) => format!("StateValue.scalar({})", v),
406            StateValue::FourBox(slice) => {
407                format!(
408                    "StateValue.four_box(FourBoxSlice(northern_ocean={}, northern_land={}, southern_ocean={}, southern_land={}))",
409                    slice.get(FourBoxRegion::NorthernOcean),
410                    slice.get(FourBoxRegion::NorthernLand),
411                    slice.get(FourBoxRegion::SouthernOcean),
412                    slice.get(FourBoxRegion::SouthernLand)
413                )
414            }
415            StateValue::Hemispheric(slice) => {
416                format!(
417                    "StateValue.hemispheric(HemisphericSlice(northern={}, southern={}))",
418                    slice.get(HemisphericRegion::Northern),
419                    slice.get(HemisphericRegion::Southern)
420                )
421            }
422        }
423    }
424}
425
426// =============================================================================
427// Python TimeseriesWindow
428// =============================================================================
429
430/// Python wrapper for TimeseriesWindow
431///
432/// Provides zero-cost view-like access to scalar timeseries data.
433/// Since Python can't handle Rust lifetimes, this stores a copy of the
434/// values and current position.
435///
436/// Example:
437///     window = TimeseriesWindow(values=[1.0, 2.0, 3.0], current_index=2)
438///     print(window.current)  # 3.0
439///     print(window.previous)  # 2.0
440#[pyclass]
441#[pyo3(name = "TimeseriesWindow")]
442#[derive(Debug, Clone)]
443pub struct PyTimeseriesWindow {
444    values: Vec<FloatValue>,
445    current_index: usize,
446}
447
448#[pymethods]
449impl PyTimeseriesWindow {
450    #[new]
451    #[pyo3(signature = (values, current_index))]
452    pub fn new(values: Vec<FloatValue>, current_index: usize) -> PyResult<Self> {
453        if current_index >= values.len() && !values.is_empty() {
454            return Err(PyValueError::new_err(format!(
455                "current_index {} out of bounds for values of length {}",
456                current_index,
457                values.len()
458            )));
459        }
460        Ok(Self {
461            values,
462            current_index,
463        })
464    }
465
466    /// Get the value at the start of the timestep (index N).
467    ///
468    /// Use this for:
469    /// - State variables (reading your own previous state)
470    /// - Exogenous inputs (external forcing data)
471    fn at_start(&self) -> PyResult<FloatValue> {
472        self.values
473            .get(self.current_index)
474            .copied()
475            .ok_or_else(|| PyValueError::new_err("No value available at current index"))
476    }
477
478    /// Get the value at the end of the timestep (index N+1), if available.
479    ///
480    /// Use this for:
481    /// - Upstream component outputs (values written during this timestep)
482    ///
483    /// Returns None if at the last timestep.
484    fn at_end(&self) -> Option<FloatValue> {
485        let next_index = self.current_index + 1;
486        self.values.get(next_index).copied()
487    }
488
489    /// Get the previous value (at current_index - 1)
490    ///
491    /// Deprecated: Use `at_offset(-1)` for consistent timestep-relative access.
492    #[getter]
493    fn previous(&self) -> PyResult<FloatValue> {
494        if self.current_index == 0 {
495            return Err(PyValueError::new_err("No previous value available"));
496        }
497        self.values
498            .get(self.current_index - 1)
499            .copied()
500            .ok_or_else(|| PyValueError::new_err("No previous value available"))
501    }
502
503    /// Get value at offset from current (negative = past, positive = future)
504    fn at_offset(&self, offset: isize) -> PyResult<FloatValue> {
505        let index = self.current_index as isize + offset;
506        if index < 0 || index >= self.values.len() as isize {
507            return Err(PyValueError::new_err(format!(
508                "Offset {} results in index {} which is out of bounds",
509                offset, index
510            )));
511        }
512        Ok(self.values[index as usize])
513    }
514
515    /// Get last n values as numpy array (including current)
516    fn last_n<'py>(&self, py: Python<'py>, n: usize) -> PyResult<Bound<'py, PyArray1<FloatValue>>> {
517        if n == 0 {
518            return Ok(Vec::<FloatValue>::new().to_pyarray(py));
519        }
520        let start = (self.current_index + 1).saturating_sub(n);
521        let end = self.current_index + 1;
522        let slice: Vec<FloatValue> = self.values[start..end].to_vec();
523        Ok(slice.to_pyarray(py))
524    }
525
526    /// Get all values as numpy array
527    fn to_array<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<FloatValue>> {
528        self.values.clone().to_pyarray(py)
529    }
530
531    /// Get the current index
532    #[getter]
533    fn current_index(&self) -> usize {
534        self.current_index
535    }
536
537    /// Get the number of values
538    fn __len__(&self) -> usize {
539        self.values.len()
540    }
541
542    fn __repr__(&self) -> String {
543        format!(
544            "TimeseriesWindow(len={}, current_index={}, current={})",
545            self.values.len(),
546            self.current_index,
547            self.values
548                .get(self.current_index)
549                .map(|v| v.to_string())
550                .unwrap_or_else(|| "N/A".to_string())
551        )
552    }
553}
554
555/// Python wrapper for GridTimeseriesWindow (FourBox)
556///
557/// Provides view-like access to grid timeseries data with regional access.
558///
559/// Example:
560///     window = FourBoxTimeseriesWindow(values=\[\[1,2,3,4\], \[5,6,7,8\]\], current_index=1)
561///     print(window.current)  # [5, 6, 7, 8]
562///     print(window.region(0))  # TimeseriesWindow for northern ocean
563#[pyclass]
564#[pyo3(name = "FourBoxTimeseriesWindow")]
565#[derive(Debug, Clone)]
566pub struct PyFourBoxTimeseriesWindow {
567    /// Values stored as [timestep][region]
568    values: Vec<[FloatValue; 4]>,
569    current_index: usize,
570}
571
572#[pymethods]
573impl PyFourBoxTimeseriesWindow {
574    #[new]
575    #[pyo3(signature = (values, current_index))]
576    pub fn new(values: Vec<[FloatValue; 4]>, current_index: usize) -> PyResult<Self> {
577        if current_index >= values.len() && !values.is_empty() {
578            return Err(PyValueError::new_err(format!(
579                "current_index {} out of bounds for values of length {}",
580                current_index,
581                values.len()
582            )));
583        }
584        Ok(Self {
585            values,
586            current_index,
587        })
588    }
589
590    /// Get a single region's value at the start of the timestep (index N).
591    ///
592    /// Use this for:
593    /// - State variables (reading your own previous state)
594    /// - Exogenous inputs (external forcing data)
595    fn at_start(&self, region: usize) -> PyResult<FloatValue> {
596        if region >= 4 {
597            return Err(PyValueError::new_err(format!(
598                "Invalid region index: {}. Must be 0-3.",
599                region
600            )));
601        }
602        self.values
603            .get(self.current_index)
604            .map(|v| v[region])
605            .ok_or_else(|| PyValueError::new_err("No value available at current index"))
606    }
607
608    /// Get a single region's value at the end of the timestep (index N+1), if available.
609    ///
610    /// Use this for:
611    /// - Upstream component outputs (values written during this timestep)
612    ///
613    /// Returns None if at the last timestep.
614    fn at_end(&self, region: usize) -> PyResult<Option<FloatValue>> {
615        if region >= 4 {
616            return Err(PyValueError::new_err(format!(
617                "Invalid region index: {}. Must be 0-3.",
618                region
619            )));
620        }
621        let next_index = self.current_index + 1;
622        Ok(self.values.get(next_index).map(|v| v[region]))
623    }
624
625    /// Get all regional values at the start of the timestep (index N).
626    fn at_start_all(&self) -> PyResult<PyFourBoxSlice> {
627        self.values
628            .get(self.current_index)
629            .map(|v| PyFourBoxSlice(FourBoxSlice::from_array(*v)))
630            .ok_or_else(|| PyValueError::new_err("No value available at current index"))
631    }
632
633    /// Get all regional values at the end of the timestep (index N+1), if available.
634    ///
635    /// Returns None if at the last timestep.
636    fn at_end_all(&self) -> Option<PyFourBoxSlice> {
637        let next_index = self.current_index + 1;
638        self.values
639            .get(next_index)
640            .map(|v| PyFourBoxSlice(FourBoxSlice::from_array(*v)))
641    }
642
643    /// Get the previous slice as FourBoxSlice
644    ///
645    /// Deprecated: Use `region(i).at_offset(-1)` for consistent timestep-relative access.
646    #[getter]
647    fn previous(&self) -> PyResult<PyFourBoxSlice> {
648        if self.current_index == 0 {
649            return Err(PyValueError::new_err("No previous value available"));
650        }
651        self.values
652            .get(self.current_index - 1)
653            .map(|v| PyFourBoxSlice(FourBoxSlice::from_array(*v)))
654            .ok_or_else(|| PyValueError::new_err("No previous value available"))
655    }
656
657    /// Get a single region's timeseries as a TimeseriesWindow
658    fn region(&self, region: usize) -> PyResult<PyTimeseriesWindow> {
659        if region >= 4 {
660            return Err(PyValueError::new_err(format!(
661                "Invalid region index: {}. Must be 0-3.",
662                region
663            )));
664        }
665        let values: Vec<FloatValue> = self.values.iter().map(|v| v[region]).collect();
666        Ok(PyTimeseriesWindow {
667            values,
668            current_index: self.current_index,
669        })
670    }
671
672    /// Get current index
673    #[getter]
674    fn current_index(&self) -> usize {
675        self.current_index
676    }
677
678    fn __len__(&self) -> usize {
679        self.values.len()
680    }
681
682    fn __repr__(&self) -> String {
683        format!(
684            "FourBoxTimeseriesWindow(len={}, current_index={})",
685            self.values.len(),
686            self.current_index
687        )
688    }
689}
690
691/// Python wrapper for GridTimeseriesWindow (Hemispheric)
692#[pyclass]
693#[pyo3(name = "HemisphericTimeseriesWindow")]
694#[derive(Debug, Clone)]
695pub struct PyHemisphericTimeseriesWindow {
696    values: Vec<[FloatValue; 2]>,
697    current_index: usize,
698}
699
700#[pymethods]
701impl PyHemisphericTimeseriesWindow {
702    #[new]
703    #[pyo3(signature = (values, current_index))]
704    pub fn new(values: Vec<[FloatValue; 2]>, current_index: usize) -> PyResult<Self> {
705        if current_index >= values.len() && !values.is_empty() {
706            return Err(PyValueError::new_err(format!(
707                "current_index {} out of bounds for values of length {}",
708                current_index,
709                values.len()
710            )));
711        }
712        Ok(Self {
713            values,
714            current_index,
715        })
716    }
717
718    /// Get a single region's value at the start of the timestep (index N).
719    ///
720    /// Use this for:
721    /// - State variables (reading your own previous state)
722    /// - Exogenous inputs (external forcing data)
723    fn at_start(&self, region: usize) -> PyResult<FloatValue> {
724        if region >= 2 {
725            return Err(PyValueError::new_err(format!(
726                "Invalid region index: {}. Must be 0-1.",
727                region
728            )));
729        }
730        self.values
731            .get(self.current_index)
732            .map(|v| v[region])
733            .ok_or_else(|| PyValueError::new_err("No value available at current index"))
734    }
735
736    /// Get a single region's value at the end of the timestep (index N+1), if available.
737    ///
738    /// Use this for:
739    /// - Upstream component outputs (values written during this timestep)
740    ///
741    /// Returns None if at the last timestep.
742    fn at_end(&self, region: usize) -> PyResult<Option<FloatValue>> {
743        if region >= 2 {
744            return Err(PyValueError::new_err(format!(
745                "Invalid region index: {}. Must be 0-1.",
746                region
747            )));
748        }
749        let next_index = self.current_index + 1;
750        Ok(self.values.get(next_index).map(|v| v[region]))
751    }
752
753    /// Get all regional values at the start of the timestep (index N).
754    fn at_start_all(&self) -> PyResult<PyHemisphericSlice> {
755        self.values
756            .get(self.current_index)
757            .map(|v| PyHemisphericSlice(HemisphericSlice::from_array(*v)))
758            .ok_or_else(|| PyValueError::new_err("No value available at current index"))
759    }
760
761    /// Get all regional values at the end of the timestep (index N+1), if available.
762    ///
763    /// Returns None if at the last timestep.
764    fn at_end_all(&self) -> Option<PyHemisphericSlice> {
765        let next_index = self.current_index + 1;
766        self.values
767            .get(next_index)
768            .map(|v| PyHemisphericSlice(HemisphericSlice::from_array(*v)))
769    }
770
771    /// Get the previous slice as HemisphericSlice
772    ///
773    /// Deprecated: Use `region(i).at_offset(-1)` for consistent timestep-relative access.
774    #[getter]
775    fn previous(&self) -> PyResult<PyHemisphericSlice> {
776        if self.current_index == 0 {
777            return Err(PyValueError::new_err("No previous value available"));
778        }
779        self.values
780            .get(self.current_index - 1)
781            .map(|v| PyHemisphericSlice(HemisphericSlice::from_array(*v)))
782            .ok_or_else(|| PyValueError::new_err("No previous value available"))
783    }
784
785    /// Get a single region's timeseries as a TimeseriesWindow
786    fn region(&self, region: usize) -> PyResult<PyTimeseriesWindow> {
787        if region >= 2 {
788            return Err(PyValueError::new_err(format!(
789                "Invalid region index: {}. Must be 0-1.",
790                region
791            )));
792        }
793        let values: Vec<FloatValue> = self.values.iter().map(|v| v[region]).collect();
794        Ok(PyTimeseriesWindow {
795            values,
796            current_index: self.current_index,
797        })
798    }
799
800    /// Get current index
801    #[getter]
802    fn current_index(&self) -> usize {
803        self.current_index
804    }
805
806    fn __len__(&self) -> usize {
807        self.values.len()
808    }
809
810    fn __repr__(&self) -> String {
811        format!(
812            "HemisphericTimeseriesWindow(len={}, current_index={})",
813            self.values.len(),
814            self.current_index
815        )
816    }
817}
818
819// =============================================================================
820// GridType Python Enum
821// =============================================================================
822
823/// Re-export GridType for Python
824pub use crate::component::GridType;
825
826#[pymodule]
827pub fn state(m: &Bound<'_, PyModule>) -> PyResult<()> {
828    m.add_class::<PyFourBoxSlice>()?;
829    m.add_class::<PyHemisphericSlice>()?;
830    m.add_class::<PyStateValue>()?;
831    m.add_class::<PyTimeseriesWindow>()?;
832    m.add_class::<PyFourBoxTimeseriesWindow>()?;
833    m.add_class::<PyHemisphericTimeseriesWindow>()?;
834    Ok(())
835}
rscm_core/python/state.rs

rscm_core/python/
state.rs
