// nn_extra.mbt

///|
/// A simple ordered container for Linear layers.
///
/// Sequential deliberately keeps the layer type small and explicit. It is
/// useful for examples and for models whose activation functions are inserted
/// between calls to `forward`.
pub struct Sequential {
  layers : Array[Linear]
}

///|
/// Create a sequential model from an ordered list of layers.
pub fn Sequential::new(layers : Array[Linear]) -> Sequential {
  { layers, }
}

///|
/// Create a sequential model containing one linear layer.
pub fn Sequential::single(layer : Linear) -> Sequential {
  { layers: [layer] }
}

///|
/// Apply every stored layer in order.
pub fn Sequential::forward(self : Sequential, input : Tensor) -> Tensor {
  let mut output = input
  for layer in self.layers {
    output = layer.forward(output)
  }
  output
}

///|
/// Return all trainable tensors in layer order.
pub fn Sequential::parameters(self : Sequential) -> Array[Tensor] {
  let params : Array[Tensor] = []
  for layer in self.layers {
    let layer_params = layer.parameters()
    for param in layer_params {
      params.push(param)
    }
  }
  params
}

///|
/// Return the number of layers in the model.
pub fn Sequential::layer_count(self : Sequential) -> Int {
  self.layers.length()
}

///|
/// Return the total number of scalar trainable values.
pub fn Sequential::parameter_count(self : Sequential) -> Int {
  let mut count = 0
  for layer in self.layers {
    for param in layer.parameters() {
      count = count + param.data.length()
    }
  }
  count
}

///|
/// Borrow a layer by index, panicking when the index is outside the model.
pub fn Sequential::layer(self : Sequential, index : Int) -> Linear {
  if index < 0 || index >= self.layers.length() {
    panic()
  }
  self.layers[index]
}

///|
/// Count scalar values in a collection of tensors.
pub fn parameter_count(params : Array[Tensor]) -> Int {
  let mut count = 0
  for param in params {
    count = count + param.data.length()
  }
  count
}

///|
/// Compute the fraction of correct class predictions.
///
/// `logits` must have shape `[batch, classes]`; labels are integer-valued
/// tensors with shape `[batch]`. This metric is intentionally detached from
/// autograd because it is intended for reporting rather than optimization.
pub fn classification_accuracy(logits : Tensor, labels : Tensor) -> Double {
  if logits.shape.length() != 2 ||
    labels.shape.length() != 1 ||
    logits.shape[0] != labels.shape[0] {
    panic()
  }
  let batch = logits.shape[0]
  let classes = logits.shape[1]
  if batch == 0 || classes == 0 {
    panic()
  }
  let mut correct = 0
  for row in 0.. logits.data[offset + best] {
        best = col
      }
    }
    if labels.data[row].to_int() == best {
      correct = correct + 1
    }
  }
  correct.to_double() / batch.to_double()
}

///|
/// Compute the mean absolute error as a detached metric.
pub fn mean_absolute_error(pred : Tensor, target : Tensor) -> Double {
  if pred.shape != target.shape || pred.data.length() == 0 {
    panic()
  }
  let mut total = 0.0
  for i in 0..