-- We use the same source code for the lazy and strict versions, however the types are different.
-- We can do this using pattern synonyms

--------------------------------------------------------------------------------

-- We maintain the invariant that @(Z Nil)@ never appears. We could use a different
-- representation and the type system to eliminate this possibility, however, that
-- would have a runtime cost (an extra indirection and extra 2 words per element),
-- so we do this instead.
-- 
pattern Z xs <- ZZ xs where 
  Z xs = case xs of
   Nil -> Nil
   _   -> ZZ xs

-- | We maintain the invariant that @(Z Nil)@ never appears. This function 
-- checks whether this is satisfied. Used only for testing.
checkInvariant :: Seq a -> Bool
checkInvariant = go where
  go :: Seq b -> Bool
  go seq = case seq of
    Nil    -> True
    Z   ys -> case ys of { Nil -> False ; _ -> go ys }
    O _ ys -> go ys

fst,snd :: Pair a -> a
fst (Pair x _) = x
snd (Pair _ y) = y

--------------------------------------------------------------------------------
-- instances

instance Monoid (Seq a) where
  mempty  = empty
  mappend = append

instance Eq a => Eq (Seq a) where
  (==) (Nil   ) seq = null seq
  (==) (Z   xs) seq = case seq of
    Z ys   -> xs == ys
    Nil    -> null xs
    _      -> False
  (==) (O x xs) seq = case seq of
    O y ys -> x == y && xs == ys
    _      -> False

instance Ord a => Ord (Seq a) where
  compare = comparing toList

#ifndef TESTING

instance Show a => Show (Seq a) where
  showsPrec p xs = showParen (p > 10) $
    showString "fromList " . shows (toList xs)

#endif

instance Functor Seq where

  fmap f seq = 
    case seq of
      Nil    -> Nil
      Z   ys -> Z (fmap g ys)  
      O y ys -> O (f y) (fmap g ys)
    where 
      g (Pair x y) = Pair (f x) (f y)

instance F.Foldable Seq where

  foldr f s seq = 
    case seq of  
      Nil    -> s
      Z   ys ->      foldr g s ys
      O y ys -> f y (foldr g s ys)
    where
      g (Pair x y) s = f x (f y s)

  foldl f s seq = 
    case seq of  
      Nil    -> s
      Z   ys -> foldl g    s    ys
      O y ys -> foldl g (f s y) ys
    where
      g s (Pair x y) = f (f s x) y

  foldMap f seq = 
    case seq  of
      Nil    -> mempty
      Z   ys -> foldMap g ys
      O y ys -> mappend (f y) (foldMap g ys)
    where
      g (Pair x y) = mappend (f x) (f y)

  null   = unsafeNull
  length = seqLength
  toList = seqToList

--------------------------------------------------------------------------------
-- * Accessing the left end of the sequence

-- | Prepending an element. Worst case @O(log(n))@, but amortized @O(1)@.
cons :: a -> Seq a -> Seq a
cons x seq = case seq of
  Nil    -> O x Nil
  Z   ys -> O x ys
  O y ys -> Z $ cons (Pair x y) ys

-- | Worst case @O(log(n))@, amortized @O(1)@
unCons :: Seq a -> Maybe (a, Seq a)
unCons seq = case seq of
  Nil    -> Nothing
  O y ys -> Just $ (y , Z ys)
  Z   ys -> case unCons ys of
    Just (Pair x y, rest) -> Just (x , cons y $ Z rest)
    Nothing               -> Nothing

--------------------------------------------------------------------------------
-- * Basic queries 

-- | Checks whether the sequence is empty. This is @O(1)@.
null :: Seq a -> Bool
null = unsafeNull

-- | The length of a sequence. @O(log(n))@.
length :: Seq a -> Int
length = seqLength 

-- we do NOT assume that @(Z Nil)@ does not appear
safeNull :: Seq a -> Bool
safeNull = go where
  go :: Seq b -> Bool
  go seq = case seq of
    Nil  -> True
    Z xs -> go xs
    O {} -> False        

-- we assume that @(Z Nil)@ does not appear
unsafeNull :: Seq a -> Bool
unsafeNull seq = case seq of
  Nil -> True
  _   -> False

