plum

module std
import std/Str
import std/Bool
import std/Result

MIN_VALUE = -0x8000_0000_0000_0000 # Lowest value of Int
MAX_VALUE = 0x7FFF_FFFF_FFFF_FFFF  # Highest value of Int
LARGE = 1 << 28 # 2**28

E = 2.718281828459045f # Euler's number, the base of natural logarithms, e, https://oeis.org/A001113
LN10 = 2.302585092994046f # The natural logarithm of 10, https://oeis.org/A002392
LN2 = 0.6931471805599453f # The natural logarithm of 2, https://oeis.org/A002162
LOG10E = 0.4342944819032518f # The base 10 logarithm of e, formula: 1 / LN10
LOG2E = 1.4426950408889634f # The base 2 logarithm of e, formula: 1 / LN2
PI = 3.141592653589793f # The ratio of the circumference of a circle to its diameter, https://oeis.org/A000796
PHI = 1.618033988749895f # https://oeis.org/A001622
SQRT1_2 = 0.7071067811865476f # The square root of 1/2
SQRT2 = 1.4142135623730951f # The square root of 2, https://oeis.org/A002193
SQRT_E = 1.6487212707001282f # https://oeis.org/A019774
SQRT_PI = 1.7724538509055159f # https://oeis.org/A002161
SQRT_PHI = 1.272019649514069f # https://oeis.org/A139339
EPSILON = 2.220446049250313e-16f # The difference between 1 and the smallest floating point number greater than 1, formula: 7/3 - 4/3 - 1
MIN_FLOAT_VALUE = 4.9406564584124654417656879286822137236505980e-324 # Lowest value of float
MAX_FLOAT_VALUE = 1.79769313486231570814527423731704356798070e+308 # Highest value of float
HALF_PI = 1.5707963267948966f # PI / 2
TAU = 6.283185307179586f # 2 * PI

