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 noLoc)

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 -> Loc -> Dim d
getVar eqs i loc = fromMaybe (B i loc) $ 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 loc) = K e loc
get eqs (B i loc) = getVar eqs i loc


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 noLoc 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 -> Loc -> DimEq' d -> DimEq d
setConst VZ e loc (eqs :< Nothing) =
  C $ eqs :< Just (K e loc)
setConst VZ e _ (eqs :< Just (K f loc)) =
  checkConst e f $ eqs :< Just (K f loc)
setConst VZ e loc (eqs :< Just (B i _)) =
  setConst i e loc eqs :<? Just (K e loc)
setConst (VS i) e loc (eqs :< p) =
  setConst i e loc eqs :<? ifVar i (K e loc) p

mutual
  private
  setVar' : (i, j : Var d) -> (0 _ : i `LT` j) -> Loc -> DimEq' d -> DimEq d
  setVar' VZ (VS i) LTZ loc (eqs :< Nothing) =
    C eqs :<? Just (B i loc)
  setVar' VZ (VS i) LTZ loc (eqs :< Just (K e eloc)) =
    setConst i e loc eqs :<? Just (K e eloc)
  setVar' VZ (VS i) LTZ loc (eqs :< Just (B j jloc)) =
    setVar i j loc jloc eqs :<? Just (if j > i then B j jloc else B i loc)
  setVar' (VS i) (VS j) (LTS lt) loc (eqs :< p) =
    setVar' i j lt loc eqs :<? ifVar i (B j loc) p

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


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


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

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

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

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

||| the Loc is put into each of the DimConsts
export %inline
splits : Loc -> DimEq d -> List (DSubst d 0)
splits _   ZeroIsOne = []
splits loc (C eqs)   = splits' loc 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 : {opts : _} -> BContext d -> Eff Pretty (SnocList (Doc opts))
prettyDVars = traverse prettyDBind . toSnocList'

private
prettyCst : {opts : _} -> BContext d -> Dim d -> Dim d -> Eff Pretty (Doc opts)
prettyCst dnames p q =
  hsep <$> sequence [prettyDim dnames p, cstD, prettyDim dnames q]

private
prettyCsts : {opts : _} -> BContext d -> DimEq' d ->
             Eff Pretty (SnocList (Doc opts))
prettyCsts [<] [<] = pure [<]
prettyCsts dnames (eqs :< Nothing) = prettyCsts (tail dnames) eqs
prettyCsts dnames (eqs :< Just q)  =
  [|prettyCsts (tail dnames) eqs :< prettyCst dnames (BV 0 noLoc) (weakD 1 q)|]

export
prettyDimEq' : {opts : _} -> BContext d -> DimEq' d -> Eff Pretty (Doc opts)
prettyDimEq' dnames eqs = do
  vars <- prettyDVars dnames
  eqs  <- prettyCsts dnames eqs
  let prec = if length vars <= 1 && null eqs then Arg else Outer
  parensIfM prec $ fillSeparateTight !commaD $ toList vars ++ toList eqs

export
prettyDimEq : {opts : _} -> BContext d -> DimEq d -> Eff Pretty (Doc opts)
prettyDimEq dnames ZeroIsOne = do
  vars <- prettyDVars dnames
  cst  <- prettyCst [<] (K Zero noLoc) (K One noLoc)
  pure $ separateTight !commaD $ vars :< cst
prettyDimEq dnames (C eqs) = prettyDimEq' dnames eqs


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