||| Abstract syntax for terms.
module Obs.Term

import Data.Fin

import Obs.Pretty
import Obs.Universe

import Text.Bounded

%default total

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

infix 5 ::<

data Term : Nat -> Type

public export
record DeclForm (n : Nat) where
  constructor MkDecl
  var : WithBounds String
  type : WithBounds (Term n)

public export
record LambdaForm (n : Nat) where
  constructor MkLambda
  var : WithBounds String
  body : WithBounds (Term (S n))

public export
data Term : Nat -> Type where
  Var : (var : String) -> (i : Fin n) -> Term n
  -- Universes
  Universe : (s : Universe) -> Term n
  -- Dependent Product
  Pi : (domain : DeclForm n) ->
    (codomain : WithBounds (Term (S n))) ->
    Term n
  Lambda : (body : LambdaForm n) -> Term n
  App : (fun, arg : WithBounds (Term n)) -> Term n
  -- Dependent Sums
  Sigma : (index : DeclForm n) ->
    (element : WithBounds (Term (S n))) ->
    Term n
  Pair : (first, second : WithBounds (Term n)) -> Term n
  First : (arg : WithBounds (Term n)) -> Term n
  Second : (arg : WithBounds (Term n)) -> Term n
  -- Containers
  Container : (input, output : WithBounds (Term n)) ->
    (shapeSort, positionSort : WithBounds Universe) ->
    Term n
  MkContainer : (shape, position, next : WithBounds (Term n)) -> Term n
  Shape : (arg : WithBounds (Term n)) -> Term n
  Position : (arg : WithBounds (Term n)) -> Term n
  Next : (arg : WithBounds (Term n)) -> Term n
  Sem : (pred : LambdaForm n) -> (arg : WithBounds (Term n)) -> Term n
  -- Booleans
  Bool : Term n
  True : Term n
  False : Term n
  If : (returnType : LambdaForm n) ->
    (discriminant, true, false : WithBounds (Term n)) ->
    Term n
  -- True
  Top : Term n
  Point : Term n
  -- False
  Bottom : Term n
  Absurd : (contradiction : WithBounds (Term n)) -> Term n
  -- Equality
  Equal : (left, right : WithBounds (Term n)) -> Term n
  Refl : (value : WithBounds (Term n)) -> Term n
  Transp : (indexedType, oldIndex, value, newIndex, equality : WithBounds (Term n)) -> Term n
  Cast : (oldType, newType, equality, value : WithBounds (Term n)) -> Term n
  CastId : (value : WithBounds (Term n)) -> Term n

public export
record Parameter (n : Nat) where
  constructor MkParameter
  name : WithBounds String
  ty : WithBounds (Term n)

public export
record Definition (n : Nat) where
  constructor MkDefinition
  name : WithBounds String
  ty   : WithBounds (Term n)
  tm   : WithBounds (Term n)

public export
data Block : Nat -> Type where
  Nil  : Block 0
  (:<) : Block n -> Definition n -> Block (S n)
  (::<) : Block n -> Parameter n -> Block (S n)

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

prettyPrecBounds : Prec -> WithBounds (Term n) -> Doc ann

prettyPrecDecl : DeclForm n -> Doc ann
prettyPrecDecl (MkDecl var type) = prettyDecl (Just var.val) (prettyPrecBounds Open type)

prettyPrecLambda : Prec -> LambdaForm n -> Doc ann
prettyPrecLambda d (MkLambda var body) = prettyLambda d var.val (prettyPrecBounds Open body)

prettyPrec : Prec -> Term n -> Doc ann
prettyPrec d (Var {var, i}) = pretty var
prettyPrec d (Universe {s}) = prettyPrec d s
prettyPrec d (Pi {domain, codomain}) =
  prettyDeclForm d [prettyPrecDecl domain] "->" (prettyPrecBounds Open codomain)
prettyPrec d (Lambda {body}) = prettyPrecLambda d body
prettyPrec d (App {fun, arg}) = prettyApp d (prettyPrecBounds Open fun) [prettyPrecBounds App arg]
prettyPrec d (Sigma {index, element}) =
  prettyDeclForm d [prettyPrecDecl index] "**" (prettyPrecBounds Open element)