# `Int`/`Float` bare-wrap into `Number` with no wrapper syntax at all (see
# `plum-checker`'s `unify`/`monomorphize::wrapPrimitiveAgainstExpected`) —
# assigning/returning/passing a bare `Int`/`Float` value wherever `Number` is
# expected just works. A method not found on `Int`/`Float` directly (they no
# longer have their own method tables at all) falls back to `Number`'s,
# boxing `self` first (see `plum-checker`/`plum-wasm-codegen`'s matching
# `AttrKind::Method` dispatch fallback) — so `x.abs()` for a bare `Int`/
# `Float` `x` dispatches here exactly as if `Int`/`Float` still had their own
# `abs` method.
enum Number =
  | Int
  | Float

  fun kind(self) -> Str =
    match self
      Int(_) => "Int"
      Float(_) => "Float"

  fun toFloatValue(self) -> Float =
    match self
      Int(i) => Float(i)
      Float(f) => f

  fun abs(self) -> Number =
    match self
      Int(i) => Int(i < 0 ? -i : i)
      Float(f) => Float(f < 0.0f ? -f : f)

  # Preserves 0.0 / -0.0 / NaN for the `Float` case, matching Go's
  # `math.Signbit`-adjacent `Copysign`/`sign` conventions loosely.
  fun sign(self) -> Number =
    match self
      Int(i) => Int(i > 0 ? 1 : i < 0 ? -1 : 0)
      Float(f) =>
        if f > 0.0f
          return Float(1.0f)
        if f < 0.0f
          return Float(-1.0f)
        return Float(f)

  # An `Int` already IS its own well-distributed bit pattern — used by
  # `Map[K: Hashable, V]` to bucket `Int` keys. A `Float` hashes via a
  # truncating conversion (no bit-level float hashing attempted).
  fun hash(self) -> Int =
    match self
      Int(i) => i
      Float(f) => Int(f)

  # An Int is always already whole, so trunc/floor/ceil/round are all just
  # the value itself widened to Float.
  fun trunc(self) -> Float =
    match self
      Int(i) => Float(i)
      # `Int(f)` truncates toward zero (see `plum-wasm-codegen`'s `Int(x)`
      # cross-type conversion) — exactly `trunc`'s definition.
      Float(f) => Float(Int(f))

  fun floor(self) -> Float =
    match self
      Int(i) => Float(i)
      Float(f) =>
        t := Float(Int(f))
        if f < 0.0f && t != f
          return t - 1.0f
        return t

  fun ceil(self) -> Float =
    match self
      Int(i) => Float(i)
      Float(f) =>
        t := Float(Int(f))
        if f > 0.0f && t != f
          return t + 1.0f
        return t

  # Round-half-away-from-zero (matching Go's `math.Round`), not
  # round-half-to-even.
  fun round(self) -> Float =
    match self
      Int(i) => Float(i)
      Float(f) =>
        neg := f < 0.0f
        v := neg ? -f : f
        r := {v + 0.5f}.floor()
        return neg ? -r : r

  fun log(self) -> Float =
    ln(self.toFloatValue())

  fun log2(self) -> Float =
    ln(self.toFloatValue()) * LOG2E

  fun log10(self) -> Float =
    ln(self.toFloatValue()) * LOG10E

  # floor(log2(|self|)). An `Int` self computes this exactly with integer
  # division rather than through Float rounding error; a `Float` self falls
  # back to the real logarithm.
  fun logb(self) -> Float =
    match self
      Int(i) =>
        if i == 0
          return -1.0f / 0.0f
        n := i < 0 ? -i : i
        k := 0
        while n >= 2
          n = n / 2
          k = k + 1
        return Float(k)
      Float(f) =>
        if f == 0.0f
          return -1.0f / 0.0f
        return {ln(f < 0.0f ? -f : f) * LOG2E}.floor()

  fun sqrt(self) -> Float =
    sqrt(self.toFloatValue())

  # self^y. Mirrors the original `Int.pow`/`Float.pow`: a whole-number `y`
  # (including negative) is computed exactly by repeated squaring; a
  # fractional `y` falls back to self^y = exp(y * ln(self)), which requires
  # self > 0.
  fun pow(self, y: Float) -> Float =
    powFloat(self.toFloatValue(), y)

  # Check whether this number is finite, ie not +/-infinity and not NaN.
  fun isFinite(self) -> Bool =
    !self.isNaN() && !self.isInfinite()

  fun isInfinite(self) -> Bool =
    match self
      Int(_) => False
      Float(f) => f > MAX_FLOAT_VALUE || f < -MAX_FLOAT_VALUE

  fun isNaN(self) -> Bool =
    match self
      Int(_) => False
      Float(f) => f != f

  fun min(self, other: Number) -> Number =
    self.toFloatValue() < other.toFloatValue() ? self : other

  fun max(self, other: Number) -> Number =
    self.toFloatValue() > other.toFloatValue() ? self : other

  # Inverse hyperbolic cosine, via acosh(x) = ln(x + sqrt(x^2 - 1)), x >= 1.
  fun acosh(self) -> Float =
    v := self.toFloatValue()
    if v < 1.0f
      return 0.0f / 0.0f
    return ln(v + sqrt(v * v - 1.0f))

  fun sinh(self) -> Float =
    v := self.toFloatValue()
    return {exp(v) - exp(-v)} / 2.0f

  fun cosh(self) -> Float =
    v := self.toFloatValue()
    return {exp(v) + exp(-v)} / 2.0f

  fun tanh(self) -> Float =
    return self.sinh() / self.cosh()

  # Inverse hyperbolic sine: asinh(x) = ln(x + sqrt(x^2 + 1)).
  fun asinh(self) -> Float =
    v := self.toFloatValue()
    return ln(v + sqrt(v * v + 1.0f))

  # Inverse hyperbolic tangent: atanh(x) = 0.5*ln((1+x)/(1-x)), |x| < 1.
  fun atanh(self) -> Float =
    v := self.toFloatValue()
    if v <= -1.0f || v >= 1.0f
      return 0.0f / 0.0f
    return 0.5f * ln({1.0f + v} / {1.0f - v})

  # Rounded to 6 fractional digits, with trailing zeros trimmed (but always
  # at least one digit after the point, e.g. `3.0` not `3.`), for a `Float`;
  # plain decimal digits for an `Int`.
  fun toStr(self) -> Str =
    match self
      Int(i) => intToStr(i)
      Float(f) =>
        if f != f
          return "NaN"
        if f > MAX_FLOAT_VALUE || f < -MAX_FLOAT_VALUE
          return f < 0.0f ? "-Infinity" : "Infinity"
        neg := f < 0.0f
        abs_val := neg ? -f : f
        int_part := Int(abs_val)
        frac_val := abs_val - Float(int_part)
        scaled := Int(frac_val * 1000000.0f + 0.5f)
        carry := scaled >= 1000000
        whole_part := carry ? int_part + 1 : int_part
        frac := carry ? 0 : scaled
        s := intToStr(whole_part) + "." + trimmedFracDigits(frac, 6)
        return neg ? "-" + s : s

