perm.idr

import Data.Fin
import Data.Vect
import Prelude.WellFounded

%default total

data Tree : Nat -> Type -> Type where
  Leaf : Tree Z a
  Node : a -> (n : Nat) -> Tree n a -> Tree m a -> Tree (S (n + m)) a

data Parity : Nat -> Type where
  Even : {n:Nat} -> Parity (n + n)
  Odd : {n:Nat} -> Parity (S (n + n))

parity : (n : Nat) -> Parity n
parity Z     = Even {n=Z}
parity (S Z) = Odd {n=Z}
parity (S (S k)) with (parity k)
  parity (S (S (j + j))) | Even =
      rewrite plusSuccRightSucc (S j) j
      in Even {n=S j}
  parity (S (S (S (j + j)))) | Odd  = prf j (Odd {n=S j})
    where
      prf : (j : Nat) -> Parity (S (S j + S j)) -> Parity (S (S (S (j + j))))
      prf j = rewrite plusSuccRightSucc (S j) j in id

accIndLt : {P : Nat -> Type}
        -> (step : (x : Nat)
        -> ((y : Nat) -> LT y x -> P y) -> P x)
        -> (z : Nat)
        -> Accessible LT z
        -> P z
accIndLt {P} step z (Access f) = step z $ \y, lt => accIndLt {P} step y (f y lt)

wfIndLt : {P : Nat -> Type}
       -> (step : (x : Nat)
       -> ((y : Nat) -> LT y x -> P y) -> P x)
       -> (x : Nat)
       -> P x
wfIndLt step x = accIndLt step x (ltAccessible x)

spanTree : (sz : Nat) -> Tree sz (Fin sz)
spanTree Z = Leaf
spanTree (S sz) = wfIndLt go (S sz) (the (Fin (S Z)) FZ)
  where

    ltDouble : (n : Nat) -> Not (n = Z) -> LT n (n + n)
    ltDouble Z p = void (p Refl)
    ltDouble (S k) p =
        LTESucc (rewrite plusCommutative k (S k) in lteAddRight (S k))

    stepRelEven : (n : Nat) -> LT (S n) (S (S n + S n))
    stepRelEven n = lteSuccRight (ltDouble (S n) SIsNotZ)

    stepRelOdd : (n : Nat) -> LT (S n) (S (S (n + n)))
    stepRelOdd n = rewrite plusSuccRightSucc (S n) n in ltDouble (S n) SIsNotZ

    stepRearrange1 : (n : Nat)
                  -> (l : Nat)
                  -> l + S (S (n + S n)) = S (n + S l) + S n
    stepRearrange1 n l =
           rewrite plusCommutative n (S l)
        in rewrite sym (plusAssociative l n (S n))
        in rewrite sym (plusCommutative (S (S (n + S n))) l)
        in rewrite plusCommutative (n + S n) l
        in Refl

    stepRearrange2 : (n : Nat)
                  -> (l : Nat)
                  -> l + S (S (n + S n)) = (l + S (S n)) + (S n)
    stepRearrange2 n l =
           rewrite sym (plusAssociative l (S (S n)) (S n))
        in rewrite sym (plusSuccRightSucc (S (S n)) n)
        in Refl

    stepRearrange3 : (n : Nat)
                  -> (l : Nat)
                  -> l + (S (S (n + n))) = (S (n + S l)) + n
    stepRearrange3 n l =
           rewrite plusCommutative n (S l)
        in rewrite sym (plusAssociative l n n)
        in rewrite sym (plusCommutative (S (S (n + n))) l)
        in rewrite plusCommutative (n + n) l
        in Refl

    stepRearrange4 : (n : Nat)
                  -> (l : Nat)
                  -> l + (S (S (n + n))) = (l + S n) + (S n)
    stepRearrange4 n l =
           rewrite sym (plusAssociative l (S n) (S n))
        in rewrite plusSuccRightSucc (S n) n
        in Refl

    stepRearrange5 : (n : Nat)
                  -> (l : Nat)
                  -> l + (S (S (n + n))) = (S n + S l) + n
    stepRearrange5 n l =
              rewrite plusCommutative n (S l)
           in rewrite sym (plusAssociative l n n)
           in rewrite plusSuccRightSucc l (n + n)
           in rewrite plusSuccRightSucc l (S (n + n))
           in Refl

    go : (sz : Nat)
      -> (step : (rsz : Nat)
              -> LT rsz sz
              -> {rl : Nat}
              -> Fin (S rl)
              -> Tree rsz (Fin (rl + rsz)))
      -> {l: Nat} -> Fin (S l) -> Tree sz (Fin (l + sz))
    go Z step {l} b = Leaf
    go (S sz) step {l} b = case parity sz of
      Even {n=Z}   => Node (rewrite plusCommutative l 1 in b) Z Leaf Leaf
      Even {n=S n} =>
          Node (rewrite stepRearrange1 n l in weakenN (S n) (shift (S n) b))
               (S n)
               (rewrite stepRearrange2 n l
                in step (S n) (stepRelEven n) (weakenN (S (S n)) b))
               (rewrite stepRearrange1 n l
                in step (S n) (stepRelEven n) (shift (S (S n)) b))
      Odd {n} =>
          Node (rewrite stepRearrange5 n l in weakenN n (shift (S n) b))
               (S n)
               (rewrite stepRearrange4 n l
                in step (S n) (stepRelOdd n) (weakenN (S n) b))
               (rewrite stepRearrange3 n l
                in step n (lteSuccLeft (stepRelOdd n)) (shift (S (S n)) b))

