/*
    Appellation: wolfram <module>
    Contrib: @FL03
*/
use super::{ComputationalEngine, Engine, RawEngine};
use crate::types::{Rulespace, Snapshot};

use eryon::traits::Symbolic;
use eryon::{Direction, Head, RawState, State, Tail};

/// An implementation of the Wolfram [2, 3] UTM that consideres two states and three symbols;
/// the machine has proven to be capable of universal computation
#[derive(Clone, Debug, Default, PartialEq)]
pub struct WolframUTM<Q = usize, A = usize>
where
    Q: RawState + Eq + core::hash::Hash,
    A: Symbolic,
{
    pub(crate) alphabet: [A; 3],
    pub(crate) execution_history: Vec<(Head<Q, A>, Tail<Q, A>)>,
    pub(crate) position: usize,
    pub(crate) ruleset: Rulespace<Q, A>,
    pub(crate) state: State<Q>,
    pub(crate) tape: Vec<A>,
}

impl<Q, A> WolframUTM<Q, A>
where
    Q: RawState + Eq + core::hash::Hash,
    A: Symbolic,
{
    pub const STATES: usize = 2;
    pub const SYMBOLS: usize = 3;
    // Create a new UTM with a specific program
    pub fn new(alphabet: [A; 3], State(state): State<Q>) -> Self {
        WolframUTM {
            alphabet,
            execution_history: Vec::new(),
            position: 0,
            state: State(state),
            tape: Vec::new(),
            ruleset: Rulespace::new(),
        }
    }
    /// create a new instance with the given alphabet and default state
    pub fn from_alphabet(alphabet: [A; 3]) -> Self
    where
        Q: Default,
    {
        Self::new(alphabet, State::default())
    }
    /// returns an immutable reference to the machine's alphabet
    pub const fn alphabet(&self) -> &[A; 3] {
        &self.alphabet
    }
    /// returns a copy of the machines current position
    #[inline]
    pub fn position(&self) -> usize {
        self.position
    }
    /// returns a reference to the machine's ruleset
    pub const fn ruleset(&self) -> &Rulespace<Q, A> {
        &self.ruleset
    }
    /// returns a mutable reference to the machine's ruleset
    pub fn ruleset_mut(&mut self) -> &mut Rulespace<Q, A> {
        &mut self.ruleset
    }
    /// returns a copy of the machine's current state
    #[inline]
    pub fn state(&self) -> State<&Q> {
        self.state.view()
    }
    /// returns a mutable reference to the machine's state
    pub fn state_mut(&mut self) -> State<&mut Q> {
        self.state.view_mut()
    }
    /// returns an immutable reference to the tape
    pub const fn tape(&self) -> &Vec<A> {
        &self.tape
    }
    /// returns a mutable reference to the tape
    pub fn tape_mut(&mut self) -> &mut Vec<A> {
        &mut self.tape
    }
    /// returns the head of the machine
    pub fn head(&self) -> Head<&Q, &A> {
        // self.get_current_symbol().cloned().unwrap_or_default()
        Head::new(self.state(), &self.tape[self.position])
    }
    /// check if a symbol is in the alphabet
    pub fn contains(&self, symbol: &A) -> bool {
        self.alphabet.contains(symbol)
    }
    /// get the current symbol on the tape
    pub fn get_current_symbol(&self) -> Option<&A> {
        self.tape.get(self.position)
    }
    /// get the tail for the current head
    pub fn get_tail(&self) -> Option<&Tail<Q, A>>
    where
        Q: Clone,
    {
        self.ruleset.get(&self.head().cloned())
    }
    /// get the tail for a given head
    pub fn get_tail_for(&self, head: Head<Q, A>) -> &Tail<Q, A> {
        &self.ruleset[&head]
    }
    /// set the alphabet of the machine
    pub fn set_alphabet(&mut self, alphabet: [A; 3]) {
        self.alphabet = alphabet;
    }
    /// set the position of the head of the machine
    pub fn set_position(&mut self, position: usize) -> &mut Self {
        self.position = position;
        self
    }
    /// set the ruleset of the machine
    pub fn set_ruleset(&mut self, ruleset: Rulespace<Q, A>) -> &mut Self {
        self.ruleset = ruleset;
        self
    }
    /// set the state of the machine
    pub fn set_state(&mut self, State(state): State<Q>) -> &mut Self {
        self.state = State(state);
        self
    }
    /// set the tape of the machine
    pub fn set_tape(&mut self, tape: Vec<A>) -> &mut Self {
        self.tape = tape;
        self
    }
    /// consumes this instance to create another with the given alphabet
    pub fn with_alphabet(self, alphabet: [A; 3]) -> Self {
        Self { alphabet, ..self }
    }
    /// consumes this instance to create another with the given ruleset
    pub fn with_ruleset(self, ruleset: Rulespace<Q, A>) -> Self {
        Self { ruleset, ..self }
    }
    /// consumes this instance to create another with the given state
    pub fn with_state(self, State(state): State<Q>) -> Self {
        Self {
            state: State(state),
            ..self
        }
    }
    /// consumes this instance to create another with the given tape
    pub fn with_tape(self, tape: Vec<A>) -> Self {
        Self {
            position: 0,
            tape,
            ..self
        }
    }
    /// reset the position of the machine's head to `0`
    pub fn reset_position(&mut self) -> &mut Self {
        self.set_position(0);
        self
    }
    /// use the given input as the tape and reset the position of the machine's head
    pub fn reset_with_tape(&mut self, tape: Vec<A>) -> &mut Self {
        self.set_tape(tape).reset_position()
    }
    /// process some input
    pub fn run(&mut self, input: Vec<A>) -> crate::Result<Vec<Snapshot<Q, [A; 3], Vec<A>>>>
    where
        Q: Clone,
    {
        self.run_for(input, usize::MAX)
    }
    /// process some input
    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(
            fields(max),
            skip(self, input),
            level = "trace",
            name = "run_for",
            target = "utm"
        )
    )]
    pub fn run_for(
        &mut self,
        input: Vec<A>,
        max: usize,
    ) -> crate::Result<Vec<Snapshot<Q, [A; 3], Vec<A>>>>
    where
        Q: Clone,
    {
        // Initialize the tape with the input
        self.reset_with_tape(input);

        let mut history = vec![self.snapshot().cloned()];

        // Run until the machine halts (i.e., step() returns false)
        let mut halt = false;
        while self.step().is_ok() && !halt {
            // Record each state after a step
            history.push(self.snapshot().cloned());
            // prevent infinite loops; halt if the history exceeds the maximum
            if history.len() > max {
                halt = true;
                #[cfg(feature = "tracing")]
                tracing::warn!(
                    "Infinite loop detected; check the system and it's ruleset to ensure convergence within {} steps",
                    max
                )
            }
        }

        Ok(history)
    }

    /// execute one step of the machine
    pub fn step(&mut self) -> crate::Result<()>
    where
        Q: Clone,
    {
        let symbol = self.get_current_symbol().cloned().unwrap_or_default();

        // Check if the current symbol is in our current triad context
        if !self.contains(&symbol) {
            panic!("symbol {} not in any accessible triad context", symbol);
        }
        let head = Head::new(self.state.clone(), symbol);
        // Find the transition rule for the current state and symbol
        self.handle(&head)
    }

    pub fn write(&mut self, symbol: A) {
        self.tape[self.position] = symbol;
    }
    /// given the head of the machine, find the corresponding tail and update the machine
    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(
            skip(self, head),
            level = "trace",
            name = "handle",
            target = "wolfram"
        )
    )]
    fn handle(&mut self, head: &Head<Q, A>) -> crate::Result<()>
    where
        Q: Clone,
    {
        if let Some(tail) = self.ruleset.get(head) {
            let Tail {
                state: new_state,
                symbol: new_symbol,
                direction,
            } = tail.clone();
            // Update the tape
            if self.position >= self.tape.len() {
                self.tape.resize(self.position + 1, A::default());
            }
            self.tape[self.position] = new_symbol;

            // Update state
            self.state = new_state;

            // Move the head
            match direction {
                Direction::Left => {
                    if self.position > 0 {
                        self.position -= 1;
                    } else {
                        self.tape.insert(0, A::default());
                    }
                }
                Direction::Right => {
                    self.position += 1;
                    if self.position >= self.tape.len() {
                        self.tape.push(A::default());
                    }
                }
                Direction::Stay => {}
            }
            return Ok(());
        }
        Err(crate::ActorError::DriverError(
            "No rule found for the current head...".to_string(),
        )) // No transition rule found
    }
}

