-----------------------------------------------------------------------------
-- Copyright 2020, Ideas project team. This file is distributed under the
-- terms of the Apache License 2.0. For more information, see the files
-- "LICENSE.txt" and "NOTICE.txt", which are included in the distribution.
-----------------------------------------------------------------------------
module Domain.Hypothesis.Strategies (hypothesisStrategy) where

import Domain.Hypothesis.Rules
import Domain.Statistics.Rules
import Domain.Statistics.ComponentSet
import Ideas.Common.Strategy.Combinators hiding (not, while)
import Ideas.Common.Library hiding (while)
import Prelude hiding (until, sequence, repeat)


----------------------------------------------------------
-- Template for the hypothesis testing strategy

hypothesisStrategy :: LabeledStrategy ComponentSet
hypothesisStrategy = label "Hypothesis testing" $
   label "Preparation" (whileNotReady $ choice $
      [ addHypothesesRule, addH0FromHARule, addH0FromHAEqualSignRule, addHARule
      , addHypothesesChiSquaredRule
      , addAlphaRule, determineSided, chooseTTestRule
      , chooseTTestTwoRule, chooseTTestPairedRule, chooseZTestRule
      , chooseRPearsonRule, chooseAnovaRule, chooseChiSquaredRule
      ] ++ sampleStatistics)
   .*.
   check (\cs -> all (derived cs `contains`) [NullHypothesis, AlternativeHypothesis])
   .*.
   label "Computation" (whileNotReady $
      (check (\cs -> derived cs `doesNotContain` TestValue) .*. 
         (addTestFormulaRule .|. choice sampleStatistics))
      .|. (check allowCriticalRoute .*. choice
         [ addTestValueRule, addRejectionRule
         , lookupZValueRule, lookupTValueRule, lookupRValueRule, lookupFValueRule, lookupChiValueRule
         ])
      .|. (check allowPValueRoute .*. choice
         [ computePValueZTest, computePValueTTest
         ])
      )
   .*.
   check (\cs -> derived cs `contains` TestValue &&
                 derived cs `contains` Critical || derived cs `contains` PValue)
   .*.
   label "Conclusion" (
      whileNotReady (criticalConclusionRule .|. addConclusionPValueRule)
      .*.
      (hypothesesConclusionCriticalRule .|. hypothesesConclusionPValueRule))
 where
   sampleStatistics = 
      [ addNRule, addAverageRule, addVarianceRule, addStandardDeviationRule
      , addStandardErrorRule, addStandardErrorSD, addDfRule, addDfBetweenWithin
      , addStandardErrorSigma, addObservedTotals, addExpectedFrequencies
      ]

-- customized while (combinator in library is greedy and uses repeat instead of many)
while :: IsStrategy f => (a -> Bool) -> f a -> Strategy a
while p s = many (check p .*. s)

whileNotReady :: Strategy ComponentSet -> Strategy ComponentSet
whileNotReady = while (`doesNotContain` ConclusionHypotheses)

allowCriticalRoute :: ComponentSet -> Bool
allowCriticalRoute cs = criticalRoute cs || not (pvalueRoute cs)

allowPValueRoute :: ComponentSet -> Bool
allowPValueRoute cs = pvalueRoute cs || not (criticalRoute cs)

criticalRoute :: ComponentSet -> Bool
criticalRoute cs = derived cs `contains` RejectionCritical || derived cs `contains` Critical

pvalueRoute :: ComponentSet -> Bool
pvalueRoute cs = derived cs `contains` PValue