module Quox.Syntax.Dim

import Quox.Loc
import Quox.Name
import Quox.Var
import Quox.Syntax.Subst
import Quox.Pretty
import Quox.Context
import Quox.PrettyValExtra

import Decidable.Equality
import Control.Function
import Derive.Prelude

%default total
%language ElabReflection


public export
data DimConst = Zero | One
%name DimConst e
%runElab derive "DimConst" [Eq, Ord, Show, PrettyVal]

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

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

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
data Dim : Nat -> Type where
  K : DimConst -> Loc -> Dim d
  B : Var d -> Loc -> Dim d
%name Dim.Dim p, q
%runElab deriveIndexed "Dim" [Eq, Ord, Show]


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


export
Located (Dim d) where
  (K _ loc).loc = loc
  (B _ loc).loc = loc

export
Relocatable (Dim d) where
  setLoc loc (K e _) = K e loc
  setLoc loc (B i _) = B i loc

export
prettyDimConst : {opts : _} -> DimConst -> Eff Pretty (Doc opts)
prettyDimConst = hl Dim . text . ends "0" "1"

export
prettyDim : {opts : _} -> BContext d -> Dim d -> Eff Pretty (Doc opts)
prettyDim names (K e _) = prettyDimConst e
prettyDim names (B i _) = prettyDBind $ names !!! i


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


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


public export
FromVarR Dim where fromVarR = B

public export
FromVar Dim where fromVar = B; fromVarSame _ _ = Refl


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

private
substDim : Dim from -> Lazy (DSubst from to) -> Dim to
substDim (K e loc) _  = K e loc
substDim (B i loc) th = get th i loc

export
CanSubstSelfR Dim where (//?) = substDim

export
CanSubstSelf Dim where (//) = substDim; substSame _ _ = Refl


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

public export
data Eqv : Dim d1 -> Dim d2 -> Type where
  EK : K e _ `Eqv` K e _
  EB : i `Eqv` j -> B i _ `Eqv` B j _

export Uninhabited (K e l1 `Eqv` B i l2) where uninhabited _ impossible
export Uninhabited (B i l1 `Eqv` K e l2) where uninhabited _ impossible

export
injectiveB : B i loc1 `Eqv` B j loc2 -> i `Eqv` j
injectiveB (EB e) = e

export
injectiveK : K e loc1 `Eqv` K f loc2 -> e = f
injectiveK EK = Refl

public export
decEqv : Dec2 Dim.Eqv
decEqv (K e _) (K f _) = case decEq e f of
  Yes Refl => Yes EK
  No  n    => No $ n . injectiveK
decEqv (B i _) (B j _) = case decEqv i j of
  Yes y => Yes $ EB y
  No  n => No  $ \(EB y) => n y
decEqv (B _ _) (K _ _) = No absurd
decEqv (K _ _) (B _ _) = No absurd

||| 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) => (loc : Loc) -> Dim d
BV i loc = B (V i) loc


export
weakD : (by : Nat) -> Dim d -> Dim (by + d)
weakD by p = p // shift by