module Data.Stream
  ( Stream
  , mkStream
  , head
  , tail
  , cons
  , elemAt
  , toSeq
  , toList
  , mapStreamSt
  , fibs
  ) where

import Prelude hiding (head, tail, drop)
import Control.Comonad
import Control.Arrow ((&&&))

data Stream a = forall b.  S b (b -> a) (b -> b)

mkStream :: b -> (b -> a) -> (b -> b) -> Stream a
mkStream x f g = S x f g

instance Functor Stream where
  fmap = liftW

instance Comonad Stream where
  extract (S s f g)  = f s
  extend h (S s f g) = S s (\s' -> h (S s' f g)) g

head :: Stream a -> a
head (S x f g) = f x

tail :: Stream a -> Stream a
tail (S x f g) = S (g x) f g

cons :: a -> Stream a -> Stream a
cons x s = mkStream (x,s) fst ((head &&& tail) . snd)

elemAt :: Integral i => i -> Stream a -> a
elemAt n = head . drop n

drop :: Integral i => i -> Stream a -> Stream a
drop 0 = id
drop n = drop (n - 1) . tail

toSeq :: Stream a -> Int -> a
toSeq = flip elemAt

toList :: Stream a -> [a]
toList = map head . iterate tail

mapStreamSt :: (a -> s -> b) -> (a -> s -> s) -> s -> Stream a -> Stream b
mapStreamSt f1 f2 s0 xs
  = mkStream (xs,s0) 
             (\(x,s) -> f1 (head x) s)
             (\(x,s) -> (tail x, f2 (head x) s))

fibs :: Stream Integer
fibs = mkStream (1,1) fst (\(i,j) -> (j,i+j))

{-
------
-- A stream can be pulled from any comonad

parallelW :: Comonad w => w (Stream a) -> Stream (w a)
parallelW w = mkStream w (fmap head) (fmap tail)

-- in fact, from any functor
parallelF :: (Functor f) => f (Stream a) -> Stream (f a)
parallelF f = mkStream f (fmap head) (fmap tail)


------
-- What good are these?

mapW :: Comonad w => (w a -> b) -> w (Stream a) -> Stream b
mapW f = fmap f . parallelW

unfold :: (b -> (a,b)) -> b -> Stream a
unfold f i = mkStream i (fst . f) (snd . f)

unfoldW :: Comonad w => (w b -> (a,b)) -> w b -> Stream a
unfoldW f w = mkStream w (fst . f) (=>> snd . f)

mkStreamW :: Comonad w => w b -> (w b -> a) -> (w b -> b) -> Stream a
mkStreamW w f g = mkStream w f (=>> g)
-}