# tinygrad/shape.bend -- pure shape arithmetic for tinygrad-bend.
#
# Port of the shape contracts of upstream tinygrad:
#   - broadcasting: tensor.py `_broadcasted` / UOp broadcast (align right, 1s expand)
#   - matmul:       mixin/op.py `dot`        ((...,m,k) @ (...,k,n) -> (...,m,n))
#   - reduce:       mixin/reduce.py `_reduce` (drop the reduced dim)
#   - pad/shrink:   mixin/movement.py         (pad adds, shrink removes; grad duality)
#
# Everything in this file is total, terminating, and free of F32, so LAWS.bend
# can state laws about it and PROOF.bend can prove them. Dims are >= 1: empty
# (size-0) tensors are not modeled (documented deviation, see README).

import Base

# Shapes are lists of dims, row-major (C order), like tinygrad.
# def Shape: List<&2, Nat>  (kept inline; List<Nat> below means List<&2, Nat>)

# numel -- product of dims. Mirrors UOp.shape prod / Tensor numel.
def numel(+s: List<&2, Nat>) -> Nat:
  match s:
    case Nil{}:
      1n
    case Con{h, t}:
      Nat.mul(h, numel(t))

# maxn -- structural max of two Nats. Own def (not Base's Nat.max, which is
# pick/cmp-based and resists induction); laws talk about this one.
def maxn(+a: Nat, +b: Nat) -> Nat:
  match a b:
    case 0n, _:
      b
    case 1n+a2, 0n:
      1n+a2
    case 1n+a2, 1n+b2:
      1n+maxn(a2, b2)

# bcast -- common broadcast shape of s and t (numpy-style: ranks align right,
# each dim is the max; a dim of 1 stretches). Mirrors tinygrad `_broadcasted`.
def bcast(+s: List<&2, Nat>, +t: List<&2, Nat>) -> List<&2, Nat>:
  match s t:
    case Nil{}, Nil{}:
      Nil{}
    case Nil{}, Con{b, bs}:
      b <> bs
    case Con{a, as}, Nil{}:
      a <> as
    case Con{a, as}, Con{b, bs}:
      maxn(a, b) <> bcast(as, bs)

# bcast_ok -- can s and t broadcast against each other?
def dim_ok(c: Cmp, d: Nat, e: Nat) -> Bool:
  match c:
    case LT{}:
      Nat.is_eq(e, 1n)
    case EQ{}:
      True{}
    case GT{}:
      Nat.is_eq(d, 1n)

def bcast_ok(+s: List<&2, Nat>, +t: List<&2, Nat>) -> Bool:
  match s t:
    case Nil{}, _:
      True{}
    case _, Nil{}:
      True{}
    case Con{a, as}, Con{b, bs}:
      Bool.and(dim_ok(Nat.cmp(a, b), a, b), bcast_ok(as, bs))

# bdim2 -- the dominance core: d is the peeled dim. d2 == 0n means the source
# dim was 1 (stretches to anything); otherwise the dims must shrink together.
def bdim2(d2: Nat, e2: Nat) -> Type:
  match d2 e2:
    case 0n, _:
      Unit
    case 1n+d3, 0n:
      Empty
    case 1n+d3, 1n+e3:
      bdim2(d3, e3)

def bdim(d: Nat, e: Nat) -> Type:
  match d e:
    case 0n, _:
      Empty
    case 1n+d2, 0n:
      Empty
    case 1n+d2, 1n+e2:
      bdim2(d2, e2)

def bc(+s: List<&2, Nat>, +t: List<&2, Nat>) -> Type:
  match s t:
    case Nil{}, _:
      Unit
    case Con{_, _}, Nil{}:
      {numel(s) == 1n : Nat}
    case Con{a, as}, Con{b, bs}:
      bdim(a, b) & bc(as, bs)

# matmul_out -- shape law of dot: (m,k) @ (k,n) -> (m,n).
def matmul_out(+m: Nat, +k: Nat, +n: Nat) -> List<&2, Nat>:
  m <> n <> Nil{}

