module Quox.NatExtra

import public Data.Nat
import public Data.Nat.Views
import Data.Nat.Division
import Data.SnocList
import Data.Vect
import Syntax.PreorderReasoning

%default total

infixl 8 `shiftL`, `shiftR`
infixl 7 .&.
infixl 6 `xor`
infixl 5 .|.


public export
data LTE' : Nat -> Nat -> Type where
  LTERefl  : LTE' n n
  LTESuccR : LTE' m n -> LTE' m (S n)
%builtin Natural LTE'

public export %hint
lteZero' : {n : Nat} -> LTE' 0 n
lteZero' {n = 0}   = LTERefl
lteZero' {n = S n} = LTESuccR lteZero'

public export %hint
lteSucc' : LTE' m n -> LTE' (S m) (S n)
lteSucc' LTERefl      = LTERefl
lteSucc' (LTESuccR p) = LTESuccR $ lteSucc' p

public export
fromLte : {n : Nat} -> LTE m n -> LTE' m n
fromLte LTEZero     = lteZero'
fromLte (LTESucc p) = lteSucc' $ fromLte p

public export
toLte : {n : Nat} -> m `LTE'` n -> m `LTE` n
toLte LTERefl      = reflexive
toLte (LTESuccR p) = lteSuccRight (toLte p)


private
0 baseNZ : n `GTE` 2 => NonZero n
baseNZ @{LTESucc _} = SIsNonZero

parameters {base : Nat} {auto 0 _ : base `GTE` 2} (chars : Vect base Char)
  private
  showAtBase' : List Char -> Nat -> List Char
  showAtBase' acc 0 = acc
  showAtBase' acc k =
    let dig = natToFinLT (modNatNZ k base baseNZ) @{boundModNatNZ {}} in
    showAtBase' (index dig chars :: acc)
                (assert_smaller k $ divNatNZ k base baseNZ)

  export
  showAtBase : Nat -> String
  showAtBase = pack . showAtBase' []

export
showHex : Nat -> String
showHex = showAtBase $ fromList $ unpack "0123456789ABCDEF"


export
0 notEvenOdd : (a, b : Nat) -> Not (a + a = S (b + b))
notEvenOdd 0     b prf = absurd prf
notEvenOdd (S a) b prf =
  notEvenOdd b a $ Calc $
    |~ b + b
    ~~ a + S a   ..<(inj S prf)
    ~~ S (a + a) ..<(plusSuccRightSucc {})

export
0 doubleInj : (m, n : Nat) -> m + m = n + n -> m = n
doubleInj 0     0     _   = Refl
doubleInj (S m) (S n) prf =
  cong S $ doubleInj m n $
  inj S $ Calc $
    |~ S (m + m)
    ~~ m + S m   ...(plusSuccRightSucc {})
    ~~ n + S n   ...(inj S prf)
    ~~ S (n + n) ..<(plusSuccRightSucc {})

export
0 halfDouble : (n : Nat) -> half (n + n) = HalfEven n
halfDouble n with (half (n + n)) | (n + n) proof nn
  _ | HalfOdd  k | S (k + k) = void $ notEvenOdd n k nn
  _ | HalfEven k | k + k     = rewrite doubleInj n k nn in Refl

export
floorHalf : Nat -> Nat
floorHalf k = case half k of
  HalfOdd  n => n
  HalfEven n => n


||| like in intercal ☺
|||
||| take all the bits of `subj` that are set in `mask`, and squish them down
||| towards the lsb
public export
select : (mask, subj : Nat) -> Nat
select mask subj = go 1 (halfRec mask) subj 0 where
  go : forall mask. Nat -> HalfRec mask -> Nat -> Nat -> Nat
  go bit HalfRecZ            subj res = res
  go bit (HalfRecEven _ rec) subj res = go bit rec (floorHalf subj) res
  go bit (HalfRecOdd  _ rec) subj res = case half subj of
    HalfOdd  subj => go (bit + bit) rec subj (res + bit)
    HalfEven subj => go (bit + bit) rec subj res

||| take the i least significant bits of subj (where i = popCount mask),
||| and place them where mask's set bits are
|||
||| left inverse of select if mask .|. subj = mask
public export
spread : (mask, subj : Nat) -> Nat
spread mask subj = go 1 (halfRec mask) subj 0 where
  go : forall mask. Nat -> HalfRec mask -> Nat -> Nat -> Nat
  go bit HalfRecZ            subj res = res
  go bit (HalfRecEven _ rec) subj res = go (bit + bit) rec subj res
  go bit (HalfRecOdd  _ rec) subj res = case half subj of
    HalfOdd  subj => go (bit + bit) rec subj (res + bit)
    HalfEven subj => go (bit + bit) rec subj res



