module Encoded.Arith

import Data.Nat
import Encoded.Bool
import Encoded.Pair
import Syntax.PreorderReasoning
import Term.Semantics
import Term.Syntax

-- Utilities -------------------------------------------------------------------

export
rec : {ty : Ty} -> Term (N ~> ty ~> (N * ty ~> ty) ~> ty) ctx
rec = Abs $ Abs $ Abs $
  let n = Var $ There $ There Here in
  let z = Var $ There Here in
  let s = Var Here in
  let 
    p : Term (N * ty) (ctx :< N :< ty :< (N * ty ~> ty))
    p = Rec n
      (App pair [<Zero, z])
      (Abs' (\p =>
        App pair
          [<Suc (App fst [<p])
          , App (shift s) [<p]]))
  in
  App snd [<p]

export
plus : Term (N ~> N ~> N) ctx
plus = Abs' (\n => Rec n Id (Abs' (Abs' Suc .)))

export
pred : Term (N ~> N) ctx
-- pred = Abs' (\n => App rec [<n, Zero, fst])
pred = Abs'
  (\n => App
    (Rec n
      (Const Zero)
      (Abs $ Abs $
        let f = Var (There Here) in
        let k = Var Here in
        Rec k
          (Suc Zero)
          (Const $ Rec (App f [<Zero])
            Zero
            (Const $ Suc $ App f [<Lit 1]))))
    [<Lit 1])

export
minus : Term (N ~> N ~> N) ctx
minus = Abs $ Abs $
  let m = Var $ There Here in
  let n = Var Here in
  Rec n m pred

export
mult : Term (N ~> N ~> N) ctx
mult = Abs' (\n =>
  Rec n
    (Const Zero)
    (Abs $ Abs $
      let f = Var $ There Here in
      let m = Var Here in
      App plus [<m, App f [<m]]))

export
lte : Term (N ~> N ~> B) ctx
lte = Abs' (\m => isZero . App minus [<m])

export
equal : Term (N ~> N ~> B) ctx
equal = Abs $ Abs $
  let m = Var $ There Here in
  let n = Var Here in
  App and [<App lte [<m, n], App lte [<n, m]]

export
divmod : Term (N ~> N ~> N * N) ctx
divmod = Abs $ Abs $
  let n = Var (There Here) in
  let d = Var Here in
  Rec
    n
    (App pair [<Zero, Zero])
    (Abs' (\p =>
      App if'
        [<App lte [<shift d, App (Abs' Suc . snd) [<p]]
        , App pair [<Suc (App fst [<p]), Zero]
        , App mapSnd [<Abs' Suc, p]]))

-- Properties ------------------------------------------------------------------

-- Redefinition in Idris

namespace Sem
  public export
  recHelper : Nat -> a -> ((Nat, a) -> a) -> (Nat, a)
  recHelper k z s = rec k (0, z) (\p => (S (fst p), s p))

  public export
  rec' : Nat -> a -> ((Nat, a) -> a) -> a
  rec' k z s = snd (recHelper k z s)

  public export
  slowPred : Nat -> Nat
  slowPred n = rec' n 0 fst

  public export
  predHelper : Nat -> Nat -> Nat
  predHelper n =
    rec n
      (const 0)
      (\f, k =>
        rec k
          1
          (const $ rec (f 0)
            0
            (const $ S $ f 1)))

  public export
  pred' : Nat -> Nat
  pred' n = predHelper n 1

  public export
  divmodStep : Nat -> (Nat, Nat) -> (Nat, Nat)
  divmodStep d (div, mod) =
    if d <= S mod
    then
      (S div, 0)
    else
      (div, S mod)

  public export
  divmod : Nat -> Nat -> (Nat, Nat)
  divmod n d = rec n (0, 0) (divmodStep d)

-- Proofs

export
fstRecHelper : (k : Nat) -> (0 z : a) -> (0 s : (Nat, a) -> a) -> fst (recHelper k z s) = k
fstRecHelper 0 z s = Refl
fstRecHelper (S k) z s = cong S $ fstRecHelper k z s

export
slowPredCorrect : (k : Nat) -> slowPred k = pred k
slowPredCorrect 0 = Refl
slowPredCorrect (S k) = fstRecHelper k 0 fst

export
predCorrect : (k : Nat) -> pred' k = pred k
predCorrect 0 = Refl
predCorrect 1 = Refl
predCorrect (S (S k)) = cong S $ predCorrect (S k)

-- divmodSmall :
--   (d, div, mod : Nat) ->
--   {auto 0 ok : NonZero d} ->
  -- mod <= d

export
divmodCorrect :
  (n, d : Nat) ->
  {auto 0 ok : NonZero d} ->
  n = fst (divmod n d) * d + snd (divmod n d)
divmodCorrect 0 d = Refl
divmodCorrect (S k) d = ?divmodCorrect_rhs_1