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

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

import Control.Arrow (second)
import Control.Monad (replicateM)
import Data.Hashable (Hashable)
import Data.Ratio
import GHC.Exts (IsList(..))
import Prelude
import qualified Data.Set as Set

import Hjugement
import Types

quickchecks :: TestTree
quickchecks =
	testGroup "QuickCheck"
	 [ testProperty "arbitraryJudgments" $ \(SameLength (x::[(G6,Share)],y)) ->
		let (gx, cx) = unzip x in
		let (gy, cy) = unzip y in
		gx == gy && sum cx == sum cy
	 , testGroup "MajorityValue"
		 [ testProperty "compare" $ \(SameLength (x::MajorityValue 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 'MajorityValue'.
expandValue :: MajorityValue a -> [a]
expandValue (MajorityValue as) = 
	let m = foldr lcm 1 (denominator . snd <$> as) in
	go $ second (\s -> numerator s * m) <$> as
	where
	go [] = []
	go ((x,c):xs) = replicate (fromIntegral c) x ++ go xs

-- | @arbitraryJudgments n@ arbitrarily generates 'n' lists of pairs of grade and 'Share'
-- for the same arbitrary grades,
-- and with the same total 'Share' of individual judgments.
arbitraryJudgments :: forall g. (Bounded g, Enum g) => Int -> Gen [[(g, Share)]]
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 $ fromIntegral <$> 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 :: (Num i, Integral i) => i -> Gen a -> [a] -> Gen [a]
arbitraryPad n pad [] = replicateM (fromIntegral 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 c, Bounded c, Enum c, Eq c, Hashable c, Show c
 , Arbitrary g, Bounded g, Enum g, Ord g, Show g
 ) => Arbitrary (MeritByChoice c g) where
	arbitrary = do
		minP <- choose (fromEnum(minBound::c), fromEnum(maxBound::c))
		maxP <- choose (minP, fromEnum(maxBound::c))
		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 (MajorityValue g) where
	arbitrary = head . (MajorityValue <$>) <$> arbitraryJudgments 1
	shrink (MajorityValue vs) = MajorityValue <$> 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 (MajorityValue g, MajorityValue g)) where
	arbitrary = do
		SameLength (x,y) <- arbitrary
		return $ SameLength (MajorityValue x, MajorityValue 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,Share)], [(g,Share)])) where
	arbitrary = do
		vs <- arbitraryJudgments 2
		case vs of
		 [x,y] -> return $ SameLength (x,y)
		 _ -> undefined