use std::collections::HashMap; use anyhow::Context; use chrono::Month; use serde::{Deserialize, Serialize}; use crate::{ resident::{Resident, ResidentDTO, ResidentId}, schedule::ShiftType, slot::{Day, Slot}, }; const MONTH: u8 = 4; const YEAR: i32 = 2026; #[derive(Debug, Clone)] pub struct ToxicPair((ResidentId, ResidentId)); impl ToxicPair { pub fn new(r_id_1: u8, r_id_2: u8) -> Self { Self((ResidentId(r_id_1), ResidentId(r_id_2))) } pub fn matches(&self, other: &ToxicPair) -> bool { let p1 = &self.0; let p2 = &other.0; (p1.0 == p2.0 && p1.1 == p2.1) || (p1.0 == p2.1 && p1.1 == p2.0) } } impl From<(ResidentId, ResidentId)> for ToxicPair { fn from(value: (ResidentId, ResidentId)) -> Self { Self((value.0, value.1)) } } #[derive(Serialize, Deserialize, Debug, Clone)] pub struct UserConfigDTO { month: u8, year: i32, holidays: Vec, residents: Vec, toxic_pairs: Vec<(u8, u8)>, } #[derive(Debug, Clone)] pub struct UserConfig { pub month: Month, pub year: i32, pub holidays: Vec, pub residents: Vec, pub toxic_pairs: Vec, pub total_days: u8, pub total_slots: u8, pub total_holiday_slots: u8, } impl UserConfig { pub fn new(month: u8, year: i32) -> Self { let month = Month::try_from(month).unwrap(); let total_days = month.num_days(year).unwrap(); let total_slots = (1..=total_days) .map(|d| if Day(d).is_open_shift() { 2 } else { 1 }) .sum(); let total_holiday_slots = (1..=total_days) .filter(|&d| Day(d).is_weekend(month.number_from_month(), year)) .map(|d| if Day(d).is_open_shift() { 2 } else { 1 }) .sum(); Self { month, year, holidays: vec![], residents: vec![], toxic_pairs: vec![], total_days, total_slots, total_holiday_slots, } } pub fn with_holidays(mut self, holidays: Vec) -> Self { self.holidays = holidays; self.total_holiday_slots = self.total_holiday_slots(); self } pub fn with_residents(mut self, residents: Vec) -> Self { self.residents = residents; self } pub fn add(&mut self, resident: Resident) { self.residents.push(resident); } pub fn with_toxic_pairs(mut self, toxic_pairs: Vec) -> Self { self.toxic_pairs = toxic_pairs; self } pub fn update_month(&mut self, month: u8) { self.month = Month::try_from(month).unwrap(); self.total_days = self.month.num_days(self.year).unwrap(); self.total_slots = (1..=self.total_days) .map(|d| if Day(d).is_open_shift() { 2 } else { 1 }) .sum(); self.total_holiday_slots = self.total_holiday_slots() } fn total_holiday_slots(&self) -> u8 { (1..=self.total_days) .filter(|&d| self.is_holiday_or_weekend(Day(d))) .map(|d| if Day(d).is_open_shift() { 2 } else { 1 }) .sum() } pub fn is_holiday_or_weekend(&self, day: Day) -> bool { let month = self.month.number_from_month(); day.is_weekend(month, self.year) || self.holidays.contains(&(day.0)) } pub fn is_holiday_or_weekend_slot(&self, slot: Slot) -> bool { self.is_holiday_or_weekend(slot.day) } pub fn get_initial_supply(&self) -> HashMap { let mut supply = HashMap::new(); let total_days = self.total_days; for d in 1..=total_days { if Day(d).is_open_shift() { *supply.entry(ShiftType::OpenFirst).or_insert(0) += 1; *supply.entry(ShiftType::OpenSecond).or_insert(0) += 1; } else { *supply.entry(ShiftType::Closed).or_insert(0) += 1; } } supply } pub fn flexibility_map(&self) -> HashMap { let mut map = HashMap::new(); for r in &self.residents { map.insert(r.id, r.allowed_types.len() as u8); } map } } impl Default for UserConfig { fn default() -> Self { let month = Month::try_from(MONTH).unwrap(); let total_days = month.num_days(YEAR).unwrap(); let total_slots = (1..=total_days) .map(|d| if Day(d).is_open_shift() { 2 } else { 1 }) .sum(); let total_holiday_slots = (1..=total_days) .filter(|&d| Day(d).is_weekend(month.number_from_month(), YEAR)) .map(|d| if Day(d).is_open_shift() { 2 } else { 1 }) .sum(); Self { month, year: YEAR, holidays: vec![], residents: vec![], toxic_pairs: vec![], total_days, total_slots, total_holiday_slots, } } } impl TryFrom for UserConfig { type Error = anyhow::Error; fn try_from(value: UserConfigDTO) -> Result { let month = Month::try_from(value.month)?; let total_days = month.num_days(value.year).context("Failed to parse")?; let total_slots = (1..=total_days) .map(|d| if Day(d).is_open_shift() { 2 } else { 1 }) .sum(); let total_holiday_slots = (1..=total_days) .filter(|&d| { Day(d).is_weekend(month.number_from_month(), value.year) || value.holidays.contains(&d) }) .map(|d| if Day(d).is_open_shift() { 2 } else { 1 }) .sum(); Ok(Self { month, year: value.year, holidays: value.holidays, residents: value.residents.into_iter().map(Resident::from).collect(), toxic_pairs: value .toxic_pairs .into_iter() .map(|p| ToxicPair::new(p.0, p.1)) .collect(), total_days, total_slots, total_holiday_slots, }) } } #[cfg(test)] mod tests { use crate::{ config::{ToxicPair, UserConfig}, fixtures::complex_config, schedule::ShiftType, slot::{Day, ShiftPosition, Slot}, }; use rstest::rstest; #[rstest] fn test_get_initial_supply(complex_config: UserConfig) { let supply = complex_config.get_initial_supply(); assert_eq!(15, *supply.get(&ShiftType::OpenFirst).unwrap()); assert_eq!(15, *supply.get(&ShiftType::OpenSecond).unwrap()); assert_eq!(15, *supply.get(&ShiftType::Closed).unwrap()); } #[rstest] fn test_is_holiday_or_weekend(complex_config: UserConfig) { assert!(!complex_config.is_holiday_or_weekend(Day(1))); assert!(complex_config.is_holiday_or_weekend(Day(4))); assert!(complex_config.is_holiday_or_weekend(Day(5))); assert!(complex_config.is_holiday_or_weekend(Day(10))); } #[rstest] fn test_is_holiday_or_weekend_slot(complex_config: UserConfig) { let weekday = Slot::new(Day(1), ShiftPosition::First); let sat = Slot::new(Day(4), ShiftPosition::First); let sun = Slot::new(Day(5), ShiftPosition::First); let manual_holiday = Slot::new(Day(10), ShiftPosition::First); assert!(!complex_config.is_holiday_or_weekend_slot(weekday)); assert!(complex_config.is_holiday_or_weekend_slot(sat)); assert!(complex_config.is_holiday_or_weekend_slot(sun)); assert!(complex_config.is_holiday_or_weekend_slot(manual_holiday)); } #[rstest] fn test_toxic_pair_matches() { let pair_1 = ToxicPair::new(1, 2); let pair_1_r = ToxicPair::new(2, 1); let pair_2 = ToxicPair::new(3, 1); assert!(pair_1.matches(&pair_1_r)); assert!(pair_1_r.matches(&pair_1)); assert!(!pair_1.matches(&pair_2)); assert!(!pair_2.matches(&pair_1)); assert!(!pair_1_r.matches(&pair_2)); assert!(!pair_2.matches(&pair_1_r)); } }