module Quox.Syntax.DimEq

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

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

public export
toMaybe : IfConsistent eqs a -> Maybe a
toMaybe Nothing  = Nothing
toMaybe (Just x) = Just x


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) -> (0 _ : 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


private
prettyDVars : Pretty.HasEnv m => m (SnocList (Doc HL))
prettyDVars = map (pretty0 False . DV) <$> asks dnames

private
prettyCst : (PrettyHL a, PrettyHL b, Pretty.HasEnv m) => a -> b -> m (Doc HL)
prettyCst p q = do
  p <- pretty0M p
  q <- pretty0M q
  pure $ hsep [p, hl Syntax "=", q]

export
PrettyHL (DimEq' d) where
  prettyM eqs {m} = do
    vars <- prettyDVars
    eqs  <- go eqs
    let prec = if length vars <= 1 && null eqs then Arg else Outer
    parensIfM prec $ align $ fillSep $ punctuate comma $ toList $ vars ++ eqs
  where
    tail : SnocList a -> SnocList a
    tail [<] = [<]
    tail (xs :< _) = xs

    go : DimEq' d' -> m (SnocList (Doc HL))
    go [<]              = pure [<]
    go (eqs :< Nothing) = local {dnames $= tail} $ go eqs
    go (eqs :< Just p)  = do
      eq  <- prettyCst (BV {d = 1} 0) (weakD 1 p)
      eqs <- local {dnames $= tail} $ go eqs
      pure $ eqs :< eq

export
PrettyHL (DimEq d) where
  prettyM ZeroIsOne = parensIfM Outer $
    align $ fillSep $ punctuate comma $ toList $
    !prettyDVars :< !(prettyCst Zero One)
  prettyM (C eqs)   = prettyM eqs

export
prettyDimEq : NContext d -> DimEq d -> Doc HL
prettyDimEq ds = pretty0With False (toSnocList' ds) [<]


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