module Obs.NormalForm

import Data.List.Elem

import Obs.Logging
import Obs.Pretty
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
record DeclForm (ctx : List Relevance) where
  constructor MkDecl
  var : Maybe String
  type : TypeNormalForm ctx

public export
data Constructor : Unsorted.Family Relevance where
  Universe : (s : Universe) -> Constructor ctx
  Pi : (domainSort, codomainSort : Universe) ->
    (domain : DeclForm ctx) ->
    (codomain : TypeNormalForm (relevance domainSort :: ctx)) ->
    Constructor ctx
  Lambda : (domainRel : Relevance) ->
    (var : Maybe String) ->
    (body : NormalForm Relevant (domainRel :: ctx)) ->
    Constructor ctx
  Sigma : (indexSort, elementSort : Universe) ->
    (index : DeclForm ctx) ->
    (element : TypeNormalForm (relevance indexSort :: ctx)) ->
    Constructor ctx
  Pair : (indexRel, elementRel : Relevance) ->
    (prf : IsRelevant (pair indexRel elementRel)) ->
    (first : NormalForm indexRel ctx) ->
    (second : NormalForm elementRel ctx) ->
    Constructor ctx
  Bool : Constructor ctx
  True : Constructor ctx
  False : Constructor ctx
  Box : (prop : TypeNormalForm ctx) -> Constructor ctx
  MkBox : Constructor ctx
  Top : Constructor ctx
  Bottom : Constructor ctx

