{-# OPTIONS_GHC -fglasgow-exts #-}

import Text.PrettyPrint.HughesPJ
import Control.Monad.Reader
import Control.Monad
import Data.List
import Data.Char
import Data.Monoid
import Data.IORef
import System.IO

hasUnicode = True

data Term
    = Bound !Int
    | Constant String
    | TVar !Int
    | TUVar (UVar Term)
    deriving(Eq)

data Formula
    = Forall Formula
    | Formula :-> Formula
    | Not Formula
    | Pred String [Term]
    | FVar !Int
    | FUVar (UVar Formula)
    deriving(Eq)

data HState = HState { stateFormula :: [Formula], stateAxioms :: [Formula] } deriving(Eq)
stateEmpty = HState { stateAxioms = reverse (axioms ++ hypo), stateFormula = []  }

main = do
    hSetBuffering stdin NoBuffering
    loop (repeat stateEmpty) ""

perform 'h' ss = do
    putStr clearScreen
    putStr help
    putStrLn "-----"
    getChar
    loop ss ""
perform 'u' (_:ss) = loop ss ""
perform c (s:ss) | isUpper c = case lookup c (zip ['A' ..] (reverse $ stateAxioms s)) of
    Just a -> loop (s { stateFormula = a:stateFormula s }:s:ss) ""
    Nothing -> loop (s:ss) "Unknown Formula"
perform c (s:ss) | isDigit c && c > '0' = do
    let idx = ord c - ord '0';
        f 0 (x:xs) rs = loop (s { stateFormula = x:(reverse rs ++ xs) }:s:ss) ""
        f n (x:xs) rs = f (n - 1) xs (x:rs)
        f n [] _ = loop (s:ss) "Not in stack"
    f idx (stateFormula s) []
perform c (s:ss) = f c s where
--    f 's' s@HState { stateFormula = (x:y:rs) } = next s { stateFormula = (y:x:rs) }
    f '0' s@HState { stateFormula = (x:rs) } = next s { stateFormula = (x:x:rs) }
    f '-' s@HState { stateFormula = (x:rs) } = next s { stateFormula = rs }
    f 'd' s@HState { stateFormula = (x:rs) } = degen s rs x
    f 'g' s@HState { stateFormula = (x:rs) } = gen s rs x
    f 'p' s@HState { stateFormula = (x:rs) } | x `notElem` stateAxioms s = next s { stateFormula = rs, stateAxioms = x:stateAxioms s }
                                             | otherwise = err "Formula Already Exists."
    f 'm' s@HState { stateFormula = (x:y:rs) } = modus_p s rs x y
    f '!' _ = return ()
    f '\x04' _ = return ()
    f c s = err $ "Unknown Command " ++ show c
    next x = if x /= s then loop (x:s:ss) "" else loop (s:ss) ""
    err x = loop (s:ss) x
    modus_p s rs (p :-> q) p' | p' == p = next s { stateFormula = q:rs }
    modus_p s rs (p :-> q) p' = do
        mp <- modus_pones p q p'
        case mp of
            Just q' ->  next s { stateFormula = q':rs }
            Nothing -> err "Modus Pones didn't unify"
    modus_p _ _ _ _ = err "Could not perform modus_pones"
    degen s rs (Forall x) = next  s { stateFormula = rmap 1 (TVar nt) x:rs }  where
        (_,ts) = freeVars x
        (nt:_) = [ c | c <- [ 0 .. ], c `notElem` ts]
    degen _ _ _ = err "Could not degeneralize"
    gen s rs x | not (null ts) = next  s { stateFormula = foldr (\_ -> Forall) (bmap (zip ts [1 ..]) x) ts:rs }  where
        ts = nub $ snd $ freeVars x
    gen _ _ _ = err "nothing to generalize"
    bmap m x = f m x where
        f m (Forall x) = Forall (f [ (x,y + 1) | (x,y) <- m ] x)
        f m (Pred s ts) = Pred s (map g ts) where
            g (TVar z) | Just x <- lookup z m  = Bound x
            g x = x
        f m x = formMap (f m) x
    rmap n w x = f n x where
        f n (Pred s ts) = Pred s (map g ts) where
            g (Bound z) | z == n = w
            g x = x
        f n (Forall x) = Forall $ f (n + 1) x
        f n x = formMap (f n) x

clearScreen = "\x1b[H\x1b[2J"

