{-# LANGUAGE GADTs, DeriveFunctor, DeriveFoldable, TupleSections #-}

import Prelude hiding (reverse, splitAt, head, concat, foldl, concatMap)
import Control.Applicative
import Data.Foldable
import Data.Traversable
import GHC.Exts (Down(..))
import Control.Arrow ((***))
import Data.Maybe (listToMaybe)
import Data.Either (partitionEithers)
import Data.Monoid
import Data.Void

import qualified Data.Map as Map
import Data.Map (Map, (!))
import qualified Data.Stream.Infinite as Stream
import Data.Stream.Infinite (Stream(..))


-- Reactive model a la reactive banana
-- The Stream represents time, with a list of events for simultaneous events.
type Events a = Stream [a]


data Change k a = Insert k a | Delete k deriving (Show, Eq, Ord, Functor, Foldable)

deletedKeys :: Ord k => [Change k a] -> [k]
deletedKeys = concatMap f
  where
    f Insert{} = []
    f (Delete k) = [k]

mapKey :: (k1 -> k2) -> Change k1 a -> Change k2 a
mapKey f (Insert k a) = Insert (f k) a
mapKey f (Delete k)   = Delete (f k)

change :: Ord k => Change k a -> Map k a -> Map k a
change (Insert k a) = Map.insert k a
change (Delete k)   = Map.delete k

changes :: Ord k => [Change k a] -> Map k a -> Map k a
changes = flip (foldl (flip change))


data RList a where
  RList :: Ord k => Events (Change k a) -> RList a

instance Functor RList where
  fmap f (RList css) = RList $ fmap (map (fmap f)) css

instance Applicative RList where
  pure a = RList $ [Insert () a] :> pure []
  RList fss <*> RList ass = RList . snd $ mapAccumL doAcc (Map.empty, Map.empty) (Stream.zip fss ass)
    where
      doAcc (mf, ma) (fs, as) =
        let mf' = fs `changes` mf
            ma' = as `changes` ma
        in  ((mf', ma'),
             ((\c (k, a) -> mapKey (,k) (fmap ($ a) c)) <$> fs <*> Map.toList ma') ++
             ((\(k, f) c -> mapKey (k,) (fmap    f  c)) <$> Map.toList mf' <*> as) ++
             ((\k1 k2 -> Delete (k1, k2)) <$> deletedKeys fs <*> deletedKeys as))

instance Alternative RList where
  empty = RList $ pure ([] :: [Change Void a])
  RList lss <|> RList rss = RList ((\ls rs -> map (mapKey Left) ls ++ map (mapKey Right) rs) <$> lss <*> rss)



rFoldMap :: Monoid m => (a -> m) -> RList a -> Events m
rFoldMap f (RList css) = snd $ mapAccumL (\a cs -> let m = cs `changes` a in (m, [foldMap f m])) Map.empty css

toRVal :: RList a -> Events [a]
toRVal = rFoldMap (:[])

fromList :: [a] -> RList a
fromList = RList . (:> pure []) . zipWith Insert [(0::Int)..]


cons :: a -> RList a -> RList a
cons a (RList (cs :> css)) = RList ((Insert Nothing a : map (mapKey Just) cs) :> fmap (map (mapKey Just)) css)

sort :: Ord a => RList a -> RList a
sort (RList css) = RList . snd $ mapAccumL (mapAccumL (\m c -> (change c m, f c m))) Map.empty css
  where
    f (Insert _ a) _ = Insert a a
    f (Delete k)   m = Delete (m ! k)

reverse :: RList a -> RList a
reverse (RList css) = RList $ fmap (map (mapKey Down)) css

head :: RList a -> Events (Maybe a)
head = fmap (map listToMaybe) . toRVal . fst . splitAt 1

splitAt :: Int -> RList a -> (RList a, RList a)
splitAt n (RList css) = (RList *** RList) . Stream.unzip . Stream.map (partitionEithers . concat) . snd $
  mapAccumL (mapAccumL (\m c -> (change c m, f c m))) Map.empty css
  where
    f c _ | n == 0 = [Right c]
    f c@Insert{} m | Map.size m < n = [Left c]
    f c@(Insert k _) m = if k > kn_1 then [Right c] else [Left c, Left $ Delete kn_1, Right $ Insert kn_1 an_1]
      where (kn_1, an_1) = Map.elemAt (n - 1) m
    f c@Delete{} m | Map.size m <= n = [Left c]
    f c@(Delete k  ) m = if k >= kn  then [Right c] else [Left c, Left $ Insert kn an, Right $ Delete kn]
      where (kn, an) = Map.elemAt n m



test :: RList Char
test = RList $ [Insert (0::Int) 'a'] :> [Insert 1 'b'] :> [Delete 0] :> [Delete 1] :> [Insert (-1) '0'] :> [Insert 0 'A'] :> pure []