    -- |
    -- Module    : Statistics.Test.MannWhitneyU
    -- Copyright : (c) 2010 Neil Brown
    -- License   : BSD3
    --
    -- Maintainer  : bos@serpentine.com
    -- Stability   : experimental
    -- Portability : portable
    --
    -- Mann-Whitney U test (also know as Mann-Whitney-Wilcoxon and
    -- Wilcoxon rank sum test) is a non-parametric test for assesing
    -- whether two samples of independent observations have different
    -- mean.
    module Statistics.Test.MannWhitneyU (
        -- * Mann-Whitney U test
        mannWhitneyUtest
      , mannWhitneyU
      , mannWhitneyUCriticalValue
      , mannWhitneyUSignificant
        -- ** Wilcoxon rank sum test
      , wilcoxonRankSums
        -- * Data types
      , TestType(..)
      , TestResult(..)
        -- * References
        -- $references
      ) where

    import Control.Applicative ((<$>))
    import Data.List           (findIndex)
    import Data.Ord            (comparing)
    import qualified Data.Vector.Unboxed as U

    import Numeric.SpecFunctions          (choose)

    import Statistics.Distribution        (quantile)
    import Statistics.Distribution.Normal (standard)
    import Statistics.Types               (Sample)
    import Statistics.Function            (sortBy)
    import Statistics.Test.Types
    import Statistics.Test.Internal

    -- | The Wilcoxon Rank Sums Test.
    --
    -- This test calculates the sum of ranks for the given two samples.  The samples
    -- are ordered, and assigned ranks (ties are given their average rank), then these
    -- ranks are summed for each sample.
    --
    -- The return value is (W&#8321;, W&#8322;) where W&#8321; is the sum of ranks of the first sample
    -- and W&#8322; is the sum of ranks of the second sample.  This test is trivially transformed
    -- into the Mann-Whitney U test.  You will probably want to use 'mannWhitneyU'
    -- and the related functions for testing significance, but this function is exposed
    -- for completeness.
    wilcoxonRankSums :: Sample -> Sample -> (Double, Double)
    wilcoxonRankSums xs1 xs2 = ( U.sum ranks1 , U.sum ranks2
                               )
      where
        -- Ranks for each sample
        (ranks1,ranks2) = splitByTags $ U.zip tags (rank (==) joinSample)
        -- Sorted and tagged sample
        (tags,joinSample) = U.unzip
                          $ sortBy (comparing snd)
                          $ tagSample True xs1 U.++ tagSample False xs2
        -- Add tag to a sample
        tagSample t = U.map ((,) t)

    -- | The Mann-Whitney U Test.
    --
    -- This is sometimes known as the Mann-Whitney-Wilcoxon U test, and
    -- confusingly many sources state that the Mann-Whitney U test is the same as
    -- the Wilcoxon's rank sum test (which is provided as 'wilcoxonRankSums').
    -- The Mann-Whitney U is a simple transform of Wilcoxon's rank sum test.
    --
    -- Again confusingly, different sources state reversed definitions for U&#8321;
    -- and U&#8322;, so it is worth being explicit about what this function returns.
    -- Given two samples, the first, xs&#8321;, of size n&#8321; and the second, xs&#8322;,
    -- of size n&#8322;, this function returns (U&#8321;, U&#8322;)
    -- where U&#8321; = W&#8321; - (n&#8321;(n&#8321;+1))\/2
    -- and U&#8322; = W&#8322; - (n&#8322;(n&#8322;+1))\/2,
    -- where (W&#8321;, W&#8322;) is the return value of @wilcoxonRankSums xs1 xs2@.
    --
    -- Some sources instead state that U&#8321; and U&#8322; should be the other way round, often
    -- expressing this using U&#8321;' = n&#8321;n&#8322; - U&#8321; (since U&#8321; + U&#8322; = n&#8321;n&#8322;).
    --
    -- All of which you probably don't care about if you just feed this into 'mannWhitneyUSignificant'.
    mannWhitneyU :: Sample -> Sample -> (Double, Double)
    mannWhitneyU xs1 xs2
      = (fst summedRanks - (n1*(n1 + 1))/2
        ,snd summedRanks - (n2*(n2 + 1))/2)
      where
        n1 = fromIntegral $ U.length xs1
        n2 = fromIntegral $ U.length xs2

        summedRanks = wilcoxonRankSums xs1 xs2

    -- | Calculates the critical value of Mann-Whitney U for the given sample
    -- sizes and significance level.
    --
    -- This function returns the exact calculated value of U for all sample sizes;
    -- it does not use the normal approximation at all.  Above sample size 20 it is
    -- generally recommended to use the normal approximation instead, but this function
    -- will calculate the higher critical values if you need them.
    --
    -- The algorithm to generate these values is a faster, memoised version of the
    -- simple unoptimised generating function given in section 2 of \"The Mann Whitney
    -- Wilcoxon Distribution Using Linked Lists\"
    mannWhitneyUCriticalValue :: (Int, Int) -- ^ The sample size
                              -> Double     -- ^ The p-value (e.g. 0.05) for which you want the critical value.
                              -> Maybe Int  -- ^ The critical value (of U).
    mannWhitneyUCriticalValue (m, n) p
      | m < 1 || n < 1 = Nothing    -- Sample must be nonempty
      | p  >= 1        = Nothing    -- Nonsensical p-value
      | p' <= 1        = Nothing    -- p-value is too small. Null hypothesys couln't be disproved
      | otherwise      = findIndex (>= p')
                       $ take (m*n)
                       $ tail
                       $ alookup !! (m+n-2) !! (min m n - 1)
      where
        mnCn = (m+n) `choose` n
        p'   = mnCn * p

    {-
    -- Original function, without memoisation, from Cheung and Klotz:
    -- Double is needed to avoid integer overflows.
    a :: Int -> Int -> Int -> Double
    a u bigN m
      | u < 0            = 0
      | u >= m * n       = bigN `choose` m
      | m == 1 || n == 1 = fromIntegral (u + 1)
      | otherwise        = a  u      (bigN - 1)  m
                         + a (u - n) (bigN - 1) (m-1)
      where
        n = bigN - m
    -}

    -- Memoised version of the original a function, above.
    --
    -- Doubles are stored to avoid integer overflow. 32-bit Ints begin to
    -- overflow for bigN as small as 33 (64-bit one at 66) while Double to
    -- go to infinity till bigN=1029
    --
    --
    -- outer list is indexed by big N - 2
    -- inner list by (m-1) (we know m < bigN)
    -- innermost list by u
    --
    -- So: (alookup !! (bigN - 2) !! (m - 1) ! u) == a u bigN m
    alookup :: [[[Double]]]
    alookup = gen 2 [1 : repeat 2]
      where
        gen bigN predBigNList
           = let bigNlist = [ [ amemoed u m
                              | u <- [0 .. m*(bigN-m)]
                              ] ++ repeat (bigN `choose` m)
                            | m <- [1 .. (bigN-1)]] -- has bigN-1 elements
             in bigNlist : gen (bigN+1) bigNlist
          where
            amemoed :: Int -> Int -> Double
            amemoed u m
              | m == 1 || n == 1 = fromIntegral (u + 1)
              | otherwise        = mList !! u
                                 + if u < n then 0 else predmList !! (u-n)
              where
                n = bigN - m
                (predmList : mList : _) = drop (m-2) predBigNList
                -- Lists for m-1 and m respectively. i-th list correspond to m=i+1
                --
                -- We know that predBigNList has bigN - 2 elements
                -- (and we know that n > 1 therefore bigN > m + 1)
                -- So bigN - 2 >= m, i.e. predBigNList must have at least m elements
                -- elements, so dropping (m-2) must leave at least 2

    -- | Calculates whether the Mann Whitney U test is significant.
    --
    -- If both sample sizes are less than or equal to 20, the exact U critical value
    -- (as calculated by 'mannWhitneyUCriticalValue') is used.  If either sample is
    -- larger than 20, the normal approximation is used instead.
    --
    -- If you use a one-tailed test, the test indicates whether the first sample is
    -- significantly larger than the second.  If you want the opposite, simply reverse
    -- the order in both the sample size and the (U&#8321;, U&#8322;) pairs.
    mannWhitneyUSignificant ::
         TestType         -- ^ Perform one-tailed test (see description above).
      -> (Int, Int)       -- ^ The samples' size from which the (U&#8321;,U&#8322;) values were derived.
      -> Double           -- ^ The p-value at which to test (e.g. 0.05)
      -> (Double, Double) -- ^ The (U&#8321;, U&#8322;) values from 'mannWhitneyU'.
      -> Maybe TestResult -- ^ Return 'Nothing' if the sample was too
                          --   small to make a decision.
    mannWhitneyUSignificant test (in1, in2) p (u1, u2)
       --Use normal approximation
      | in1 > 20 || in2 > 20 =
        let mean  = n1 * n2 / 2
            sigma = sqrt $ n1*n2*(n1 + n2 + 1) / 12
            z     = (mean - u1) / sigma
        in Just $ case test of
                    OneTailed -> significant $ z     < quantile standard  p
                    TwoTailed -> significant $ abs z > abs (quantile standard (p/2))
      -- Use exact critical value
      | otherwise = do crit <- fromIntegral <$> mannWhitneyUCriticalValue (in1, in2) p
                       return $ case test of
                                  OneTailed -> significant $ u2        <= crit
                                  TwoTailed -> significant $ min u1 u2 <= crit
      where
        n1 = fromIntegral in1
        n2 = fromIntegral in2

    -- | Perform Mann-Whitney U Test for two samples and required
    -- significance. For additional information check documentation of
    -- 'mannWhitneyU' and 'mannWhitneyUSignificant'. This is just a helper
    -- function.
    --
    -- One-tailed test checks whether first sample is significantly larger
    -- than second. Two-tailed whether they are significantly different.
    mannWhitneyUtest :: TestType    -- ^ Perform one-tailed test (see description above).
                     -> Double      -- ^ The p-value at which to test (e.g. 0.05)
                     -> Sample      -- ^ First sample
                     -> Sample      -- ^ Second sample
                     -> Maybe TestResult
                     -- ^ Return 'Nothing' if the sample was too small to
                     --   make a decision.
    mannWhitneyUtest ontTail p smp1 smp2 =
      mannWhitneyUSignificant ontTail (n1,n2) p $ mannWhitneyU smp1 smp2
        where
          n1 = U.length smp1
          n2 = U.length smp2

    -- $references
    --
    -- * Cheung, Y.K.; Klotz, J.H. (1997) The Mann Whitney Wilcoxon
    --   distribution using linked lists. /Statistica Sinica/
    --   7:805&#8211;813.
    -- <http://www3.stat.sinica.edu.tw/statistica/oldpdf/A7n316.pdf>.