public export
data Neutral : Unsorted.Family Relevance where
  Var : (var : String) -> (i : Elem Relevant ctx) -> Neutral ctx
  App : (argRel : Relevance) -> (fun : Neutral ctx) -> (arg : NormalForm argRel ctx) -> Neutral ctx
  First : (secondRel : Relevance) -> (arg : Neutral ctx) -> Neutral ctx
  Second : (firstRel : Relevance) -> (arg : Neutral ctx) -> Neutral ctx
  If : (discriminant : Neutral ctx) -> (true, false : NormalForm Relevant ctx) -> Neutral ctx
  Absurd : Neutral ctx
  Equal : (type : Neutral ctx) -> (left, right : NormalForm Relevant ctx) -> Neutral ctx
  EqualL : (left : Neutral ctx) -> (right : TypeNormalForm ctx) -> Neutral ctx
  EqualR : (left : Constructor ctx) -> (right : Neutral ctx) -> Neutral ctx
  EqualStuck : (left, right : Constructor ctx) -> Neutral ctx
  CastL : (oldType : Neutral ctx) ->
    (newType : TypeNormalForm ctx) ->
    (value : NormalForm Relevant ctx) ->
    Neutral ctx
  CastR : (oldType : Constructor ctx) ->
    (newType : Neutral ctx) ->
    (value : NormalForm Relevant ctx) ->
    Neutral ctx
  CastStuck : (oldType, newType : Constructor ctx) ->
    (value : 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 Parameter (ctx : List Relevance) where
  constructor MkParameter
  name : WithBounds String
  sort : Universe
  ty : TypeNormalForm ctx

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)
  (::<) : Context ctx -> (param : Parameter ctx) -> Context (relevance param.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

prettyPrecDecl : DeclForm ctx -> Doc ann
prettyPrecDecl (MkDecl var type) = prettyDecl var (prettyPrec Open type)

prettyPrecCnstr : Prec -> Constructor ctx -> Doc ann
prettyPrecCnstr d (Universe {s}) = prettyPrec d s
prettyPrecCnstr d (Pi {domainSort, codomainSort, domain, codomain}) =
  prettyDeclForm d [prettyPrecDecl domain] "->" (prettyPrec Open codomain)
prettyPrecCnstr d (Lambda {domainRel, var, body}) =
  prettyLambda d (maybe "_" id var) (prettyPrec Open body)
prettyPrecCnstr d (Sigma {indexSort, elementSort, index, element}) =
  prettyDeclForm d [prettyPrecDecl index] "**" (prettyPrec Open element)
prettyPrecCnstr d (Pair {indexRel, elementRel, prf, first, second}) =
  prettyPair (prettyPrec Open first) (prettyPrec Open second)
prettyPrecCnstr d Bool = pretty "Bool"
prettyPrecCnstr d True = pretty "True"
prettyPrecCnstr d False = pretty "False"
prettyPrecCnstr d (Box {prop}) = prettyApp d (pretty "Box") [prettyPrec App prop]
prettyPrecCnstr d MkBox = prettyApp d (pretty "box") [pretty "_"]
prettyPrecCnstr d Top = pretty "()"
prettyPrecCnstr d Bottom = pretty "Void"

prettyPrecNtrl : Prec -> Neutral ctx -> Doc ann
prettyPrecNtrl d (Var {var, i}) = pretty "\{var}@\{show $ elemToNat i}"
prettyPrecNtrl d (App {argRel, fun, arg}) =
  prettyApp d (prettyPrecNtrl Open fun) [prettyPrec App arg]
prettyPrecNtrl d (First {secondRel, arg}) =
  prettyApp d (pretty "fst") [prettyPrecNtrl App arg]
prettyPrecNtrl d (Second {firstRel, arg}) =
  prettyApp d (pretty "snd") [prettyPrecNtrl App arg]
prettyPrecNtrl d (If {discriminant, true, false}) =
  prettyApp d (pretty "if") $
  [pretty "_", prettyPrecNtrl App discriminant, prettyPrec App true, prettyPrec App false]
prettyPrecNtrl d Absurd = prettyApp d (pretty "absurd") [pretty "_"]
prettyPrecNtrl d (Equal {type, left, right}) =
  prettyEqual d (prettyPrec Equal left) (prettyPrec Equal right)
prettyPrecNtrl d (EqualL {left, right}) =
  prettyEqual d (prettyPrecNtrl Equal left) (prettyPrec Equal right)
prettyPrecNtrl d (EqualR {left, right}) =
  prettyEqual d (prettyPrecCnstr Equal left) (prettyPrecNtrl Equal right)
prettyPrecNtrl d (EqualStuck {left, right}) =
  prettyEqual d (prettyPrecCnstr Equal left) (prettyPrecCnstr Equal right)
prettyPrecNtrl d (CastL {oldType, newType, value}) =
  prettyApp d (pretty "cast") $
  [prettyPrecNtrl App oldType, prettyPrec App newType, pretty "_", prettyPrec App value]
prettyPrecNtrl d (CastR {oldType, newType, value}) =
  prettyApp d (pretty "cast") $
  [prettyPrecCnstr App oldType, prettyPrecNtrl App newType, pretty "_", prettyPrec App value]
prettyPrecNtrl d (CastStuck {oldType, newType, value}) =
  prettyApp d (pretty "cast") $
  [prettyPrecCnstr App oldType, prettyPrecCnstr App newType, pretty "_", prettyPrec App value]

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

renameDecl : DeclForm ~|> Hom Elem DeclForm
renameDecl (MkDecl var type) f = MkDecl var (rename type f)

renameCnstr : Constructor ~|> Hom Elem Constructor
renameCnstr (Universe {s}) f = Universe {s}
renameCnstr (Pi {domainSort, codomainSort, domain, codomain}) f =
  Pi
    { domainSort
    , codomainSort
    , domain = renameDecl domain f
    , codomain = rename codomain (lift [(_ ** ())] f)
    }
renameCnstr (Lambda {domainRel, var, body}) f =
  Lambda {domainRel, var, body = rename body (lift [(_ ** ())] f)}
renameCnstr (Sigma {indexSort, elementSort, index, element}) f =
  Sigma
    { indexSort
    , elementSort
    , index = renameDecl index f
    , element = rename element (lift [(_ ** ())] f)
    }
renameCnstr (Pair {indexRel, elementRel, prf, first, second}) f =
  Pair
    { indexRel
    , elementRel
    , prf
    , first = rename first f
    , second = rename second f
    }
renameCnstr Bool f = Bool
renameCnstr True f = True
renameCnstr False f = False
renameCnstr (Box {prop}) f = Box {prop = rename prop f}
renameCnstr MkBox f = MkBox
renameCnstr Top f = Top
renameCnstr Bottom f = Bottom

renameNtrl : Neutral ~|> Hom Elem Neutral
renameNtrl (Var {var, i}) f = Var {var, i = f i}
renameNtrl (App {argRel, fun, arg}) f = App {argRel, fun = renameNtrl fun f, arg = rename arg f}
renameNtrl (First {secondRel, arg}) f = First {secondRel, arg = renameNtrl arg f}
renameNtrl (Second {firstRel, arg}) f = Second {firstRel, arg = renameNtrl arg f}
renameNtrl (If {discriminant, true, false}) f =
  If {discriminant = renameNtrl discriminant f, true = rename true f, false = rename false f}
renameNtrl Absurd f = Absurd
renameNtrl (Equal {type, left, right}) f =
  Equal {type = renameNtrl type f, left = rename left f, right = rename right f}
renameNtrl (EqualL {left, right}) f = EqualL {left = renameNtrl left f, right = rename right f}
renameNtrl (EqualR {left, right}) f =
  EqualR {left = renameCnstr left f, right = renameNtrl right f}
renameNtrl (EqualStuck {left, right}) f =
  EqualStuck {left = renameCnstr left f, right = renameCnstr right f}
renameNtrl (CastL {oldType, newType, value}) f =
 CastL {oldType = renameNtrl oldType f, newType = rename newType f, value = rename value f}
renameNtrl (CastR {oldType, newType, value}) f =
 CastR {oldType = renameCnstr oldType f, newType = renameNtrl newType f, value = rename value f}
renameNtrl (CastStuck {oldType, newType, value}) f =
 CastStuck
   { oldType = renameCnstr oldType f
   , newType = renameCnstr newType f
   , value = rename value 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 => Maybe String) NormalForm where
  point {s = Irrelevant} _ _ = Irrel
  point {s = Relevant} var i =
    Ntrl (Var {var = maybe "" id var, i})

-- Useful Constructors ---------------------------------------------------------

public export
record ContainerTy (ctx : List Relevance) where
  constructor MkContainerTy
  inputSort, shapeSort, positionSort, outputSort : Universe
  input, output : TypeNormalForm ctx