# pad_shape -- pad each dim: new = lo + dim + hi (arg = los, his lists).
def pad_shape(+s: List<&2, Nat>, +los: List<&2, Nat>, +his: List<&2, Nat>) -> List<&2, Nat>:
  match s los his:
    case Nil{}, _, _:
      Nil{}
    case _, Nil{}, _:
      s
    case _, _, Nil{}:
      s
    case Con{h, t}, Con{lo, los2}, Con{hi, his2}:
      Nat.add(lo, Nat.add(h, hi)) <> pad_shape(t, los2, his2)

# shrink_shape -- inverse view: new = dim - lo - hi.
def shrink_shape(+s: List<&2, Nat>, +los: List<&2, Nat>, +his: List<&2, Nat>) -> List<&2, Nat>:
  match s los his:
    case Nil{}, _, _:
      Nil{}
    case _, Nil{}, _:
      s
    case _, _, Nil{}:
      s
    case Con{h, t}, Con{lo, los2}, Con{hi, his2}:
      Nat.sub(Nat.sub(h, lo), hi) <> shrink_shape(t, los2, his2)

# drop_at -- shape of a reduction over axis k (0-based): removes dim k.
# Total: k >= rank leaves the shape untouched (callers never do that).
def drop_at(+s: List<&2, Nat>, +k: Nat) -> List<&2, Nat>:
  match s k:
    case Nil{}, _:
      Nil{}
    case Con{h, t}, 0n:
      t
    case Con{h, t}, 1n+k2:
      h <> drop_at(t, k2)

# at -- dim k of s (1 when out of range), so reduce_numel below is total.
def at(+s: List<&2, Nat>, +k: Nat) -> Nat:
  match s k:
    case Nil{}, _:
      1n
    case Con{h, t}, 0n:
      h
    case Con{h, t}, 1n+k2:
      at(t, k2)

def bdim_idx_if(+d: Nat, +i: Nat, z: Bool) -> Nat:
  match z:
    case True{}:
      0n
    case False{}:
      i

# flat -- row-major flat index of a multi-index (multi-index lists are short;
# rank is bounded by nesting depth of user code, so O(rank) walks are fine).
def flat(+mi: List<&2, Nat>, +s: List<&2, Nat>) -> Nat:
  match mi s:
    case Nil{}, _:
      0n
    case _, Nil{}:
      0n
    case Con{i, it}, Con{d, ds}:
      Nat.add(Nat.mul(i, numel(ds)), flat(it, ds))

def bdim_idx(+d: Nat, +i: Nat) -> Nat:
  bdim_idx_if(d, i, Nat.is_eq(d, 1n))

# bdim_idx -- index into a broadcast dim: a stretched dim (size 1) reads 0.
# bcast_flat -- the source flat index that output multi-index `mi` reads from,
# when out shape is `os` and source shape is `s` (same rank; s dims may be 1).
def bcast_flat_skip(+ods: List<&2, Nat>) -> Nat:
  match ods:
    case Nil{}:
      0n
    case Con{_, t}:
      bcast_flat_skip(t)

def bcast_flat_go(+mi: List<&2, Nat>, +s: List<&2, Nat>, +os: List<&2, Nat>) -> Nat:
  match mi s os:
    case Nil{}, _, _:
      0n
    case _, Nil{}, Nil{}:
      0n
    case _, Nil{}, Con{_, ods}:
      # right-aligned rank difference: leading output dims read past the
      # (implicitly 1-sized) source dims -- nothing to add, keep walking
      bcast_flat_skip(ods)
    case _, _, Nil{}:
      0n
    case Con{i, it}, Con{d, ds}, Con{od, ods}:
      Nat.add(Nat.mul(bdim_idx(d, i), numel(ds)), bcast_flat_go(it, ds, ods))

def bcast_flat(+mi: List<&2, Nat>, +s: List<&2, Nat>, +os: List<&2, Nat>) -> Nat:
  bcast_flat_go(mi, s, os)

# unflat -- multi-index of flat index i in shape s. Head index = i / numel(rest),
# tail recurses on the remainder (row-major; wraps into range, so total).
def unflat_go(+ds: List<&2, Nat>, +i: Nat) -> List<&2, Nat>:
  match ds:
    case Nil{}:
      Nil{}
    case Con{d2, ds2}:
      Nat.div(i, numel(ds2)) <> unflat_go(ds2, Nat.mod(i, numel(ds2)))
def unflat(+i: Nat, +s: List<&2, Nat>) -> List<&2, Nat>:
  unflat_go(s, i)