module Obs.NormalForm

import Data.Fin

import Obs.Sort
import Obs.Substitution

import Text.Bounded
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 -> String -> NormalForm n -> NormalForm (S n) -> Constructor n
  Lambda : String -> NormalForm (S n) -> Constructor n
  Sigma  : Sort -> Sort -> String -> NormalForm n -> NormalForm (S n) -> Constructor n
  Pair   : NormalForm n -> NormalForm n -> Constructor n
  Top    : Constructor n
  Bottom : Constructor n

public export
data Neutral : Nat -> Type where
  Var    : String -> Fin n -> Neutral n
  App    : Neutral n -> NormalForm n -> Neutral n
  Fst    : Neutral n -> Neutral n
  Snd    : Neutral n -> Neutral n
  Absurd : Neutral n
  Equal  : Neutral n -> NormalForm n -> NormalForm n -> Neutral n
  EqualL : Nat -> Neutral n -> NormalForm n -> Neutral n
  EqualR : Nat -> Constructor n -> Neutral n -> Neutral n
  EqualU : Nat -> Constructor n -> Constructor n -> Neutral n
  CastL  : Neutral n -> NormalForm n -> NormalForm n -> Neutral n
  CastR  : Constructor n -> Neutral n -> NormalForm n -> Neutral n
  CastU  : Constructor n -> Constructor n -> NormalForm 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   : WithBounds 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 (Sigma s s' _ a b) (Sigma l l' _ a' b') = s == l && s' == l' && eqWhnf a a' && eqWhnf b b'
