-----------------------------------------------------------------------------
-- |
-- Module      :  Generics.Putlenses.Putlens
-- Copyright   :  (C) 2013 Hugo Pacheco
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Hugo Pacheco <hpacheco@nii.ac.jp>
-- Stability   :  provisional
--
-- General framework for put-based lenses.
--
-- 
--
--
----------------------------------------------------------------------------

module Generics.Putlenses.Putlens where

import Control.Monad.Reader
import Control.Monad.Writer.Lazy
import Control.Monad.State (State,MonadState,StateT)
import qualified Control.Monad.State as State
import Data.Maybe
import qualified Control.Lens as L
import Control.Monad.Identity
import Control.Monad.Trans
import Data.Monoid
import Control.Monad.Trans.Maybe

newtype GetPut = GetPut Bool -- boolean flag to test the GetPut law
newtype PutGet = PutGet Bool -- boolean flag to test the PutGet law

testGetPut (GetPut gp) = gp
testPutGet (PutGet pg) = pg

instance Monoid PutGet where
	mempty = PutGet False								-- the flag is initialized as false
	mappend (PutGet x) (PutGet y) = PutGet (x || y)		-- and may be turned on by any putlens

-- | Interface for normal lenses
data LensM m s v = LensM { get :: s -> v, put :: s -> v -> m s, create :: v -> m s }

-- | Interface for lenses without monadic effects
type Lens s v = LensM Identity s v

-- | Non-monadic @put@ function
simpleput :: Lens s v -> (s -> v -> s)
simpleput l s = runIdentity . put l s

-- | Non-monadic @create@ function
simplecreate :: Lens s v -> (v -> s)
simplecreate l = runIdentity . create l

-- | Monad for put-based lenses
-- includes an environment, state, and boolean tags that our system will use to ensure GetPut and PutGet
type PutM m a = ReaderT GetPut (WriterT PutGet m) a

type Get s v = s -> Maybe v
type Put m s v = Maybe s -> v -> PutM m s
type Create m s v = v -> PutM m s

-- | Framework for put-based lenses
data PutlensM m s v = PutlensM {
    getputM :: s -> (Maybe v,Create m s v)
  , createM :: Create m s v }
-- tupled framework for efficiency

type PutlensMaybeM m s v = PutlensM (MaybeT m) s v

type PutlensReaderM m e s v = PutlensM (ReaderT e m) s v

type PutlensStateM m st s v = PutlensM (StateT st m) s v

type Putlens s v = PutlensM Identity s v

-- | Changes the resulting monad of a put computation

mapPutM' :: (Monad m,Monad n) => (m a -> n b) -> (forall c. m c -> n c) -> PutM m a -> PutM n b
mapPutM' f g = mapReaderT (mapWriterT mapT) where
	mapT m = do
		b <- f (liftM fst m)
		w <- g (liftM snd m)
		return (b,w)

mapPutM :: (forall a. m a -> n a) -> PutM m a -> PutM n a
mapPutM f = mapReaderT (mapWriterT f)
	
