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
  , genNaturalMemoryState
  , genGrowableState
  , genElapsedDays
  , genRetrievability
  , 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 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))
    ]

-- | A whole memory state. The fast trace is usually a modest multiple of the
-- slow one, as it is in practice, but sometimes wildly different so that
-- nothing may assume a relationship between them.
genMemoryState :: Gen MemoryState
genMemoryState = do
  s <- genStability
  d <- genDifficulty
  sFast <-
    frequency
      [ (6, pure (clampStability (s * fastTraceRatio)))
      , (2, (\k -> clampStability (s * k)) <$> choose (0.05, 20))
      , (2, genStability)
      ]
  pure (MemoryState s d sFast)
  where
    clampStability = min stabilityMax . max stabilityMin

-- | A state whose fast trace sits at exactly the ratio to the slow one that
-- the model itself establishes for a new card.
--
-- Properties about what /more stability/ or /more difficulty/ does need this:
-- the FSRS-7 curve weights its two components by their own stabilities, so for
-- a wildly lopsided pair of traces those questions have no monotone answer —
-- raising difficulty slows the slow component but also shifts mixture weight
-- onto the fast one, and which effect wins depends on the ratio.
genNaturalMemoryState :: Gen MemoryState
genNaturalMemoryState = do
  s <- choose (log naturalStabilityMin, log stabilityMax)
  d <- genDifficulty
  pure (naturalState (exp s) d)

-- | A 'genNaturalMemoryState' together with a factor by which both of its
-- traces can be multiplied without either leaving its range — so that scaling
-- really does scale, instead of saturating one trace and skewing the ratio.
genGrowableState :: Gen (MemoryState, Double)
genGrowableState = do
  growth <- exp <$> choose (1.0e-6, 5)
  s <- choose (log naturalStabilityMin, log (stabilityMax / growth))
  d <- genDifficulty
  pure (naturalState (exp s) d, growth)

-- | The smallest slow stability whose fast trace is still above the floor, so
-- that the ratio is exact rather than clamped.
naturalStabilityMin :: Stability
naturalStabilityMin = stabilityMin / fastTraceRatio

naturalState :: Stability -> Difficulty -> MemoryState
naturalState s d = MemoryState s d (fastTraceRatio * s)

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

-- | Anything the forgetting curve could plausibly have returned, including the
-- extremes it is rescaled into.
genRetrievability :: Gen Retrievability
genRetrievability =
  frequency
    [ (1, pure retrievabilityFloor)
    , (1, pure (1 - retrievabilityFloor))
    , (8, choose (retrievabilityFloor, 1 - retrievabilityFloor))
    ]

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