nexosim 1.0.0

use std::collections::HashMap;
use std::hash::{BuildHasherDefault, DefaultHasher};
use std::time::Duration;

use serde::{Deserialize, Serialize, de::DeserializeOwned};

use nexosim::model::{Context, Model, schedulable};
use nexosim::ports::{
    EventQueueReader, EventSinkReader, EventSource, Output, QuerySource, SinkState, event_queue,
};
use nexosim::simulation::{EventId, EventKey, Mailbox, QueryId, SimInit};
use nexosim::time::MonotonicTime;

#[derive(Default, Serialize, Deserialize)]
struct ModelWithOutput {
    output: Output<String>,
}
#[Model]
impl ModelWithOutput {
    #[nexosim(schedulable)]
    pub async fn send(&mut self, input: u32) {
        self.output.send(format!("{input}")).await;
    }
}

#[derive(Default, Serialize, Deserialize)]
struct ModelWithKey {
    key: Option<EventKey>,
}
#[Model]
impl ModelWithKey {
    #[nexosim(schedulable)]
    pub async fn process(&mut self) {
        if let Some(key) = self.key.take() {
            key.cancel();
        }
    }
    pub fn set_key(&mut self, key: EventKey) {
        self.key = Some(key);
    }
}

#[derive(Serialize, Deserialize)]
struct ModelWithState {
    state: u32,
}
#[Model]
impl ModelWithState {
    pub fn new(state: u32) -> Self {
        Self { state }
    }
    #[nexosim(init)]
    async fn init(&mut self, _: &Context<Self>, _: &mut ()) {
        self.state *= 11;
    }
    pub async fn query(&mut self) -> u32 {
        self.state
    }
    pub async fn add(&mut self, arg: u32) {
        self.state += arg;
    }
    pub async fn mul(&mut self, arg: u32) {
        self.state *= arg;
    }
}

#[derive(Serialize, Deserialize)]
struct ModelWithSchedule {
    output: Output<u32>,
}
#[Model]
impl ModelWithSchedule {
    pub fn new() -> Self {
        Self {
            output: Output::default(),
        }
    }
    #[nexosim(init)]
    async fn init(&mut self, cx: &Context<Self>, _: &mut ()) {
        cx.schedule_periodic_event(
            Duration::from_secs(1),
            Duration::from_secs(2),
            schedulable!(Self::send),
            7,
        )
        .unwrap();
    }
    #[nexosim(schedulable)]
    async fn send(&mut self, arg: u32, cx: &Context<Self>) {
        self.output.send(cx.time().as_secs() as u32 * arg).await;
    }
}

#[derive(Serialize, Deserialize)]
pub struct ModelWithGenerics<T, U, const N: usize>
where
    U: Clone + Send + Sync + 'static,
{
    value: T,
    output: Output<U>,
}

#[Model]
impl<T, U, const N: usize> ModelWithGenerics<T, U, N>
where
    T: Serialize + DeserializeOwned + Clone + Send + 'static,
    U: Serialize + DeserializeOwned + Clone + Send + Sync + 'static,
{
    #[nexosim(schedulable)]
    async fn input(&mut self, arg: (T, U)) {
        self.value = arg.0;
        self.output.send(arg.1).await;
    }

    #[nexosim(schedulable)]
    async fn process(&mut self) {}

    #[nexosim(init)]
    async fn init(&mut self, cx: &Context<Self>, _: &mut ()) {
        // Test if schedulable! macro works with generics.
        cx.schedule_event(
            Duration::from_secs(255),
            schedulable!(ModelWithGenerics::<T, U, N>::process),
            (),
        )
        .unwrap();
    }
}

#[derive(Serialize, Deserialize)]
pub struct ModelWithGenericEnv<T>(T);
#[Model(type Env=T)]
impl<T> ModelWithGenericEnv<T>
where
    T: Serialize + DeserializeOwned + Clone + Send + 'static,
{
    #[nexosim(schedulable)]
    async fn input(&mut self, arg: T, _: &Context<Self>, env: &mut T) {
        *env = arg;
    }

    #[nexosim(init)]
    async fn init(&mut self, _: &Context<Self>, _: &mut T) {}
}

#[derive(Serialize, Deserialize)]
pub struct ModelWithEnvWithGeneric<T>(T);
#[derive(Default)]
pub struct EnvWithGeneric<T>(T);

