-- | This module contains lenses for common structures in Haskell.
-- It also contains the lens combinators 'mergeL' and '***'.
module Lens.Family.Stock
  ( -- * Lens Combinators
    mergeL
  , (***)
  -- * Stock Lenses
  , fstL, sndL
  , funL
  , mapL, intMapL
  , setL, intSetL
  -- * Types
  , RefFamily, Ref
  ) where

import Lens.Family.Unchecked (RefFamily, Ref, mkLens)
import Lens.Family ((^.), (<~))
import Lens.Family.Clone (ClonerFamily, clone)
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import qualified Data.Set as Set
import qualified Data.IntSet as IntSet

-- | Given two lens\/getter\/setter families with the same substructure, make a new lens\/getter\/setter on 'Either'.
mergeL :: Functor f => RefFamily f a a' c c' -> RefFamily f b b' c c' -> RefFamily f (Either a b) (Either a' b') c c'
mergeL la _  f (Left a)  = Left  `fmap` la f a
mergeL _  lb f (Right b) = Right `fmap` lb f b

-- I suspect there is a more clever way to define this function.
-- | Given two lens families, make a new lens on their product.
(***) :: Functor f => ClonerFamily a1 a1' b1 b1' -> ClonerFamily a2 a2' b2 b2' -> RefFamily f (a1, a2) (a1', a2') (b1, b2) (b1', b2')
(***) l1 l2 f (a1, a2) = (\(v'1, v'2) -> (cl1 <~ v'1 $ a1, cl2 <~ v'2 $ a2)) `fmap` f (a1 ^. cl1, a2 ^. cl2)
  where
    cl1 x = clone l1 x
    cl2 x = clone l2 x

-- | Lens on the first element of a pair.
fstL :: Functor f => RefFamily f (a, b) (a', b) a a'
fstL f (a, b) = (\a' -> (a', b)) `fmap` f a

-- | Lens on the second element of a pair.
sndL :: Functor f => RefFamily f (a, b) (a, b') b b'
sndL f (a, b) = (\b' -> (a, b')) `fmap` f b

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

-- | Lens on a given point of a 'Map.Map'.
mapL :: (Ord k, Functor f) => k -> Ref f (Map.Map k v) (Maybe v)
mapL k = mkLens (Map.lookup k) (\m -> maybe (Map.delete k m) (\v' -> Map.insert k v' m))

-- | Lens on a given point of a 'IntMap.IntMap'.
intMapL :: (Functor f) => Int -> Ref f (IntMap.IntMap v) (Maybe v)
intMapL k = mkLens (IntMap.lookup k) (\m -> maybe (IntMap.delete k m) (\v' -> IntMap.insert k v' m))

-- | Lens on a given point of a 'Set.Set'.
setL :: (Ord k, Functor f) => k -> Ref f (Set.Set k) Bool
setL k = mkLens (Set.member k) (\m nv -> if nv then Set.insert k m else Set.delete k m)

-- | Lens on a given point of a 'IntSet.IntSet'.
intSetL :: (Functor f) => Int -> Ref f IntSet.IntSet Bool
intSetL k = mkLens (IntSet.member k) (\m nv -> if nv then IntSet.insert k m else IntSet.delete k m)