{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS -fno-warn-orphans #-}
module QuickCheck.Election where

import Test.Tasty.QuickCheck
import qualified Data.List as List
import qualified Data.Text as Text
import Data.Eq (Eq(..))
import Data.Int (Int)
import Data.Ord (Ord(..))
import Prelude (undefined)

import Voting.Protocol

import Utils

-- Hardcoded limits to avoid keep a reasonable testing time.
maxArbitraryChoices :: Natural
maxArbitraryChoices = 5
maxArbitraryQuestions :: Natural
maxArbitraryQuestions = 5

quickcheck :: TestTree
quickcheck =
	testGroup "Election"
	 [ testGroup "verifyBallot" $
		 [ testElection @WeakParams
		 , testElection @BeleniosParams
		 ]
	 ]

testElection :: forall q. Params q => TestTree
testElection  =
	testGroup (paramsName @q)
	 [ testProperty "verifyBallot" $ \(seed, (elec::Election q) :> votes) ->
		isRight $ runExcept $
			(`evalStateT` mkStdGen seed) $ do
				-- ballotSecKey :: SecretKey q <- randomSecretKey
				ballot <- encryptBallot elec Nothing votes
				unless (verifyBallot elec ballot) $
					lift $ throwE $ ErrorBallot_Wrong
	 ]

instance PrimeField p => Arbitrary (F p) where
	arbitrary = choose (zero, F (fieldCharac @p) - one)
instance SubGroup q => Arbitrary (G q) where
	arbitrary = do
		m <- arbitrary
		return (groupGen ^ m)
instance SubGroup q => Arbitrary (E q) where
	arbitrary = E <$> choose (zero, groupOrder @q - one)
instance Arbitrary UUID where
	arbitrary = do
		seed <- arbitrary
		(`evalStateT` mkStdGen seed) $
			randomUUID
instance SubGroup q => Arbitrary (Proof q) where
	arbitrary = do
		proof_challenge <- arbitrary
		proof_response  <- arbitrary
		return Proof{..}
instance SubGroup q => Arbitrary (Question q) where
	arbitrary = do
		let question_text = "question"
		choices :: Natural <- choose (1, maxArbitraryChoices)
		let question_choices = [Text.pack ("c"<>show c) | c <- [1..choices]]
		question_mini <- fromNatural <$> choose (0, choices)
		question_maxi <- fromNatural <$> choose (nat question_mini, choices)
		return Question{..}
	shrink quest =
		[ quest{question_choices, question_mini, question_maxi}
		| question_choices <- shrinkList pure $ question_choices quest
		, let nChoices = fromNatural $ fromIntegral $ List.length question_choices
		, question_mini <- shrink $ min nChoices $ max zero $ question_mini quest
		, question_maxi <- shrink $ min nChoices $ max question_mini $ question_maxi quest
		]
instance SubGroup q => Arbitrary (Election q) where
	arbitrary = do
		let election_name = "election"
		let election_description = "description"
		election_PublicKey <- arbitrary
		election_questions <- resize (fromIntegral maxArbitraryQuestions) $ listOf1 arbitrary
		election_uuid <- arbitrary
		let election_hash = Hash ""
		return Election{..}
	shrink elec =
		[ elec{election_questions}
		| election_questions <- shrink $ election_questions elec
		]

-- | A type to declare an 'Arbitrary' instance where @b@ depends on @a@.
data (:>) a b = a :> b
 deriving (Eq,Show)
instance SubGroup q => Arbitrary (Question q :> [Bool]) where
	arbitrary = do
		quest@Question{..} <- arbitrary
		votes <- do
			let numChoices = List.length question_choices
			numTrue <- fromIntegral <$> choose (nat question_mini, nat question_maxi)
			rank <- choose (0, nCk numChoices numTrue - 1)
			return $ boolsOfCombin numChoices numTrue rank
		return (quest :> votes)
	shrink (quest :> votes) =
		[ q :> shrinkVotes q votes
		| q <- shrink quest
		]
instance SubGroup q => Arbitrary (Election q :> [[Bool]]) where
	arbitrary = do
		elec@Election{..} <- arbitrary
		votes <- forM election_questions $ \Question{..} -> do
			let numChoices = List.length question_choices
			numTrue <- fromIntegral <$> choose (nat question_mini, nat question_maxi)
			rank <- choose (0, nCk numChoices numTrue - 1)
			return $ boolsOfCombin numChoices numTrue rank
		return (elec :> votes)
	shrink (elec :> votes) =
		[ e :> List.zipWith shrinkVotes (election_questions e) votes
		| e <- shrink elec
		]

-- | @('boolsOfCombin' nBits nTrue rank)@ returns the 'rank'-th sequence of 'nBits'-bits possible
-- with 'nTrue' bits set at 'True' and @(nBits-nTrue)@ set at 'False'.
-- @rank@ has to be in @[0 .. 'nCk' nBits nTrue '-' 1]@
boolsOfCombin :: Int -> Int -> Int -> [Bool]
boolsOfCombin nBits nTrue rank
 | rank < 0 = undefined
 | nTrue == 0 = List.replicate nBits False
 | otherwise = go 0 cs <> List.replicate (nBits-List.last cs) False
	where
	cs = combinOfRank nBits nTrue rank
	go _d [] = []
	go curr (next:ns) =
		List.replicate (next-1-curr) False <> [True]
		 <> go next ns

-- | @('shrinkVotes' quest votes)@
-- returns a reduced version of the given @votes@
-- to fit the requirement of the given @quest@.
shrinkVotes :: Question q -> [Bool] -> [Bool]
shrinkVotes Question{..} votes =
	(\(nTrue, b) -> nTrue <= nat question_maxi && b)
	 <$> List.zip (countTrue votes) votes
	where
	countTrue :: [Bool] -> [Natural]
	countTrue = List.foldl' (\ns b -> if b then inc ns else ns) []
		where
		inc [] = [zero]
		inc (n:ns) = (n+one):n:ns