#[Model(type Env=EnvWithGeneric<T>)]
impl<T> ModelWithEnvWithGeneric<T>
where
    T: Serialize + DeserializeOwned + Clone + Send + 'static,
{
    #[nexosim(schedulable)]
    async fn input(&mut self, arg: T, _: &Context<Self>, env: &mut EnvWithGeneric<T>) {
        env.0 = arg;
    }

    #[nexosim(init)]
    async fn init(&mut self, _: &Context<Self>, _: &mut EnvWithGeneric<T>) {}
}

/// This model uses a non random hasher to allow consistent (de)serialization
/// order, needed by outputs.
#[derive(Serialize, Deserialize)]
pub struct ModelWithHashMap {
    outputs: HashMap<usize, Output<usize>, BuildHasherDefault<DefaultHasher>>,
}
#[Model]
impl ModelWithHashMap {
    pub async fn send(&mut self, arg: usize) {
        self.outputs.get_mut(&arg).unwrap().send(arg).await;
    }
}

#[test]
fn model_with_output() {
    fn get_bench() -> (SimInit, EventId<u32>, EventQueueReader<String>) {
        let mbox = Mailbox::new();
        let mut model = ModelWithOutput::default();

        let (sink, msg) = event_queue(SinkState::Enabled);
        model.output.connect_sink(sink);

        let mut bench = SimInit::new();
        let event = EventSource::new()
            .connect(ModelWithOutput::send, &mbox)
            .register(&mut bench);
        bench = bench.add_model(model, mbox, "modelWithOutput");

        (bench, event, msg)
    }

    let (bench, event, _) = get_bench();
    let t0 = MonotonicTime::EPOCH;

    let mut simu = bench.init(t0).unwrap();

    // Schedule event on an initialized sim.
    let _ = simu
        .scheduler()
        .schedule_event(Duration::from_secs(5), &event, 5);

    // Store state with an event scheduled.
    let mut state = Vec::new();
    simu.save(&mut state).unwrap();

    // Recreate the bench with the state restored.
    let (bench, _, mut msg) = get_bench();
    let mut simu = bench.restore(&state[..]).unwrap();

    // Verify that the scheduled event gets fired.
    simu.step().unwrap();
    assert_eq!(msg.try_read(), Some("5".to_string()));
}

#[test]
fn model_with_key() {
    fn get_bench() -> (
        SimInit,
        EventId<u32>,
        EventId<EventKey>,
        EventId<()>,
        EventQueueReader<String>,
    ) {
        let output_mbox = Mailbox::new();
        let mut output_model = ModelWithOutput::default();

        let key_mbox = Mailbox::new();
        let key_model = ModelWithKey { key: None };

        let (sink, msg) = event_queue(SinkState::Enabled);
        output_model.output.connect_sink(sink);

        let mut bench = SimInit::new();

        let output = EventSource::new()
            .connect(ModelWithOutput::send, &output_mbox)
            .register(&mut bench);

        let set_key = EventSource::new()
            .connect(ModelWithKey::set_key, &key_mbox)
            .register(&mut bench);

        let process = EventSource::new()
            .connect(ModelWithKey::process, &key_mbox)
            .register(&mut bench);

        bench = bench
            .add_model(output_model, output_mbox, "modelWithOutput")
            .add_model(key_model, key_mbox, "modelWithKey");

        (bench, output, set_key, process, msg)
    }

    let (bench, output, set_key, _, _) = get_bench();
    let t0 = MonotonicTime::EPOCH;

    let mut simu = bench.init(t0).unwrap();
    // Schedule event on an initialized sim.
    let key = simu
        .scheduler()
        .schedule_keyed_event(Duration::from_secs(5), &output, 5)
        .unwrap();

    let _ = simu.process_event(&set_key, key);

    // Store state with an event scheduled and key set.
    let mut state = Vec::new();
    simu.save(&mut state).unwrap();

    // Recreate the bench with the state restored.
    let (bench, _, _, process, mut msg) = get_bench();
    let mut simu = bench.restore(&state[..]).unwrap();

    // Cancel the serialized key.
    let _ = simu.process_event(&process, ());

    // Verify that the scheduled event does not fire.
    simu.step().unwrap();
    assert_eq!(msg.try_read(), None);
}

