fuzzytags-sim/src/oracle.rs

// todo would be neat to just make this a subscriber of the tracing log, but there doesn't
// seem to be a nice way for that API to do what we want.. so until then...

use hashbrown::{HashMap, HashSet};
use rand::{thread_rng, Rng};
use rand_distr::num_traits::FloatConst;
use std::fs::File;
use std::io::Write;

#[derive(Clone)]
pub struct Event {
    sender: String,
    intended_receiver: String,
    entangled_receiver: Option<String>,
    confidence: f64,
}

#[derive(Clone)]
pub struct Oracle {
    parties: Vec<String>,
    actual_events: Vec<Event>,
    suspect: HashMap<(String, String), f64>,
}

impl Oracle {
    pub fn new() -> Oracle {
        Oracle {
            parties: vec![],
            actual_events: vec![],
            suspect: HashMap::new(),
        }
    }

    pub fn register_party(&mut self, party: String) {
        self.parties.push(party);
    }

    pub fn add_event(&mut self, sender: String, intended_receiver: String, entangled_receiver: Option<String>, confidence: f64) {
        self.actual_events.push(Event {
            sender,
            intended_receiver,
            entangled_receiver,
            confidence,
        });
    }

    pub fn add_suspect(&mut self, sender: &String, intended_receiver: &String, weight: f64) {
        self.suspect.insert((sender.clone(), intended_receiver.clone()), weight);
    }

    pub fn compile_to_dot(&self, filename: &str, strict: bool, inverse: bool, max: f64) -> f64 {
        let mut output = File::create(filename).unwrap();
        write!(output, "strict digraph {{\n");
        write!(
            output,
            r##"
                        K=1.3;
                        repulsiveforce=0.00001;
                        dpi=2400;
                        bgcolor = "#111111ff";
                        "##
        );

        for (i, party) in self.parties.iter().enumerate() {
            let x = (((f64::PI() * 2.0) / self.parties.len() as f64) * (i as f64)).cos() * 10.0;
            let y = (((f64::PI() * 2.0) / self.parties.len() as f64) * (i as f64)).sin() * 10.0;

            let r = hex::decode(party).unwrap()[0];
            let g = hex::decode(party).unwrap()[1];
            let b = hex::decode(party).unwrap()[2];
            writeln!(
                output,
                "\"{}\" [tooltip=\"{} ({})\", shape=point, penwidth=0, fixedsize=true, width=0.005,height=0.005,peripheries=0,style=\"filled,setlinewidth(0)\", color=\"#{:02x}{:02x}{:02x}\"]",
                party, party, i, r, g, b
            );
            //writeln!(output, "\"{}\" [shape=point, color=\"#{:x}{:x}{:x}\",pos=\"{},{}!\"]", party, r, g, b,x,y);
        }

        let mut real_connection_map: HashMap<(String, String), f64> = HashMap::new();
        let mut entangled_connection_map: HashMap<(String, String), f64> = HashMap::new();
        let mut max_conn = 1.0;
        for event in self.actual_events.iter() {
            let key = (event.sender.clone(), event.intended_receiver.clone());
            if real_connection_map.contains_key(&key) {
                *real_connection_map.get_mut(&key).unwrap() += (1.0 * (1.0 - event.confidence));
                if real_connection_map[&key] > max_conn {
                    max_conn = real_connection_map[&key];
                }
            } else {
                real_connection_map.insert(key, 1.0 * (1.0 - event.confidence));
            }
            match &event.entangled_receiver {
                Some(entangled_receiver) => {
                    let key = (event.sender.clone(), entangled_receiver.clone());
                    if entangled_connection_map.contains_key(&key) {
                        *entangled_connection_map.get_mut(&key).unwrap() += 1.0;
                        if entangled_connection_map[&key] > max_conn {
                            max_conn = entangled_connection_map[&key];
                        }
                    } else {
                        entangled_connection_map.insert(key, 1.0);
                    }
                }
                _ => {}
            };
        }

        for ((sender, receiver), size) in real_connection_map.iter() {
            let normalized = (*size as f64 / max_conn as f64);
            let mut transparency = (normalized * 64.0) as u8 + 172;
            let mut penwidth = (normalized * 0.01) as f64;
            match self.suspect.get(&(sender.clone(), receiver.clone())) {
                Some(weight) => {
                    let normalized = ((*weight as f64 * 1000.0).max(999.0) - 999.0) / 1.0;
                    let mut penwidth = (normalized * 0.01) as f64;
                    let mut transparency = (normalized * 64.0) as u8 + 172;
                    writeln!(
                        output,
                        "\"{}\" -> \"{}\" [arrowhead=none, penwidth={}, color=\"#{:02x}0000{:02x}\", weight={}]",
                        sender,
                        receiver,
                        f64::max(0.005, penwidth),
                        transparency,
                        transparency,
                        penwidth
                    );
                }
                _ => {
                    if self.suspect.is_empty() {
                        writeln!(
                            output,
                            "\"{}\" -> \"{}\" [arrowhead=none, penwidth={}, color=\"#ffffff{:02x}\", weight={}]",
                            sender,
                            receiver,
                            f64::max(0.005, penwidth),
                            transparency,
                            penwidth
                        );
                    }
                }
            }
        }
        for ((sender, receiver), size) in entangled_connection_map.iter() {
            let normalized = (*size as f64 / max_conn as f64);
            let transparency = (normalized * 172.0) as u8 + 64;
            let penwidth = ((normalized * 32.0) as f64).log2();
            // writeln!(
            //     output,
            //     "\"{}\" -> \"{}\" [arrowhead=none, style=dashed, penwidth={}, color=\"#ffffff{:02x}\"]",
            //    sender, receiver, penwidth, transparency
            //);
        }
        write!(output, "}}");
        max_conn
    }
}