oxilean-std 0.1.2

//! # Process Algebra — Functions and Environment Builder
//!
//! Algorithms for process algebra: LTS generation from CCS, bisimulation checking,
//! HML model checking, trace computation, and the Lean4-kernel environment builder.

use std::collections::{HashMap, HashSet, VecDeque};

use oxilean_kernel::{Declaration, Environment, Expr, Level, Name};

use super::types::{
    Action, BehavioralEquivalence, BisimulationRelation, CcsProcess, CspProcess, Failure,
    FailuresModel, HmlFormula, Lts, State, StructuralCongruenceClass, TestOutcome, Trace, TraceSet,
    Transition,
};

// ─── Kernel Expression Helpers ────────────────────────────────────────────────

fn cst(s: &str) -> Expr {
    Expr::Const(Name::str(s), vec![])
}

fn type0() -> Expr {
    Expr::Sort(Level::succ(Level::zero()))
}

fn prop() -> Expr {
    Expr::Sort(Level::zero())
}

fn arrow(a: Expr, b: Expr) -> Expr {
    Expr::Pi(
        oxilean_kernel::BinderInfo::Default,
        Name::str("_"),
        Box::new(a),
        Box::new(b),
    )
}

fn app(f: Expr, a: Expr) -> Expr {
    Expr::App(Box::new(f), Box::new(a))
}

fn nat_ty() -> Expr {
    cst("Nat")
}

fn bool_ty() -> Expr {
    cst("Bool")
}

// ─── Lean4-Type Declarations ──────────────────────────────────────────────────

/// `Action : Type` — process algebra action (visible or τ).
pub fn action_ty() -> Expr {
    type0()
}

/// `CcsProcess : Type` — a CCS process term.
pub fn ccs_process_ty() -> Expr {
    type0()
}

/// `CspProcess : Type` — a CSP process term.
pub fn csp_process_ty() -> Expr {
    type0()
}

/// `Lts : Type` — a labeled transition system.
pub fn lts_ty() -> Expr {
    type0()
}

/// `Bisimilar : CcsProcess → CcsProcess → Prop` — strong bisimilarity.
pub fn bisimilar_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(ccs_process_ty(), prop()))
}

/// `WeakBisimilar : CcsProcess → CcsProcess → Prop` — weak bisimilarity.
pub fn weak_bisimilar_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(ccs_process_ty(), prop()))
}

/// `HmlFormula : Type` — Hennessy-Milner logic formula.
pub fn hml_formula_ty() -> Expr {
    type0()
}

/// `hml_satisfies : CcsProcess → HmlFormula → Prop` — HML satisfaction.
pub fn hml_satisfies_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(hml_formula_ty(), prop()))
}

/// `TraceSet : Type` — the set of traces of a process.
pub fn trace_set_ty() -> Expr {
    type0()
}

/// `FailuresModel : Type` — the failures model of a process.
pub fn failures_model_ty() -> Expr {
    type0()
}

// ─── Axioms and Theorems ──────────────────────────────────────────────────────

/// `bisim_reflexive : ∀ P, P ~ P` — reflexivity of bisimilarity.
pub fn bisim_reflexive_ty() -> Expr {
    arrow(ccs_process_ty(), prop())
}

/// `bisim_symmetric : ∀ P Q, P ~ Q → Q ~ P` — symmetry.
pub fn bisim_symmetric_ty() -> Expr {
    arrow(
        ccs_process_ty(),
        arrow(ccs_process_ty(), arrow(prop(), prop())),
    )
}

/// `bisim_transitive : ∀ P Q R, P ~ Q → Q ~ R → P ~ R` — transitivity.
pub fn bisim_transitive_ty() -> Expr {
    arrow(
        ccs_process_ty(),
        arrow(
            ccs_process_ty(),
            arrow(ccs_process_ty(), arrow(prop(), arrow(prop(), prop()))),
        ),
    )
}

/// `bisim_congruence : bisimilarity is a congruence for all CCS operators`.
pub fn bisim_congruence_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(ccs_process_ty(), prop()))
}

/// `hml_bisim_characterization : P ~ Q ↔ ∀ φ : HML, P ⊨ φ ↔ Q ⊨ φ`
/// (Hennessy-Milner theorem for image-finite processes).
pub fn hml_bisim_characterization_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(ccs_process_ty(), prop()))
}

/// `tau_law_1 : τ.P ~ P` — a process that can only do τ then P is equivalent to P.
pub fn tau_law_1_ty() -> Expr {
    arrow(ccs_process_ty(), prop())
}

/// `tau_law_2 : α.(τ.P + Q) ~ α.(P + Q)` — absorption of τ.
pub fn tau_law_2_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(ccs_process_ty(), prop()))
}

/// `expansion_theorem : P | Q ~ Σ α.P'|Q when P →α P', + Σ τ.(P'|Q') when sync`.
pub fn expansion_theorem_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(ccs_process_ty(), prop()))
}

/// `trace_preorder_from_bisim : P ~ Q → traces(P) = traces(Q)`.
pub fn trace_preorder_from_bisim_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(ccs_process_ty(), prop()))
}

/// `failures_from_bisim : P ~ Q → failures(P) = failures(Q)`.
pub fn failures_from_bisim_ty() -> Expr {
    arrow(ccs_process_ty(), arrow(ccs_process_ty(), prop()))
}

