{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies      #-}

{- | Fair implementation of the 'Treap' data structure that uses random
generator for priorities.
-}

module Treap.Rand
       ( -- * Data structure
         RTreap (..)

         -- * Smart constructors
       , emptyWithGen
       , oneWithGen
       , empty
       , one

         -- * Query functions
       , size
       , at
       , query

       -- * Cuts and joins
       , splitAt
       , merge
       , take
       , drop
       , rotate

       -- * Modification functions
       , insert
       , delete

         -- * General purpose functions
       , withTreap
       , overTreap

         -- * Pretty printing functions
       , prettyPrint
       ) where

import Prelude hiding (drop, lookup, splitAt, take)

import Control.DeepSeq (NFData (..))
import Data.Coerce (Coercible)
import Data.Foldable (foldl')
import GHC.Exts (IsList (..))
import GHC.Generics (Generic)

import Treap.Measured (Measured (..))
import Treap.Pure (Priority (..), Size (..), Treap)

import qualified System.Random.Mersenne.Pure64 as Random
import qualified Treap.Pretty as Treap
import qualified Treap.Pure as Treap

-- $setup
-- >>> import Data.Monoid

----------------------------------------------------------------------------
-- Data structure and instances
----------------------------------------------------------------------------

{- | Specialized version of 'Treap' where priority is
generated by the stored random generator.
-}
data RTreap m a = RTreap
    { rTreapGen  :: !Random.PureMT
    , rTreapTree :: !(Treap m a)
    } deriving (Show, Generic, Foldable)

-- | \( O(n) \). This instance doesn't compare random generators inside trees.
instance (Eq m, Eq a) => Eq (RTreap m a) where
    (==) :: RTreap m a -> RTreap m a -> Bool
    RTreap _ t1 == RTreap _ t2 = t1 == t2

-- | \( O(1) \). Takes cached value from the root.
instance Monoid m => Measured m (RTreap m a) where
    measure :: RTreap m a -> m
    measure = withTreap measure
    {-# INLINE measure #-}

{- | Pure implementation of 'RTreap' construction functions. Uses
@'empty' :: RTreap k a@ as a starting point. Functions have the following
time complexity:

1. 'fromList': \( O(n\ \log \ n) \)
2. 'toList': \( O(n) \)

>>> prettyPrint $ fromList @(RTreap (Sum Int) Int) [1..5]
   5,15:2
      ╱╲
     ╱  ╲
    ╱    ╲
   ╱      ╲
1,1:1   3,12:4
          ╱╲
         ╱  ╲
        ╱    ╲
      1,3:3 1,5:5
-}
instance Measured m a => IsList (RTreap m a) where
    type Item (RTreap m a) = a

    fromList :: [a] -> RTreap m a
    fromList = foldl' (\t (i, a) -> insert i a t) empty . zip [0..]
    {-# INLINE fromList #-}

    toList :: RTreap m a -> [a]
    toList = map snd . toList . rTreapTree
    {-# INLINE toList #-}

instance (NFData m, NFData a) => NFData (RTreap m a) where
    rnf RTreap{..} = rnf rTreapTree `seq` ()

----------------------------------------------------------------------------
-- Smart constructors
----------------------------------------------------------------------------

defaultRandomGenerator :: Random.PureMT
defaultRandomGenerator = Random.pureMT 0

-- | \( O(1) \). Create empty 'RTreap' with given random generator.
emptyWithGen :: Random.PureMT -> RTreap m a
emptyWithGen gen = RTreap gen Treap.Empty
{-# INLINE emptyWithGen #-}

-- | \( O(1) \). Create empty 'RTreap' using random generator with seed @0@.
empty :: RTreap m a
empty = emptyWithGen defaultRandomGenerator
{-# INLINE empty #-}

-- | \( O(1) \). Create singleton 'RTreap' with given random generator.
oneWithGen :: Measured m a => Random.PureMT -> a -> RTreap m a
oneWithGen gen a =
    let (priority, newGen) = Random.randomWord64 gen
    in RTreap newGen $ Treap.one (Priority priority) a
{-# INLINE oneWithGen #-}

-- | \( O(1) \). Create singleton 'RTreap' using random generator with seed @0@.
one :: Measured m a => a -> RTreap m a
one = oneWithGen defaultRandomGenerator
{-# INLINE one #-}

----------------------------------------------------------------------------
-- Query functions
----------------------------------------------------------------------------

{- | \( O(1) \). Returns the size of the 'RTreap'.

__Properties:__

* \( \forall (t\ ::\ \mathrm{Treap}\ m\ a)\ .\ \mathrm{size}\ t \geqslant 0 \)
-}
size :: RTreap m a -> Int
size = unSize . withTreap Treap.size
{-# INLINE size #-}

-- | \( O(\log \ n) \). Lookup a value by a given key inside 'RTreap'.
at :: Int -> RTreap m a -> Maybe a
at i = withTreap $ Treap.at i
{-# INLINE at #-}

-- | \( O(\log \ n) \). Return value of monoidal accumulator on a segment @[l, r)@.
query :: forall m a . Measured m a => Int -> Int -> RTreap m a -> m
query l r = withTreap (Treap.query l r)
{-# INLINE query #-}

----------------------------------------------------------------------------
-- Cuts and joins
----------------------------------------------------------------------------

-- | \( O(\log \ n) \). Lifted to 'RTreap' version of 'Treap.splitAt'.
splitAt :: forall m a . Measured m a => Int -> RTreap m a -> (RTreap m a, RTreap m a)
splitAt i (RTreap gen t) = let (l, r) = Treap.splitAt i t in (RTreap gen l, RTreap gen r)
{-# INLINE splitAt #-}

-- | \( O(\log \ n) \). Lifted to 'RTreap' version of 'Treap.merge'.
merge :: Measured m a => RTreap m a -> RTreap m a -> RTreap m a
merge (RTreap gen t1) (RTreap _ t2) = RTreap gen (Treap.merge t1 t2)
{-# INLINE merge #-}

-- | \( O(\log \ n) \). Lifted to 'RTreap' version of 'Treap.take'.
take :: forall m a . Measured m a => Int -> RTreap m a -> RTreap m a
take n = overTreap (Treap.take n)
{-# INLINE take #-}

-- | \( O(\log \ n) \). Lifted to 'RTreap' version of 'Treap.drop'.
drop :: forall m a . Measured m a => Int -> RTreap m a -> RTreap m a
drop n = overTreap (Treap.drop n)
{-# INLINE drop #-}

-- | \( O(\log \ n) \). Lifted to 'RTreap' version of 'Treap.rotate'.
rotate :: forall m a . Measured m a => Int -> RTreap m a -> RTreap m a
rotate n = overTreap (Treap.rotate n)
{-# INLINE rotate #-}

----------------------------------------------------------------------------
-- Modification functions
----------------------------------------------------------------------------

-- | \( O(\log \ n) \). Insert a value into 'RTreap' by given key.
insert :: forall m a . Measured m a => Int -> a -> RTreap m a -> RTreap m a
insert i a (RTreap gen t) =
    let (priority, newGen) = Random.randomWord64 gen
    in RTreap newGen $ Treap.insert i (Priority priority) a t
{-# INLINE insert #-}

{- | \( O(\log \ n) \). Delete 'RTreap' node that contains given key. If there is no
such key, 'RTreap' remains unchanged.
-}
delete :: forall m a . Measured m a => Int -> RTreap m a -> RTreap m a
delete i (RTreap gen t) = RTreap gen $ Treap.delete i t
{-# INLINE delete #-}

----------------------------------------------------------------------------
-- Generic functions
----------------------------------------------------------------------------

-- | Lift a function that works with 'Treap' to 'RTreap'.
withTreap :: (Treap m a -> r) -> (RTreap m a -> r)
withTreap f = f . rTreapTree
{-# INLINE withTreap #-}

-- | Lift a function that works with 'Treap' to 'RTreap'.
overTreap :: (Treap m a -> Treap m a) -> (RTreap m a -> RTreap m a)
overTreap set t = t { rTreapTree = set $ rTreapTree t }
{-# INLINE overTreap #-}

----------------------------------------------------------------------------
-- Pretty printing functions
----------------------------------------------------------------------------

-- | Pretty prints 'RTreap' without printing random generator.
prettyPrint :: forall m a . (Coercible m a, Show a) => RTreap m a -> IO ()
prettyPrint = withTreap Treap.prettyPrint