{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_GHC -Wall #-}

module Main (main) where

import Data.Word
import Test.Tasty
import Test.Tasty.SmallCheck
import Test.SmallCheck.Series
import Control.Monad.ST
import qualified ArrayList as AL
import ArrayList (ArrayList)
import Data.Primitive
import Foreign.Storable
import Data.Foldable
import System.IO.Unsafe
import qualified Data.List.NonEmpty as NE
import qualified Data.List as L
import Data.Hashable
import Data.Bifunctor
import Control.Monad.Random.Strict hiding (fromList)

main :: IO ()
main = defaultMain (testGroup "arrayList" arraylistTests)

singletonPrimArray :: forall a. Prim a => a -> PrimArray a
singletonPrimArray x = runST sing where
  sing :: forall s. ST s (PrimArray a)
  sing = do
    arr <- newPrimArray 1
    writePrimArray arr 0 x
    unsafeFreezePrimArray arr

arrayListInsertions :: (Eq a, Show a, Prim a, Storable a) => [a] -> Either String String
arrayListInsertions xs = unsafePerformIO $ AL.with $ \a0 -> do
  a1 <- foldlM AL.pushR a0 xs
  (a2,ys) <- dumpList a1
  return $ (,) a2 $ if xs == ys
    then Right "good"
    else Left ("expected " ++ show xs ++ "but got " ++ show ys)

pushPop :: forall a. (Eq a, Show a, Prim a, Storable a) => [a] -> Either String String
pushPop xs = unsafePerformIO $ AL.with $ \a0 -> do
  a1 <- foldlM AL.pushR a0 xs
  let go :: AL.ArrayList a -> IO (AL.ArrayList a, [a])
      go al = do
        (al',m) <- AL.popL al
        case m of
          Nothing -> return (al',[])
          Just a -> fmap (second (a:)) (go al')
  (a2,ys) <- go a1
  return $ (,) a2 $ if xs == ys
    then Right "good"
    else Left $ "expected " ++ show xs ++ " but got " ++ show ys

arrayListDropWhile :: forall a. (Hashable a, Eq a, Show a, Prim a, Storable a) => [a] -> Either String String
arrayListDropWhile xs = unsafePerformIO $ AL.with $ \a0 ->
  case deterministicShuffle xs of
    [] -> return (a0, Right "good")
    x : _ -> do
     a1 <- foldlM AL.pushR a0 xs
     (a2,_) <- AL.dropWhileL a1 (\y -> return (y /= x))
     (a3,ys) <- dumpList a2
     let expected = L.dropWhile (/= x) xs
     return $ (,) a3 $ if expected == ys
       then Right "good"
       else Left ("expected " ++ show expected ++ " but got " ++ show ys ++ " using pivot of " ++ show x)
  
arrayListInsertArray :: forall a. (Hashable a, Eq a, Show a, Prim a, Storable a)
  => [a] -> Either String String
arrayListInsertArray xs = unsafePerformIO $ AL.with $ \a0 -> do
  a1 <- foldlM AL.pushArrayR a0 (map singletonPrimArray xs)
  let go :: AL.ArrayList a -> IO (AL.ArrayList a, [a])
      go al = do
        (al',m) <- AL.popL al
        case m of
          Nothing -> return (al',[])
          Just a -> fmap (second (a:)) (go al')
  (a2,ys) <- go a1
  return $ (,) a2 $ if xs == ys
    then Right "good"
    else Left $ "expected " ++ show xs ++ " but got " ++ show ys
  
arrayListInsertBigArray :: forall a. (Hashable a, Eq a, Show a, Prim a, Storable a)
  => [a] -> Either String String
arrayListInsertBigArray xs = unsafePerformIO $ AL.with $ \a0 -> do
  a1 <- AL.pushArrayR a0 (fromList xs)
  let go :: AL.ArrayList a -> IO (AL.ArrayList a, [a])
      go al = do
        (al',m) <- AL.popL al
        case m of
          Nothing -> return (al',[])
          Just a -> fmap (second (a:)) (go al')
  (a2,ys) <- go a1
  return $ (,) a2 $ if xs == ys
    then Right "good"
    else Left $ "expected " ++ show xs ++ " but got " ++ show ys

arrayListInsertArrays :: forall a. (Hashable a, Eq a, Show a, Prim a, Storable a)
  => [a] -> Either String String
arrayListInsertArrays xs = unsafePerformIO $ AL.with $ \a0 -> do
  a1 <- AL.pushArrayR a0 (fromList xs)
  a2 <- AL.pushArrayR a1 (fromList xs)
  let go :: AL.ArrayList a -> IO (AL.ArrayList a, [a])
      go al = do
        (al',m) <- AL.popL al
        case m of
          Nothing -> return (al',[])
          Just a -> fmap (second (a:)) (go al')
  (a3,zs) <- go a2
  return $ (,) a3 $ if zs == (xs ++ xs)
    then Right "good"
    else Left $ "expected " ++ show (xs ++ xs) ++ " but got " ++ show zs

-- | This should not be used in production code.
dumpList :: (Prim a, Storable a) => ArrayList a -> IO (ArrayList a, [a])
dumpList xs = do
  let len = AL.size xs
  marr <- newPrimArray len
  newXs <- AL.dump xs marr 0
  arr <- unsafeFreezePrimArray marr
  return (newXs,primArrayToListN len arr)

deterministicShuffle :: Hashable a => [a] -> [a]
deterministicShuffle xs = evalRand (shuffle xs) (mkStdGen (hash xs))

shuffle :: [a] -> Rand StdGen [a]
shuffle [] = return []
shuffle xs = do
  randomPosition <- getRandomR (0, length xs - 1)
  let (left, (a:right)) = L.splitAt randomPosition xs
  fmap (a:) (shuffle (left ++ right))

primArrayToListN :: forall a. Prim a => Int -> PrimArray a -> [a]
primArrayToListN len arr = go 0
  where
  go :: Int -> [a]
  go !ix = if ix < len
    then indexPrimArray arr ix : go (ix + 1)
    else []

arraylistTests :: [TestTree]
arraylistTests =
  [ testPropDepth 10 "arraylist inserts followed by dump (short)" (over word16Series arrayListInsertions)
  , testPropDepth 150 "arraylist inserts followed by dump (long)" (over word32Series arrayListInsertions)
  , testPropDepth 150 "arraylist inserts followed by repeated pop (long)" (over word32Series pushPop)
  , testPropDepth 50 "arraylist dropWhile" (over word32Series arrayListDropWhile)
  , testPropDepth 50 "insert array" (over word32Series arrayListInsertArray)
  , testPropDepth 100 "insert big array" (over word32Series arrayListInsertBigArray)
  , testPropDepth 100 "insert big arrays" (over word32Series arrayListInsertArrays)
--  , testPropDepth 150 "arraylist push, pop, twice (long)" (over word32Series pushPopTwice)
  ]

testPropDepth :: Testable IO a => Int -> String -> a -> TestTree
testPropDepth n name = localOption (SmallCheckDepth n) . testProperty name

scanSeries :: forall m a. (a -> [a]) -> a -> Series m [a]
scanSeries f x0 = generate $ \n ->
  map toList $ concat $ take n $ iterate
    (\ys -> ys >>= \xs@(x NE.:| _) -> f x >>= \z -> [z NE.:| (toList xs)])
    [x0 NE.:| []]

word16Series :: Series m [Word16]
word16Series = (scanSeries (\n -> [n + 89, n + 71]) 0)

word32Series :: Series m [Word32]
word32Series = (scanSeries (\n -> [n + 73]) 0)