# bend-blas: the shaped layer. A Mat<rows, cols> is one row-major
# Array<F32> with its shape in the type, a Vec<n> likewise, and every
# routine below states the shape rule in its signature, so a wrong
# operand is a checker error naming both shapes, not a wrong answer:
#   Mat.gemm(rows, k, cols, alpha, a, b, beta, c) needs a Mat<rows, k>, a
#   Mat<k, cols> and a Mat<rows, cols>; a transpose is a separate def
#   (gemm_t, gemm_tn, gemv_t), since it changes the shape.
# The sizes are Nat arguments that double as the type indices. The raw
# routines of blas.bend and lapack.bend run underneath, unchanged; the
# shaped forms answer the same Result, with .try twins, and consume
# their operands the same way (Mat.clone keeps a copy). This is the
# scipy.linalg to blas.bend's scipy.linalg.blas, plus the shapes.
import Base
import ./blas.bend as B
import ./lapack.bend as L

type Mat<-rows: Nat, -cols: Nat> is Type:
  MatA{a: Array<F32>}

type Vec<-n: Nat> is Type:
  VecA{a: Array<F32>}

# Construction
# ------------

# the depth of the smallest power-of-two block holding count floats
def depth.go(fuel: Nat, +n: U32, +d: Nat, +cap: U32) -> Nat:
  match fuel:
    case 0n:
      d
    case 1n+p:
      Bool.pick(Nat, U32.is_lt(cap, n), depth.go(p, n, 1n+d, (cap * 2 : U32)), d)

def depth(+count: Nat) -> Nat:
  depth.go(40n, U32.from_nat(count), 0n, 1)

