{-# LANGUAGE TypeOperators, TypeFamilies, FlexibleContexts, CPP #-}
module NoSlow.Util.Computation (
  DeepSeq(..), TestData(..), -- Computation(..),

  ListLike(..), testList,

  Len(..), WithSqrtLen(..), Sorted(..), ParenTree(..), Edges(..),
  ConnectedEdges(..),
  -- Nil(..), (:>)(..),

  generateData
) where

import Statistics.RandomVariate
import qualified Data.Array.ST as STA
import Control.Monad.ST ( ST, runST )
import Control.Monad
import Data.List ( sort )


class DeepSeq a where
  deepSeq :: a -> b -> b
  deepSeq = seq

  -- NOTE: The dummy method is necessary because 6.13 doesn't optimise newtype
  -- DFuns correctly. See http://hackage.haskell.org/trac/ghc/ticket/3772.
  deepSeqDummy :: a
  deepSeqDummy = undefined

instance DeepSeq Int
instance DeepSeq Float
instance DeepSeq Double where
  deepSeq = seq
instance DeepSeq Bool

instance (DeepSeq a, DeepSeq b) => DeepSeq (a,b) where
  {-# INLINE deepSeq #-}
  deepSeq (a,b) x = deepSeq a $ deepSeq b x

instance (DeepSeq a, DeepSeq b, DeepSeq c) => DeepSeq (a,b,c) where
  {-# INLINE deepSeq #-}
  deepSeq (a,b,c) x = deepSeq a $ deepSeq b $ deepSeq c x

instance (DeepSeq a, DeepSeq b, DeepSeq c, DeepSeq d) => DeepSeq (a,b,c,d) where
  {-# INLINE deepSeq #-}
  deepSeq (a,b,c,d) x = deepSeq a $ deepSeq b $ deepSeq c $ deepSeq d x

instance (DeepSeq a, DeepSeq b, DeepSeq c, DeepSeq d, DeepSeq e)
                => DeepSeq (a,b,c,d,e) where
  {-# INLINE deepSeq #-}
  deepSeq (a,b,c,d,e) x = deepSeq a $ deepSeq b $ deepSeq c $ deepSeq d
                        $ deepSeq e x

instance DeepSeq a => DeepSeq [a] where
  -- NOTE: We can't use foldr deepSeq b xs here because GHC 6.12 won't
  -- optimise it
  {-# INLINE deepSeq #-}
  deepSeq xs b = go xs
    where
      go [] = b
      go (x:xs) = x `deepSeq` go xs

class DeepSeq a => TestData a where
  testData   :: Int -> Gen s -> ST s a

generateData :: TestData a => Int -> a
generateData n = runST (testData n =<< create)

instance TestData Int where
  testData _ = uniform

instance TestData Float where
  testData _ = uniform

instance TestData Double where
  testData _ = uniform

instance (TestData a, TestData b) => TestData (a,b) where
  testData n g = liftM2 (,) (testData n g) (testData n g)

instance (TestData a, TestData b, TestData c) => TestData (a,b,c) where
  testData n g = liftM3 (,,) (testData n g) (testData n g) (testData n g)

instance (TestData a, TestData b, TestData c, TestData d)
                => TestData (a,b,c,d) where
  testData n g = liftM4 (,,,) (testData n g) (testData n g) (testData n g)
                              (testData n g)

instance (TestData a, TestData b, TestData c, TestData d, TestData e)
                => TestData (a,b,c,d,e) where
  testData n g = liftM5 (,,,,) (testData n g) (testData n g) (testData n g)
                               (testData n g) (testData n g)

instance TestData a => TestData [a] where
  testData n g = sequence $ replicate n $ testData n g


newtype Len = Len Int

instance DeepSeq Len

instance TestData Len where
  testData n _ = return (Len n)


newtype WithSqrtLen a = WithSqrtLen a

instance DeepSeq a => DeepSeq (WithSqrtLen a) where
  deepSeq (WithSqrtLen xs) = deepSeq xs

instance TestData a => TestData (WithSqrtLen a) where
  testData n g = WithSqrtLen `fmap` testData n' g
    where
      n' = floor $ sqrt $ fromIntegral n


class ListLike l a where
  fromList :: [a] -> l a

instance ListLike [] a where
  fromList = id

testList :: (TestData a, ListLike l a) => Int -> Gen s -> ST s (l a)
testList n g = fromList `fmap` testData n g

newtype Sorted l a = Sorted (l a)

instance DeepSeq (l a) => DeepSeq (Sorted l a) where
  deepSeq (Sorted xs) y = deepSeq xs y

instance (Ord a, TestData a, DeepSeq (l a), ListLike l a)
        => TestData (Sorted l a) where
  testData n g = do
                   xs <- testData n g
                   return $ Sorted $ fromList $ sort xs


parenTree :: Int -> ([Int],[Int])
parenTree n = case go ([],[]) 0 (if even n then n else n+1) of
               (ls,rs) -> (reverse ls, reverse rs)
  where
    go (ls,rs) i j = case j-i of
                       0 -> (ls,rs)
                       2 -> (ls',rs')
                       d -> let k = ((d-2) `div` 4) * 2
                            in
                            go (go (ls',rs') (i+1) (i+1+k)) (i+1+k) (j-1)
      where
        ls' = i:ls
        rs' = j-1:rs

data ParenTree l a = ParenTree (l a) (l a)

instance DeepSeq (l a) => DeepSeq (ParenTree l a) where
  deepSeq (ParenTree t u) x = deepSeq t (deepSeq u x)

instance (DeepSeq (l Int), ListLike l Int) => TestData (ParenTree l Int) where
  testData n _ = case parenTree n of
                   (ls,rs) -> return $ ParenTree (fromList ls) (fromList rs)


randomGraph :: Int -> Int -> Gen s -> ST s ([Int], [Int])
randomGraph n e g
  = do
      arr <- STA.newArray (0,n-1) [] :: ST s (STA.STArray s Int [Int])
      addRandomEdges n g arr e
      xs <- STA.getAssocs arr
      return $ unzip $ [(i,j) | (i,js) <- xs, j <- js ]

randomConnectedGraph :: Int -> Int  -> Gen s -> ST s ([Int], [Int])
randomConnectedGraph n e g
  = do
      arr <- STA.newListArray (0,n-1) $ [[n] | n <- [1 .. n-1]] ++ [[]]
                 :: ST s (STA.STArray s Int [Int])
      addRandomEdges n g arr e
      xs <- STA.getAssocs arr
      return $ unzip $ [(i,j) | (i,js) <- xs, j <- js ]


addRandomEdges :: Int -> Gen s -> STA.STArray s Int [Int] -> Int -> ST s ()
addRandomEdges n g arr = fill
  where
    fill 0 = return ()
    fill e
      = do
          m <- random_index
          n <- random_index
          let lo = min m n
              hi = max m n
          ns <- STA.readArray arr lo
          if lo == hi || hi `elem` ns
            then fill e
            else do
                   STA.writeArray arr lo (hi:ns)
                   fill (e-1)

    random_index = do
                     x <- uniform g
                     let i = floor ((x::Double) * toEnum n)
                     if i == n then return 0 else return i


data Edges l a = Edges Int (l a) (l a)

instance DeepSeq (l a) => DeepSeq (Edges l a) where
  deepSeq (Edges n t u) x = n `seq` deepSeq t (deepSeq u x)

instance (DeepSeq (l Int), ListLike l Int) => TestData (Edges l Int) where
  testData n g = do
                   (xs,ys) <- randomGraph nodes n g
                   return $ Edges nodes (fromList xs) (fromList ys)
    where
      nodes = n `div` 10


data ConnectedEdges l a = ConnectedEdges Int (l a) (l a)

instance DeepSeq (l a) => DeepSeq (ConnectedEdges l a) where
  deepSeq (ConnectedEdges n t u) x = n `seq` deepSeq t (deepSeq u x)

instance (DeepSeq (l Int), ListLike l Int)
                => TestData (ConnectedEdges l Int) where
  testData n g = do
                   (xs,ys) <- randomGraph nodes n g
                   return $ ConnectedEdges nodes (fromList xs) (fromList ys)
    where
      nodes = n `div` 10