rota/src-tauri/src/schedule.rs

use serde::{ser::SerializeMap, Deserialize, Serialize};
use std::collections::HashMap;

use crate::{
    bounds::WorkloadBounds,
    config::UserConfig,
    resident::Resident,
    slot::{weekday_to_greek, Day, ShiftPosition, Slot},
};

use serde::Serializer;

/// each slot has one resident
/// a day can span between 1 or 2 slots depending on if it is open(odd) or closed(even)
#[derive(Deserialize, Debug, Clone)]
pub struct MonthlySchedule(pub HashMap<Slot, String>);

impl MonthlySchedule {
    pub fn new() -> Self {
        Self(HashMap::new())
    }

    pub fn prefill(&mut self, config: &UserConfig) {
        for r in &config.residents {
            for s in &r.manual_shifts {
                self.insert(*s, &r.id);
            }
        }
    }

    pub fn get_resident_id(&self, slot: &Slot) -> Option<&String> {
        self.0.get(slot)
    }

    pub fn current_workload(&self, resident_id: &str) -> usize {
        self.0
            .values()
            .filter(|res_id| res_id == &resident_id)
            .count()
    }

    pub fn current_holiday_workload(&self, resident: &Resident, config: &UserConfig) -> usize {
        self.0
            .iter()
            .filter(|(slot, res_id)| {
                res_id == &&resident.id && config.is_holiday_or_weekend_slot(slot.day.0)
            })
            .count()
    }

    pub fn count_shifts(&self, resident_id: &str, shift_type: Option<ShiftType>) -> usize {
        self.0
            .iter()
            .filter(|(slot, id)| {
                if *id != resident_id {
                    return false;
                }

                match &shift_type {
                    None => true,
                    Some(target) => {
                        let actual_type = if slot.is_open_shift() {
                            match slot.position {
                                ShiftPosition::First => ShiftType::OpenFirst,
                                ShiftPosition::Second => ShiftType::OpenSecond,
                            }
                        } else {
                            ShiftType::Closed
                        };
                        actual_type == *target
                    }
                }
            })
            .count()
    }

    pub fn is_per_shift_threshold_met(&self, config: &UserConfig, bounds: &WorkloadBounds) -> bool {
        for res in &config.residents {
            for stype in [
                ShiftType::OpenFirst,
                ShiftType::OpenSecond,
                ShiftType::Closed,
            ] {
                let count = self.count_shifts(&res.id, Some(stype.clone()));
                if let Some(&threshold) = bounds.min_by_shift_type.get(&(res.id.clone(), stype)) {
                    if count < threshold as usize {
                        return false;
                    }
                }
            }
        }
        true
    }

    pub fn insert(&mut self, slot: Slot, resident_id: &str) {
        self.0.insert(slot, resident_id.to_string());
    }

    pub fn remove(&mut self, slot: Slot) {
        self.0.remove(&slot);
    }

    pub fn is_slot_manually_assigned(&self, slot: &Slot) -> bool {
        self.0.contains_key(slot)
    }

    /// if any restriction is violated => we return true (leading to pruning in the backtracking algorithm)
    /// 1) no same person in consecutive days
    /// 2) avoid input toxic pairs
    /// 3) apply fairness on total shifts split residents also take into account reduced (-1) workload
    ///
    /// @slot points to an occupied slot
    /// @config info manually set on the GUI by the user
    pub fn restrictions_violated(
        &self,
        slot: &Slot,
        config: &UserConfig,
        bounds: &WorkloadBounds,
    ) -> bool {
        self.same_resident_in_consecutive_days(slot)
            || self.has_toxic_pair(slot, config)
            || self.is_workload_unbalanced(slot, config, bounds)
            || self.is_holiday_workload_imbalanced(slot, config, bounds)
            || self.is_shift_type_distribution_unfair(slot, bounds)
    }

