use crate::client::Client;
use crate::{ErrorCode, HstpError};
use ndarray::{Array1, ArrayBase, Dim, OwnedRepr};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(untagged)]
pub(crate) enum Tensor {
Float(f64),
Vector(Vec<Tensor>),
}
#[derive(Debug, Deserialize, Serialize)]
pub(crate) struct Factor {
pub swid: String,
pub schema: Vec<String>,
pub variables: Vec<String>,
pub tensor: Vec<Tensor>,
pub messages: HashMap<String, Vec<f64>>,
pub update: Option<bool>,
pub learnable: Option<bool>,
pub cpd: Option<bool>,
pub sharable: Option<bool>,
pub last_updated: Option<String>,
}
impl TryFrom<&Value> for Factor {
type Error = HstpError;
fn try_from(factor: &Value) -> Result<Self, HstpError> {
let deserialized: Factor =
serde_json::from_value(factor.clone()).map_err(|e| HstpError {
code: ErrorCode::EntityParseError as i32,
message: format!("Entity parsing error: {} Error: {}", factor, e),
swid: "".to_string(),
})?;
Ok(deserialized)
}
}
#[derive(Deserialize, Serialize, Clone)]
struct Variable {
swid: String,
schema: Vec<String>,
name: String,
elements: Vec<String>,
messages: HashMap<String, Vec<f64>>,
observation: Option<Observation>,
}
impl TryFrom<Value> for Variable {
type Error = HstpError;
fn try_from(var: Value) -> Result<Self, HstpError> {
let deserialized: Variable =
serde_json::from_value(var.clone()).map_err(|e| HstpError {
code: ErrorCode::EntityParseError as i32,
message: format!("Entity parsing error: {} Error: {}", var, e),
swid: "".to_string(),
})?;
Ok(deserialized)
}
}
impl TryFrom<&Value> for Variable {
type Error = HstpError;
fn try_from(var: &Value) -> Result<Self, HstpError> {
let deserialized: Variable =
serde_json::from_value(var.clone()).map_err(|e| HstpError {
code: ErrorCode::EntityParseError as i32,
message: format!("Entity parsing error: {} Error: {}", var, e),
swid: "".to_string(),
})?;
Ok(deserialized)
}
}
#[derive(Deserialize, Serialize, Clone)]
#[serde(untagged)]
pub enum ObservationValue {
Element(String),
Distribution(Vec<f64>),
None,
}
#[derive(Deserialize, Serialize, Clone)]
struct Observation {
value: ObservationValue,
timestamp: f32,
consumed_by: Option<HashMap<String, bool>>,
}
pub(crate) async fn observe_variables(
client: &mut Client,
observations: HashMap<String, ObservationValue>,
allow_learning: Option<bool>,
remove_observations: Option<bool>,
) -> Result<Vec<Value>, HstpError> {
let mut upsert_results = Vec::new();
let allow_learning = allow_learning.unwrap_or(true);
let remove_observations = remove_observations.unwrap_or(false);
for (var_identifier, observation_value) in &observations {
let existing_var: Variable = client.query(var_identifier).await?.try_into()?;
let obs_array = match observation_value {
ObservationValue::Element(string) => {
let var_elements = existing_var.elements.clone();
let mut observation_array = vec![0.0f64; var_elements.len()];
if let Some(found_position) = var_elements.iter().position(|x| x == string) {
observation_array[found_position] = 1.0;
observation_array
} else {
return Err(HstpError {
code: ErrorCode::EntityParseError as i32,
message: format!(
"Element {} not found in variable {}",
string, var_identifier
),
swid: "".to_string(),
});
}
}
ObservationValue::Distribution(vec) => vec.clone(),
ObservationValue::None => {
vec![]
}
};
let consumed_by = if allow_learning {
Some(
existing_var
.messages
.keys()
.map(|k| (k.clone(), false))
.collect(),
)
} else {
None
};
let observation = if !obs_array.is_empty() {
Some(Observation {
value: ObservationValue::Distribution(obs_array),
timestamp: 1.0,
consumed_by,
})
} else {
None
};
let mut var_with_observation = existing_var.clone();
var_with_observation.observation = observation;
upsert_results.push(_upsert(client, _get_query::<Variable>(var_with_observation)?).await?);
if remove_observations {
clear_observations(client, Some(observations.keys().cloned().collect())).await?;
}
}
Ok(upsert_results)
}
pub async fn clear_observations(
client: &mut Client,
variables: Option<Vec<String>>,
) -> Result<Vec<Value>, HstpError> {
let var_query = match variables {
Some(_) => _get_query(variables)?,
None => "#variable".to_string(),
};
let mut upsert_results: Vec<Value> = vec![];
let response = &client.query(var_query).await?["entities"];
match response {
Value::Array(vec) => {
for v in vec.iter() {
let mut variable: Variable = v.try_into()?;
variable.observation.take();
upsert_results.push(_upsert(client, _get_query::<Variable>(variable)?).await?);
}
}
_ => {
return Err(HstpError {
code: ErrorCode::InsufficientPermissions as i32,
message: "No token provided".to_string(),
swid: "".to_string(),
})
}
};
Ok(upsert_results)
}
fn _get_query<T: Serialize>(entities: T) -> Result<String, HstpError> {
let result = serde_json::to_string(&entities);
match result {
Ok(s) => Ok(s),
_ => Err(HstpError {
code: ErrorCode::EntityParseError as i32,
message: "Entity parsing error".to_string(),
swid: "".to_string(),
}),
}
}
pub(crate) fn get_marginal(variable: Value) -> Result<Vec<f64>, HstpError> {
let messages = variable.get("messages");
let elements = variable.get("elements");
match (messages, elements) {
(Some(messages), Some(elements)) => match messages.as_object() {
Some(object) => {
if let Some(dim) = elements.as_array().map(|arr| arr.len()) {
let mut out = Array1::from_vec(vec![1.0; dim]);
let observation = variable.get("observation");
if let Some(obs) = observation {
if let Some(obs_value) = obs.get("value").and_then(|v| v.as_array()) {
let obs_vec: ArrayBase<OwnedRepr<f64>, Dim<[usize; 1]>> = Array1::from(
obs_value
.iter()
.map(|x| x.as_f64().unwrap_or(1.0))
.collect::<Vec<_>>(),
);
out = out * obs_vec;
}
}
for (_k, value) in object.iter() {
match value {
Value::Array(arr) => {
let arr_vec: ArrayBase<OwnedRepr<f64>, Dim<[usize; 1]>> =
Array1::from(
arr.iter()
.map(|x| x.as_f64().unwrap_or(1.0))
.collect::<Vec<_>>(),
);
out = out * arr_vec;
}
_ => {
return Err(HstpError {
code: ErrorCode::EntityParseError as i32,
message: "Unexpected value in marginal computation".to_string(),
swid: "".to_string(),
})
}
}
}
let out = out.clone() / out.clone().iter().sum::<f64>();
Ok(out.into_raw_vec())
} else {
Err(HstpError {
code: ErrorCode::EntityParseError as i32,
message: "Elements array missing on variable".to_string(),
swid: "".to_string(),
})
}
}
_ => Err(HstpError {
code: ErrorCode::EntityParseError as i32,
message: "Entity parsing error".to_string(),
swid: "".to_string(),
}),
},
_ => Err(HstpError {
code: ErrorCode::EntityParseError as i32,
message: "Entity parsing error".to_string(),
swid: "".to_string(),
}),
}
}
pub(crate) async fn get_marginal_from_core(
client: &mut Client,
variable: &str,
) -> Result<Vec<f64>, HstpError> {
let entity_response = client.query(variable).await?;
get_marginal(entity_response)
}
async fn _upsert(client: &mut Client, data: String) -> Result<Value, HstpError> {
Ok(client
.query(format!(
"%upsert
{:}
on collide overwrite",
data
))
.await?)
}
pub async fn get_probability(
client: &mut Client,
variables: Vec<String>,
evidence: Option<HashMap<String, ObservationValue>>,
) -> Result<HashMap<String, Vec<f64>>, HstpError> {
if let Some(evidence) = evidence {
observe_variables(client, evidence.clone(), Some(false), None).await?;
}
let mut response_dict = HashMap::new();
for v in variables.iter() {
let marginal: Vec<f64> = get_marginal_from_core(client, v).await?;
response_dict.insert(v.clone(), marginal);
}
Ok(response_dict)
}