-- because Data.Foldable also calls its member function @length@...
seqLength :: Seq a -> Int
seqLength = go where
  go :: Seq b -> Int
  go seq = case seq of 
    Nil    -> 0
    Z   xs ->     2 * go xs
    O _ xs -> 1 + 2 * go xs

--------------------------------------------------------------------------------
-- * Basic construction

-- | The empty sequence.
empty :: Seq a 
empty = Nil

-- | Conversion from a list. @O(n)@.
fromList :: [a] -> Seq a
fromList = foldr cons Nil

{- 

-- | Builds a sequence of the given size from the input list. Unfortunately it's slower than 'fromList' 
-- for some reason... Example usage would be: @build n [1..]@
build :: Int -> [a] -> Seq a
build = go where

  go :: Int -> [b] -> Seq b
  go  0 _  = Nil
  go !n xs = 
    case (n .&. 1) of
      0 -> Z $ go halfn $ pairs xs
      1 -> case xs of
        (this:rest) -> O this (go halfn $ pairs rest)
        []          -> error "Seq.build: not enough data"
    where
      halfn = shiftR n 1
  
  pairs :: [a] -> [Pair a]
  pairs (x:y:rest) = (Pair x y) : pairs rest
  pairs []         = []
  pairs _          = error "Seq.build: fatal error (shouldn't happen)"

-}

-- | Conversion to a list. @O(n)@.
toList :: Seq a -> [a]
toList = seqToList
-- toList = foldr (:) []          -- it's not fully clear if one is better than the other or not...

-- sincee Data.Foldable also calls its member function @toList@...
seqToList :: Seq a -> [a]
seqToList = go where

  go :: Seq b -> [b]
  go seq = case seq of
    Nil    -> []
    Z   xs ->     concatPairs (go xs)
    O x xs -> x : concatPairs (go xs)

  concatPairs :: [Pair a] -> [a]
  concatPairs pairs = case pairs of
    Pair x y : rest -> x : y : concatPairs rest
    [] -> []

-- | Naive implementation of 'toList'
toListNaive :: Seq a -> [a]
toListNaive = go where
  go seq = case unCons seq of
    Nothing     -> []
    Just (x,xs) -> x : go xs

--------------------------------------------------------------------------------
-- * Short sequences

singleton :: a -> Seq a
singleton x = O x Nil

pair :: a -> a -> Seq a
pair x y = Z (O (Pair x y) Nil)

triple :: a -> a -> a -> Seq a
triple x y z = O x (O (Pair y z) Nil)

quad :: a -> a -> a -> a -> Seq a
quad x y z w = Z (Z (O (Pair (Pair x y) (Pair z w)) Nil))

--------------------------------------------------------------------------------
-- * Unsafe head and tail

-- | First element of the sequence. Worst case @O(log(n))@, amortized @O(1)@.
head :: Seq a -> a
head seq = case mbHead seq of
  Just y  -> y
  Nothing -> error "Seq.head"

-- | Tail of the sequence. Worst case @O(log(n))@, amortized @O(1)@.
tail :: Seq a -> Seq a
tail seq = case mbTail seq of
  Just ys -> ys
  Nothing -> error "Seq.tail"

-- | Last element of the sequence. @O(log(n))@.
last :: Seq a -> a
last seq = case mbLast seq of
  Just y  -> y
  Nothing -> error "Seq.last"

--------------------------------------------------------------------------------
-- * Safe head and tail

-- | First element of the sequence. Worst case @O(log(n))@, amortized @O(1)@.
mbHead :: Seq a -> Maybe a
mbHead seq = case seq of
  Nil    -> Nothing
  O y ys -> Just y
  Z   ys -> fst <$> mbHead ys

-- | Last element of the sequence. @O(log(n))@
mbLast :: Seq a -> Maybe a
mbLast seq = case seq of
  Nil    -> Nothing
  Z   ys -> snd <$> mbLast ys
  O y ys -> case mbLast ys of
    Just ab -> Just (snd ab)
    Nothing -> Just y

