{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-} -- for Reifies instances
{-# OPTIONS -fno-warn-orphans #-}
module QuickCheck.Election where

import Data.Eq (Eq(..))
import Data.Int (Int)
import Data.Maybe (fromJust)
import Data.Ord (Ord(..))
import GHC.Natural (minusNaturalMaybe)
import Prelude (undefined)
import Test.Tasty.QuickCheck
import qualified Data.List as List
import qualified Data.Text as Text

import Voting.Protocol

import Utils

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

quickcheck :: TestTree
quickcheck =
	testGroup "Election"
	 [ testGroup "verifyBallot" $
		 [ testElection weakFFC
		 , testElection beleniosFFC
		 ]
	 ]

testElection :: FFC -> TestTree
testElection ffc =
	reify ffc $ \(Proxy::Proxy c) ->
	testGroup (Text.unpack $ ffc_name ffc)
	 [ testProperty "verifyBallot" $ \(seed, (elec::Election c) :> votes) ->
		isRight $ runExcept $
			(`evalStateT` mkStdGen seed) $ do
				-- ballotSecKey :: SecretKey c <- randomSecretKey
				ballot <- encryptBallot elec Nothing votes
				unless (verifyBallot elec ballot) $
					lift $ throwE $ ErrorBallot_Wrong
	 ]

instance Reifies c FFC => Arbitrary (F c) where
	arbitrary = F <$> choose (zero, fromJust $ fieldCharac @c `minusNaturalMaybe` one)
instance Reifies c FFC => Arbitrary (G c) where
	arbitrary = do
		m <- arbitrary
		return (groupGen ^ m)
instance Reifies c FFC => Arbitrary (E c) where
	arbitrary = E <$> choose (zero, fromJust $ groupOrder @c `minusNaturalMaybe` one)
instance Arbitrary UUID where
	arbitrary = do
		seed <- arbitrary
		(`evalStateT` mkStdGen seed) $
			randomUUID
instance Reifies c FFC => Arbitrary (Proof c) where
	arbitrary = do
		proof_challenge <- arbitrary
		proof_response  <- arbitrary
		return Proof{..}
instance Arbitrary Question 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 <- choose (0, choices)
		question_maxi <- 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 = fromIntegral $ List.length question_choices
		, question_mini <- shrinkIntegral $ min nChoices $ max zero $ question_mini quest
		, question_maxi <- shrinkIntegral $ min nChoices $ max question_mini $ question_maxi quest
		]
instance Reifies c FFC => Arbitrary (Election c) where
	arbitrary = do
		let election_name = "election"
		let election_description = "description"
		election_crypto <- arbitrary
		election_questions <- resize (fromIntegral maxArbitraryQuestions) $ listOf1 arbitrary
		election_uuid <- arbitrary
		let elec =
			Election
			 { election_hash = hashElection elec
			 , ..
			 }
		return elec
	shrink elec =
		[ elec{election_questions}
		| election_questions <- shrink $ election_questions elec
		]
instance Reifies c FFC => Arbitrary (ElectionCrypto c) where
	arbitrary = do
		let electionCrypto_FFC_params = reflect (Proxy::Proxy c)
		electionCrypto_FFC_PublicKey <- arbitrary
		return ElectionCrypto_FFC{..}

-- | A type to declare an 'Arbitrary' instance where @b@ depends on @a@.
data (:>) a b = a :> b
 deriving (Eq,Show)
instance Arbitrary (Question :> [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 Reifies c FFC => Arbitrary (Election c :> [[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 -> [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