module Obs.NormalForm

import Data.List.Elem

import Obs.Substitution
import Obs.Universe

import Text.Bounded
import Text.PrettyPrint.Prettyprinter

%default total

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

data NormalForm  : Sorted.Family Relevance

public export
TypeNormalForm : Unsorted.Family Relevance
TypeNormalForm = NormalForm Relevant

public export
data Constructor : Unsorted.Family Relevance where
  Universe : (s : Universe) -> Constructor ctx
  Pi : (s, s' : Universe)
    -> (var : String)
    -> (a : TypeNormalForm ctx)
    -> (b : TypeNormalForm (relevance s :: ctx))
    -> Constructor ctx
  Lambda : (s : Universe)
    -> (var : String)
    -> NormalForm Relevant (relevance s :: ctx)
    -> Constructor ctx
  Sigma : (s, s' : Universe)
    -> (var : String)
    -> (a : TypeNormalForm ctx)
    -> (b : TypeNormalForm (relevance s :: ctx))
    -> Constructor ctx
  Pair : (s, s' : Universe)
    -> (prf : IsRelevant (max s s'))
    -> NormalForm (relevance s) ctx
    -> NormalForm (relevance 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 Relevance where
  Var : (var : String)
    -> (sort : Universe)
    -> (prf : IsRelevant sort)
    -> (i : Elem Relevant ctx)
    -> Neutral ctx
  App : (r : Relevance) -> Neutral ctx -> NormalForm r ctx -> Neutral ctx
  Fst : (r : Relevance) -> Neutral ctx -> Neutral ctx
  Snd : (r : Relevance) -> Neutral ctx -> Neutral ctx
  If : Neutral ctx -> (f, g : NormalForm Relevant ctx) -> Neutral ctx
  Absurd : Neutral ctx
  Equal : Neutral ctx -> NormalForm Relevant ctx -> NormalForm Relevant ctx -> Neutral ctx
  EqualL : (a : Neutral ctx) -> (b : TypeNormalForm ctx) -> Neutral ctx
  EqualR : (a : Constructor ctx) -> (b : Neutral ctx) -> Neutral ctx
  EqualStuck : (a, b : Constructor ctx) -> Neutral ctx
  CastL : (a : Neutral ctx) ->
    (b : TypeNormalForm ctx) ->
    NormalForm Relevant ctx ->
    Neutral ctx
  CastR : (a : Constructor ctx) -> (b : Neutral ctx) -> NormalForm Relevant ctx -> Neutral ctx
  CastStuck : (a, b : Constructor ctx) -> NormalForm Relevant ctx -> Neutral ctx

public export
data NormalForm : Sorted.Family Relevance where
  Ntrl : Neutral ~|> NormalForm Relevant
  Cnstr : Constructor ~|> NormalForm Relevant
  Irrel : NormalForm Irrelevant ctx

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

public export
record Definition (ctx : List Relevance) where
  constructor MkDefinition
  name : WithBounds String
  sort : Universe
  ty   : TypeNormalForm ctx
  tm   : NormalForm (relevance sort) ctx

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

%name Context ctx, ctx', ctx''

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

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

public export
Cast Universe (NormalForm Relevant ctx) where
  cast s = Cnstr $ cast s

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

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

prettyPrecCnstr : Prec -> Constructor ctx -> Doc ann
prettyPrecCnstr d (Universe 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 (Universe s) f = Universe 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 Relevance Constructor where
  rename = renameCnstr

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

export
PointedRename Relevance (\r => (String, (s : Universe ** relevance s = r))) NormalForm where
  point {s = Irrelevant} _ _ = Irrel
  point {s = Relevant} (var, (s@(Set _) ** prf)) i = Ntrl (Var var s Set i)