// ─── Environment Builder ──────────────────────────────────────────────────────

/// Build the process algebra environment with all type and axiom declarations.
pub fn build_process_algebra_env(env: &mut Environment) {
    let axioms: &[(&str, Expr)] = &[
        ("Action", action_ty()),
        ("CcsProcess", ccs_process_ty()),
        ("CspProcess", csp_process_ty()),
        ("Lts", lts_ty()),
        ("Bisimilar", bisimilar_ty()),
        ("WeakBisimilar", weak_bisimilar_ty()),
        ("HmlFormula", hml_formula_ty()),
        ("HmlSatisfies", hml_satisfies_ty()),
        ("TraceSet", trace_set_ty()),
        ("FailuresModel", failures_model_ty()),
        ("BisimReflexive", bisim_reflexive_ty()),
        ("BisimSymmetric", bisim_symmetric_ty()),
        ("BisimTransitive", bisim_transitive_ty()),
        ("BisimCongruence", bisim_congruence_ty()),
        ("HmlBisimCharacterization", hml_bisim_characterization_ty()),
        ("TauLaw1", tau_law_1_ty()),
        ("TauLaw2", tau_law_2_ty()),
        ("ExpansionTheorem", expansion_theorem_ty()),
        ("TracePreorderFromBisim", trace_preorder_from_bisim_ty()),
        ("FailuresFromBisim", failures_from_bisim_ty()),
    ];
    for (name, ty) in axioms {
        env.add(Declaration::Axiom {
            name: Name::str(*name),
            univ_params: vec![],
            ty: ty.clone(),
        })
        .ok();
    }
}

// ─── LTS Generation from CCS ─────────────────────────────────────────────────

/// Generate the LTS of a CCS process (up to `max_states` states, `max_unfolds` recursion depth).
///
/// Uses a BFS exploration: each state is a CCS process term (up to structural equivalence).
/// Returns `None` if the state space exceeds `max_states`.
pub fn generate_lts(
    initial: &CcsProcess,
    max_states: usize,
    max_unfolds: usize,
) -> Option<(Lts, Vec<CcsProcess>)> {
    let mut state_map: HashMap<String, State> = HashMap::new();
    let mut state_procs: Vec<CcsProcess> = Vec::new();
    let mut transitions: Vec<Transition> = Vec::new();
    let mut queue: VecDeque<(State, CcsProcess, usize)> = VecDeque::new();

    let init_key = format!("{}", initial);
    state_map.insert(init_key, 0);
    state_procs.push(initial.clone());
    queue.push_back((0, initial.clone(), 0));

    while let Some((state_id, proc, depth)) = queue.pop_front() {
        if state_procs.len() > max_states {
            return None;
        }

        // Compute one-step transitions
        let steps = one_step_ccs(&proc, depth, max_unfolds);
        for (action, next_proc) in steps {
            let next_key = format!("{}", next_proc);
            let next_id = if let Some(&id) = state_map.get(&next_key) {
                id
            } else {
                let id = state_procs.len();
                state_map.insert(next_key, id);
                state_procs.push(next_proc.clone());
                queue.push_back((id, next_proc, depth + 1));
                id
            };
            transitions.push(Transition::new(state_id, action, next_id));
        }
    }

    let lts = Lts::new(state_procs.len(), 0, transitions);
    Some((lts, state_procs))
}

/// Compute all immediate (one-step) transitions of a CCS process.
pub fn one_step_ccs(
    proc: &CcsProcess,
    depth: usize,
    max_unfolds: usize,
) -> Vec<(Action, CcsProcess)> {
    match proc {
        CcsProcess::Nil => vec![],

        CcsProcess::Prefix(a, p) => {
            vec![(a.clone(), *p.clone())]
        }

        CcsProcess::Choice(p, q) => {
            let mut steps = one_step_ccs(p, depth, max_unfolds);
            steps.extend(one_step_ccs(q, depth, max_unfolds));
            steps
        }

        CcsProcess::Parallel(p, q) => {
            let mut steps = Vec::new();

            // P moves independently
            for (a, p_prime) in one_step_ccs(p, depth, max_unfolds) {
                steps.push((
                    a.clone(),
                    CcsProcess::parallel(*p.clone(), q.as_ref().clone()),
                ));
                // Actually we want (a, p' | q)
                steps.pop();
                steps.push((a, CcsProcess::parallel(p_prime, *q.clone())));
            }

            // Q moves independently
            for (a, q_prime) in one_step_ccs(q, depth, max_unfolds) {
                steps.push((a, CcsProcess::parallel(*p.clone(), q_prime)));
            }

            // Synchronization: P does ᾱ and Q does α (or vice versa) → τ
            let p_steps = one_step_ccs(p, depth, max_unfolds);
            let q_steps = one_step_ccs(q, depth, max_unfolds);
            for (a1, p_prime) in &p_steps {
                for (a2, q_prime) in &q_steps {
                    if *a1 == a2.complement() && a1.is_visible() {
                        steps.push((
                            Action::Tau,
                            CcsProcess::parallel(p_prime.clone(), q_prime.clone()),
                        ));
                    }
                }
            }

            steps
        }

        CcsProcess::Restriction(p, labels) => {
            one_step_ccs(p, depth, max_unfolds)
                .into_iter()
                .filter(|(a, _)| {
                    // Keep only actions not in restriction set
                    match a {
                        Action::Tau => true,
                        Action::Input(s) | Action::Output(s) => !labels.contains(s),
                    }
                })
                .map(|(a, p_prime)| {
                    (
                        a,
                        CcsProcess::Restriction(Box::new(p_prime), labels.clone()),
                    )
                })
                .collect()
        }

        CcsProcess::Relabeling(p, map) => one_step_ccs(p, depth, max_unfolds)
            .into_iter()
            .map(|(a, p_prime)| {
                let new_a = apply_relabeling(&a, map);
                (
                    new_a,
                    CcsProcess::Relabeling(Box::new(p_prime), map.clone()),
                )
            })
            .collect(),

        CcsProcess::Var(_) => vec![], // Unbound variable: no transitions

        CcsProcess::Rec(x, body) => {
            if depth >= max_unfolds {
                return vec![]; // Guard against infinite unfolding
            }
            let unfolded = body.substitute(x, proc);
            one_step_ccs(&unfolded, depth + 1, max_unfolds)
        }
    }
}

