declare symbol x, y, z

def f1 := function (x)
def g1 := function (x)
def f2 := function (x, y)
def g2 := function (x, y)
def f3 := function (x, y, z)
def g3 := function (x, y, z)
def h3 := function (x, y, z)

--
-- Tensor Arithmetics
--
assertEqual "scalar + tensor"
  (1 + [| 1, 2, 3 |])
  [| 2, 3, 4 |]

assertEqual "tensor + scalar"
  ([| 1, 2, 3 |] + 1)
  [| 2, 3, 4 |]

assertEqual "tensor + tensor (same index)"
  ([| 1, 2, 3 |]_i + [| 1, 2, 3 |]_i)
  [| 2, 4, 6 |]_i

assertEqual "tensor + tensor (outer product)"
  ([| 10, 20, 30 |]_i + [| 1, 2, 3 |]_j)
  [| [| 11, 12, 13 |], [| 21, 22, 23 |], [| 31, 32, 33 |] |]

assertEqual "tensor + 2D tensor"
  ([| 100, 200, 300 |]_i + [|[| 1, 2, 3 |], [| 10, 20, 30 |]|]_j_i)
  [| [| 101, 110 |], [| 202, 220 |], [| 303, 330 |] |]_i_j

assertEqual "2D tensor + 1D tensor"
  ([|[| 11, 12 |], [| 21, 22 |], [| 31, 32 |]|]_i_j + [| 100, 200, 300 |]_i)
  [| [| 111, 112 |], [| 221, 222 |], [| 331, 332 |] |]_i_j

--
-- Derivative
--
assertEqual "partial derivative of f(x,y,z)"
  (∂/∂ f3 x)
  (∂/∂ f3 x)

assertEqual "derivative of vector function"
  (∂/∂ [| f1, g1 |] x)
  [| ∂/∂ f1 x, ∂/∂ g1 x |]

assertEqual "gradient of f(x,y,z)"
  (∂/∂ f3 [| x, y, z |])
  [| ∂/∂ f3 x, ∂/∂ f3 y, ∂/∂ f3 z |]

assertEqual "apply partial derivatives"
  ([| (\e -> ∂/∂ e x), (\e -> ∂/∂ e y) |] f2)
  [| ∂/∂ f2 x, ∂/∂ f2 y |]

assert "Jacobian matrix"
  (show ([| (\e -> ∂/∂ e x), (\e -> ∂/∂ e y) |] [| f2, g2 |]) = show [| [| ∂/∂ f2 x, ∂/∂ g2 x |], [| ∂/∂ f2 y, ∂/∂ g2 y |] |])

--
-- Nabla (uses ∇ from lib/math/analysis/derivative.egi)
--
assertEqual "nabla f"
  (∇ f2 [| x, y |])
  [| ∂/∂ f2 x, ∂/∂ f2 y |]

assertEqual "nabla vector"
  [| ∂/∂ f2 [| x, y |], ∂/∂ g2 [| x, y |] |]
  [| [| ∂/∂ f2 x, ∂/∂ f2 y |], [| ∂/∂ g2 x, ∂/∂ g2 y |] |]

--
-- Contraction (uses trace from lib/math/algebra/vector.egi)
--
assertEqual "element-wise product"
  (contract ([|1, 2, 3|]~i * [|10, 20, 30|]_i))
  [10, 40, 90]

assertEqual "trace of matrix"
  (trace [|[|10, 20, 30|], [|20, 40, 60|], [|30, 60, 90|]|])
  140

--
-- Divergence (uses div from lib/math/algebra/vector.egi)
--
assertEqual "divergence"
  (div [| f3, g3, h3 |] [| x, y, z |])
  (∂/∂ f3 x + ∂/∂ g3 y + ∂/∂ h3 z)

--
-- Taylor Expansion
--
def multivariateTaylorExpansion fexpr xs ys :=
  withSymbols [h]
    let hs := generateTensor (\[x] -> h_x) (tensorShape xs)
     in map2
          (*)
          (map (\n -> 1 / fact n) nats0)
          (map
             (compose
                (\e -> V.substitute xs ys e)
                (\e -> V.substitute hs (withSymbols [i] xs_i - ys_i) e))
             (iterate (compose (\e -> ∇ e xs) (\e -> V.* hs e)) fexpr))

def taylorExpansion fexpr x a := multivariateTaylorExpansion fexpr [|x|] [|a|]

assert "Taylor expansion of f(x)"
  (show (take 3 (taylorExpansion f1 x 0)) = "[f1 0, x * f1|1 0, x^2 * f1|1|1 0 / 2]")

assert "Multivariate Taylor expansion"
  (show (take 3 (multivariateTaylorExpansion f2 [| x, y |] [| 0, 0 |])) = "[f2 0 0, x * f2|1 0 0 + y * f2|2 0 0, (x^2 * f2|1|1 0 0 + x * y * f2|2|1 0 0 + y * x * f2|1|2 0 0 + y^2 * f2|2|2 0 0) / 2]")