#define ATS_MAINATSFLAG 1

#include "share/atspre_staload.hats"

staload "libats/libc/SATS/math.sats"

fnx fact {n : nat} .<n>. (k : int(n)) :<> int =
  case+ k of
    | 0 => 1
    | k =>> fact(k - 1) * k

fnx dfact {n : nat} .<n>. (k : int(n)) :<> int =
  case+ k of
    | 0 => 1
    | 1 => 1
    | k =>> k * dfact(k - 2)

// TODO make this more versatile?
fn choose {n : nat}{ m : nat | m <= n } (n : int(n), k : int(m)) : int =
  let
    fun numerator_loop { m : nat | m > 1 } .<m>. (i : int(m)) : int =
      case+ i of
        | 1 => n
        | 2 => (n - 1) * n
        | i =>> (n + 1 - i) * numerator_loop(i - 1)
  in
    case+ k of
      | 0 => 1
      | 1 => n
      | k =>> numerator_loop(k) / fact(k)
  end

// FIXME
fun bad(n : int) : [ m : nat ] int(m) =
  case+ n of
    | 0 => 0
    | n => 1 + bad(n - 1)

fun is_prime(k : intGt(0)) : bool =
  case+ k of
    | 1 => false
    | k => 
      begin
        let
          var pre_bound: int = g0float2int(sqrt_float(g0int2float_int_float(k)))
          var bound: [ m : nat ] int(m) = bad(pre_bound)
          
          fun loop {n : nat}{m : nat} .<max(0,m-n)>. ( i : int(n)
                                                     , bound : int(m)
                                                     ) :<> bool =
            if i < bound then
              if k % i = 0 then
                false
              else
                true && loop(i + 1, bound)
            else
              if i = bound then
                if k % i = 0 then
                  false
                else
                  true
              else
                true
        in
          loop(2, bound)
        end
      end

extern
fun choose_ats {n : nat}{ m : nat | m <= n } : (int(n), int(m)) -> int =
  "mac#"

extern
fun double_factorial {n : nat} : int(n) -> int =
  "mac#"

extern
fun factorial_ats {n : nat} : int(n) -> int =
  "mac#"

extern
fun is_prime_ats { n : nat | n > 0 } : int(n) -> bool =
  "mac#"

implement choose_ats (n, k) =
  choose(n, k)

implement double_factorial (m) =
  dfact(m)

implement is_prime_ats (n) =
  is_prime(n)

implement factorial_ats (m) =
  fact(m)