loop ss@(s:_) msg = do
    putStr clearScreen
    putStrLn "press 'h' for help\n----"
    let pa c a = [c] ++ ". " ++ show a
    mapM_ putStrLn (zipWith pa ['a' .. ] (reverse $ stateAxioms s))
    putStr "----\n"
    mapM_ putStrLn (reverse $ zipWith pa ['0' ..] (stateFormula s))
    putStr $ "----\n" ++ msg
    hFlush stdout
    gl <- getChar
    putStrLn []
    perform gl ss


freeVars x = f x where
    f (FVar v) = ([v],[])
    f (Pred _ ts) = ([],[ v | TVar v <- ts])
    f x = formCollect f x

modus_pones :: Formula -> Formula -> Formula -> IO (Maybe Formula)
modus_pones p q p' = ans where
    ans = do
        (p :-> q) <- uvarIze (p :-> q)
        p' <- uvarIze p'
        let fn = do
            unify p p'
            deUVar q
        catch (fmap Just fn) (const (return Nothing))

    unify :: Formula -> Formula -> IO ()
    unify x y = join $ return f `ap` (findVal x) `ap` (findVal y) where
        f :: Formula -> Formula -> IO ()
        f (FVar _) _ = error "fvarfound"
        f _ (FVar _) = error "fvarfound"
        f (FUVar x) y = varBind x y
        f x (FUVar y) = varBind y x
        f x y = match_ unify g' x y
        g,g' :: Term -> Term -> IO ()
        g' x y = join $ return g `ap` findVal x `ap` findVal y
        g (TUVar x) y = varBind x y
        g x (TUVar y) = varBind y x
        g x y = if x == y then return () else fail "terms don't match"

    varBind uv t | Just uv == getUVar t = return ()
    varBind uv t | uv `elem` collectUVar t = fail "occurs check"
    varBind (UVar v) t = writeIORef v (Just t)

    deUVar x = do
        dm <- newIORef (map FVar [0..])
        tm <- newIORef (map TVar [0..])
        let f' v = findVal v >>= f
            f (FUVar v) = duvar dm v
            f (Pred s ts) = return (Pred s) `ap` mapM g' ts
            f x = formMapM f' x
            duvar dm uv = do
                (nn:nns) <- readIORef dm
                writeIORef dm nns
                varBind uv nn
                return nn
            g' v = findVal v >>= g
            g (TUVar v) = duvar tm v
            g x = return x
        f' x

    uvarIze x = do
        fs <- newUVarSource
        ts <- newUVarSource
        let f (FVar v) = do FUVar `fmap` fs v
            f (Pred s ts) = return (Pred s) `ap` mapM g ts
            f x = formMapM f x
            g (TVar v) = TUVar `fmap` ts v
            g x = return x
        f x

---------------
-- unification
---------------

newtype UVar t = UVar (IORef (Maybe t))
    deriving(Eq)

newUVar = fmap UVar $ newIORef Nothing

newUVarSource :: IO (Int -> IO (UVar a))
newUVarSource = do
    uv <- newIORef []
    let f n = readIORef uv >>= g n []
        g 0 _ (x:_) = return x
        g n rs (x:xs) = g (n - 1) (x:rs) xs
        g n rs [] = do
            u <- newUVar
            if n == 0 then do
                writeIORef uv (reverse $ u:rs)
                return u
             else g (n - 1) (u:rs) []
    return f

class GetUVar a where
    getUVar :: a -> Maybe (UVar a)
    collectUVar :: a -> [UVar a]

instance GetUVar Term where
    getUVar (TUVar v) = Just v
    getUVar _ = Nothing
    collectUVar (TUVar v) = [v]
    collectUVar _ = []

instance GetUVar Formula where
    getUVar (FUVar v) = Just v
    getUVar _ = Nothing
    collectUVar (FUVar v) = [v]
    collectUVar f = formCollect collectUVar f