    /// same_resident_in_consecutive_days
    pub fn same_resident_in_consecutive_days(&self, slot: &Slot) -> bool {
        if slot.day == Day(1) {
            return false;
        }

        let previous_slots = if slot.is_open_second() {
            vec![slot.previous().previous()]
        } else if slot.is_open_first() {
            vec![slot.previous()]
        } else {
            // if current shift is closed, we need to check both residents in the previous day
            vec![slot.previous(), slot.previous().previous()]
        };

        previous_slots
            .iter()
            .any(|s| self.get_resident_id(s) == self.get_resident_id(slot))
    }

    /// has_toxic_pair
    pub fn has_toxic_pair(&self, slot: &Slot, config: &UserConfig) -> bool {
        // can only have caused a toxic pair violation if we just added a 2nd resident in an open shift
        if !slot.is_open_second() {
            return false;
        }

        let first_id = self.get_resident_id(&slot.previous());
        let second_id = self.get_resident_id(slot);

        if let (Some(f), Some(s)) = (first_id, second_id) {
            return config
                .toxic_pairs
                .iter()
                .any(|(r1, r2)| (r1 == f && r2 == s) || (r1 == s && r2 == f));
        }

        false
    }

    /// is_workload_unbalanced
    pub fn is_workload_unbalanced(
        &self,
        slot: &Slot,
        config: &UserConfig,
        bounds: &WorkloadBounds,
    ) -> bool {
        let res_id = match self.get_resident_id(slot) {
            Some(id) => id,
            None => return false,
        };

        if let Some(resident) = config.residents.iter().find(|r| &r.id == res_id) {
            let current_workload = self.current_workload(&resident.id);

            if let Some(&limit) = bounds.max_workloads.get(res_id) {
                let mut workload_limit = limit;
                if resident.reduced_load {
                    workload_limit -= 1;
                }

                if current_workload > workload_limit as usize {
                    return true;
                }
            }
        }

        false
    }

    /// is_holiday_workload_imbalanced
    pub fn is_holiday_workload_imbalanced(
        &self,
        slot: &Slot,
        config: &UserConfig,
        bounds: &WorkloadBounds,
    ) -> bool {
        if !config.is_holiday_or_weekend_slot(slot.day.0) {
            return false;
        }

        let res_id = match self.get_resident_id(slot) {
            Some(id) => id,
            None => return false,
        };

        if let Some(resident) = config.residents.iter().find(|r| &r.id == res_id) {
            let current_holiday_workload = self.current_holiday_workload(resident, config);

            if let Some(&holiday_limit) = bounds.max_holiday_shifts.get(res_id) {
                if current_holiday_workload > holiday_limit as usize {
                    return true;
                }
            }
        }

        false
    }

    /// is_shift_type_distribution_unfair
    pub fn is_shift_type_distribution_unfair(&self, slot: &Slot, bounds: &WorkloadBounds) -> bool {
        let resident_id = match self.get_resident_id(slot) {
            Some(id) => id,
            None => return false,
        };

        let current_shift_type = if slot.is_open_shift() {
            match slot.position {
                ShiftPosition::First => ShiftType::OpenFirst,
                ShiftPosition::Second => ShiftType::OpenSecond,
            }
        } else {
            ShiftType::Closed
        };

        let current_count = self.count_shifts(resident_id, Some(current_shift_type.clone()));

        if let Some(&limit) = bounds
            .max_by_shift_type
            .get(&(resident_id.clone(), current_shift_type.clone()))
        {
            return current_count > limit as usize;
        }

        false
    }

    pub fn pretty_print(&self, config: &UserConfig) -> String {
        let mut sorted: Vec<_> = self.0.iter().collect();
        sorted.sort_by_key(|(slot, _)| (slot.day, slot.position));

        let mut output = String::from("Μηνιαίο Πρόγραμμα Εφημεριών\n");
        for (slot, res_id) in sorted {
            let res_name = config
                .residents
                .iter()
                .find(|r| &r.id == res_id)
                .map(|r| r.name.as_str())
                .unwrap();

            output.push_str(&format!(
                "Ημέρα {:2} - {:9} - {:11}: {},\n",
                slot.day.0,
                weekday_to_greek(
                    slot.day
                        .weekday(config.month.number_from_month(), config.year)
                ),
                slot.shift_type_str(),
                res_name
            ));
        }
        output
    }