use crate::mem::TopoLedger;

impl WolframUTM {
    /// returns a string representation of the UTM
    pub fn pretty_print(&self) -> String {
        let mut result = String::new();
        // insert the state and alphabet
        result.push_str(&format!(
            "(State: {:?}, Alphabet: {:?})",
            self.state, self.alphabet
        ));
        // Print the tape with the current position marked
        for (i, &symbol) in self.tape.iter().enumerate() {
            if i == self.position {
                result.push_str(&format!("[{}]", symbol));
            } else {
                result.push_str(&format!(" {} ", symbol));
            }
        }
        // return the result
        result
    }
    /// Update with a learned rule
    pub fn update_rule(&mut self, head: Head<usize, usize>, tail: Tail<usize, usize>) {
        // Store the rule in the ruleset
        let entry = self.ruleset.entry(head).or_default();
        *entry = tail;
    }

    /// Learn from a successful execution sequence
    pub fn learn_from_execution(&mut self, memory: &mut TopoLedger, steps: usize, reward: f32) {
        // Record the last 'steps' state transitions
        if self.execution_history.len() < steps {
            return;
        }

        // Calculate reinforcement amount based on reward
        let reinforcement = reward / steps as f32;

        // Get the most recent steps
        let history = self.execution_history.clone();
        let relevant_history = history
            .iter()
            .rev()
            .take(steps)
            .collect::<Vec<_>>()
            .into_iter()
            .rev();

        // Reinforce each rule in the sequence
        for &(head, tail) in relevant_history {
            let Head { state, symbol } = head;
            let Tail {
                state: next_state,
                symbol: next_symbol,
                direction,
            } = tail;
            // Store or update the rule in memory
            memory.store_learned_rule(
                state,
                symbol,
                next_state,
                next_symbol,
                direction,
                reinforcement,
            );

            // Also update the UTM's rules directly
            self.update_rule(head, tail);
        }
    }
}

