module Quox.Syntax.Var

import Quox.Name
import Quox.Pretty
import Quox.OPE

import Data.Nat
import Data.List
import Decidable.Equality
import Data.Bool.Decidable

%default total


public export
data Var : Nat -> Type where
  VZ : Var (S n)
  VS : Var n -> Var (S n)
%name Var i, j
%builtin Natural Var

public export
(.nat) : Var n -> Nat
(VZ).nat   = 0
(VS i).nat = S i.nat
%transform "Var.(.nat)" Var.(.nat) i = believe_me i

public export %inline Cast (Var n) Nat     where cast = (.nat)
public export %inline Cast (Var n) Integer where cast = cast . cast {to = Nat}

export %inline Eq   (Var n) where i == j = i.nat == j.nat
export %inline Ord  (Var n) where compare i j = compare i.nat j.nat
export %inline Show (Var n) where showPrec d i = showCon d "V" $ showArg i.nat

public export %inline Injective VS where injective Refl = Refl


parameters {auto _ : Pretty.HasEnv m}
  private
  prettyIndex : Nat -> m (Doc a)
  prettyIndex i =
    ifUnicode (pretty $ pack $ map sup $ unpack $ show i) (":" <+> pretty i)
   where
    sup : Char -> Char
    sup c = case c of
      '0' => '⁰'; '1' => '¹'; '2' => '²'; '3' => '³'; '4' => '⁴'
      '5' => '⁵'; '6' => '⁶'; '7' => '⁷'; '8' => '⁸'; '9' => '⁹'; _ => c

  ||| `prettyVar hlok hlerr names i` pretty prints the de Bruijn index `i`.
  |||
  ||| If it is within the bounds of `names`, then it uses the name at that index,
  ||| highlighted as `hlok`. Otherwise it is just printed as a number highlighted
  ||| as `hlerr`.
  export
  prettyVar' : HL -> HL -> List Name -> Nat -> m (Doc HL)
  prettyVar' hlok hlerr names i =
    case inBounds i names of
         Yes _ => hlF' hlok [|prettyM (index i names) <+> prettyIndex i|]
         No  _ => pure $ hl hlerr $ pretty i

  export %inline
  prettyVar : HL -> HL -> List Name -> Var n -> m (Doc HL)
  prettyVar hlok hlerr names i = prettyVar' hlok hlerr names i.nat


public export
fromNatWith : (i : Nat) -> (0 p : i `LT` n) -> Var n
fromNatWith Z     (LTESucc _) = VZ
fromNatWith (S i) (LTESucc p) = VS $ fromNatWith i p
%transform "Var.fromNatWith" fromNatWith i p = believe_me i

public export %inline
V : (i : Nat) -> {auto 0 p : i `LT` n} -> Var n
V i {p} = fromNatWith i p

export %inline
tryFromNat : Alternative f => (n : Nat) -> Nat -> f (Var n)
tryFromNat n i =
  case i `isLT` n of
       Yes p => pure $ fromNatWith i p
       No  _ => empty

export
0 toNatLT : (i : Var n) -> i.nat `LT` n
toNatLT VZ     = LTESucc LTEZero
toNatLT (VS i) = LTESucc $ toNatLT i

public export
toNatInj : {i, j : Var n} -> i.nat = j.nat -> i = j
toNatInj {i = VZ}     {j = VZ}     Refl = Refl
toNatInj {i = VZ}     {j = (VS i)} Refl impossible
toNatInj {i = (VS i)} {j = VZ}     Refl impossible
toNatInj {i = (VS i)} {j = (VS j)} prf  = cong VS $ toNatInj $ injective prf

public export %inline Injective (.nat) where injective = toNatInj

export
0 fromToNat : (i : Var n) -> (p : i.nat `LT` n) -> fromNatWith i.nat p = i
fromToNat VZ     (LTESucc p) = Refl
fromToNat (VS i) (LTESucc p) = rewrite fromToNat i p in Refl

export
0 toFromNat : (i : Nat) -> (p : i `LT` n) -> (fromNatWith i p).nat = i
toFromNat 0     (LTESucc x) = Refl
toFromNat (S k) (LTESucc x) = cong S $ toFromNat k x


-- not using %transform like other things because weakSpec requires the proof
-- to be relevant. but since only `LTESucc` is ever possible that seems
-- to be an instance of <https://github.com/idris-lang/Idris2/issues/1259>?
export
weak : (0 p : m `LTE` n) -> Var m -> Var n
weak p i = fromNatWith i.nat $ transitive (toNatLT i) p

public export
0 weakSpec : m `LTE` n -> Var m -> Var n
weakSpec LTEZero     _      impossible
weakSpec (LTESucc p) VZ     = VZ
weakSpec (LTESucc p) (VS i) = VS $ weakSpec p i

export
0 weakSpecCorrect : (p : m `LTE` n) -> (i : Var m) -> (weakSpec p i).nat = i.nat
weakSpecCorrect LTEZero     _      impossible
weakSpecCorrect (LTESucc x) VZ     = Refl
weakSpecCorrect (LTESucc x) (VS i) = cong S $ weakSpecCorrect x i

export
0 weakCorrect : (p : m `LTE` n) -> (i : Var m) -> (weak p i).nat = i.nat
weakCorrect LTEZero     _      impossible
weakCorrect (LTESucc p) VZ     = Refl
weakCorrect (LTESucc p) (VS i) = cong S $ weakCorrect p i

