#![forbid(unsafe_code)]
use async_trait::async_trait;
use std::sync::{Arc, Mutex};
use serde_json::{Value, json};
use wm_core::{Context, EffectRow, Gana, Resource, Tool, ToolStats};
use wm_selfmodel::{MetricKind, SelfModel};
use wm_simulation::{
CalibrationStore, CounterfactualEstimator, Distribution, ForecastMethod, Forecaster, McConfig,
MonteCarloSimulator,
};
pub struct SimMcTool {
stats: ToolStats,
effects: EffectRow,
}
impl SimMcTool {
#[must_use]
pub fn new() -> Self {
Self {
stats: ToolStats::default(),
effects: EffectRow::read_only(vec![Resource::Galaxy("simulation".into())]),
}
}
}
impl Default for SimMcTool {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl Tool for SimMcTool {
fn name(&self) -> &str {
"sim.mc"
}
fn gana(&self) -> Gana {
Gana::Mound
}
fn effects(&self) -> &EffectRow {
&self.effects
}
fn description(&self) -> &str {
"Run a Monte Carlo simulation (args: n_samples, seed, quasi_mc, distributions, model)"
}
async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
let n_samples = args
.get("n_samples")
.and_then(Value::as_u64)
.unwrap_or(5000) as usize;
let seed = args.get("seed").and_then(Value::as_u64).unwrap_or(42);
let quasi_mc = args
.get("quasi_mc")
.and_then(Value::as_bool)
.unwrap_or(false);
let dists_json = args
.get("distributions")
.and_then(Value::as_array)
.ok_or_else(|| {
wm_core::CoreError::InvalidArgs("distributions array required".into())
})?;
let distributions: Vec<Distribution> = dists_json
.iter()
.map(parse_distribution)
.collect::<Result<_, _>>()?;
let model_str = args.get("model").and_then(Value::as_str).unwrap_or("sum");
let mut sim = MonteCarloSimulator::new(McConfig {
n_samples,
seed,
quasi_mc,
});
let result = sim.simulate(&distributions, |inputs| match model_str {
"sum" => inputs.iter().sum(),
"product" => inputs.iter().product(),
"mean" => inputs.iter().sum::<f64>() / inputs.len() as f64,
s if s.starts_with("identity:") => {
let idx: usize = s[9..].parse().unwrap_or(0);
inputs.get(idx).copied().unwrap_or(0.0)
}
_ => inputs.iter().sum(),
});
Ok(json!({
"status": "success",
"result": result.to_json(),
}))
}
fn stats(&self) -> &ToolStats {
&self.stats
}
}
pub struct SimForecastTool {
stats: ToolStats,
effects: EffectRow,
}
impl SimForecastTool {
#[must_use]
pub fn new() -> Self {
Self {
stats: ToolStats::default(),
effects: EffectRow::read_only(vec![Resource::Galaxy("simulation".into())]),
}
}
}
impl Default for SimForecastTool {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl Tool for SimForecastTool {
fn name(&self) -> &str {
"sim.forecast"
}
fn gana(&self) -> Gana {
Gana::Mound
}
fn effects(&self) -> &EffectRow {
&self.effects
}
fn description(&self) -> &str {
"Forecast a time series (args: data, horizon, method=moving_average|exponential_smoothing|linear_trend)"
}
async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
let data: Vec<f64> = args
.get("data")
.and_then(Value::as_array)
.ok_or_else(|| wm_core::CoreError::InvalidArgs("data array required".into()))?
.iter()
.map(|v| v.as_f64().unwrap_or(0.0))
.collect();
if data.is_empty() {
return Err(wm_core::CoreError::InvalidArgs(
"data must not be empty".into(),
));
}
let horizon = args.get("horizon").and_then(Value::as_u64).unwrap_or(5) as usize;
let method_str = args
.get("method")
.and_then(Value::as_str)
.unwrap_or("exponential_smoothing");
let method = match method_str {
"moving_average" => ForecastMethod::MovingAverage,
"exponential_smoothing" => ForecastMethod::ExponentialSmoothing,
"linear_trend" => ForecastMethod::LinearTrend,
_ => {
return Err(wm_core::CoreError::InvalidArgs(format!(
"unknown method: {method_str}"
)));
}
};
let forecaster = Forecaster::default();
let result = forecaster.forecast(&data, horizon, method);
Ok(json!({
"status": "success",
"result": result.to_json(),
}))
}
fn stats(&self) -> &ToolStats {
&self.stats
}
}
pub struct SimCounterfactualTool {
stats: ToolStats,
effects: EffectRow,
}
impl SimCounterfactualTool {
#[must_use]
pub fn new() -> Self {
Self {
stats: ToolStats::default(),
effects: EffectRow::read_only(vec![Resource::Galaxy("simulation".into())]),
}
}
}
impl Default for SimCounterfactualTool {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl Tool for SimCounterfactualTool {
fn name(&self) -> &str {
"sim.counterfactual"
}
fn gana(&self) -> Gana {
Gana::Mound
}
fn effects(&self) -> &EffectRow {
&self.effects
}
fn description(&self) -> &str {
"Estimate causal impact of an intervention (args: pre, post)"
}
async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
let pre: Vec<f64> = args
.get("pre")
.and_then(Value::as_array)
.ok_or_else(|| wm_core::CoreError::InvalidArgs("pre array required".into()))?
.iter()
.map(|v| v.as_f64().unwrap_or(0.0))
.collect();
let post: Vec<f64> = args
.get("post")
.and_then(Value::as_array)
.ok_or_else(|| wm_core::CoreError::InvalidArgs("post array required".into()))?
.iter()
.map(|v| v.as_f64().unwrap_or(0.0))
.collect();
if pre.is_empty() || post.is_empty() {
return Err(wm_core::CoreError::InvalidArgs(
"pre and post must not be empty".into(),
));
}
let estimator = CounterfactualEstimator::default();
let result = estimator.estimate(&pre, &post);
Ok(json!({
"status": "success",
"result": result.to_json(),
}))
}
fn stats(&self) -> &ToolStats {
&self.stats
}
}
fn parse_distribution(v: &Value) -> Result<Distribution, wm_core::CoreError> {
let kind = v
.get("kind")
.and_then(Value::as_str)
.ok_or_else(|| wm_core::CoreError::InvalidArgs("distribution kind required".into()))?;
match kind {
"uniform" => {
let min = v.get("min").and_then(Value::as_f64).unwrap_or(0.0);
let max = v.get("max").and_then(Value::as_f64).unwrap_or(1.0);
Ok(Distribution::Uniform { min, max })
}
"normal" => {
let mean = v.get("mean").and_then(Value::as_f64).unwrap_or(0.0);
let std_dev = v.get("std_dev").and_then(Value::as_f64).unwrap_or(1.0);
Ok(Distribution::Normal { mean, std_dev })
}
"exponential" => {
let lambda = v.get("lambda").and_then(Value::as_f64).unwrap_or(1.0);
Ok(Distribution::Exponential { lambda })
}
"triangular" => {
let min = v.get("min").and_then(Value::as_f64).unwrap_or(0.0);
let mode = v.get("mode").and_then(Value::as_f64).unwrap_or(0.5);
let max = v.get("max").and_then(Value::as_f64).unwrap_or(1.0);
Ok(Distribution::Triangular { min, mode, max })
}
"constant" => {
let val = v.get("value").and_then(Value::as_f64).unwrap_or(0.0);
Ok(Distribution::Constant(val))
}
_ => Err(wm_core::CoreError::InvalidArgs(format!(
"unknown distribution kind: {kind}"
))),
}
}
pub struct SimulationCalibrateTool {
store: Arc<Mutex<CalibrationStore>>,
self_model: Option<Arc<Mutex<SelfModel>>>,
stats: ToolStats,
effects: EffectRow,
}
impl SimulationCalibrateTool {
#[must_use]
pub fn new(
store: Arc<Mutex<CalibrationStore>>,
self_model: Option<Arc<Mutex<SelfModel>>>,
) -> Self {
Self {
store,
self_model,
stats: ToolStats::default(),
effects: EffectRow::read_only(vec![Resource::Galaxy("simulation".into())]),
}
}
}
impl Default for SimulationCalibrateTool {
fn default() -> Self {
Self::new(Arc::new(Mutex::new(CalibrationStore::new())), None)
}
}
#[async_trait]
impl Tool for SimulationCalibrateTool {
fn name(&self) -> &str {
"simulation.calibrate"
}
fn gana(&self) -> Gana {
Gana::Mound
}
fn effects(&self) -> &EffectRow {
&self.effects
}
fn stats(&self) -> &ToolStats {
&self.stats
}
fn description(&self) -> &str {
"Record, resolve, and scorecard calibrated predictions — Brier score decomposition (reliability, resolution, uncertainty) for probabilistic forecasts (args: action=record|resolve|scorecard, ...)"
}
async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
let action = args
.get("action")
.and_then(Value::as_str)
.unwrap_or("scorecard");
let mut store = self
.store
.lock()
.map_err(|e| wm_core::CoreError::Tool(format!("calibration store lock: {e}")))?;
match action {
"record" => {
let statement = args
.get("statement")
.and_then(Value::as_str)
.ok_or_else(|| {
wm_core::CoreError::InvalidArgs("statement is required for record".into())
})?;
let probability = args
.get("probability")
.and_then(Value::as_f64)
.unwrap_or(0.5)
.clamp(0.0, 1.0);
let confidence = args
.get("confidence")
.and_then(Value::as_f64)
.unwrap_or(0.5)
.clamp(0.0, 1.0);
let scenario = args
.get("scenario")
.and_then(Value::as_str)
.unwrap_or("default");
let pred = store.record(statement, probability, confidence, scenario);
Ok(json!({
"status": "success",
"prediction_id": pred.id,
"probability": pred.probability,
"adjusted_probability": pred.adjusted_probability,
"calibration_adjustment": store.calibration_gap(),
}))
}
"resolve" => {
let pred_id = args
.get("prediction_id")
.and_then(Value::as_str)
.ok_or_else(|| {
wm_core::CoreError::InvalidArgs(
"prediction_id is required for resolve".into(),
)
})?;
let outcome = args
.get("outcome")
.and_then(Value::as_bool)
.ok_or_else(|| {
wm_core::CoreError::InvalidArgs(
"outcome (boolean) is required for resolve".into(),
)
})?;
match store.resolve(pred_id, outcome) {
Ok((brier, gap)) => Ok(json!({
"status": "success",
"prediction_id": pred_id,
"brier_score": brier,
"calibration_gap": gap,
})),
Err(e) => Err(wm_core::CoreError::InvalidArgs(e)),
}
}
"scorecard" => {
let card = store.scorecard();
let mut result = json!({
"status": "success",
"total_predictions": card.total_predictions,
"resolved": card.resolved,
"unresolved": card.unresolved,
"avg_brier_score": card.avg_brier_score,
"reliability": card.reliability,
"resolution": card.resolution,
"uncertainty": card.uncertainty,
"skill_score": card.skill_score,
"calibration_gap": card.calibration_gap,
"perfect_calibration": card.perfect_calibration,
"good_calibration": card.good_calibration,
"calibration_bins": card.calibration_bins,
});
if let Some(model) = &self.self_model {
if let Ok(model) = model.lock() {
model.record(MetricKind::BrierScore, card.avg_brier_score as f32);
let alerts = model
.check_alerts()
.into_iter()
.filter(|a| a.metric == MetricKind::BrierScore)
.map(|a| json!({"level": format!("{:?}", a.level), "message": a.message}))
.collect::<Vec<_>>();
result["alerts"] = json!(alerts);
}
}
Ok(result)
}
other => Err(wm_core::CoreError::InvalidArgs(format!(
"unknown action '{other}' (expected record | resolve | scorecard)"
))),
}
}
}
#[must_use]
pub fn register_simulation(
registry: &wm_dispatch::ToolRegistry,
calibration_store: Option<Arc<Mutex<CalibrationStore>>>,
self_model: Option<Arc<Mutex<SelfModel>>>,
) -> wm_dispatch::ToolRegistry {
let calibrate = match calibration_store {
Some(store) => SimulationCalibrateTool::new(store, self_model),
None => SimulationCalibrateTool::default(),
};
registry
.register(Arc::new(SimMcTool::new()))
.register(Arc::new(SimForecastTool::new()))
.register(Arc::new(SimCounterfactualTool::new()))
.register(Arc::new(calibrate))
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn sim_mc_runs_simulation() {
let tool = SimMcTool::new();
let mut ctx = Context::default();
let v = tool
.call(
&mut ctx,
json!({
"n_samples": 1000,
"distributions": [{"kind": "uniform", "min": 0.0, "max": 10.0}],
"model": "sum"
}),
)
.await
.unwrap();
assert_eq!(v["status"], "success");
assert!(v["result"]["mean"].is_number());
}
#[tokio::test]
async fn sim_mc_missing_distributions_errors() {
let tool = SimMcTool::new();
let mut ctx = Context::default();
let result = tool.call(&mut ctx, json!({})).await;
assert!(result.is_err());
}
#[tokio::test]
async fn sim_forecast_runs() {
let tool = SimForecastTool::new();
let mut ctx = Context::default();
let v = tool
.call(
&mut ctx,
json!({
"data": [1.0, 2.0, 3.0, 4.0, 5.0],
"horizon": 3,
"method": "linear_trend"
}),
)
.await
.unwrap();
assert_eq!(v["status"], "success");
assert!(v["result"]["forecast"].is_array());
}
#[tokio::test]
async fn sim_forecast_empty_data_errors() {
let tool = SimForecastTool::new();
let mut ctx = Context::default();
let result = tool.call(&mut ctx, json!({"data": []})).await;
assert!(result.is_err());
}
#[tokio::test]
async fn sim_counterfactual_runs() {
let tool = SimCounterfactualTool::new();
let mut ctx = Context::default();
let v = tool
.call(
&mut ctx,
json!({
"pre": [10.0, 10.0, 10.0, 10.0, 10.0],
"post": [15.0, 15.0, 15.0, 15.0, 15.0]
}),
)
.await
.unwrap();
assert_eq!(v["status"], "success");
assert!(v["result"]["impact"].is_number());
}
#[tokio::test]
async fn sim_counterfactual_missing_pre_errors() {
let tool = SimCounterfactualTool::new();
let mut ctx = Context::default();
let result = tool.call(&mut ctx, json!({"post": [1.0]})).await;
assert!(result.is_err());
}
#[tokio::test]
async fn sim_tools_are_mound_gana() {
assert_eq!(SimMcTool::new().gana(), Gana::Mound);
assert_eq!(SimForecastTool::new().gana(), Gana::Mound);
assert_eq!(SimCounterfactualTool::new().gana(), Gana::Mound);
}
#[tokio::test]
async fn parse_distribution_uniform() {
let d = parse_distribution(&json!({"kind": "uniform", "min": 0.0, "max": 10.0})).unwrap();
assert!(matches!(
d,
Distribution::Uniform {
min: 0.0,
max: 10.0
}
));
}
#[tokio::test]
async fn parse_distribution_normal() {
let d =
parse_distribution(&json!({"kind": "normal", "mean": 5.0, "std_dev": 2.0})).unwrap();
assert!(matches!(
d,
Distribution::Normal {
mean: 5.0,
std_dev: 2.0
}
));
}
#[tokio::test]
async fn parse_distribution_unknown_errors() {
let result = parse_distribution(&json!({"kind": "unknown"}));
assert!(result.is_err());
}
#[tokio::test]
async fn calibrate_record_resolve_scorecard_flow() {
let store = Arc::new(Mutex::new(CalibrationStore::new()));
let tool = SimulationCalibrateTool::new(Arc::clone(&store), None);
let mut ctx = Context::default();
let v = tool
.call(
&mut ctx,
json!({"action": "record", "statement": "It will rain", "probability": 0.8, "confidence": 0.6, "scenario": "weather"}),
)
.await
.unwrap();
assert_eq!(v["status"], "success");
let pred_id = v["prediction_id"].as_str().unwrap().to_string();
let v = tool
.call(
&mut ctx,
json!({"action": "resolve", "prediction_id": pred_id, "outcome": true}),
)
.await
.unwrap();
assert_eq!(v["status"], "success");
assert!((v["brier_score"].as_f64().unwrap() - 0.04).abs() < 1e-9);
let v = tool
.call(&mut ctx, json!({"action": "scorecard"}))
.await
.unwrap();
assert_eq!(v["resolved"], 1);
assert_eq!(v["unresolved"], 0);
assert!((v["avg_brier_score"].as_f64().unwrap() - 0.04).abs() < 1e-9);
assert_eq!(v["calibration_bins"].as_array().unwrap().len(), 10);
}
#[tokio::test]
async fn calibrate_requires_statement_and_outcome() {
let tool = SimulationCalibrateTool::default();
let mut ctx = Context::default();
assert!(
tool.call(&mut ctx, json!({"action": "record"}))
.await
.is_err()
);
assert!(
tool.call(&mut ctx, json!({"action": "resolve"}))
.await
.is_err()
);
assert!(
tool.call(
&mut ctx,
json!({"action": "resolve", "prediction_id": "nope", "outcome": true})
)
.await
.is_err()
);
assert!(
tool.call(&mut ctx, json!({"action": "bogus"}))
.await
.is_err()
);
}
#[tokio::test]
async fn calibrate_scorecard_defaults_to_scorecard_action() {
let tool = SimulationCalibrateTool::default();
let mut ctx = Context::default();
let v = tool.call(&mut ctx, json!({})).await.unwrap();
assert_eq!(v["status"], "success");
assert_eq!(v["resolved"], 0);
assert_eq!(v["total_predictions"], 0);
}
#[tokio::test]
async fn calibrate_feeds_brier_into_self_model() {
let store = Arc::new(Mutex::new(CalibrationStore::new()));
let model = Arc::new(Mutex::new(wm_selfmodel::SelfModel::new()));
let tool = SimulationCalibrateTool::new(Arc::clone(&store), Some(Arc::clone(&model)));
let mut ctx = Context::default();
for i in 0..14 {
let v = tool
.call(
&mut ctx,
json!({"action": "record", "statement": format!("pred {i}"), "probability": 0.9}),
)
.await
.unwrap();
let id = v["prediction_id"].as_str().unwrap().to_string();
let _ = tool
.call(
&mut ctx,
json!({"action": "resolve", "prediction_id": id, "outcome": false}),
)
.await
.unwrap();
}
let mut saw_alert = false;
for _ in 0..20 {
let v = tool
.call(&mut ctx, json!({"action": "scorecard"}))
.await
.unwrap();
assert!(v["avg_brier_score"].as_f64().unwrap() > 0.6);
if let Some(alerts) = v["alerts"].as_array() {
if alerts
.iter()
.any(|a| a["level"] == "Warning" || a["level"] == "Critical")
{
saw_alert = true;
break;
}
}
}
assert!(saw_alert, "drift should produce Brier alerts");
let model = model.lock().unwrap();
assert!(model.sample_count(MetricKind::BrierScore) >= 1);
}
}