{-# 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 (first)
import Control.Monad (replicateM)
import Data.Hashable (Hashable)
import qualified Data.Map.Strict as Map
import Data.Ratio
import GHC.Exts (IsList(..))
import Prelude
import System.Random (Random(..))
import qualified Data.Set as Set

import Majority.Judgment
import Types

quickchecks :: TestTree
quickchecks =
	testGroup "QuickCheck"
	 [ testProperty "arbitraryMerits" $ \(SameLength (Merit x::Merit G6,Merit y::Merit G6)) ->
		Map.keys x == Map.keys y &&
		sum x == sum y
	 , 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 ms) = 
	let d = foldr lcm 1 (denominator . middleShare <$> ms) in
	go $ (\m -> (numerator (middleShare m) * d, lowGrade m, highGrade m)) <$> ms
	where
	go [] = []
	go ((s,l,h):xs) = concat (replicate (fromIntegral s) [l, h]) ++ go xs

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

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

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

instance Random Rational where
	randomR (minR, maxR) g =
		if d - b == 0
		then first (% b) $ randomR (a, c) g
		else first (bd2ac . nat2bd) $ randomR (0, toInteger (maxBound::Int)) g
		where
		a = numerator   minR
		b = denominator minR
		c = numerator   maxR
		d = denominator maxR
		nat2bd x = ((d - b) % toInteger (maxBound::Int)) * (x%1) + (b%1)
		bd2ac x = alpha * x + beta
			where
			alpha = (c-a) % (d-b)
			beta = (a%1) - alpha * (b%1)
		
	random = randomR (toInteger (minBound::Int)%1, toInteger (maxBound::Int)%1)

-- | @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 = 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 (Bounded g, Eq g, Integral g, Arbitrary g) => Arbitrary (MajorityValue g) where
	arbitrary = head . (majorityValue <$>) <$> arbitraryMerits 1
	shrink (MajorityValue vs) = MajorityValue <$> shrink vs
instance (Bounded g, Enum g) => Arbitrary (Middle g) where
	arbitrary = do
		lowG  <- choose (fromEnum(minBound::g), fromEnum(maxBound::g))
		highG <- choose (lowG, fromEnum(maxBound::g))
		share <- choose (0, 1)
		return $ Middle share (toEnum lowG) (toEnum highG)

-- * 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, Ord 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) => Arbitrary (SameLength (Merit g, Merit g)) where
	arbitrary = do
		vs <- arbitraryMerits 2
		case vs of
		 [x,y] -> return $ SameLength (x,y)
		 _ -> undefined
instance (Arbitrary g, Bounded g, Enum g, Ord g) => Arbitrary (SameLength ([Middle g], [Middle g])) where
	arbitrary = do
		SameLength (m0, m1) <- arbitrary
		return $ SameLength
		 ( unMajorityValue $ majorityValue m0
		 , unMajorityValue $ majorityValue m1 )