/// Apply a relabeling function to an action.
fn apply_relabeling(a: &Action, map: &HashMap<String, String>) -> Action {
    match a {
        Action::Input(s) => {
            let new_s = map.get(s).cloned().unwrap_or_else(|| s.clone());
            Action::Input(new_s)
        }
        Action::Output(s) => {
            let new_s = map.get(s).cloned().unwrap_or_else(|| s.clone());
            Action::Output(new_s)
        }
        Action::Tau => Action::Tau,
    }
}

// ─── Strong Bisimulation ─────────────────────────────────────────────────────

/// Check **strong bisimilarity** between two states of an LTS using
/// the **signature-based partition refinement** algorithm.
///
/// Each state's signature is `(current_block, sorted_multiset_of_(action_idx, succ_block))`.
/// Two states are in the same block iff they have the same signature.
/// Iterate until stable. Guaranteed to converge in ≤ n iterations.
pub fn compute_bisimilarity(lts: &Lts) -> Vec<Vec<State>> {
    let n = lts.num_states;
    if n == 0 {
        return vec![];
    }

    // partition[s] = block id for state s; start with all in block 0
    let mut partition: Vec<usize> = vec![0; n];

    // Sort actions for determinism
    let mut actions: Vec<Action> = {
        let mut set: HashSet<Action> = HashSet::new();
        for t in &lts.transitions {
            set.insert(t.action.clone());
        }
        set.into_iter().collect()
    };
    actions.sort();

    // Iteratively refine until stable; max n iterations
    for _ in 0..n + 1 {
        // Compute signature for each state
        // Signature: (current_block, Vec<(action_idx, succ_block)>) sorted
        let mut sig_to_id: HashMap<Vec<usize>, usize> = HashMap::new();
        let mut new_partition = vec![0usize; n];
        let mut next_id = 0usize;

        for s in 0..n {
            // Build flat signature vector: [block, a0, sb0, a0, sb1, ..., a1, sb0, ...]
            let mut sig: Vec<usize> = vec![partition[s]];
            for (ai, action) in actions.iter().enumerate() {
                let mut succs: Vec<usize> = lts
                    .transitions_by_action(s, action)
                    .iter()
                    .map(|&t| partition[t])
                    .collect();
                succs.sort();
                for sb in succs {
                    sig.push(ai);
                    sig.push(sb);
                }
                sig.push(usize::MAX); // separator between actions
            }

            let id = if let Some(&existing) = sig_to_id.get(&sig) {
                existing
            } else {
                let id = next_id;
                next_id += 1;
                sig_to_id.insert(sig, id);
                id
            };
            new_partition[s] = id;
        }

        if new_partition == partition {
            break; // stable
        }
        partition = new_partition;
    }

    // Collect partition into blocks
    let mut blocks: HashMap<usize, Vec<State>> = HashMap::new();
    for s in 0..n {
        blocks.entry(partition[s]).or_default().push(s);
    }
    blocks.into_values().collect()
}

/// Check if two states `p` and `q` are strongly bisimilar in the LTS.

pub fn are_bisimilar(lts: &Lts, p: State, q: State) -> bool {
    let partition = compute_bisimilarity(lts);
    for block in &partition {
        let has_p = block.contains(&p);
        let has_q = block.contains(&q);
        if has_p && has_q {
            return true;
        }
        if has_p || has_q {
            return false;
        }
    }
    false
}

