{-# OPTIONS_GHC -XMultiParamTypeClasses -XFlexibleContexts -XUndecidableInstances #-}
module Fix where

import Control.Applicative
import Control.Monad
import Data.Monoid
import qualified Data.Foldable as F
import qualified Data.Traversable as F


-- The basic 'Fix' type, It creates a simple recursive type.

newtype Fix f = F (f (Fix f))

instance Show (f (Fix f)) => Show (Fix f) where
    showsPrec n (F x) = showsPrec n x

instance Eq (f (Fix f)) => Eq (Fix f) where
    F x == F y = x == y

instance Ord (f (Fix f)) => Ord (Fix f) where
    F x `compare` F y = x `compare` y

foldFix :: Functor f => (f w -> w) -> Fix f -> w
foldFix f (F ji) =  f (fmap (foldFix f) ji)

foldFixM :: (Monad m,F.Traversable f) => (f w -> m w) -> Fix f -> m w
foldFixM f (F ji) =  f =<< (F.mapM (foldFixM f) ji)

foldFixM' :: (Monad m,F.Traversable f) => (Fix f -> m (Fix f)) -> (f w -> m w) -> Fix f -> m w
foldFixM' fd f x = do
    F x <- fd x
    f =<< (F.mapM (foldFixM' fd f) x)

-- A recursive type that attaches some memoized data to each subterm

data FixM f a = FM a (f (FixM f a))

instance Functor f => Functor (FixM f) where
    fmap f (FM x y) = FM (f x) (fmap (fmap f) y)

instance F.Foldable f => F.Foldable (FixM f) where
    foldMap f (FM x y) = f x `mappend` F.foldMap (F.foldMap f) y

instance (Functor f, F.Traversable f) => F.Traversable (FixM f) where
    traverse f (FM x y) = FM <$> f x <*> (F.traverse (F.traverse f) y)

--instance Eq (f (FixM f a)) => Eq (FixM f a) where
--    FM _ x == FM _ y = x == y

--instance Ord (f (FixM f a)) => Ord (FixM f a) where
--    FM _ x `compare` FM _ y = x `compare` y

fixMemo :: FixM f a -> a
fixMemo (FM a _) = a

fromFixMemo :: FixM f a -> (a,f (FixM f a))
fromFixMemo (FM a x) = (a,x)

toFixMemo ::  (f (FixM f a) -> a) -> f (FixM f a) -> FixM f a
toFixMemo f x = FM (f x) x


fixDeMemoize :: Functor f => FixM f a -> Fix f
fixDeMemoize ja = f ja where
    f (FM _ j) = F (fmap f j)


fixMemoize :: Functor f => (f (FixM f a) -> a) -> Fix f -> FixM f a
fixMemoize f (F ji) =  foldFix f' (F ji) where
    f' x = FM (f x) x

-- relys on laziness
fixMemoizeKnot :: Functor f => (c -> f (FixM f a) -> (c,a)) -> c -> Fix f -> FixM f a
fixMemoizeKnot f c fji = g c fji where
    g c (F ji) = FM a nji where
        (c',a) = f c nji
        nji = fmap (g c') ji

fixMemoizeM :: (Monad m,F.Traversable f) => (f (FixM f a) -> m a) -> Fix f -> m (FixM f a)
fixMemoizeM f (F ji) =  foldFixM f' (F ji) where
    f' x = do fx <- f x; return $ FM fx x

-- like fixMemoize, but lets you examine nodes on the way down as well as annotate them on the way up
fixMemoizeM' :: (Monad m,F.Traversable f) => (Fix f -> m (Fix f)) -> (f (FixM f a) -> m a) -> Fix f -> m (FixM f a)
fixMemoizeM' fd f x =  foldFixM' fd f' x where
    f' x = do fx <- f x; return $ FM fx x

-- hash cons

--hashCons :: Ord (f Int) => Fix f -> [f Int]
--hashCons x = runState foldFix f x where

class SelfFunctor a where
    sfmap :: (a -> a) -> a -> a
    sfmapM :: Monad m => (a -> m a) -> a -> m a

class HasMemo a where
    type Memo a
    memo :: a -> Memo a

class HasContents a where
    type Contents a
    open :: a -> Contents a

instance HasContents (Fix t) where
    type Contents (Fix t) = t (Fix t)
    open (F x) = x

instance HasContents (FixM t a) where
    type Contents (FixM  t a) = t (FixM t a)
    open (FM _  x) = x

instance HasMemo (FixM t a) where
    type Memo (FixM t a) = a
    memo (FM a _) = a