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