-- | Tail of the sequence. Worst case @O(log(n))@, amortized @O(1)@.
mbTail :: Seq a -> Maybe (Seq a)
mbTail seq = case seq of
  Nil    -> Nothing
  O y ys -> Just (Z ys)
  Z   ys -> case unCons ys of
    Just ((Pair a b),rest) -> Just (cons b (Z rest))
    Nothing                -> Nothing

-- | All tails of the sequence (starting with the sequence itself)
tails :: Seq a -> [Seq a]
tails = go where
  go :: Seq a -> [Seq a]
  go seq = seq : case mbTail seq of
    Just ys -> go ys
    Nothing -> []
                
--------------------------------------------------------------------------------
-- * indexing

-- | Lookup the @k@-th element of a sequence. This is worst case @O(log(n))@ and amortized @O(log(k))@, and quite efficient.
lookup :: Int -> Seq a -> a
lookup !n
  | n < 0     = error "Seq.lookup: negative index"
  | otherwise = go n 
  where
    go :: Int -> Seq a -> a
    go !k seq = case seq of
      Nil    -> error "Seq.lookup"
      Z   xs -> cont k xs 
      O x xs -> if k==0 then x else cont (k-1) xs              

    cont !k xs = select (k .&. 1) $ go (shiftR k 1) xs
    select !bit (Pair x y) = if bit==0 then x else y

-- | Update the @k@-th element of a sequence. 
update :: (a -> a) -> Int -> Seq a -> Seq a
update f n 
  | n < 0     = error "Seq.update: negative index"
  | otherwise = go f n
  where
    go :: (a -> a) -> Int -> Seq a -> Seq a
    go f !k seq = 
      case seq of
        Nil    -> error "Seq.update"
        Z   xs -> Z (cont k xs)
        O x xs -> if k==0 then O (f x) xs else O x (cont (k-1) xs)
      where     
        cont !k xs = go (mkF (k .&. 1)) (shiftR k 1) xs where
          mkF bit = case bit of
            0 -> \(Pair x y) -> Pair (f x) y
            _ -> \(Pair x y) -> Pair x (f y)

-- | Replace the @k@-th element. @replace n x == update (const x) n@
replace :: Int -> a -> Seq a -> Seq a
replace n x = update (const x) n

mbLookup :: Int -> Seq a -> Maybe a
mbLookup !n seq
  | n < 0     = Nothing
  | otherwise = go n seq
  where
    go :: Int -> Seq a -> Maybe a
    go !k seq = case seq of
      Nil    -> Nothing
      Z   xs -> cont k xs 
      O x xs -> if k==0 then Just x else cont (k-1) xs              

    cont !k xs = select (k .&. 1) <$> go (shiftR k 1) xs
    select !bit (Pair x y) = if bit==0 then x else y

-- | Drop is efficient: @drop k@ is amortized @O(log(k))@, worst case maybe @O(log(n)^2)@ ?
drop :: Int -> Seq a -> Seq a
drop = go where

  go :: Int -> Seq b -> Seq b
  go  0 seq = seq
  go !k seq = case seq of
    Nil    -> Nil
    O x xs -> cont (k-1) xs
    Z   xs -> cont  k    xs

  cont :: Int -> Seq (Pair b) -> Seq b
  cont !k xs = 
    let ps = go (shiftR k 1) xs
    in  case (k .&. 1) of
          0 -> Z ps
          _ -> case unCons ps of  
            Just (Pair a b, rest) -> cons b (Z rest)
            Nothing               -> Nil

--------------------------------------------------------------------------------
-- * Slow operations, for completeness

-- | @O(n)@ (for large @n@ at least), where @n@ is the length of the first sequence.
append :: Seq a -> Seq a -> Seq a
append seq1 seq2
  | Nil <- seq1  = seq2
  | Nil <- seq2  = seq1
  | otherwise    = foldr cons seq2 (toList seq1)

-- | Take is slow: @O(n)@
take :: Int -> Seq a -> Seq a
take n = fromList . Prelude.take n . toList

-- | The sequence without the last element. Warning, this is slow, @O(n)@
init :: Seq a -> Seq a
init = fromList . Prelude.init . toList

