{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module QuickCheck where

import Test.QuickCheck
import Test.Tasty
import Test.Tasty.QuickCheck

import Control.Monad (replicateM)
import qualified Data.Set as Set

import Hjugement
import Types

quickchecks :: TestTree
quickchecks =
	testGroup "QuickCheck"
	 [ testProperty "arbitraryJudgments" $ \(SameLength (x::[(G6,Count)],y)) ->
		let (gx, cx) = unzip x in
		let (gy, cy) = unzip y in
		gx == gy && sum cx == sum cy
	 , testGroup "Value"
		 [ testProperty "compare" $ \(SameLength (x::Value G6,y)) ->
			expandValue x`compare` expandValue y == x`compare`y
		 ]
	 {-
	 , testProperty "majorityGauge and majorityValue consistency" $
		 \(SameLength (x@(Merit xs)::Merit G6,y@(Merit ys))) ->
			not (all (==0) xs || all (==0) ys) ==>
			case majorityGauge x`compare`majorityGauge y of
			 LT -> majorityValue x < majorityValue y
			 GT -> majorityValue x > majorityValue y
			 EQ -> True
	 -}
	 ]

-- | Decompress a 'Value'.
expandValue :: Value a -> [a]
expandValue (Value []) = []
expandValue (Value ((x,c):xs)) = replicate c x ++ expandValue (Value xs)

-- | @arbitraryJudgments n@ arbitrarily generates 'n' lists of pairs of grade and 'Count'
-- for the same arbitrary grades,
-- and with the same total 'Count' of individual judgments.
arbitraryJudgments :: forall g. (Bounded g, Enum g) => Int -> Gen [[(g, Count)]]
arbitraryJudgments n = sized $ \s -> do
	minG <- choose (fromEnum(minBound::g), fromEnum(maxBound::g))
	maxG <- choose (minG, fromEnum(maxBound::g))
	let gs::[g] = toEnum minG`enumFromTo`toEnum maxG
	let lg = maxG - minG + 1
	replicateM n $ do
		cs  <- resize s $ arbitrarySizedNaturalSum lg
		cs' <- arbitraryPad (lg - length cs) (return 0) cs
		return $ zip gs cs'

-- | @arbitrarySizedNaturalSum maxLen@
-- arbitrarily chooses a list of 'length' at most 'maxLen',
-- containing 'Int's summing up to 'sized'.
arbitrarySizedNaturalSum :: Int -> Gen [Int]
arbitrarySizedNaturalSum maxLen = sized (go maxLen)
	where
	go :: Int -> Int -> Gen [Int]
	go len tot | len <= 0 = return []
	           | len == 1 = return [tot]
	           | tot <= 0 = return [tot]
	go len tot = do
		d <- choose (0, tot)
		(d:) <$> go (len-1) (tot - d)

-- | @arbitraryPad n pad xs@
-- arbitrarily grows list 'xs' with 'pad' elements
-- up to length 'n'.
arbitraryPad :: Int -> Gen a -> [a] -> Gen [a]
arbitraryPad n pad [] = replicateM n pad
arbitraryPad n pad xs = do
	(r, xs') <- go n xs
	if r > 0
	 then arbitraryPad r pad xs'
	 else return xs'
	where
	go r xs' | r <= 0 = return (0,xs')
	go r [] =  arbitrary >>= \b ->
		if b then pad >>= \p -> ((p:)<$>) <$> go (r-1) []
		     else return (r,[])
	go r (x:xs') = arbitrary >>= \b ->
		if b then pad >>= \p -> (([p,x]++)<$>) <$> go (r-1) xs'
		     else ((x:)<$>) <$> go r xs'

-- | Like 'nub', but O(n * log n).
nubList :: Ord a => [a] -> [a]
nubList = go Set.empty where
	go _ [] = []
	go s (x:xs) | x`Set.member`s = go s xs
	            | otherwise      = x:go (Set.insert x s) xs

instance Arbitrary G6 where
	arbitrary = arbitraryBoundedEnum
instance (Arbitrary g, Bounded g, Enum g, Ord g, Show g) => Arbitrary (Merit g) where
	arbitrary = fromList . head <$> arbitraryJudgments 1
	shrink (Merit m) = Merit <$> shrink m
instance
 ( Arbitrary p, Bounded p, Enum p, Ord p, Show p
 , Arbitrary g, Bounded g, Enum g, Ord g, Show g
 ) => Arbitrary (Merits p g) where
	arbitrary = do
		minP <- choose (fromEnum(minBound::p), fromEnum(maxBound::p))
		maxP <- choose (minP, fromEnum(maxBound::p))
		let ps = toEnum minP`enumFromTo`toEnum maxP
		let ms = (fromList <$>) <$> arbitraryJudgments (maxP - minP + 1)
		fromList . zip ps <$> ms
instance (Bounded g, Eq g, Integral g, Arbitrary g) => Arbitrary (Value g) where
	arbitrary = head . (Value <$>) <$> arbitraryJudgments 1
	shrink (Value vs) = Value <$> shrink vs

-- * Type 'SameLength'
newtype SameLength a = SameLength a
 deriving (Eq, Show)
instance Functor SameLength where
	fmap f (SameLength x) = SameLength (f x)
instance (Arbitrary g, Bounded g, Enum g) => Arbitrary (SameLength (Value g, Value g)) where
	arbitrary = do
		SameLength (x,y) <- arbitrary
		return $ SameLength (Value x, Value y)
instance (Arbitrary g, Bounded g, Enum g, Ord g, Show g) => Arbitrary (SameLength (Merit g, Merit g)) where
	arbitrary = do
		SameLength (x,y) <- arbitrary
		return $ SameLength (fromList x, fromList y)
instance (Arbitrary g, Bounded g, Enum g) => Arbitrary (SameLength ([(g,Count)], [(g,Count)])) where
	arbitrary = do
		vs <- arbitraryJudgments 2
		case vs of
		 [x,y] -> return $ SameLength (x,y)
		 _ -> undefined