impl<Q, S> RawEngine for WolframUTM<Q, S>
where
    Q: RawState + Clone + Eq + core::hash::Hash,
    S: Symbolic,
{
    type Store = Vec<S>;

    seal!();
}

impl<Q, S> Engine for WolframUTM<Q, S>
where
    Q: RawState + Clone + Default + Eq + core::hash::Hash + Send + Sync,
    S: Symbolic,
{
    seal!();
}

impl<Q, S> ComputationalEngine<Q, [S; 3]> for WolframUTM<Q, S>
where
    Q: RawState + Clone + Eq + core::hash::Hash,
    S: Symbolic,
{
    fn new(alphabet: [S; 3], initial_state: State<Q>) -> Self {
        WolframUTM::new(alphabet, initial_state)
    }

    fn alphabet(&self) -> &[S; 3] {
        &self.alphabet
    }

    fn alphabet_mut(&mut self) -> &mut [S; 3] {
        &mut self.alphabet
    }

    fn head(&self) -> Head<&Q, &S> {
        self.head()
    }

    fn head_mut(&mut self) -> Head<&mut Q, &mut S> {
        Head::new(self.state.view_mut(), &mut self.tape[self.position])
    }

    fn position(&self) -> usize {
        self.position
    }

    fn position_mut(&mut self) -> &mut usize {
        &mut self.position
    }

    fn state(&self) -> State<&Q> {
        self.state.view()
    }

    fn state_mut(&mut self) -> State<&mut Q> {
        self.state.view_mut()
    }

    fn tape(&self) -> &Vec<S> {
        &self.tape
    }

    fn tape_mut(&mut self) -> &mut Vec<S> {
        &mut self.tape
    }

    fn step(&mut self) -> Result<(), crate::ActorError> {
        self.step()
    }

    fn set_alphabet(&mut self, alphabet: [S; 3]) {
        self.alphabet = alphabet;
    }

    fn set_position(&mut self, position: usize) {
        self.position = position;
    }

    fn set_state(&mut self, state: State<Q>) {
        self.state = state;
    }

    fn set_tape<I>(&mut self, tape: I)
    where
        I: IntoIterator<Item = S>,
    {
        self.tape = Vec::from_iter(tape);
    }
}

impl<Q, S> core::fmt::Display for WolframUTM<Q, S>
where
    Q: RawState + Eq + core::hash::Hash,
    S: Symbolic,
{
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(
            f,
            "{head:?} @ {pos}: {tape:?}",
            head = self.head(),
            pos = self.position(),
            tape = self.tape()
        )
    }
}

impl<Q, S> core::ops::Index<usize> for WolframUTM<Q, S>
where
    Q: RawState + Eq + core::hash::Hash,
    S: Symbolic,
{
    type Output = S;

    fn index(&self, index: usize) -> &Self::Output {
        &self.tape[index]
    }
}

impl<Q, S> core::ops::Index<Head<Q, S>> for WolframUTM<Q, S>
where
    Q: RawState + Eq + core::hash::Hash,
    S: Symbolic,
{
    type Output = Tail<Q, S>;

    fn index(&self, index: Head<Q, S>) -> &Self::Output {
        &self.ruleset[&index]
    }
}