/// Compute a **bisimulation relation** between two states.
///
/// Returns a `BisimulationRelation` if `p ~ q`, otherwise a witness to non-bisimilarity
/// as an error message.
pub fn compute_bisimulation_relation(
    lts: &Lts,
    p: State,
    q: State,
) -> Result<BisimulationRelation, String> {
    let mut rel = BisimulationRelation::empty();
    let mut to_check: VecDeque<(State, State)> = VecDeque::new();
    to_check.push_back((p, q));
    rel.add(p, q);

    let actions: HashSet<Action> = lts.alphabet();

    while let Some((s, t)) = to_check.pop_front() {
        for action in &actions {
            // Check s's transitions
            let s_succs = lts.transitions_by_action(s, action);
            for s_prime in &s_succs {
                let t_succs = lts.transitions_by_action(t, action);
                // Find t' such that (s', t') can be added to relation
                let matched = t_succs.iter().any(|&t_prime| {
                    rel.contains(*s_prime, t_prime) || {
                        // Provisional: check if they would work
                        true // simplified: greedy add
                    }
                });
                if !matched {
                    return Err(format!(
                        "Non-bisimilarity witness: state {} can do {} to {} but state {} cannot match",
                        s, action, s_prime, t
                    ));
                }
                // Add first match to relation
                for &t_prime in &t_succs {
                    if !rel.contains(*s_prime, t_prime) {
                        rel.add(*s_prime, t_prime);
                        to_check.push_back((*s_prime, t_prime));
                        break;
                    }
                }
            }

            // Check t's transitions (symmetric)
            let t_succs = lts.transitions_by_action(t, action);
            for t_prime in &t_succs {
                let s_succs = lts.transitions_by_action(s, action);
                if s_succs.is_empty() {
                    return Err(format!(
                        "Non-bisimilarity witness: state {} can do {} to {} but state {} cannot",
                        t, action, t_prime, s
                    ));
                }
                for &s_prime in &s_succs {
                    if !rel.contains(s_prime, *t_prime) {
                        rel.add(s_prime, *t_prime);
                        to_check.push_back((s_prime, *t_prime));
                        break;
                    }
                }
            }
        }
    }

    Ok(rel)
}

// ─── Weak Bisimilarity ────────────────────────────────────────────────────────

/// Check **weak bisimilarity** (≈) between two states.
///
/// A relation R is a weak bisimulation if for all (p, q) ∈ R and visible α:
/// - If `p →α p'`, then ∃ q' with `q ⇒α q'` and `(p', q') ∈ R`
/// - If `p →τ p'`, then ∃ q' with `q ⇒ε q'` and `(p', q') ∈ R`
/// (and symmetrically for q)
pub fn are_weakly_bisimilar(lts: &Lts, p: State, q: State) -> bool {
    // Compute weak bisimilarity by saturation
    let n = lts.num_states;
    let mut rel: HashSet<(State, State)> = HashSet::new();
    let mut worklist: VecDeque<(State, State)> = VecDeque::new();

    // Start with reflexivity
    for s in 0..n {
        rel.insert((s, s));
    }
    worklist.push_back((p, q));
    worklist.push_back((q, p));
    rel.insert((p, q));
    rel.insert((q, p));

    // Saturate: if (s, t) ∈ R and s →α s' then ∃ t' with t ⇒α t' and (s', t') ∈ R
    let actions: Vec<Action> = lts.alphabet().into_iter().collect();
    let mut changed = true;
    while changed {
        changed = false;
        let pairs: Vec<(State, State)> = rel.iter().copied().collect();
        for (s, t) in pairs {
            for action in &actions {
                let s_succs = lts.transitions_by_action(s, action);
                let t_weak = lts.weak_transitions(t, action);
                for s_prime in s_succs {
                    let matched = t_weak
                        .iter()
                        .any(|&t_prime| rel.contains(&(s_prime, t_prime)));
                    if !matched {
                        // Try to add a new pair
                        if let Some(&t_prime) = t_weak.iter().next() {
                            if rel.insert((s_prime, t_prime)) {
                                changed = true;
                            }
                        } else {
                            return false; // no match found
                        }
                    }
                }
            }
        }
    }

    rel.contains(&(p, q))
}

// ─── HML Model Checking ───────────────────────────────────────────────────────

/// Check if a state `s` in an LTS satisfies an HML formula `φ`.
///
/// Inductively:
/// - `s ⊨ tt` always
/// - `s ⊨ ff` never
/// - `s ⊨ φ ∧ ψ` iff `s ⊨ φ` and `s ⊨ ψ`
/// - `s ⊨ φ ∨ ψ` iff `s ⊨ φ` or `s ⊨ ψ`
/// - `s ⊨ ¬φ` iff `s ⊭ φ`
/// - `s ⊨ ⟨α⟩φ` iff ∃ `t` with `s →α t` and `t ⊨ φ`
/// - `s ⊨ \[α\]φ` iff ∀ `t` with `s →α t`, `t ⊨ φ`
pub fn hml_check(lts: &Lts, state: State, formula: &HmlFormula) -> bool {
    match formula {
        HmlFormula::True => true,
        HmlFormula::False => false,
        HmlFormula::And(phi, psi) => hml_check(lts, state, phi) && hml_check(lts, state, psi),
        HmlFormula::Or(phi, psi) => hml_check(lts, state, phi) || hml_check(lts, state, psi),
        HmlFormula::Not(phi) => !hml_check(lts, state, phi),
        HmlFormula::Diamond(action, phi) => lts
            .transitions_by_action(state, action)
            .iter()
            .any(|&t| hml_check(lts, t, phi)),
        HmlFormula::Box_(action, phi) => lts
            .transitions_by_action(state, action)
            .iter()
            .all(|&t| hml_check(lts, t, phi)),
    }
}

