module Obs.Abstract

import Data.Fin
import Data.String

import Obs.Logging
import Obs.Syntax
import Obs.Term

import System
import Text.Bounded

%default total

Context : Nat -> Type
Context n = List (String, Fin n)

bind : Context n -> String -> Context (S n)
bind ctx str = (str, 0) :: map (mapSnd FS) ctx

abstractSyntaxBounds : Context n -> WithBounds Syntax -> Logging ann (WithBounds (Term n))

abstractDecl : Context n -> DeclForm -> Logging ann (DeclForm n)
abstractDecl ctx (MkDecl var type) = do
  type <- abstractSyntaxBounds ctx type
  pure (MkDecl var type)

abstractLambda : Context n -> LambdaForm -> Logging ann (LambdaForm n)
abstractLambda ctx (MkLambda var body) = do
  body <- abstractSyntaxBounds (bind ctx var.val) body
  pure (MkLambda var body)

abstractSyntax : Context n -> Syntax -> Logging ann (Term n)
abstractSyntax ctx (Var {var}) = do
  let Just i = lookup var ctx
    | Nothing => inScope "unbound variable" $ fatal var
  pure (Var {var, i})
abstractSyntax ctx (Universe {s}) = pure (Universe {s})
abstractSyntax ctx (Pi {domain, codomain}) = do
  domain <- abstractDecl ctx domain

  let ctx = bind ctx domain.var.val
  codomain <- abstractSyntaxBounds ctx codomain

  pure (Pi {domain, codomain})
abstractSyntax ctx (Lambda {body}) = do
  body <- abstractLambda ctx body
  pure (Lambda {body})
abstractSyntax ctx (App {fun, arg}) = do
  fun <- abstractSyntaxBounds ctx fun
  arg <- abstractSyntaxBounds ctx arg
  pure (App {fun, arg})
abstractSyntax ctx (Sigma {index, element}) = do
  index <- abstractDecl ctx index

  let ctx = bind ctx index.var.val
  element <- abstractSyntaxBounds ctx element

  pure (Sigma {index, element})
abstractSyntax ctx (Pair {first, second}) = do
  first <- abstractSyntaxBounds ctx first
  second <- abstractSyntaxBounds ctx second
  pure (Pair {first, second})
abstractSyntax ctx (First {arg}) = do
  arg <- abstractSyntaxBounds ctx arg
  pure (First {arg})
abstractSyntax ctx (Second {arg}) = do
  arg <- abstractSyntaxBounds ctx arg
  pure (Second {arg})
abstractSyntax ctx Bool = pure Bool
abstractSyntax ctx True = pure True
abstractSyntax ctx False = pure False
abstractSyntax ctx (If {returnType, discriminant, true, false}) = do
  returnType <- abstractLambda ctx returnType
  discriminant <- abstractSyntaxBounds ctx discriminant
  true <- abstractSyntaxBounds ctx true
  false <- abstractSyntaxBounds ctx false
  pure (If {returnType, discriminant, true, false})
abstractSyntax ctx Top = pure Top
abstractSyntax ctx Point = pure Point
abstractSyntax ctx Bottom = pure Bottom
abstractSyntax ctx (Absurd {contradiction}) = do
  contradiction <- abstractSyntaxBounds ctx contradiction
  pure (Absurd {contradiction})
abstractSyntax ctx (Equal {left, right}) = do
  left <- abstractSyntaxBounds ctx left
  right <- abstractSyntaxBounds ctx right
  pure (Equal {left, right})
abstractSyntax ctx (Refl {value}) = do
  value <- abstractSyntaxBounds ctx value
  pure (Refl {value})
abstractSyntax ctx (Transp {indexedType, oldIndex, value, newIndex, equality}) = do
  indexedType <- abstractSyntaxBounds ctx indexedType
  oldIndex <- abstractSyntaxBounds ctx oldIndex
  value <- abstractSyntaxBounds ctx value
  newIndex <- abstractSyntaxBounds ctx newIndex
  equality <- abstractSyntaxBounds ctx equality
  pure (Transp {indexedType, oldIndex, value, newIndex, equality})
abstractSyntax ctx (Cast {oldType, newType, equality, value}) = do
  oldType <- abstractSyntaxBounds ctx oldType
  newType <- abstractSyntaxBounds ctx newType
  equality <- abstractSyntaxBounds ctx equality
  value <- abstractSyntaxBounds ctx value
  pure (Cast {oldType, newType, equality, value})
abstractSyntax ctx (CastId {value}) = do
  value <- abstractSyntaxBounds ctx value
  pure (CastId {value})

abstractSyntaxBounds ctx val@(MkBounded v irrel bnds) = do
  term <- inBounds $ (MkBounded (abstractSyntax ctx v) irrel bnds)
  pure $ map (const term) val

abstractDefinition : Context n -> Definition -> Logging ann (Definition n)
abstractDefinition ctx def = pure $ MkDefinition
  { name = def.name
  , ty = !(abstractSyntaxBounds ctx def.ty)
  , tm = !(abstractSyntaxBounds ctx def.tm)
  }

export
abstractBlock : (defs : List Definition) -> Logging ann (Block (length defs))
abstractBlock defs =
  let asContext : Block n -> Context n
      asContext []            = []
      asContext (defs :< def) = bind (asContext defs) def.name.val
  in let helper : Block n -> (defs : List Definition) -> Logging ann (Block (length defs + n))
         helper blk []            = pure blk
         helper blk (def :: defs) = do
           def <- abstractDefinition (asContext blk) def
           (rewrite plusSuccRightSucc (length defs) n in helper (blk :< def) defs)
  in rewrite sym $ plusZeroRightNeutral (length defs) in helper [] defs