extern crate rand;
extern crate smyl;

use rand::thread_rng;
use smyl::network::LayerGradient;
use smyl::prelude::*;

fn main() {
    let mut rng = thread_rng();

    let input = Matrix::from([[1.0, 0.0, 1.0, 0.5]]);
    let expected = Matrix::from([[1.0, 0.0, 0.0, 0.0]]);
    let learning_rate: f64 = 0.01;

    let mut synapses1 = SynapseLayer::<f64, ReLU>::random(4, 2, &mut rng);
    let mut synapses2 = SynapseLayer::<f64, ReLU>::random(2, 2, &mut rng);
    let mut synapses3 = SynapseLayer::<f64, ReLU>::random(2, 4, &mut rng);

    let batch = 100;

    for _ in 0..10_000 / batch {
        let mut gradient1 = LayerGradient::zero(4, 2);
        let mut gradient2 = LayerGradient::zero(2, 2);
        let mut gradient3 = LayerGradient::zero(2, 4);

        for _ in 0..batch {
            let output1 = synapses1.forward(input.clone());
            let output2 = synapses2.forward(output1.clone());
            let output3 = synapses3.forward(output2.clone());

            let output_gradient3 = output3.clone() - expected.clone();
            let output_gradient2 =
                synapses3.backward(&output2, output3.clone(), output_gradient3, &mut gradient3);
            let output_gradient1 =
                synapses2.backward(&output1, output2.clone(), output_gradient2, &mut gradient2);
            synapses1.backward(&input, output1.clone(), output_gradient1, &mut gradient1);

            dbg!(output3);
        }

        synapses3.apply_gradient(&gradient3, learning_rate);
        synapses2.apply_gradient(&gradient2, learning_rate);
        synapses1.apply_gradient(&gradient1, learning_rate);
    }
}