module Quox.Syntax.DimEq

import public Quox.Syntax.Var
import public Quox.Syntax.Dim
import public Quox.Syntax.Subst
import public Quox.Context

import Data.Maybe
import Data.Nat
import Data.DPair
import Data.Fun.Graph
import Decidable.Decidable
import Decidable.Equality
import Derive.Prelude

%language ElabReflection
%default total


public export
DimEq' : Nat -> Type
DimEq' = Context (Maybe . Dim)


public export
data DimEq : Nat -> Type where
  ZeroIsOne : DimEq d
  C         : (eqs : DimEq' d) -> DimEq d
%name DimEq eqs
%runElab deriveIndexed "DimEq" [Eq, Ord, Show]


public export
consistent : DimEq d -> Bool
consistent ZeroIsOne = False
consistent (C eqs)   = True


public export
data IfConsistent : DimEq d -> Type -> Type where
  Nothing :      IfConsistent ZeroIsOne a
  Just    : a -> IfConsistent (C eqs)   a

export
Functor (IfConsistent eqs) where
  map f Nothing  = Nothing
  map f (Just x) = Just (f x)

export
Foldable (IfConsistent eqs) where
  foldr f z Nothing  = z
  foldr f z (Just x) = f x z

export
Traversable (IfConsistent eqs) where
  traverse f Nothing  = pure Nothing
  traverse f (Just x) = Just <$> f x

public export
ifConsistent : Applicative f => (eqs : DimEq d) -> f a -> f (IfConsistent eqs a)
ifConsistent ZeroIsOne act = pure Nothing
ifConsistent (C _)     act = Just <$> act


export
fromGround' : Context' DimConst d -> DimEq' d
fromGround' [<]        = [<]
fromGround' (ctx :< e) = fromGround' ctx :< Just (K e)

export
fromGround : Context' DimConst d -> DimEq d
fromGround = C . fromGround'


public export %inline
zeroEq : DimEq 0
zeroEq = C [<]

public export %inline
new' : {d : Nat} -> DimEq' d
new' {d = 0}   = [<]
new' {d = S d} = new' :< Nothing

public export %inline
new : {d : Nat} -> DimEq d
new = C new'


public export %inline
get' : DimEq' d -> Var d -> Maybe (Dim d)
get' = getWith $ \p, by => map (// by) p

public export %inline
getVar : DimEq' d -> Var d -> Dim d
getVar eqs i = fromMaybe (B i) $ get' eqs i

public export %inline
getShift' : Shift len out -> DimEq' len -> Var len -> Maybe (Dim out)
getShift' = getShiftWith $ \p, by => map (// by) p

public export %inline
get : DimEq' d -> Dim d -> Dim d
get _   (K e) = K e
get eqs (B i) = getVar eqs i


public export %inline
equal : DimEq d -> (p, q : Dim d) -> Bool
equal ZeroIsOne p q = True
equal (C eqs)   p q = get eqs p == get eqs q


infixl 7 :<?
export %inline
(:<?) : DimEq d -> Maybe (Dim d) -> DimEq (S d)
ZeroIsOne :<? d = ZeroIsOne
C eqs     :<? d = C $ eqs :< map (get eqs) d


private %inline
ifVar : Var d -> Dim d -> Maybe (Dim d) -> Maybe (Dim d)
ifVar i p = map $ \q => if q == B i then p else q

-- (using decEq instead of (==) because of the proofs below)
private %inline
checkConst : (e, f : DimConst) -> (eqs : Lazy (DimEq' d)) -> DimEq d
checkConst e f eqs = if isYes $ e `decEq` f then C eqs else ZeroIsOne


export
setConst : Var d -> DimConst -> DimEq' d -> DimEq d
setConst VZ e (eqs :< Nothing)    = C $ eqs :< Just (K e)
setConst VZ e (eqs :< Just (K f)) = checkConst e f $ eqs :< Just (K f)
setConst VZ e (eqs :< Just (B i)) = setConst i e eqs :<? Just (K e)
setConst (VS i) e (eqs :< p)      = setConst i e eqs :<? ifVar i (K e) p

mutual
  private
  setVar' : (i, j : Var d) -> i `LT` j -> DimEq' d -> DimEq d
  setVar' VZ (VS i) LTZ (eqs :< Nothing) =
    C $ eqs :< Just (B i)
  setVar' VZ (VS i) LTZ (eqs :< Just (K e)) =
    setConst i e eqs :<? Just (K e)
  setVar' VZ (VS i) LTZ (eqs :< Just (B j)) =
    setVar i j eqs :<? Just (B (max i j))
  setVar' (VS i) (VS j) (LTS lt) (eqs :< p) =
    setVar' i j lt eqs :<? ifVar i (B j) p

  export %inline
  setVar : (i, j : Var d) -> DimEq' d -> DimEq d
  setVar i j eqs with (compareP i j) | (compare i.nat j.nat)
    setVar i j eqs | IsLT lt | LT = setVar' i j lt eqs
    setVar i i eqs | IsEQ    | EQ = C eqs
    setVar i j eqs | IsGT gt | GT = setVar' j i gt eqs


export %inline
set : (p, q : Dim d) -> DimEq d -> DimEq d
set _ _ ZeroIsOne = ZeroIsOne
set (K e) (K f) (C eqs) = checkConst e f eqs
set (K e) (B i) (C eqs) = setConst i e eqs
set (B i) (K e) (C eqs) = setConst i e eqs
set (B i) (B j) (C eqs) = setVar   i j eqs


public export %inline
Split : Nat -> Type
Split d = (DimEq' d, DSubst (S d) d)

export %inline
split1 : DimConst -> DimEq' (S d) -> Maybe (Split d)
split1 e eqs = case setConst VZ e eqs of
  ZeroIsOne    => Nothing
  C (eqs :< _) => Just (eqs, K e ::: id)

export %inline
split : DimEq' (S d) -> List (Split d)
split eqs = toList (split1 Zero eqs) <+> toList (split1 One eqs)


export
splits' : DimEq' d -> List (DSubst d 0)
splits' [<]          = [id]
splits' eqs@(_ :< _) = [th . ph | (eqs', th) <- split eqs, ph <- splits' eqs']

export %inline
splits : DimEq d -> List (DSubst d 0)
splits ZeroIsOne = []
splits (C eqs)   = splits' eqs


private
0 newGetShift : (d : Nat) -> (i : Var d) -> (by : Shift d d') ->
                getShift' by (new' {d}) i = Nothing
newGetShift (S d) VZ     by = Refl
newGetShift (S d) (VS i) by = newGetShift d i (ssDown by)

export
0 newGet' : (d : Nat) -> (i : Var d) -> get' (new' {d}) i = Nothing
newGet' d i = newGetShift d i SZ

export
0 newGet : (d : Nat) -> (p : Dim d) -> get (new' {d}) p = p
newGet d (K e) = Refl
newGet d (B i) = rewrite newGet' d i in Refl


export
0 setSelf : (p : Dim d) -> (eqs : DimEq d) -> set p p eqs = eqs
setSelf p ZeroIsOne = Refl
setSelf (K Zero) (C eqs) = Refl
setSelf (K One)  (C eqs) = Refl
setSelf (B i)    (C eqs) with (compareP i i) | (compare i.nat i.nat)
  _ | IsLT lt | LT = absurd lt
  _ | IsEQ    | EQ = Refl
  _ | IsGT gt | GT = absurd gt


public export
wf' : DimEq' d -> Bool
wf' [<]                 = True
wf' (eqs :< Nothing)    = wf' eqs
wf' (eqs :< Just (K e)) = wf' eqs
wf' (eqs :< Just (B i)) = isNothing (get' eqs i) && wf' eqs

public export
wf : DimEq d -> Bool
wf ZeroIsOne = True
wf (C eqs)   = wf' eqs


-- [todo] "well formed" dimeqs
-- [todo] operations maintain well-formedness
-- [todo] if 'Wf eqs' then 'equal eqs' is an equivalence
-- [todo] 'set' never breaks existing equalities