module Obs.NormalForm

import Data.List.Elem
import Data.So

import Obs.Sort
import Obs.Substitution

import Text.Bounded
import Text.PrettyPrint.Prettyprinter

%default total

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

data NormalForm  : Sorted.Family Bool

public export
data Constructor : Unsorted.Family Bool where
  Sort : (s : Sort) -> Constructor ctx
  Pi : (s, s' : Sort)
    -> (var : String)
    -> (a : NormalForm True ctx)
    -> (b : NormalForm True (isSet s :: ctx))
    -> Constructor ctx
  Lambda : (s : Sort)
    -> (var : String)
    -> NormalForm True (isSet s :: ctx)
    -> Constructor ctx
  Sigma : (s, s' : Sort)
    -> (var : String)
    -> (a : NormalForm True ctx)
    -> (b : NormalForm True (isSet s :: ctx))
    -> Constructor ctx
  Pair : (s, s' : Sort)
    -> (prf : So (isSet (max s s')))
    -> NormalForm (isSet s) ctx
    -> NormalForm (isSet s') ctx
    -> Constructor ctx
  Bool : Constructor ctx
  True : Constructor ctx
  False : Constructor ctx
  Top : Constructor ctx
  Bottom : Constructor ctx

public export
data Neutral : Unsorted.Family Bool where
  Var : (var : String)
    -> (sort : Sort)
    -> (prf : So (isSet sort))
    -> (i : Elem True ctx)
    -> Neutral ctx
  App : (b : Bool) -> Neutral ctx -> NormalForm b ctx -> Neutral ctx
  Fst : (b : Bool) -> Neutral ctx -> Neutral ctx
  Snd : (b : Bool) -> Neutral ctx -> Neutral ctx
  If : Neutral ctx -> (f, g : NormalForm True ctx) -> Neutral ctx
  Absurd : Neutral ctx
  Equal : Neutral ctx -> NormalForm True ctx -> NormalForm True ctx -> Neutral ctx
  EqualL : (a : Neutral ctx) -> (b : NormalForm True ctx) -> Neutral ctx
  EqualR : (a : Constructor ctx) -> (b : Neutral ctx) -> Neutral ctx
  EqualStuck : (a, b : Constructor ctx) -> Neutral ctx
  CastL : (a : Neutral ctx) ->
    (b : NormalForm True ctx) ->
    NormalForm True ctx ->
    Neutral ctx
  CastR : (a : Constructor ctx) -> (b : Neutral ctx) -> NormalForm True ctx -> Neutral ctx
  CastStuck : (a, b : Constructor ctx) -> NormalForm True ctx -> Neutral ctx

public export
data NormalForm : Sorted.Family Bool where
  Ntrl : Neutral ~|> NormalForm True
  Cnstr : Constructor ~|> NormalForm True
  Irrel : NormalForm False ctx

%name Constructor t, u, v
%name Neutral t, u, v
%name NormalForm t, u, v

public export
record Definition (ctx : List Bool) where
  constructor MkDefinition
  name : WithBounds String
  sort : Sort
  ty   : NormalForm True ctx
  tm   : NormalForm (isSet sort) ctx

public export
data Context : List Bool -> Type where
  Nil  : Context []
  (:<) : Context ctx -> (def : Definition ctx) -> Context (isSet def.sort :: ctx)

%name Context ctx, ctx', ctx''

-- Convenience Casts -----------------------------------------------------------

public export
Cast Sort (Constructor ctx) where
  cast s = Sort s

public export
Cast Sort (NormalForm True ctx) where
  cast s = Cnstr $ cast s

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

prettyPrec : Prec -> NormalForm b ctx -> Doc ann

prettyPrecCnstr : Prec -> Constructor ctx -> Doc ann
prettyPrecCnstr d (Sort s) = prettyPrec d s
prettyPrecCnstr d (Pi s s' var a b) =
  let lhs = case var of
              "" => align (prettyPrec App a)
              _ => parens (pretty var <++> colon <+> softline <+> align (prettyPrec Open a))
  in
  parenthesise (d > Open) $
  group $
  lhs <++> pretty "->" <+> line <+> align (prettyPrec Open b)
prettyPrecCnstr d (Lambda _ var t) =
  parenthesise (d > Open) $
  group $
  backslash <+> pretty var <++>
  pretty "=>" <+> line <+>
  align (prettyPrec Open t)
prettyPrecCnstr d (Sigma s s' var a b) =
  let lhs = case var of
              "" => align (prettyPrec App a)
              _ => parens (pretty var <++> colon <+> softline <+> align (prettyPrec Open a))
  in
  parenthesise (d > Open) $
  group $
  lhs <++> pretty "**" <+> line <+> align (prettyPrec Open b)
prettyPrecCnstr d (Pair s s' prf t u) =
  angles $
  group $
  neutral <++> prettyPrec Open t <+> comma <+> line <+> prettyPrec Open u <++> neutral
prettyPrecCnstr d Bool = pretty "Bool"
prettyPrecCnstr d True = pretty "True"
prettyPrecCnstr d False = pretty "False"
prettyPrecCnstr d Top = pretty "()"
prettyPrecCnstr d Bottom = pretty "Void"

prettyPrecNtrl : Prec -> Neutral ctx -> Doc ann
prettyPrecNtrl d (Var var sort prf i) = pretty "\{var}@\{show $ elemToNat i}"
prettyPrecNtrl d (App _ t u) =
  parenthesise (d >= App) $
  group $
  vsep [prettyPrecNtrl Open t, prettyPrec App u]
prettyPrecNtrl d (Fst _ t) =
  parenthesise (d >= App) $
  group $
  vsep [pretty "fst", prettyPrecNtrl App t]
prettyPrecNtrl d (Snd _ t) =
  parenthesise (d >= App) $
  group $
  vsep [pretty "snd", prettyPrecNtrl App t]
prettyPrecNtrl d (If t f g) =
  parenthesise (d >= App) $
  group $
  vsep [pretty "if", prettyPrecNtrl App t, prettyPrec App f, prettyPrec App g]
prettyPrecNtrl d Absurd = pretty "absurd"
prettyPrecNtrl d (Equal a t u) =
  parenthesise (d >= Equal) $
  group $
  prettyPrec Equal t <++> pretty "~" <+> line <+> prettyPrec Equal u
prettyPrecNtrl d (EqualL a b) =
  parenthesise (d >= Equal) $
  group $
  prettyPrecNtrl Equal a <++> pretty "~" <+> line <+> prettyPrec Equal b
prettyPrecNtrl d (EqualR a b) =
  parenthesise (d >= Equal) $
  group $
  prettyPrecCnstr Equal a <++> pretty "~" <+> line <+> prettyPrecNtrl Equal b
prettyPrecNtrl d (EqualStuck a b) =
  parenthesise (d >= Equal) $
  group $
  prettyPrecCnstr Equal a <++> pretty "~" <+> line <+> prettyPrecCnstr Equal b
prettyPrecNtrl d (CastL a b t) =
  parenthesise (d >= App) $
  group $
  vsep [pretty "cast", prettyPrecNtrl App a, prettyPrec App b, prettyPrec App t]
prettyPrecNtrl d (CastR a b t) =
  parenthesise (d >= App) $
  group $
  vsep [pretty "cast", prettyPrecCnstr App a, prettyPrecNtrl App b, prettyPrec App t]
prettyPrecNtrl d (CastStuck a b t) =
  parenthesise (d >= App) $
  group $
  vsep [pretty "cast", prettyPrecCnstr App a, prettyPrecCnstr App b, prettyPrec App t]

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

export
Pretty (Constructor ctx) where
  prettyPrec = prettyPrecCnstr

export
Pretty (Neutral ctx) where
  prettyPrec = prettyPrecNtrl

export
Pretty (NormalForm b ctx) where
  prettyPrec = NormalForm.prettyPrec

-- Renaming --------------------------------------------------------------------

rename : NormalForm ~|> Hom Elem NormalForm

renameCnstr : Constructor ~|> Hom Elem Constructor
renameCnstr (Sort s) f = Sort s
renameCnstr (Pi s s' var a b) f = Pi s s' var (rename a f) (rename b (lift [(_ ** ())] f))
renameCnstr (Lambda s var t) f = Lambda s var (rename t (lift [(_ ** ())] f))
renameCnstr (Sigma s s' var a b) f =
  Sigma s s' var (rename a f) (rename b (lift [(_ ** ())] f))
renameCnstr (Pair s s' prf t u) f = Pair s s' prf (rename t f) (rename u f)
renameCnstr Bool f = Bool
renameCnstr True f = True
renameCnstr False f = False
renameCnstr Top f = Top
renameCnstr Bottom f = Bottom

renameNtrl : Neutral ~|> Hom Elem Neutral
renameNtrl (Var var sort prf i) f = Var var sort prf (f i)
renameNtrl (App b t u) f = App b (renameNtrl t f) (rename u f)
renameNtrl (Fst b t) f = Fst b (renameNtrl t f)
renameNtrl (Snd b t) f = Snd b (renameNtrl t f)
renameNtrl (If t u v) f = If (renameNtrl t f) (rename u f) (rename v f)
renameNtrl Absurd f = Absurd
renameNtrl (Equal a t u) f = Equal (renameNtrl a f) (rename t f) (rename u f)
renameNtrl (EqualL a b) f = EqualL (renameNtrl a f) (rename b f)
renameNtrl (EqualR a b) f = EqualR (renameCnstr a f) (renameNtrl b f)
renameNtrl (EqualStuck a b) f = EqualStuck (renameCnstr a f) (renameCnstr b f)
renameNtrl (CastL a b t) f = CastL (renameNtrl a f) (rename b f) (rename t f)
renameNtrl (CastR a b t) f = CastR (renameCnstr a f) (renameNtrl b f) (rename t f)
renameNtrl (CastStuck a b t) f = CastStuck (renameCnstr a f) (renameCnstr b f) (rename t f)

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

export
Unsorted.Rename Bool Constructor where
  rename = renameCnstr

export
Unsorted.Rename Bool Neutral where
  rename = renameNtrl
  
export
Sorted.Rename Bool NormalForm where
  rename = NormalForm.rename

export
PointedRename Bool (\b => (String, (s : Sort ** isSet s = b))) NormalForm where
  point {s = False} _ _ = Irrel
  point {s = True} (var, (s ** prf)) i = Ntrl (Var var s (eqToSo prf) i)