// reduce.mbt

///|
fn validate_reduce_dim(shape : Array[Int], dim : Int) -> Unit {
  if dim < 0 || dim >= shape.length() {
    panic()
  }
  if shape[dim] == 0 {
    panic()
  }
}

///|
fn reduction_shape(shape : Array[Int], dim : Int, keepdim : Bool) -> Array[Int] {
  validate_reduce_dim(shape, dim)
  if keepdim {
    let result = shape.copy()
    result[dim] = 1
    result
  } else {
    let result : Array[Int] = []
    for i in 0.. Int {
  let mut remaining = input_flat
  let mut output_index = 0
  let mut output_dim = 0
  for i in 0.. (Array[Double], Array[Int]) {
  let out_shape = reduction_shape(shape, dim, keepdim)
  let out_strides = shape_to_strides(out_shape)
  let out_size = if out_shape.length() == 0 {
    1
  } else {
    let mut size = 1
    for d in out_shape {
      size = size * d
    }
    size
  }
  let out = Array::make(out_size, 0.0)
  for i in 0.. Tensor {
  let mut total = 0.0
  for value in self.data {
    total = total + value
  }
  let out = Tensor::new([total], [], requires_grad=self.requires_grad)
  if self.requires_grad {
    out.creator = Some(ReduceSum(self, None, false, false))
  }
  out
}

///|
/// Compute the arithmetic mean of all elements.
pub fn Tensor::mean(self : Tensor) -> Tensor {
  if self.data.length() == 0 {
    panic()
  }
  let total = self.sum()
  let out = Tensor::new(
    [total.data[0] / self.data.length().to_double()],
    [],
    requires_grad=self.requires_grad,
  )
  if self.requires_grad {
    out.creator = Some(ReduceSum(self, None, false, true))
  }
  out
}

///|
/// Sum elements along one dimension.
pub fn Tensor::sum_dim(
  self : Tensor,
  dim : Int,
  keepdim? : Bool = false,
) -> Tensor {
  let (data, shape) = reduce_sum_data(
    self.data,
    self.shape,
    self.strides,
    dim,
    keepdim,
  )
  let out = Tensor::new(data, shape, requires_grad=self.requires_grad)
  if self.requires_grad {
    out.creator = Some(ReduceSum(self, Some(dim), keepdim, false))
  }
  out
}

///|
/// Compute means along one dimension.
pub fn Tensor::mean_dim(
  self : Tensor,
  dim : Int,
  keepdim? : Bool = false,
) -> Tensor {
  let (data, shape) = reduce_sum_data(
    self.data,
    self.shape,
    self.strides,
    dim,
    keepdim,
  )
  let count = self.shape[dim].to_double()
  for i in 0.. (Array[Double], Array[Double], Array[Int]) {
  let out_shape = reduction_shape(tensor.shape, dim, keepdim)
  let out_strides = shape_to_strides(out_shape)
  let out_size = if out_shape.length() == 0 {
    1
  } else {
    let mut size = 1
    for d in out_shape {
      size = size * d
    }
    size
  }
  let values = Array::make(out_size, if find_max { -1.0e300 } else { 1.0e300 })
  let indices = Array::make(out_size, 0.0)
  for i in 0.. values[index]
    } else {
      value < values[index]
    }
    if better {
      values[index] = value
      indices[index] = coord.to_double()
    }
  }
  (values, indices, out_shape)
}

///|
/// Return the largest scalar value.
pub fn Tensor::max(self : Tensor) -> Tensor {
  if self.data.length() == 0 {
    panic()
  }
  let mut result = self.data[0]
  for value in self.data {
    if value > result {
      result = value
    }
  }
  Tensor::new([result], [], requires_grad=false)
}

///|
/// Return the smallest scalar value.
pub fn Tensor::min(self : Tensor) -> Tensor {
  if self.data.length() == 0 {
    panic()
  }
  let mut result = self.data[0]
  for value in self.data {
    if value < result {
      result = value
    }
  }
  Tensor::new([result], [], requires_grad=false)
}

///|
/// Return the flat index of the largest scalar value.
pub fn Tensor::argmax(self : Tensor) -> Tensor {
  if self.data.length() == 0 {
    panic()
  }
  let mut index = 0
  for i in 1.. self.data[index] {
      index = i
    }
  }
  Tensor::new([index.to_double()], [], requires_grad=false)
}

///|
/// Reduce maximum values along a dimension.
pub fn Tensor::max_dim(
  self : Tensor,
  dim : Int,
  keepdim? : Bool = false,
) -> Tensor {
  let (values, _, shape) = extreme_dim(self, dim, keepdim, true)
  Tensor::new(values, shape, requires_grad=false)
}

///|
/// Reduce minimum values along a dimension.
pub fn Tensor::min_dim(
  self : Tensor,
  dim : Int,
  keepdim? : Bool = false,
) -> Tensor {
  let (values, _, shape) = extreme_dim(self, dim, keepdim, false)
  Tensor::new(values, shape, requires_grad=false)
}

///|
/// Return the index of the maximum value along a dimension.
pub fn Tensor::argmax_dim(
  self : Tensor,
  dim : Int,
  keepdim? : Bool = false,
) -> Tensor {
  let (_, indices, shape) = extreme_dim(self, dim, keepdim, true)
  Tensor::new(indices, shape, requires_grad=false)
}