content: -- |
  -- Module        : Test_SDP
  -- Description   : This module contains property-based tests for the SDP model.
  -- Tool          : QuickCheck
  -- Key Components:
  -- - Arbitrary   : To enable random generation of test cases.
  -- - `dummyProb2`: Helper to create a probability distribution over a list of elements.
  
  module Test_SDP where
    
  import qualified Data.Map as Map
  
  import SDPCompute
  import Test.QuickCheck
  import Prob
  
  -- | Adjust this import based on the SDP to test. 
  import GHGCase as Case
  -- =====================================================================
  -- Arbitrary Instances
  -- =====================================================================
  
  instance Arbitrary Action where
    arbitrary = elements sdpGetActions
  
  instance Arbitrary State where 
    arbitrary = elements sdpStates
  
  
  -- | Specific implementation to test. Exchange depending on model for testing. 
  
  sdpInstance :: SDP State Action
  sdpInstance   = ghgcase
  sdpGetActions :: [Action]
  sdpGetActions = Case.getActions
  sdpNext :: Int -> State -> Action -> Prob State
  sdpNext       = Case.next
  sdpActions :: Int -> State -> [Action]
  sdpActions    = Case.actions
  sdpStates :: [State]
  sdpStates     = Case.getStates
  
  name :: String
  name = Case.moduleName
  
  -- =====================================================================
  -- Test val
  -- =====================================================================
  
  -- | Test that when the policy sequence is empty, the value function `val` returns 0.
  prop_valEmptyPolicy :: State -> Property
  prop_valEmptyPolicy x =
    forAll genValidInt $ \t -> 
      val sdpInstance t [] x == 0
  
  
  -- | Monotonicity of val with respect to policy sequence extension.
  prop_valMonotonic :: State -> Property
  prop_valMonotonic x =
    forAll genValidIntZero $ \t ->
    forAll genValidInt $ \n ->
      let ps = bi sdpInstance t n
          ps' = ps ++ take 1 ps  -- extend by repeating last or one more valid step
      in not (null ps) ==>
        val sdpInstance t ps x <= val sdpInstance t ps' x
  
  -- TODO, error checks? Import Control.Exception
  
  
  -- =====================================================================
  -- Test bi
  -- =====================================================================
  
  -- | Test that the length of the policy sequence generated by `bi` matches the specified horizon `n`.
  prop_biLength :: Property
  prop_biLength =
    forAll genValidInt $ \t -> 
    forAll genValidInt $ \n -> 
    length (bi sdpInstance t n) == n
  
  -- | Test that `bi` returns an empty policy sequence when n <= 0, regardless of timestep (value of t).
  prop_biEmptyPolicy :: Property
  prop_biEmptyPolicy =
    forAll genValidIntZero $ \t ->
      bi sdpInstance t 0 == []
  
  -- | All actions in the policy are feasible for the given state.
  prop_policyOnlyFeasibleActions :: Property
  prop_policyOnlyFeasibleActions =
    forAll genValidIntZero $ \t ->
    forAll genValidInt $ \n ->
      all (\p -> all (\(s, a) -> a `elem` sdpActions t s) (Map.toList p)) (bi sdpInstance t n)
  
  -- | Test that bi generates optimal policy sequences (by comparison with random policies)
  prop_biOptimal ::State -> Property
  prop_biOptimal x = 
    forAll genValidIntZero $ \t ->
    forAll genValidInt $ \n ->
    forAll (genPs t n) $ \ps' ->
      let ps = bi sdpInstance t n
      in val sdpInstance t ps' x <= val sdpInstance t ps x
  
  
  -- =====================================================================
  -- Test bestExt, worstExt
  -- =====================================================================
  
  -- | Test that the length of the list returned by `bestExt` matches the number of states.
  prop_bestExtWorstCorrectLength :: Property
  prop_bestExtWorstCorrectLength =
    forAll genValidInt $ \t ->
    forAll genValidInt $ \n ->
      let ps = bi sdpInstance t n
      in length (bestExt sdpInstance t ps)  == length sdpStates 
      && length (worstExt sdpInstance t ps) == length sdpStates
  
  -- | Test that the policy from `bestExt` is "better" (higher val) than that of `worstExt`.
  prop_bestExtGEQworst :: State -> Property
  prop_bestExtGEQworst s =
    forAll genValidInt $ \n ->
    forAll genValidInt $ \t ->
      let ps       = bi ghgcase t n
          bestPol  = bestExt sdpInstance t ps
          worstPol = worstExt sdpInstance t ps
          bestVal  = val sdpInstance t (bestPol : ps) s
          worstVal = val sdpInstance t (worstPol : ps) s
      in bestVal  >= worstVal
  
          
  -- =====================================================================
  -- Test best
  -- =====================================================================
  
  -- | Test that the action returned by `best` is one of the possible actions in that state at time t.
  prop_bestActionIsValid :: State -> Property
  prop_bestActionIsValid s =
    forAll genValidInt $ \t ->
    forAll genValidInt $ \n ->
      let (a, _) = best sdpInstance t n s
      in a `elem` (sdpActions t s)
  
  -- | Test that `best` returns the same value as computing `val` with the returned action and policy.
  prop_bestValMatchesVal :: State -> Property
  prop_bestValMatchesVal s =
    forAll genValidInt $ \t ->
    forAll genValidInt $ \n ->
      let (a, v) = best sdpInstance t n s
          ps = bi sdpInstance (t + 1) (n - 1)
          p = Map.singleton s a
          v' = val sdpInstance t (p : ps) s
      in v === v'
  
  -- | Test that the value returned by `best` is greater than or equal to any other possible single-action value.
  prop_bestValIsOptimal :: State -> Property
  prop_bestValIsOptimal s =
    forAll genValidInt $ \t ->
    forAll genValidInt $ \n ->
      let (_, vBest) = best sdpInstance t n s
          ps = bi sdpInstance (t + 1) (n - 1)
          otherActions = sdpActions t s
          altVals = [val sdpInstance t (Map.singleton s a' : ps) s | a' <- otherActions]
      in all (<= vBest) altVals
  
  -- =====================================================================
  -- Test meas
  -- =====================================================================
  
  -- | Test that the value returned from `mMeas` is larger or equal to 0.
  prop_valueLTOne :: State -> Property
  prop_valueLTOne x =
    forAll genValidInt $ \t ->
    forAll genValidInt $ \n ->
      mMeas sdpInstance t n x >= 0.0
  
  -- =====================================================================
  -- Test next
  -- =====================================================================
  
  -- | Test that probabilities returned by `next` sum to 1.
  prop_sumEqualToOne :: State -> Property
  prop_sumEqualToOne x = 
    forAll genValidInt $ \t ->
    forAll (genValidAction t x) $ \s ->
      sum (Prob.weights (sdpNext t x s)) >= 0.99 &&
      sum (Prob.weights (sdpNext t x s)) <= 1.01
  
  
  -- =====================================================================
  -- Test actions
  -- =====================================================================
  
  -- | the set of valid controls in any state has to be non-empty,  
  prop_actionsNotEmpty :: State -> Property
  prop_actionsNotEmpty x = 
    forAll genValidInt $ \t ->
      length (Case.actions t x) >= 1 
  
  
  -- =====================================================================
  -- Test Val
  -- =====================================================================
  
  -- | Total preorders on Val
  prop_valTotalPreorder :: Val -> Val -> Val -> Property
  prop_valTotalPreorder x y z =
    (x <= y || y <= x) &&    
    x <= x &&                
    (x <= y && y <= z) ==> x <= z 
  
  -- =====================================================================
  -- Test Runners
  -- =====================================================================
  
  -- | Run tests for `val`.
  testVal :: IO ()
  testVal = do
    quickCheck prop_valEmptyPolicy
    quickCheck prop_valMonotonic
  
  
  -- | Run tests for `bi`.
  testBi :: IO ()
  testBi = do
    quickCheck prop_biLength
    quickCheck prop_biEmptyPolicy
    quickCheck prop_biOptimal
    quickCheck prop_policyOnlyFeasibleActions
  
  -- | Run tests for `bestExt` and `worstExt`.
  testBestExt :: IO ()
  testBestExt = do
    quickCheck prop_bestExtGEQworst
    quickCheck prop_bestExtWorstCorrectLength
  
  -- | Run tests for `best`.
  testBest :: IO ()
  testBest = do
    quickCheck prop_bestActionIsValid
    quickCheck prop_bestValMatchesVal
    quickCheck prop_bestValIsOptimal
  
  -- | Run tests for `next`.
  testNext :: IO ()
  testNext = do
    quickCheck prop_sumEqualToOne
  
  -- | Run tests for `mMeas`.
  testmMeas :: IO ()
  testmMeas = do
    quickCheck prop_valueLTOne
  
  testValType :: IO()
  testValType = do
    quickCheck prop_valTotalPreorder
  
  testActions :: IO()
  testActions = do
    quickCheck prop_actionsNotEmpty
  
  -- | Run all tests.
  testAll :: IO ()
  testAll = do
    putStrLn $ "Currently running tests for: " ++ name
    testVal
    testBi
    testBestExt
    testBest
    testNext
    testmMeas
    testValType
    testActions
  
  -- =====================================================================
  -- Utility Functions
  -- =====================================================================
  
  -- | Run `bi` with the given parameters.
  runBi :: Int -> Int -> PolicySeq State Action
  runBi t n = bi sdpInstance t n
  
  -- | Generate a valid integer in the range [1, 7].
  genValidInt :: Gen Int
  genValidInt = choose (1, 7)
  
  -- | Generate a valid integer in the range [0, 7].
  genValidIntZero :: Gen Int
  genValidIntZero = choose (0, 7)
  
  -- | Generate a valid action for a given state and time.
  genValidAction :: Int -> State -> Gen Action
  genValidAction t x = elements (Case.actions t x)
  
  genPs :: Int -> Int -> Gen (PolicySeq State Action)
  genPs t n = do
    return $ bi sdpInstance t n