-- 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 :: Seq (Triple a) -> Seq a 
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
    T _ _ ys -> go ys

fst,snd,thd :: Triple a -> a
fst (Triple x _ _) = x
snd (Triple _ y _) = y
thd (Triple _ _ z) = z

--------------------------------------------------------------------------------
-- 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
  (==) (T x x2 xs) seq = case seq of
    T y y2 ys -> x == y && x2 == y2 && 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)
      T x y ys -> T (f x) (f y) (fmap g ys)
    where 
      g (Triple x y z) = Triple (f x) (f y) (f z)

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)
      T x y ys -> f x (f y (foldr g s ys))
    where
      g (Triple x y z) s = f x (f y (f z s))

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

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

  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 -> T x y ys
  T y z ys -> Z $ cons (Triple x y z) ys

cons2 :: a -> a -> Seq a -> Seq a
cons2 x y seq = case seq of
  Nil      -> T x y Nil
  Z     ys -> T x y ys
  O   z ys -> cons x (T y z ys)
  T {}     -> cons x $ cons y $ seq

-- | 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)
  T x y ys -> Just $ (x , O y ys)
  Z     ys -> case unCons ys of
    Just (Triple x y z, rest) -> Just (x , cons2 y z $ 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
    _    -> 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 ->     3 * go xs
    O   _ xs -> 1 + 3 * go xs
    T _ _ xs -> 2 + 3 * 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

-- | 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 ->         concatTriples (go xs)
    O   x xs ->     x : concatTriples (go xs)
    T x y xs -> x : y : concatTriples (go xs)

  concatTriples :: [Triple a] -> [a]
  concatTriples triples = case triples of
    Triple x y z : rest -> x : y : z : concatTriples 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 = T x y Nil

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

quad :: a -> a -> a -> a -> Seq a
quad x y z w = O x (O (Triple y 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
  T x _ ys -> Just x
  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 -> thd <$> mbLast ys
  O y ys -> case mbLast ys of
    Just abc -> Just (thd abc)
    Nothing  -> Just y
  T x y ys -> case mbLast ys of
    Just abc -> Just (thd abc)
    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)
  T x y ys -> Just (O y ys)
  Z     ys -> case unCons ys of
    Just (Triple a b c, rest) -> Just (cons2 b c (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 -> case k of { 0 -> x ;          _ -> cont (k-1) xs }             
      T x y xs -> case k of { 0 -> x ; 1 -> y ; _ -> cont (k-2) xs }

    cont !k xs = select r (go q xs) where (q,r) = divMod k 3
    select !rem (Triple x y z) = case rem of { 0 -> x ; 1 -> y ; _ -> z }

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 -> case k of { 0 -> Just x ;               _ -> cont (k-1) xs }             
      T x y xs -> case k of { 0 -> Just x ; 1 -> Just y ; _ -> cont (k-2) xs }

    cont !k xs = select r <$> (go q xs) where (q,r) = divMod k 3
    select !rem (Triple x y z) = case rem of { 0 -> x ; 1 -> y ; _ -> z }

-- | 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 -> case k of { 0 -> O (f x)   xs ;                     _ -> O   x (cont (k-1) xs) }             
        T x y xs -> case k of { 0 -> T (f x) y xs ; 1 -> T x (f y) xs ; _ -> T x y (cont (k-2) xs) }
      where     
        cont !k xs = go (mkF r) q xs where
          (q,r) = divMod k 3
          mkF !rem = case rem of
            0 -> \(Triple x y z) -> Triple (f x) y z
            1 -> \(Triple x y z) -> Triple x (f y) z
            _ -> \(Triple x y z) -> Triple x y (f z)

-- | 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

-- | 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
    T x y xs -> case k of
      1 -> O y xs
      _ -> cont (k-2) xs

  cont :: Int -> Seq (Triple b) -> Seq b
  cont !k xs = 
    let (q,r) = divMod k 3
        ts    = go q xs
    in  case r of
          0 -> Z ts
          1 -> case unCons ts of  
            Just (Triple a b c, rest) -> cons2 b c (Z rest)
            Nothing                   -> Nil      
          2 -> case unCons ts of  
            Just (Triple a b c, rest) -> cons c (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","T")

showsPrecInternal' 
  :: Show a 
  => (String,String,String,String)  -- ^ name of the foure constuctors: @("Nil","Zero","One","Two")@
  -> Int                            -- ^ precedence
  -> Seq a -> ShowS
showsPrecInternal' (nil,zero,one,two) = go where
  go :: Show b => Int -> Seq b -> ShowS
  go d seq =  case seq of
    Nil      -> showString nil
    Z     xs -> showParen (d > 10) $ showString zero1 . go 11 xs
    O   x xs -> showParen (d > 10) $ showString one1  . showsPrec 11 x . showChar ' ' . go 11 xs
    T x y xs -> showParen (d > 10) $ showString two1  . showsPrec 11 x . showChar ' ' 
                                                      . showsPrec 11 y . showChar ' ' . go 11 xs
  zero1 = zero ++ " "
  one1  = one ++ " "
  two1  = two ++ " "

--------------------------------------------------------------------------------
-- * 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 "Zero"
      next <- go xs
      newArrow node next
      return node

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

    T x y xs -> do
      node <- newNode True "Two"
      this1 <- newNode False (show x)
      this2 <- newNode False (show y)
      next <- go xs
      newArrow node this1
      newArrow node this2
      newArrow node next
      return node

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

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