# ---- free-function helpers (no `self` to dispatch on) ----

# Plain decimal rendering of an `Int` — factored out of `Number.toStr` so the
# `Float` branch can reuse it (via `intToStr(whole_part)`) without
# re-dispatching back through `Number.toStr` itself.
fun intToStr(n: Int) -> Str =
  if n == 0
    return "0"
  neg := n < 0
  v := neg ? -n : n
  s := digitsToStr(v)
  return neg ? "-" + s : s

# Builds the decimal digits of a positive Int, most significant first, by
# peeling off one digit at a time from the least significant end (so the
# recursion emits digits in reverse call order — the leading, smallest-place
# digit is appended last).
fun digitsToStr(n: Int) -> Str =
  if n == 0
    return ""
  return digitsToStr(n / 10) + digitChar(n % 10)

fun digitChar(d: Int) -> Str =
  match d
    0 => "0"
    1 => "1"
    2 => "2"
    3 => "3"
    4 => "4"
    5 => "5"
    6 => "6"
    7 => "7"
    8 => "8"
    9 => "9"
    _ => "0"

# Renders `n`'s digits most-significant-first, always emitting exactly as many
# digits as `place_value` has powers of ten (leading zeros included) — used to
# render a fractional part at a fixed width regardless of its own leading zeros.
fun fixedFracDigits(n: Int, place_value: Int) -> Str =
  if place_value == 0
    return ""
  return digitChar(n / place_value % 10) + fixedFracDigits(n, place_value / 10)

fun pow10(n: Int) -> Int =
  if n <= 0
    return 1
  return 10 * pow10(n - 1)

# `fixedFracDigits(n, pow10(digits - 1))`, minus however many trailing-zero
# digits `n` has — but always at least one digit ("0" if `n` is 0 outright).
fun trimmedFracDigits(n: Int, digits: Int) -> Str =
  if n == 0 || digits <= 1
    return digitChar(n % 10)
  if n % 10 == 0
    return trimmedFracDigits(n / 10, digits - 1)
  return fixedFracDigits(n, pow10(digits - 1))

# Shared `base^y` algorithm behind `Number.pow`, extracted so both the `Int`
# and `Float` self cases (which only differ in how `base` was obtained) share
# one body.
fun powFloat(base: Float, y: Float) -> Float =
  if y == 0.0f
    return 1.0f
  yi := Int(y)
  if Float(yi) == y
    neg := yi < 0
    n := neg ? -yi : yi
    result := 1.0f
    b := base
    m := n
    while m > 0
      if m % 2 == 1
        result = result * b
      b = b * b
      m = m / 2
    return neg ? 1.0f / result : result
  if base < 0.0f
    return 0.0f / 0.0f
  return exp(y * ln(base))