export
0 weakIsSpec : (p : m `LTE` n) -> (i : Var m) -> weak p i = weakSpec p i
weakIsSpec p i = toNatInj $ trans (weakCorrect p i) (sym $ weakSpecCorrect p i)


public export
interface FromVar f where %inline fromVar : Var n -> f n

public export FromVar Var where fromVar = id


public export
data LT : Var n -> Var n -> Type where
  LTZ : VZ `LT` VS i
  LTS : i `LT` j -> VS i `LT` VS j
%builtin Natural Var.LT
%name Var.LT lt

public export %inline
GT : Var n -> Var n -> Type
i `GT` j = j `LT` i

export
Transitive (Var n) LT where
  transitive LTZ     (LTS _) = LTZ
  transitive (LTS p) (LTS q) = LTS $ transitive p q

export Uninhabited (i    `Var.LT` i)  where uninhabited (LTS p) = uninhabited p
export Uninhabited (VS i `LT`     VZ) where uninhabited _ impossible

export
isLT : (i, j : Var n) -> Dec (i `LT` j)
isLT VZ     VZ     = No uninhabited
isLT VZ     (VS j) = Yes LTZ
isLT (VS i) VZ     = No uninhabited
isLT (VS i) (VS j) with (isLT i j)
  _ | Yes prf   = Yes (LTS prf)
  _ | No contra = No  (\case LTS p => contra p)


public export
data Compare : (i, j : Var n) -> Type where
  IsLT : (lt : i `LT` j) -> Compare i j
  IsEQ :                    Compare i i
  IsGT : (gt : i `GT` j) -> Compare i j
%name Compare cmp

export
compareS : Compare i j -> Compare (VS i) (VS j)
compareS (IsLT lt) = IsLT (LTS lt)
compareS IsEQ      = IsEQ
compareS (IsGT gt) = IsGT (LTS gt)

export
compareP : (i, j : Var n) -> Compare i j
compareP VZ     VZ     = IsEQ
compareP VZ     (VS j) = IsLT LTZ
compareP (VS i) VZ     = IsGT LTZ
compareP (VS i) (VS j) = compareS $ compareP i j

export
0 compareSelf : (c : Compare i i) -> c = IsEQ
compareSelf (IsLT lt) = absurd lt
compareSelf IsEQ      = Refl
compareSelf (IsGT gt) = absurd gt

export
0 comparePSelf : (i : Var n) -> compareP i i = IsEQ
comparePSelf i = compareSelf {}


public export
data LTE : Var n -> Var n -> Type where
  LTEZ : VZ `LTE` j
  LTES : i  `LTE` j -> VS i `LTE` VS j

export
Reflexive (Var n) LTE where
  reflexive {x = VZ}   = LTEZ
  reflexive {x = VS i} = LTES reflexive

export
Transitive (Var n) LTE where
  transitive LTEZ     q        = LTEZ
  transitive (LTES p) (LTES q) = LTES $ transitive p q

export
Antisymmetric (Var n) LTE where
  antisymmetric LTEZ     LTEZ     = Refl
  antisymmetric (LTES p) (LTES q) = cong VS $ antisymmetric p q

export
splitLTE : {j : Var n} -> i `LTE` j -> Either (i = j) (i `LT` j)
splitLTE {j = VZ}   LTEZ     = Left Refl
splitLTE {j = VS _} LTEZ     = Right LTZ
splitLTE            (LTES p) with (splitLTE p)
  _ | (Left  eq) = Left $ cong VS eq
  _ | (Right lt) = Right $ LTS lt


export Uninhabited (VZ   = VS i) where uninhabited _ impossible
export Uninhabited (VS i = VZ)   where uninhabited _ impossible


public export
eqReflect : (i, j : Var n) -> (i = j) `Reflects` (i == j)
eqReflect VZ     VZ     = RTrue  Refl
eqReflect VZ     (VS i) = RFalse absurd
eqReflect (VS i) VZ     = RFalse absurd
eqReflect (VS i) (VS j) with (eqReflect i j)
  eqReflect (VS i) (VS j) | r with (i == j)
    eqReflect (VS i) (VS j) | RTrue  yes | True   =  RTrue  $ cong VS yes
    eqReflect (VS i) (VS j) | RFalse no  | False  =  RFalse $ no . injective

public export
reflectToDec : p `Reflects` b -> Dec p
reflectToDec (RTrue  y) = Yes y
reflectToDec (RFalse n) = No  n

public export %inline
varDecEq : (i, j : Var n) -> Dec (i = j)
varDecEq i j = reflectToDec $ eqReflect i j

-- justified by eqReflect [citation needed]
private %inline
decEqFromBool : (i, j : Var n) -> Dec (i = j)
decEqFromBool i j =
  if i == j then Yes $ believe_me $ Refl {x = 0}
            else No  $ id . believe_me

%transform "Var.decEq" varDecEq = decEqFromBool

public export %inline DecEq (Var n) where decEq = varDecEq

export
Tighten Var where
  tighten Id       i      = pure i
  tighten (Drop q) VZ     = empty
  tighten (Drop q) (VS i) = tighten q i
  tighten (Keep q) VZ     = pure VZ
  tighten (Keep q) (VS i) = VS <$> tighten q i