import Base
import ./math.bend as Math

# Rat.bend — exact small rational arithmetic on U32 sign-magnitude.
#
# Representation: R{neg, num, den} means (-1)^neg * num / den.
# mk_rat always divides out Math.gcd(num, den) and canonicalizes zero to
# R{False{}, 0, 1}, so every value built through mk_rat is in reduced form.
#
# Setoid-vs-decidable note: rational equality in general is a setoid
# (many representations per value), but here reduced forms are canonical,
# so req is a DECIDABLE Boolean equality: it cross-multiplies
# (n1*d2 == n2*d1) and treats any zero magnitude as equal regardless of
# sign (mirroring pyval int_eq_raw). req is sound even on unnormalized
# R values, provided denominators are nonzero.
#
# Bounds / preconditions (exactness is for SMALL values only):
# - den != 0 is a precondition everywhere; den = 0 behavior is UNSPECIFIED
#   (mk_rat maps (n != 0, 0) to R{neg, 1, 0} and (0, 0) to R{False{}, 0, 1};
#   rdiv by a zero numerator inverts to den = 0 — all degenerate).
# - Every intermediate product (n1*d2, n2*d1, d1*d2, num_same) must not
#   wrap mod 2^32; results are exact only under that bound.
# - Cost: each mk_rat pays one Math.gcd (64 steps, one U32.div/mod per
#   step, ~32 sub-steps each).
#
# Selection rule vs fixed.bend: exact small fractions here (rat); sustained
# fractional or performance-sensitive work there (fixed Q16.16).
#
# NOTE (radd_comm): a general commutativity law was attempted and DROPPED
# (see note at the end of this file). Closed instances below are kept.

type Rat is Data:
  R{neg: Bool, num: U32, den: U32}

# mk_rat: reduce by gcd, canonicalize zero. Straight-line, no match.
def mk_rat(neg: Bool, +num: U32, +den: U32) -> Rat:
  +g = Math.gcd(num, den)
  +n2 = U32.div(num, g)
  +d2 = U32.div(den, g)
  +zero = U32.is_zero(num)
  R{Bool.pick(Bool, zero, False{}, neg),
    Bool.pick(U32, zero, 0, n2),
    Bool.pick(U32, zero, 1, d2)}

# radd: signed addition, straight-line picks (no match on computed).
# Same sign: mk_rat(s1, p1 + p2, den). Different signs: the larger
# magnitude wins (cmp on the cross products, exact absent wrap); a tie
# is zero. Follows the pyval int_add_raw shape.
def radd(a: Rat, b: Rat) -> Rat:
  match a b:
    case R{+s1, n1, +d1} R{+s2, n2, +d2}:
      +p1 = U32.mul(n1, d2)
      +p2 = U32.mul(n2, d1)
      +same = Bool.not(Bool.xor(s1, s2))
      +c = U32.cmp(p1, p2)
      +num_same = U32.add(p1, p2)
      +den = U32.mul(d1, d2)
      +dif1 = U32.sub(p1, p2)
      +dif2 = U32.sub(p2, p1)
      Bool.pick(Rat, same, mk_rat(s1, num_same, den),
        Bool.pick(Rat, Cmp.is_lt(c), mk_rat(s2, dif2, den),
          Bool.pick(Rat, Cmp.is_gt(c), mk_rat(s1, dif1, den),
            mk_rat(False{}, 0, 1))))

# rsub: negate b's sign, then radd (rebuild is a constructor, not a match
# on computed; radd matches on its own parameters).
def rsub(a: Rat, b: Rat) -> Rat:
  match b:
    case R{s2, n2, d2}:
      radd(a, R{Bool.not(s2), n2, d2})

# rmul: xor signs, multiply through, normalize.
def rmul(a: Rat, b: Rat) -> Rat:
  match a b:
    case R{s1, n1, d1} R{s2, n2, d2}:
      mk_rat(Bool.xor(s1, s2), U32.mul(n1, n2), U32.mul(d1, d2))

# rdiv: invert b (swap num/den; requires b's num != 0), then rmul.
def rdiv(a: Rat, b: Rat) -> Rat:
  match b:
    case R{s2, n2, d2}:
      rmul(a, R{s2, d2, n2})

# req: decidable equality (cross-multiply + both-zero collapse).
def req(a: Rat, b: Rat) -> Bool:
  match a b:
    case R{s1, +n1, d1} R{s2, +n2, d2}:
      Bool.or(
        Bool.and(Bool.not(Bool.xor(s1, s2)),
          U32.is_eq(U32.mul(n1, d2), U32.mul(n2, d1))),
        Bool.and(U32.is_zero(n1), U32.is_zero(n2)))

# --- extractors ---

def neg_of(r: Rat) -> Bool:
  match r:
    case R{neg, n, d}:
      neg

def num_of(r: Rat) -> U32:
  match r:
    case R{neg, n, d}:
      n

def den_of(r: Rat) -> U32:
  match r:
    case R{neg, n, d}:
      d

# --- proven laws (closed instances, cross-checked with fractions.Fraction:
# 1/2+1/3=5/6, 2/3*3/4=1/2, 3/4-1/4=1/2, (1/2)/(1/4)=2) ---

law radd_1_2_1_3:
  {radd(mk_rat(False{}, 1, 2), mk_rat(False{}, 1, 3))
    == mk_rat(False{}, 5, 6) : Rat}

def radd_1_2_1_3():
  {==}

law rmul_2_3_3_4:
  {rmul(mk_rat(False{}, 2, 3), mk_rat(False{}, 3, 4))
    == mk_rat(False{}, 1, 2) : Rat}

def rmul_2_3_3_4():
  {==}

law rsub_3_4_1_4:
  {rsub(mk_rat(False{}, 3, 4), mk_rat(False{}, 1, 4))
    == mk_rat(False{}, 1, 2) : Rat}

def rsub_3_4_1_4():
  {==}

law req_sym_1_2_1_3:
  {req(mk_rat(False{}, 1, 2), mk_rat(False{}, 1, 3))
    == req(mk_rat(False{}, 1, 3), mk_rat(False{}, 1, 2)) : Bool}

def req_sym_1_2_1_3():
  {==}

law rdiv_1_2_1_4:
  {rdiv(mk_rat(False{}, 1, 2), mk_rat(False{}, 1, 4))
    == mk_rat(False{}, 2, 1) : Rat}

def rdiv_1_2_1_4():
  {==}

# NOTE (radd_comm, general law, DROPPED after 2 failed attempts):
# {radd(a, b) == radd(b, a)} is kept as closed instances only.
# Attempt 1 (general law, nested sign-split, definitional {==}): hard
# checker error — a match is allowed only on def parameters, never on
# pattern-bound binders (s1 s2), so the signs cannot be split after
# matching the Rats.
# Attempt 2 (same-sign helper law radd_comm_ff + %U32.add_comm rewrite):
# the check never completes (worker killed, no output) — symbolically
# unfolding radd/mk_rat/gcd under a rewrite ascription is too heavy;
# and mathematically it would stall one step later at the denominator
# swap d1*d2 == d2*d1, which needs U32.mul_comm. Base proves U32.add_comm
# (via Word.add_comm) but lists NO U32.mul_comm (full U32 def listing
# verified), and the different-sign arms would additionally need
# subtraction lemmas. Keeping the closed instances above; the radd
# function itself is unaffected.