use crate::models::{
CapabilityExt, CapabilityType, NodeConfig, NodeConfigExt, NodeState, NodeStateExt, Operator,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum CellRole {
Leader,
Sensor,
Compute,
Relay,
Strike,
Support,
Follower,
}
impl CellRole {
pub fn assignable_roles() -> Vec<Self> {
vec![
Self::Sensor,
Self::Compute,
Self::Relay,
Self::Strike,
Self::Support,
Self::Follower,
]
}
pub fn description(&self) -> &'static str {
match self {
Self::Leader => "Cell leader - coordinates operations and makes tactical decisions",
Self::Sensor => "Sensor/scout - provides long-range detection and reconnaissance",
Self::Compute => "Compute node - processes sensor data and runs analysis algorithms",
Self::Relay => "Communications relay - extends network range and connectivity",
Self::Strike => "Strike platform - engages targets with weapons systems",
Self::Support => {
"Support platform - provides logistics, medical, or maintenance support"
}
Self::Follower => "General squad member - performs assigned tasks",
}
}
pub fn required_capabilities(&self) -> Vec<CapabilityType> {
match self {
Self::Leader => vec![CapabilityType::Communication],
Self::Sensor => vec![CapabilityType::Sensor],
Self::Compute => vec![CapabilityType::Compute],
Self::Relay => vec![CapabilityType::Communication],
Self::Strike => vec![CapabilityType::Payload],
Self::Support => vec![],
Self::Follower => vec![],
}
}
pub fn preferred_capabilities(&self) -> Vec<CapabilityType> {
match self {
Self::Leader => vec![CapabilityType::Compute, CapabilityType::Sensor],
Self::Sensor => vec![CapabilityType::Communication],
Self::Compute => vec![CapabilityType::Communication],
Self::Relay => vec![CapabilityType::Sensor],
Self::Strike => vec![CapabilityType::Sensor, CapabilityType::Compute],
Self::Support => vec![CapabilityType::Mobility],
Self::Follower => vec![],
}
}
pub fn relevant_mos(&self) -> Vec<&'static str> {
match self {
Self::Leader => vec!["11B", "11C", "19D"], Self::Sensor => vec!["19D", "35M", "35N"], Self::Compute => vec!["35F", "35N", "17C"], Self::Relay => vec!["25U", "25B", "25Q"], Self::Strike => vec!["11B", "11C", "19K"], Self::Support => vec!["68W", "88M", "91B"], Self::Follower => vec![],
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoleAssignment {
pub platform_id: String,
pub role: CellRole,
pub score: f64,
pub is_primary_choice: bool,
}
impl RoleAssignment {
pub fn new(platform_id: String, role: CellRole, score: f64, is_primary_choice: bool) -> Self {
Self {
platform_id,
role,
score,
is_primary_choice,
}
}
}
pub struct RoleScorer;
impl RoleScorer {
pub fn score_platform_for_role(
config: &NodeConfig,
state: &NodeState,
role: CellRole,
) -> Option<f64> {
let operator = config.get_primary_operator();
let mut score = 0.0;
let mut weight_sum = 0.0;
for required_cap_type in role.required_capabilities() {
let has_required = config
.capabilities
.iter()
.any(|c| c.get_capability_type() == required_cap_type);
if !has_required {
return None; }
}
let required_score = Self::score_required_capabilities(config, &role);
score += required_score * 0.3;
weight_sum += 0.3;
let preferred_score = Self::score_preferred_capabilities(config, &role);
score += preferred_score * 0.2;
weight_sum += 0.2;
if let Some(op) = operator {
let mos_score = Self::score_operator_mos(op, &role);
score += mos_score * 0.3;
weight_sum += 0.3;
}
let health_score = Self::score_platform_health(state);
score += health_score * 0.2;
weight_sum += 0.2;
if weight_sum < 1.0 {
score /= weight_sum;
}
Some(score.clamp(0.0, 1.0))
}
fn score_required_capabilities(config: &NodeConfig, role: &CellRole) -> f64 {
let required = role.required_capabilities();
if required.is_empty() {
return 1.0;
}
let mut total_score = 0.0;
for req_type in &required {
let best_capability = config
.capabilities
.iter()
.filter(|c| c.get_capability_type() == *req_type)
.max_by(|a, b| {
a.confidence
.partial_cmp(&b.confidence)
.unwrap_or(std::cmp::Ordering::Equal)
});
if let Some(cap) = best_capability {
total_score += cap.confidence as f64;
}
}
total_score / required.len() as f64
}
fn score_preferred_capabilities(config: &NodeConfig, role: &CellRole) -> f64 {
let preferred = role.preferred_capabilities();
if preferred.is_empty() {
return 1.0;
}
let mut total_score = 0.0;
let mut count = 0;
for pref_type in preferred {
if let Some(best_cap) = config
.capabilities
.iter()
.filter(|c| c.get_capability_type() == pref_type)
.max_by(|a, b| {
a.confidence
.partial_cmp(&b.confidence)
.unwrap_or(std::cmp::Ordering::Equal)
})
{
total_score += best_cap.confidence as f64;
count += 1;
}
}
if count > 0 {
total_score / count as f64
} else {
0.5 }
}
fn score_operator_mos(operator: &Operator, role: &CellRole) -> f64 {
let relevant_mos = role.relevant_mos();
if relevant_mos.is_empty() {
return 0.5; }
if relevant_mos.contains(&operator.mos.as_str()) {
0.9 } else {
0.3 }
}
fn score_platform_health(state: &NodeState) -> f64 {
match state.get_health() {
crate::models::HealthStatus::Nominal => 1.0,
crate::models::HealthStatus::Degraded => 0.6,
crate::models::HealthStatus::Critical => 0.3,
crate::models::HealthStatus::Failed => 0.0,
crate::models::HealthStatus::Unspecified => 0.5,
}
}
pub fn score_all_roles(config: &NodeConfig, state: &NodeState) -> HashMap<CellRole, f64> {
let mut scores = HashMap::new();
for role in CellRole::assignable_roles() {
if let Some(score) = Self::score_platform_for_role(config, state, role) {
scores.insert(role, score);
}
}
scores
}
pub fn best_role_for_platform(
config: &NodeConfig,
state: &NodeState,
) -> Option<(CellRole, f64)> {
Self::score_all_roles(config, state)
.into_iter()
.max_by(|(_, score_a), (_, score_b)| {
score_a
.partial_cmp(score_b)
.unwrap_or(std::cmp::Ordering::Equal)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::models::{
AuthorityLevel, BindingType, Capability, HumanMachinePair, HumanMachinePairExt, NodeConfig,
NodeConfigExt, NodeStateExt, OperatorExt, OperatorRank,
};
fn create_test_platform_with_capabilities(caps: Vec<Capability>) -> (NodeConfig, NodeState) {
let mut config = NodeConfig::new("test_platform".to_string());
for cap in caps {
config.add_capability(cap);
}
let state = NodeState::new((0.0, 0.0, 0.0));
(config, state)
}
fn create_test_operator(mos: &str, rank: OperatorRank) -> Operator {
Operator::new(
"op_1".to_string(),
"Test Operator".to_string(),
rank,
AuthorityLevel::Commander,
mos.to_string(),
)
}
#[test]
fn test_role_required_capabilities() {
assert_eq!(
CellRole::Sensor.required_capabilities(),
vec![CapabilityType::Sensor]
);
assert_eq!(
CellRole::Strike.required_capabilities(),
vec![CapabilityType::Payload]
);
assert!(CellRole::Follower.required_capabilities().is_empty());
}
#[test]
fn test_role_relevant_mos() {
let sensor_mos = CellRole::Sensor.relevant_mos();
assert!(sensor_mos.contains(&"19D"));
let relay_mos = CellRole::Relay.relevant_mos();
assert!(relay_mos.contains(&"25U")); }
#[test]
fn test_score_platform_without_required_capability() {
let (config, state) = create_test_platform_with_capabilities(vec![Capability::new(
"cpu_1".to_string(),
"CPU".to_string(),
CapabilityType::Compute,
0.8,
)]);
let score = RoleScorer::score_platform_for_role(&config, &state, CellRole::Sensor);
assert!(score.is_none());
}
#[test]
fn test_score_platform_with_required_capability() {
let (config, state) = create_test_platform_with_capabilities(vec![Capability::new(
"radar_1".to_string(),
"Radar".to_string(),
CapabilityType::Sensor,
0.9,
)]);
let score = RoleScorer::score_platform_for_role(&config, &state, CellRole::Sensor);
assert!(score.is_some());
assert!(score.unwrap() > 0.5);
}
#[test]
fn test_score_with_operator_mos_match() {
let (mut config, state) = create_test_platform_with_capabilities(vec![Capability::new(
"camera_1".to_string(),
"Camera".to_string(),
CapabilityType::Sensor,
0.8,
)]);
let operator = create_test_operator("19D", OperatorRank::E5); let binding = crate::models::HumanMachinePair::new(
vec![operator],
vec![config.id.clone()],
crate::models::BindingType::OneToOne,
);
config.set_operator_binding(Some(binding));
let score_with_match =
RoleScorer::score_platform_for_role(&config, &state, CellRole::Sensor).unwrap();
let (config_no_op, state_no_op) =
create_test_platform_with_capabilities(vec![Capability::new(
"camera_2".to_string(),
"Camera".to_string(),
CapabilityType::Sensor,
0.8,
)]);
let score_without_operator =
RoleScorer::score_platform_for_role(&config_no_op, &state_no_op, CellRole::Sensor)
.unwrap();
assert!(score_with_match > score_without_operator);
}
#[test]
fn test_score_with_operator_mos_mismatch() {
let (mut config, state) = create_test_platform_with_capabilities(vec![Capability::new(
"camera_1".to_string(),
"Camera".to_string(),
CapabilityType::Sensor,
0.8,
)]);
let operator = create_test_operator("68W", OperatorRank::E4); let binding = crate::models::HumanMachinePair::new(
vec![operator],
vec![config.id.clone()],
crate::models::BindingType::OneToOne,
);
config.set_operator_binding(Some(binding));
let score_with_mismatch =
RoleScorer::score_platform_for_role(&config, &state, CellRole::Sensor).unwrap();
assert!(score_with_mismatch > 0.0);
assert!(score_with_mismatch < 1.0);
}
#[test]
fn test_score_all_roles() {
let (config, state) = create_test_platform_with_capabilities(vec![
Capability::new(
"camera_1".to_string(),
"Camera".to_string(),
CapabilityType::Sensor,
0.9,
),
Capability::new(
"radio_1".to_string(),
"Radio".to_string(),
CapabilityType::Communication,
0.7,
),
]);
let scores = RoleScorer::score_all_roles(&config, &state);
assert!(scores.contains_key(&CellRole::Sensor));
assert!(scores.contains_key(&CellRole::Relay));
assert!(scores.contains_key(&CellRole::Follower));
assert!(!scores.contains_key(&CellRole::Strike));
assert!(!scores.contains_key(&CellRole::Compute));
}
#[test]
fn test_best_role_for_platform() {
let mut config = NodeConfig::new("test_platform".to_string());
config.add_capability(Capability::new(
"radar_1".to_string(),
"Radar".to_string(),
CapabilityType::Sensor,
0.95,
));
config.add_capability(Capability::new(
"radio_1".to_string(),
"Radio".to_string(),
CapabilityType::Communication,
0.5,
));
let operator = create_test_operator("19D", OperatorRank::E4); let platform_id = config.id.clone();
config.operator_binding = Some(HumanMachinePair::new(
vec![operator],
vec![platform_id],
BindingType::OneToOne,
));
let state = NodeState::new((0.0, 0.0, 0.0));
let (best_role, score) = RoleScorer::best_role_for_platform(&config, &state).unwrap();
assert_eq!(best_role, CellRole::Sensor);
assert!(score > 0.5);
}
#[test]
fn test_role_assignment_creation() {
let assignment = RoleAssignment::new("node_1".to_string(), CellRole::Sensor, 0.85, true);
assert_eq!(assignment.platform_id, "node_1");
assert_eq!(assignment.role, CellRole::Sensor);
assert_eq!(assignment.score, 0.85);
assert!(assignment.is_primary_choice);
}
#[test]
fn test_assignable_roles() {
let roles = CellRole::assignable_roles();
assert!(!roles.contains(&CellRole::Leader));
assert!(roles.contains(&CellRole::Sensor));
assert!(roles.contains(&CellRole::Compute));
assert!(roles.contains(&CellRole::Relay));
assert!(roles.contains(&CellRole::Strike));
assert!(roles.contains(&CellRole::Support));
assert!(roles.contains(&CellRole::Follower));
}
#[test]
fn test_degraded_platform_role_scoring() {
let (config_nominal, state_nominal) =
create_test_platform_with_capabilities(vec![Capability::new(
"sensor_1".to_string(),
"Sensor".to_string(),
CapabilityType::Sensor,
0.9,
)]);
let (config_degraded, mut state_degraded) =
create_test_platform_with_capabilities(vec![Capability::new(
"sensor_2".to_string(),
"Sensor".to_string(),
CapabilityType::Sensor,
0.9,
)]);
state_degraded.update_health(crate::models::HealthStatus::Degraded);
let score_nominal =
RoleScorer::score_platform_for_role(&config_nominal, &state_nominal, CellRole::Sensor)
.unwrap();
let score_degraded = RoleScorer::score_platform_for_role(
&config_degraded,
&state_degraded,
CellRole::Sensor,
)
.unwrap();
assert!(score_degraded < score_nominal);
assert!(score_degraded > 0.4);
}
#[test]
fn test_critical_platform_role_scoring() {
let (config, mut state) = create_test_platform_with_capabilities(vec![Capability::new(
"sensor_1".to_string(),
"Sensor".to_string(),
CapabilityType::Sensor,
0.9,
)]);
state.update_health(crate::models::HealthStatus::Critical);
let score = RoleScorer::score_platform_for_role(&config, &state, CellRole::Sensor).unwrap();
assert!(score > 0.0);
assert!(score < 0.7); }
#[test]
fn test_failed_platform_role_scoring() {
let (config, mut state) = create_test_platform_with_capabilities(vec![Capability::new(
"sensor_1".to_string(),
"Sensor".to_string(),
CapabilityType::Sensor,
0.9,
)]);
state.update_health(crate::models::HealthStatus::Failed);
let score = RoleScorer::score_platform_for_role(&config, &state, CellRole::Sensor).unwrap();
assert!(score > 0.2);
assert!(score < 0.6); }
}