findVal :: GetUVar a => a -> IO a
findVal a = case getUVar a of
    Nothing -> return a
    Just (UVar ref) -> do
        v <- readIORef ref
        case v of
            Nothing -> return a
            Just t -> do
                t' <- findVal t
                writeIORef ref (Just t')
                return t'

match_ fn pm f1 f2 = f f1 f2 where
    f (Forall x) (Forall y) =  fn x y
    f (x :-> y) (a :-> b) = fn x a >> fn y b
    f (Not f1)  (Not f2)  = fn f1 f2
    f (Pred s1 ts1) (Pred s2 ts2) | s1 == s2 = sequence_ (zipWith pm ts1 ts2)
    f _ _ = fail "match: formulas don't match"

formCollect :: Monoid a => (Formula -> a) -> Formula -> a
formCollect fn x = f x where
    f (Forall a) = fn a
    f (a :-> b) = fn a `mappend` fn b
    f (Not a)   = fn a
    f _ = mempty

formMapM :: Monad m => (Formula -> m Formula) -> Formula -> m Formula
formMapM f (Forall a) = liftM Forall (f a)
formMapM f (a :-> b) = liftM2 (:->) (f a) (f b)
formMapM f (Not a)   = liftM Not (f a)
formMapM _ x = return x

formMap :: (Formula -> Formula) -> Formula -> Formula
formMap f (Forall a) = Forall (f a)
formMap f (a :-> b) = (f a :-> f b)
formMap f (Not a)   = Not $ f a
formMap _ x = x

peq a b = Pred "EQ" [a,b]
b1 = Bound 1
b2 = Bound 2
b3 = Bound 3

fA = FVar 0
fB = FVar 1
fC = FVar 2

axioms = [
    fA :-> (fB :-> fA),
    (fA :-> (fB :-> fC)) :-> ((fA :-> fB) :-> (fA :-> fC)),
    (Not fA :-> Not fB) :-> (fB :-> fA),
    Forall (fA :-> fB) :-> (Forall fA :-> Forall fB),
    fA :-> Forall fA
    ]

hypo =
    [Forall $ peq b1 b1
    ,Forall $ Forall $ Forall $ peq b1 b2 :-> (peq b2 b3 :-> peq b1 b3)
    ,Forall $ Forall $ peq b1 b2 :-> peq b2 b1
    ,Forall $ Forall $ Not (peq b1 b2) :-> Not (peq b2 b1)
    ]


-- axiom schema
--
-- 1) A -> (B -> A)
-- 2) (A -> (B -> C)) -> ((A -> B) -> (A -> C))
-- 3) (!A -> !B) -> (B -> A)
--
-- 4) ∀xA -> A[x := t]
-- 5) ∀x(A -> B) -> (∀xA -> ∀xB)
-- 6) A -> ∀xA    (where x is not free in A)
--


help = unlines
    ["---- stack operations ----"
    ,"0 duplicate top of lemma stack"
    ,"1-9 move the specified formula to the top of the stack"
    ,"shift A-Z copy the specified formula from the theorem list to the top of the stack"
    ,"- delete top of lemma stack"
    ,"p promote the top of the lemma stack to a theorem"
    ,"---- rules of inference ----"
    ,"d (degeneralize) replace a quantifier with an unbound term"
    ,"g (generalize) universally quantify all unbound terms"
    ,"m use modus pones to apply the top of the stack to the second item in the stack"
    ,"---- utilities ----"
    ,"h show this help"
    ,"u undo last operation"
    ,"! quit"]

--------------------
-- pretty printing
--------------------

data TPEnv = TPEnv { tpChar :: !Char, tpParen :: !Bool }

instance Show Formula where
    show fl = show $ runReader tp TPEnv { tpChar = 'a', tpParen = False} where
        TP tp = formPretty fl


newtype TP a = TP (Reader TPEnv a)
    deriving(Monad,MonadReader TPEnv)

termPretty :: Term -> TP Doc
termPretty (TVar t) = return $ text $ 't':show t
termPretty (Bound i) = do
    ch <- asks tpChar
    return $ char (chr (ord ch - i))
termPretty (Constant s) = return $ text s

needParen = local (\env -> env { tpParen = True })


codes = (if hasUnicode then map fst else map snd)
    [("\xe2\x88\x80","@")
    ,("\xc2\xac","!")
    ,(" \xe2\x86\x92 "," -> ")
    ]

formPretty :: Formula -> TP Doc
formPretty h = f h where
    f (a :-> b) = do
        np <- asks tpParen
        a <- needParen $ f a
        b <- needParen $ f b
        let px = (a <> text (codes !! 2) <> b)
        return $ if np then paren px else px
    f (Not h)  = do
        h <- needParen $ f h
        return $ text (codes !! 1) <> h
    f (Pred s ts) = do
        ts <- mapM termPretty ts
        return $ text s <> paren (hcat $ punctuate (char ',') ts)
    f (Forall h) = do
        c <- asks tpChar
        local (\e -> e { tpChar = succ c }) $ do
            h <- needParen (f h)
            return $ text (codes !! 0) <> char c <> h
    f (FVar ub) = return $ text ('p':show ub)
    paren x = char '(' <> x <> char ')'