-- | This module contains lenses and traversals for common structures in Haskell.
-- It also contains the combinators for lenses and traversals.
module Lens.Family.Stock (
-- * Lens Combinators
    choosing
  , alongside
  , beside
-- * Stock Lenses
  , _1, _2, both
  , chosen
  , ix
  , at, intAt
  , contains, intContains
-- * Stock Traversals
  , _Left, _Right
  , _Just, _Nothing
  , ignored
-- * Types
  , AlongsideLeft, AlongsideRight
-- * Re-exports
  , LensLike, LensLike'
  , Applicative
  ) where

import Control.Arrow (first, second)
import Control.Applicative (Applicative, pure, (<$>), (<*>))
import Lens.Family (LensLike, LensLike')
import Lens.Family.Unchecked (lens)
import Lens.Family.Phantom (Phantom, coerce)
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import qualified Data.Set as Set
import qualified Data.IntSet as IntSet

choosing :: Functor f => LensLike f a a' c c' -> LensLike f b b' c c' -> LensLike f (Either a b) (Either a' b') c c'
-- ^ @
-- choosing :: Lens a a' c c' -> Lens b b' c c' -> Lens (Either a b) (Either a' b') c c'
-- @
--
-- @
-- choosing :: Traversal a a' c c' -> Traversal b b' c c' -> Traversal (Either a b) (Either a' b') c c'
-- @
--
-- @
-- choosing :: Getter a a' c c' -> Getter b b' c c' -> Getter (Either a b) (Either a' b') c c'
-- @
--
-- @
-- choosing :: Fold a a' c c' -> Fold b b' c c' -> Fold (Either a b) (Either a' b') c c'
-- @
--
-- @
-- choosing :: Setter a a' c c' -> Setter b b' c c' -> Setter (Either a b) (Either a' b') c c'
-- @
--
-- Given two lens\/traversal\/getter\/fold\/setter families with the same substructure, make a new lens\/traversal\/getter\/fold\/setter on 'Either'.
choosing la _  f (Left a)  = Left  `fmap` la f a
choosing _  lb f (Right b) = Right `fmap` lb f b

_1 :: Functor f => LensLike f (a, b) (a', b) a a'
-- ^ @
-- _1 :: Lens (a, b) (a', b) a a'
-- @
--
-- Lens on the first element of a pair.
_1 f (a, b) = (\a' -> (a', b)) `fmap` f a

_2 :: Functor f => LensLike f (a, b) (a, b') b b'
-- ^ @
-- _2 :: Lens (a, b) (a, b') b b'
-- @
--
-- Lens on the second element of a pair.
_2 f (a, b) = (\b' -> (a, b')) `fmap` f b

chosen :: Functor f => LensLike f (Either a a) (Either b b) a b
-- ^ @
-- chosen :: Lens (Either a a) (Either b b) a b
-- @
--
-- Lens on the Left or Right element of an ('Either' a a).
chosen = choosing id id

ix :: (Eq k, Functor f) => k -> LensLike' f (k -> v) v
-- ^ @
-- ix :: Eq k => k -> Lens' (k -> v) v
-- @
--
-- Lens on a given point of a function.
ix k f g = (\v' x -> if (k == x) then v' else g x) `fmap` f (g k)

at :: (Ord k, Functor f) => k -> LensLike' f (Map.Map k v) (Maybe v)
-- ^ @
-- at :: Ord k => k -> Lens' (Map.Map k v) (Maybe v)
-- @
--
-- Lens on a given point of a 'Map.Map'.
at k = lens (Map.lookup k) (\m -> maybe (Map.delete k m) (\v' -> Map.insert k v' m))

intAt :: Functor f => Int -> LensLike' f (IntMap.IntMap v) (Maybe v)
-- ^ @
-- intAt :: Int -> Lens (IntMap.IntMap v) (Maybe v)
-- @
--
-- Lens on a given point of a 'IntMap.IntMap'.
intAt k = lens (IntMap.lookup k) (\m -> maybe (IntMap.delete k m) (\v' -> IntMap.insert k v' m))

contains :: (Ord k, Functor f) => k -> LensLike' f (Set.Set k) Bool
-- ^ @
-- contains :: Ord => k -> Lens' (Set.Set k) Bool
-- @
--
-- Lens on a given point of a 'Set.Set'.
contains k = lens (Set.member k) (\m nv -> if nv then Set.insert k m else Set.delete k m)

intContains :: Functor f => Int -> LensLike' f IntSet.IntSet Bool
-- ^ @
-- intContains :: Int -> Lens' IntSet.IntSet Bool
-- @
--
-- Lens on a given point of a 'IntSet.IntSet'.
intContains k = lens (IntSet.member k) (\m nv -> if nv then IntSet.insert k m else IntSet.delete k m)

_Left :: Applicative f => LensLike f (Either a b) (Either a' b) a a'
-- ^ @
-- _Left :: Traversal (Either a b) (Either a' b) a a'
-- @
--
-- Traversal on the 'Left' element of an 'Either'.
_Left f (Left a) = Left <$> f a
_Left _ (Right b) = pure (Right b)

_Right :: Applicative f => LensLike f (Either a b) (Either a b') b b'
-- ^ @
-- _Right :: Traversal (Either a b) (Either a b') b b'
-- @
--
-- Traversal on the 'Right' element of an 'Either'.
_Right f (Right b) = Right <$> f b
_Right _ (Left a) = pure (Left a)

_Just :: Applicative f => LensLike f (Maybe a) (Maybe a') a a'
-- ^ @
-- _Just :: Traversal (Maybe a) (Maybe a') a a'
-- @
--
-- Traversal on the 'Just' element of a 'Maybe'.
_Just f (Just a) = Just <$> f a
_Just _ Nothing = pure Nothing

_Nothing :: Applicative f => LensLike' f (Maybe a) ()
-- ^ @
-- _Nothing :: Traversal' (Maybe a) ()
-- @
--
-- Traversal on the 'Nothing' element of a 'Maybe'.
_Nothing f Nothing = const Nothing <$> f ()
_Nothing _ j = pure j

both :: Applicative f => LensLike f (a,a) (b,b) a b
-- ^ @
-- both :: Traversal (a,a) (b,b) a b
-- @
--
-- Traversals on both elements of a pair @(a,a)@.
both f (x,y) = (,) <$> f x <*> f y

beside :: Applicative f => LensLike f a a' c c' -> LensLike f b b' c c' -> LensLike f (a,b) (a',b') c c'
-- ^ @
-- beside :: Traversal a a' c c' -> Traversal b' b' c c' -> Traversal (a,b) (a',b') c c'
-- @
--
-- @
-- beside :: Fold a a' c c' -> Fold b' b' c c' -> Fold (a,b) (a',b') c c'
-- @
--
-- @
-- beside :: Setter a a' c c' -> Setter b' b' c c' -> Setter (a,b) (a',b') c c'
-- @
--
-- Given two traversals\/folds\/setters referencing a type 'c', create a traversal\/fold\/setter on the pair referencing 'c'.
beside la lb f (x,y) = (,) <$> la f x <*> lb f y

ignored :: Applicative f => null -> a -> f a
-- ^ @
-- ignored :: Traversal a a b b'
-- @
--
-- The empty traversal on any type.
ignored _ = pure

{- Alongside -}

newtype AlongsideLeft f b a = AlongsideLeft (f (a, b))

instance Functor f => Functor (AlongsideLeft f a) where
  fmap f (AlongsideLeft x) = AlongsideLeft (fmap (first f) x)

instance Phantom f => Phantom (AlongsideLeft f a) where
  coerce (AlongsideLeft x) = AlongsideLeft (coerce x)

newtype AlongsideRight f a b = AlongsideRight (f (a, b))

instance Functor f => Functor (AlongsideRight f a) where
  fmap f (AlongsideRight x) = AlongsideRight (fmap (second f) x)

instance Phantom f => Phantom (AlongsideRight f a) where
  coerce (AlongsideRight x) = AlongsideRight (coerce x)

alongside :: Functor f => LensLike (AlongsideLeft f b2') a1 a1' b1 b1'
                       -> LensLike (AlongsideRight f a1') a2 a2' b2 b2'
                       -> LensLike f (a1, a2) (a1', a2') (b1, b2) (b1', b2')
-- ^ @
-- alongside :: Lens a1 a1' b1 b1' -> Lens a2 a2' b2 b2' -> Lens (a1, a2) (a1', a2') (b1, b2) (b1', b2')
-- @
--
-- @
-- alongside :: Getter a1 a1' b1 b1' -> Getter a2 a2' b2 b2' -> Getter (a1, a2) (a1', a2') (b1, b2) (b1', b2')
-- @
--
-- Given two lens\/getter families, make a new lens\/getter on their product.
alongside l1 l2 f (a1, a2) = fa1'a2'
 where
  AlongsideRight fa1'a2' = l2 f2 a2
  f2 b2 = AlongsideRight fa1'b2'
   where
    AlongsideLeft fa1'b2' = l1 f1 a1
    f1 b1 = AlongsideLeft (f (b1, b2))