#[test]
fn model_init() {
    fn get_bench() -> (SimInit, QueryId<(), u32>) {
        let mbox = Mailbox::new();
        let model = ModelWithState::new(1);

        let mut bench = SimInit::new();

        let query = QuerySource::new()
            .connect(ModelWithState::query, &mbox)
            .register(&mut bench);

        bench = bench.add_model(model, mbox, "modelWithKey");

        (bench, query)
    }

    let (bench, query) = get_bench();
    let t0 = MonotonicTime::EPOCH;

    let mut simu = bench.init(t0).unwrap();
    // Verify `init` called.
    let model_state = simu.process_query(&query, ()).unwrap();
    assert_eq!(model_state, 11);

    // // Store state with an initialized model.
    let mut state = Vec::new();
    simu.save(&mut state).unwrap();

    // // Recreate the bench with the state restored.
    let (bench, _) = get_bench();
    let mut simu = bench.restore(&state[..]).unwrap();

    // Verify that `init` has not been called again.
    let model_state = simu.process_query(&query, ()).unwrap();
    assert_eq!(model_state, 11);
}

#[test]
fn model_with_schedule() {
    fn get_bench() -> (SimInit, EventQueueReader<u32>) {
        let mbox = Mailbox::new();
        let mut model = ModelWithSchedule::new();

        let (sink, msg) = event_queue(SinkState::Enabled);
        model.output.connect_sink(sink);

        let bench = SimInit::new().add_model(model, mbox, "modelWithSchedule");

        (bench, msg)
    }

    let (bench, mut msg) = get_bench();
    let t0 = MonotonicTime::EPOCH;

    let mut simu = bench.init(t0).unwrap();
    simu.step().unwrap();
    assert_eq!(msg.try_read(), Some(7));

    // Store state with a scheduled model after one step.
    let mut state = Vec::new();
    simu.save(&mut state).unwrap();

    // Recreate the bench with the state restored.
    let (bench, mut msg) = get_bench();
    let mut simu = bench.restore(&state[..]).unwrap();

    // Verify that the scheduled event gets fired as step two.
    simu.step().unwrap();
    assert_eq!(msg.try_read(), Some(21));
}

#[test]
fn model_with_generics() {
    fn get_bench() -> (SimInit, EventQueueReader<f64>, EventId<(i32, f64)>) {
        let mbox = Mailbox::new();
        let mut model = ModelWithGenerics::<i32, f64, 13> {
            value: 8,
            output: Output::<f64>::default(),
        };

        let mut bench = SimInit::new();

        let (sink, msg) = event_queue(SinkState::Enabled);
        model.output.connect_sink(sink);

        let input = EventSource::new()
            .connect(ModelWithGenerics::input, &mbox)
            .register(&mut bench);

        bench = bench.add_model(model, mbox, "modelWithGenerics");

        (bench, msg, input)
    }

    let (bench, mut msg, input) = get_bench();
    let t0 = MonotonicTime::EPOCH;

    let mut simu = bench.init(t0).unwrap();
    let scheduler = simu.scheduler();
    scheduler
        .schedule_event(Duration::from_secs(2), &input, (-5, 5.14))
        .unwrap();
    scheduler
        .schedule_event(Duration::from_secs(5), &input, (-5, 7.14))
        .unwrap();
    simu.step().unwrap();

    assert_eq!(msg.try_read(), Some(5.14));

    // Store state with a scheduled model after one step.
    let mut state = Vec::new();
    simu.save(&mut state).unwrap();

    // Recreate the bench with the state restored.
    let (bench, mut msg, _) = get_bench();
    let mut simu = bench.restore(&state[..]).unwrap();

    // Verify that the scheduled event gets fired as step two.
    simu.step().unwrap();
    assert_eq!(msg.try_read(), Some(7.14));
}

#[test]
fn model_with_generic_env() {
    fn get_bench() -> SimInit {
        let mbox_a = Mailbox::new();
        let mbox_b = Mailbox::new();
        let model_a = ModelWithGenericEnv(13);
        let model_b = ModelWithEnvWithGeneric(27);

        SimInit::new()
            .add_model(model_a, mbox_a, "model_a")
            .add_model(model_b, mbox_b, "model_b")
    }

    let bench = get_bench();
    let t0 = MonotonicTime::EPOCH;

    bench.init(t0).unwrap();
}

