module Obs.NormalForm

import Data.Fin

import Obs.Sort
import Obs.Substitution

import Text.PrettyPrint.Prettyprinter

%default total

-- Definition ------------------------------------------------------------------

data Constructor : Nat -> Type
data Neutral     : Nat -> Type
data NormalForm  : Nat -> Type

public export
data Constructor : Nat -> Type where
  Sort : Sort -> Constructor n
  Top  : Constructor n

public export
data Neutral : Nat -> Type where
  Var : Fin n -> Neutral n

public export
data NormalForm : Nat -> Type where
  Ntrl  : Neutral n -> NormalForm n
  Cnstr : Constructor n -> NormalForm n
  Irrel : NormalForm n

public export
record Definition (n : Nat) where
  constructor MkDefinition
  name   : Maybe String
  sort   : Sort
  ty     : NormalForm n
  tm     : NormalForm n

public export
data Context : Nat -> Type where
  Nil  : Context 0
  (:<) : Context n -> Definition n -> Context (S n)

-- Interfaces ------------------------------------------------------------------

-- Naive equality tests
eqCnstr : Constructor n -> Constructor n -> Bool
eqNtrl  : Neutral n -> Neutral n -> Bool
eqWhnf  : NormalForm n -> NormalForm n -> Bool

eqCnstr (Sort s) (Sort s') = s == s'
eqCnstr Top      Top       = True
eqCnstr _ _ = False

eqNtrl (Var i) (Var j) = i == j

eqWhnf (Ntrl t)  (Ntrl u)  = eqNtrl t u
eqWhnf (Cnstr t) (Cnstr u) = eqCnstr t u
eqWhnf Irrel     Irrel     = True
eqWhnf _ _ = False

export
Eq (Constructor n) where
  t == u = eqCnstr t u

export
Eq (Neutral n) where
  t == u = eqNtrl t u

export
Eq (NormalForm n) where
  t == u = eqWhnf t u

-- Pretty Print ----------------------------------------------------------------

prettyPrecCnstr : Prec -> Constructor n -> Doc ann
prettyPrecNtrl  : Prec -> Neutral n -> Doc ann
prettyPrecWhnf  : Prec -> NormalForm n -> Doc ann

prettyPrecCnstr d (Sort s) = prettyPrec d s
prettyPrecCnstr d Top = pretty "()"

prettyPrecNtrl d (Var i) = pretty "$\{show i}"

prettyPrecWhnf d (Ntrl t)  = prettyPrecNtrl d t
prettyPrecWhnf d (Cnstr t) = prettyPrecCnstr d t
prettyPrecWhnf d Irrel     = pretty "_"

export
Pretty (Constructor n) where
  prettyPrec = prettyPrecCnstr

export
Pretty (Neutral n) where
  prettyPrec = prettyPrecNtrl

export
Pretty (NormalForm n) where
  prettyPrec = prettyPrecWhnf

export
Pretty (Definition n) where
  pretty def = group $
    pretty def.name <++> colon <+> softline <+> pretty def.ty <+> softline <+> colon <++> pretty def.sort <+> hardline <+>
    pretty def.name <++> equals <+> softline <+> pretty def.tm

export
Pretty (Context n) where
  pretty []           = neutral
  pretty ([] :< def)  = pretty def
  pretty (ctx :< def) = pretty ctx <+> hardline <+> hardline <+> pretty def

-- Operations ------------------------------------------------------------------

-- Renaming

renameCnstr : Constructor n -> (Fin n -> Fin m) -> Constructor m
renameNtrl  : Neutral n -> (Fin n -> Fin m) -> Neutral m
renameWhnf  : NormalForm n -> (Fin n -> Fin m) -> NormalForm m

renameCnstr (Sort s) f = Sort s
renameCnstr Top      f = Top

renameNtrl (Var i) f = Var (f i)

renameWhnf (Ntrl t)  f = Ntrl $ renameNtrl t f
renameWhnf (Cnstr t) f = Cnstr $ renameCnstr t f
renameWhnf Irrel     f = Irrel

export
Rename Constructor where
  rename = renameCnstr

export
Rename Neutral where
  rename = renameNtrl

export
Rename NormalForm where
  rename = renameWhnf

-- Substitution

substCnstr : Constructor n -> (Fin n -> NormalForm m) -> Constructor m
substNtrl  : Neutral n -> (Fin n -> NormalForm m) -> NormalForm m
substWhnf  : NormalForm n -> (Fin n -> NormalForm m) -> NormalForm m

substCnstr (Sort s) f = Sort s
substCnstr Top      f = Top

substNtrl (Var i) f = f i

substWhnf (Ntrl t)  f = substNtrl t f
substWhnf (Cnstr t) f = Cnstr $ substCnstr t f
substWhnf Irrel     f = Irrel

export
Subst Constructor NormalForm Constructor where
  subst = substCnstr

export
Subst Neutral NormalForm NormalForm where
  subst = substNtrl

export
Subst NormalForm NormalForm NormalForm where
  subst = substWhnf