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
  Pi     : Sort -> Sort -> NormalForm n -> NormalForm (S n) -> Constructor n
  Lambda : NormalForm (S n) -> Constructor n
  Top    : Constructor n
  Bottom : Constructor n

public export
data Neutral : Nat -> Type where
  Var    : Fin n -> Neutral n
  App    : Neutral n -> NormalForm n -> Neutral n
  Absurd : 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   : 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 (Pi s s' a b) (Pi l l' a' b') = s == l && s' == l' && eqWhnf a a' && eqWhnf b b'
eqCnstr (Lambda t) (Lambda u) = eqWhnf t u
eqCnstr Top Top = True
eqCnstr Bottom Bottom = True
eqCnstr _ _ = False

eqNtrl (Var i) (Var j) = i == j
eqNtrl (App t u) (App t' u') = eqNtrl t t' && eqWhnf u u'
eqNtrl Absurd Absurd = True
eqNtrl _ _ = False

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

export
Cast Sort (Constructor n) where
  cast = Sort

export
Cast Sort (NormalForm n) where
  cast = Cnstr . cast

-- 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 (Pi _ _ a b) =
  parenthesise (d > Open) $
  group $
  parens (prettyPrecWhnf Open a) <++>
  pretty "->" <+> softline <+>
  prettyPrecWhnf Open b
prettyPrecCnstr d (Lambda t) =
  parenthesise (d > Open) $
  group $
  backslash <+> softline <+> prettyPrecWhnf Open t
prettyPrecCnstr d Top = pretty "()"
prettyPrecCnstr d Bottom = pretty "Void"

prettyPrecNtrl d (Var i) = pretty "$\{show i}"
prettyPrecNtrl d (App t u) =
  parenthesise (d >= App) $
  group $
  fillSep [prettyPrecNtrl Open t, prettyPrecWhnf App u]
prettyPrecNtrl d Absurd = pretty "absurd"

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 (Pi s s' a b) f = Pi s s' (renameWhnf a f) (renameWhnf b $ lift 1 f)
renameCnstr (Lambda t) f = Lambda (renameWhnf t $ lift 1 f)
renameCnstr Top f = Top
renameCnstr Bottom f = Bottom

renameNtrl (Var i) f = Var (f i)
renameNtrl (App t u) f = App (renameNtrl t f) (renameWhnf u f)
renameNtrl Absurd f = Absurd

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

export
PointedRename Neutral where
  point = Var

export
PointedRename NormalForm where
  point = Ntrl . point