def fill.go(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
  match xs:
    case []:
      a
    case x <> rest:
      fill.go(rest, (i + 1 : U32), Array.set(F32, a, i, x))

# an array of the smallest power-of-two size holding count floats,
# the first ones from xs (cut or padded with 0.0), the rest 0.0
def array.of(+count: Nat, xs: List<&2, F32>) -> Array<F32>:
  fill.go(List.take(&2, F32, xs, count), 0, Array.new(F32, depth(count), 0.0))

def take.go(n: Nat, +i: U32, r: Array<F32> & F32, acc: List<&2, F32>) -> List<&2, F32>:
  match n:
    case 0n:
      (a, x) = r
      x <> acc
    case 1n+p:
      (a, x) = r
      take.go(p, (i - 1 : U32), Array.get(F32, a, (i - 1 : U32)), x <> acc)

# the first count floats of an array, as a list
def list.of(+count: Nat, a: Array<F32>) -> List<&2, F32>:
  match count:
    case 0n:
      []
    case 1n+p:
      take.go(p, (U32.from_nat(p) : U32), Array.get(F32, a, U32.from_nat(p)), [])

# xs cut into rows of cols
def rows.of(rows: Nat, +cols: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>:
  match rows:
    case 0n:
      []
    case 1n+p:
      List.take(&2, F32, xs, cols) <> rows.of(p, cols, List.drop(&2, F32, xs, cols))

def F32.show.of(x: F32) -> String:
  F32.show(x)

def row.show(xs: List<&2, F32>) -> String:
  List.show(~&2, ~F32, ~F32.show.of, xs)

# rows * cols floats, row-major, cut or padded with 0.0
def Mat.from_list(+rows: Nat, +cols: Nat, xs: List<&2, F32>) -> Mat<rows, cols>:
  MatA{array.of(Nat.mul(rows, cols), xs)}

def Mat.zeros(+rows: Nat, +cols: Nat) -> Mat<rows, cols>:
  MatA{Array.new(F32, depth(Nat.mul(rows, cols)), 0.0)}

def Mat.to_list(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, F32>:
  MatA{a} = m
  list.of(Nat.mul(rows, cols), a)

def Mat.to_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, List<&2, F32>>:
  rows.of(rows, cols, Mat.to_list(rows, cols, m))

def Mat.show(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> String:
  List.show(~&2, ~List<&2, F32>, ~row.show, Mat.to_rows(rows, cols, m))

def Mat.clone.fin(-rows: Nat, -cols: Nat, r: Array<F32> & Array<F32>) -> Mat<rows, cols> & Mat<rows, cols>:
  (x, y) = r
  (MatA{x}, MatA{y})

# a Mat is affine: two copies, to be opened by matching in a def of its own
def Mat.clone(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols> & Mat<rows, cols>:
  MatA{a} = m
  Mat.clone.fin(rows, cols, Array.clone(F32, a))

def Vec.from_list(+n: Nat, xs: List<&2, F32>) -> Vec<n>:
  VecA{array.of(n, xs)}

def Vec.zeros(+n: Nat) -> Vec<n>:
  VecA{Array.new(F32, depth(n), 0.0)}

def Vec.to_list(+n: Nat, v: Vec<n>) -> List<&2, F32>:
  VecA{a} = v
  list.of(n, a)

def Vec.show(+n: Nat, v: Vec<n>) -> String:
  row.show(Vec.to_list(n, v))

def Vec.clone.fin(-n: Nat, r: Array<F32> & Array<F32>) -> Vec<n> & Vec<n>:
  (x, y) = r
  (VecA{x}, VecA{y})

def Vec.clone(+n: Nat, v: Vec<n>) -> Vec<n> & Vec<n>:
  VecA{a} = v
  Vec.clone.fin(n, Array.clone(F32, a))

# the raw answers, shaped
def Mat.wrap(-rows: Nat, -cols: Nat, r: Result<&1, &1, U32 & String, Array<F32>>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
  IO.pure(Result<&1, &1, U32 & String, Mat<rows, cols>>, Result.map(&1, &1, U32 & String, Array<F32>, Mat<rows, cols>, x => MatA{x}, r))

def Vec.wrap(-n: Nat, r: Result<&1, &1, U32 & String, Array<F32>>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
  IO.pure(Result<&1, &1, U32 & String, Vec<n>>, Result.map(&1, &1, U32 & String, Array<F32>, Vec<n>, x => VecA{x}, r))

# BLAS
# ----

# C = alpha A B + beta C
def Mat.gemm(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
  MatA{x} = a
  MatA{y} = b
  MatA{z} = c
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(0, 0, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))

def Mat.gemm.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
  IO.try(Mat<rows, cols>, Mat.gemm(rows, k, cols, alpha, a, b, beta, c))

# C = alpha A B^T + beta C, B stored cols x k
def Mat.gemm_t(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<cols, k>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
  MatA{x} = a
  MatA{y} = b
  MatA{z} = c
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(0, 1, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))

def Mat.gemm_t.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<cols, k>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
  IO.try(Mat<rows, cols>, Mat.gemm_t(rows, k, cols, alpha, a, b, beta, c))

# C = alpha A^T B + beta C, A stored k x rows
def Mat.gemm_tn(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<k, rows>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
  MatA{x} = a
  MatA{y} = b
  MatA{z} = c
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(1, 0, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))

def Mat.gemm_tn.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<k, rows>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
  IO.try(Mat<rows, cols>, Mat.gemm_tn(rows, k, cols, alpha, a, b, beta, c))

# y = alpha A x + beta y
def Mat.gemv(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<cols>, beta: F32, y: Vec<rows>) -> IO(Result<&1, &1, U32 & String, Vec<rows>>):
  MatA{m} = a
  VecA{xs} = x
  VecA{ys} = y
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<rows>>, B.Blas.sgemv(0, U32.from_nat(rows), U32.from_nat(cols), alpha, m, xs, beta, ys), r => Vec.wrap(rows, r))

def Mat.gemv.try(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<cols>, beta: F32, y: Vec<rows>) -> IO(Vec<rows>):
  IO.try(Vec<rows>, Mat.gemv(rows, cols, alpha, a, x, beta, y))

# y = alpha A^T x + beta y
def Mat.gemv_t(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<rows>, beta: F32, y: Vec<cols>) -> IO(Result<&1, &1, U32 & String, Vec<cols>>):
  MatA{m} = a
  VecA{xs} = x
  VecA{ys} = y
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<cols>>, B.Blas.sgemv(1, U32.from_nat(rows), U32.from_nat(cols), alpha, m, xs, beta, ys), r => Vec.wrap(cols, r))

def Mat.gemv_t.try(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<rows>, beta: F32, y: Vec<cols>) -> IO(Vec<cols>):
  IO.try(Vec<cols>, Mat.gemv_t(rows, cols, alpha, a, x, beta, y))

# A = alpha x y^T + A
def Mat.ger(+rows: Nat, +cols: Nat, alpha: F32, x: Vec<rows>, y: Vec<cols>, a: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
  VecA{xs} = x
  VecA{ys} = y
  MatA{m} = a
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sger(U32.from_nat(rows), U32.from_nat(cols), alpha, xs, ys, m), r => Mat.wrap(rows, cols, r))

def Mat.ger.try(+rows: Nat, +cols: Nat, alpha: F32, x: Vec<rows>, y: Vec<cols>, a: Mat<rows, cols>) -> IO(Mat<rows, cols>):
  IO.try(Mat<rows, cols>, Mat.ger(rows, cols, alpha, x, y, a))

def Vec.dot(+n: Nat, x: Vec<n>, y: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
  VecA{xs} = x
  VecA{ys} = y
  B.Blas.sdot(U32.from_nat(n), xs, ys)

def Vec.dot.try(+n: Nat, x: Vec<n>, y: Vec<n>) -> IO(F32):
  IO.try(F32, Vec.dot(n, x, y))

def Vec.nrm2(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
  VecA{xs} = x
  B.Blas.snrm2(U32.from_nat(n), xs)

def Vec.nrm2.try(+n: Nat, x: Vec<n>) -> IO(F32):
  IO.try(F32, Vec.nrm2(n, x))

def Vec.asum(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
  VecA{xs} = x
  B.Blas.sasum(U32.from_nat(n), xs)

def Vec.asum.try(+n: Nat, x: Vec<n>) -> IO(F32):
  IO.try(F32, Vec.asum(n, x))

def Vec.iamax(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, U32>):
  VecA{xs} = x
  B.Blas.isamax(U32.from_nat(n), xs)

def Vec.iamax.try(+n: Nat, x: Vec<n>) -> IO(U32):
  IO.try(U32, Vec.iamax(n, x))

# y = alpha x + y
def Vec.axpy(+n: Nat, alpha: F32, x: Vec<n>, y: Vec<n>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
  VecA{xs} = x
  VecA{ys} = y
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<n>>, B.Blas.saxpy(U32.from_nat(n), alpha, xs, ys), r => Vec.wrap(n, r))

def Vec.axpy.try(+n: Nat, alpha: F32, x: Vec<n>, y: Vec<n>) -> IO(Vec<n>):
  IO.try(Vec<n>, Vec.axpy(n, alpha, x, y))

def Vec.scal(+n: Nat, alpha: F32, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
  VecA{xs} = x
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<n>>, B.Blas.sscal(U32.from_nat(n), alpha, xs), r => Vec.wrap(n, r))

def Vec.scal.try(+n: Nat, alpha: F32, x: Vec<n>) -> IO(Vec<n>):
  IO.try(Vec<n>, Vec.scal(n, alpha, x))

# LAPACK
# ------

# X = A^-1 B; Fail with the zero pivot as the code when A is singular
def Mat.solve(+n: Nat, +nrhs: Nat, a: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Result<&1, &1, U32 & String, Mat<n, nrhs>>):
  MatA{x} = a
  MatA{y} = b
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, nrhs>>, L.Lapack.sgesv(U32.from_nat(n), U32.from_nat(nrhs), x, y), r => Mat.wrap(n, nrhs, r))

def Mat.solve.try(+n: Nat, +nrhs: Nat, a: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Mat<n, nrhs>):
  IO.try(Mat<n, nrhs>, Mat.solve(n, nrhs, a, b))

# the lower Cholesky factor L of a symmetric positive definite A
def Mat.cholesky(+n: Nat, a: Mat<n, n>) -> IO(Result<&1, &1, U32 & String, Mat<n, n>>):
  MatA{x} = a
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, n>>, L.Lapack.spotrf(U32.from_nat(n), x), r => Mat.wrap(n, n, r))

def Mat.cholesky.try(+n: Nat, a: Mat<n, n>) -> IO(Mat<n, n>):
  IO.try(Mat<n, n>, Mat.cholesky(n, a))

# X = A^-1 B given L from Mat.cholesky
def Mat.cholesky_solve(+n: Nat, +nrhs: Nat, l: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Result<&1, &1, U32 & String, Mat<n, nrhs>>):
  MatA{x} = l
  MatA{y} = b
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, nrhs>>, L.Lapack.spotrs(U32.from_nat(n), U32.from_nat(nrhs), x, y), r => Mat.wrap(n, nrhs, r))

def Mat.cholesky_solve.try(+n: Nat, +nrhs: Nat, l: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Mat<n, nrhs>):
  IO.try(Mat<n, nrhs>, Mat.cholesky_solve(n, nrhs, l, b))

# the singular values, largest first
def Mat.svd_values(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)>>):
  MatA{x} = a
  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<Nat.min(m, n)>>, L.Lapack.sgesvd_s(U32.from_nat(m), U32.from_nat(n), x), r => Vec.wrap(Nat.min(m, n), r))

def Mat.svd_values.try(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Vec<Nat.min(m, n)>):
  IO.try(Vec<Nat.min(m, n)>, Mat.svd_values(m, n, a))

def Mat.svd.fin(-m: Nat, -n: Nat, s: Array<F32>, uv: Array<F32> & Array<F32>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
  match uv:
    case (u, vt):
      IO.pure(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, Done{(VecA{s}, (MatA{u}, MatA{vt}))})

def Mat.svd.wrap(-m: Nat, -n: Nat, r: Result<&1, &1, U32 & String, Array<F32> & (Array<F32> & Array<F32>)>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
  match r:
    case Fail{e}:
      IO.pure(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, Fail{e})
    case Done{v}:
      match v:
        case (s, uv):
          Mat.svd.fin(m, n, s, uv)

# the thin SVD: (s, (u, vt)) with A = u diag(s) vt
def Mat.svd(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
  MatA{x} = a
  IO.bind(Result<&1, &1, U32 & String, Array<F32> & (Array<F32> & Array<F32>)>, Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, L.Lapack.sgesvd(U32.from_nat(m), U32.from_nat(n), x), r => Mat.svd.wrap(m, n, r))

def Mat.svd.try(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)):
  IO.try(Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>), Mat.svd(m, n, a))

def Mat.eigh.wrap(-n: Nat, r: Result<&1, &1, U32 & String, Array<F32> & Array<F32>>) -> IO(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>):
  match r:
    case Fail{e}:
      IO.pure(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, Fail{e})
    case Done{v}:
      match v:
        case (w, vs):
          IO.pure(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, Done{(VecA{w}, MatA{vs})})

# the eigenvalues ascending and the eigenvectors as the rows of the Mat
def Mat.eigh(+n: Nat, a: Mat<n, n>) -> IO(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>):
  MatA{x} = a
  IO.bind(Result<&1, &1, U32 & String, Array<F32> & Array<F32>>, Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, L.Lapack.ssyev(U32.from_nat(n), x), r => Mat.eigh.wrap(n, r))

def Mat.eigh.try(+n: Nat, a: Mat<n, n>) -> IO(Vec<n> & Mat<n, n>):
  IO.try(Vec<n> & Mat<n, n>, Mat.eigh(n, a))