module Data.Term.Zipper

import Data.DPair
import Data.Fin.Occurs
import Data.Fin.Strong
import Data.Vect.Properties.Map
import Data.Term
import Syntax.PreorderReasoning

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

public export
insertAt' : SFin k -> a -> Vect (pred k) a -> Vect k a
insertAt' FZ x xs = x :: xs
insertAt' (FS i) x (y :: xs) = y :: insertAt' i x xs

public export
deleteAt' : SFin k -> Vect k a -> Vect (pred k) a
deleteAt' FZ (x :: xs) = xs
deleteAt' (FS i) (x :: xs) = x :: deleteAt' i xs

mapInsertAt :
  (0 f : a -> b) ->
  (i : SFin k) ->
  (x : a) ->
  (xs : Vect (pred k) a) ->
  map f (insertAt' i x xs) = insertAt' i (f x) (map f xs)
mapInsertAt f FZ x xs = Refl
mapInsertAt f (FS i) x (y :: xs) = cong (f y ::) $ mapInsertAt f i x xs

insertAtDeleteAt :
  (i : SFin k) ->
  (xs : Vect k a) ->
  insertAt' i (index (strongToFin i) xs) (deleteAt' i xs) = xs
insertAtDeleteAt FZ (x :: xs) = Refl
insertAtDeleteAt (FS i) (x :: xs) = cong (x ::) $ insertAtDeleteAt i xs

indexInsertAt :
  (i : SFin k) ->
  (x : a) ->
  (xs : Vect (pred k) a) ->
  index (strongToFin i) (insertAt' i x xs) = x
indexInsertAt FZ x xs = Refl
indexInsertAt (FS i) x (y :: xs) = indexInsertAt i x xs

-- Definition ------------------------------------------------------------------

public export
data Zipper : Signature -> Nat -> Type where
  Top : Zipper sig n
  Op : sig.Operator k -> SFin k -> Vect (pred k) (Term sig n) -> Zipper sig n -> Zipper sig n

%name Zipper zip

-- Monoid ----------------------------------------------------------------------

export
(++) : Zipper sig n -> Zipper sig n -> Zipper sig n
Top ++ zip2 = zip2
Op op i ts zip1 ++ zip2 = Op op i ts (zip1 ++ zip2)

-- Properties

export
leftUnit : (zip : Zipper sig n) -> Top ++ zip = zip
leftUnit zip = Refl

export
rightUnit : (zip : Zipper sig n) -> zip ++ Top = zip
rightUnit Top = Refl
rightUnit (Op op i ts zip) = cong (Op op i ts) (rightUnit zip)

export
assoc : (zip1, zip2, zip3 : Zipper sig n) -> (zip1 ++ zip2) ++ zip3 = zip1 ++ (zip2 ++ zip3)
assoc Top zip2 zip3 = Refl
assoc (Op op i ts zip1) zip2 zip3 = cong (Op op i ts) (assoc zip1 zip2 zip3)

-- Action ----------------------------------------------------------------------

-- NOTE: should this be public?
public export
(+) : Zipper sig n -> Term sig n -> Term sig n
Top + t = t
Op op i ts zip + t = Op op (insertAt' i (zip + t) ts)

-- Properties

export
actionAssoc :
  (zip1, zip2 : Zipper sig n) ->
  (t : Term sig n) ->
  (zip1 ++ zip2) + t = zip1 + (zip2 + t)
actionAssoc Top zip2 t = Refl
actionAssoc (Op op i ts zip1) zip2 t =
  cong (\t => Op op (insertAt' i t ts)) $ actionAssoc zip1 zip2 t

-- Substitution ----------------------------------------------------------------

export
(<$>) : (SFin k -> Term sig n) -> Zipper sig k -> Zipper sig n
f <$> Top = Top
f <$> Op op i ts zip = Op op i (map (f <$>) ts) (f <$> zip)

-- Properties

export
actionHomo :
  (0 f : SFin k -> Term sig n) ->
  (zip : Zipper sig k) ->
  (0 t : Term sig k) ->
  f <$> zip + t = (f <$> zip) + (f <$> t)
actionHomo f Top t = Refl
actionHomo f (Op op i ts zip) t = cong (Op op) $ Calc $
  |~ map (f <$>) (insertAt' i (zip + t) ts)
  ~~ insertAt' i (f <$> zip + t) (map (f <$>) ts)           ...(mapInsertAt (f <$>) i (zip + t) ts)
  ~~ insertAt' i ((f <$> zip) + (f <$> t)) (map (f <$>) ts) ...(cong (\t => insertAt' i t (map (f <$>) ts)) $ actionHomo f zip t)

export
invertActionTop : (zip : Zipper sig k) -> (0 _ : f <$> zip = Top) -> zip = Top
invertActionTop Top prf = Refl

-- Cycles ----------------------------------------------------------------------

pivotAt : sig.Operator k -> SFin k -> Vect k (Term sig n) -> Zipper sig n
pivotAt op i ts = Op op i (deleteAt' i ts) Top

actionPivotAt :
  (op : sig.Operator k) ->
  (i : SFin k) ->
  (ts : Vect k (Term sig n)) ->
  pivotAt op i ts + index (strongToFin i) ts = Op op ts
actionPivotAt op i ts = cong (Op op) $ insertAtDeleteAt i ts

pivotNotTop : (zip : Zipper sig k) -> Not (zip ++ pivotAt op i ts = Top)
pivotNotTop (Op op' j us zip) prf impossible

export
noCycles : (zip : Zipper sig k) -> (t : Term sig k) -> zip + t = t -> zip = Top
noCycles Top t prf = Refl
noCycles (Op {k} op i ts zip') t@(Op {k = j} op' us) prf =
  void $ pivotNotTop zip' prf''
  where
  i' : Fin k
  i' = strongToFin i

  I : SFin j
  I = replace {p = SFin} (opInjectiveLen prf) i

  I' : Fin j
  I' = strongToFin I

  prf' : ((zip' ++ pivotAt op' I us) + index I' us = index I' us)
  prf' = Calc $
    |~ (zip' ++ pivotAt op' I us) + index I' us
    ~~ zip' + (pivotAt op' I us + index I' us)      ...(actionAssoc zip' (pivotAt op' I us) (index I' us))
    ~~ zip' + Op op' us                             ...(cong (zip' +) $ actionPivotAt op' I us)
    ~~ index i' (insertAt' i (zip' + Op op' us) ts) ..<(indexInsertAt i (zip' + Op op' us) ts)
    ~~ index i' (rewrite opInjectiveLen prf in us)  ...(cong (index i') $ opInjectiveTs prf)
    ~~ index I' us                                  ...(rewrite opInjectiveLen prf in Refl)

  prf'' : zip' ++ pivotAt op' I us = Top
  prf'' = noCycles (zip' ++ pivotAt op' I us) (assert_smaller t (index I' us)) prf'