membership.rs

use super::equivalence::equivalence_trait::CounterExample;
use crate::{
    ads::{
        traits::{AdaptiveDistinguishingSequence, AdsStatus},
        tree::threaded::Ads as OtADS,
    },
    automatadefs::{
        mealy::{InputSymbol, Mealy, OutputSymbol, State},
        traits::{FiniteStateMachine, ObservationTree},
    },
    learner::{
        apartness::{compute_witness, states_are_apart},
        obs_tree::array_tree::TreeErr,
    },
    sul::{simulator::Simulator, system_under_learning::SystemUnderLearning},
    util::{
        learning_config::{Rule2, Rule3},
        toolbox,
    },
};
use itertools::Itertools;
use rand::prelude::SliceRandom;
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use std::collections::VecDeque;

pub struct Oracle<'a, T> {
    sul: &'a mut Simulator,
    _input_alphabet: Vec<InputSymbol>,
    obs_tree: T,
    rule2: Rule2,
    rule3: Rule3,
}

impl<'a, T> Oracle<'a, T>
where
    T: ObservationTree + Sync + Send,
{
    pub fn new(sul: &'a mut Simulator, rule2: Rule2, rule3: Rule3) -> Self {
        Self {
            _input_alphabet: sul.get_alphabet(),
            obs_tree: T::new(sul.get_alphabet().len()),
            sul,
            rule2,
            rule3,
        }
    }

    pub fn get_counts(&mut self) -> (usize, usize) {
        self.sul.get_counts()
    }

    pub fn borrow_sul(&mut self) -> &mut Simulator {
        self.sul
    }

    /// Immutably borrow the underlying observation tree.
    pub fn borrow_tree(&self) -> &T {
        &self.obs_tree
    }

    /// Mutably borrow the underlying observation tree.
    pub fn borrow_mut_tree(&mut self) -> &mut T {
        &mut self.obs_tree
    }

    /// Identifies state `fs` amongst the basis candidates to just one basis candidate.
    pub fn identify_frontier(&mut self, fs: State, candidates: &mut Vec<State>) {
        candidates.retain(|&b| !states_are_apart(&self.obs_tree, fs, b));
        if candidates.len() < 2 {
            return;
        }
        let mut prefix = self.obs_tree.get_access_seq(fs);
        // let mut candidates = org_candidates;
        let (input_seq, output_seq) = match self.rule3 {
            Rule3::Ads => {
                if candidates.len() == 2 {
                    let mut wit = compute_witness(&self.obs_tree, candidates[0], candidates[1])
                        .expect("No witness found between basis states!");
                    let mut input_seq = prefix;
                    input_seq.append(&mut wit);
                    let output_seq = self.output_query(&input_seq);
                    (input_seq, output_seq)
                } else {
                    let suffix = &mut OtADS::new(&self.obs_tree, candidates);
                    self.adaptive_output_query(&mut prefix, None, suffix)
                }
            }
            Rule3::SepSeq => {
                let mut local_rng = rand::thread_rng();
                let mut basis_pair = candidates.choose_multiple(&mut local_rng, 2).copied();
                let (b1, b2) = (
                    basis_pair.next().expect("Safe"),
                    basis_pair.next().expect("Safe"),
                );
                let mut wit = compute_witness(&self.obs_tree, b1, b2)
                    .expect("No witness found between basis states!");
                let mut input_seq = prefix;
                input_seq.append(&mut wit);
                let output_seq = self.output_query(&input_seq);
                (input_seq, output_seq)
            }
        };
        let _ = self
            .obs_tree
            .insert_observation(None, &input_seq, &output_seq);
        candidates.retain(|&b| !states_are_apart(&self.obs_tree, fs, b));
        if candidates.len() > 1 {
            self.identify_frontier(fs, candidates);
        }
    }

    /// Explores the frontier for all basis states and returns a vector of the frontier states
    /// paired with their corresponding basis candidates.
    #[must_use]
    pub fn explore_frontier(&mut self, basis: &[State]) -> Vec<(State, Vec<State>)> {
        self.obs_tree
            .no_succ_defined(basis, true)
            .iter()
            .inspect(|(q, i)| log::debug!("{:?} has {:?} undefined", *q, *i))
            .map(|(q, i)| self._explore_frontier(*q, *i, basis))
            .inspect(|(q, cands)| log::debug!("Frontier {:?} has candidates {:?}", q, cands))
            .collect()
    }

    /// Explore the frontier for a single (state, input) pair.
    fn _explore_frontier(
        &mut self,
        q: State,
        i: InputSymbol,
        basis: &[State],
    ) -> (State, Vec<State>) {
        let mut access_q = self.obs_tree.get_access_seq(q);
        // let prefix = concat_slices(&[&access_q, &[i]]);
        let (input_seq, output_seq) = match self.rule2 {
            Rule2::Ads => {
                log::debug!("Constructing ADS for {:?} at {:?}", q, i);
                let suffix = &mut OtADS::new(&self.obs_tree, basis);
                self.adaptive_output_query(&mut access_q, Some(i), suffix)
            }
            Rule2::Nothing => {
                let mut prefix = access_q;
                prefix.push(i);
                let o_seq = self.output_query(&prefix);
                (prefix, o_seq)
            }
            Rule2::SepSeq => {
                let mut local_rng = rand::thread_rng();
                let mut wit = {
                    if basis.len() > 1 {
                        let (b1, b2) = basis
                            .choose_multiple(&mut local_rng, 2)
                            .copied()
                            .collect_tuple()
                            .expect("Safe");
                        compute_witness(&self.obs_tree, b1, b2).expect("Safe")
                    } else {
                        Vec::new()
                    }
                };
                let mut input_seq = access_q;
                input_seq.push(i);
                input_seq.append(&mut wit);
                let output_seq = self.output_query(&input_seq);
                (input_seq, output_seq)
            }
        };
        let _ = self
            .obs_tree
            .insert_observation(None, &input_seq, &output_seq);
        let fs = self.obs_tree.get_succ(q, &[i]).expect("Safe");
        log::debug!("Added frontier {:?}.", fs);
        let bs_not_sep = basis
            .par_iter()
            .copied()
            .filter(|&b| !states_are_apart(&self.obs_tree, fs, b))
            .collect();
        (fs, bs_not_sep)
    }

    /// Make a non-adaptive output query.
    pub fn output_query(&mut self, input_seq: &[InputSymbol]) -> Vec<OutputSymbol> {
        let out_seq = self
            .obs_tree
            .get_observation(None, input_seq)
            .unwrap_or_else(|| self.sul.trace(input_seq).to_vec());
        let _ = self.obs_tree.insert_observation(None, input_seq, &out_seq);
        out_seq.to_vec()
    }

    /// Make an adaptive output query of the form (prefix -> infix -> ADS).
    pub(crate) fn adaptive_output_query<ADS: AdaptiveDistinguishingSequence>(
        &mut self,
        prefix: &mut Vec<InputSymbol>,
        infix: Option<InputSymbol>,
        suffix: &mut ADS,
    ) -> (Vec<InputSymbol>, Vec<OutputSymbol>) {
        if let Some(i) = infix {
            prefix.push(i);
        }
        self._adaptive_output_query(prefix, suffix)
    }

    /// Make an adaptive query, where the infix has been moved into the prefix.
    fn _adaptive_output_query<ADS: AdaptiveDistinguishingSequence>(
        &mut self,
        prefix: &mut Vec<InputSymbol>,
        suffix: &mut ADS,
    ) -> (Vec<InputSymbol>, Vec<OutputSymbol>) {
        let tree_reply = self
            .obs_tree
            .get_succ(State::new(0), prefix)
            .and_then(|curr_state| self.answer_ads_from_tree(suffix, curr_state).ok());
        suffix.reset_to_root();
        if tree_reply.is_some() {
            unreachable!("ADS is not increasing the norm, we already knew this information.");
        }
        self.sul.reset();
        let prefix_out = self.sul.step(prefix);
        let (mut suffix_inputs, suffix_outputs) = self.sul_adaptive_query(suffix);
        let input_seq = prefix;
        input_seq.append(&mut suffix_inputs); //concat_slices(&[prefix, &suffix_inputs]);
        let output_seq = toolbox::concat_slices(&[&prefix_out, &suffix_outputs]);
        let _ = self.add_observation(input_seq, &output_seq);
        (input_seq.to_vec(), output_seq)
    }

    // Assuming the prefix has been sent to the SUL, perform the adaptive query.
    fn sul_adaptive_query<ADS: AdaptiveDistinguishingSequence>(
        &mut self,
        ads: &mut ADS,
    ) -> (Vec<InputSymbol>, Box<[OutputSymbol]>) {
        let mut inputs_sent = Vec::new();
        let mut outputs_received = Vec::new();
        let mut last_output = None;
        loop {
            let next_input = ads.next_input(last_output);
            if let Ok(next_input) = next_input {
                log::debug!("Next input: {:?}", next_input);
                inputs_sent.push(next_input);
                let o = self.sul.step(&[next_input])[0];
                last_output = Some(o);
                outputs_received.push(o);
            } else {
                log::debug!("Next input undefined.");
                break;
            }
        }
        (inputs_sent, outputs_received.into_boxed_slice())
    }

    pub fn ads_equiv_test(
        &mut self,
        ads: &mut super::equivalence::incomplete::impl_ads::AdsTree,
        prefix: &[InputSymbol],
        fsm: &Mealy,
    ) -> CounterExample {
        if let Some(tree_prefix_out) = self.obs_tree.get_observation(None, prefix) {
            let hyp_out = fsm.trace(prefix).1.to_vec();
            if tree_prefix_out != hyp_out {
                return Some((prefix.to_vec(), tree_prefix_out));
            }
            let ts = self.obs_tree.get_succ(State::new(0), prefix).expect("Safe");
            assert!(ads.next_input(None).is_ok());
            let ads_res = self.answer_ads_from_tree(ads, ts);
            match ads_res {
                Ok((ads_inputs, ads_outputs)) => {
                    println!("Tree had a reply, duplicate query!");
                    // assert!(self.check_consistency(fsm).is_none());
                    // println!("ADS: {:#?}", ads);
                    let mut input_seq = Vec::from(prefix);
                    input_seq.extend(ads_inputs.iter());
                    let mut output_seq = tree_prefix_out;
                    output_seq.extend(ads_outputs.iter());
                    let hyp_out = fsm.trace(&input_seq).1;
                    if output_seq != hyp_out.to_vec() {
                        return Some((input_seq, output_seq));
                    }
                }
                Err(_) => ads.reset_to_root(),
            }
        }

        self.sul.reset();
        let prefix_out = self.sul.step(prefix);
        let (mut hyp_curr, hyp_out) = fsm.trace(prefix);
        let hyp_prefix_state = hyp_curr;
        if prefix_out != hyp_out {
            log::info!("Found CE in the random prefix.");
            return Some((prefix.to_vec(), prefix_out.to_vec()));
        }
        let mut inputs_sent = Vec::from_iter(prefix.iter().copied());
        let mut outputs_received = Vec::from_iter(prefix_out.iter().copied());
        let mut hyp_out;
        let mut prev_output = None;
        loop {
            let next_input = ads.next_input(prev_output);
            match next_input {
                Ok(input) => {
                    log::debug!("Input sent: {:?}", input);
                    inputs_sent.push(input);
                    let o = self.sul.step(&[input])[0];
                    (hyp_curr, hyp_out) = fsm.step_from(hyp_curr, input);
                    log::debug!("Output Received: {:?}", o);
                    prev_output = Some(o);
                    outputs_received.push(o);
                    if o != hyp_out {
                        self.add_observation(&inputs_sent, &outputs_received);
                        return Some((inputs_sent, outputs_received));
                    }
                }
                Err(ads_err) => match ads_err {
                    AdsStatus::Done => {
                        // if !ads
                        //     .tree
                        //     .ref_at(ads.curr_idx)
                        //     .initial
                        //     .contains(&hyp_prefix_state)
                        // {
                        //     panic!("Probably not the state we wanted!");
                        // }
                        self.add_observation(&inputs_sent, &outputs_received);
                        return None;
                    }
                    AdsStatus::Unexpected => {
                        unreachable!("IADS root must always contain the hyp state to identify.");
                    }
                },
            }
        }
    }

    pub fn check_consistency(&self, hypothesis: &Mealy) -> CounterExample {
        let mut queue = VecDeque::from([(State::new(0), State::new(0))]);
        while !queue.is_empty() {
            let (q, r) = queue.pop_front().expect("Always safe!");
            for i in (0..hypothesis.input_alphabet().len())
                .into_iter()
                .map(|x| x as u16)
                .map(InputSymbol::from)
            {
                if let Some(((out_tree, dest_tree), (dest_hyp, out_hyp))) = self
                    .obs_tree
                    .get_out_succ(q, i)
                    .zip(Some(hypothesis.step_from(r, i)))
                {
                    if out_hyp == out_tree {
                        queue.push_back((dest_tree, dest_hyp));
                    } else {
                        let mut inputs = self.obs_tree.get_access_seq(q);
                        inputs.push(i);
                        let outputs = self.obs_tree.get_observation(None, &inputs).expect("Safe");
                        return Some((inputs, outputs));
                    }
                }
            }
        }
        None
    }

    #[allow(clippy::type_complexity)]
    fn answer_ads_from_tree(
        &self,
        ads: &mut impl AdaptiveDistinguishingSequence,
        from_state: State,
    ) -> Result<(Box<[InputSymbol]>, Box<[OutputSymbol]>), TreeErr> {
        let mut prev_output = None;
        let mut outputs_received = vec![];
        let mut inputs_sent = vec![];
        let mut curr_state = from_state;
        loop {
            let next_input = ads.next_input(prev_output);
            if let Ok(next_input) = next_input {
                inputs_sent.push(next_input);
                let (output, dest) = self
                    .obs_tree
                    .get_out_succ(curr_state, next_input)
                    .ok_or(TreeErr::AbsentEntry)?;
                prev_output = Some(output);
                outputs_received.push(output);
                curr_state = dest;
            } else {
                ads.reset_to_root();
                return Ok((
                    inputs_sent.into_boxed_slice(),
                    outputs_received.into_boxed_slice(),
                ));
            }
        }
    }

    pub fn add_observation(
        &mut self,
        input_seq: &[InputSymbol],
        output_seq: &[OutputSymbol],
    ) -> State {
        self.obs_tree
            .insert_observation(None, input_seq, output_seq)
    }
}