@no_std
@settings(defaultLengthUnit = mm, kclVersion = 1.0, experimentalFeatures = allow)

import assert from "std"
import * from "std::array"
import pow, sqrt from "std::math"

/// Adds every element of u to its corresponding element in v.
/// Both vectors must have the same length. Returns a new vector of the same length.
/// In other words, component-wise addition.
/// ```kcl
/// u = [1, 2, 3]
/// v = [10, 10, 10]
/// v2 = vector::add(u, v)
/// assert(v2[0], isEqualTo = 11)
/// assert(v2[1], isEqualTo = 12)
/// assert(v2[2], isEqualTo = 13)
/// ```
@(feature_tree = false)
export fn add(@u: [number], v: [number]): [number] {
  l0 = count(u)
  l1 = count(v)
  assert(l0, isEqualTo = l1, error = "The two input vectors must be the same length")
  arrayRange = [0..<l0]
  return reduce(arrayRange, initial = [], f = fn(@i, accum) {
    return push(accum, item = u[i] + v[i])
  })
}

/// Subtracts from every element of u its corresponding element in v.
/// Both vectors must have the same length. Returns a new vector of the same length.
/// In other words, component-wise subtraction.
/// ```kcl
/// u = [10, 10, 10]
/// v = [1, 2, 3]
/// v2 = vector::sub(u, v)
/// assert(v2[0], isEqualTo = 9)
/// assert(v2[1], isEqualTo = 8)
/// assert(v2[2], isEqualTo = 7)
/// ```
@(feature_tree = false)
export fn sub(@u: [number], v: [number]): [number] {
  l0 = count(u)
  l1 = count(v)
  assert(l0, isEqualTo = l1, error = "The two input vectors must be the same length")
  arrayRange = [0..<l0]
  return reduce(arrayRange, initial = [], f = fn(@i, accum) {
    return push(accum, item = u[i] - v[i])
  })
}


/// Multiplies every element of u by its corresponding element in v.
/// Both vectors must have the same length. Returns a new vector of the same length.
/// In other words, component-wise multiplication.
/// ```kcl
/// u = [10, 10, 10]
/// v = [1, 2, 3]
/// v2 = vector::mul(u, v)
/// assert(v2[0], isEqualTo = 10)
/// assert(v2[1], isEqualTo = 20)
/// assert(v2[2], isEqualTo = 30)
/// ```
@(feature_tree = false)
export fn mul(@u: [number], v: [number]): [number] {
  l0 = count(u)
  l1 = count(v)
  assert(l0, isEqualTo = l1, error = "The two input vectors must be the same length")
  arrayRange = [0..<l0]
  return reduce(arrayRange, initial = [], f = fn(@i, accum) {
    return push(accum, item = u[i] * v[i])
  })
}


/// Divides every element of u by its corresponding element in v.
/// Both vectors must have the same length. Returns a new vector of the same length.
/// In other words, component-wise division.
/// ```kcl
/// u = [10, 10, 10]
/// v = [1, 2, 3]
/// v2 = vector::div(u, v)
/// assert(v2[0], isEqualTo = 10)
/// assert(v2[1], isEqualTo = 5)
/// assert(v2[2], isEqualTo = 3.333, tolerance = 0.01)
/// ```
@(feature_tree = false)
export fn div(@u: [number], v: [number]): [number] {
  l0 = count(u)
  l1 = count(v)
  assert(l0, isEqualTo = l1, error = "The two input vectors must be the same length")
  arrayRange = [0..<l0]
  return reduce(arrayRange, initial = [], f = fn(@i, accum) {
    return push(accum, item = u[i] / v[i])
  })
}

/// Find the cross product of two 3D points or vectors.
/// ```kcl
/// vx = [1, 0, 0]
/// vy = [0, 1, 0]
/// vz = vector::cross(vx, v = vy)
/// assert(vz[0], isEqualTo = 0)
/// assert(vz[1], isEqualTo = 0)
/// assert(vz[2], isEqualTo = 1)
/// ```
@(feature_tree = false)
export fn cross(@u: Point3d, v: Point3d) {
  return [
    u[1] * v[2] - (u[2] * v[1]),
    u[2] * v[0] - (u[0] * v[2]),
    u[0] * v[1] - (u[1] * v[0])
  ]
}

/// Find the dot product of two points or vectors of any dimension.
/// ```kcl
/// u = [1, 2, 3]
/// v = [4, -5, 6]
/// dotprod = vector::dot(u, v)
/// assert(dotprod, isEqualTo = 12)
/// ```
@(feature_tree = false)
export fn dot(@u: [number], v: [number]): number {
  l0 = count(u)
  l1 = count(v)
  assert(l0, isEqualTo = l1, error = "The two input vectors must be the same length")
  return reduce(
    [0..<l0],
    initial = 0,
    f = fn(@i, accum) {
      item = u[i] * v[i]
      return accum + item
    },
  )
}

/// Find the Euclidean distance of a vector.
/// ```kcl
/// v = [3, 4]
/// m = vector::magnitude(v)
/// assert(m, isEqualTo = 5)
/// ```
@(feature_tree = false)
export fn magnitude(@v: [number]): number {
  return reduce(
    [0..<count(v)],
    initial = 0,
    f = fn(@i, accum) {
      item = pow(v[i], exp = 2) 
      return accum + item
    },
  )
  |> sqrt()
}

/// Normalize a vector (with any number of dimensions)
/// ```kcl
/// v = [3, 4]
/// normed = vector::normalize(v)
/// assert(normed[0], isEqualTo = 0.6)
/// assert(normed[1], isEqualTo = 0.8)
/// ```
@(feature_tree = false)
export fn normalize(@v: [number]): [number] {
  mag = magnitude(v)
  return map(v, f = fn(@i) { return i/mag })
}
