use crate::{
    core::compile_passes::{circuit_pass::CircuitPass, Gate},
    utils::circuit_utils::is_input_really_input,
    AsyncMPCCircuit,
};

#[derive(Debug, Default)]
pub struct DeadGateEliminator;

impl CircuitPass<AsyncMPCCircuit> for DeadGateEliminator {
    fn run(&mut self, old: AsyncMPCCircuit) -> AsyncMPCCircuit {
        let outputs: Vec<usize> = old.output_indices().map(|x| *x as usize).collect();
        let gates = old.ops;
        let mut used = vec![false; gates.len()];
        for &output in &outputs {
            used[output] = true;
        }
        for (idx, gate) in gates.iter().enumerate().rev() {
            if !used[idx] {
                if let Gate::Input { input_type } = gate {
                    used[idx] = is_input_really_input(input_type);
                }
                continue;
            }
            gate.for_each_gate_index(|label| {
                used[label as usize] = true;
            });
        }
        let mut old_id_to_new_id = vec![None; gates.len()];
        let mut res = AsyncMPCCircuit::default();
        for (idx, gate) in gates.into_iter().enumerate() {
            if !used[idx] {
                continue;
            }
            let new_gate = gate
                .map_gate_indices(|label| old_id_to_new_id[label as usize].expect("unknown label"));
            let label = res.add_gate(new_gate);
            old_id_to_new_id[idx] = Some(label);
        }
        for output in outputs {
            res.output_gate(old_id_to_new_id[output].expect("unknown output"));
        }

        res
    }
}