    pub fn report(&self, config: &UserConfig) -> String {
        let mut output = String::new();
        output.push_str("\n--- Αναφορά ---\n");
        // Using standard widths for Greek characters and alignment
        output.push_str(&format!(
            "{:<15} | {:<6} | {:<10} | {:<10} | {:<7} | {:<10}\n",
            "Ειδικευόμενος", "Σύνολο", "Ανοιχτή(1)", "Ανοιχτή(2)", "Κλειστή", "ΣΚ/Αργίες"
        ));
        output.push_str("-".repeat(75).as_str());
        output.push('\n');

        let mut residents: Vec<_> = config.residents.iter().collect();
        residents.sort_by_key(|r| &r.name);

        for res in residents {
            let total = self.current_workload(&res.id);
            let o1 = self.count_shifts(&res.id, Some(ShiftType::OpenFirst));
            let o2 = self.count_shifts(&res.id, Some(ShiftType::OpenSecond));
            let cl = self.count_shifts(&res.id, Some(ShiftType::Closed));
            let sun = self.current_holiday_workload(res, config);

            output.push_str(&format!(
                "{:<15} | {:<6} | {:<10} | {:<10} | {:<7} | {:<10}\n",
                res.name, total, o1, o2, cl, sun
            ));
        }
        output.push_str("-".repeat(75).as_str());
        output.push('\n');
        output
    }
}

impl Serialize for MonthlySchedule {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut map = serializer.serialize_map(Some(self.0.len()))?;
        for (slot, name) in &self.0 {
            let pos_str = match slot.position {
                ShiftPosition::First => "First",
                ShiftPosition::Second => "Second",
            };
            let key = format!("{}-{}", slot.day.0, pos_str);
            map.serialize_entry(&key, name)?;
        }
        map.end()
    }
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)]
pub enum ShiftType {
    Closed,
    OpenFirst,
    OpenSecond,
}

#[cfg(test)]
mod tests {
    use rstest::{fixture, rstest};

    use crate::{
        bounds::WorkloadBounds,
        config::UserConfig,
        resident::Resident,
        schedule::{Day, MonthlySchedule, Slot},
        slot::ShiftPosition,
    };

    #[fixture]
    fn schedule() -> MonthlySchedule {
        MonthlySchedule::new()
    }

    #[fixture]
    fn resident() -> Resident {
        Resident::new("1", "Stefanos")
    }

    #[fixture]
    fn toxic_config() -> UserConfig {
        UserConfig::default()
            .with_residents(vec![
                Resident::new("1", "Stefanos"),
                Resident::new("2", "Iordanis"),
            ])
            .with_toxic_pairs(vec![(("1".to_string(), "2".to_string()))])
    }

    #[fixture]
    fn config() -> UserConfig {
        UserConfig::default().with_residents(vec![
            Resident::new("1", "Stefanos"),
            Resident::new("2", "Iordanis"),
        ])
    }

    #[rstest]
    fn test_insert_resident(mut schedule: MonthlySchedule, resident: Resident) {
        let slot_1 = Slot::new(Day(1), ShiftPosition::First);
        let slot_2 = Slot::new(Day(1), ShiftPosition::Second);

        schedule.insert(slot_1, &resident.id);

        assert_eq!(schedule.get_resident_id(&slot_1), Some(&"1".to_string()));
        assert_eq!(schedule.current_workload(&resident.id), 1);
        assert_eq!(schedule.get_resident_id(&slot_2), None);
    }

