rust_autograd_01

use std::time::Instant;

// Define the NaiveVar struct
#[derive(Debug, Clone, Copy)]
struct NaiveVar {
    val: f64,
    grad: f64,
}

// Define operation types
#[derive(Clone, Copy)]
enum OpType {
    Sum,
    Prod,
    Softplus,
}

// Define the Operation struct
struct Operation {
    op_type: OpType,
    input_indices: Vec<usize>,
    output_index: usize,
}

// Define the NaiveTape struct
struct NaiveTape {
    ops: Vec<Operation>,
    vars: Vec<NaiveVar>,
}

impl NaiveTape {
    fn new() -> Self {
        NaiveTape {
            ops: Vec::with_capacity(16),
            vars: Vec::with_capacity(16),
        }
    }

    fn add_var(&mut self, val: f64) -> usize {
        let idx = self.vars.len();
        self.vars.push(NaiveVar { val, grad: 0.0 });
        idx
    }

    fn sum(&mut self, input_indices: Vec<usize>) -> usize {
        let sum_val: f64 = input_indices.iter()
            .map(|&idx| self.vars[idx].val)
            .sum();
        let output_index = self.add_var(sum_val);

        self.ops.push(Operation {
            op_type: OpType::Sum,
            input_indices,
            output_index,
        });
        output_index
    }

    fn prod(&mut self, input1_idx: usize, input2_idx: usize) -> usize {
        let prod_val = self.vars[input1_idx].val * self.vars[input2_idx].val;
        let output_index = self.add_var(prod_val);

        self.ops.push(Operation {
            op_type: OpType::Prod,
            input_indices: vec![input1_idx, input2_idx],
            output_index,
        });
        output_index
    }

    fn softplus(&mut self, input_idx: usize) -> usize {
        let softplus_val = (self.vars[input_idx].val.exp() + 1.0).ln();
        let output_index = self.add_var(softplus_val);

        self.ops.push(Operation {
            op_type: OpType::Softplus,
            input_indices: vec![input_idx],
            output_index,
        });
        output_index
    }

    fn backward(&mut self, output_idx: usize) {
        self.vars[output_idx].grad = 1.0;

        for op in self.ops.iter().rev() {
            let output_grad = self.vars[op.output_index].grad;

            match op.op_type {
                OpType::Sum => {
                    for &input_idx in &op.input_indices {
                        self.vars[input_idx].grad += output_grad;
                    }
                }
                OpType::Prod => {
                    let input1_val = self.vars[op.input_indices[0]].val;
                    let input2_val = self.vars[op.input_indices[1]].val;
                    self.vars[op.input_indices[0]].grad += input2_val * output_grad;
                    self.vars[op.input_indices[1]].grad += input1_val * output_grad;
                }
                OpType::Softplus => {
                    let input_val = self.vars[op.input_indices[0]].val;
                    self.vars[op.input_indices[0]].grad += output_grad / (1.0 + (-input_val).exp());
                }
            }
        }
    }
}

fn main() {
    let iterations: usize = 1000000;

    let start = Instant::now();
    for i in 0..iterations {
        let mut tape = NaiveTape::new();

        // Create variables
        let var1_idx = tape.add_var(1.0);
        let var2_idx = tape.add_var(2.0);

        // Build computational graph
        let sum_idx = tape.sum(vec![var1_idx, var2_idx]);
        let prod_idx = tape.prod(sum_idx, sum_idx);
        let softplus_idx = tape.softplus(prod_idx);

        // Perform backward pass
        tape.backward(softplus_idx);

        if i == iterations - 1 {
            println!("sum_var val: {}", tape.vars[sum_idx].val);
            println!("prod_var val: {}", tape.vars[prod_idx].val);
            println!("softplus_var val: {}", tape.vars[softplus_idx].val);
            println!("sum_var grad: {}", tape.vars[sum_idx].grad);
            println!("var1 grad: {}", tape.vars[var1_idx].grad);
            println!("var2 grad: {}", tape.vars[var2_idx].grad);
        }
    }
    let duration = start.elapsed();
    println!("\nElapsed time: {:?}\n", duration);
}