# Parses a decimal integer, with an optional leading `+`/`-`.
fun parseInt(s: Str) -> Result[Int, Str] =
  len := s.length()
  if len == 0
    return Err("empty string")
  neg := s.byteAt(0) == 45
  start := neg || s.byteAt(0) == 43 ? 1 : 0
  if start >= len
    return Err("invalid integer: '{s}'")
  value := 0
  i := start
  while i < len
    b := s.byteAt(i)
    if b < 48 || b > 57
      return Err("invalid integer: '{s}'")
    value = value * 10 + b - 48
    i = i + 1
  return Ok(neg ? -value : value)

# Parses a decimal float, with an optional leading `+`/`-` and an optional
# `.` fractional part (no exponent notation).
fun parseFloat(s: Str) -> Result[Float, Str] =
  len := s.length()
  if len == 0
    return Err("empty string")
  neg := s.byteAt(0) == 45
  start := neg || s.byteAt(0) == 43 ? 1 : 0
  if start >= len
    return Err("invalid float: '{s}'")
  int_part := 0.0f
  saw_digit := False
  i := start
  while i < len && s.byteAt(i) != 46
    b := s.byteAt(i)
    if b < 48 || b > 57
      return Err("invalid float: '{s}'")
    int_part = int_part * 10.0f + Float(b - 48)
    saw_digit = True
    i = i + 1
  frac_part := 0.0f
  frac_scale := 1.0f
  if i < len && s.byteAt(i) == 46
    i = i + 1
    while i < len
      b = s.byteAt(i)
      if b < 48 || b > 57
        return Err("invalid float: '{s}'")
      frac_scale = frac_scale / 10.0f
      frac_part = frac_part + Float(b - 48) * frac_scale
      saw_digit = True
      i = i + 1
  if !saw_digit
    return Err("invalid float: '{s}'")
  value := int_part + frac_part
  return Ok(neg ? -value : value)

# A raw, host-provided pseudo-random 64-bit Int (xorshift64*, seeded from the
# wall clock — not cryptographically secure). Every other random-number
# function in std is built on top of this one host extern.
extern fun rawRandomInt() -> Int

# A random Float uniformly distributed in [0.0, 1.0).
fun random() -> Float =
  return Float(rawRandomInt() & MAX_VALUE) / Float(MAX_VALUE)

# A random Int uniformly distributed in [0, n). Returns 0 for n <= 0.
fun randomInt(n: Int) -> Int =
  if n <= 0
    return 0
  return { rawRandomInt() & MAX_VALUE } % n

# Natural exponential, e^x, via range reduction (halve x until |x| <= 0.5,
# run the Taylor series there where it converges fast, then square the
# result back up the same number of halvings).
fun exp(x: Float) -> Float =
  if x != x
    return x
  if x > 700.0f
    return 1.0f / 0.0f
  if x < -700.0f
    return 0.0f
  v := x
  k := 0
  while v > 0.5f || v < -0.5f
    v = v / 2.0f
    k = k + 1
  term := 1.0f
  sum := 1.0f
  n := 1
  while n < 25
    term = term * v / Float(n)
    sum = sum + term
    n = n + 1
  result := sum
  i := 0
  while i < k
    result = result * result
    i = i + 1
  return result

# Natural logarithm, via range reduction to v in [1, 2) plus the
# fast-converging series ln(v) = 2*atanh((v-1)/(v+1)).
fun ln(x: Float) -> Float =
  if x != x || x < 0.0f
    return 0.0f / 0.0f
  if x == 0.0f
    return -1.0f / 0.0f
  if x > MAX_FLOAT_VALUE
    return x
  v := x
  k := 0
  while v >= 2.0f
    v = v / 2.0f
    k = k + 1
  while v < 1.0f
    v = v * 2.0f
    k = k - 1
  t := {v - 1.0f} / {v + 1.0f}
  t2 := t * t
  term := t
  sum := t
  n := 1
  while n < 30
    term = term * t2
    sum = sum + term / Float(2 * n + 1)
    n = n + 1
  return Float(k) * LN2 + 2.0f * sum