eqCnstr (Pair t u) (Pair t' u') = eqWhnf t t' && eqWhnf u 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 (Fst t) (Fst t') = eqNtrl t t'
eqNtrl (Snd t) (Snd t') = eqNtrl t t'
eqNtrl Absurd Absurd = True
eqNtrl (Equal a t u) (Equal a' t' u') = eqNtrl a a' && eqWhnf t t' && eqWhnf u u'
eqNtrl (EqualL i t u) (EqualL j t' u') = i == j && eqNtrl t t' && eqWhnf u u'
eqNtrl (EqualR i t u) (EqualR j t' u') = i == j && eqCnstr t t' && eqNtrl u u'
eqNtrl (EqualU i t u) (EqualU j t' u') = i == j && eqCnstr t t' && eqCnstr u u'
eqNtrl (CastL a b t) (CastL a' b' t') = eqNtrl a a' && eqWhnf b b' && eqWhnf t t'
eqNtrl (CastR a b t) (CastR a' b' t') = eqCnstr a a' && eqNtrl b b' && eqWhnf t t'
eqNtrl (CastU a b t) (CastU a' b' t') = eqCnstr a a' && eqCnstr b b' && eqWhnf t t'
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 _ _ var a b) =
  parenthesise (d > Open) $
  group $
  parens (pretty var <++> colon <+> softline <+> prettyPrecWhnf Open a) <++>
  pretty "->" <+> softline <+>
  prettyPrecWhnf Open b
prettyPrecCnstr d (Lambda var t) =
  parenthesise (d > Open) $
  group $
  backslash <+> pretty var <++>
  pretty "=>" <+> softline <+>
  prettyPrecWhnf Open t
prettyPrecCnstr d (Sigma _ _ var a b) =
  parenthesise (d > Open) $
  group $
  parens (pretty var <++> colon <+> softline <+> prettyPrecWhnf Open a) <++>
  pretty "**" <+> softline <+>
  prettyPrecWhnf Open b
prettyPrecCnstr d (Pair t u) =
  angles $
  group $
  neutral <++> prettyPrecWhnf Open t <+> comma <+> softline <+> prettyPrecWhnf Open u <++> neutral
prettyPrecCnstr d Top = pretty "()"
prettyPrecCnstr d Bottom = pretty "Void"

prettyPrecNtrl d (Var var i) = pretty "\{show var}@\{show i}"
prettyPrecNtrl d (App t u) =
  parenthesise (d >= App) $
  group $
  fillSep [prettyPrecNtrl Open t, prettyPrecWhnf App u]
prettyPrecNtrl d (Fst t) =
  parenthesise (d >= App) $
  group $
  fillSep [pretty "fst", prettyPrecNtrl App t]
prettyPrecNtrl d (Snd t) =
  parenthesise (d >= App) $
  group $
  fillSep [pretty "snd", prettyPrecNtrl App t]
prettyPrecNtrl d Absurd = pretty "absurd"
prettyPrecNtrl d (Equal _ t u) =
  parenthesise (d >= Equal) $
  group $
  prettyPrecWhnf Equal t <++> pretty "~" <+> softline <+> prettyPrecWhnf Equal u
prettyPrecNtrl d (EqualL _ t u) =
  parenthesise (d >= Equal) $
  group $
  prettyPrecNtrl Equal t <++> pretty "~" <+> softline <+> prettyPrecWhnf Equal u
prettyPrecNtrl d (EqualR _ t u) =
  parenthesise (d >= Equal) $
  group $
  prettyPrecCnstr Equal t <++> pretty "~" <+> softline <+> prettyPrecNtrl Equal u
prettyPrecNtrl d (EqualU _ t u) =
  parenthesise (d >= Equal) $
  group $
  prettyPrecCnstr Equal t <++> pretty "~" <+> softline <+> prettyPrecCnstr Equal u
prettyPrecNtrl d (CastL a b t) =
  parenthesise (d >= App) $
  group $
  fillSep [pretty "cast", prettyPrecNtrl App a, prettyPrecWhnf App b, prettyPrecWhnf App t]
prettyPrecNtrl d (CastR a b t) =
  parenthesise (d >= App) $
  group $
  fillSep [pretty "cast", prettyPrecCnstr App a, prettyPrecNtrl App b, prettyPrecWhnf App t]
prettyPrecNtrl d (CastU a b t) =
  parenthesise (d >= App) $
  group $
  fillSep [pretty "cast", prettyPrecCnstr App a, prettyPrecCnstr App b, prettyPrecWhnf App t]

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.val <++> colon <+> softline <+> pretty def.ty <+> softline <+> colon <++> pretty def.sort <+> hardline <+>
    pretty def.name.val <++> 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' var a b) f = Pi s s' var (renameWhnf a f) (renameWhnf b $ lift 1 f)
renameCnstr (Lambda var t) f = Lambda var (renameWhnf t $ lift 1 f)
renameCnstr (Sigma s s' var a b) f = Sigma s s' var (renameWhnf a f) (renameWhnf b $ lift 1 f)
renameCnstr (Pair t u) f = Pair (renameWhnf t f) (renameWhnf u f)
renameCnstr Top f = Top
renameCnstr Bottom f = Bottom

renameNtrl (Var var i) f = Var var (f i)
renameNtrl (App t u) f = App (renameNtrl t f) (renameWhnf u f)
renameNtrl (Fst t) f = Fst (renameNtrl t f)
renameNtrl (Snd t) f = Snd (renameNtrl t f)
renameNtrl Absurd f = Absurd
renameNtrl (Equal a t u) f = Equal (renameNtrl a f) (renameWhnf t f) (renameWhnf u f)
renameNtrl (EqualL i t u) f = EqualL i (renameNtrl t f) (renameWhnf u f)
renameNtrl (EqualR i t u) f = EqualR i (renameCnstr t f) (renameNtrl u f)
renameNtrl (EqualU i t u) f = EqualU i (renameCnstr t f) (renameCnstr u f)
renameNtrl (CastL a b t) f = CastL (renameNtrl a f) (renameWhnf b f) (renameWhnf t f)
renameNtrl (CastR a b t) f = CastR (renameCnstr a f) (renameNtrl b f) (renameWhnf t f)
renameNtrl (CastU a b t) f = CastU (renameCnstr a f) (renameCnstr b f) (renameWhnf t f)

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