// 3D & 4D Tensor operations, matrix decomposition and Taylor expansion.

///|
/// 3D Complex tensor structure
pub struct Tensor3D {
  dim0 : Int
  dim1 : Int
  dim2 : Int
  data : Array[Complex]
} derive(Debug, Eq)

///|
/// Create 3D Complex Tensor
pub fn Tensor3D::new(dim0 : Int, dim1 : Int, dim2 : Int) -> Tensor3D {
  let size = dim0 * dim1 * dim2
  let data = Array::make(size, Complex::zero())
  { dim0, dim1, dim2, data }
}

///|
/// Get tensor element at (i, j, k)
pub fn Tensor3D::get(self : Tensor3D, i : Int, j : Int, k : Int) -> Complex {
  let idx = i * self.dim1 * self.dim2 + j * self.dim2 + k
  self.data[idx]
}

///|
/// Set tensor element at (i, j, k)
pub fn Tensor3D::set(
  self : Tensor3D,
  i : Int,
  j : Int,
  k : Int,
  val : Complex,
) -> Unit {
  let idx = i * self.dim1 * self.dim2 + j * self.dim2 + k
  self.data[idx] = val
}

///|
/// Compute matrix exponential exp(M) for 2x2 matrix via Taylor series
pub fn matrix2x2_exp(m : Matrix2x2, terms : Int) -> Matrix2x2 {
  let mut res = Matrix2x2::identity()
  let mut term = Matrix2x2::identity()

  for k = 1; k <= terms; k = k + 1 {
    term = mul_matrix2x2(term, m)
    let inv_k = 1.0 / k.to_double()
    let scaled_term = {
      m00: term.m00.scale(inv_k),
      m01: term.m01.scale(inv_k),
      m10: term.m10.scale(inv_k),
      m11: term.m11.scale(inv_k),
    }
    term = scaled_term
    res = {
      m00: res.m00.add(term.m00),
      m01: res.m01.add(term.m01),
      m10: res.m10.add(term.m10),
      m11: res.m11.add(term.m11),
    }
  }
  res
}