/*
    Appellation: neural_engine <example>
    Contrib: @FL03
*/
extern crate eryon_core as eryon;

use eryon::rstm::{RuleSet, State, ruleset};
use eryon_actors::NeuralEngine;
use eryon_mem::TopoLedger;

fn main() -> anyhow::Result<()> {
    let ruleset: RuleSet<usize, usize> = ruleset! {
        (0, 0) -> Right(0, 1),
        (0, 1) -> Right(1, 2),
        (0, 2) -> Right(0, 0),
        (1, 0) -> Right(1, 0),
        (1, 1) -> Left(0, 1),
        (1, 2) -> Right(1, 1),
    };
    let mut rules = ruleset
        .iter()
        .map(|rule| eryon::LearnedRule::from_rule(*rule, 1f32))
        .collect::<Vec<_>>();
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::DEBUG)
        .init();

    tracing::info!("Initializing Wolfram [2,3] NTM implementation...");

    // Create NTM with the 3-element alphabet [0,1,2]
    let alphabet = [0, 1, 2];
    let mut ntm = NeuralEngine::<f32>::new(alphabet, State(0)).init();

    // Create a topological memory
    let mut memory = TopoLedger::new();

    // Add rules that mimic a binary counter (a classic Wolfram UTM behavior)
    memory.rulespace_mut().append(&mut rules);

    // Initialize tape with 0s (starting a binary counter)
    ntm.tape_mut().extend_from_slice(&[0, 0, 0, 0, 0, 0]);

    tracing::info!(
        "Initial configuration:\nState: {:?}\nTape: {:?}\nPosition: {}",
        ntm.state(),
        ntm.tape(),
        ntm.position()
    );

    // Run the machine with learning from the rule set
    tracing::info!("\nRunning with learned rules...");
    for i in 0..15 {
        let current_state = ntm.state();
        let current_symbol = if ntm.position() < ntm.tape().len() {
            ntm.tape()[ntm.position()]
        } else {
            0
        };

        tracing::info!(
            "Step {}: State={}, Symbol={}, Position={}",
            i,
            current_state,
            current_symbol,
            ntm.position()
        );

        ntm.execute_with_stored(&memory)?;

        // Print tape after each step to visualize binary counting
        let tape_str = ntm
            .tape()
            .iter()
            .enumerate()
            .map(|(i, &s)| {
                if i == ntm.position() {
                    format!("[{}]", s)
                } else {
                    format!(" {} ", s)
                }
            })
            .collect::<Vec<_>>()
            .join("");

        tracing::info!("Tape: {}", tape_str);
    }

    // Now test if NTM has learned the pattern by training
    tracing::info!("Training the NTM to learn the binary counter pattern...");

    // Reset the NTM
    ntm = NeuralEngine::new(alphabet, State(0)).init();

    // Set up tape
    ntm.tape_mut().extend_from_slice(&[0, 0, 0, 0, 0, 0]);

    // Train by running with explicit rules multiple times
    for _ in 0..10 {
        for _ in 0..15 {
            ntm.execute_with_stored(&memory)?;
        }

        // Reset position but keep weights
        ntm.reset_position();
        ntm.set_state(State(0));
        ntm.tape_mut().fill(0);
    }

    // Test execution without rules
    tracing::info!("Testing what the NTM has learned (execution without rules)...");
    for i in 0..15 {
        let current_state = ntm.state();
        let current_symbol = if ntm.position() < ntm.tape().len() {
            ntm.tape()[ntm.position()]
        } else {
            0
        };

        tracing::info!(
            "Step {}: State={}, Symbol={}, Position={}",
            i,
            current_state,
            current_symbol,
            ntm.position()
        );

        // Use neural weights only (step instead of execute_with_learning)
        ntm.step()?;

        // Print tape
        let tape_str = ntm
            .tape()
            .iter()
            .enumerate()
            .map(|(i, &s)| {
                if i == ntm.position() {
                    format!("[{}]", s)
                } else {
                    format!(" {} ", s)
                }
            })
            .collect::<Vec<_>>()
            .join("");

        tracing::info!("Tape: {}", tape_str);
    }

    Ok(())
}