toList : Tree sz a -> List a
toList Leaf = []
toList (Node x n y z) = toList y ++ [x] ++ toList z

example10 : map Data.Fin.finToNat (toList (spanTree 10)) = the (List Nat) [0..9]
example10 = Refl

example9 : map Data.Fin.finToNat (toList (spanTree 9)) = the (List Nat) [0..8]
example9 = Refl

maxView : a -> (n:Nat) -> {m : Nat} -> Tree n a -> Tree m a -> (a, Tree (m+n) a)
maxView y n l tr = case tr of
  Leaf => (y, l)
  (Node x xs {m} xl xr) =>
    case maxView x xs xl xr of
      (ny,nr) =>
           rewrite plusCommutative xs m
        in rewrite plusCommutative (m + xs) n
        in (ny, Node y n l nr)

merge : Tree n a -> Tree m a -> Tree (n+m) a
merge Leaf Leaf = Leaf
merge {n} l Leaf = rewrite plusZeroRightNeutral n in l
merge Leaf r = r
merge (Node y ys yl {m} yr) r =
    rewrite plusCommutative ys m
    in case maxView y ys yl yr of
      (key, l') => Node key (m+ys) l' r

pop : (i : Nat) -> Tree (S n) a -> {prf:LTE i n} -> (Tree n a, a)
pop {n} i tr {prf} = wfIndLt {P=P} go n i tr prf
  where
    prfLTE : (x : Nat) -> (y : Nat) -> LTE x (x + y)
    prfLTE Z y = LTEZero
    prfLTE (S k) y = LTESucc (prfLTE k y)

    treeConv : Tree (x + S y) a -> Tree (S (x + y)) a
    treeConv {x} {y} tr = rewrite plusSuccRightSucc x y in tr

    unTreeConv : Tree (S (x + y)) a -> Tree (x + S y) a
    unTreeConv {x} {y} tr = rewrite sym (plusSuccRightSucc x y) in tr

    lteBack : LTE (S x) (S y) -> LTE x y
    lteBack (LTESucc prf) = prf

    prfWalk : LTE (y + S x) (y + m) -> LTE (S x) m
    prfWalk {y = Z} prf = prf
    prfWalk {y = (S k)} (LTESucc x) = prfWalk {y=k} x

    prfSyn : (x : Nat) -> (y : Nat) -> LTE (S (x + y)) ((x + S y) + m)
    prfSyn x y = rewrite sym (plusSuccRightSucc x y) in lteAddRight (S (x + y))

    prfSyn2 : {y : Nat} -> {right : Nat} -> LTE (S right) (y + S right)
    prfSyn2 {y} {right} =
        rewrite plusCommutative y (S right)
        in lteAddRight (S right)

    P : Nat -> Type
    P x = (i : Nat) -> Tree (S x) a -> (LTE i x) -> (Tree x a, a)

    go : (x : Nat)
      -> (step : (y : Nat)
              -> LT y x
              -> (i' : Nat)
              -> Tree (S y) a
              -> (LTE i' y)
              -> (Tree y a, a))
      -> (i : Nat)
      -> Tree (S x) a
      -> (LTE i x)
      -> (Tree x a, a)
    go (s + m) step i (Node j s l r) prflt with (cmp i s)
      go n step x (Node j (x + S y) l r) prflt | CmpLT y =
          case step (x+y) (prfSyn x y) x (treeConv l) (prfLTE x y) of
            (l',i') =>
              rewrite sym (plusSuccRightSucc x y)
              in (Node j (x+y) l' r, i')
      go n step i (Node j i l r) prflt | CmpEQ = (merge l r, j)
      go n step (y + S x) (Node j y l r) prflt | CmpGT x =
          case prfWalk prflt of
            LTESucc nprf => case step _ prfSyn2 x r nprf of
              (r',i') => (unTreeConv (Node j _ l r'), i')

popFin : Tree (S n) a -> Fin (S n) -> (Tree n a, a)
popFin tr fn = pop (finToNat fn) tr {prf=finToNatLT fn}
  where
    finToNatLT : (x: Fin (S n)) -> LTE (finToNat x) n
    finToNatLT FZ = LTEZero
    finToNatLT {n = Z} (FS x) = FinZElim x
    finToNatLT {n = (S k)} (FS x) = LTESucc (finToNatLT x)

data Permutation : Nat -> Type where
  Nil : Permutation Z
  (::) : Fin (S n) -> Permutation n -> Permutation (S n)

toVector : Permutation n -> Vect n (Fin n)
toVector {n} = go n (spanTree n)
  where
    go : (n:Nat) -> Tree n a -> Permutation n -> Vect n a
    go Z tr [] = []
    go (S n) tr (x::xs) = case popFin tr x of
      (tr',y) => y :: go n tr' xs

example3 : map Data.Fin.finToNat (toVector [1,0,0]) = [1,0,2]
example3 = Refl