prettyPrec d (Pair {first, second}) =
  prettyPair (prettyPrecBounds Open first) (prettyPrecBounds Open second)
prettyPrec d (First {arg}) = prettyApp d (pretty "fst") [prettyPrecBounds App arg]
prettyPrec d (Second {arg}) = prettyApp d (pretty "snd") [prettyPrecBounds App arg]
prettyPrec d (Container {input, output, shapeSort, positionSort}) =
  prettyApp d (pretty "Container") $
  [ prettyPrecBounds App input
  , prettyPrecBounds App output
  , prettyPrec App shapeSort.val
  , prettyPrec App positionSort.val
  ]
prettyPrec d (MkContainer {shape, position, next}) =
  prettyApp d (pretty "MkContainer") $
  [ prettyPrecBounds App shape
  , prettyPrecBounds App position
  , prettyPrecBounds App next
  ]
prettyPrec d (Shape {arg}) = prettyApp d (pretty "Shape") [prettyPrecBounds App arg]
prettyPrec d (Position {arg}) = prettyApp d (pretty "Position") [prettyPrecBounds App arg]
prettyPrec d (Next {arg}) = prettyApp d (pretty "nextIndex") [prettyPrecBounds App arg]
prettyPrec d (Sem {pred, arg}) =
  prettyApp d (pretty "extension") [prettyPrecLambda App pred, prettyPrecBounds App arg]
prettyPrec d Bool = pretty "Bool"
prettyPrec d True = pretty "True"
prettyPrec d False = pretty "False"
prettyPrec d (If {returnType, discriminant, true, false}) =
  prettyApp d (pretty "if") $
  [ prettyPrecLambda App returnType
  , prettyPrecBounds App discriminant
  , prettyPrecBounds App true
  , prettyPrecBounds App false
  ]
prettyPrec d Top = pretty "()"
prettyPrec d Point = pretty "tt"
prettyPrec d Bottom = pretty "Void"
prettyPrec d (Absurd {contradiction}) =
  prettyApp d (pretty "absurd") [prettyPrecBounds App contradiction]
prettyPrec d (Equal {left, right}) =
  prettyEqual d (prettyPrecBounds Equal left) (prettyPrecBounds Equal right)
prettyPrec d (Refl {value}) = prettyApp d (pretty "refl") [prettyPrecBounds App value]
prettyPrec d (Transp {indexedType, oldIndex, value, newIndex, equality}) =
  prettyApp d (pretty "transp") $
  [ prettyPrecBounds App indexedType
  , prettyPrecBounds App oldIndex
  , prettyPrecBounds App value
  , prettyPrecBounds App newIndex
  , prettyPrecBounds App equality
  ]
prettyPrec d (Cast {oldType, newType, equality, value}) =
  prettyApp d (pretty "cast") $
  [ prettyPrecBounds App oldType
  , prettyPrecBounds App newType
  , prettyPrecBounds App equality
  , prettyPrecBounds App value
  ]
prettyPrec d (CastId {value}) = prettyApp d (pretty "castRefl") [prettyPrecBounds App value]

prettyPrecBounds d (MkBounded v _ _) = prettyPrec d v

export
Pretty (Term n) where
  prettyPrec = Term.prettyPrec

export
Pretty (Parameter n) where
  pretty param = group $
    pretty "parameter" <++> pretty param.name.val <++> colon <+> softline <+> pretty param.ty.val

export
Pretty (Definition n) where
  pretty def = group $
    pretty def.name.val <++> colon <+> softline <+> pretty def.ty.val <+> hardline <+>
    pretty def.name.val <++> equals <+> softline <+> pretty def.tm.val

export
Pretty (Block n) where
  pretty []           = neutral
  pretty ([] :< def)  = pretty def
  pretty ([] ::< param)  = pretty param
  pretty (blk :< def) = pretty blk <+> hardline <+> hardline <+> pretty def
  pretty (blk ::< param) = pretty blk <+> hardline <+> hardline <+> pretty param