/// Find a **distinguishing formula** for two non-bisimilar states.
///
/// If `p ≁ q`, returns an HML formula that `p` satisfies but `q` does not (or vice versa).
/// Uses bounded depth search.
pub fn distinguishing_formula(lts: &Lts, p: State, q: State, depth: usize) -> Option<HmlFormula> {
    if depth == 0 {
        return None;
    }

    // Check if p can do something q cannot
    let actions: HashSet<Action> = lts.alphabet();
    for action in &actions {
        let p_succs = lts.transitions_by_action(p, action);
        let q_succs = lts.transitions_by_action(q, action);

        // p can do action but q cannot
        if !p_succs.is_empty() && q_succs.is_empty() {
            return Some(HmlFormula::diamond(action.clone(), HmlFormula::True));
        }

        // q can do action but p cannot
        if p_succs.is_empty() && !q_succs.is_empty() {
            return Some(HmlFormula::not(HmlFormula::diamond(
                action.clone(),
                HmlFormula::True,
            )));
        }

        // Both can do action — look for distinguishing post-state formula
        for &p_prime in &p_succs {
            let distinguished = q_succs.iter().all(|&q_prime| {
                // Try to find formula distinguishing p_prime from q_prime
                distinguishing_formula(lts, p_prime, q_prime, depth - 1).is_some()
            });
            if distinguished {
                if let Some(&q_prime) = q_succs.first() {
                    if let Some(phi) = distinguishing_formula(lts, p_prime, q_prime, depth - 1) {
                        return Some(HmlFormula::diamond(action.clone(), phi));
                    }
                }
            }
        }
    }

    None
}

// ─── Trace Computation ────────────────────────────────────────────────────────

/// Compute the **traces** of a state in an LTS up to length `max_len`.
pub fn compute_traces(lts: &Lts, initial: State, max_len: usize) -> TraceSet {
    let mut traces = TraceSet::new();
    let mut worklist: VecDeque<(State, Vec<Action>)> = VecDeque::new();
    worklist.push_back((initial, vec![]));

    while let Some((state, trace)) = worklist.pop_front() {
        traces.traces.insert(trace.clone());
        if trace.len() >= max_len {
            continue;
        }
        for t in &lts.transitions {
            if t.source == state {
                let mut new_trace = trace.clone();
                if t.action.is_visible() {
                    new_trace.push(t.action.clone());
                    worklist.push_back((t.target, new_trace));
                } else {
                    // τ transition: don't add to trace, but follow
                    worklist.push_back((t.target, trace.clone()));
                }
            }
        }
    }

    traces
}

/// Check **trace equivalence** between two states.
pub fn are_trace_equivalent(lts: &Lts, p: State, q: State, max_len: usize) -> bool {
    let traces_p = compute_traces(lts, p, max_len);
    let traces_q = compute_traces(lts, q, max_len);
    traces_p.traces == traces_q.traces
}

// ─── Failures Semantics ───────────────────────────────────────────────────────

/// Compute the **failures** of a state up to trace length `max_len`.
pub fn compute_failures(lts: &Lts, initial: State, max_len: usize) -> FailuresModel {
    let mut model = FailuresModel::new();
    let all_visible: HashSet<Action> = lts.visible_alphabet();

    let mut worklist: VecDeque<(State, Vec<Action>)> = VecDeque::new();
    worklist.push_back((initial, vec![]));

    while let Some((state, trace)) = worklist.pop_front() {
        // Compute refusal set: actions that cannot be done from this state (after τ-closure)
        let tau_closure = {
            let mut cl = lts.tau_closure(state);
            cl.insert(state);
            cl
        };

        let enabled: HashSet<Action> = tau_closure
            .iter()
            .flat_map(|&s| {
                lts.transitions_from(s)
                    .into_iter()
                    .map(|t| t.action.clone())
            })
            .filter(|a| a.is_visible())
            .collect();

        let refusal: HashSet<Action> = all_visible.difference(&enabled).cloned().collect();
        model.failures.push(Failure::new(trace.clone(), refusal));
        model.traces.traces.insert(trace.clone());

        if trace.len() >= max_len {
            continue;
        }

        // Follow transitions
        for s in &tau_closure {
            for t in lts.transitions_from(*s) {
                if t.action.is_visible() {
                    let mut new_trace = trace.clone();
                    new_trace.push(t.action.clone());
                    worklist.push_back((t.target, new_trace));
                }
            }
        }
    }

    model
}

// ─── CCS Structural Congruence ────────────────────────────────────────────────

/// Normalize a CCS process using structural congruence rules.
///
/// Rules applied:
/// - `P + 0 → P`
/// - `0 + P → P`
/// - `P | 0 → P`
/// - `0 | P → P`
pub fn normalize_ccs(proc: &CcsProcess) -> StructuralCongruenceClass {
    let (normalized, reductions) = normalize_inner(proc, 0);
    StructuralCongruenceClass {
        representative: normalized,
        reductions,
    }
}