#[test]
fn model_with_hashmap() {
    const COUNT: usize = 10;
    const ITERATIONS: usize = 30;

    fn get_bench() -> (SimInit, EventId<usize>, Vec<EventQueueReader<usize>>) {
        let mut sinks = vec![];

        let mbox = Mailbox::new();
        let mut model = ModelWithHashMap {
            outputs: HashMap::with_hasher(BuildHasherDefault::new()),
        };

        let mut bench = SimInit::new();

        for idx in 0..COUNT {
            let mut output = Output::new();
            let (sink, msg) = event_queue(SinkState::Enabled);
            output.connect_sink(sink);
            model.outputs.insert(idx, output);
            sinks.push(msg);
        }

        let event = EventSource::new()
            .connect(ModelWithHashMap::send, &mbox)
            .register(&mut bench);

        bench = bench.add_model(model, mbox, "model");

        (bench, event, sinks)
    }

    let (bench, _, _) = get_bench();
    let t0 = MonotonicTime::EPOCH;

    let mut simu = bench.init(t0).unwrap();

    let mut state = Vec::new();
    simu.save(&mut state).unwrap();

    let (bench, event, mut sinks) = get_bench();
    let mut simu = bench.restore(&state[..]).unwrap();

    for _ in 0..ITERATIONS {
        // Verify that after the deserialization output connections still point to
        // assigned sinks. (fails when a standard RandomState HashMap is used)
        #[allow(clippy::needless_range_loop)]
        for idx in 0..COUNT {
            simu.process_event(&event, idx).unwrap();
            assert_eq!(sinks[idx].try_read(), Some(idx));
        }
    }
}

#[test]
fn model_relative_order() {
    fn get_bench() -> (SimInit, EventId<u32>, EventId<u32>, QueryId<(), u32>) {
        let mbox = Mailbox::new();
        let model = ModelWithState::new(1);

        let mut bench = SimInit::new();
        let add = EventSource::new()
            .connect(ModelWithState::add, &mbox)
            .register(&mut bench);

        let mul = EventSource::new()
            .connect(ModelWithState::mul, &mbox)
            .register(&mut bench);

        let query = QuerySource::new()
            .connect(ModelWithState::query, &mbox)
            .register(&mut bench);

        bench = bench.add_model(model, mbox, "modelWithKey");

        (bench, add, mul, query)
    }

    // Test mul -> add order.

    let (bench, add, mul, _) = get_bench();
    let t0 = MonotonicTime::EPOCH;

    let mut simu = bench.init(t0).unwrap();

    // Schedule two events at the same time
    simu.scheduler()
        .schedule_event(Duration::from_secs(1), &mul, 7)
        .unwrap();
    simu.scheduler()
        .schedule_event(Duration::from_secs(1), &add, 19)
        .unwrap();

    // Store state with an initialized model and events scheduled.
    let mut state = Vec::new();
    simu.save(&mut state).unwrap();

    // Recreate the bench with the state restored.
    let (bench, _, _, query) = get_bench();
    let mut simu = bench.restore(&state[..]).unwrap();

    // Verify events have been called in the right order.
    simu.step().unwrap();
    assert_eq!(11 * 7 + 19, simu.process_query(&query, ()).unwrap());

    // Test add -> mul order.

    let (bench, add, mul, _) = get_bench();
    let t0 = MonotonicTime::EPOCH;

    let mut simu = bench.init(t0).unwrap();

    // Schedule two events at the same time
    simu.scheduler()
        .schedule_event(Duration::from_secs(1), &add, 19)
        .unwrap();
    simu.scheduler()
        .schedule_event(Duration::from_secs(1), &mul, 7)
        .unwrap();

    // Store state with an initialized model and events scheduled.
    let mut state = Vec::new();
    simu.save(&mut state).unwrap();

    // Recreate the bench with the state restored.
    let (bench, _, _, query) = get_bench();
    let mut simu = bench.restore(&state[..]).unwrap();

    // Verify events have been called in the right order.
    simu.step().unwrap();
    assert_eq!((11 + 19) * 7, simu.process_query(&query, ()).unwrap());
}