// tensor.mbt

// Global tensor ID counter to identify tensors in the computation graph

///|
let next_tensor_id : Ref[Int] = { val: 0 }

///|
fn get_next_tensor_id() -> Int {
  let id = next_tensor_id.val
  next_tensor_id.val = id + 1
  id
}

///|
/// Multi-dimensional array representing a Tensor.
pub struct Tensor {
  id : Int
  shape : Array[Int]
  strides : Array[Int]
  data : Array[Double]
  mut grad : Array[Double]?
  requires_grad : Bool
  mut creator : Op?
}

///|
/// Helper to calculate strides for a given shape.
pub fn shape_to_strides(shape : Array[Int]) -> Array[Int] {
  let len = shape.length()
  if len == 0 {
    return []
  }
  let strides = Array::make(len, 1)
  let mut s = 1
  for i in 0.. @int.MAX_VALUE / dim) {
      panic()
    }
    strides[idx] = s
    s = s * dim
  }
  strides
}

///|
/// Return the number of elements described by a shape.
/// Panics when a dimension is negative.
fn checked_shape_size(shape : Array[Int]) -> Int {
  let mut size = 1
  for dim in shape {
    if dim < 0 || (dim != 0 && size > @int.MAX_VALUE / dim) {
      panic()
    }
    size = size * dim
  }
  size
}

///|
/// Create a new Tensor.
pub fn Tensor::new(
  data : Array[Double],
  shape : Array[Int],
  requires_grad? : Bool = false,
) -> Tensor {
  let size = checked_shape_size(shape)
  if data.length() != size {
    panic()
  }
  let strides = shape_to_strides(shape)
  let id = get_next_tensor_id()
  { id, shape, strides, data, grad: None, requires_grad, creator: None }
}

///|
/// Create a Tensor of zeros with the specified shape.
pub fn zeros(shape : Array[Int], requires_grad? : Bool = false) -> Tensor {
  let size = checked_shape_size(shape)
  let data = Array::make(size, 0.0)
  Tensor::new(data, shape, requires_grad~)
}

///|
/// Create a Tensor of ones with the specified shape.
pub fn ones(shape : Array[Int], requires_grad? : Bool = false) -> Tensor {
  let size = checked_shape_size(shape)
  let data = Array::make(size, 1.0)
  Tensor::new(data, shape, requires_grad~)
}

///|
/// Simple LCG pseudo-random number generator for reproducible initialization.
priv struct Random {
  mut state : Int64
}

///|
fn Random::new(seed : Int) -> Random {
  { state: seed.to_int64() }
}

///|
fn Random::next_double(self : Random) -> Double {
  // LCG parameters (Numerical Recipes)
  let a : Int64 = 1664525
  let c : Int64 = 1013904223
  let m : Int64 = 4294967296
  self.state = (a * self.state + c) % m
  let d = self.state.to_double() / m.to_double()
  d
}

// Approximate normal distribution using Central Limit Theorem

///|
fn Random::next_randn(self : Random) -> Double {
  let mut sum = 0.0
  for _ in 0..<12 {
    sum = sum + self.next_double()
  }
  sum - 6.0
}

///|
/// Create a Tensor of random values from a normal distribution.
pub fn randn(
  shape : Array[Int],
  seed? : Int = 42,
  requires_grad? : Bool = false,
) -> Tensor {
  let size = checked_shape_size(shape)
  let rng = Random::new(seed)
  let data = Array::make(size, 0.0)
  for i in 0.. Tensor {
  Tensor::new(data, shape, requires_grad~)
}

///|
/// Validate that indices address exactly one element in a shape.
fn validate_indices(shape : Array[Int], indices : Array[Int]) -> Unit {
  if indices.length() != shape.length() {
    panic()
  }
  for i in 0..= shape[i] {
      panic()
    }
  }
}

///|
/// Get element at specified indices.
pub fn Tensor::op_get(self : Tensor, indices : Array[Int]) -> Double {
  validate_indices(self.shape, indices)
  let mut flat_idx = 0
  let len = indices.length()
  for i in 0.. Unit {
  validate_indices(self.shape, indices)
  let mut flat_idx = 0
  let len = indices.length()
  for i in 0.. String {
  if dim == shape.length() {
    return data[offset].to_string()
  }
  if dim == shape.length() - 1 {
    let mut s = "["
    let size = shape[dim]
    let stride = strides[dim]
    for i in 0.. 0 {
        s = s + ", "
      }
      s = s + data[offset + i * stride].to_string()
    }
    s = s + "]"
    return s
  }
  let mut s = "["
  let size = shape[dim]
  let stride = strides[dim]
  for i in 0.. 0 {
      s = s + ",\n"
      let indent_limit = dim + 1
      for _ in 0.. String {
  let shape_len = self.shape.length()
  let data_str = if shape_len == 0 {
    self.data[0].to_string()
  } else {
    format_tensor_data(self.data, self.shape, self.strides, 0, 0)
  }
  let grad_info = if self.requires_grad { ", requires_grad=true" } else { "" }
  "tensor(" + data_str + grad_info + ")"
}

///|
/// Implementation of Show trait for Tensor.
pub fn Tensor::output(self : Tensor, logger : &Logger) -> Unit {
  logger.write_string(self.to_string())
}