{-# LANGUAGE FlexibleContexts #-}
module PTS.Quote where

import Prelude hiding (div)

import PTS.AST
import PTS.Core hiding (nat, debug)
import PTS.Parser (parseTerm)
import Parametric.Error
import Tools.Errors.Class
import Tools.Errors(runErrors, runErrorsT)
import Parametric.Pretty(multiLine)
import Control.Monad.Trans
import PTS.Pretty
import PTS.Debug
import System.IO.Unsafe

lang = fomegastar
(axioms, sorts, relations, _) = lang

rep = (!! 0)
lam = (!! 1)
app = (!! 2)
tlam = (!! 3)
tapp = (!! 4)
ifzero = (!! 5)
nat = (!! 6)
add = (!! 7)
sub = (!! 8)
mul = (!! 9)
div = (!! 10)

varnames = ["R", "lam", "app", "tlam", "tapp", "ifzero", "nat", "add", "sub", "mul", "div"]

unsafeParse :: String -> Term String
unsafeParse text = 
  case runErrors (parseTerm fomegastarConstants "<Quote.hs>" text) of
    Right x -> x
    Left e -> error (showErrors e)

interface r = 
  [ p "* -> *"
  , p "Pi S : * . Pi T : * . (R S -> R T) -> R (S -> T)"
  , p "Pi S : * . Pi T : * . R (S -> T) -> R S -> R T"
  , p "Pi S : ** . Pi T : S -> * . (Pi X : S . R (T X)) -> R (Pi X : S . T X)"
  , p "Pi S : ** . Pi T : S -> * . R (Pi X : S . T X) -> (Pi X : S . R (T X))"
  , p "Pi T : * . R Nat -> R T -> R T -> R T"
  , p "Nat -> R Nat"
  , p "R Nat -> R Nat -> R Nat"
  , p "R Nat -> R Nat -> R Nat"
  , p "R Nat -> R Nat -> R Nat"
  , p "R Nat -> R Nat -> R Nat"
  ] where
  
  p text = unsafeParse $ text >>= \x -> if x == 'R' then r else return x

unsafeQuote :: Term String -> Term String
unsafeQuote t = unsafePerformIO $ do
  res <- runErrorsT (quote (map Var varnames) 0 [] t)
  case res of 
    Left e -> fail (showErrors e)
    Right q -> return q 

unsafeQuoteQuote :: Term String -> Term String
unsafeQuoteQuote t = unsafePerformIO $ do
  res <- runErrorsT (quotequote 0 [] t)
  case res of
    Left e -> fail (showErrors e)
    Right q -> return q 


quotequote :: (MonadIO m, MonadErrors Errors m, Functor m) => Int -> [(String, Term String)] -> Term String -> m (Term String)
quotequote d ctx t = do
  let vars = map (freshvarl (allvars t)) varnames 
  q <- quote (map Var vars) 0 [] t
  return $ foldr (uncurry Lam) q (zip vars (interface (rep vars)))  

debug :: MonadIO m => Int -> String -> [(String, Term String)] -> Term String -> m (Term String) -> m (Term String) 
debug d n ctx t result | DebugQuote `elem` options = do
  liftIO $ putStrLn $ 
    replicate (d+d) ' ' ++ 
    "--> " ++ n ++ ": " ++
    (let gamma = showCtx ctx in if null gamma then "" else gamma ++ " ") ++
    "|- <" ++ show t ++ "> = ???"
  x <- result
  liftIO $ putStrLn $ 
    replicate (d+d) ' ' ++ 
    "<-- " ++ n ++ ": " ++
    (let gamma = showCtx ctx in if null gamma then "" else gamma ++ " ") ++
    "|- <" ++ show t ++ "> = " ++ show x
  return x
  
debug d n ctx t result = result


quote :: (MonadIO m, MonadErrors Errors m, Functor m) => [Term String] -> Int -> [(String, Term String)] -> Term String -> m (Term String)

quote vars d ctx q@(Nat n) = debug d "QuoteNat" ctx q $ 
  return $ nat vars `App` Nat n

quote vars d ctx q@(Var x) = debug d "QuoteVar" ctx q $
  return q

quote vars d ctx q@(Lam x a b) = debug d "QuoteLam" ctx q $ do
  let (newx, newb, newctx) = bind ctx x a b
  Const s1 <- normalform `fmap` typecheck lang (succ d) ctx a
  tb <- typecheck lang (succ d) newctx newb
  Const s2 <- normalform `fmap` typecheck lang (succ d) newctx tb
  b' <- quote vars (succ d) newctx newb
  
  r <- case (s1, s2) of
    ("*", "*") -> return $ lam vars `App` a `App` tb `App` Lam newx (rep vars `App` a) b' 
    ("**", "*") -> return $ tlam vars `App` a `App` Lam newx a tb `App` Lam newx a b' 
    _ -> fail $ "cannot quote non-value-level term " ++ show q
  
  return r

quote vars d ctx q@(App f x) = debug d "QuoteApp" ctx q $do
  Pi v a b <- normalform `fmap` typecheck lang (succ d) ctx f
  Const s1 <- normalform `fmap` typecheck lang (succ d) ctx a
  let (newv, newb, newctx) = bind ctx v a b
  Const s2 <- normalform `fmap` typecheck lang (succ d) newctx newb
  f' <- quote vars (succ d) ctx f
  case (s1, s2) of
    ("*", "*") -> do 
      x' <- quote vars (succ d) ctx x
      return $ app vars `App` a `App` b `App` f' `App` x'
    ("**", "*") -> return $ tapp vars `App` a `App` Lam v a b `App` f' `App` x
    _ -> fail $ "cannot quote non-value-level term" ++ show q

quote vars d ctx q@(NatOp n f x y) = debug d "QuoteNatOp" ctx q $ do
  x' <- quote vars d ctx x
  y' <- quote vars d ctx y
  case n of
    "add" -> return $ add vars `App` x' `App` y'
    "sub" -> return $ sub vars `App` x' `App` y'
    "mul" -> return $ mul vars `App` x' `App` y'
    "div" -> return $ div vars `App` x' `App` y'

quote vars d ctx q@(IfZero c t e) = debug d "QuoteIfZero" ctx q $ do
  c' <- quote vars d ctx c
  t' <- quote vars d ctx t
  e' <- quote vars d ctx e
  return $ ifzero vars `App` c' `App` t' `App` e'

quote vars d ctx q@(Pos p t) = do
  t' <- quote vars d ctx t
  return $ Pos p t'

quote vars d ctx q = fail $ "cannot quote non-value-level term" ++ show q
