use crate::{Goal, GoalStatus, GoalType};
use std::collections::{HashMap, HashSet};
pub type GoalId = String;
pub struct HierarchicalGoalSolver {
goals: HashMap<GoalId, Goal>,
goal_tree: GoalTree,
decomposition_rules: Vec<DecompositionRule>,
active_goals: HashSet<GoalId>,
}
#[derive(Debug, Clone)]
pub struct GoalTree {
root_goals: Vec<GoalId>,
children: HashMap<GoalId, Vec<GoalId>>,
parents: HashMap<GoalId, GoalId>,
}
pub type ConditionFn = Box<dyn Fn(&Goal) -> bool + Send + Sync>;
pub type DecomposeFn = Box<dyn Fn(&Goal) -> Vec<Goal> + Send + Sync>;
pub struct DecompositionRule {
pub name: String,
pub goal_type_filter: Option<GoalTypeFilter>,
pub condition: ConditionFn,
pub decompose: DecomposeFn,
}
#[derive(Clone)]
pub enum GoalTypeFilter {
Achieve,
Maintain,
Maximize,
Minimize,
Avoid,
Perform,
Respond,
Custom,
}
pub struct DecompositionResult {
pub parent_id: GoalId,
pub subgoals: Vec<Goal>,
pub dependencies: Vec<(GoalId, GoalId)>,
}
#[derive(Debug, Clone)]
pub struct GoalConflict {
pub goal1: GoalId,
pub goal2: GoalId,
pub conflict_type: ConflictType,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ConflictType {
ResourceContention,
MutuallyExclusive,
TemporalOverlap,
PriorityConflict,
}
pub enum ConflictResolution {
PrioritizeFirst,
PrioritizeSecond,
Sequential(GoalId, GoalId),
Merge(Goal),
Abandon(GoalId),
}
impl HierarchicalGoalSolver {
pub fn new() -> Self {
Self {
goals: HashMap::new(),
goal_tree: GoalTree::new(),
decomposition_rules: Vec::new(),
active_goals: HashSet::new(),
}
}
pub fn add_goal(&mut self, goal: Goal) -> GoalId {
let id = goal.id.clone();
if goal.parent.is_none() {
self.goal_tree.add_root(id.clone());
}
self.goals.insert(id.clone(), goal);
id
}
pub fn decompose(&mut self, goal_id: &GoalId) -> Result<DecompositionResult, String> {
let goal = self
.goals
.get(goal_id)
.ok_or_else(|| format!("Goal {} not found", goal_id))?
.clone();
let rule = self
.decomposition_rules
.iter()
.find(|r| {
if let Some(ref filter) = r.goal_type_filter {
if !matches_goal_type(&goal.goal_type, filter) {
return false;
}
}
(r.condition)(&goal)
})
.ok_or_else(|| format!("No decomposition rule found for goal {}", goal_id))?;
let mut subgoals = (rule.decompose)(&goal);
for subgoal in &mut subgoals {
subgoal.parent = Some(goal_id.clone());
let subgoal_id = subgoal.id.clone();
self.goal_tree
.add_child(goal_id.clone(), subgoal_id.clone());
self.goals.insert(subgoal_id.clone(), subgoal.clone());
if let Some(parent) = self.goals.get_mut(goal_id) {
parent.add_subgoal(&subgoal.id);
}
}
let mut dependencies = Vec::new();
for i in 1..subgoals.len() {
dependencies.push((subgoals[i].id.clone(), subgoals[i - 1].id.clone()));
}
Ok(DecompositionResult {
parent_id: goal_id.clone(),
subgoals,
dependencies,
})
}
pub fn decompose_all(&mut self) -> Vec<DecompositionResult> {
let mut results = Vec::new();
let goal_ids: Vec<GoalId> = self.goals.keys().cloned().collect();
for goal_id in goal_ids {
if let Some(goal) = self.goals.get(&goal_id) {
if !goal.subgoals.is_empty() {
continue;
}
}
if let Ok(result) = self.decompose(&goal_id) {
results.push(result);
}
}
results
}
pub fn register_rule(&mut self, rule: DecompositionRule) {
self.decomposition_rules.push(rule);
}
pub fn get_executable_goals(&self) -> Vec<&Goal> {
self.goals
.values()
.filter(|g| {
matches!(g.status, GoalStatus::Pending | GoalStatus::Active)
&& g.subgoals.is_empty()
|| g.subgoals.iter().all(|sg_id| {
self.goals
.get(sg_id)
.map(|sg| sg.is_complete())
.unwrap_or(false)
})
})
.collect()
}
pub fn mark_achieved(&mut self, goal_id: &GoalId) -> Vec<GoalId> {
let mut affected = Vec::new();
if let Some(goal) = self.goals.get_mut(goal_id) {
goal.mark_achieved();
affected.push(goal_id.clone());
self.active_goals.remove(goal_id);
if let Some(parent_id) = goal.parent.clone() {
self.propagate_progress(&parent_id);
if let Some(parent) = self.goals.get(&parent_id) {
let all_complete = !parent.subgoals.is_empty()
&& parent.subgoals.iter().all(|sg_id| {
self.goals
.get(sg_id)
.map(|sg| sg.status == GoalStatus::Achieved)
.unwrap_or(false)
});
if all_complete {
let parent_affected = self.mark_achieved(&parent_id);
affected.extend(parent_affected);
}
}
}
}
affected
}
pub fn mark_failed(&mut self, goal_id: &GoalId, _reason: String) {
if let Some(goal) = self.goals.get_mut(goal_id) {
goal.fail();
self.active_goals.remove(goal_id);
if let Some(parent_id) = goal.parent.clone() {
self.propagate_progress(&parent_id);
}
}
}
pub fn detect_conflicts(&self) -> Vec<GoalConflict> {
let mut conflicts = Vec::new();
let active: Vec<&Goal> = self.goals.values().filter(|g| g.is_active()).collect();
for i in 0..active.len() {
for j in (i + 1)..active.len() {
let g1 = active[i];
let g2 = active[j];
if g1.priority == g2.priority && g1.priority >= crate::types::Priority::High {
conflicts.push(GoalConflict {
goal1: g1.id.clone(),
goal2: g2.id.clone(),
conflict_type: ConflictType::PriorityConflict,
});
}
if let (Some(d1), Some(d2)) = (g1.deadline, g2.deadline) {
if d1 == d2 {
conflicts.push(GoalConflict {
goal1: g1.id.clone(),
goal2: g2.id.clone(),
conflict_type: ConflictType::TemporalOverlap,
});
}
}
if is_mutually_exclusive(&g1.goal_type, &g2.goal_type) {
conflicts.push(GoalConflict {
goal1: g1.id.clone(),
goal2: g2.id.clone(),
conflict_type: ConflictType::MutuallyExclusive,
});
}
}
}
conflicts
}
pub fn resolve_conflict(&mut self, conflict: &GoalConflict, resolution: ConflictResolution) {
match resolution {
ConflictResolution::PrioritizeFirst => {
if let Some(goal) = self.goals.get_mut(&conflict.goal2) {
goal.status = GoalStatus::OnHold;
}
}
ConflictResolution::PrioritizeSecond => {
if let Some(goal) = self.goals.get_mut(&conflict.goal1) {
goal.status = GoalStatus::OnHold;
}
}
ConflictResolution::Sequential(_first, second) => {
if let Some(goal) = self.goals.get_mut(&second) {
goal.status = GoalStatus::OnHold;
}
}
ConflictResolution::Merge(_merged_goal) => {
if let Some(goal) = self.goals.get_mut(&conflict.goal1) {
goal.cancel();
}
if let Some(goal) = self.goals.get_mut(&conflict.goal2) {
goal.cancel();
}
}
ConflictResolution::Abandon(goal_id) => {
if let Some(goal) = self.goals.get_mut(&goal_id) {
goal.cancel();
}
}
}
}
pub fn get_progress(&self, goal_id: &GoalId) -> f32 {
if let Some(goal) = self.goals.get(goal_id) {
if goal.subgoals.is_empty() {
goal.progress
} else {
let total: f32 = goal
.subgoals
.iter()
.filter_map(|sg_id| self.goals.get(sg_id))
.map(|sg| sg.progress)
.sum();
if goal.subgoals.is_empty() {
0.0
} else {
total / goal.subgoals.len() as f32
}
}
} else {
0.0
}
}
fn propagate_progress(&mut self, goal_id: &GoalId) {
let progress = self.get_progress(goal_id);
if let Some(goal) = self.goals.get_mut(goal_id) {
goal.set_progress(progress);
}
}
pub fn get_subgoals(&self, goal_id: &GoalId) -> Vec<&Goal> {
if let Some(goal) = self.goals.get(goal_id) {
goal.subgoals
.iter()
.filter_map(|sg_id| self.goals.get(sg_id))
.collect()
} else {
Vec::new()
}
}
pub fn get_dependency_graph(&self) -> &GoalTree {
&self.goal_tree
}
pub fn get_goal(&self, goal_id: &GoalId) -> Option<&Goal> {
self.goals.get(goal_id)
}
pub fn get_goal_mut(&mut self, goal_id: &GoalId) -> Option<&mut Goal> {
self.goals.get_mut(goal_id)
}
pub fn activate_goal(&mut self, goal_id: &GoalId) {
if let Some(goal) = self.goals.get_mut(goal_id) {
goal.activate();
self.active_goals.insert(goal_id.clone());
}
}
}
impl Default for HierarchicalGoalSolver {
fn default() -> Self {
Self::new()
}
}
impl GoalTree {
pub fn new() -> Self {
Self {
root_goals: Vec::new(),
children: HashMap::new(),
parents: HashMap::new(),
}
}
pub fn add_root(&mut self, goal_id: GoalId) {
if !self.root_goals.contains(&goal_id) {
self.root_goals.push(goal_id);
}
}
pub fn add_child(&mut self, parent: GoalId, child: GoalId) {
self.children
.entry(parent.clone())
.or_default()
.push(child.clone());
self.parents.insert(child, parent);
}
pub fn get_children(&self, goal_id: &GoalId) -> Option<&Vec<GoalId>> {
self.children.get(goal_id)
}
pub fn get_parent(&self, goal_id: &GoalId) -> Option<&GoalId> {
self.parents.get(goal_id)
}
pub fn root_goals(&self) -> &[GoalId] {
&self.root_goals
}
pub fn topological_sort(&self) -> Vec<GoalId> {
let mut result = Vec::new();
let mut visited = HashSet::new();
let mut temp_mark = HashSet::new();
for root in &self.root_goals {
self.topological_visit(root, &mut visited, &mut temp_mark, &mut result);
}
result
}
fn topological_visit(
&self,
node: &GoalId,
visited: &mut HashSet<GoalId>,
temp_mark: &mut HashSet<GoalId>,
result: &mut Vec<GoalId>,
) {
if visited.contains(node) {
return;
}
if temp_mark.contains(node) {
return;
}
temp_mark.insert(node.clone());
if let Some(children) = self.children.get(node) {
for child in children {
self.topological_visit(child, visited, temp_mark, result);
}
}
temp_mark.remove(node);
visited.insert(node.clone());
result.push(node.clone());
}
}
impl Default for GoalTree {
fn default() -> Self {
Self::new()
}
}
fn matches_goal_type(goal_type: &GoalType, filter: &GoalTypeFilter) -> bool {
matches!(
(goal_type, filter),
(GoalType::Achieve { .. }, GoalTypeFilter::Achieve)
| (GoalType::Maintain { .. }, GoalTypeFilter::Maintain)
| (GoalType::Maximize { .. }, GoalTypeFilter::Maximize)
| (GoalType::Minimize { .. }, GoalTypeFilter::Minimize)
| (GoalType::Avoid { .. }, GoalTypeFilter::Avoid)
| (GoalType::Perform { .. }, GoalTypeFilter::Perform)
| (GoalType::Respond { .. }, GoalTypeFilter::Respond)
| (GoalType::Custom { .. }, GoalTypeFilter::Custom)
)
}
fn is_mutually_exclusive(g1: &GoalType, g2: &GoalType) -> bool {
match (g1, g2) {
(GoalType::Maximize { target: t1 }, GoalType::Minimize { target: t2 }) => t1 == t2,
(GoalType::Minimize { target: t1 }, GoalType::Maximize { target: t2 }) => t1 == t2,
_ => false,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::Timestamp;
#[test]
fn test_goal_solver_creation() {
let solver = HierarchicalGoalSolver::new();
assert!(solver.goals.is_empty());
assert!(solver.decomposition_rules.is_empty());
}
#[test]
fn test_add_goal() {
let mut solver = HierarchicalGoalSolver::new();
let goal = Goal::maintain("temperature", 20.0..25.0);
let id = solver.add_goal(goal);
assert!(solver.get_goal(&id).is_some());
}
#[test]
fn test_goal_tree() {
let mut tree = GoalTree::new();
tree.add_root("goal1".to_string());
tree.add_child("goal1".to_string(), "goal2".to_string());
tree.add_child("goal1".to_string(), "goal3".to_string());
assert_eq!(tree.root_goals.len(), 1);
assert_eq!(tree.get_children(&"goal1".to_string()).unwrap().len(), 2);
assert_eq!(tree.get_parent(&"goal2".to_string()).unwrap(), "goal1");
}
#[test]
fn test_mark_achieved() {
let mut solver = HierarchicalGoalSolver::new();
let goal = Goal::maintain("temp", 20.0..25.0);
let id = solver.add_goal(goal);
solver.activate_goal(&id);
let affected = solver.mark_achieved(&id);
assert!(!affected.is_empty());
assert_eq!(solver.get_goal(&id).unwrap().status, GoalStatus::Achieved);
}
#[test]
fn test_conflict_detection() {
let mut solver = HierarchicalGoalSolver::new();
let mut goal1 = Goal::maximize("efficiency");
goal1.activate();
let id1 = solver.add_goal(goal1);
solver.activate_goal(&id1);
let mut goal2 = Goal::minimize("efficiency");
goal2.activate();
let id2 = solver.add_goal(goal2);
solver.activate_goal(&id2);
let conflicts = solver.detect_conflicts();
assert!(!conflicts.is_empty());
assert_eq!(conflicts[0].conflict_type, ConflictType::MutuallyExclusive);
}
}