module Quox.Syntax.Dim

import Quox.Syntax.Var
import Quox.Syntax.Subst
import Quox.Pretty

import Decidable.Equality
import Control.Function
import Generics.Derive

%default total
%language ElabReflection
%hide SOP.from; %hide SOP.to


public export
data DimConst = Zero | One
%name DimConst e

%runElab derive "DimConst" [Generic, Meta, Eq, Ord, DecEq, Show]

||| `ends l r e` returns `l` if `e` is `Zero`, or `r` if it is `One`.
public export
ends : a -> a -> DimConst -> a
ends l r Zero = l
ends l r One  = r


public export
data Dim : Nat -> Type where
  K : DimConst -> Dim d
  B : Var d -> Dim d
%name Dim.Dim p, q

private %inline
drepr : Dim n -> Either DimConst (Var n)
drepr (K k) = Left k
drepr (B x) = Right x

export Eq  (Dim n) where (==) = (==) `on` drepr
export Ord (Dim n) where compare = compare `on` drepr

export
Show (Dim n) where
  show (K k) = showCon App "K" $ show k
  show (B i) = showCon App "B" $ show i

export
PrettyHL DimConst where
  prettyM Zero = hl Dim <$> ifUnicode "𝟬" "0"
  prettyM One  = hl Dim <$> ifUnicode "𝟭" "1"

export
PrettyHL (Dim n) where
  prettyM (K e) = prettyM e
  prettyM (B i) = prettyVar DVar DVarErr (!ask).dnames i


||| `endsOr l r x e` returns:
||| - `l` if `p` is `K Zero`;
||| - `r` if `p` is `K One`;
||| - `x` otherwise.
public export
endsOr : a -> a -> Lazy a -> Dim n -> a
endsOr l r x (K e) = ends l r e
endsOr l r x (B _) = x


public export %inline
toConst : Dim 0 -> DimConst
toConst (K e) = e


public export
DSubst : Nat -> Nat -> Type
DSubst = Subst Dim


export %inline
prettyDSubst : Pretty.HasEnv m => DSubst from to -> m (Doc HL)
prettyDSubst th =
  prettySubstM prettyM (!ask).dnames DVar
    !(ifUnicode "⟨" "<") !(ifUnicode "⟩" ">") th


public export FromVar Dim where fromVar = B

export
CanShift Dim where
  K e // _  = K e
  B i // by = B (i // by)

export
CanSubst Dim Dim where
  K e // _  = K e
  B i // th = th !! i


export Uninhabited (Zero = One)  where uninhabited _ impossible
export Uninhabited (One  = Zero) where uninhabited _ impossible

export Uninhabited (B i = K e) where uninhabited _ impossible
export Uninhabited (K e = B i) where uninhabited _ impossible

public export %inline Injective Dim.B where injective Refl = Refl
public export %inline Injective Dim.K where injective Refl = Refl

public export
DecEq DimConst where
  decEq Zero Zero = Yes Refl
  decEq Zero One  = No absurd
  decEq One  Zero = No absurd
  decEq One  One  = Yes Refl

public export
DecEq (Dim d) where
  decEq (K e) (K f) with (decEq e f)
    _ | Yes prf    = Yes $ cong K prf
    _ | No  contra = No  $ contra . injective
  decEq (K e) (B j) = No absurd
  decEq (B i) (K f) = No absurd
  decEq (B i) (B j) with (decEq i j)
    _ | Yes prf    = Yes $ cong B prf
    _ | No  contra = No  $ contra . injective

||| abbreviation for a bound variable like `BV 4` instead of
||| `B (VS (VS (VS (VS VZ))))`
public export %inline
BV : (i : Nat) -> (0 _ : LT i d) => Dim d
BV i = B $ V i