# Square root via Newton's method, iterating to a fixed point.
fun sqrt(x: Float) -> Float =
  if x < 0.0f
    return 0.0f / 0.0f
  if x == 0.0f || x != x
    return x
  guess := x
  prev := 0.0f
  i := 0
  while guess != prev && i < 100
    prev = guess
    guess = 0.5f * {guess + x / guess}
    i = i + 1
  return guess

# Reduces `x` into roughly `[-PI, PI]` by subtracting the nearest multiple of
# `TAU` — the range `sin`/`cos`'s Taylor series below actually converge
# quickly over. For very large `|x|` (many multiples of `TAU`), floating-point
# cancellation in `x - k*TAU` loses precision the same way any naive
# range-reduction by subtraction does; a real libm uses extended-precision
# constants to avoid this, which isn't attempted here.
fun reduceToPi(x: Float) -> Float =
  k := Int(x / TAU + {x >= 0.0f ? 0.5f : -0.5f})
  return x - Float(k) * TAU

fun sin(x: Float) -> Float =
  if x != x || x > MAX_FLOAT_VALUE || x < -MAX_FLOAT_VALUE
    return 0.0f / 0.0f
  r := reduceToPi(x)
  r2 := r * r
  term := r
  sum := r
  n := 1
  while n < 10
    term = term * {-r2} / Float({2 * n} * {2 * n + 1})
    sum = sum + term
    n = n + 1
  return sum

fun cos(x: Float) -> Float =
  if x != x || x > MAX_FLOAT_VALUE || x < -MAX_FLOAT_VALUE
    return 0.0f / 0.0f
  return sin(x + HALF_PI)

fun tan(x: Float) -> Float =
  return sin(x) / cos(x)

# Arctangent, via repeated argument-halving (`atan(x) = 2*atan(x / (1 +
# sqrt(1+x^2)))`) until `|x| <= 0.5` (where the Taylor series below converges
# quickly), then doubling the result back up the same number of times.
fun atan(x: Float) -> Float =
  if x != x
    return x
  neg := x < 0.0f
  v := neg ? -x : x
  k := 0
  while v > 0.5f && k < 8
    v = v / {1.0f + sqrt(1.0f + v * v)}
    k = k + 1
  v2 := v * v
  term := v
  sum := v
  n := 1
  while n < 20
    term = term * {-v2}
    sum = sum + term / Float(2 * n + 1)
    n = n + 1
  scale := Float(1 << k)
  result := sum * scale
  return neg ? -result : result

fun asin(x: Float) -> Float =
  if x != x || x < -1.0f || x > 1.0f
    return 0.0f / 0.0f
  if x == 1.0f
    return HALF_PI
  if x == -1.0f
    return -HALF_PI
  return atan(x / sqrt(1.0f - x * x))

fun acos(x: Float) -> Float =
  return HALF_PI - asin(x)

# Angle (in radians) of the point `(x, y)` from the origin, in the correct
# quadrant for any sign combination of `x`/`y` (unlike plain `atan(y/x)`,
# which can't distinguish opposite quadrants).
fun atan2(y: Float, x: Float) -> Float =
  if x > 0.0f
    return atan(y / x)
  if x < 0.0f && y >= 0.0f
    return atan(y / x) + PI
  if x < 0.0f && y < 0.0f
    return atan(y / x) - PI
  if x == 0.0f && y > 0.0f
    return HALF_PI
  if x == 0.0f && y < 0.0f
    return -HALF_PI
  return 0.0f

fun hypot(a: Float, b: Float) -> Float =
  return sqrt(a * a + b * b)

# Cube root via Newton's method (fixed iteration count, unlike `sqrt`'s
# converge-to-a-fixed-point loop, since the cubic update step doesn't reach an
# exact fixed point in float precision as reliably as the quadratic one does).
fun cbrt(x: Float) -> Float =
  if x == 0.0f || x != x
    return x
  neg := x < 0.0f
  v := neg ? -x : x
  guess := v
  i := 0
  while i < 60
    guess = {2.0f * guess + v / {guess * guess}} / 3.0f
    i = i + 1
  return neg ? -guess : guess