{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-warnings-deprecations #-}
module Optics.Zoom
  (
    -- * Zoom
    Zoom(..)
    -- * Magnify
  , Magnify(..)
  , MagnifyMany(..)
  ) where

import Control.Monad.Reader (ReaderT (..), MonadReader)
import Control.Monad.State (MonadState (..))
import Control.Monad.Trans.Error (Error, ErrorT (..))
import Control.Monad.Trans.Except (ExceptT (..), runExceptT)
import Control.Monad.Trans.Identity (IdentityT (..))
import Control.Monad.Trans.List (ListT (..))
import Control.Monad.Trans.Maybe (MaybeT (..))
import qualified Control.Monad.Trans.RWS.Lazy as L
import qualified Control.Monad.Trans.RWS.Strict as S
import qualified Control.Monad.Trans.State.Lazy as L
import qualified Control.Monad.Trans.State.Strict as S
import qualified Control.Monad.Trans.Writer.Lazy as L
import qualified Control.Monad.Trans.Writer.Strict as S

import Optics.Core
import Optics.Internal.Utils
import Optics.Extra.Internal.Zoom

-- $setup
-- >>> import Data.Monoid
-- >>> import Control.Monad.Reader (runReader, ask)
-- >>> import Optics.State
-- >>> import Optics.State.Operators
-- >>> import Optics.View

-- Chosen so that they have lower fixity than ('%=').
infixr 2 `zoom`, `zoomMaybe`, `zoomMany`
infixr 2 `magnify`, `magnifyMaybe`, `magnifyMany`

------------------------------------------------------------------------------
-- Zoom
------------------------------------------------------------------------------

-- | This class allows us to 'zoom' in, changing the 'State' supplied by many
-- different monad transformers, potentially quite deep in a monad transformer
-- stack.
--
-- Its functions can be used to run a monadic action in a larger 'State' than it
-- was defined in, using a 'Lens'', an 'AffineTraversal'' or a 'Traversal''.
--
-- This is commonly used to lift actions in a simpler 'State' 'Monad' into a
-- 'State' 'Monad' with a larger 'State' type.
--
-- When used with a 'Traversal'' over multiple values, the actions for each
-- target are executed sequentially and the results are aggregated.
--
-- This can be used to edit pretty much any 'Monad' transformer stack with a
-- 'State' in it!
--
-- >>> flip L.evalState ('a','b') $ zoom _1 $ use equality
-- 'a'
--
-- >>> flip S.execState ('a','b') $ zoom _1 $ equality .= 'c'
-- ('c','b')
--
-- >>> flip L.execState [(1,2),(3,4)] $ zoomMany traversed $ _2 %= (*10)
-- [(1,20),(3,40)]
--
-- >>> flip S.runState [('a',"b"),('c',"d")] $ zoomMany traversed $ _2 <%= (\x -> x <> x)
-- ("bbdd",[('a',"bb"),('c',"dd")])
--
-- >>> flip S.evalState ("a","b") $ zoomMany each (use equality)
-- "ab"
--
class
  (MonadState s m, MonadState t n
  ) => Zoom m n s t | m -> s, n -> t, m t -> n, n s -> m where
  zoom
    :: Is k A_Lens
    => Optic' k is t s
    -> m c
    -> n c

  zoomMaybe
    :: Is k An_AffineTraversal
    => Optic' k is t s
    -> m c
    -> n (Maybe c)

  zoomMany
    :: (Is k A_Traversal, Monoid c)
    => Optic' k is t s
    -> m c
    -> n c

instance Monad m => Zoom (S.StateT s m) (S.StateT t m) s t where
  zoom      o = \(S.StateT m) -> S.StateT $ stateZoom      o m
  zoomMaybe o = \(S.StateT m) -> S.StateT $ stateZoomMaybe o m
  zoomMany  o = \(S.StateT m) -> S.StateT $ stateZoomMany  o m
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance Monad m => Zoom (L.StateT s m) (L.StateT t m) s t where
  zoom      o = \(L.StateT m) -> L.StateT $ stateZoom      o m
  zoomMaybe o = \(L.StateT m) -> L.StateT $ stateZoomMaybe o m
  zoomMany  o = \(L.StateT m) -> L.StateT $ stateZoomMany  o m
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance Zoom m n s t => Zoom (ReaderT e m) (ReaderT e n) s t where
  zoom      o = \(ReaderT m) -> ReaderT (zoom      o . m)
  zoomMaybe o = \(ReaderT m) -> ReaderT (zoomMaybe o . m)
  zoomMany  o = \(ReaderT m) -> ReaderT (zoomMany  o . m)
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance Zoom m n s t => Zoom (IdentityT m) (IdentityT n) s t where
  zoom      o = \(IdentityT m) -> IdentityT (zoom      o m)
  zoomMaybe o = \(IdentityT m) -> IdentityT (zoomMaybe o m)
  zoomMany  o = \(IdentityT m) -> IdentityT (zoomMany  o m)
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance (Monoid w, Monad m) => Zoom (S.RWST r w s m) (S.RWST r w t m) s t where
  zoom      o = \(S.RWST m) -> S.RWST $ rwsZoom      o m
  zoomMaybe o = \(S.RWST m) -> S.RWST $ rwsZoomMaybe o m
  zoomMany  o = \(S.RWST m) -> S.RWST $ rwsZoomMany  o m
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance (Monoid w, Monad m) => Zoom (L.RWST r w s m) (L.RWST r w t m) s t where
  zoom      o = \(L.RWST m) -> L.RWST $ rwsZoom      o m
  zoomMaybe o = \(L.RWST m) -> L.RWST $ rwsZoomMaybe o m
  zoomMany  o = \(L.RWST m) -> L.RWST $ rwsZoomMany  o m
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance (Monoid w, Zoom m n s t) => Zoom (S.WriterT w m) (S.WriterT w n) s t where
  zoom      o = S.WriterT #.                 zoom      o .# S.runWriterT
  zoomMaybe o = S.WriterT #. fmap shuffleW . zoomMaybe o .# S.runWriterT
  zoomMany  o = S.WriterT #.                 zoomMany  o .# S.runWriterT
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance (Monoid w, Zoom m n s t) => Zoom (L.WriterT w m) (L.WriterT w n) s t where
  zoom      o = L.WriterT #.                 zoom      o .# L.runWriterT
  zoomMaybe o = L.WriterT #. fmap shuffleW . zoomMaybe o .# L.runWriterT
  zoomMany  o = L.WriterT #.                 zoomMany  o .# L.runWriterT
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance Zoom m n s t => Zoom (ListT m) (ListT n) s t where
  zoom      o = ListT #.                  zoom      o .# runListT
  zoomMaybe o = ListT #. fmap sequenceA . zoomMaybe o .# runListT
  zoomMany  o = ListT #.                  zoomMany  o .# runListT
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance Zoom m n s t => Zoom (MaybeT m) (MaybeT n) s t where
  zoom o =
    MaybeT #. zoom o .# runMaybeT
  zoomMaybe o =
    MaybeT #. fmap (getMay . shuffleMay) . zoomMaybe o . fmap May .# runMaybeT
  zoomMany o =
    MaybeT #. fmap getMay . zoomMany o . fmap May .# runMaybeT
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance (Error e, Zoom m n s t) => Zoom (ErrorT e m) (ErrorT e n) s t where
  zoom o =
    ErrorT #. zoom o .# runErrorT
  zoomMaybe o =
    ErrorT #. fmap (getErr . shuffleErr) . zoomMaybe o . fmap Err .# runErrorT
  zoomMany o =
    ErrorT #. fmap getErr . zoomMany o . fmap Err .# runErrorT
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

instance Zoom m n s t => Zoom (ExceptT e m) (ExceptT e n) s t where
  zoom o =
    ExceptT #. zoom o .# runExceptT
  zoomMaybe o =
    ExceptT #. fmap (getErr . shuffleErr) . zoomMaybe o . fmap Err .# runExceptT
  zoomMany o =
    ExceptT #. fmap getErr . zoomMany o . fmap Err .# runExceptT
  {-# INLINE zoom #-}
  {-# INLINE zoomMaybe #-}
  {-# INLINE zoomMany #-}

------------------------------------------------------------------------------
-- Magnify
------------------------------------------------------------------------------

-- | This class allows us to 'magnify' part of the environment, changing the
-- environment supplied by many different 'Monad' transformers. Unlike 'zoom'
-- this can change the environment of a deeply nested 'Monad' transformer.
--
-- Its functions can be used to run a monadic action in a larger environment
-- than it was defined in, using a 'Getter' or an 'AffineFold'.
--
-- They act like 'Control.Monad.Reader.Class.local', but can in many cases
-- change the type of the environment as well.
--
-- They're commonly used to lift actions in a simpler 'Reader' 'Monad' into a
-- 'Monad' with a larger environment type.
--
-- They can be used to edit pretty much any 'Monad' transformer stack with an
-- environment in it:
--
-- >>> (1,2) & magnify _2 (+1)
-- 3
--
-- >>> flip runReader (1,2) $ magnify _1 ask
-- 1
--
-- >>> flip runReader (1,2,[10..20]) $ magnifyMaybe (_3 % _tail) ask
-- Just [11,12,13,14,15,16,17,18,19,20]
--
class
  (MonadReader b m, MonadReader a n
  ) => Magnify m n b a | m -> b, n -> a, m a -> n, n b -> m where
  magnify
    :: Is k A_Getter
    => Optic' k is a b
    -> m c
    -> n c

  magnifyMaybe
    :: Is k An_AffineFold
    => Optic' k is a b
    -> m c
    -> n (Maybe c)

-- | Extends 'Magnify' with an ability to magnify using a 'Fold' over multiple
-- targets so that actions for each one are executed sequentially and the
-- results are aggregated.
--
-- There is however no sensible instance of 'MagnifyMany' for 'StateT'.
class
  (MonadReader b m, MonadReader a n, Magnify m n b a
  ) => MagnifyMany m n b a | m -> b, n -> a, m a -> n, n b -> m where
  magnifyMany
    :: (Is k A_Fold, Monoid c)
    => Optic' k is a b
    -> m c
    -> n c

-- | @
-- 'magnify'      = 'views'
-- 'magnifyMaybe' = 'previews'
-- @
instance Magnify ((->) b) ((->) a) b a where
  magnify      = views
  magnifyMaybe = previews
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

-- | @
-- 'magnifyMany' = 'foldMapOf'
-- @
instance MagnifyMany ((->) b) ((->) a) b a where
  magnifyMany = foldMapOf
  {-# INLINE magnifyMany #-}

instance Monad m => Magnify (ReaderT b m) (ReaderT a m) b a where
  magnify o = \(ReaderT m) ->
    ReaderT $ \r -> getEffect (views o (Effect #. m) r)
  magnifyMaybe o = \(ReaderT m) ->
    ReaderT $ \r -> traverse getEffect (previews o (Effect #. m) r)
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance Monad m => MagnifyMany (ReaderT b m) (ReaderT a m) b a where
  magnifyMany o = \(ReaderT m) ->
    ReaderT $ \r -> getEffect (foldMapOf o (Effect #. m) r)
  {-# INLINE magnifyMany #-}

instance (Monad m, Monoid w) => Magnify (S.RWST b w s m) (S.RWST a w s m) b a where
  magnify      o = \(S.RWST m) -> S.RWST $ rwsMagnify      o m
  magnifyMaybe o = \(S.RWST m) -> S.RWST $ rwsMagnifyMaybe o m
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance
  (Monad m, Monoid w
  ) => MagnifyMany (S.RWST b w s m) (S.RWST a w s m) b a where
  magnifyMany o = \(S.RWST m) -> S.RWST $ rwsMagnifyMany  o m
  {-# INLINE magnifyMany #-}

instance (Monad m, Monoid w) => Magnify (L.RWST b w s m) (L.RWST a w s m) b a where
  magnify      o = \(L.RWST m) -> L.RWST $ rwsMagnify      o m
  magnifyMaybe o = \(L.RWST m) -> L.RWST $ rwsMagnifyMaybe o m
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance
  (Monad m, Monoid w
  ) => MagnifyMany (L.RWST b w s m) (L.RWST a w s m) b a where
  magnifyMany o = \(L.RWST m) -> L.RWST $ rwsMagnifyMany  o m
  {-# INLINE magnifyMany #-}

instance Magnify m n b a => Magnify (IdentityT m) (IdentityT n) b a where
  magnify      o = \(IdentityT m) -> IdentityT (magnify      o m)
  magnifyMaybe o = \(IdentityT m) -> IdentityT (magnifyMaybe o m)
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance MagnifyMany m n b a => MagnifyMany (IdentityT m) (IdentityT n) b a where
  magnifyMany o = \(IdentityT m) -> IdentityT (magnifyMany  o m)
  {-# INLINE magnifyMany #-}

instance Magnify m n b a => Magnify (S.StateT s m) (S.StateT s n) b a where
  magnify      o = \(S.StateT m) -> S.StateT $ magnify o . m
  magnifyMaybe o = \(S.StateT m) -> S.StateT $ \s ->
    fmap (shuffleS s) $ magnifyMaybe o (m s)
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

-- No instance for MagnifyMany (S.StateT s m) (S.StateT s n) b a

instance Magnify m n b a => Magnify (L.StateT s m) (L.StateT s n) b a where
  magnify      o = \(L.StateT m) -> L.StateT $ magnify o . m
  magnifyMaybe o = \(L.StateT m) -> L.StateT $ \s ->
    fmap (shuffleS s) $ magnifyMaybe o (m s)
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

-- No instance for MagnifyMany (L.StateT s m) (L.StateT s n) b a

instance
  (Monoid w, Magnify m n b a
  ) => Magnify (S.WriterT w m) (S.WriterT w n) b a where
  magnify      o = S.WriterT #.                 magnify      o .# S.runWriterT
  magnifyMaybe o = S.WriterT #. fmap shuffleW . magnifyMaybe o .# S.runWriterT
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance
  (Monoid w, MagnifyMany m n b a
  ) => MagnifyMany (S.WriterT w m) (S.WriterT w n) b a where
  magnifyMany o = S.WriterT #. magnifyMany o .# S.runWriterT
  {-# INLINE magnifyMany #-}

instance
  (Monoid w, Magnify m n b a
  ) => Magnify (L.WriterT w m) (L.WriterT w n) b a where
  magnify      o = L.WriterT #.                 magnify      o .# L.runWriterT
  magnifyMaybe o = L.WriterT #. fmap shuffleW . magnifyMaybe o .# L.runWriterT
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance
  (Monoid w, MagnifyMany m n b a
  ) => MagnifyMany (L.WriterT w m) (L.WriterT w n) b a where
  magnifyMany o = L.WriterT #. magnifyMany o .# L.runWriterT
  {-# INLINE magnifyMany #-}

instance Magnify m n b a => Magnify (ListT m) (ListT n) b a where
  magnify      o = ListT #.                  magnify      o .# runListT
  magnifyMaybe o = ListT #. fmap sequenceA . magnifyMaybe o .# runListT
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance MagnifyMany m n b a => MagnifyMany (ListT m) (ListT n) b a where
  magnifyMany o = ListT #. magnifyMany o .# runListT
  {-# INLINE magnifyMany #-}

instance Magnify m n b a => Magnify (MaybeT m) (MaybeT n) b a where
  magnify o = MaybeT #. magnify o .# runMaybeT
  magnifyMaybe o =
    MaybeT #. fmap (getMay . shuffleMay) . magnifyMaybe o . fmap May .# runMaybeT
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance MagnifyMany m n b a => MagnifyMany (MaybeT m) (MaybeT n) b a where
  magnifyMany o = MaybeT #. fmap getMay . magnifyMany o . fmap May .# runMaybeT
  {-# INLINE magnifyMany #-}

instance (Error e, Magnify m n b a) => Magnify (ErrorT e m) (ErrorT e n) b a where
  magnify o = ErrorT #. magnify o .# runErrorT
  magnifyMaybe o =
    ErrorT #. fmap (getErr . shuffleErr) . magnifyMaybe o . fmap Err .# runErrorT
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance
  (Error e, MagnifyMany m n b a
  ) => MagnifyMany (ErrorT e m) (ErrorT e n) b a where
  magnifyMany o = ErrorT #. fmap getErr . magnifyMany o . fmap Err .# runErrorT
  {-# INLINE magnifyMany #-}

instance Magnify m n b a => Magnify (ExceptT e m) (ExceptT e n) b a where
  magnify o = ExceptT #. magnify o .# runExceptT
  magnifyMaybe o =
    ExceptT #. fmap (getErr . shuffleErr) . magnifyMaybe o . fmap Err .# runExceptT
  {-# INLINE magnify #-}
  {-# INLINE magnifyMaybe #-}

instance MagnifyMany m n b a => MagnifyMany (ExceptT e m) (ExceptT e n) b a where
  magnifyMany o = ExceptT #. fmap getErr . magnifyMany o . fmap Err .# runExceptT
  {-# INLINE magnifyMany #-}