module Rope
    ( Rope
    , balanced
    , length
    , substr
    , rebalance
    , toList
    , fromList
    ) where

import Prelude hiding (length)
import qualified Data.List as List
import Data.Array
import Data.Monoid
import qualified Data.IntMap as IM
import Control.Exception (assert)

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

data Rope a
    = Flat !(Array Int a)
    | Branch {-# UNPACK #-} !Length {-# UNPACK #-} !Depth (Rope a) (Rope a)

-- smart constructor
branch :: Rope a -> Rope a -> Rope a
branch x y = Branch (length x + length y) (max (depth x) (depth y) + 1) x y

type Length = Int
type Depth = Int

depth :: Rope a -> Depth
depth (Flat _) = 0
depth (Branch _ d _ _) = d

length :: Rope a -> Length
length (Flat a) = alength a
length (Branch len _ _ _) = len

balanced :: Rope a -> Bool
balanced t = length t >= fib (depth t + 2)

fib :: Int -> Int
fib n = fibl1 !! (n-1)

fibl1, fibl3 :: [Int]
fibl1 = 1 : 1 : zipWith (+) fibl1 (tail fibl1)
fibl3 = drop 2 fibl1

flatLim :: Length
flatLim = 1024

depthLim :: Depth
depthLim = last $ takeWhile (\n -> v (fib (n+2))) [0..]
  where v n = 0 <= n && n < maxBound

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

instance Monoid (Rope a) where
    mempty = Flat (listArray (0,-1) [])
    mappend a b
        | depth r > depthLim && not (balanced r) = rebalance r
        | otherwise = r
        where r = append' a b

instance Eq a => Eq (Rope a) where
    a == b = toList a == toList b

instance Show a => Show (Rope a) where
    showsPrec p r = showParen (p > app_prec) $
                    showString "fromList " . showList (toList r)
        where app_prec = 10

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

-- append without rebalance
append' :: Rope a -> Rope a -> Rope a
append' x y
    | length x == 0 = y
    | length y == 0 = x
append' (Branch _ _ l (Flat a1)) (Branch _ _ (Flat a2) r)
    | alength a1 + alength a2 <= flatLim =
        branch (branch l (Flat (aappend a1 a2))) r
append' (Branch _ _ l (Flat a1)) (Flat a2)
    | alength a1 + alength a2 <= flatLim = branch l (Flat (aappend a1 a2))
append' (Flat a1) (Branch _ _ (Flat a2) r)
    | alength a1 + alength a2 <= flatLim = branch (Flat (aappend a1 a2)) r
append' a b = branch a b

substr :: (Int,Int) -> Rope a -> Rope a
substr (beg, end) r
    | end < 0 || length r <= end = mempty
substr (beg, end) (Flat a) =
    let xs = [v | (i,v) <- assocs a, beg <= i, i <= end]
    in Flat $ listArray (0, List.length xs - 1) xs
substr (beg, end) (Branch len _ l r) = append' l' r'
  where
    r_beg = length l
    l_end = r_beg - 1
    l' = if beg <= 0 && l_end <= end
         then l
         else substr (beg, end) l
    r' = if beg <= r_beg && (len - 1) <= end
         then r
         else substr (beg - r_beg, end - r_beg) r

rebalance :: Rope a -> Rope a
rebalance r = r'
  where
    r' = foldl (\a b -> append' b a) mempty [v | (_,v) <- IM.toAscList bs]
    bs = foldl (f fibl3) IM.empty (balancedUnits r)
    f (x:xs) m b = assert (balanced b) $
        case IM.lookup x m of
        Nothing
            | length b < x -> IM.insert x b m
            | otherwise -> f xs m b
        Just b' ->
            f xs (IM.delete x m) (append' b' b)

balancedUnits :: Rope a -> [Rope a]
balancedUnits x | balanced x = [x]
balancedUnits (Flat a) = assert (alength a == 0) $ []
balancedUnits (Branch _ _ l r) = balancedUnits l ++ balancedUnits r

toList :: Rope a -> [a]
toList (Flat a) = elems a
toList (Branch _ _ l r) = toList l ++ toList r

fromList :: [a] -> Rope a
fromList = rebalance . foldl branch mempty . splits
  where
    splits xs =
        case List.splitAt flatLim xs of
        (ys, []) -> [f (List.length ys) ys]
        (ys, zs) -> f flatLim ys : splits zs
    f n ys = Flat (listArray (0, n - 1) ys)

-- ---------------------------------------------------------------------------
-- array utilities

type Array0 a = Array Int a

alength :: Array0 a -> Length
alength = rangeSize . bounds

aappend :: Array0 a -> Array0 a -> Array0 a
aappend a b = listArray (0, alength a + alength b - 1) (elems a ++ elems b)