{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

module Support where

import Control.Monad.ST
import Data.Bit
import qualified Data.Bit.ThreadSafe as TS
import Data.Bits
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Generic.New as N
import qualified Data.Vector.Unboxed as U
import Test.QuickCheck.Classes
import Test.Tasty
import Test.Tasty.QuickCheck

instance Arbitrary Bit where
  arbitrary = Bit <$> arbitrary
  shrink    = fmap Bit . shrink . unBit

instance CoArbitrary Bit where
  coarbitrary = coarbitrary . unBit

instance Function Bit where
  function f = functionMap unBit Bit f

instance Arbitrary TS.Bit where
  arbitrary = TS.Bit <$> arbitrary
  shrink    = fmap TS.Bit . shrink . TS.unBit

instance CoArbitrary TS.Bit where
  coarbitrary = coarbitrary . TS.unBit

instance Function TS.Bit where
  function f = functionMap TS.unBit TS.Bit f

instance (Arbitrary a, U.Unbox a) => Arbitrary (U.Vector a) where
  arbitrary = (\v -> runST (N.run (v :: N.New U.Vector a) >>= U.freeze)) <$> arbitrary
  shrink v = let len = U.length v in
    [ U.slice s l v
    | s <- [0 .. len - 1]
    , l <- [0 .. len - s]
    , (s, l) /= (0, len)
    ]

instance Arbitrary F2Poly where
  arbitrary = toF2Poly <$> arbitrary
  shrink v = toF2Poly <$> shrink (unF2Poly v)

instance (Show (v a), V.Vector v a) => Show (N.New v a) where
  showsPrec p = showsPrec p . V.new

newFromList :: forall a v . V.Vector v a => [a] -> N.New v a
newFromList xs = N.create (V.thaw (V.fromList xs :: v a))

-- this instance is designed to make sure that the arbitrary vectors we work with are not all nicely aligned; we need to deal with cases where the vector is a weird slice of some other vector.
instance (V.Vector v a, Arbitrary a) => Arbitrary (N.New v a) where
  arbitrary = frequency
    [ (10, newFromList <$> arbitrary)
    , (1 , N.drop <$> arbitrary <*> arbitrary)
    , (1 , N.take <$> arbitrary <*> arbitrary)
    , (1 , slice <$> arbitrary <*> arbitrary <*> arbitrary)
    ]
   where
    slice s n = N.apply
      $ \v -> let (s', n') = trimSlice s n (M.length v) in M.slice s' n' v
  shrink v =
    [ N.slice s l v
    | s <- [0 .. len - 1]
    , l <- [0 .. len - s]
    , (s, l) /= (0, len)
    ]
    where len = runST (M.length <$> N.run v)

trimSlice :: Integral a => a -> a -> a -> (a, a)
trimSlice s n l = (s', n')
 where
  s' | l == 0    = 0
     | otherwise = s `mod` l
  n' | s' == 0   = 0
     | otherwise = n `mod` (l - s')

sliceList :: Int -> Int -> [a] -> [a]
sliceList s n = take n . drop s

wordSize :: Int
wordSize = finiteBitSize (0 :: Word)

packBitsToWord :: [Bit] -> (Word, [Bit])
packBitsToWord = loop 0 0
 where
  loop _ w [] = (w, [])
  loop i w (x : xs)
    | i >= wordSize = (w, x : xs)
    | otherwise     = loop (i + 1) (if unBit x then setBit w i else w) xs

readWordL :: [Bit] -> Int -> Word
readWordL xs 0 = fst (packBitsToWord xs)
readWordL xs n = readWordL (drop n xs) 0

wordToBitList :: Word -> [Bit]
wordToBitList w = [ Bit (testBit w i) | i <- [0 .. wordSize - 1] ]

writeWordL :: [Bit] -> Int -> Word -> [Bit]
writeWordL xs 0 w = zipWith const (wordToBitList w) xs ++ drop wordSize xs
writeWordL xs n w = pre ++ writeWordL post 0 w
  where (pre, post) = splitAt n xs

prop_writeWordL_preserves_length :: [Bit] -> NonNegative Int -> Word -> Property
prop_writeWordL_preserves_length xs (NonNegative n) w =
  length (writeWordL xs n w) === length xs

prop_writeWordL_preserves_prefix :: [Bit] -> NonNegative Int -> Word -> Property
prop_writeWordL_preserves_prefix xs (NonNegative n) w =
  take n (writeWordL xs n w) === take n xs

prop_writeWordL_preserves_suffix :: [Bit] -> NonNegative Int -> Word -> Property
prop_writeWordL_preserves_suffix xs (NonNegative n) w =
  drop (n + wordSize) (writeWordL xs n w) === drop (n + wordSize) xs

prop_writeWordL_readWordL :: [Bit] -> Int -> Property
prop_writeWordL_readWordL xs n = writeWordL xs n (readWordL xs n) === xs

withNonEmptyMVec
  :: (Eq t, Show t)
  => (U.Vector Bit -> t)
  -> (forall s . U.MVector s Bit -> ST s t)
  -> Property
withNonEmptyMVec f g = forAll arbitrary $ \xs ->
  let xs' = V.new xs in not (U.null xs') ==> f xs' === runST (N.run xs >>= g)

tenTimesLess :: TestTree -> TestTree
tenTimesLess = adjustOption $
  \(QuickCheckTests n) -> QuickCheckTests (max 100 (n `div` 10))

lawsToTest :: Laws -> TestTree
lawsToTest (Laws name props) =
  testGroup name $ map (uncurry testProperty) props