{-# OPTIONS_GHC -fno-warn-orphans #-}
module QuickCheck.Merit where
import Control.Monad (Monad(..), replicateM)
import Data.Bool
import Data.Eq (Eq(..))
import Data.Foldable (Foldable(..))
import Data.Function (($), (.))
import Data.Functor ((<$>))
import Data.Hashable (Hashable)
import Data.Int (Int)
import Data.List ((++), head, zip)
import Data.Ord (Ord(..))
import Data.Ratio (Rational)
import GHC.Exts (IsList(..))
import Majority.Merit
import Prelude (Enum(..), Num(..), Integral(..), Bounded(..), fromIntegral, undefined)
import Test.QuickCheck
import Test.Tasty
import Test.Tasty.QuickCheck
import Text.Show (Show(..))
import qualified Data.Map.Strict as Map

import QuickCheck.Utils
import Types

quickcheck :: TestTree
quickcheck =
	testGroup "Merit"
	 [ testProperty "arbitraryMerits" $ \(SameLength (Merit x::Merit SchoolGrade,Merit y::Merit SchoolGrade)) ->
		Map.keys x == Map.keys y &&
		sum x == sum y
	 ]

-- | @arbitraryMerits n@ arbitrarily generates 'n' lists of 'Merit'
-- for the same arbitrary grades,
-- and with the same total 'Share' of individual judgments.
arbitraryMerits :: forall g. (Bounded g, Enum g, Ord g) => Int -> Gen [Merit g]
arbitraryMerits n = sized $ \shareSum -> do
	minG <- choose (fromEnum(minBound::g), fromEnum(maxBound::g))
	maxG <- choose (minG, fromEnum(maxBound::g))
	let gs::[g] = toEnum minG`enumFromTo`toEnum maxG
	let lenGrades = maxG - minG + 1
	replicateM n $ do
		shares  <- resize shareSum $ arbitrarySizedPositiveRationalSum lenGrades
		shares' :: [Share] <- arbitraryPad (lenGrades - length shares) (return 0) shares
		return $ Merit $ fromList $ zip gs shares'

-- | @arbitrarySizedPositiveRationalSum maxLen@
-- arbitrarily chooses a list of 'length' at most 'maxLen',
-- containing positive 'Rational's summing up to 'sized'.
arbitrarySizedPositiveRationalSum :: Int -> Gen [Rational]
arbitrarySizedPositiveRationalSum maxLen = sized (go maxLen . fromIntegral)
	where
	go :: Int -> Rational -> Gen [Rational]
	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)

-- | @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'

instance
 (Arbitrary g, Bounded g, Enum g, Ord g, Show g) =>
 Arbitrary (Merit g) where
	arbitrary = head <$> arbitraryMerits 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 = arbitraryMerits (maxP - minP + 1)
		fromList . zip ps <$> ms
instance
 (Arbitrary g, Bounded g, Enum g, Ord g) =>
 Arbitrary (SameLength (Merit g, Merit g)) where
	arbitrary = do
		vs <- arbitraryMerits 2
		case vs of
		 [x,y] -> return $ SameLength (x,y)
		 _ -> undefined