public export
data BitwiseRec : Nat -> Nat -> Type where
  BwDone : BitwiseRec 0 0
  Bw00   : (m, n : Nat) -> Lazy (BitwiseRec m n) ->
           BitwiseRec (m + m) (n + n)
  Bw01   : (m, n : Nat) -> Lazy (BitwiseRec m n) ->
           BitwiseRec (m + m) (S (n + n))
  Bw10   : (m, n : Nat) -> Lazy (BitwiseRec m n) ->
           BitwiseRec (S (m + m)) (n + n)
  Bw11   : (m, n : Nat) -> Lazy (BitwiseRec m n) ->
           BitwiseRec (S (m + m)) (S (n + n))

export
bitwiseRec : (m, n : Nat) -> BitwiseRec m n
bitwiseRec m n = go (halfRec m) (halfRec n) where
  go : forall m, n. HalfRec m -> HalfRec n -> BitwiseRec m n
  go HalfRecZ           HalfRecZ           = BwDone
  go HalfRecZ           (HalfRecEven n nr) = Bw00 0 n $ go HalfRecZ nr
  go HalfRecZ           (HalfRecOdd  n nr) = Bw01 0 n $ go HalfRecZ nr
  go (HalfRecEven m mr) HalfRecZ           = Bw00 m 0 $ go mr HalfRecZ
  go (HalfRecEven m mr) (HalfRecEven n nr) = Bw00 m n $ go mr nr
  go (HalfRecEven m mr) (HalfRecOdd  n nr) = Bw01 m n $ go mr nr
  go (HalfRecOdd  m mr) HalfRecZ           = Bw10 m 0 $ go mr HalfRecZ
  go (HalfRecOdd  m mr) (HalfRecEven n nr) = Bw10 m n $ go mr nr
  go (HalfRecOdd  m mr) (HalfRecOdd  n nr) = Bw11 m n $ go mr nr

public export
bitwise : (Bool -> Bool -> Bool) -> Nat -> Nat -> Nat
bitwise f m n = go 1 (bitwiseRec m n) 0 where
  one : Bool -> Bool -> Nat -> Nat -> Nat
  one p q bit res = if f p q then bit + res else res
  go : forall m, n. Nat -> BitwiseRec m n -> Nat -> Nat
  go bit BwDone         res = res
  go bit (Bw00 m n rec) res = go (bit + bit) rec $ one False False bit res
  go bit (Bw01 m n rec) res = go (bit + bit) rec $ one False True  bit res
  go bit (Bw10 m n rec) res = go (bit + bit) rec $ one True  False bit res
  go bit (Bw11 m n rec) res = go (bit + bit) rec $ one True  True  bit res

public export
(.&.) : Nat -> Nat -> Nat
(.&.) = bitwise $ \p, q => p && q

private %foreign "scheme:blodwen-and"
primAnd : Nat -> Nat -> Nat
%transform "NatExtra.(.&.)" NatExtra.(.&.) m n = primAnd m n

public export
(.|.) : Nat -> Nat -> Nat
(.|.) = bitwise $ \p, q => p || q

private %foreign "scheme:blodwen-or"
primOr : Nat -> Nat -> Nat
%transform "NatExtra.(.|.)" NatExtra.(.|.) m n = primOr m n

public export
xor : Nat -> Nat -> Nat
xor = bitwise (/=)

private %foreign "scheme:blodwen-xor"
primXor : Nat -> Nat -> Nat
%transform "NatExtra.xor" NatExtra.xor m n = primXor m n


public export
shiftL : Nat -> Nat -> Nat
shiftL n 0     = n
shiftL n (S i) = shiftL (n + n) i

private %foreign "scheme:blodwen-shl"
primShiftL : Nat -> Nat -> Nat
%transform "NatExtra.shiftL" NatExtra.shiftL n i = primShiftL n i

public export
shiftR : Nat -> Nat -> Nat
shiftR n 0     = n
shiftR n (S i) = shiftL (floorHalf n) i

private %foreign "scheme:blodwen-shr"
primShiftR : Nat -> Nat -> Nat
%transform "NatExtra.shiftR" NatExtra.shiftR n i = primShiftR n i