-- basics
def matmul [n][m][p] 'a
           (add: a -> a -> a) (mul: a -> a -> a) (zero: a)
           (A: [n][m]a) (B: [m][p]a) : [n][p]a =
  map (\A_row ->
         map (\B_col ->
                reduce add zero (map2 mul A_row B_col))
             (transpose B))
      A
def matmul_f64 = matmul (+) (*) 0f64
def sum = reduce (+) 0f64
def dot u v = sum (map2 (*) u v)

-- sets the last row to 1
def unitlst [n][m] (X: *[n][m]f64): *[n][m]f64 =
  X with [n-1] = tabulate m (const 1)
def unitlstT [n][m] (X: *[n][m]f64): *[n][m]f64 = transpose (unitlst (transpose X))

-- Bregman generators
def phi (z: f64) : f64 = z * f64.log z - z
def psi [d] (x: [d]f64) : f64 = sum (map ((*x[d-1]) <-< phi <-< (/(f64.epsilon+x[d-1]))) x)

def embed [n][d][k] (B: [n][k]f64) (U: [k][d]f64) : [n][d]f64 =
  let C = tabulate_2d k k (\i j -> (if i == j then 1f64 else 0f64) - f64.from_fraction 1 k)
  in (matmul_f64 (matmul_f64 (unitlstT (copy B)) C) U)

def loss [n][d][k] (X: [n][d]f64) (B: [n][k]f64) (U: [k][d]f64) : f64 =
  let Y = map (map f64.exp) (embed B U)
  in sum (map2 (\x y -> - psi y - dot (map2 (-) x y) (vjp psi y 1)) X Y)

def dloss [n][d][k] (X: [n][d]f64) (B: [n][k]f64) (U: [k][d]f64) : ([n][k]f64, [k][d]f64) = vjp (uncurry (loss X)) (B, U) 1

-- Approximate quantiles
def quantile [n] (q: f64) (xs: [n]f64) : f64 =
  let (p, _, _) = loop (_, xs, idx) = (0/0, xs, f64.to_i64 (q*f64.from_fraction n 1)) while length xs > 1 do
    let pivot = head xs
    let (left, right) = partition (<=pivot) (tail xs)
    let m = 1+length left
    in if m > idx
      then (pivot, left, idx)
      else (pivot, right, idx - m)
  in p

-- Optimiser
def rmsprop [n] (iters: i32) (beta: f64) (par: [n]f64) (df: [n]f64 -> [n]f64) : [n]f64 =
  let step (i, par, s) =
    let d = df par
    let s' = map2 (\s x -> beta * s + (1-beta)*x**2) s d
    let par' = map3 (\x s d -> x - 1e-3 * d / (f64.epsilon + f64.sqrt s)) par s' d
    in (i-1, par', s')
  let (_, par, _) = iterate_until (\(i,_,_) -> i <= 0) step (iters, par, replicate n 0)
  in par

def pca [n][d] (iters: i32) (k: i64) (X: [n][d]f64) : ([n][k]f64, [k][d]f64, f64) =
  let mX = sum (flatten X)
  let X' = map (map (/ mX)) X
  let fp x = x - f64.floor x
  let B = tabulate_2d n k (\i j -> if j == k-1 then 1 else 1e-3 * (fp (0.5 + 1.3247 * f64.from_fraction (i * k + j) 1) - 0.5))
  let U = tabulate_2d k d (\i j -> 1e-3 * (fp (0.5 + 1.7548 * f64.from_fraction (i * k + j) 1) - 0.5))
  let pack [n][d][k] (B: [n][k]f64) (U: [k][d]f64) : []f64 = flatten B ++ flatten U
  let unpack xs =
    let (a, b) = split (n*k) xs
    in (unflatten n k a, unflatten k d b)
  let df [m] (par: [m]f64) : [m]f64 = (uncurry pack <-< uncurry (dloss X') <-< unpack) par :> [m]f64
  let init = flatten B ++ flatten U
  let par = rmsprop iters 0.999 init df
  let (Bf, Uf ) = unpack par
  in (Bf, Uf, loss X' Bf Uf)

entry pcaWithQuantile [n][d] (iters: i32) (q: f64) (k: i64) (X: [n][d]f64) : ([n][k]f64, [d][k]f64, [n][d]f64, f64) =
  let qs = map (quantile q) X
  let Y = tabulate_2d n (d+1) (\i j -> if j == d then qs[i] else X[i][j])
  let (B, U, l) = pca iters (1+k) Y
  let Y' = embed B U
  in (map (\x -> x[0:k]) B, (transpose U[0:k])[0:d], map (\x -> map (f64.- x[d]) x[0:d]) Y', l)