-- | The sequence without the last element. Warning, this is slow, @O(n)@
mbInit :: Seq a -> Maybe (Seq a)
mbInit seq = case null seq of
  False -> Just $ (fromList . Prelude.init . toList) seq
  True  -> Nothing

-- | Warning, this is slow: @O(n)@ (with bad constant factor).
snoc :: Seq a -> a -> Seq a 
snoc seq y = f seq where
  f = fromList . (++[y]) . toList

-- | Stripping the last element from a sequence is a slow operation, @O(n)@. 
-- If you only need extracting the last element, use 'mbLast' instead,
-- which is fast. 
unSnoc :: Seq a -> Maybe (Seq a, a)
unSnoc seq = case mbLast seq of
  Just y  -> Just (init seq, y)
  Nothing -> Nothing

--------------------------------------------------------------------------------

-- | Show the internal structure of the sequence. The constructor names
-- @Z@ and @O@ come from \"zero\" and \"one\", respectively.
showInternal :: Show a => Seq a -> String
showInternal seq = showsPrecInternal 0 seq ""

showsPrecInternal :: Show a => Int -> Seq a -> ShowS
showsPrecInternal = showsPrecInternal' ("Nil","Z","O")

showsPrecInternal' 
  :: Show a 
  => (String,String,String)  -- ^ name of the three constuctors: @("Nil","Even","Odd")@
  -> Int                     -- ^ precedence
  -> Seq a -> ShowS
showsPrecInternal' (nil,even,odd) = go where
  go :: Show b => Int -> Seq b -> ShowS
  go d seq =  case seq of
    Nil    -> showString nil
    Z   xs -> showParen (d > 10) $ showString even1 . go 11 xs
    O x xs -> showParen (d > 10) $ showString odd1  . showsPrec 11 x . showChar ' ' . go 11 xs
  even1 = even ++ " "
  odd1  = odd  ++ " "

--------------------------------------------------------------------------------
-- * Graphviz

-- so that we don't need to depend on mtl or transformers
data State s a = State { runState :: s -> (s,a) }

instance Functor (State t) where
  fmap f (State run) = State $ \s -> let (s',x) = run s in (s',f x)

instance Monad (State s) where
  return x = State $ \s -> (s,x)
  a >>= u  = State $ \s -> case runState a s of (s',x) -> runState (u x) s' 

instance Applicative (State s) where
  pure  = return
  (<*>) = ap

-- | Generates a graphviz @DOT@ file, showing the internal structure of a sequence
graphviz :: Show a => Seq a -> String
graphviz = graphviz' "plaintext"

-- | You can furthermore set the shape of the nodes
graphviz' :: Show a => String -> Seq a -> String
graphviz' nodeShape seq = unlines stuff where

  stuff = prologue ++ meat ++ epilogue  
  meat  = reverse $ Prelude.snd $ Prelude.fst $ runState (go seq) (0,[])

  node    k      = "node" ++ show k
  arrow   a b    = a ++ " -> " ++ b
  nodeDef flag k text = node k ++ " [label=\"" ++ text ++ "\"" ++ shape ++ "];" where 
    shape = if flag
      then " shape=" ++ nodeShape
      else "" 

  newNode :: Bool -> String -> State (Int,[String]) String
  newNode flag text = State $ \(k,ls) -> let def = nodeDef flag k text
                                         in  ( (k+1,def:ls) , node k ) 

  newArrow :: String -> String -> State (Int,[String]) ()
  newArrow a b = State $ \(k,ls) -> let arr = arrow a b 
                                    in  ( (k  ,arr:ls) , () ) 

  go :: Show b => Seq b -> State (Int,[String]) String
  go seq = case seq of

    Nil    -> do 
      newNode True "Nil"

    ZZ  xs -> do
      node <- newNode True "Even"
      next <- go xs
      newArrow node next
      return node

    O x xs -> do
      node <- newNode True "Odd"
      this <- newNode False (show x)
      next <- go xs
      newArrow node this
      newArrow node next
      return node

  prologue = [ "digraph sequence {" 
             , "node [shape=plaintext];" 
             ]
  epilogue = [ "}" ]

--------------------------------------------------------------------------------