module MProver.Checker where

import MProver.Syntax
import MProver.Monad
import MProver.Eval
import MProver.PPrint

import Control.Monad.Reader
import Control.Monad.Error
import Control.Monad.Identity

import qualified Data.Map as Map
import Data.List (nub)

import Unbound.LocallyNameless

import Text.PrettyPrint.HughesPJ

--import Debug.Trace (trace)

check :: (Monad m) => Proof -> MPT m Formula
check (ForallIExpr ty b) = do
                             (x,pr) <- unbind b
                             fopr   <- localEE (bindEnv x (Just ty,Nothing)) (check pr)
                             return (ForallExpr ty (bind x fopr))
check (ProofImpl fo b)   = do
                             (x,pr) <- unbind b
                             fopr <- localPE (bindEnv x fo) (check pr)
                             return (ForallProof fo fopr)
check (ProofVar x)          = do
                                env <- askPE
                                case Map.lookup x env of
                                  (Just fo) -> return fo
                                  _         -> throwError $ "unbound proof variable " ++ name2String x
check (ProofAppProof pr1 pr2) = do
                                  fopr1 <- check pr1
                                  fopr2 <- check pr2
                                  case fopr1 of
                                    (ForallProof fo1 fo2) -> if fo1 `aeq` fopr2
                                                               then return fo2
                                                               else throwError "argument type mismatch when applying proof"
                                    _                     -> throwError "attempt to apply a proof of a formula whose classifier is not universal quantification or implication"
check (ProofAppExpr pr e)   = do
                                fopr <- check pr
                                -- FIXME: type check e?
                                case fopr of
                                  (ForallExpr _ b) -> do
                                                        (x,fo) <- unbind b
                                                        return (subst x e fo)
                                  _                   -> throwError "attempt to apply a proof of a formula whose classifier is not universal quantification or implication"
check (ProofAnno Eval fo)   = case fo of
                                (FormulaEq e1 e2) -> do
                                                       env <- ask
                                                       ok  <- tryJoin e1 e2                                                       
                                                       if ok
                                                         then return (FormulaEq e1 e2)
                                                         else do
                                                                e1_ <- ppr e1
                                                                e2_ <- ppr e2
                                                                throwError $ show $ text "Error: could not join terms:" $+$ nest 4 e1_ $+$ text "" $+$ text "and:" $+$ nest 4 e2_
                                _                 -> throwError "eval formula must be an equality"
check Eval                  = throwError "eval must be annotated with ::: formula"
check (Trans pr1 pr2)       = do
                                fopr1 <- check pr1
                                fopr2 <- check pr2
                                case fopr1 of
                                  (FormulaEq e1 e2) -> case fopr2 of
                                                         (FormulaEq e3 e4) -> if e2 `aeq` e3
                                                                                then return (FormulaEq e1 e4)
                                                                                else do
                                                                                       e2_ <- ppr e2
                                                                                       e3_ <- ppr e3
                                                                                       throwError $ show $ text "left and right sides of trans proof are not compatible:" $+$ nest 8 e2_ $+$ text "and" $+$ nest 8 e3_
                                                         _                 -> throwError "left side of trans proof is not an equality proof"
check (Subst b)             = do
                                (x,(pr,e)) <- unbind b
                                fopr <- check pr
                                case fopr of
                                  (FormulaEq e1 e2) -> return (FormulaEq (subst x e1 e) (subst x e2 e))
                                  _                 -> throwError "\"by\" clause of a case proof is not an equality proof"
check (Symm pr)             = do
                                fopr <- check pr
                                case fopr of
                                  (FormulaEq e1 e2) -> return (FormulaEq e2 e1)
                                  _                 -> throwError "symm must be applied to an equality proof"
check (ProofCase e fo b) = do
                             (x,alts) <- unbind b
                             tot <- checkTotality alts
                             if not tot
                               then throwError "Pattern matching is not exhaustive"
                               else do
                                 mapM_ (checkAlt e fo x) alts
                                 return fo
