Design a more advanced neural network in rust
Provide the full code. Make this thorough.
We will be training the neural network on text. fyi

use ndarray::{Array1, Array2, Axis};
use ndarray::prelude::*;
use ndarray_rand::RandomExt;
use ndarray_rand::rand_distr::Uniform;
use rand::thread_rng;

/// Trait for activation functions
pub trait ActivationFunction {
    fn activate(x: &Array2<f64>) -> Array2<f64>;
    fn derivative(x: &Array2<f64>) -> Array2<f64>;
}

/// Sigmoid activation function
pub struct Sigmoid;

impl ActivationFunction for Sigmoid {
    fn activate(x: &Array2<f64>) -> Array2<f64> {
        x.mapv(|v| 1.0 / (1.0 + (-v).exp()))
    }

    fn derivative(x: &Array2<f64>) -> Array2<f64> {
        let activated = Self::activate(x);
        &activated * &(1.0 - &activated)
    }
}

/// ReLU activation function
pub struct ReLU;

impl ActivationFunction for ReLU {
    fn activate(x: &Array2<f64>) -> Array2<f64> {
        x.mapv(|v| v.max(0.0))
    }

    fn derivative(x: &Array2<f64>) -> Array2<f64> {
        x.mapv(|v| if v > 0.0 { 1.0 } else { 0.0 })
    }
}

/// Layer of the neural network
pub struct Layer {
    weights: Array2<f64>,
    biases: Array1<f64>,
}

impl Layer {
    /// Create a new layer with random weights and biases
    pub fn new(input_size: usize, output_size: usize) -> Self {
        let mut rng = thread_rng();
        let weights = Array::random_using((input_size, output_size), Uniform::new(-1.0, 1.0), &mut rng);
        let biases = Array::random_using(output_size, Uniform::new(-1.0, 1.0), &mut rng);
        Layer { weights, biases }
    }

    /// Forward pass through this layer
    pub fn forward(&self, input: &Array2<f64>) -> Array2<f64> {
        input.dot(&self.weights) + &self.biases
    }
}

/// Neural network struct
pub struct NeuralNetwork<A: ActivationFunction> {
    layers: Vec<Layer>,
    activation: A,
    learning_rate: f64,
}

impl<A: ActivationFunction> NeuralNetwork<A> {
    /// Create a new neural network
    pub fn new(layer_sizes: &[usize], activation: A, learning_rate: f64) -> Self {
        let layers = layer_sizes.windows(2)
            .map(|w| Layer::new(w[0], w[1]))
            .collect();
        NeuralNetwork {
            layers,
            activation,
            learning_rate,
        }
    }

    /// Predict output for given input
    pub fn predict(&self, input: &Array2<f64>) -> Array2<f64> {
        let mut output = input.clone();
        for layer in &self.layers {
            output = layer.forward(&output);
            output = A::activate(&output);
        }
        output
    }

    /// Train the network with input and target output
    pub fn train(&mut self, input: &Array2<f64>, target: &Array2<f64>, epochs: usize) {
        for epoch in 0..epochs {
            let outputs = self.forward_pass(input);
            self.backward_pass(input, target, &outputs);
            if epoch % 1000 == 0 {
                let loss = self.compute_loss(&outputs.last().unwrap(), target);
                println!("Epoch {}: Loss = {}", epoch, loss);
            }
        }
    }

    /// Forward pass through the network, returning outputs at each layer
    fn forward_pass(&self, input: &Array2<f64>) -> Vec<Array2<f64>> {
        let mut outputs = Vec::new();
        let mut output = input.clone();
        outputs.push(output.clone());
        for layer in &self.layers {
            output = layer.forward(&output);
            output = A::activate(&output);
            outputs.push(output.clone());
        }
        outputs
    }

    /// Backward pass using backpropagation
    fn backward_pass(&mut self, input: &Array2<f64>, target: &Array2<f64>, outputs: &Vec<Array2<f64>>) {
        let mut delta = outputs.last().unwrap() - target;
        for i in (0..self.layers.len()).rev() {
            let z = &outputs[i + 1]; // Activated output at layer i+1
            let a = &outputs[i]; // Activated output at layer i

            let derivative = A::derivative(z);
            delta = delta * &derivative;

            let weight_gradient = a.t().dot(&delta);
            let bias_gradient = delta.sum_axis(Axis(0));

            // Update weights and biases
            self.layers[i].weights -= &(weight_gradient * self.learning_rate);
            self.layers[i].biases -= &(bias_gradient * self.learning_rate);

            if i > 0 {
                delta = delta.dot(&self.layers[i].weights.t());
            }
        }
    }

    /// Compute mean squared error loss
    fn compute_loss(&self, output: &Array2<f64>, target: &Array2<f64>) -> f64 {
        ((output - target).mapv(|x| x.powi(2))).mean().unwrap()
    }
}

fn main() {
    // Example usage with XOR problem
    let inputs = array![
        [0.0, 0.0],
        [0.0, 1.0],
        [1.0, 0.0],
        [1.0, 1.0],
    ];

    let targets = array![
        [0.0],
        [1.0],
        [1.0],
        [0.0],
    ];

    // Create a neural network with 2 input neurons, one hidden layer with 2 neurons, and 1 output neuron
    let mut nn = NeuralNetwork::new(&[2, 2, 1], Sigmoid, 0.5);

    // Train the network
    nn.train(&inputs, &targets, 10000);

    // Predict output
    let output = nn.predict(&inputs);
    println!("Predictions:\n{:?}", output);
}