Gen.hs

-- | Generators and comparison helpers shared by the test modules.
--
-- Everything here produces /valid/ inputs: parameter vectors inside the box
-- the upstream optimiser clips to, stabilities inside
-- @['stabilityMin', 'stabilityMax']@, difficulties inside @[1, 10]@ and
-- non-negative elapsed times. Properties that should hold for nonsense inputs
-- too say so explicitly.
module Test.FSRS.Gen
  ( -- * Generators
    genParameters
  , genRating
  , genStability
  , genDifficulty
  , genMemoryState
  , genElapsedDays
  , genDesiredRetention
  , genReviewHistory
  , genUTCTime
  , genScheduler

    -- * Approximate comparison
  , approxEqual
  , relativeError
  ) where

import Data.Time.Calendar (addDays, fromGregorian)
import Data.Time.Clock (UTCTime (..), secondsToDiffTime)
import Test.QuickCheck

import FSRS

-- | A parameter vector inside the valid box, biased towards the defaults.
genParameters :: Gen Parameters
genParameters =
  frequency
    [ (1, pure defaultParameters)
    , (3, genRandomParameters)
    ]

genRandomParameters :: Gen Parameters
genRandomParameters = do
  ws <- traverse choose parameterBounds
  -- `clampParameters` only re-establishes the three ordering constraints; the
  -- weights are already inside their individual bounds.
  case clampParameters ws of
    Right p -> pure p
    Left errs -> error ("genRandomParameters: " <> show errs)

genRating :: Gen Rating
genRating = elements allRatings

-- | Log-uniform across the whole legal range, with the endpoints thrown in.
genStability :: Gen Stability
genStability =
  frequency
    [ (1, pure stabilityMin)
    , (1, pure stabilityMax)
    , (1, pure 1)
    , (9, exp <$> choose (log stabilityMin, log stabilityMax))
    ]

genDifficulty :: Gen Difficulty
genDifficulty =
  frequency
    [ (1, pure difficultyMin)
    , (1, pure difficultyMax)
    , (8, choose (difficultyMin, difficultyMax))
    ]

genMemoryState :: Gen MemoryState
genMemoryState = MemoryState <$> genStability <*> genDifficulty

-- | Anything from "the same instant" to ten years, log-uniform in between.
genElapsedDays :: Gen Days
genElapsedDays =
  frequency
    [ (2, pure 0)
    , (8, exp <$> choose (log (1 / 86400), log 3650))
    ]

genDesiredRetention :: Gen Retrievability
genDesiredRetention =
  frequency
    [ (1, choose (0.5, 0.7))
    , (8, choose (0.7, 0.98))
    , (1, choose (0.98, 0.999))
    ]

genReviewHistory :: Gen [(Days, Rating)]
genReviewHistory = sized $ \n -> do
  k <- choose (0, min n 30)
  vectorOf k ((,) <$> genElapsedDays <*> genRating)

genUTCTime :: Gen UTCTime
genUTCTime = do
  day <- choose (0, 3650)
  seconds <- choose (0, 86399)
  pure (UTCTime (addDays day (fromGregorian 2026 1 1)) (secondsToDiffTime seconds))

genScheduler :: Gen Scheduler
genScheduler = do
  params <- genParameters
  retention <- genDesiredRetention
  learning <- genSteps
  relearning <- genSteps
  maxIvl <- choose (1, 36500)
  minIvl <- choose (1 / 86400, maxIvl)
  pure
    Scheduler
      { schedulerParameters = params
      , schedulerDesiredRetention = retention
      , schedulerLearningSteps = learning
      , schedulerRelearningSteps = relearning
      , schedulerMinimumInterval = minIvl
      , schedulerMaximumInterval = maxIvl
      }
  where
    genSteps = do
      k <- choose (0 :: Int, 4)
      vectorOf k (fromInteger <$> choose (60, 86400))

-- | @abs (a - b) <= atol + rtol * abs b@.
approxEqual
  :: Double
  -- ^ Absolute tolerance.
  -> Double
  -- ^ Relative tolerance.
  -> Double
  -- ^ Actual.
  -> Double
  -- ^ Expected.
  -> Bool
approxEqual atol rtol actual expected =
  abs (actual - expected) <= atol + rtol * abs expected

relativeError :: Double -> Double -> Double
relativeError actual expected
  | expected == 0 = abs actual
  | otherwise = abs (actual - expected) / abs expected