fn normalize_inner(proc: &CcsProcess, count: usize) -> (CcsProcess, usize) {
    match proc {
        CcsProcess::Nil => (CcsProcess::Nil, count),
        CcsProcess::Prefix(a, p) => {
            let (p_norm, c) = normalize_inner(p, count);
            (CcsProcess::Prefix(a.clone(), Box::new(p_norm)), c)
        }
        CcsProcess::Choice(p, q) => {
            let (p_norm, c1) = normalize_inner(p, count);
            let (q_norm, c2) = normalize_inner(q, c1);
            // P + 0 → P; 0 + P → P
            match (&p_norm, &q_norm) {
                (_, CcsProcess::Nil) => (p_norm, c2 + 1),
                (CcsProcess::Nil, _) => (q_norm, c2 + 1),
                _ => (CcsProcess::Choice(Box::new(p_norm), Box::new(q_norm)), c2),
            }
        }
        CcsProcess::Parallel(p, q) => {
            let (p_norm, c1) = normalize_inner(p, count);
            let (q_norm, c2) = normalize_inner(q, c1);
            // P | 0 → P; 0 | P → P
            match (&p_norm, &q_norm) {
                (_, CcsProcess::Nil) => (p_norm, c2 + 1),
                (CcsProcess::Nil, _) => (q_norm, c2 + 1),
                _ => (CcsProcess::Parallel(Box::new(p_norm), Box::new(q_norm)), c2),
            }
        }
        CcsProcess::Restriction(p, l) => {
            let (p_norm, c) = normalize_inner(p, count);
            if l.is_empty() {
                (p_norm, c + 1) // P \ {} = P
            } else {
                (CcsProcess::Restriction(Box::new(p_norm), l.clone()), c)
            }
        }
        CcsProcess::Relabeling(p, f) => {
            let (p_norm, c) = normalize_inner(p, count);
            if f.is_empty() {
                (p_norm, c + 1) // P[id] = P
            } else {
                (CcsProcess::Relabeling(Box::new(p_norm), f.clone()), c)
            }
        }
        CcsProcess::Var(x) => (CcsProcess::Var(x.clone()), count),
        CcsProcess::Rec(x, p) => {
            let (p_norm, c) = normalize_inner(p, count);
            (CcsProcess::Rec(x.clone(), Box::new(p_norm)), c)
        }
    }
}

// ─── Process Equivalence Checking ────────────────────────────────────────────

/// Check if two CCS processes are equivalent under the specified behavioral equivalence.
///
/// Generates LTSs for both and checks the appropriate equivalence.
pub fn check_process_equivalence(
    p: &CcsProcess,
    q: &CcsProcess,
    equiv: BehavioralEquivalence,
    max_states: usize,
) -> Result<bool, String> {
    // Generate a combined LTS with states for both P and Q
    let (p_lts, _) =
        generate_lts(p, max_states, 10).ok_or_else(|| "LTS too large for P".to_string())?;
    let (q_lts, _) =
        generate_lts(q, max_states, 10).ok_or_else(|| "LTS too large for Q".to_string())?;

    // Combine LTSs: offset Q states by p_lts.num_states
    let offset = p_lts.num_states;
    let combined_states = p_lts.num_states + q_lts.num_states;
    let mut combined_transitions = p_lts.transitions.clone();
    for t in &q_lts.transitions {
        combined_transitions.push(Transition::new(
            t.source + offset,
            t.action.clone(),
            t.target + offset,
        ));
    }

    let combined = Lts::new(combined_states, 0, combined_transitions);
    let p_init = p_lts.initial;
    let q_init = q_lts.initial + offset;

    match equiv {
        BehavioralEquivalence::StrongBisimilarity => Ok(are_bisimilar(&combined, p_init, q_init)),
        BehavioralEquivalence::WeakBisimilarity => {
            Ok(are_weakly_bisimilar(&combined, p_init, q_init))
        }
        BehavioralEquivalence::TraceEquivalence => {
            Ok(are_trace_equivalent(&combined, p_init, q_init, 20))
        }
        _ => {
            // Fallback to trace equivalence for unimplemented ones
            Ok(are_trace_equivalent(&combined, p_init, q_init, 20))
        }
    }
}

