module Data.Profunctor.Extraction where

import Data.Profunctor
import Control.Comonad
import Data.Distributive
import qualified Data.List.NonEmpty as NE
import Control.Comonad.Store
import Data.Pair

class Profunctor p => Extraction p where
  extractions :: Comonad w => p (w a) b -> p (w a) (w b)

instance Extraction (Forget r) where
  extractions :: Forget r (w a) b -> Forget r (w a) (w b)
extractions (Forget f :: w a -> r
f) = (w a -> r) -> Forget r (w a) (w b)
forall r a b. (a -> r) -> Forget r a b
Forget w a -> r
f

instance Extraction (->) where
  extractions :: (w a -> b) -> w a -> w b
extractions f :: w a -> b
f = (w a -> b) -> w a -> w b
forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend w a -> b
f

instance Distributive f => Extraction (Star f) where
  extractions :: Star f (w a) b -> Star f (w a) (w b)
extractions (Star f :: w a -> f b
f) = (w a -> f (w b)) -> Star f (w a) (w b)
forall (f :: * -> *) d c. (d -> f c) -> Star f d c
Star (w (f b) -> f (w b)
forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute (w (f b) -> f (w b)) -> (w a -> w (f b)) -> w a -> f (w b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (w a -> f b) -> w a -> w (f b)
forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend w a -> f b
f)

act :: (Star f a b -> Star f s t) -> (a -> f b) -> s -> f t
act :: (Star f a b -> Star f s t) -> (a -> f b) -> s -> f t
act o :: Star f a b -> Star f s t
o f :: a -> f b
f = Star f s t -> s -> f t
forall (f :: * -> *) d c. Star f d c -> d -> f c
runStar (Star f a b -> Star f s t
o ((a -> f b) -> Star f a b
forall (f :: * -> *) d c. (d -> f c) -> Star f d c
Star a -> f b
f))

t :: (NE.NonEmpty a) -> Pair (Maybe a)
t :: NonEmpty a -> Pair (Maybe a)
t (a :: a
a NE.:| (b :: a
b : _)) = Maybe a -> Maybe a -> Pair (Maybe a)
forall a. a -> a -> Pair a
Pair (a -> Maybe a
forall a. a -> Maybe a
Just a
a) (a -> Maybe a
forall a. a -> Maybe a
Just a
b)
t (a :: a
a NE.:| _) = Maybe a -> Maybe a -> Pair (Maybe a)
forall a. a -> a -> Pair a
Pair (a -> Maybe a
forall a. a -> Maybe a
Just a
a) Maybe a
forall a. Maybe a
Nothing

u :: (NE.NonEmpty Int) -> Pair Int
u :: NonEmpty Int -> Pair Int
u xs :: NonEmpty Int
xs = Int -> Int -> Pair Int
forall a. a -> a -> Pair a
Pair (NonEmpty Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum NonEmpty Int
xs) (NonEmpty Int -> Int
forall a. NonEmpty a -> Int
NE.length NonEmpty Int
xs)




-- thingy :: forall w a b. Comonad w => (a -> w (Either a b)) -> (w a -> w b)
-- thingy f = extend loop . fmap Left
--     where
--       loop :: w (Either a b) -> b
--       loop w = case extract w of
--           Left a' -> loop $ f a'
--           Right b -> b

home :: Int -> Store Int Int -> Either Int Int
home :: Int -> Store Int Int -> Either Int Int
home n :: Int
n s :: Store Int Int
s | Store Int Int -> Int
forall (w :: * -> *) a. Comonad w => w a -> a
extract Store Int Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = Int -> Either Int Int
forall a b. a -> Either a b
Left (Int -> Either Int Int) -> Int -> Either Int Int
forall a b. (a -> b) -> a -> b
$ Store Int Int -> Int
forall s (w :: * -> *) a. ComonadStore s w => w a -> s
pos Store Int Int
s
  | Int -> Int
forall a. Num a => a -> a
abs ((Int -> Int) -> Store Int Int -> Int
forall s (w :: * -> *) a. ComonadStore s w => (s -> s) -> w a -> a
peeks (Int -> Int -> Int
forall a. Num a => a -> a -> a
+1) Store Int Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Int
forall a. Num a => a -> a
abs ((Int -> Int) -> Store Int Int -> Int
forall s (w :: * -> *) a. ComonadStore s w => (s -> s) -> w a -> a
peeks (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract 1) Store Int Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) = Int -> Either Int Int
forall a b. b -> Either a b
Right (Int -> Either Int Int) -> Int -> Either Int Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Store Int Int -> Int
forall s (w :: * -> *) a. ComonadStore s w => (s -> s) -> w a -> a
peeks (Int -> Int -> Int
forall a. Num a => a -> a -> a
+1) Store Int Int
s
  | Bool
otherwise = Int -> Either Int Int
forall a b. b -> Either a b
Right (Int -> Either Int Int) -> Int -> Either Int Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Store Int Int -> Int
forall s (w :: * -> *) a. ComonadStore s w => (s -> s) -> w a -> a
peeks (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract 1) Store Int Int
s

looper :: NE.NonEmpty Int -> Either [Int] Int
looper :: NonEmpty Int -> Either [Int] Int
looper xs :: NonEmpty Int
xs | NonEmpty Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum NonEmpty Int
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 10 = [Int] -> Either [Int] Int
forall a b. a -> Either a b
Left (NonEmpty Int -> [Int]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Int
xs)
  | Bool
otherwise = Int -> Either [Int] Int
forall a b. b -> Either a b
Right (Int -> Either [Int] Int) -> Int -> Either [Int] Int
forall a b. (a -> b) -> a -> b
$ (NonEmpty Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum NonEmpty Int
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1)

coiterate :: forall w a b.
         (Traversable w, Comonad w)
         => (w a -> Either b a)
         -> (w a -> w b)
coiterate :: (w a -> Either b a) -> w a -> w b
coiterate f :: w a -> Either b a
f = (w a -> b) -> w a -> w b
forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend w a -> b
loop
    where
      loop :: w a -> b
      loop :: w a -> b
loop = (b -> b) -> (w a -> b) -> Either b (w a) -> b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either b -> b
forall a. a -> a
id w a -> b
loop (Either b (w a) -> b) -> (w a -> Either b (w a)) -> w a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. w (Either b a) -> Either b (w a)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (w (Either b a) -> Either b (w a))
-> (w a -> w (Either b a)) -> w a -> Either b (w a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (w a -> Either b a) -> w a -> w (Either b a)
forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend w a -> Either b a
f

-- thingy :: (Choice p, Cochoice p) => p a (Either b a) -> p (w a) (w b)
-- thingy = _ . unright . _ . right'