{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE AllowAmbiguousTypes #-}

module MicroKanren where

-- based on MikroKanren.hs (https://github.com/rntz/ukanren.git)
-- modified to accomodate general data structures.

import Control.Applicative
import Control.Monad
import Control.Monad.State.Strict hiding (State)

import Data.Dynamic
import Data.Maybe (fromMaybe, mapMaybe)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM

newtype Var a = Var Int deriving (Eq, Show)
unVar (Var a) = a

type Subst = IntMap Dynamic
type State = (Int, Subst)

--  A monad that can generate, bind, and look-up variables.
class Monad m => VarGen m where
    newVar :: m (Var a)
    assign :: Typeable a => Var a -> Either (Var a) a -> m ()
    deref  :: Typeable a => Var a -> m (Either (Var a) a)

fresh :: VarGen m => m (Var a)
fresh = newVar

assignVar :: (VarGen m, Typeable a) => Var a -> Var a -> m ()
assignVar v w = assign v (Left w)

assignValue :: (VarGen m, Typeable a) => Var a -> a -> m ()
assignValue v w = assign v (Right w)

lv :: (VarGen m, Typeable a) => a -> m (Var a)
-- lv :: (VarGen m) => Int -> m (Var Int)
lv a = do v <- fresh
          assignValue v a
          return v

-- Fair interleaving of finitely many lists.
interleaves :: [[a]] -> [a]
interleaves [] = []
interleaves l = [x | x:_ <- l] ++ interleaves [xs | _:xs <- l]

interleave a b = interleaves [a,b]


-- Search trees, so we can define the search algorithm separately.
data Tree a = Empty | Single a | Node (Tree a) (Tree a) deriving Show

instance Functor Tree where fmap = liftM
instance Applicative Tree where pure = return; (<*>) = ap
instance Monad Tree where
    return = Single
    Empty >>= _ = Empty
    Single x >>= f = f x
    Node l r >>= f = Node (l >>= f) (r >>= f)

-- NB. not a law-abiding Alternative/MonadPlus instance: not associative.
instance MonadPlus Tree where mzero = empty; mplus = (<|>)
instance Alternative Tree where empty = Empty; (<|>) = Node


-- Search strategies over Trees.
listify :: ([a] -> [a] -> [a]) -> Tree a -> [a]
listify f Empty = []
listify f (Single a) = [a]
listify f (Node l r) = f (listify f l) (listify f r)

dfs = listify (++)

-- Not sure if there's a standard name for this search strategy.
ifs = listify interleave

-- Unfortunately we can't write iterated deepening as a function on Trees,
-- because of Haskell's memoizing call-by-need evaluation. So we'll just do a
-- plain old memory-hogging BFS.
bfs t = search [] [t]
    -- we use the usual trick of encoding a queue as two stacks.
    where search [] [] = []
          search a [] = search [] (reverse a)
          search a (Empty : b) = search a b
          search a (Single x : b) = x : search a b
          search a (Node l r : b) = search (r:l:a) b


-- Implementation of the Kanren interface
type K = StateT State Tree
type Goal = K ()

instance Monad m => VarGen (StateT State m) where
   newVar = state (\(v,s) -> (Var v, (v + 1, s)))
   assign v t = modify (fmap (IM.insert (unVar v) (toDyn t)))
   deref v = do either <-  fmap (join . fmap fromDynamic)
                                (gets (\(_,sub) -> IM.lookup (unVar v) sub))
                case either of
                  Just (Right a) -> return (Right a)
                  Just (Left v') -> deref v'
                  Nothing        -> return (Left v)

start :: State
start = (0, IM.empty)

runK :: K a -> State -> [(a, State)]
runK k st = bfs $ runStateT k st
-- NB. if we change bfs to ifs, then fivesR fails!
-- with dfs you get prolog-like behavior, and even more things fail.

evalK :: K a -> State -> [a]
evalK k st = map fst (runK k st)

execK :: K a -> State -> [State]
execK k st = map snd (runK k st)


-- Basic operators.
ok :: Goal
ok = pure ()

class Typeable a => Unifiable a where
  unify :: a -> a -> Goal

(===) :: Unifiable a => a -> a -> Goal
(===) = unify

unifyVar :: Unifiable a => Var a -> Var a -> Goal
unifyVar x y  = do x' <- deref x
                   y' <- deref y
                   unifyVarAux x' y'
    where unifyVarAux (Left x) (Left y) | x == y = ok
          unifyVarAux (Left x) t                 = assign x t
          unifyVarAux t (Left y)                 = assign y t
          unifyVarAux (Right x) (Right y)        = unify x y

instance Unifiable a => Unifiable (Var a) where
  unify =  unifyVar

data LVList a = LVNil | LVCons (Var a) (Var (LVList a)) deriving (Show,Typeable)

instance Unifiable a => Unifiable (LVList a) where
  unify LVNil LVNil                 = ok
  unify (LVCons x xs) (LVCons y ys) = do { unifyVar x y; unifyVar xs ys }
  unify _ _                       = mzero


-- cons :: Typeable a => Var a -> Var (LVList a) -> K (Var (LVList a))
-- To avoid ambiguity, the element type is fixed to ``Int''.
cons :: Var Int -> Var (LVList Int) -> K (Var (LVList Int))
cons x xs = do v <- fresh
               assignValue v (LVCons x xs)
               return v

-- nil :: Typeable a => K (Var (LVList a))
nil :: K (Var (LVList Int))
nil       = do v <- fresh
               assignValue v LVNil
               return v

-- toUList :: Typeable a => [a] -> K (Var (LVList a))
toUList :: [Int] -> K (Var (LVList Int))
toUList [] = nil
toUList (x:xs) = do ys <- toUList xs
                    y  <- fresh
                    assignValue y x
                    cons y ys

-- fromUList :: Typeable a => a -> Var (LVList a) -> K [a]
fromUList :: Int -> Var (LVList Int) -> K [Int]
fromUList dflt v = do x <- deref v
                      case x of
                        Right LVNil -> return []
                        Right (LVCons y ys) ->
                          do z <- deref y
                             let z' = case z of { Right z1 -> z1 ; _ -> dflt }
                             zs <- fromUList dflt ys
                             return (z':zs)
                        Left _     -> return []    -- default list is []


toLVList :: [Int] -> K (Var (LVList Int))
toLVList = toUList

fromLVList :: Int -> Var (LVList Int) -> K [Int]
fromLVList = fromUList

instance Unifiable Int where
  unify x y | x == y = ok
  unify _ _          = mzero

instance Unifiable Double where
  unify x y | x == y = ok
  unify _ _          = mzero

data MyTree a = MyEmpty | MyBranch (MyTree a) a (MyTree a)
              deriving (Show, Typeable)

data LVTree a = LVEmpty | LVBranch (Var (LVTree a)) (Var a) (Var (LVTree a))
             deriving (Show, Typeable)

instance Unifiable a => Unifiable (LVTree a) where
  unify LVEmpty LVEmpty             = ok
  unify (LVBranch xsL x xsR) (LVBranch ysL y ysR) = do {
         unifyVar xsL ysL; unifyVar x y; unifyVar xsR ysR
       }
  unify _ _                       = mzero

branch :: Typeable a => Var (LVTree a) -> Var a -> Var (LVTree a)
                         -> K (Var (LVTree a))
branch xsL x xsR = do v <- fresh
                      assignValue v (LVBranch xsL x xsR)
                      return v

empT :: Typeable a => K (Var (LVTree a))
empT  = do v <- fresh
           assignValue v LVEmpty
           return v

toUTree :: Typeable a => MyTree a -> K (Var (LVTree a))
toUTree MyEmpty = empT
toUTree (MyBranch xsL x xsR)
                = do ysL <- toUTree xsL
                     y  <- fresh
                     ysR <- toUTree xsR
                     assignValue y x
                     branch ysL y ysR

fromUTree :: Typeable a => a -> Var (LVTree a) -> K (MyTree a)
fromUTree dflt v = do x <- deref v
                      case x of
                        Right LVEmpty -> return MyEmpty
                        Right (LVBranch ysL y ysR) ->
                          do z <- deref y
                             let z' = case z of { Right z1 -> z1 ; _ -> dflt }
                             zsL <- fromUTree dflt ysL
                             zsR <- fromUTree dflt ysR
                             return (MyBranch zsL z' zsR)
                        Left _     -> return MyEmpty    -- default tree

toLVTree :: Typeable a => MyTree a -> K (Var (LVTree a))
toLVTree = toUTree
fromLVTree :: Typeable a => a -> Var (LVTree a) -> K (MyTree a)
fromLVTree = fromUTree

pair x y     = (x,y)
triple x y z = (x,y,z)

{-
expand :: Typeable a => Var a -> K (Either (Var a) a)
expand v = deref v
-}

disj, conj :: Goal -> Goal -> Goal
disj = (<|>)
conj = (>>)

-- equivalent of disj+, conj+
disjs, conjs :: [Goal] -> Goal
disjs = msum
conjs = sequence_

{-
myAppend a b ab =
  do h <- fresh
     t <- fresh
     a' <- cons h t
     res <- fresh
     ab' <- cons h res
     unify a a'
     unify ab ab'
     myAppend t b res
  `mplus`
  do a' <- nil
     unify a a'
     unify b ab

a, b :: [Int]
a = [1,2,3]
b = [1,2,3,5]

-- try: runK myAppendTest1 start
myAppendTest1 =
   do a <- toLVList a
      b <- toLVList [8,9]
      ab <- fresh
      myAppend a b ab
      fromLVList 999 ab


-- try: runK myAppendTest2 start
myAppendTest2 =
   do xs <- fresh
      ys <- fresh
      zs <- toLVList a
      myAppend xs ys zs
      xs' <- fromLVList 99 xs
      ys' <- fromLVList 99 ys
      return (xs', ys')
-}