omega_brain/
sleep_system.rs

1//! Sleep System - Self-contained sleep stages and circadian rhythm
2
3use crate::{BrainConfig, Result};
4use serde::{Deserialize, Serialize};
5
6/// Sleep stages
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
8pub enum SleepStage { Awake, N1, N2, N3, REM }
9
10impl SleepStage {
11    pub fn name(&self) -> &str {
12        match self { Self::Awake => "Awake", Self::N1 => "N1", Self::N2 => "N2", Self::N3 => "N3", Self::REM => "REM" }
13    }
14}
15
16/// Sleep controller
17#[derive(Debug, Clone)]
18pub struct SleepController {
19    stage: SleepStage,
20    cycle_position: usize,
21    cycle_length: usize,
22    sws_ratio: f64,
23    rem_ratio: f64,
24}
25
26impl SleepController {
27    pub fn new() -> Self {
28        Self { stage: SleepStage::Awake, cycle_position: 0, cycle_length: 1000, sws_ratio: 0.6, rem_ratio: 0.25 }
29    }
30    pub fn start_sleep(&mut self) { self.stage = SleepStage::N1; self.cycle_position = 0; }
31    pub fn end_sleep(&mut self) { self.stage = SleepStage::Awake; }
32    pub fn advance(&mut self) -> SleepStage {
33        self.cycle_position += 1;
34        let p = self.cycle_position as f64 / self.cycle_length as f64;
35        self.stage = if p < 0.1 { SleepStage::N1 }
36            else if p < 0.3 { SleepStage::N2 }
37            else if p < 0.3 + self.sws_ratio * 0.5 { SleepStage::N3 }
38            else if p < 1.0 - self.rem_ratio { SleepStage::N2 }
39            else { SleepStage::REM };
40        self.stage
41    }
42    pub fn current_stage(&self) -> SleepStage { self.stage }
43    pub fn should_wake(&self) -> bool { self.cycle_position >= self.cycle_length }
44    pub fn sleep_depth(&self) -> f64 {
45        match self.stage { SleepStage::Awake => 0.0, SleepStage::N1 => 0.2, SleepStage::N2 => 0.4, SleepStage::N3 => 0.9, SleepStage::REM => 0.3 }
46    }
47}
48
49impl Default for SleepController {
50    fn default() -> Self { Self::new() }
51}
52
53/// Circadian rhythm
54#[derive(Debug, Clone)]
55pub struct CircadianRhythm {
56    phase: f64,
57}
58
59impl CircadianRhythm {
60    pub fn new() -> Self { Self { phase: 0.5 } }
61    pub fn current_phase(&self) -> f64 { self.phase }
62    pub fn advance(&mut self, hours: f64) { self.phase = (self.phase + hours / 24.0) % 1.0; }
63    pub fn should_sleep(&self) -> bool { self.phase > 0.75 || self.phase < 0.25 }
64}
65
66impl Default for CircadianRhythm {
67    fn default() -> Self { Self::new() }
68}
69
70/// Sleep output
71#[derive(Debug, Clone)]
72pub struct SleepOutput {
73    pub is_sws: bool,
74    pub is_rem: bool,
75    pub replay_count: usize,
76    pub depth: f64,
77    pub stage_name: String,
78}
79
80/// Sleep system
81pub struct SleepSystem {
82    controller: SleepController,
83    circadian: CircadianRhythm,
84    is_sleeping: bool,
85    cycles_completed: usize,
86}
87
88impl SleepSystem {
89    pub fn new(_config: &BrainConfig) -> Self {
90        Self { controller: SleepController::new(), circadian: CircadianRhythm::new(), is_sleeping: false, cycles_completed: 0 }
91    }
92    pub fn should_sleep(&self) -> bool { self.is_sleeping }
93    pub fn initiate_sleep(&mut self) -> Result<()> {
94        if !self.is_sleeping { self.is_sleeping = true; self.controller.start_sleep(); }
95        Ok(())
96    }
97    pub fn wake_up(&mut self) -> Result<()> {
98        if self.is_sleeping { self.is_sleeping = false; self.controller.end_sleep(); self.cycles_completed += 1; }
99        Ok(())
100    }
101    pub fn process_cycle(&mut self) -> Result<SleepOutput> {
102        if !self.is_sleeping {
103            return Ok(SleepOutput { is_sws: false, is_rem: false, replay_count: 0, depth: 0.0, stage_name: "Awake".to_string() });
104        }
105        let stage = self.controller.advance();
106        let (is_sws, is_rem, replay_count, depth) = match stage {
107            SleepStage::N1 => (false, false, 0, 0.2),
108            SleepStage::N2 => (false, false, 5, 0.4),
109            SleepStage::N3 => (true, false, 20, 0.9),
110            SleepStage::REM => (false, true, 10, 0.3),
111            SleepStage::Awake => (false, false, 0, 0.0),
112        };
113        if self.controller.should_wake() { self.wake_up()?; }
114        Ok(SleepOutput { is_sws, is_rem, replay_count, depth, stage_name: stage.name().to_string() })
115    }
116    pub fn current_stage_name(&self) -> Option<String> {
117        if self.is_sleeping { Some(self.controller.current_stage().name().to_string()) } else { None }
118    }
119    pub fn is_rem(&self) -> bool { self.is_sleeping && self.controller.current_stage() == SleepStage::REM }
120    pub fn is_sws(&self) -> bool { self.is_sleeping && self.controller.current_stage() == SleepStage::N3 }
121    pub fn sleep_depth(&self) -> f64 { if !self.is_sleeping { 0.0 } else { self.controller.sleep_depth() } }
122    pub fn cycles_completed(&self) -> usize { self.cycles_completed }
123    pub fn circadian_phase(&self) -> f64 { self.circadian.current_phase() }
124    pub fn update_circadian(&mut self, hours: f64) { self.circadian.advance(hours); }
125    pub fn circadian_suggests_sleep(&self) -> bool { self.circadian.should_sleep() }
126    pub fn reset(&mut self) {
127        self.controller = SleepController::new();
128        self.circadian = CircadianRhythm::new();
129        self.is_sleeping = false;
130        self.cycles_completed = 0;
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    #[test]
138    fn test_sleep_controller() {
139        let mut ctrl = SleepController::new();
140        ctrl.start_sleep();
141        assert_eq!(ctrl.current_stage(), SleepStage::N1);
142    }
143}
omega_brain/sleep_system.rs

omega_brain/
sleep_system.rs
