{-# LANGUAGE Rank2Types #-}
module Clean.Lens where

import Clean.Core
import Clean.Functor
import Clean.Applicative
import Clean.Foldable

type LensLike f s t a b = (s -> f t) -> (a -> f b)
type LensLike' f a b = LensLike f b b a a
type Lens s t a b = forall f.Functor f => LensLike f s t a b
type Lens' a b = Lens b b a a

lens :: (a -> s) -> (a -> t -> b) -> Lens s t a b
lens f g = \k a -> g a <$> k (f a) 
iso :: (a -> s) -> (t -> b) -> Lens s t a b
iso f g = lens f (const g)
iso' :: (a -> b) -> (b -> a) -> Lens' a b
iso' = iso
lam f = lens f const

(^.) :: a -> Lens' a b -> b
infixl 2 ^.
x^.l = getConst (l Const x)

type Traversal s t a b = forall f. Applicative f => LensLike f s t a b
type Traversal' a b = Traversal b b a a

(%~) :: Traversal' a b -> (b -> b) -> (a -> a)
(l %~ f) a = getId (l (pure . f) a)
(.~) :: Traversal' a b -> b -> (a -> a)
l .~ x = l %~ const x

prism :: (a -> (b:+:s)) -> (a -> t -> b) -> Traversal s t a b 
prism f g = \k a -> (pure <|> map (g a) . k) (f a)
prism' :: (a -> (a:+:b)) -> (a -> b -> a) -> Traversal' a b
prism' = prism

_1 = lens fst (\(_,b) a -> (a,b))
_2 :: Lens' (a:*:b) b
_2 = lens snd (\(a,_) b -> (a,b))
_l :: Traversal' (a:+:b) a
_l = prism (\e -> (Right <|> const (Left e)) e) (const Left)
_r :: Traversal' (a:+:b) b
_r = prism (\e -> (const (Left e) <|> Right) e) (const Right)

_list :: Lens' [a] (():+:(a:*:[a]))
_list = iso (\l -> case l of
                [] -> Left ()
                (x:t) -> Right (x,t)) (const [] <|> uncurry (:))

_head :: Traversal' [a] a
_head = _list._r._1
_tail :: Traversal' [a] [a]
_tail = _list._r._2

_drop :: Int -> Traversal' [a] [a]
_drop n = foldr (.) id (_tail<$[1..n])