{-# LANGUAGE LambdaCase       #-}
{-# LANGUAGE TypeApplications #-}

import qualified Control.Arrow                    as A
import           Control.Monad                    (replicateM)
import           Control.Monad.Trans.State.Strict (State, evalState)
import qualified Control.Monad.Trans.State.Strict as S
import           Data.List                        (elemIndex, sort)
import           Data.Maybe                       (fromJust)
import qualified System.Random.Shuffle            as RS
import           Test.Hspec
import           Test.Hspec.QuickCheck
import qualified Test.QuickCheck                  as QC

import           Algorithms.Random.Shuffle.Pure

main :: IO ()
main = hspec $ do
  describe "shuffle" $
    let gen = do
          xs <- (++)
            <$> (map QC.getPrintableString <$> replicateM 2 QC.arbitrary)
            <*> (map QC.getPrintableString . QC.getNonEmpty <$> QC.arbitrary)
          let len = length xs
          is <- mapM (\i -> QC.choose (0, i)) [len - 2, len - 3 .. 0]
          return (xs, is)
     in
      prop "behaves as a monadic version of System.Random.Shuffle.shuffle" $ QC.forAll gen $ \(xs, is) -> do
        let getRandR = const genSt
        evalState (shuffle getRandR xs) is `shouldBe` RS.shuffle xs is

  describe "sampleOne and sampleSplitOne" $ do
    it "returns Nothing given an empty list" $ do
      sampleOne undefined "" `shouldReturn` Nothing
      sampleSplitOne undefined "" `shouldReturn` Nothing

    prop "returns the only element given a singleton list" $ \x -> do
      sampleOne undefined [x :: Char] `shouldReturn` Just x
      sampleSplitOne undefined [x :: Char] `shouldReturn` Just (x, [])

    let gen = do
          len <- (+ 2) <$> QC.arbitrarySizedNatural
          let xs = take len ['0' ..]
          -- Indices list containing a zero somewhere
          -- NOTE: Why the length is two shorter than xs:
          --   - Minus one: Supplied by prepending by "(0 :)"
          --   - Minus one: getRandR is not actually called when picking the first element. See below.
          is <- (0 :) <$> QC.vectorOf (len - 2) (QC.arbitrarySizedNatural @Int)
          return (xs, is)
     in
      prop "sampleOne returns the element at the last index where the generator returns 0, and sampleSplitOne returns the other elements in addition" $
        QC.forAll gen $ \(xs, is) -> do
          let getRandR = const genSt
              -- NOTE: Because getRandR is not called when picking the first element,
              --       The actual index is shifted by one.
              chosen = xs !! (fromJust (elemRIndex 0 is) + 1)
              others = sort $ filter (/= chosen) xs

          evalState (sampleOne getRandR xs) is `shouldBe` Just chosen
          A.second sort <$> evalState (sampleSplitOne getRandR xs) is `shouldBe` Just (chosen, others)


genSt :: State [Int] Int
genSt =
  S.get >>=
    \case
      [] -> error "No more elements!"
      (hd : tl) -> do
        S.put tl
        return hd


elemRIndex :: Eq a => a -> [a] -> Maybe Int
elemRIndex x xs = abs . subtract (length xs - 1) <$> elemIndex x (reverse xs)