module Control.Monad.Unify.Trans where

import Prelude (($), (<<<), const, flip, unit)

import Control.Plus (class Plus)
import Control.Monad (class Monad)
import Control.Monad.Rec.Class (class MonadRec)
import Control.Monad.State.Trans (StateT, runStateT)
import Control.Monad.State.Class (gets, modify)
import Control.Monad.Unify.Class (class MonadUnify, class Unifiable, class UnificationVariable, Stream, pop, squash, variable, variables, reify, unify)
import Control.Apply (class Apply)
import Control.Alt (class Alt)
import Control.Alternative (class Alternative)
import Control.Applicative (class Applicative, pure)
import Control.Bind (class Bind, bind, (>>=))
import Control.MonadPlus (class MonadPlus)
import Data.Functor (class Functor, (<$>))
import Data.Foldable (foldMap, foldr)
import Data.Map (Map, lookup, insert)
import Data.Map as Map
import Data.Set (Set, singleton, union)
import Data.Tuple (Tuple(..))
import Data.Tuple.Nested ((/\))
import Data.Newtype (class Newtype, un)
import Data.Maybe (Maybe(..), fromMaybe)
import Data.Maybe.First (First(..))

type UnificationState k f =
    { boundVariables :: Map k { equivalence :: Set k, boundTo :: Maybe (f k) }
    , currentVariables :: Stream k
    }

newtype UnifyT k f m a = MkUnifyT (StateT (UnificationState k f) m a)

derive instance Newtype (UnifyT k f m a) _
derive instance Functor m => Functor (UnifyT k f m)
derive newtype instance Monad m => Apply (UnifyT k f m)
derive newtype instance Monad m => Applicative (UnifyT k f m)
derive newtype instance Monad m => Bind (UnifyT k f m)
instance Monad m => Monad (UnifyT k f m)
derive newtype instance (Monad m, Alt m) => Alt (UnifyT k f m)
derive newtype instance (Monad m, Plus m) => Plus (UnifyT k f m)
instance (Monad m, Alternative m) => Alternative (UnifyT k f m)
instance (MonadPlus m) => MonadPlus (UnifyT k f m)
derive newtype instance MonadRec m => MonadRec (UnifyT k f m)

instance (Unifiable k f, MonadPlus m) => MonadUnify k f (UnifyT k f m) where
    fresh = MkUnifyT $ do
        Tuple k restVariables <- gets (pop <<< _.currentVariables)
        _ <- modify $ _ { currentVariables = restVariables }
        pure k
    merge k1 k2 = do
        boundVariables <- MkUnifyT $ gets _.boundVariables
        let
            equivalence k = fromMaybe (singleton k) (_.equivalence <$> lookup k boundVariables)
            boundTo k = lookup k boundVariables >>= _.boundTo
            fullSet = equivalence k1 `union` equivalence k2
        newTerm <-
            case boundTo k1 /\ boundTo k2 of
                Just t1 /\ Just t2 -> unify t1 t2 >>= const (Just <$> reify t1)
                Just t1 /\ Nothing -> pure $ Just t1
                Nothing /\ Just t2 -> pure $ Just t2
                Nothing /\ Nothing -> pure $ Nothing
        let newMapValue = {equivalence: fullSet, boundTo: newTerm}
        _ <- MkUnifyT $ modify $ _ { boundVariables = foldr (flip insert newMapValue) boundVariables fullSet }
        pure unit
    store k f = do
        boundVariables <- MkUnifyT $ gets _.boundVariables
        let fullSet = fromMaybe (singleton k) (_.equivalence <$> lookup k boundVariables)
        let firstBound = un First $ foldMap (First <<< (_>>=(_.boundTo)) <<< (flip lookup boundVariables)) fullSet
        case firstBound of
            Just f' -> unify f f'
            Nothing -> do
                let newMapValue = {equivalence: fullSet, boundTo: Just f}
                _ <- MkUnifyT $ modify $ _ { boundVariables = foldr (flip insert newMapValue) boundVariables fullSet }
                pure unit
    reify f =
        MkUnifyT $ do
            boundVariables <- gets _.boundVariables
            let
                reify' f' = squash $ process <$> f'
                process k = fromMaybe (variable k) (reify' <$> (lookup k boundVariables >>= _.boundTo))
            pure $ reify' f

runUnifyT :: forall k f m a. Monad m => UnificationVariable k => UnifyT k f m a -> m (Tuple a (UnificationState k f))
runUnifyT m = runStateT (un MkUnifyT m) { boundVariables: Map.empty, currentVariables: variables }