/// Compute the **equational axioms** relating two bisimilar processes.
///
/// Checks basic CCS laws:
/// - Commutativity of choice: `P + Q ~ Q + P`
/// - Associativity of choice: `(P + Q) + R ~ P + (Q + R)`
/// - Idempotency: `P + P ~ P`
/// - Unit: `P + 0 ~ P`
pub fn verify_ccs_laws(p: &CcsProcess, q: &CcsProcess, _r: &CcsProcess) -> HashMap<String, bool> {
    let max_s = 50;
    let mut results = HashMap::new();

    // P + Q ~ Q + P
    let p_plus_q = CcsProcess::choice(p.clone(), q.clone());
    let q_plus_p = CcsProcess::choice(q.clone(), p.clone());
    results.insert(
        "choice_commutative".to_string(),
        check_process_equivalence(
            &p_plus_q,
            &q_plus_p,
            BehavioralEquivalence::StrongBisimilarity,
            max_s,
        )
        .unwrap_or(false),
    );

    // P + 0 ~ P
    let p_plus_0 = CcsProcess::choice(p.clone(), CcsProcess::Nil);
    results.insert(
        "choice_unit".to_string(),
        check_process_equivalence(
            &p_plus_0,
            p,
            BehavioralEquivalence::StrongBisimilarity,
            max_s,
        )
        .unwrap_or(false),
    );

    // P + P ~ P (idempotency of choice)
    let p_plus_p = CcsProcess::choice(p.clone(), p.clone());
    results.insert(
        "choice_idempotent".to_string(),
        check_process_equivalence(
            &p_plus_p,
            p,
            BehavioralEquivalence::StrongBisimilarity,
            max_s,
        )
        .unwrap_or(false),
    );

    // P | 0 ~ P
    let p_par_0 = CcsProcess::parallel(p.clone(), CcsProcess::Nil);
    results.insert(
        "parallel_unit".to_string(),
        check_process_equivalence(
            &p_par_0,
            p,
            BehavioralEquivalence::StrongBisimilarity,
            max_s,
        )
        .unwrap_or(false),
    );

    results
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::process_algebra::types::*;

    fn make_simple_lts() -> Lts {
        // States: 0, 1, 2
        // Transitions: 0 --a--> 1, 0 --b--> 2, 1 --c--> 0, 2 --c--> 0
        Lts::new(
            3,
            0,
            vec![
                Transition::new(0, Action::input("a"), 1),
                Transition::new(0, Action::input("b"), 2),
                Transition::new(1, Action::input("c"), 0),
                Transition::new(2, Action::input("c"), 0),
            ],
        )
    }

    fn make_bisimilar_lts() -> Lts {
        // Two processes that are bisimilar:
        // P = a.0 + a.0  and  Q = a.0
        // They are NOT strongly bisimilar (P has two a-transitions, Q has one)
        // but let's make truly bisimilar ones.
        // States 0-1: P  (0 --a--> 1)
        // States 2-3: Q  (2 --a--> 3)
        Lts::new(
            4,
            0,
            vec![
                Transition::new(0, Action::input("a"), 1),
                Transition::new(2, Action::input("a"), 3),
            ],
        )
    }

    #[test]
    fn test_action_complement() {
        let a = Action::input("x");
        let abar = Action::output("x");
        assert_eq!(a.complement(), abar);
        assert_eq!(abar.complement(), a);
        assert_eq!(Action::Tau.complement(), Action::Tau);
    }

    #[test]
    fn test_action_is_tau() {
        assert!(Action::Tau.is_tau());
        assert!(!Action::input("x").is_tau());
    }

    #[test]
    fn test_ccs_prefix() {
        let p = CcsProcess::prefix(Action::input("a"), CcsProcess::Nil);
        assert_eq!(format!("{}", p), "a.0");
    }

    #[test]
    fn test_ccs_choice() {
        let p = CcsProcess::choice(
            CcsProcess::prefix(Action::input("a"), CcsProcess::Nil),
            CcsProcess::prefix(Action::input("b"), CcsProcess::Nil),
        );
        let s = format!("{}", p);
        assert!(s.contains("+"));
    }

    #[test]
    fn test_ccs_parallel() {
        let p = CcsProcess::parallel(
            CcsProcess::prefix(Action::input("a"), CcsProcess::Nil),
            CcsProcess::prefix(Action::output("a"), CcsProcess::Nil),
        );
        let s = format!("{}", p);
        assert!(s.contains("|"));
    }

    #[test]
    fn test_ccs_substitute() {
        let x = CcsProcess::Var("X".to_string());
        let a_prefix = CcsProcess::prefix(Action::input("a"), CcsProcess::Nil);
        let result = x.substitute("X", &a_prefix);
        assert_eq!(result, a_prefix);
    }

    #[test]
    fn test_ccs_unfold() {
        // μX.(a.X) → a.(μX.(a.X))
        let body = CcsProcess::prefix(Action::input("a"), CcsProcess::Var("X".to_string()));
        let rec = CcsProcess::rec("X", body);
        let unfolded = rec.unfold();
        // Should be a.μX.a.X
        match &unfolded {
            CcsProcess::Prefix(a, _) => assert_eq!(*a, Action::input("a")),
            _ => panic!("expected prefix after unfolding"),
        }
    }

    #[test]
    fn test_one_step_ccs_nil() {
        let nil = CcsProcess::Nil;
        let steps = one_step_ccs(&nil, 0, 5);
        assert!(steps.is_empty());
    }

    #[test]
    fn test_one_step_ccs_prefix() {
        let p = CcsProcess::prefix(Action::input("a"), CcsProcess::Nil);
        let steps = one_step_ccs(&p, 0, 5);
        assert_eq!(steps.len(), 1);
        assert_eq!(steps[0].0, Action::input("a"));
        assert_eq!(steps[0].1, CcsProcess::Nil);
    }

    #[test]
    fn test_one_step_ccs_choice() {
        let p = CcsProcess::choice(
            CcsProcess::prefix(Action::input("a"), CcsProcess::Nil),
            CcsProcess::prefix(Action::input("b"), CcsProcess::Nil),
        );
        let steps = one_step_ccs(&p, 0, 5);
        assert_eq!(steps.len(), 2);
    }

    #[test]
    fn test_one_step_ccs_parallel_sync() {
        // a.0 | ā.0 should produce a τ transition
        let p = CcsProcess::parallel(
            CcsProcess::prefix(Action::input("a"), CcsProcess::Nil),
            CcsProcess::prefix(Action::output("a"), CcsProcess::Nil),
        );
        let steps = one_step_ccs(&p, 0, 5);
        let tau_steps: Vec<_> = steps.iter().filter(|(a, _)| a.is_tau()).collect();
        assert!(!tau_steps.is_empty(), "should have a τ sync transition");
    }

    #[test]
    fn test_lts_transitions_from() {
        let lts = make_simple_lts();
        let trans = lts.transitions_from(0);
        assert_eq!(trans.len(), 2);
    }

    #[test]
    fn test_lts_tau_closure() {
        let mut transitions = vec![
            Transition::new(0, Action::Tau, 1),
            Transition::new(1, Action::Tau, 2),
            Transition::new(2, Action::input("a"), 3),
        ];
        let lts = Lts::new(4, 0, transitions);
        let cl = lts.tau_closure(0);
        assert!(cl.contains(&1));
        assert!(cl.contains(&2));
        assert!(!cl.contains(&3)); // 3 reachable only via 'a', not τ
    }

    #[test]
    fn test_hml_true_false() {
        let lts = make_simple_lts();
        assert!(hml_check(&lts, 0, &HmlFormula::True));
        assert!(!hml_check(&lts, 0, &HmlFormula::False));
    }

    #[test]
    fn test_hml_diamond() {
        let lts = make_simple_lts();
        // State 0 can do 'a' — ⟨a⟩tt should hold
        let phi = HmlFormula::diamond(Action::input("a"), HmlFormula::True);
        assert!(hml_check(&lts, 0, &phi));

        // State 1 cannot do 'a' from state 0 — [a]ff should hold for state 1? No.
        // State 1 has transition 1 --c--> 0, not 'a'
        let phi2 = HmlFormula::diamond(Action::input("a"), HmlFormula::True);
        assert!(!hml_check(&lts, 1, &phi2)); // state 1 cannot do 'a'
    }

    #[test]
    fn test_hml_box() {
        let lts = make_simple_lts();
        // [b]tt at state 0: all b-successors satisfy tt (trivially true)
        let phi = HmlFormula::box_(Action::input("b"), HmlFormula::True);
        assert!(hml_check(&lts, 0, &phi));

        // [a]ff at state 0: all a-successors satisfy ff? state 1 does not satisfy ff
        let phi2 = HmlFormula::box_(Action::input("a"), HmlFormula::False);
        assert!(!hml_check(&lts, 0, &phi2)); // state 1 does not satisfy ff
    }

    #[test]
    fn test_hml_and_or() {
        let lts = make_simple_lts();
        let phi_a = HmlFormula::diamond(Action::input("a"), HmlFormula::True);
        let phi_b = HmlFormula::diamond(Action::input("b"), HmlFormula::True);

        // ⟨a⟩tt ∧ ⟨b⟩tt at state 0: both hold
        let phi_and = HmlFormula::and(phi_a.clone(), phi_b.clone());
        assert!(hml_check(&lts, 0, &phi_and));

        // ⟨a⟩tt ∨ ff: should hold
        let phi_or = HmlFormula::or(phi_a, HmlFormula::False);
        assert!(hml_check(&lts, 0, &phi_or));
    }

    #[test]
    fn test_compute_traces() {
        let lts = make_simple_lts();
        let traces = compute_traces(&lts, 0, 2);
        // Should contain: [], [a], [b], [a,c], [b,c]
        assert!(traces.contains(&[]));
        assert!(traces.contains(&[Action::input("a")]));
        assert!(traces.contains(&[Action::input("b")]));
    }

    #[test]
    fn test_normalize_choice_unit() {
        let p = CcsProcess::prefix(Action::input("a"), CcsProcess::Nil);
        let p_plus_0 = CcsProcess::choice(p.clone(), CcsProcess::Nil);
        let result = normalize_ccs(&p_plus_0);
        assert_eq!(result.representative, p);
        assert!(result.reductions > 0);
    }

    #[test]
    fn test_normalize_parallel_unit() {
        let p = CcsProcess::prefix(Action::input("a"), CcsProcess::Nil);
        let p_par_0 = CcsProcess::parallel(p.clone(), CcsProcess::Nil);
        let result = normalize_ccs(&p_par_0);
        assert_eq!(result.representative, p);
        assert!(result.reductions > 0);
    }

    #[test]
    fn test_bisimilarity_reflexive() {
        let lts = make_simple_lts();
        // Any state is bisimilar to itself
        assert!(are_bisimilar(&lts, 0, 0));
        assert!(are_bisimilar(&lts, 1, 1));
    }

    #[test]
    fn test_bisimilarity_non_bisimilar() {
        let lts = make_simple_lts();
        // States 0 and 1 have different transition capabilities
        assert!(!are_bisimilar(&lts, 0, 1));
    }

    #[test]
    fn test_generate_lts_simple() {
        let p = CcsProcess::prefix(
            Action::input("a"),
            CcsProcess::prefix(Action::input("b"), CcsProcess::Nil),
        );
        let result = generate_lts(&p, 10, 5);
        assert!(result.is_some());
        let (lts, _) = result.expect("lts generated");
        assert!(lts.num_states >= 2); // at least start and after 'a'
    }

    #[test]
    fn test_csp_display() {
        let p = CspProcess::Sequential(
            Box::new(CspProcess::Prefix(
                "a".to_string(),
                Box::new(CspProcess::Stop),
            )),
            Box::new(CspProcess::Skip),
        );
        let s = format!("{}", p);
        assert!(s.contains(";"));
    }

    #[test]
    fn test_behavioral_equivalence_ordering() {
        assert!(
            BehavioralEquivalence::StrongBisimilarity < BehavioralEquivalence::TraceEquivalence
        );
        assert!(BehavioralEquivalence::WeakBisimilarity < BehavioralEquivalence::TraceEquivalence);
    }

    #[test]
    fn test_build_env() {
        let mut env = Environment::new();
        build_process_algebra_env(&mut env);
    }
}