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

%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


export
new' : (d : Nat) -> DimEq' d
new' 0     = [<]
new' (S d) = new' d :< Nothing

export
new : (d : Nat) -> DimEq d
new d = C (new' d)


export
get : DimEq' d -> Dim d -> Dim d
get _   (K e) = K e
get eqs (B i) = fromMaybe (B i) $ (eqs !! i) @{Compose}

export
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 5 :<?
export
(:<?) : DimEq d -> Maybe (Dim d) -> DimEq (S d)
ZeroIsOne :<? d = ZeroIsOne
C eqs     :<? d = C $ eqs :< d

private
checkConst : (e, f : DimConst) -> (eqs : Lazy (DimEq' d)) -> DimEq d
checkConst Zero Zero eqs = C eqs
checkConst One  One  eqs = C eqs
checkConst _    _    _   = 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 :< x) =
  setConst i e eqs :<? (if x == Just (B i) then Just (K e) else x)

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 :<? map (\k => if k == B i then B j else k) p

  export
  setVar : (i, j : Var d) -> DimEq' d -> DimEq d
  setVar i j eqs =
    case compareP i j of
      IsLT lt => setVar' i j lt eqs
      IsEQ    => C eqs
      IsGT gt => setVar' j i gt eqs


export
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


export
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)
  _ | IsLT lt = absurd lt
  _ | IsEQ    = Refl
  _ | IsGT gt = absurd gt



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

export
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
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@(_ :< _) = do
  (eqs, th) <- split eqs
  ph <- splits' eqs
  pure $ th . ph

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