-- | Changes the resulting monad of a putlens
mapPutlensM :: Monad m => (forall a. m a -> n a) -> PutlensM m s v -> PutlensM n s v
mapPutlensM f l = PutlensM getput' create'
    where getput' s = let (v,put') = getputM l s in (v,\v' -> mapPutM f (put' v'))
          create' v' = mapPutM f (createM l v')

-- | Lifts a function on a given monad to a put computation
liftPutM :: Monad m => m s -> PutM m s
liftPutM m = lift (lift m)

-- | Function @lift@ applied to the monadic argument of a putlens
liftPutlensM :: (MonadTrans t,Monad m) => PutlensM m s v -> PutlensM (t m) s v
liftPutlensM = mapPutlensM lift

-- | Forward |get| function
getM :: PutlensM m s v -> Get s v
getM l s = let (v,create) = getputM l s in v

-- | Domain of a putlens (the domains of its get function)
dom ::  PutlensM m s v -> (s -> Bool)
dom f = isJust . getM f

-- | Backward |put| function
putM :: PutlensM m s v -> Put m s v
putM l Nothing v' = createM l v'
putM l (Just s) v' = let (v,create) = getputM l s in create v'

-- | Runs a putlens for a particular environment
evalPutM :: Monad m => PutM m s -> GetPut -> m s
evalPutM putm e = liftM fst (runPutM putm e)

-- Runs a putlens for a particular environment and GetPut flag, returning also the TestPutGet flag
runPutM :: Monad m => PutM m s -> GetPut -> m (s,PutGet)
runPutM putm e = runWriterT (runReaderT putm e)

-- | Converts a normal lens to a putlens (without explicit failure)
lens2put :: Monad m => LensM m s v -> PutlensM m s v
lens2put (LensM get put create) = PutlensM getputM' createM'
    where getputM' s = (Just (get s),\v -> liftPutM $ put s v)
          createM' v = liftPutM $ create v

-- | Converts a simple lens of the Haskell lens package into a putlens
simplelens2put :: Monad m => L.Lens' s v -> PutlensM m s v
simplelens2put l = PutlensM getput' create'
    where get' s = L.view l s
          put' s v' =  return $ L.set l v' s
          getput' s = (return (get' s),put' s)
          create' v' = put' (error "simplelens2put: no original source") v'

-- | Converts a putlens to a normal lens, enforcing well-behavedness
put2lens :: (Eq v) => Putlens s v -> Lens s v
put2lens = put2lensM' True True

-- | Converts a putlens to a normal lens but without enforcing PutGet (for unsafe casts)
put2quicklens :: (Eq v) => Putlens s v -> Lens s v
put2quicklens = put2lensM' False True

-- | Converts a putlens to a normal lens, without enforcing GetPut
put2quicklensNoGetPut :: (Eq v) => Putlens s v -> Lens s v
put2quicklensNoGetPut = put2lensM' False False

-- | Converts a putlens to a normal lens, without enforcing GetPut
put2lensNoGetPut :: (Eq v) => Putlens s v -> Lens s v
put2lensNoGetPut = put2lensM' False False

-- | Converts a monadic putlens to a monadic lens, enforcing well-behavedness
put2lensM :: (Monad m,Eq v) => PutlensM m s v -> LensM m s v
put2lensM = put2lensM' True True

-- | Converts a monadic putlens to a monadic lens but without enforcing PutGet (for unsafe casts)
put2quicklensM :: (Monad m,Eq v) => PutlensM m s v -> LensM m s v
put2quicklensM = put2lensM' False True

-- | Converts a monadic putlens to a monadic lens, without enforcing GetPut
put2lensNoGetPutM :: (Monad m,Eq v) => PutlensM m s v -> LensM m s v
put2lensNoGetPutM = put2lensM' True False

-- | Converts a monadic putlens to a monadic lens, without enforcing GetPut
put2quicklensNoGetPutM :: (Monad m,Eq v) => PutlensM m s v -> LensM m s v
put2quicklensNoGetPutM = put2lensM' False False

-- | Converts a putlens to a normal lens with a parameter telling to ensure well-behavedness or not (for unsafe casts)
-- Initializes the environment as the original source, the state as empty, the GetPut tag as True and the PutGet tag as False
put2lensM' :: (Monad m,Eq v) => Bool -> Bool -> PutlensM m s v -> LensM m s v
put2lensM' checkPutGet checkGetPut l = LensM get put create
	where get = liftM fst $ getput' checkPutGet checkGetPut l
	      put = liftM snd $ getput' checkPutGet checkGetPut l
	      create = create' checkPutGet l

-- | The get function of a putlens
get' :: (Eq v) => PutlensM m s v -> (s -> v)
get' l s = case getM l s of
	Just v -> v
	Nothing -> error "get fails"

-- The put function of a putlens (with PutGet and GetPut flags)
put' :: (Monad m,Eq v) => Bool -> Bool -> PutlensM m s v -> (s -> v -> m s)
put' checkPutGet checkGetPut l = put (put2lensM' checkPutGet checkGetPut l)

-- | The tupled get/put function of a putlens (with PutGet and GetPut flags)
getput' :: (Monad m,Eq v) => Bool -> Bool -> PutlensM m s v -> (s -> (v,v -> m s))
getput' checkPutGet checkGetPut l s =
    let (mbv,put) = getputM l s
        put' v' = do
            (s',pg) <- runPutM (put v') (GetPut checkGetPut)
            if checkPutGet && testPutGet pg && getM l s' /= Just v'
                then fail "put2lens (unsafe casts violate PutGet)"
                else return s'
        v = case mbv of { Just x -> x ; otherwise -> error "get function fails"}
    in (v,put')

-- | The create function of a putlens (with PutGet flag)
create' :: (Monad m,Eq v) => Bool -> PutlensM m s v -> (v -> m s)
create' checkPutGet l v' = do
	(s',pg) <- runPutM (createM l v') (GetPut False)
	if checkPutGet && testPutGet pg && getM l s' /= Just v'
		then fail "put2lens (unsafe casts violate PutGet)"
		else return s'

-- | Converts a putlens into another putlens that only uses the create function
put2create :: (Monad m,Eq v) => PutlensM m s v -> PutlensM m s v
put2create l = checkGetPut $ l { getputM = getput }
    where getput s = let (mbv,put) = getputM l s
                     in (mbv,createM l)

-- internal
checkGetPut :: (Monad m,Eq v) => PutlensM m s v -> PutlensM m s v
checkGetPut l = l { getputM = getput' } where
	getput' s = let (v,put) = getputM l s
	                put' v' = ask >>= \gp -> if testGetPut gp && v == Just v' then return s else put v'
	            in (v,put')
          
-- internal
checkPutGet :: Monad m => PutlensM m s v -> PutlensM m s v
checkPutGet l = PutlensM getput' (create' (createM l))
    where getput' s = let (v,put) = getputM l s in (v,create' put)
          create' put v' = lift (tell $ PutGet True) >> put v'

-- internal
offGetPut :: Monad m => PutM m s -> PutM m s
offGetPut m = withReaderT (\gp -> GetPut False) m

-- internal
onPutGet :: Monad m => PutM m s -> PutM m s
onPutGet m = lift (tell $ PutGet True) >> m