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.DPair

%default total

mutual
  ||| consistent (0≠1) set of constraints between dimension variables
  public export
  data DimEq' : Nat -> Type where
    ||| empty context
    Nil : DimEq' 0
    ||| Ψ, 𝑖, 𝑖=ε
    Const : (eqs : DimEq' d) -> (e : DimConst) -> DimEq' (S d)
    ||| Ψ, 𝑖, 𝑖=𝑗  (Ψ ⊢ 𝑗 and 𝑗 is unassigned)
    Var : (eqs : DimEq' d) -> (i : Var d) -> (0 un : Unassigned eqs i) ->
          DimEq' (S d)
    ||| Ψ, 𝑖  (𝑖 unassigned)
    None : (eqs : DimEq' d) -> DimEq' (S d)
  %name DimEq' eqs

  public export
  data Unassigned : DimEq' d -> Var d -> Type where
    UZ  : Unassigned (None eqs) VZ
    USK : Unassigned eqs i -> Unassigned (Const eqs e)    (VS i)
    USV : Unassigned eqs i -> Unassigned (Var   eqs j un) (VS i)
    USN : Unassigned eqs i -> Unassigned (None  eqs )     (VS i)
  %name Unassigned un


||| set of constraints that might be inconsistent
public export
data DimEq : Nat -> Type where
  ||| 0=1
  ZeroIsOne : DimEq d
  ||| 0≠1, plus other constraints
  C : (eqs : DimEq' d) -> DimEq d
%name DimEq eqs


||| contains a value iff the dim ctx is consistent
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

||| performs an action if the dim ctx is consistent
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 %inline
weakD : Dim d -> Dim (S d)
weakD p = p // SS SZ


public export
tail' : DimEq' (S d) -> DimEq' d
tail' (Const eqs e)    = eqs
tail' (Var   eqs i un) = eqs
tail' (None  eqs )     = eqs

public export
tail : DimEq (S d) -> DimEq d
tail ZeroIsOne = ZeroIsOne
tail (C eqs)   = C $ tail' eqs

public export
head' : DimEq' (S d) -> Maybe (Dim d)
head' (Const _ e)   = Just $ K e
head' (Var   _ i _) = Just $ B i
head' (None  _)     = Nothing

export
tailU : Unassigned eqs (VS i) -> Unassigned (tail' eqs) i
tailU (USK un) = un
tailU (USV un) = un
tailU (USN un) = un


||| make a dim ctx where each variable has a constant assignment
public export
fromGround' : Context' DimConst d -> DimEq' d
fromGround' [<]        = Nil
fromGround' (ctx :< e) = Const (fromGround' ctx) e

||| make a dim ctx where each variable has a constant assignment
public export
fromGround : Context' DimConst d -> DimEq d
fromGround = C . fromGround'


||| make a dim ctx where each variable is unassigned
public export
new' : (d : Nat) -> DimEq' d
new' 0     = Nil
new' (S d) = None (new' d)

||| make a dim ctx where each variable is unassigned
public export
new : (d : Nat) -> DimEq d
new d = C $ new' d


||| if the dim is a variable, then it is unassigned
public export
data UnassignedDim : DimEq' d -> Dim d -> Type where
  UDK : UnassignedDim eqs (K e)
  UDB : Unassigned eqs i -> UnassignedDim eqs (B i)

export
weakUD : {eqs : DimEq' (S d)} -> {p : Dim d} ->
         UnassignedDim (tail' eqs) p -> UnassignedDim eqs (weakD p)
weakUD UDK = UDK
weakUD (UDB un) {eqs = Const eqs e}   = UDB $ USK un
weakUD (UDB un) {eqs = Var   eqs _ _} = UDB $ USV un
weakUD (UDB un) {eqs = None  eqs}     = UDB $ USN un


||| get the constraint on a variable 𝑖. if it is equal to another var 𝑗,
||| then 𝑗 is not further constrained
public export
getVarPrf : (eqs : DimEq' d) -> Var d -> Subset (Dim d) (UnassignedDim eqs)
getVarPrf (Const eqs e)    VZ     = Element (K e)      UDK
getVarPrf (Var   eqs i un) VZ     = Element (B $ VS i) (UDB $ USV un)
getVarPrf (None  eqs)      VZ     = Element (B VZ)     (UDB UZ)
getVarPrf (Const eqs _)    (VS i) = let p = getVarPrf eqs i in
                                    Element (weakD p.fst) (weakUD p.snd)
getVarPrf (Var   eqs _ _)  (VS i) = let p = getVarPrf eqs i in
                                    Element (weakD p.fst) (weakUD p.snd)
getVarPrf (None  eqs)      (VS i) = let p = getVarPrf eqs i in
                                    Element (weakD p.fst) (weakUD p.snd)

public export
getVar : (eqs : DimEq' d) -> Var d -> Dim d
getVar eqs i = fst $ getVarPrf eqs i

public export
getPrf : (eqs : DimEq' d) -> Dim d -> Subset (Dim d) (UnassignedDim eqs)
getPrf eqs (K e) = Element (K e) UDK
getPrf eqs (B i) = getVarPrf eqs i

public export
get : DimEq' d -> Dim d -> Dim d
get eqs p = fst $ getPrf eqs p


-- version of `get` that only shifts once but is even more annoying to prove
-- anything about
private
getShift' : Shift d out -> DimEq' d -> Var d -> Maybe (Dim out)
getShift' by (Const eqs e)    VZ     = Just $ K e
getShift' by (Var   eqs i un) VZ     = Just $ B $ i // ssDown by
getShift' by (None  eqs)      VZ     = Nothing
getShift' by eqs              (VS i) = getShift' (ssDown by) (tail' eqs) i

private
getShift0' : DimEq' d -> Var d -> Maybe (Dim d)
getShift0' = getShift' SZ

private
get' : DimEq' d -> Dim d -> Dim d
get' eqs (K e) = K e
get' eqs (B i) = fromMaybe (B i) $ getShift0' eqs i

%transform "DimEq.get" get = get'


public export
Equal' : DimEq' d -> Rel (Dim d)
Equal' eqs p q = get eqs p = get eqs q

||| whether two dimensions are equal under the current constraints
public export
data Equal : DimEq d -> Rel (Dim d) where
  Eq01 : Equal ZeroIsOne p q
  EqC  : Equal' eqs p q -> Equal (C eqs) p q
%name DimEq.Equal prf

export
decEqual : (eqs : DimEq d) -> Dec2 (Equal eqs)
decEqual ZeroIsOne _ _ = Yes Eq01
decEqual (C eqs)   p q = case get eqs p `decEq` get eqs q of
  Yes y => Yes $ EqC y
  No  n => No  $ \case EqC p => n p

export
equal : (eqs : DimEq d) -> Dim d -> Dim d -> Bool
equal eqs p q = isYes $ decEqual eqs p q

export
{eqs : DimEq d} -> Reflexive _ (Equal eqs) where
  reflexive = case eqs of
    ZeroIsOne => Eq01
    C eqs     => EqC Refl

export
Symmetric _ (Equal eqs) where
  symmetric Eq01     = Eq01
  symmetric (EqC eq) = EqC $ sym eq

export
Transitive _ (Equal eqs) where
  transitive Eq01    Eq01    = Eq01
  transitive (EqC p) (EqC q) = EqC $ p `trans` q

export {eqs : DimEq d} -> Equivalence _ (Equal eqs) where


||| extend the context with a new variable, possibly constrained
public export
(:<) : DimEq' d -> Maybe (Dim d) -> DimEq' (S d)
eqs :< Nothing    = None  eqs
eqs :< Just (K e) = Const eqs e
eqs :< Just (B i) with (getVarPrf eqs i)
  _ | Element (K e) _  = Const eqs e
  _ | Element (B j) un = Var eqs j $ let UDB un = un in un

infixl 7 :<?
||| extend the context with a new variable, possibly constrained
public export
(:<?) : DimEq d -> Maybe (Dim d) -> DimEq (S d)
ZeroIsOne :<? p = ZeroIsOne
C eqs     :<? p = C $ eqs :< p


public export
checkConst : DimConst -> DimConst -> DimEq' d -> DimEq d
checkConst e f eqs = case decEq e f of Yes _ => C eqs; No _ => ZeroIsOne

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

public export
setVar : Var d -> Var d -> DimEq' d -> DimEq d
setVar VZ VZ eqs = C eqs
setVar VZ     (VS j) (Const eqs e)    = setConst j e eqs :<? Just (K e)
setVar VZ     (VS j) (Var   eqs k un) = setVar j k eqs :<? Just (B k)
setVar VZ     (VS j) (None  eqs)      = C eqs :<? Just (B j)
setVar (VS i) VZ     (Const eqs e)    = setConst i e eqs :<? Just (K e)
setVar (VS i) VZ     (Var   eqs k un) = setVar i k eqs :<? Just (B k)
setVar (VS i) VZ     (None  eqs)      = C eqs :<? Just (B i)
setVar (VS i) (VS j) (Const eqs e)    = setVar i j eqs :<? Just (K e)
setVar (VS i) (VS j) (Var   eqs k un) = setVar i j eqs :<? Just (B k)
setVar (VS i) (VS j) (None  eqs)      = setVar i j eqs :<? Nothing

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


private
splitV0 : (p : Dim (S d)) -> Either (p = B VZ) (q : Dim d ** p = weakD q)
splitV0 (K e)      = Right (K e ** Refl)
splitV0 (B VZ)     = Left Refl
splitV0 (B (VS i)) = Right (B i ** Refl)


export
0 getSnoc : (eqs : DimEq' d) -> (u : Maybe (Dim d)) -> (i : Var d) ->
            getVar (eqs :< u) (VS i) = weakD (getVar eqs i)
getSnoc eqs Nothing      i = Refl
getSnoc eqs (Just (K e)) i = Refl
getSnoc eqs (Just (B j)) i with (getVarPrf eqs j)
  _ | Element (K _) _ = Refl
  _ | Element (B _) _ = Refl

export
0 snocStrengthen : (p, q : Dim d) ->
                   Equal' (eqs :< u) (weakD p) (weakD q) -> Equal' eqs p q
snocStrengthen (K e) (K e) Refl = Refl
snocStrengthen (K e) (B i) prf =
  shiftInj (SS SZ) _ _ $
  rewrite sym $ getSnoc eqs u i in prf
snocStrengthen (B i) (K e) prf =
  shiftInj (SS SZ) _ _ $
  rewrite sym $ getSnoc eqs u i in prf
snocStrengthen (B i) (B j) prf =
  shiftInj (SS SZ) _ _ $
  rewrite sym $ getSnoc eqs u i in
  rewrite sym $ getSnoc eqs u j in prf

export
0 snocStable : (eqs : DimEq d) -> (u : Maybe (Dim d)) -> (p, q : Dim d) ->
               Equal eqs p q -> Equal (eqs :<? u) (weakD p) (weakD q)
snocStable ZeroIsOne u p q Eq01 = Eq01
snocStable (C eqs) u (K e) (K e) (EqC Refl) = reflexive
snocStable (C eqs) u (K e) (B i) (EqC prf) = EqC $
  rewrite getSnoc eqs u i in rewrite sym prf in Refl
snocStable (C eqs) u (B i) (K e) (EqC prf) = EqC $
  rewrite getSnoc eqs u i in rewrite prf in Refl
snocStable (C eqs) u (B i) (B j) (EqC prf) = EqC $
  rewrite getSnoc eqs u i in
  rewrite getSnoc eqs u j in
  rewrite prf in Refl

export
0 checkConstStable : (eqs : DimEq' d) -> (e, f : DimConst) ->
                     (p, q : Dim d) -> Equal' eqs p q ->
                     Equal (checkConst e f eqs) p q
checkConstStable eqs e f p q prf with (decEq e f)
  _ | Yes _ = EqC prf
  _ | No  _ = Eq01

export
0 setConstStable : (eqs : DimEq' d) -> (i : Var d) -> (e : DimConst) ->
                   (p, q : Dim d) -> Equal' eqs p q ->
                   Equal (setConst i e eqs) p q
setConstStable (Const eqs f) VZ e p q prf with (decEq e f)
  _ | Yes _ = EqC prf
  _ | No  _ = Eq01
setConstStable (Const eqs f) (VS i) e p q prf = ?setConstStable_rhs_5
setConstStable (Var eqs j un) VZ e p q prf = ?setConstStable_rhs_6
setConstStable (Var eqs j un) (VS i) e p q prf = ?setConstStable_rhs_7
setConstStable (None eqs) VZ e p q prf = ?setConstStable_rhs_8
setConstStable (None eqs) (VS i) e p q prf = ?setConstStable_rhs_9

export
0 setVarStable : (eqs : DimEq' d) -> (i, j : Var d) ->
                 (p, q : Dim d) -> Equal' eqs p q ->
                 Equal (setVar i j eqs) p q

export
0 setStable : (eqs : DimEq d) -> (u, v, p, q : Dim d) ->
              Equal eqs p q -> Equal (set u v eqs) p q
setStable ZeroIsOne p q u v Eq01 = Eq01
setStable (C eqs) (K e) (K f) p q (EqC prf) = checkConstStable eqs e f p q prf
setStable (C eqs) (K e) (B i) p q (EqC prf) = setConstStable eqs i e p q prf
setStable (C eqs) (B i) (K e) p q (EqC prf) = setConstStable eqs i e p q prf
setStable (C eqs) (B i) (B j) p q (EqC prf) = setVarStable eqs i j p q prf