    #[rstest]
    fn test_remove_resident(mut schedule: MonthlySchedule, resident: Resident) {
        let slot_1 = Slot::new(Day(1), ShiftPosition::First);

        schedule.insert(slot_1, &resident.id);
        assert_eq!(schedule.current_workload(&resident.id), 1);

        schedule.remove(slot_1);
        assert_eq!(schedule.get_resident_id(&slot_1), None);
        assert_eq!(schedule.current_workload(&resident.id), 0);
    }

    #[rstest]
    fn test_same_resident_in_consecutive_days(mut schedule: MonthlySchedule, resident: Resident) {
        let slot_1 = Slot::new(Day(1), ShiftPosition::First);
        let slot_2 = Slot::new(Day(1), ShiftPosition::Second);
        let slot_3 = Slot::new(Day(2), ShiftPosition::First);

        schedule.insert(slot_1, &resident.id);
        schedule.insert(slot_2, &resident.id);
        schedule.insert(slot_3, &resident.id);

        assert!(!schedule.same_resident_in_consecutive_days(&slot_1));
        assert!(!schedule.same_resident_in_consecutive_days(&slot_2));
        assert!(schedule.same_resident_in_consecutive_days(&slot_3));
    }

    #[rstest]
    fn test_has_toxic_pair(mut schedule: MonthlySchedule, toxic_config: UserConfig) {
        let slot_1 = Slot::new(Day(1), ShiftPosition::First);
        let slot_2 = Slot::new(Day(1), ShiftPosition::Second);

        let stefanos = &toxic_config.residents[0];
        let iordanis = &toxic_config.residents[1];

        schedule.insert(slot_1, &stefanos.id);
        schedule.insert(slot_2, &iordanis.id);

        assert!(schedule.has_toxic_pair(&slot_2, &toxic_config))
    }

    #[rstest]
    fn test_is_workload_unbalanced(mut schedule: MonthlySchedule, config: UserConfig) {
        let slot_1 = Slot::new(Day(1), ShiftPosition::First);
        let slot_2 = Slot::new(Day(1), ShiftPosition::Second);
        let slot_3 = Slot::new(Day(2), ShiftPosition::First);

        let stefanos = &config.residents[0];
        let iordanis = &config.residents[1];

        let mut bounds = WorkloadBounds::new();

        bounds.max_workloads.insert("1".to_string(), 1);
        bounds.max_workloads.insert("2".to_string(), 2);

        schedule.insert(slot_1, &stefanos.id);
        assert!(!schedule.is_workload_unbalanced(&slot_1, &config, &bounds));

        schedule.insert(slot_2, &iordanis.id);
        assert!(!schedule.is_workload_unbalanced(&slot_2, &config, &bounds));

        schedule.insert(slot_3, &stefanos.id);
        assert!(schedule.is_workload_unbalanced(&slot_3, &config, &bounds));
    }

    #[rstest]
    fn test_is_holiday_workload_imbalanced(mut schedule: MonthlySchedule, config: UserConfig) {
        let slot_1 = Slot::new(Day(1), ShiftPosition::First);
        let slot_2 = Slot::new(Day(1), ShiftPosition::Second);
        let slot_7 = Slot::new(Day(7), ShiftPosition::First);

        let stefanos = &config.residents[0];
        let iordanis = &config.residents[1];

        let mut bounds = WorkloadBounds::new();

        bounds.max_holiday_shifts.insert("1".to_string(), 1);
        bounds.max_holiday_shifts.insert("2".to_string(), 1);

        schedule.insert(slot_1, &stefanos.id);
        assert!(!schedule.is_holiday_workload_imbalanced(&slot_1, &config, &bounds));

        schedule.insert(slot_2, &iordanis.id);
        assert!(!schedule.is_holiday_workload_imbalanced(&slot_2, &config, &bounds));

        schedule.insert(slot_7, &stefanos.id);
        assert!(schedule.is_holiday_workload_imbalanced(&slot_7, &config, &bounds));
    }
}