use miniature::functions as F;
use miniature::graph::backward;
use miniature::optimizer as S;
use miniature::parametric_functions as PF;
use miniature::variable::Variable;

use std::rc::Rc;
use std::cell::RefCell;

fn main() {
    // define layers
    let fc1 = PF::linear(28 * 28, 256);
    let fc2 = PF::linear(256, 256);
    let fc3 = PF::linear(256, 10);

    // define optimizer
    let mut optim = S::adam(0.001, (0.9, 0.999), 1e-8);
    optim.set_params(fc1.get_params());
    optim.set_params(fc2.get_params());
    optim.set_params(fc3.get_params());

    let x = Rc::new(RefCell::new(Variable::rand(vec![32, 28 * 28])));
    let t = Rc::new(RefCell::new(Variable::rand(vec![32])));

    // forward
    let h1 = F::relu(fc1.call(x));
    let h2 = F::relu(fc2.call(h1));
    let y = fc3.call(h2);

    // loss
    let loss = F::cross_entropy_loss(y, F::onehot(t, 10));

    // update
    optim.zero_grad();
    backward(loss);
    optim.update();
}