type Model = (
    AddInto<(
        // phoneme a
        Linear<1, HIDDEN_SIZE>,
        
        // phoneme b         
        Linear<1, HIDDEN_SIZE>, 
        
        // noise
        Linear<1, HIDDEN_SIZE>, 
        
        // state
        Linear<STATE_SIZE, HIDDEN_SIZE>
    )>, 
    ReLU,
    SplitInto<(
        // state
        Linear<HIDDEN_SIZE, STATE_SIZE>,
        
        // next
        Linear<HIDDEN_SIZE, 1>,        

        // sample
        Linear<HIDDEN_SIZE, 1>
    )>
);

fn main() {
    // make rng
    let mut rng = StdRng::seed_from_u64(0);

    // make model
    let mut model = Model::default();
    
    // data TODO
    let x: Tensor1D<2> = Tensor1D::randn(&mut rng);
    let y: Tensor1D<8> = Tensor1D::randn(&mut rng);

    // gradient descent
    let mut sgd = Sgd::new(SgdConfig {
        lr: 0.01,
        momentum: Some(Momentum::Nesterov(0.9)),
        weight_decay: None,
    });

    // other idea:
    // generate a voice line
    // split on phonemes
    // train on one phoneme and that way you can still do batching

    // train
    for _ in 0..10 {
        // internal state
        let mut state = Tensor1D::<STATE_SIZE>::zeros().traced();

        for _ in 0..5 {

            // input
            let phoneme_a = Tensor1D::new([0.0]).traced();
            let phoneme_b = Tensor1D::new([0.0]).traced();
            let noise = Tensor1D::new([0.0]).traced();
            
            // forward
            let (new_state, next, sample) = model.forward((phoneme_a, phoneme_b, noise, state));

            // loss
            let loss = mse_loss(sample, Tensor1D::new([0.0]));

            // gradients, breaks here
            // thread 'main' panicked at 'called `Option::unwrap()` on a `None` value', /home/username/.cargo/git/checkouts/dfdx-318e6e5ad83eea79/3fc7be4/src/gradients.rs:273:14 
            let gradients = loss.backward();

            // update
            sgd.update(&mut model, gradients).expect("nn machine broke");
            
            // keep state
            state = new_state.traced();
        }
    }

    //println!("{:?}", model);
}