rust_autograd_02

use std::time::{Instant, SystemTime, UNIX_EPOCH};
use rand::prelude::*;

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

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

struct Operation {
    op_type: OpType,
    input_indices: [usize; 2], // fixed-size array for input indices
    n_inputs: usize,
    output_index: usize,
}

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

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

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

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

        self.ops.push(Operation {
            op_type: OpType::Sum,
            input_indices: {
                let mut arr = [0; 2];
                arr[..input_indices.len()].copy_from_slice(input_indices);
                arr
            },
            n_inputs: input_indices.len(),
            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.create_var(prod_val);

        self.ops.push(Operation {
            op_type: OpType::Prod,
            input_indices: [input1_idx, input2_idx],
            n_inputs: 2, // fixed number of inputs for Prod
            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.create_var(softplus_val);

        self.ops.push(Operation {
            op_type: OpType::Softplus,
            input_indices: [input_idx, 0], // not using the second index
            n_inputs: 1,
            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[..op.n_inputs] {
                        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;
                    let exp_val = (-input_val).exp();
                    self.vars[op.input_indices[0]].grad += output_grad / (1.0 + exp_val);
                }
            }
        }
    }
}

fn seed_rng() -> impl Rng {
    let now = SystemTime::now();
    let since_epoch = now
        .duration_since(UNIX_EPOCH)
        .expect("Time went backwards");
    let seed = since_epoch.as_nanos() as u64;

    rand::rngs::StdRng::seed_from_u64(seed)
}

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

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

        let var1_idx = tape.create_var(rng.gen::<f64>());
        let var2_idx = tape.create_var(rng.gen::<f64>());

        let sum_idx = tape.sum(&[var1_idx, var2_idx]);
        let prod_idx = tape.prod(sum_idx, sum_idx);
        let softplus_idx = tape.softplus(prod_idx);

        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().as_micros() as f64 / 1000.0;
    println!("\nElapsed time: {:.2} ms\n", duration);
}