check (ProofAnno pr fo)     = do
                                fopr <- check pr
                                if fopr `aeq` fo
                                  then return fo
                                  else do
                                    fo_   <- ppr fo
                                    fopr_ <- ppr fopr
                                    throwError $ show $ text "type mismatch in annotated proof. Expected:" $+$ nest 8 fo_ $+$ text "Got:" $+$ nest 8 fopr_
check (ProofBisim ctor prs) = do
                                -- FIXME: need to typecheck prs's fos?
                                dte <- askCE
                                case Map.lookup ctor dte of
                                  (Just (b,n)) -> do
                                    (xs,ConstrDecl _ tys) <- unbind b
                                    if length prs /= length tys
                                      then throwError "wrong arity in bisim"
                                      else do
                                        fos <- mapM check prs
                                        -- FIXME: pattern exhaustiveness here (gonna crash with pmf in some cases)
                                        let eql (FormulaEq e _) = e
                                            eqr (FormulaEq _ e) = e
                                            ls = map eql fos
                                            rs = map eqr fos
                                        return (FormulaEq (foldl App (Ctor ctor) ls) (foldl App (Ctor ctor) rs))
                                  Nothing  -> throwError "invalid constructor in bisimulation proof"

checkTotality :: (Monad m) => [ProofAlt] -> MPT m Bool
checkTotality _ = return True -- FIXME FIXME

checkAlt :: (Monad m) => Expr -> Formula -> Name Proof -> ProofAlt -> MPT m ()
checkAlt e fo x a = do
                      (pa,pr)  <- unbind a
                      let mepa =  patToExpr pa
                      case mepa of
                        Nothing  -> throwError "proof pattern contains wildcard"
                        Just epa -> do
                                      -- FIXME: shadow the pat vars
                                      fopr <- localPE (bindEnv x (FormulaEq e epa)) (check pr)
                                      if fo `aeq` fopr
                                         then return ()
                                         else do
                                           fo_   <- ppr fo
                                           fopr_ <- ppr fopr
                                           throwError $ show $ text "case alternative does not prove the needed formula." $+$ text "" $+$ text "Expected:" $+$ nest 4 fo_ $+$ text "" $+$ text "Got:" $+$ nest 4 fopr_

mkInitialEE :: [Decl] -> ExprEnv
mkInitialEE []     = Map.empty
mkInitialEE (d:ds) = case d of
                       (ExprDecl mety x em) -> bindEnv x (fmap unembed mety,Just (unembed em)) (mkInitialEE ds)
                       _                    -> mkInitialEE ds

mkInitialPE :: [Decl] -> ProofEnv
mkInitialPE []     = Map.empty
mkInitialPE (d:ds) = case d of
                       (ProofDecl efo x _)  -> bindEnv x (unembed efo) (mkInitialPE ds)
                       _                    -> mkInitialPE ds

mkInitialDTE :: [Decl] -> DTEnv
mkInitialDTE []     = Map.empty
mkInitialDTE (d:ds) = case d of
                        (DataDecl n b) -> bindEnv n (unembed b) (mkInitialDTE ds)
                        _              -> mkInitialDTE ds

mkInitialCE :: (Monad m) => [Decl] -> MPT m CEnv
mkInitialCE []     = return Map.empty
mkInitialCE (d:ds) = case d of
                       (DataDecl n b) -> do
                         (xs,cds) <- unbind (unembed b)
                         ce       <- mkInitialCE ds
                         let ce'  =  foldr (\ cd@(ConstrDecl ctor tys) env -> bindEnv ctor (bind xs cd,n) env) ce cds
                         return ce'
                       _              -> mkInitialCE ds

checkDecl :: (Monad m) => Decl -> MPT m ()
checkDecl d@(ProofDecl efo x ep) = do
                                     fo <- check (unembed ep) 
                                     if unembed efo `aeq` fo
                                        then return ()
                                        else throwError $ "proof classifier and decl classifier don't match"
checkDecl _                      = return ()

checkProg :: (Monad m) => Program -> MPT m ()
checkProg (Program _ tr) = do
  ds <- untrec tr
  ce <- mkInitialCE ds
  localCE (const ce) $
   localDTE (const $ mkInitialDTE ds) $
   localPE (const $ mkInitialPE ds) $
   localEE (const $ mkInitialEE ds) $
    mapM_ checkDecl ds