{-# LANGUAGE FlexibleContexts #-}
{- |
Module      : Language.Scheme.Numerical
Copyright   : Justin Ethier
Licence     : MIT (see LICENSE in the distribution)

Maintainer  : github.com/justinethier
Stability   : experimental
Portability : portable

This module implements the numerical tower.
-}

module Language.Scheme.Numerical (
 -- * Generic functions
   numSub
 , numMul
 , numDiv
 , numAdd
 , numMod
 , numRationalize
 , numBoolBinopEq
 , numBoolBinopGt
 , numBoolBinopGte
 , numBoolBinopLt
 , numBoolBinopLte
 , numCast
 , numDenominator
 , numNumerator
 , numInexact2Exact
 , numExact2Inexact
 , num2String
 , unpackNum
 , numericBinop
 -- * Floating point functions
 , numFloor
 , numCeiling
 , numTruncate
 , numRound
 , numExpt
 , numSqrt
 , numExp
 , numLog
 -- * Trigonometric functions
 , numSin
 , numCos
 , numTan
 , numAsin
 , numAcos
 , numAtan
 -- * Complex functions
 , buildComplex
 , numMakePolar
 , numRealPart
 , numImagPart
 , numMagnitude
 , numAngle
 , numMakeRectangular
 -- * Predicates
 , isComplex
 , isReal
 , isRational
 , isInteger
 , isNumber
 , isFloatAnInteger
 , isNumNaN
 , isNumInfinite
 , isNumFinite
 , isNumExact
 , isNumInexact
) where
import Language.Scheme.Types

import Control.Monad.Error
import Data.Char hiding (isNumber)
import Data.Complex
import Data.Fixed
import Data.Ratio
import Numeric
import Text.Printf

-- |A helper function to perform a numeric operation on two values
numericBinop :: (Integer -> Integer -> Integer) -> [LispVal] -> ThrowsError LispVal
numericBinop _ singleVal@[_] = throwError $ NumArgs (Just 2) singleVal
numericBinop op aparams = mapM unpackNum aparams >>= return . Number . foldl1 op

-- - Begin GenUtil - http://repetae.net/computer/haskell/GenUtil.hs
foldlM :: Monad m => (a -> b -> m a) -> a -> [b] -> m a
foldlM f v (x : xs) = (f v x) >>= \ a -> foldlM f a xs
foldlM _ v [] = return v

foldl1M :: Monad m => (a -> a -> m a) -> [a] -> m a
foldl1M f (x : xs) = foldlM f x xs
foldl1M _ _ = error "Unexpected error in foldl1M"
-- end GenUtil

{- FUTURE: as a general comment here, operations need to be more permissive of the
numerical types they accept. Within reason, a user should not have to know
what numerical type they are passing when using these functions -}

-- |Add the given numbers
numAdd :: [LispVal] -> ThrowsError LispVal
numAdd [] = return $ Number 0
numAdd aparams = do
  foldl1M (\ a b -> doAdd =<< (numCast [a, b])) aparams
  where doAdd (List [(Number a), (Number b)]) = return $ Number $ a + b
        doAdd (List [(Float a), (Float b)]) = return $ Float $ a + b
        doAdd (List [(Rational a), (Rational b)]) = return $ Rational $ a + b
        doAdd (List [(Complex a), (Complex b)]) = return $ Complex $ a + b
        doAdd _ = throwError $ Default "Unexpected error in +"

-- |Subtract the given numbers
numSub :: [LispVal] -> ThrowsError LispVal
numSub [] = throwError $ NumArgs (Just 1) []
numSub [Number n] = return $ Number $ -1 * n
numSub [Float n] = return $ Float $ -1 * n
numSub [Rational n] = return $ Rational $ -1 * n
numSub [Complex n] = return $ Complex $ -1 * n
numSub aparams = do
  foldl1M (\ a b -> doSub =<< (numCast [a, b])) aparams
  where doSub (List [(Number a), (Number b)]) = return $ Number $ a - b
        doSub (List [(Float a), (Float b)]) = return $ Float $ a - b
        doSub (List [(Rational a), (Rational b)]) = return $ Rational $ a - b
        doSub (List [(Complex a), (Complex b)]) = return $ Complex $ a - b
        doSub _ = throwError $ Default "Unexpected error in -"

-- |Multiply the given numbers
numMul :: [LispVal] -> ThrowsError LispVal
numMul [] = return $ Number 1
numMul aparams = do
  foldl1M (\ a b -> doMul =<< (numCast [a, b])) aparams
  where doMul (List [(Number a), (Number b)]) = return $ Number $ a * b
        doMul (List [(Float a), (Float b)]) = return $ Float $ a * b
        doMul (List [(Rational a), (Rational b)]) = return $ Rational $ a * b
        doMul (List [(Complex a), (Complex b)]) = return $ Complex $ a * b
        doMul _ = throwError $ Default "Unexpected error in *"

-- |Divide the given numbers
numDiv :: [LispVal] -> ThrowsError LispVal
numDiv [] = throwError $ NumArgs (Just 1) []
numDiv [Number 0] = throwError $ DivideByZero
numDiv [Rational 0] = throwError $ DivideByZero
numDiv [Number n] = return $ Rational $ 1 / (fromInteger n)
numDiv [Float n] = return $ Float $ 1.0 / n
numDiv [Rational n] = return $ Rational $ 1 / n
numDiv [Complex n] = return $ Complex $ 1 / n
numDiv aparams = do
  foldl1M (\ a b -> doDiv =<< (numCast [a, b])) aparams
  where doDiv (List [(Number a), (Number b)])
            | b == 0 = throwError $ DivideByZero
            | (mod a b) == 0 = return $ Number $ div a b
            | otherwise = -- Not an integer
                return $ Rational $ (fromInteger a) / (fromInteger b)
        doDiv (List [(Float a), (Float b)])
            | b == 0.0 = throwError $ DivideByZero
            | otherwise = return $ Float $ a / b
        doDiv (List [(Rational a), (Rational b)])
            | b == 0 = throwError $ DivideByZero
            | otherwise = return $ Rational $ a / b
        doDiv (List [(Complex a), (Complex b)])
            | b == 0 = throwError $ DivideByZero
            | otherwise = return $ Complex $ a / b
        doDiv _ = throwError $ Default "Unexpected error in /"

-- |Take the modulus of the given numbers
numMod :: [LispVal] -> ThrowsError LispVal
numMod [] = return $ Number 1
numMod aparams = do
  foldl1M (\ a b -> doMod =<< (numCast [a, b])) aparams
  where doMod (List [(Number a), (Number b)]) = return $ Number $ mod' a b
        doMod (List [(Float a), (Float b)]) = return $ Float $ mod' a b
        doMod (List [(Rational a), (Rational b)]) = return $ Rational $ mod' a b
        doMod (List [(Complex _), (Complex _)]) = throwError $ Default "modulo not implemented for complex numbers"
        doMod _ = throwError $ Default "Unexpected error in modulo"

-- |Compare a series of numbers using a given numeric comparison
--  function and an array of lisp values
numBoolBinopCompare :: (LispVal
                    -> LispVal -> Either LispError LispVal)
                    -> LispVal -> [LispVal] -> Either LispError LispVal
numBoolBinopCompare cmp n1 (n2 : ns) = do
  (n1', n2') <- numCast' (n1, n2)
  result <- cmp n1' n2'
  case result of
    Bool True -> numBoolBinopCompare cmp n2' ns
    _ -> return $ Bool False
numBoolBinopCompare _ _ _ = return $ Bool True

-- |Numeric equals
numBoolBinopEq :: [LispVal] -> ThrowsError LispVal
numBoolBinopEq [] = throwError $ NumArgs (Just 0) []
numBoolBinopEq (n : ns) = numBoolBinopCompare cmp n ns
  where
    f a b = a == b
    cmp (Number a) (Number b) = return $ Bool $ f a b
    cmp (Float a) (Float b) = return $ Bool $ f a b
    cmp (Rational a) (Rational b) = return $ Bool $ f a b
    cmp (Complex a) (Complex b) = return $ Bool $ f a b
    cmp _ _ = throwError $ Default "Unexpected error in ="

-- |Numeric greater than
numBoolBinopGt :: [LispVal] -> ThrowsError LispVal
numBoolBinopGt [] = throwError $ NumArgs (Just 0) []
numBoolBinopGt (n : ns) = numBoolBinopCompare cmp n ns
  where
    f a b = a > b
    cmp (Number a) (Number b) = return $ Bool $ f a b
    cmp (Float a) (Float b) = return $ Bool $ f a b
    cmp (Rational a) (Rational b) = return $ Bool $ f a b
    cmp _ _ = throwError $ Default "Unexpected error in >"

-- |Numeric greater than equal
numBoolBinopGte :: [LispVal] -> ThrowsError LispVal
numBoolBinopGte [] = throwError $ NumArgs (Just 0) []
numBoolBinopGte (n : ns) = numBoolBinopCompare cmp n ns
  where
    f a b = a >= b
    cmp (Number a) (Number b) = return $ Bool $ f a b
    cmp (Float a) (Float b) = return $ Bool $ f a b
    cmp (Rational a) (Rational b) = return $ Bool $ f a b
    cmp _ _ = throwError $ Default "Unexpected error in >="

-- |Numeric less than
numBoolBinopLt :: [LispVal] -> ThrowsError LispVal
numBoolBinopLt [] = throwError $ NumArgs (Just 0) []
numBoolBinopLt (n : ns) = numBoolBinopCompare cmp n ns
  where
    f a b = a < b
    cmp (Number a) (Number b) = return $ Bool $ f a b
    cmp (Float a) (Float b) = return $ Bool $ f a b
    cmp (Rational a) (Rational b) = return $ Bool $ f a b
    cmp _ _ = throwError $ Default "Unexpected error in <"

-- |Numeric less than equal
numBoolBinopLte :: [LispVal] -> ThrowsError LispVal
numBoolBinopLte [] = throwError $ NumArgs (Just 0) []
numBoolBinopLte (n : ns) = numBoolBinopCompare cmp n ns
  where
    f a b = a <= b
    cmp (Number a) (Number b) = return $ Bool $ f a b
    cmp (Float a) (Float b) = return $ Bool $ f a b
    cmp (Rational a) (Rational b) = return $ Bool $ f a b
    cmp _ _ = throwError $ Default "Unexpected error in <="

-- |Accept two numbers and cast one of them to the appropriate type, if necessary
numCast' :: (LispVal, LispVal) -> ThrowsError (LispVal, LispVal)
numCast' (a@(Number _), b@(Number _)) = return $ (a, b)
numCast' (a@(Float _), b@(Float _)) = return $ (a, b)
numCast' (a@(Rational _), b@(Rational _)) = return $ (a, b)
numCast' (a@(Complex _), b@(Complex _)) = return $ (a, b)
numCast' ((Number a), b@(Float _)) = return $ (Float $ fromInteger a, b)
numCast' ((Number a), b@(Rational _)) = return $ (Rational $ fromInteger a, b)
numCast' ((Number a), b@(Complex _)) = return $ (Complex $ fromInteger a, b)
numCast' (a@(Float _), (Number b)) = return $ (a, Float $ fromInteger b)
numCast' (a@(Float _), (Rational b)) = return $ (a, Float $ fromRational b)
numCast' ((Float a), b@(Complex _)) = return $ (Complex $ a :+ 0, b)
numCast' (a@(Rational _), (Number b)) = return $ (a, Rational $ fromInteger b)
numCast' ((Rational a), b@(Float _)) = return $ (Float $ fromRational a, b)
numCast' ((Rational a), b@(Complex _)) = return $ (Complex $ (fromInteger $ numerator a) / (fromInteger $ denominator a), b)
numCast' (a@(Complex _), (Number b)) = return $ (a, Complex $ fromInteger b)
numCast' (a@(Complex _), (Float b)) = return $ (a, Complex $ b :+ 0)
numCast' (a@(Complex _), (Rational b)) = return $ (a, Complex $ (fromInteger $ numerator b) / (fromInteger $ denominator b))
numCast' (a, b) = case a of
               Number _ -> doThrowError b
               Float _ -> doThrowError b
               Rational _ -> doThrowError b
               Complex _ -> doThrowError b
               _ -> doThrowError a
  where doThrowError num = throwError $ TypeMismatch "number" num

-- |Accept two numbers and cast one of them to the appropriate type, if necessary
numCast :: [LispVal] -> ThrowsError LispVal
numCast [a, b] = do
  (a', b') <- numCast' (a, b)
  pure $ List [a', b']
numCast _ = throwError $ Default "Unexpected error in numCast"

-- |Convert the given number to a rational
numRationalize :: [LispVal] -> ThrowsError LispVal
numRationalize [(Number n)] = return $ Rational $ toRational n
numRationalize [(Float n)] = return $ Rational $ toRational n
numRationalize [n@(Rational _)] = return n
numRationalize [x] = throwError $ TypeMismatch "number" x
numRationalize badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Round the given number
numRound :: [LispVal] -> ThrowsError LispVal
numRound [n@(Number _)] = return n
numRound [(Rational n)] = return $ Number $ round n
numRound [(Float n)] = return $ Float $ fromInteger $ round n
numRound [(Complex n)] = do
  return $ Complex $ (fromInteger $ round $ realPart n) :+ (fromInteger $ round $ imagPart n)
numRound [x] = throwError $ TypeMismatch "number" x
numRound badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Floor the given number
numFloor :: [LispVal] -> ThrowsError LispVal
numFloor [n@(Number _)] = return n
numFloor [(Rational n)] = return $ Number $ floor n
numFloor [(Float n)] = return $ Float $ fromInteger $ floor n
numFloor [(Complex n)] = do
  return $ Complex $ (fromInteger $ floor $ realPart n) :+ (fromInteger $ floor $ imagPart n)
numFloor [x] = throwError $ TypeMismatch "number" x
numFloor badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Take the ceiling of the given number
numCeiling :: [LispVal] -> ThrowsError LispVal
numCeiling [n@(Number _)] = return n
numCeiling [(Rational n)] = return $ Number $ ceiling n
numCeiling [(Float n)] = return $ Float $ fromInteger $ ceiling n
numCeiling [(Complex n)] = do
  return $ Complex $ (fromInteger $ ceiling $ realPart n) :+ (fromInteger $ ceiling $ imagPart n)
numCeiling [x] = throwError $ TypeMismatch "number" x
numCeiling badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Truncate the given number
numTruncate :: [LispVal] -> ThrowsError LispVal
numTruncate [n@(Number _)] = return n
numTruncate [(Rational n)] = return $ Number $ truncate n
numTruncate [(Float n)] = return $ Float $ fromInteger $ truncate n
numTruncate [(Complex n)] = do
  return $ Complex $ (fromInteger $ truncate $ realPart n) :+ (fromInteger $ truncate $ imagPart n)
numTruncate [x] = throwError $ TypeMismatch "number" x
numTruncate badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Sine
numSin :: [LispVal] -> ThrowsError LispVal
numSin [(Number n)] = return $ Float $ sin $ fromInteger n
numSin [(Float n)] = return $ Float $ sin n
numSin [(Rational n)] = return $ Float $ sin $ fromRational n
numSin [(Complex n)] = return $ Complex $ sin n
numSin [x] = throwError $ TypeMismatch "number" x
numSin badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Cosine
numCos :: [LispVal] -> ThrowsError LispVal
numCos [(Number n)] = return $ Float $ cos $ fromInteger n
numCos [(Float n)] = return $ Float $ cos n
numCos [(Rational n)] = return $ Float $ cos $ fromRational n
numCos [(Complex n)] = return $ Complex $ cos n
numCos [x] = throwError $ TypeMismatch "number" x
numCos badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Tangent
numTan :: [LispVal] -> ThrowsError LispVal
numTan [(Number n)] = return $ Float $ tan $ fromInteger n
numTan [(Float n)] = return $ Float $ tan n
numTan [(Rational n)] = return $ Float $ tan $ fromRational n
numTan [(Complex n)] = return $ Complex $ tan n
numTan [x] = throwError $ TypeMismatch "number" x
numTan badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Arcsine
numAsin :: [LispVal] -> ThrowsError LispVal
numAsin [(Number n)] = return $ Float $ asin $ fromInteger n
numAsin [(Float n)] = return $ Float $ asin n
numAsin [(Rational n)] = return $ Float $ asin $ fromRational n
numAsin [(Complex n)] = return $ Complex $ asin n
numAsin [x] = throwError $ TypeMismatch "number" x
numAsin badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Arccosine
numAcos :: [LispVal] -> ThrowsError LispVal
numAcos [(Number n)] = return $ Float $ acos $ fromInteger n
numAcos [(Float n)] = return $ Float $ acos n
numAcos [(Rational n)] = return $ Float $ acos $ fromRational n
numAcos [(Complex n)] = return $ Complex $ acos n
numAcos [x] = throwError $ TypeMismatch "number" x
numAcos badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Arctangent
numAtan :: [LispVal] -> ThrowsError LispVal
numAtan [(Number n)] = return $ Float $ atan $ fromInteger n
numAtan [Number y, Number x] = return $ Float $ phase $ (fromInteger x) :+ (fromInteger y)
numAtan [(Float n)] = return $ Float $ atan n
numAtan [Float y, Float x] = return $ Float $ phase $ x :+ y
numAtan [(Rational n)] = return $ Float $ atan $ fromRational n
numAtan [Rational y, Rational x] = return $ Float $ phase $ (fromRational x) :+ (fromRational y)
numAtan [(Complex n)] = return $ Complex $ atan n
numAtan [x] = throwError $ TypeMismatch "number" x
numAtan badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Take the square root of the given number
numSqrt :: [LispVal] -> ThrowsError LispVal
numSqrt [(Number n)] = if n >= 0 then return $ Float $ sqrt $ fromInteger n
                                 else return $ Complex $ sqrt ((fromInteger n) :+ 0)
numSqrt [(Float n)] = if n >= 0 then return $ Float $ sqrt n
                                else return $ Complex $ sqrt (n :+ 0)
numSqrt [(Rational n)] = numSqrt [Float $ fromRational n]
numSqrt [(Complex n)] = return $ Complex $ sqrt n
numSqrt [x] = throwError $ TypeMismatch "number" x
numSqrt badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Raise the first number to the power of the second
numExpt :: [LispVal] -> ThrowsError LispVal
numExpt [(Number n), (Number p)] = return $ Float $ (fromInteger n) ^ p
numExpt [(Rational n), (Number p)] = return $ Float $ (fromRational n) ^ p
numExpt [(Float n), (Number p)] = return $ Float $ n ^ p
numExpt [(Complex n), (Number p)] = return $ Complex $ n ^ p
numExpt [_, y] = throwError $ TypeMismatch "integer" y
numExpt badArgList = throwError $ NumArgs (Just 2) badArgList

{- numExpt params = do
  foldl1M (\a b -> doExpt =<< (numCast [a, b])) params
  where doExpt (List [(Number a), (Number b)]) = return $ Float $ (fromInteger a) ^ (fromInteger b)
--        doExpt (List [(Rational a), (Rational b)]) = return $ Float $ fromRational $ a ^ b
        doExpt (List [(Float a), (Float b)]) = return $ Float $ a ^ b
--        doExpt (List [(Complex a), (Complex b)]) = return $ Complex $ a ^ b -}

-- |Take the exponent of the given number
numExp :: [LispVal] -> ThrowsError LispVal
numExp [(Number n)] = return $ Float $ exp $ fromInteger n
numExp [(Float n)] = return $ Float $ exp n
numExp [(Rational n)] = return $ Float $ exp $ fromRational n
numExp [(Complex n)] = return $ Complex $ exp n
numExp [x] = throwError $ TypeMismatch "number" x
numExp badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Compute the log of a given number
numLog :: [LispVal] -> ThrowsError LispVal
numLog [(Number n)] = return $ Float $ log $ fromInteger n
numLog [Number n, Number base] = return $ Float $ logBase (fromInteger base) (fromInteger n)
numLog [(Float n)] = return $ Float $ log n
numLog [Float n, Number base] = return $ Float $ logBase (fromInteger base) n
numLog [(Rational n)] = return $ Float $ log $ fromRational n
numLog [Rational n, Number base] = return $ Float $ logBase (fromInteger base) (fromRational n)
numLog [(Complex n)] = return $ Complex $ log n
numLog [Complex n, Number base] = return $ Complex $ logBase (fromInteger base) n
numLog [x] = throwError $ TypeMismatch "number" x
numLog badArgList = throwError $ NumArgs (Just 1) badArgList

-- Complex number functions

-- |Create a complex number
buildComplex :: LispVal
             -- ^ Real part
             -> LispVal
             -- ^ Imaginary part
             -> ThrowsError LispVal
             -- ^ Complex number
buildComplex (Number x) (Number y) = return $ Complex $ (fromInteger x) :+ (fromInteger y)
buildComplex (Number x) (Rational y) = return $ Complex $ (fromInteger x) :+ (fromRational y)
buildComplex (Number x) (Float y) = return $ Complex $ (fromInteger x) :+ y
buildComplex (Rational x) (Number y) = return $ Complex $ (fromRational x) :+ (fromInteger y)
buildComplex (Rational x) (Rational y) = return $ Complex $ (fromRational x) :+ (fromRational y)
buildComplex (Rational x) (Float y) = return $ Complex $ (fromRational x) :+ y
buildComplex (Float x) (Number y) = return $ Complex $ x :+ (fromInteger y)
buildComplex (Float x) (Rational y) = return $ Complex $ x :+ (fromRational y)
buildComplex (Float x) (Float y) = return $ Complex $ x :+ y
buildComplex x y = throwError $ TypeMismatch "number" $ List [x, y]

-- |Create a complex number given its real and imaginary parts
numMakeRectangular :: [LispVal] -> ThrowsError LispVal
numMakeRectangular [x, y] = buildComplex x y
numMakeRectangular badArgList = throwError $ NumArgs (Just 2) badArgList

-- |Create a complex number from its magnitude and phase (angle)
numMakePolar :: [LispVal] -> ThrowsError LispVal
numMakePolar [(Float x), (Float y)] = return $ Complex $ mkPolar x y
numMakePolar [(Float _), y] = throwError $ TypeMismatch "real" y
numMakePolar [x, (Float _)] = throwError $ TypeMismatch "real real" $ x
numMakePolar badArgList = throwError $ NumArgs (Just 2) badArgList

-- |The phase of a complex number
numAngle :: [LispVal] -> ThrowsError LispVal
numAngle [(Complex c)] = return $ Float $ phase c
numAngle [x] = throwError $ TypeMismatch "complex number" x
numAngle badArgList = throwError $ NumArgs (Just 1) badArgList

-- |The nonnegative magnitude of a complex number
numMagnitude :: [LispVal] -> ThrowsError LispVal
numMagnitude [(Complex c)] = return $ Float $ magnitude c
numMagnitude [x] = throwError $ TypeMismatch "complex number" x
numMagnitude badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Retrieve real part of a complex number
numRealPart :: [LispVal] -> ThrowsError LispVal
numRealPart [(Complex c)] = return $ Float $ realPart c
numRealPart [n@(Float _)] = return n
numRealPart [n@(Rational _)] = return n
numRealPart [n@(Number _)] = return n
numRealPart [x] = throwError $ TypeMismatch "complex number" x
numRealPart badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Retrieve imaginary part of a complex number
numImagPart :: [LispVal] -> ThrowsError LispVal
numImagPart [(Complex c)] = do
  let n = imagPart c
      f = Float n
  if isFloatAnInteger f
     then return $ Number $ floor n
     else return f
numImagPart [(Float _)] = return $ Number 0
numImagPart [(Rational _)] = return $ Number 0
numImagPart [(Number _)] = return $ Number 0
numImagPart [x] = throwError $ TypeMismatch "complex number" x
numImagPart badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Take the numerator of the given number
numNumerator :: [LispVal] -> ThrowsError LispVal
numNumerator [n@(Number _)] = return n
numNumerator [(Rational r)] = return $ Number $ numerator r
numNumerator [(Float f)] = return $ Float $ fromInteger . numerator . toRational $ f
numNumerator [x] = throwError $ TypeMismatch "rational number" x
numNumerator badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Take the denominator of the given number
numDenominator :: [LispVal] -> ThrowsError LispVal
numDenominator [Number _] = return $ Number 1
numDenominator [(Rational r)] = return $ Number $ denominator r
numDenominator [(Float f)] = return $ Float $ fromInteger $ denominator $ toRational f
numDenominator [x] = throwError $ TypeMismatch "rational number" x
numDenominator badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Convert an exact number to inexact
numExact2Inexact :: [LispVal] -> ThrowsError LispVal
numExact2Inexact [(Number n)] = return $ Float $ fromInteger n
numExact2Inexact [(Rational n)] = return $ Float $ fromRational n
numExact2Inexact [n@(Float _)] = return n
numExact2Inexact [n@(Complex _)] = return n
numExact2Inexact [badType] = throwError $ TypeMismatch "number" badType
numExact2Inexact badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Convert an inexact number to exact
numInexact2Exact :: [LispVal] -> ThrowsError LispVal
numInexact2Exact [n@(Number _)] = return n
numInexact2Exact [n@(Rational _)] = return n
numInexact2Exact [(Float n)] = return $ Number $ round n
numInexact2Exact [c@(Complex _)] = numRound [c]
numInexact2Exact [badType] = throwError $ TypeMismatch "number" badType
numInexact2Exact badArgList = throwError $ NumArgs (Just 1) badArgList

-- |Convert a number to a string; radix is optional, defaults to base 10
num2String :: [LispVal] -> ThrowsError LispVal
num2String [(Number n)] = return $ String $ show n
num2String [(Number n), (Number radix)] = do
  case radix of
    2 -> do -- Nice tip from StackOverflow question #1959715
             return $ String $ showIntAtBase 2 intToDigit n ""
    8 -> return $ String $ printf "%o" n
    10 -> return $ String $ printf "%d" n
    16 -> return $ String $ printf "%x" n
    _ -> throwError $ BadSpecialForm "Invalid radix value" $ Number radix
num2String [n@(Rational _)] = return $ String $ show n
num2String [(Float n)] = return $ String $ show n
num2String [n@(Complex _)] = return $ String $ show n
num2String [x] = throwError $ TypeMismatch "number" x
num2String badArgList = throwError $ NumArgs (Just 1) badArgList

-- | Determine if the given value is not a number
isNumNaN :: [LispVal] -> ThrowsError LispVal
isNumNaN ([Float n]) = return $ Bool $ isNaN n
isNumNaN _ = return $ Bool False

-- | Determine if number is infinite
isNumInfinite :: [LispVal] -> ThrowsError LispVal
isNumInfinite ([Float n]) = return $ Bool $ isInfinite n
isNumInfinite _ = return $ Bool False

-- | Determine if number is not infinite
isNumFinite :: [LispVal] -> ThrowsError LispVal
isNumFinite ([Number _]) = return $ Bool True
isNumFinite ([Float n]) = return $ Bool $ not $ isInfinite n
isNumFinite ([Complex _]) = return $ Bool True
isNumFinite ([Rational _]) = return $ Bool True
isNumFinite _ = return $ Bool False

-- | Determine if number is exact
isNumExact :: [LispVal] -> ThrowsError LispVal
isNumExact ([Number _]) = return $ Bool True
isNumExact ([Float _]) = return $ Bool False
isNumExact ([Complex _]) = return $ Bool False -- TODO: could be either
isNumExact ([Rational _]) = return $ Bool True
isNumExact _ = return $ Bool False

-- | Determine if number is inexact
isNumInexact :: [LispVal] -> ThrowsError LispVal
isNumInexact ([Number _])   = return $ Bool False
isNumInexact ([Float _])    = return $ Bool True
isNumInexact ([Complex _])  = return $ Bool True
isNumInexact ([Rational _]) = return $ Bool False
isNumInexact _ = return $ Bool False

-- |Predicate to determine if given value is a number
isNumber :: [LispVal] -> ThrowsError LispVal
isNumber ([Number _]) = return $ Bool True
isNumber ([Float _]) = return $ Bool True
isNumber ([Complex _]) = return $ Bool True
isNumber ([Rational _]) = return $ Bool True
isNumber _ = return $ Bool False

-- |Predicate to determine if given number is complex.
--  Keep in mind this does not just look at the types
isComplex :: [LispVal] -> ThrowsError LispVal
isComplex ([Complex _]) = return $ Bool True
isComplex ([Number _]) = return $ Bool True
isComplex ([Rational _]) = return $ Bool True
isComplex ([Float _]) = return $ Bool True
isComplex _ = return $ Bool False

-- |Predicate to determine if given number is a real.
--  Keep in mind this does not just look at the types
isReal :: [LispVal] -> ThrowsError LispVal
isReal ([Number _]) = return $ Bool True
isReal ([Rational _]) = return $ Bool True
isReal ([Float _]) = return $ Bool True
isReal ([Complex c]) = do
  imagPt <- numImagPart [(Complex c)]
  isExact <- isNumExact [imagPt]
  isZero <- numBoolBinopEq [imagPt, (Number 0)]
  case (isExact, isZero) of
    (Bool True, Bool True) -> return $ Bool True
    _ -> return $ Bool False
isReal _ = return $ Bool False

-- |Predicate to determine if given number is a rational.
--  Keep in mind this does not just look at the types
isRational :: [LispVal] -> ThrowsError LispVal
isRational ([Number _]) = return $ Bool True
isRational ([Rational _]) = return $ Bool True
isRational ([Float n]) = return $ Bool $ not $ isInfinite n
isRational _ = return $ Bool False

-- |Predicate to determine if given number is an integer.
--  Keep in mind this does not just look at the types;
--  a floating point input value can return true, for example.
isInteger :: [LispVal] -> ThrowsError LispVal
isInteger ([Number _]) = return $ Bool True
isInteger ([Complex n]) = do
  return $ Bool $ (isFloatAnInteger $ Float $ realPart n) && (isFloatAnInteger $ Float $ imagPart n)
isInteger ([Rational n]) = do
    let numer = abs $ numerator n
    let denom = abs $ denominator n
    return $ Bool $ (numer >= denom) && ((mod numer denom) == 0)
isInteger ([n@(Float _)]) = return $ Bool $ isFloatAnInteger n
isInteger _ = return $ Bool False

-- |A utility function to determine if given value is a floating point
--  number representing an whole number (integer).
isFloatAnInteger :: LispVal -> Bool
isFloatAnInteger (Float n) =
    ((floor n) :: Integer) == ((ceiling n) :: Integer)
isFloatAnInteger _ = False

-- - end Numeric operations section ---

-- |Extract an integer from the given value, throwing a type error if
--  the wrong type is passed.
unpackNum :: LispVal -> ThrowsError Integer
unpackNum (Number n) = return n
unpackNum notNum = throwError $ TypeMismatch "number" notNum


{-# LANGUAGE BangPatterns #-}
module FourierSeries where

import           Control.Monad              as M
import           Data.Array.Repa            as R
import           Data.Complex
import           Data.List                  as L
import           Data.Vector.Storable       as VS
import           Data.Vector.Unboxed        as VU
import           FokkerPlanck.FourierSeries
import           Foreign.CUDA.BLAS          as BLAS
import           Foreign.CUDA.Driver        as CUDA
import           Image.IO
import           Numeric.LinearAlgebra      as NL hiding (All)
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.BLAS
import           Utils.Distribution
import           Utils.Parallel
import           Utils.SimpsonRule
import           Utils.Time

main = do
  args@(gpuIDStr:sizeStr:deltaStr:periodStr:freqStr:angularFreqStr:radialFreqStr:stdStr:_) <-
    getArgs
  printCurrentTime . show $ args
  let gpuID = read gpuIDStr :: Int
      numPoints = read sizeStr :: Int
      delta = read deltaStr :: Double
      period = read periodStr :: Double
      freq = read freqStr :: Int
      numFreq = 2 * freq + 1
      angularFreq = read angularFreqStr :: Int
      radialFreq = read radialFreqStr :: Int
      std = read stdStr :: Double
      folderPath = "output/test/FourierSeries"
      len' = round $ (fromIntegral numPoints / delta :: Double)
      len =
        if odd len'
          then len'
          else len' + 1
      center = div len 2
      !numAngularFreq = 2 * angularFreq + 1
      !numRadialFreq = 2 * radialFreq + 1
      !numPinwheels = numAngularFreq * numRadialFreq
      pinwheel =
        fromFunction (Z :. len :. len :. numAngularFreq :. numRadialFreq) $ \(Z :. i :. j :. af :. rf) ->
          let x = delta * (fromIntegral $ i - center)
              y = delta * (fromIntegral $ j - center)
          in fourierMellin 0.5 (af - angularFreq) (rf - radialFreq) (x, y)
      rectangularHarmonics' =
        fromFunction (Z :. numFreq :. numFreq :. len :. len) $ \(Z :. xFreq :. yFreq :. i :. j) ->
          let !x = fromIntegral $ i - center
              !y = fromIntegral $ j - center
          in cis
               ((-2) * pi * delta *
                ((fromIntegral $ xFreq - freq) * x +
                 (fromIntegral $ yFreq - freq) * y) /
                period)
      simpsonWeights =
        computeUnboxedS $ computeWeightArrFromListOfShape [len, len] :: R.Array U DIM2 (Complex Double)
      inverseR2Harmonics' =
        fromFunction (Z :. numPoints :. numPoints :. numFreq :. numFreq) $ \(Z :. i :. j :. xFreq :. yFreq) ->
          cis
            (2 * pi *
             fromIntegral
               ((xFreq - freq) * (i - (div numPoints 2)) +
                (yFreq - freq) * (j - (div numPoints 2))) /
             period)
      !weight = (delta / 3) ^ 2 :+ 0
  createDirectoryIfMissing True folderPath
  createDirectoryIfMissing True (folderPath </> "Recons")
  createDirectoryIfMissing True (folderPath </> "Coefficients")
  print len
  !weightedPinwheels <-
    computeUnboxedP . R.traverse2 pinwheel simpsonWeights const $ \fP fS idx@(Z :. a :. b :. _ :. _) ->
      fP idx * fS (Z :. a :. b)
  !rectangularHarmonics <- computeUnboxedP rectangularHarmonics'
  !inverseR2Harmonics <- computeUnboxedP inverseR2Harmonics'
  -- printCurrentTime "CPU starts .."
  -- let !gaussianCoefficients =
  --       computeUnboxedS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
  --         weight *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  -- M.mapM_
  --   (\(af, rf) -> do
  --      let coefficients =
  --            sumS .
  --            sumS .
  --            R.traverse2
  --              rectangularHarmonics
  --              (R.slice weightedPinwheels (Z :. All :. All :. af :. rf))
  --              const $ \fR fP idx@(Z :. _ :. _ :. i :. j) ->
  --              fR idx * fP (Z :. i :. j)
  --          recon =
  --            sumS .
  --            sumS .
  --            R.traverse2
  --              inverseR2Harmonics
  --              (coefficients *^ gaussianCoefficients)
  --              const $ \fR fC idx@(Z :. _ :. _ :. i :. j) ->
  --              fR idx * fC (Z :. i :. j)
  --      plotImageRepaComplex
  --        (folderPath </> "Recon" </> printf "Recon%03d.png" (af * numRadialFreq + rf + 1)) .
  --        ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --        recon)
  --   [(af, rf) | af <- [0 .. numAngularFreq - 1], rf <- [0 .. numRadialFreq - 1]]
  -- printCurrentTime "Done"
  -- coefficients <-
  --   (sumP . R.traverse2 rectangularHarmonics weightedPinwheel const $ \fR fP idx@(Z :. _ :. _ :. i :. j) ->
  --      fR idx * fP (Z :. i :. j)) >>=
  --   sumP
  --     gaussianCoefficients =
  --       computeUnboxedS . R.traverse coefficients id $ \f idx@(Z :. xFreq :. yFreq) ->
  --         weight * (f idx) *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  -- recon <-
  --   (sumP . R.traverse2 inverseR2Harmonics gaussianCoefficients const $ \fR fC idx@(Z :. _ :. _ :. i :. j) ->
  --      fR idx * fC (Z :. i :. j)) >>=
  --   sumP
  -- -- Repa
  -- let coefficients =
  --       sumS . sumS . R.traverse2 rectangularHarmonics weightedPinwheel const $ \fR fP idx@(Z :. _ :. _ :. i :. j) ->
  --         fR idx * fP (Z :. i :. j)
  --     gaussianCoefficients =
  --       computeUnboxedS . R.traverse coefficients id $ \f idx@(Z :. xFreq :. yFreq) ->
  --         weight * (f idx) *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  --     recon =
  --       sumS . sumS . R.traverse2 inverseR2Harmonics gaussianCoefficients const $ \fR fC idx@(Z :. _ :. _ :. i :. j) ->
  --         fR idx * fC (Z :. i :. j)
  -- -- HMatrix
  -- let rectangularHarmonicsMat =
  --       ((numFreq ^ 2) >< (len ^ 2)) . R.toList $ rectangularHarmonics
  --     inverseR2HarmonicsMat =
  --       ((numPoints ^ 2) >< (numFreq ^ 2)) . R.toList $ inverseR2Harmonics
  --     weightedPinwheelVec = NL.fromList . R.toList $ weightedPinwheel
  --     gaussianCoefficientsVec =
  --       NL.fromList . R.toList . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
  --         weight *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  --     coefficientsVec = rectangularHarmonicsMat #> weightedPinwheelVec
  --     recon =
  --       fromListUnboxed (Z :. numPoints :. numPoints) . NL.toList $
  --       inverseR2HarmonicsMat #> (coefficientsVec * gaussianCoefficientsVec)
  -- --BLAS mxv
  -- let rectangularHarmonicsMat = VU.convert . toUnboxed $ rectangularHarmonics
  --     inverseR2HarmonicsMat = VU.convert . toUnboxed $ inverseR2Harmonics
  --     weightedPinwheelVec = VU.convert . toUnboxed $ weightedPinwheel
  --     gaussianCoefficientsVec =
  --       VU.convert .
  --       toUnboxed . computeS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
  --         weight *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  -- coefficientsVec <-
  --   mxv (numFreq ^ 2) (len ^ 2) rectangularHarmonicsMat weightedPinwheelVec
  -- reconVec <-
  --   mxv
  --     (numPoints ^ 2)
  --     (numFreq ^ 2)
  --     inverseR2HarmonicsMat
  --     (VS.zipWith (*) coefficientsVec gaussianCoefficientsVec)
  -- let recon = fromUnboxed (Z :. numPoints :. numPoints) . VS.convert $ reconVec
  --BLAS mxm
  -- let !rectangularHarmonicsMat = VU.convert . toUnboxed $ rectangularHarmonics
  --     !inverseR2HarmonicsMat = VU.convert . toUnboxed $ inverseR2Harmonics
  --     !weightedPinwheelVec = VU.convert . toUnboxed $ weightedPinwheels
  --     !gaussianCoefficients =
  --       computeUnboxedS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
  --         weight *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  -- printCurrentTime "Start.."
  -- coefficientsVec <-
  --   gemmBLAS
  --     (numFreq ^ 2)
  --     numPinwheels
  --     (len ^ 2)
  --     rectangularHarmonicsMat
  --     weightedPinwheelVec
  -- printCurrentTime "coefficients done"
  -- let coefficientsArr =
  --       fromUnboxed (Z :. numFreq :. numFreq :. numAngularFreq :. numRadialFreq) .
  --       VS.convert $
  --       coefficientsVec
  --     weightedCoefficientsArr =
  --       computeS .
  --       R.traverse2
  --         coefficientsArr
  --         gaussianCoefficients
  --         const $ \f1 f2 idx@(Z :. i :. j :. _ :. _) ->
  --         f1 idx * f2 (Z :. i :. j)
  --     weightedCoefficients = VS.convert . toUnboxed $ weightedCoefficientsArr
  -- M.mapM_
  --   (\(af, rf) ->
  --      plotImageRepaComplex
  --        (folderPath </> "Coefficients" </>
  --         printf
  --           "Coef(%d,%d).png"
  --           (af - div numAngularFreq 2)
  --           (rf - div numRadialFreq 2)) .
  --      ImageRepa 8 .
  --      computeS .
  --      extend (Z :. (1 :: Int) :. All :. All) . R.slice coefficientsArr $
  --      (Z :. All :. All :. af :. rf))
  --   [(af, rf) | af <- [0 .. numAngularFreq - 1], rf <- [0 .. numRadialFreq - 1]]
  -- printCurrentTime "Done"
  -- reconVec <-
  --   gemmBLAS
  --     (numPoints ^ 2)
  --     numPinwheels
  --     (numFreq ^ 2)
  --     inverseR2HarmonicsMat
  --     weightedCoefficients
  -- let !recon =
  --       fromUnboxed
  --         (Z :. numPoints :. numPoints :. numAngularFreq :. numRadialFreq) .
  --       VS.convert $
  --       reconVec
  -- M.mapM_
  --   (\(af, rf) ->
  --      plotImageRepaComplex
  --        (folderPath </> "Recons" </>
  --         printf
  --           "Recon(%d,%d).png"
  --           (af - div numAngularFreq 2)
  --           (rf - div numRadialFreq 2)) .
  --      ImageRepa 8 .
  --      computeS . extend (Z :. (1 :: Int) :. All :. All) . R.slice recon $
  --      (Z :. All :. All :. af :. rf))
  --   [(af, rf) | af <- [0 .. numAngularFreq - 1], rf <- [0 .. numRadialFreq - 1]]
  -- printCurrentTime "Done"
  --BLAS mxmGPU
  let !rectangularHarmonicsMat = VU.convert . toUnboxed $ rectangularHarmonics
      !inverseR2HarmonicsMat = VU.convert . toUnboxed $ inverseR2Harmonics
      !weightedPinwheelVec = VU.convert . toUnboxed $ weightedPinwheels
      !gaussianCoefficients =
        computeUnboxedS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
          weight *
          (gaussian2D
             (fromIntegral $ xFreq - freq)
             (fromIntegral $ yFreq - freq)
             std :+
           0)
  initialise []
  dev <- device gpuID
  ctx <- CUDA.create dev []
  handle <- BLAS.create
  printCurrentTime "Start.."
  coefficientsVec <-
    gemmCuBLAS11
      handle
      (numFreq ^ 2)
      numPinwheels
      (len ^ 2)
      rectangularHarmonicsMat
      weightedPinwheelVec
  let coefficientsArr =
        fromUnboxed (Z :. numFreq :. numFreq :. numAngularFreq :. numRadialFreq) .
        VS.convert $
        coefficientsVec
      weightedCoefficients =
        VU.convert .
        toUnboxed .
        computeS . R.traverse2 coefficientsArr gaussianCoefficients const $ \f1 f2 idx@(Z :. i :. j :. _ :. _) ->
          f1 idx * f2 (Z :. i :. j)
  printCurrentTime "coefficients done"
  reconVec <-
    gemmCuBLAS11
      handle
      (numPoints ^ 2)
      numPinwheels
      (numFreq ^ 2)
      inverseR2HarmonicsMat
      weightedCoefficients
  let !recon =
        fromUnboxed (Z :. numPoints :. numPoints :. numPinwheels) . VS.convert $
        reconVec
  M.mapM_
    (\i ->
       plotImageRepaComplex
         (folderPath </> "Recons" </> printf "Recon%03d.png" (i + 1)) .
       ImageRepa 8 .
       computeS . extend (Z :. (1 :: Int) :. All :. All) . R.slice recon $
       (Z :. All :. All :. i))
    [0 .. numPinwheels - 1]
  printCurrentTime "Done"
  M.mapM_
    (\(af, rf) ->
       plotImageRepaComplex
         (folderPath </> "Coefficients" </>
          printf
            "Coef(%d,%d).png"
            (af - div numAngularFreq 2)
            (rf - div numRadialFreq 2)) .
       ImageRepa 8 .
       computeS .
       extend (Z :. (1 :: Int) :. All :. All) . R.slice coefficientsArr $
       (Z :. All :. All :. af :. rf))
    [(af, rf) | af <- [0 .. numAngularFreq - 1], rf <- [0 .. numRadialFreq - 1]]
  -- plotImageRepaComplex (folderPath </> "Grating.png") .
  --   ImageRepa 8 .
  --   computeS .
  --   extend (Z :. (1 :: Int) :. All :. All) . R.slice rectangularHarmonics $
  --   (Z :. (0 :: Int) :. freq :. All :. All)
  -- plotImageRepaComplex (folderPath </> "Pinwheel.png") .
  --   ImageRepa 8 .
  --   computeS . fromFunction (Z :. (1 :: Int) :. numPoints :. numPoints) $ \(Z :. _ :. i :. j) ->
  --   fourierMellin
  --     angularFreq
  --     radialFreq
  --     ( (fromIntegral $ i - (div numPoints 2))
  --     , (fromIntegral $ j - (div numPoints 2)))


module Main where

import Control.Monad
import Data.Complex
import Data.Int
import Graphics.UI.GLUT
import qualified Streamly as St
import qualified Streamly.Prelude as S

--import qualified Streamly.Prelude as S
--import qualified Control.Monad.Parallel as MP
iterations = 400

x // y = fromIntegral x / fromIntegral y

-- Divides [a] into [[a], [a], ...] with each sublist of length n,
-- except the last sublist which has length <= n.
chunkify n [] = []
chunkify n xs =
  let (xs', rest) = splitAt n xs
  in xs' : chunkify n rest

-- Converts a coordinate in screen space to a vertex.
pix2vert (Size w h) (x, y) =
  Vertex2 ((3 // w * fromIntegral x) - 2.0) ((2 // h * fromIntegral y) - 1.0)

-- List of all of the vertices that represent screen pixels.
vertices :: IO [Vertex2 GLfloat]
vertices =
  get windowSize >>= \(Size w h) ->
    return $ [pix2vert (Size w h) (x, y) | x <- [0 .. w - 1], y <- [0 .. h - 1]]

-- Gets the color for a number of iterations.
color3 r g b = Color3 r g b

getcolor :: Int -> Color3 Float
getcolor iter
  | iter == iterations = color3 0 0 0
  | otherwise = color3 (amt * 0.5) amt (amt * 0.5)
  where
    amt = iter // iterations

-- Returns the number of iterations <= the maximum iterations of the
-- Mandelbrot set at the given vertex.
mandel (Vertex2 r i) =
  length . takeWhile (\z -> magnitude z <= 2) . take iterations $
  iterate (\z -> z ^ 2 + (r :+ i)) 0

-- plots one point.
drawVert v = do
  color . getcolor $ mandel v
  vertex v

-- draws all the vertices in slices (to update the display while drawing).
display' :: [[Vertex2 GLfloat]] -> IO ()
display' chunks =
  let foo = S.each chunks :: St.ParallelT IO ([Vertex2 GLfloat]) -- O(n)
      bar =
        (\vs -> do
           renderPrimitive Points $ mapM_ drawVert vs
           flush -- ::[Vertex2 GLfloat] -> IO ()
         )
  in do S.mapM_ bar foo
        displayCallback $= display

-- draws the whole fractal
display :: IO ()
display = do
  clear [ColorBuffer]
  displayCallback $=
    (do v <- vertices
        let c = chunkify 256 v
        let (x:y:z:_) = chunkify 3 c
                                    --display' c
        St.runStreamT $
          St.liftIO (display' x) St.<|> St.liftIO (display' y) St.<|>
          St.liftIO (display' z))
  get currentWindow >>= postRedisplay

main = do
  getArgsAndInitialize
  initialDisplayMode $= [SingleBuffered, RGBMode]
  initialWindowSize $= Size 1200 1024
  initialWindowPosition $= Position 100 100
  createWindow "Mandelbrot"
  clearColor $= Color4 0 0 0 0
  matrixMode $= Projection
  loadIdentity
  ortho (-2) 1 (-1) 1 (-1) 1
  displayCallback $= display
  mainLoop
-- import Control.Parallel.Strategies
-- type Point = (Float, Float)
-- -- z = x + i y ; p = u + iv
-- -- next p z = z^2 + p -- order of application reversed for currying
-- next :: Point -> Point -> Point
-- next (x,y) (u,v) = (x * x - y * y + u, 2 * x * y + v)
-- mandelbrot :: Point -> [Point]
-- mandelbrot p = iterate (next p) (0,0)
-- julia :: Point -> Point -> [Point]
-- julia c = iterate (next c)
-- fairlyClose :: Point -> Bool
-- fairlyClose (u,v) = (u * u + v * v) < 100
-- inMandelbrotSet :: Point -> Bool
-- inMandelbrotSet p = all fairlyClose (mandelbrot p)
-- chooseColor :: [a] -> [Point] -> a
-- chooseColor palette = (palette !!) . length . take n . takeWhile fairlyClose
--                        where n = length palette - 1
-- type Image a = Point -> a
-- fracImage :: (Point -> [Point]) -> [a] -> Image a
-- fracImage fractal palette = chooseColor palette . fractal
-- type Grid a = [[a]]
-- grid :: Int -> Int -> Point -> Point -> Grid Point
-- grid c r (xmin, ymin) (xmax, ymax)
--   = [[(x,y) | x <- for c xmin xmax] | y <- for r ymin ymax]
-- for :: Int -> Float -> Float -> [Float]
-- for n min max = take n [min, min + delta ..]
--                 where delta = (max - min) / fromIntegral (n - 1)
-- pmap = parMap rseq
-- sample :: Grid Point -> Image a -> Grid a
-- sample points image = map (map image) points
-- draw ::
--   Grid Point
--   -> (Point -> [Point]) -- mandelbrot
--   -> [a] -- [Char]
--   -> (Grid a -> b) -- the charRender IO function
--   -> b
-- draw points fractal palette render
--   = render (sample points (fracImage fractal palette))
-- charPalette :: [Char]
-- charPalette = " ,.‘\"~:;o-!|?/<>X+={^O#%&@8*$"
-- --putStr . unlines
-- charRender :: Grid Char -> IO ()
-- charRender x = return ()--putStr . unlines
-- main :: IO ()
-- main = draw points mandelbrot charPalette charRender
--           where points = grid (1790 * 300) (1370 * 300) (-2.25,-1.5) (0.75,1.5)
-- figure2 = draw points(julia (0.32,0.043)) charPalette charRender
--        where points = grid 79 37 (-1.5,-1.5) (1.5,1.5)


import Data.Complex
import Seq

avg xs = (*rec) $ sum xs
  where rec = recip $ fromIntegral $ length xs

harmonic xs = (*rec) . recip . sum . map recip $ xs
  where rec = fromIntegral $ length xs

--TODO use unfold
--
--util function to iteratively generate a list from results
--nothing about this implies a particular "zero"
--scan2 :: Eq a => a -> (a -> (b,a)) -> a -> [b]
--scan2 z f x | x == z    = []
--            | otherwise = fist:scan2 z f sound
--  where (fist, sound) = f x

scan2' = scan2 []
--semi-stable geometric mean that doesn't use too many numbers at once
stableGeometric xs = product $ map ((**rec) . product) $ scan2' (splitAt 10) xs
  where rec = recip $ fromIntegral $ length xs

agm a b = map fst $ iterate (\(x,y) -> (arith x y, geo x y)) (a,b)
  where arith a b = (/2) $ a + b
        geo a b   = sqrt $ a * b

agm1i = realPart $ flip (!!) 10 $ agm (1 :+ 0) (0 :+ 1)


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Numeric.Loss.MeanSquaredError where

import Numeric.Loss
import Dependent.Size
import qualified Numeric.Vector.Sized as Sized
import Numeric.Vector.Sized(SizedVector',SizedVector)
import Data.Proxy
import GHC.TypeLits(KnownNat,Nat,natVal)

newtype MeanSquaredError (n :: Nat) = MeanSquaredError ()

instance forall n. KnownNat n => Loss (ZZ ::. n) (MeanSquaredError n) where
  -- forward (MeanSquaredError targets) x
  --   = (err, x)
  --     where
  --       err = (normalize . Sized.foldAll (+) 0 . Sized.map square) (Sized.zipWith (-) x (Sized.use targets))
  --       square x = x ^2
  --       n = fromInteger $ natVal (Proxy :: Proxy (n :: Nat))
  --       normalize total = Sized.map (\t -> t / n) total

  lossDerivative _ x targets = dx
    where
      dx = normalize (Sized.zipWith (-) x targets)
      n = fromInteger $ natVal (Proxy :: Proxy (n :: Nat))
      normalize zs = Sized.map (\z -> 2/n * z) zs


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE UnboxedTuples         #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Numeric.Matrix.Internal.Float () where

import qualified Control.Monad.ST        as ST
import qualified Numeric.DataFrame.ST    as ST
import           Numeric.Dimensions      (KnownDim)
import           Numeric.Matrix.Internal
import           Numeric.Matrix.LU
import           Numeric.Scalar.Internal
import           Numeric.Vector.Internal

{-# INLINE mkMat #-}
mkMat ::
  Float -> Float -> Float -> Float ->
  Float -> Float -> Float -> Float ->
  Float -> Float -> Float -> Float ->
  Float -> Float -> Float -> Float ->
  Mat44f
mkMat
  _11 _12 _13 _14
  _21 _22 _23 _24
  _31 _32 _33 _34
  _41 _42 _43 _44
  = ST.runST $ do
    df <- ST.newDataFrame
    ST.writeDataFrameOff df 0  $ scalar _11
    ST.writeDataFrameOff df 1  $ scalar _12
    ST.writeDataFrameOff df 2  $ scalar _13
    ST.writeDataFrameOff df 3  $ scalar _14
    ST.writeDataFrameOff df 4  $ scalar _21
    ST.writeDataFrameOff df 5  $ scalar _22
    ST.writeDataFrameOff df 6  $ scalar _23
    ST.writeDataFrameOff df 7  $ scalar _24
    ST.writeDataFrameOff df 8  $ scalar _31
    ST.writeDataFrameOff df 9  $ scalar _32
    ST.writeDataFrameOff df 10 $ scalar _33
    ST.writeDataFrameOff df 11 $ scalar _34
    ST.writeDataFrameOff df 12 $ scalar _41
    ST.writeDataFrameOff df 13 $ scalar _42
    ST.writeDataFrameOff df 14 $ scalar _43
    ST.writeDataFrameOff df 15 $ scalar _44
    ST.unsafeFreezeDataFrame df

instance HomTransform4 Float where
  {-# INLINE translate4 #-}
  translate4 (unpackV4# -> (# x, y, z, _ #)) = mkMat
    1 0 0 0
    0 1 0 0
    0 0 1 0
    x y z 1

  {-# INLINE translate3 #-}
  translate3 (unpackV3# -> (# x, y, z #)) = mkMat
    1 0 0 0
    0 1 0 0
    0 0 1 0
    x y z 1

  {-# INLINE rotateX #-}
  rotateX a = mkMat
    1 0 0 0
    0 c s 0
    0 n c 0
    0 0 0 1
    where
      c = cos a
      s = sin a
      n = -s

  {-# INLINE rotateY #-}
  rotateY a = mkMat
    c 0 n 0
    0 1 0 0
    s 0 c 0
    0 0 0 1
    where
      c = cos a
      s = sin a
      n = -s

  {-# INLINE rotateZ #-}
  rotateZ a = mkMat
    c s 0 0
    n c 0 0
    0 0 1 0
    0 0 0 1
    where
      c = cos a
      s = sin a
      n = -s

  {-# INLINE rotate #-}
  rotate (unpackV3# -> (# x, y, z #)) a = mkMat
    (c+xxv)  (yxv+zs) (zxv-ys) 0
    (xyv-zs) (c+yyv)  (zyv+xs) 0
    (xzv+ys) (yzv-xs) (c+zzv)  0
     0        0        0       1
    where
      c = cos a
      v = 1 - c -- v for versine
      s = sin a
      xxv = x * x * v
      xyv = x * y * v
      xzv = x * z * v
      yxv = xyv
      yyv = y * y * v
      yzv = y * z * v
      zxv = xzv
      zyv = yzv
      zzv = z * z * v
      xs = x * s
      ys = y * s
      zs = z * s

  {-# INLINE rotateEuler #-}
  rotateEuler x y z = mkMat
    (cy*cz)  (cx*sz+sx*sy*cz) (sx*sz-cx*sy*cz) 0
    (-cy*sz) (cx*cz-sx*sy*sz) (sx*cz+cx*sy*sz) 0
     sy      (-sx*cy)         (cx*cy)          0
     0        0                0               1
    where
      cx = cos x
      sx = sin x
      cy = cos y
      sy = sin y
      cz = cos z
      sz = sin z

  {-# INLINE lookAt #-}
  lookAt up cam foc = mkMat
    xb1 yb1 zb1 0
    xb2 yb2 zb2 0
    xb3 yb3 zb3 0
    tx  ty  tz  1
    where
      (# xb1, xb2, xb3 #) = unpackV3# xb
      (# yb1, yb2, yb3 #) = unpackV3# yb
      (# zb1, zb2, zb3 #) = unpackV3# zb
      zb = normalized $ cam - foc -- Basis vector for "backward", since +Z is behind the camera
      xb = normalized $ up `cross` zb -- Basis vector for "right"
      yb = zb `cross` xb -- Basis vector for "up"
      ncam = -cam
      tx = unScalar $ xb `dot` ncam
      ty = unScalar $ yb `dot` ncam
      tz = unScalar $ zb `dot` ncam

  {-# INLINE perspective #-}
  perspective n f fovy aspect = mkMat
    dpw 0   0   0
    0   dph 0   0
    0   0   a (-1)
    0   0   b   0
    where
      hpd = tan (fovy * 0.5) -- height/distance
      wpd = aspect * hpd; -- width/distance
      dph = recip hpd -- distance/height
      dpw = recip wpd -- distance/width
      nmf = n - f
      a = (n + f) / nmf
      b = 2 * n * f / nmf

  {-# INLINE orthogonal #-}
  orthogonal n f w h = mkMat
    iw 0  0 0
    0  ih 0 0
    0  0  a 0
    0  0  b 1
    where
      ih = 2 / h
      iw = 2 / w
      nmf = n - f
      a = 2 / nmf
      b = (n + f) / nmf

  {-# INLINE toHomPoint #-}
  toHomPoint (unpackV3# -> (# x, y, z #)) = vec4 x y z 1

  {-# INLINE toHomVector #-}
  toHomVector (unpackV3# -> (# x, y, z #)) = vec4 x y z 0

  {-# INLINE fromHom #-}
  fromHom (unpackV4# -> (# x, y, z, w #))
    | w == 0    = vec3 x y z
    | otherwise = vec3 (x/w) (y/w) (z/w)

instance KnownDim n => MatrixInverse Float n where
    inverse = inverseViaLU

instance KnownDim n => MatrixDeterminant Float n where
    det = detViaLU


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.StudentT
-- Copyright : (c) 2011 Aleksey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Student-T distribution
module Statistics.Distribution.StudentT (
    StudentT
    -- * Constructors
  , studentT
  , studentTE
  , studentTUnstandardized
    -- * Accessors
  , studentTndf
  ) where

import Control.Applicative
import Data.Aeson          (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary         (Binary(..))
import Data.Data           (Data, Typeable)
import GHC.Generics        (Generic)
import Numeric.SpecFunctions (
  logBeta, incompleteBeta, invIncompleteBeta, digamma)

import qualified Statistics.Distribution as D
import Statistics.Distribution.Transform (LinearTransform (..))
import Statistics.Internal

-- | Student-T distribution
newtype StudentT = StudentT { studentTndf :: Double }
                   deriving (Eq, Typeable, Data, Generic)

instance Show StudentT where
  showsPrec i (StudentT ndf) = defaultShow1 "studentT" ndf i
instance Read StudentT where
  readPrec = defaultReadPrecM1 "studentT" studentTE

instance ToJSON StudentT
instance FromJSON StudentT where
  parseJSON (Object v) = do
    ndf <- v .: "studentTndf"
    maybe (fail $ errMsg ndf) return $ studentTE ndf
  parseJSON _ = empty

instance Binary StudentT where
  put = put . studentTndf
  get = do
    ndf <- get
    maybe (fail $ errMsg ndf) return $ studentTE ndf

-- | Create Student-T distribution. Number of parameters must be positive.
studentT :: Double -> StudentT
studentT ndf = maybe (error $ errMsg ndf) id $ studentTE ndf

-- | Create Student-T distribution. Number of parameters must be positive.
studentTE :: Double -> Maybe StudentT
studentTE ndf
  | ndf > 0   = Just (StudentT ndf)
  | otherwise = Nothing

errMsg :: Double -> String
errMsg _ = modErr "studentT" "non-positive number of degrees of freedom"

instance D.Distribution StudentT where
  cumulative      = cumulative
  complCumulative = complCumulative

instance D.ContDistr StudentT where
  density    d@(StudentT ndf) x = exp (logDensityUnscaled d x) / sqrt ndf
  logDensity d@(StudentT ndf) x = logDensityUnscaled d x - log (sqrt ndf)
  quantile = quantile

cumulative :: StudentT -> Double -> Double
cumulative (StudentT ndf) x
  | x > 0     = 1 - 0.5 * ibeta
  | otherwise = 0.5 * ibeta
  where
    ibeta = incompleteBeta (0.5 * ndf) 0.5 (ndf / (ndf + x*x))

complCumulative :: StudentT -> Double -> Double
complCumulative (StudentT ndf) x
  | x > 0     = 0.5 * ibeta
  | otherwise = 1 - 0.5 * ibeta
  where
    ibeta = incompleteBeta (0.5 * ndf) 0.5 (ndf / (ndf + x*x))

logDensityUnscaled :: StudentT -> Double -> Double
logDensityUnscaled (StudentT ndf) x =
    log (ndf / (ndf + x*x)) * (0.5 * (1 + ndf)) - logBeta 0.5 (0.5 * ndf)

quantile :: StudentT -> Double -> Double
quantile (StudentT ndf) p
  | p >= 0 && p <= 1 =
    let x = invIncompleteBeta (0.5 * ndf) 0.5 (2 * min p (1 - p))
    in case sqrt $ ndf * (1 - x) / x of
         r | p < 0.5   -> -r
           | otherwise -> r
  | otherwise = modErr "quantile" $ "p must be in [0,1] range. Got: "++show p

instance D.MaybeMean StudentT where
  maybeMean (StudentT ndf) | ndf > 1   = Just 0
                           | otherwise = Nothing

instance D.MaybeVariance StudentT where
  maybeVariance (StudentT ndf) | ndf > 2   = Just $! ndf / (ndf - 2)
                               | otherwise = Nothing

instance D.Entropy StudentT where
  entropy (StudentT ndf) =
    0.5 * (ndf+1) * (digamma ((1+ndf)/2) - digamma(ndf/2))
    + log (sqrt ndf)
    + logBeta (ndf/2) 0.5

instance D.MaybeEntropy StudentT where
  maybeEntropy = Just . D.entropy

instance D.ContGen StudentT where
  genContVar = D.genContinuous

-- | Create an unstandardized Student-t distribution.
studentTUnstandardized :: Double -- ^ Number of degrees of freedom
                       -> Double -- ^ Central value (0 for standard Student T distribution)
                       -> Double -- ^ Scale parameter
                       -> LinearTransform StudentT
studentTUnstandardized ndf mu sigma
  | sigma > 0 = LinearTransform mu sigma $ studentT ndf
  | otherwise = modErr "studentTUnstandardized" $ "sigma must be > 0. Got: " ++ show sigma

modErr :: String -> String -> a
modErr fun msg = error $ "Statistics.Distribution.StudentT." ++ fun ++ ": " ++ msg


module Main where

import Test.Hspec
-- import Test.Hspec.QuickCheck

-- import qualified Statistics.Matrix as M
-- import qualified Statistics.Matrix.Types as T
-- import qualified Statistics.Matrix.Algorithms as Alg

import qualified Fixtures as F
import AlgorithmsSpec
-- import Utils

spec :: Spec
spec = describe "Q-R Decomposition" $ do
        it "Q x R returns the original matrix" $ do
          qrDecompositionInvariant F.matB `shouldBe` True
        it "Matrix Q is orthogonal for the QR factorization of an invertible matrix" $ do
          qrFirstOrthoInvariant F.matD`shouldBe` True
        it "Matrix R is triangular" $ do
          qrSecondTriInvariant F.matB `shouldBe` True

main :: IO ()
main = do
  hspec spec


{- |
Description :  Rate matrix helper functions
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

Maintainer  :  dominik.schrempf@gmail.com
Stability   :  unstable
Portability :  non-portable (not tested)

Some helper functions that come handy when working with rate matrices of
continuous-time discrete-state Markov processes.

* Changelog

-}

module RateMatrix where

import           Numeric.LinearAlgebra
import qualified Tools                 as T

-- | A rate matrix is just a real matrix.
type RateMatrix     = Matrix R

-- | A state is just an integer at the moment.
newtype State       = State { fromState :: Int }
                    deriving (Show,Read,Eq,Ord)

-- | A matrix of exchangeabilities, we have q = e * pi, where q is a rate
-- matrix, e is the exchangeability matrix and pi is the diagonal matrix
-- containing the stationary frequency distribution.
type ExchMatrix     = Matrix R

-- | Stationary distribution of a rate matrix.
type StationaryDist = Vector R

-- | Normalizes a Markov process generator such that one event happens per unit time.
normalizeRates :: StationaryDist -> RateMatrix -> RateMatrix
normalizeRates f m = scale (1.0 / totalRate) m
  where totalRate = norm_1 $ f <# T.matrixSetDiagToZero m

-- | Set the diagonal entries of a matrix such that the rows sum to 0.
setDiagonal :: RateMatrix -> RateMatrix
setDiagonal m = diagZeroes - diag (fromList rowSums)
  where diagZeroes = T.matrixSetDiagToZero m
        rowSums    = map norm_1 $ toRows diagZeroes

-- | Extract the exchangeability matrix from a rate matrix.
toExchMatrix :: RateMatrix -> StationaryDist -> ExchMatrix
toExchMatrix m f = m <> diag oneOverF
  where oneOverF = cmap (1.0/) f


{-# LANGUAGE ScopedTypeVariables #-}

{-|
Module      : Numeric.LinearAlgebra.Static.Util
Description : Functions to manipulate vectors and matrices
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.LinearAlgebra.Static.Util
    ( atM
    , atV
    , foldR
    , foldR1
    )
where

import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra         as LA
import           Data.Fin
import qualified Data.Vector.Storable          as V

---------------------------------------------------------------
-- Get data from vector/matrix with compile time bound check --
---------------------------------------------------------------

-- | Get a value from a vector at a given coordinate/index with guarantee
-- that the element index is within the length of the vector at compile
-- time.
atV
    :: forall n
     . (KnownNat n)
    => R n
    -> Fin n
    -> Double
atV v p = extract v LA.! (fromFin p)

-- | Get element from matrix at a given coordinate/index with guarantee that
-- the element index is within the shape of the matrix at compile time.
atM
    :: forall m n
     . (KnownNat m, KnownNat n)
    => L m n
    -> Fin m
    -> Fin n
    -> Double
atM mat m n = extract mat `LA.atIndex` (fromFin m, fromFin n)

-- | Fold a vector.
foldR :: KnownNat n => (b -> Double -> b) -> b -> R n -> b
foldR f initial x = V.foldl' f initial (extract x)

-- | Fold a vector without an initial element.
foldR1 :: KnownNat n => (Double -> Double -> Double) -> R n -> Double
foldR1 f x = V.foldl1' f (extract x)


{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module    : Statistics.Sample
-- Copyright : (c) 2008 Don Stewart, 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Commonly used sample statistics, also known as descriptive
-- statistics.

module Statistics.Sample
    (
    -- * Types
      Sample
    , WeightedSample
    -- * Descriptive functions
    , range

    -- * Statistics of location
    , mean
    , welfordMean
    , meanWeighted
    , harmonicMean
    , geometricMean

    -- * Statistics of dispersion
    -- $variance

    -- ** Functions over central moments
    , centralMoment
    , centralMoments
    , skewness
    , kurtosis

    -- ** Two-pass functions (numerically robust)
    -- $robust
    , variance
    , varianceUnbiased
    , meanVariance
    , meanVarianceUnb
    , stdDev
    , varianceWeighted
    , stdErrMean

    -- ** Single-pass functions (faster, less safe)
    -- $cancellation
    , fastVariance
    , fastVarianceUnbiased
    , fastStdDev

    -- * Joint distributions
    , covariance
    , correlation
    , pair
    -- * References
    -- $references
    ) where

import Statistics.Function (minMax)
import Statistics.Sample.Internal (robustSumVar, sum)
import Statistics.Types.Internal  (Sample,WeightedSample)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

-- Operator ^ will be overridden
import Prelude hiding ((^), sum)

-- | /O(n)/ Range. The difference between the largest and smallest
-- elements of a sample.
range :: (G.Vector v Double) => v Double -> Double
range s = hi - lo
    where (lo , hi) = minMax s
{-# INLINE range #-}

-- | /O(n)/ Arithmetic mean.  This uses Kahan-Babuška-Neumaier
-- summation, so is more accurate than 'welfordMean' unless the input
-- values are very large.
mean :: (G.Vector v Double) => v Double -> Double
mean xs = sum xs / fromIntegral (G.length xs)
{-# SPECIALIZE mean :: U.Vector Double -> Double #-}
{-# SPECIALIZE mean :: V.Vector Double -> Double #-}

-- | /O(n)/ Arithmetic mean.  This uses Welford's algorithm to provide
-- numerical stability, using a single pass over the sample data.
--
-- Compared to 'mean', this loses a surprising amount of precision
-- unless the inputs are very large.
welfordMean :: (G.Vector v Double) => v Double -> Double
welfordMean = fini . G.foldl' go (T 0 0)
  where
    fini (T a _) = a
    go (T m n) x = T m' n'
        where m' = m + (x - m) / fromIntegral n'
              n' = n + 1
{-# SPECIALIZE welfordMean :: U.Vector Double -> Double #-}
{-# SPECIALIZE welfordMean :: V.Vector Double -> Double #-}

-- | /O(n)/ Arithmetic mean for weighted sample. It uses a single-pass
-- algorithm analogous to the one used by 'welfordMean'.
meanWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> Double
meanWeighted = fini . G.foldl' go (V 0 0)
    where
      fini (V a _) = a
      go (V m w) (x,xw) = V m' w'
          where m' | w' == 0   = 0
                   | otherwise = m + xw * (x - m) / w'
                w' = w + xw
{-# INLINE meanWeighted #-}

-- | /O(n)/ Harmonic mean.  This algorithm performs a single pass over
-- the sample.
harmonicMean :: (G.Vector v Double) => v Double -> Double
harmonicMean = fini . G.foldl' go (T 0 0)
  where
    fini (T b a) = fromIntegral a / b
    go (T x y) n = T (x + (1/n)) (y+1)
{-# INLINE harmonicMean #-}

-- | /O(n)/ Geometric mean of a sample containing no negative values.
geometricMean :: (G.Vector v Double) => v Double -> Double
geometricMean = exp . mean . G.map log
{-# INLINE geometricMean #-}

-- | Compute the /k/th central moment of a sample.  The central moment
-- is also known as the moment about the mean.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
centralMoment :: (G.Vector v Double) => Int -> v Double -> Double
centralMoment a xs
    | a < 0  = error "Statistics.Sample.centralMoment: negative input"
    | a == 0 = 1
    | a == 1 = 0
    | otherwise = sum (G.map go xs) / fromIntegral (G.length xs)
  where
    go x = (x-m) ^ a
    m    = mean xs
{-# SPECIALIZE centralMoment :: Int -> U.Vector Double -> Double #-}
{-# SPECIALIZE centralMoment :: Int -> V.Vector Double -> Double #-}

-- | Compute the /k/th and /j/th central moments of a sample.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
centralMoments :: (G.Vector v Double) => Int -> Int -> v Double -> (Double, Double)
centralMoments a b xs
    | a < 2 || b < 2 = (centralMoment a xs , centralMoment b xs)
    | otherwise      = fini . G.foldl' go (V 0 0) $ xs
  where go (V i j) x = V (i + d^a) (j + d^b)
            where d  = x - m
        fini (V i j) = (i / n , j / n)
        m            = mean xs
        n            = fromIntegral (G.length xs)
{-# SPECIALIZE
    centralMoments :: Int -> Int -> U.Vector Double -> (Double, Double) #-}
{-# SPECIALIZE
    centralMoments :: Int -> Int -> V.Vector Double -> (Double, Double) #-}

-- | Compute the skewness of a sample. This is a measure of the
-- asymmetry of its distribution.
--
-- A sample with negative skew is said to be /left-skewed/.  Most of
-- its mass is on the right of the distribution, with the tail on the
-- left.
--
-- > skewness $ U.to [1,100,101,102,103]
-- > ==> -1.497681449918257
--
-- A sample with positive skew is said to be /right-skewed/.
--
-- > skewness $ U.to [1,2,3,4,100]
-- > ==> 1.4975367033335198
--
-- A sample's skewness is not defined if its 'variance' is zero.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
skewness :: (G.Vector v Double) => v Double -> Double
skewness xs = c3 * c2 ** (-1.5)
    where (c3 , c2) = centralMoments 3 2 xs
{-# SPECIALIZE skewness :: U.Vector Double -> Double #-}
{-# SPECIALIZE skewness :: V.Vector Double -> Double #-}

-- | Compute the excess kurtosis of a sample.  This is a measure of
-- the \"peakedness\" of its distribution.  A high kurtosis indicates
-- that more of the sample's variance is due to infrequent severe
-- deviations, rather than more frequent modest deviations.
--
-- A sample's excess kurtosis is not defined if its 'variance' is
-- zero.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
kurtosis :: (G.Vector v Double) => v Double -> Double
kurtosis xs = c4 / (c2 * c2) - 3
    where (c4 , c2) = centralMoments 4 2 xs
{-# SPECIALIZE kurtosis :: U.Vector Double -> Double #-}
{-# SPECIALIZE kurtosis :: V.Vector Double -> Double #-}

-- $variance
--
-- The variance — and hence the standard deviation — of a
-- sample of fewer than two elements are both defined to be zero.

-- $robust
--
-- These functions use the compensated summation algorithm of Chan et
-- al. for numerical robustness, but require two passes over the
-- sample data as a result.
--
-- Because of the need for two passes, these functions are /not/
-- subject to stream fusion.

data V = V {-# UNPACK #-} !Double {-# UNPACK #-} !Double

-- | Maximum likelihood estimate of a sample's variance.  Also known
-- as the population variance, where the denominator is /n/.
variance :: (G.Vector v Double) => v Double -> Double
variance samp
    | n > 1     = robustSumVar (mean samp) samp / fromIntegral n
    | otherwise = 0
    where
      n = G.length samp
{-# SPECIALIZE variance :: U.Vector Double -> Double #-}
{-# SPECIALIZE variance :: V.Vector Double -> Double #-}

-- | Unbiased estimate of a sample's variance.  Also known as the
-- sample variance, where the denominator is /n/-1.
varianceUnbiased :: (G.Vector v Double) => v Double -> Double
varianceUnbiased samp
    | n > 1     = robustSumVar (mean samp) samp / fromIntegral (n-1)
    | otherwise = 0
    where
      n = G.length samp
{-# SPECIALIZE varianceUnbiased :: U.Vector Double -> Double #-}
{-# SPECIALIZE varianceUnbiased :: V.Vector Double -> Double #-}

-- | Calculate mean and maximum likelihood estimate of variance. This
-- function should be used if both mean and variance are required
-- since it will calculate mean only once.
meanVariance ::  (G.Vector v Double) => v Double -> (Double,Double)
meanVariance samp
  | n > 1     = (m, robustSumVar m samp / fromIntegral n)
  | otherwise = (m, 0)
    where
      n = G.length samp
      m = mean samp
{-# SPECIALIZE meanVariance :: U.Vector Double -> (Double,Double) #-}
{-# SPECIALIZE meanVariance :: V.Vector Double -> (Double,Double) #-}

-- | Calculate mean and unbiased estimate of variance. This
-- function should be used if both mean and variance are required
-- since it will calculate mean only once.
meanVarianceUnb :: (G.Vector v Double) => v Double -> (Double,Double)
meanVarianceUnb samp
  | n > 1     = (m, robustSumVar m samp / fromIntegral (n-1))
  | otherwise = (m, 0)
    where
      n = G.length samp
      m = mean samp
{-# SPECIALIZE meanVarianceUnb :: U.Vector Double -> (Double,Double) #-}
{-# SPECIALIZE meanVarianceUnb :: V.Vector Double -> (Double,Double) #-}

-- | Standard deviation.  This is simply the square root of the
-- unbiased estimate of the variance.
stdDev :: (G.Vector v Double) => v Double -> Double
stdDev = sqrt . varianceUnbiased
{-# SPECIALIZE stdDev :: U.Vector Double -> Double #-}
{-# SPECIALIZE stdDev :: V.Vector Double -> Double #-}

-- | Standard error of the mean. This is the standard deviation
-- divided by the square root of the sample size.
stdErrMean :: (G.Vector v Double) => v Double -> Double
stdErrMean samp = stdDev samp / (sqrt . fromIntegral . G.length) samp
{-# SPECIALIZE stdErrMean :: U.Vector Double -> Double #-}
{-# SPECIALIZE stdErrMean :: V.Vector Double -> Double #-}

robustSumVarWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> V
robustSumVarWeighted samp = G.foldl' go (V 0 0) samp
    where
      go (V s w) (x,xw) = V (s + xw*d*d) (w + xw)
          where d = x - m
      m = meanWeighted samp
{-# INLINE robustSumVarWeighted #-}

-- | Weighted variance. This is biased estimation.
varianceWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> Double
varianceWeighted samp
    | G.length samp > 1 = fini $ robustSumVarWeighted samp
    | otherwise         = 0
    where
      fini (V s w) = s / w
{-# SPECIALIZE varianceWeighted :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE varianceWeighted :: V.Vector (Double,Double) -> Double #-}

-- $cancellation
--
-- The functions prefixed with the name @fast@ below perform a single
-- pass over the sample data using Knuth's algorithm. They usually
-- work well, but see below for caveats. These functions are subject
-- to array fusion.
--
-- /Note/: in cases where most sample data is close to the sample's
-- mean, Knuth's algorithm gives inaccurate results due to
-- catastrophic cancellation.

fastVar :: (G.Vector v Double) => v Double -> T1
fastVar = G.foldl' go (T1 0 0 0)
  where
    go (T1 n m s) x = T1 n' m' s'
      where n' = n + 1
            m' = m + d / fromIntegral n'
            s' = s + d * (x - m')
            d  = x - m

-- | Maximum likelihood estimate of a sample's variance.
fastVariance :: (G.Vector v Double) => v Double -> Double
fastVariance = fini . fastVar
  where fini (T1 n _m s)
          | n > 1     = s / fromIntegral n
          | otherwise = 0
{-# INLINE fastVariance #-}

-- | Unbiased estimate of a sample's variance.
fastVarianceUnbiased :: (G.Vector v Double) => v Double -> Double
fastVarianceUnbiased = fini . fastVar
  where fini (T1 n _m s)
          | n > 1     = s / fromIntegral (n - 1)
          | otherwise = 0
{-# INLINE fastVarianceUnbiased #-}

-- | Standard deviation.  This is simply the square root of the
-- maximum likelihood estimate of the variance.
fastStdDev :: (G.Vector v Double) => v Double -> Double
fastStdDev = sqrt . fastVariance
{-# INLINE fastStdDev #-}

-- | Covariance of sample of pairs. For empty sample it's set to
--   zero
covariance :: (G.Vector v (Double,Double), G.Vector v Double)
           => v (Double,Double)
           -> Double
covariance xy
  | n == 0    = 0
  | otherwise = mean $ G.zipWith (*)
                         (G.map (\x -> x - muX) xs)
                         (G.map (\y -> y - muY) ys)
  where
    n       = G.length xy
    (xs,ys) = G.unzip xy
    muX     = mean xs
    muY     = mean ys
{-# SPECIALIZE covariance :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE covariance :: V.Vector (Double,Double) -> Double #-}

-- | Correlation coefficient for sample of pairs. Also known as
--   Pearson's correlation. For empty sample it's set to zero.
correlation :: (G.Vector v (Double,Double), G.Vector v Double)
           => v (Double,Double)
           -> Double
correlation xy
  | n == 0    = 0
  | otherwise = cov / sqrt (varX * varY)
  where
    n       = G.length xy
    (xs,ys) = G.unzip xy
    (muX,varX) = meanVariance xs
    (muY,varY) = meanVariance ys
    cov = mean $ G.zipWith (*)
            (G.map (\x -> x - muX) xs)
            (G.map (\y -> y - muY) ys)
{-# SPECIALIZE correlation :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE correlation :: V.Vector (Double,Double) -> Double #-}

-- | Pair two samples. It's like 'G.zip' but requires that both
--   samples have equal size.
pair :: (G.Vector v a, G.Vector v b, G.Vector v (a,b)) => v a -> v b -> v (a,b)
pair va vb
  | G.length va == G.length vb = G.zip va vb
  | otherwise = error "Statistics.Sample.pair: vector must have same length"
{-# INLINE pair #-}

------------------------------------------------------------------------
-- Helper code. Monomorphic unpacked accumulators.

-- (^) operator from Prelude is just slow.
(^) :: Double -> Int -> Double
x ^ 1 = x
x ^ n = x * (x ^ (n-1))
{-# INLINE (^) #-}

-- don't support polymorphism, as we can't get unboxed returns if we use it.
data T = T {-# UNPACK #-}!Double {-# UNPACK #-}!Int

data T1 = T1 {-# UNPACK #-}!Int {-# UNPACK #-}!Double {-# UNPACK #-}!Double

{-

Consider this core:

with data T a = T !a !Int

$wfold :: Double#
               -> Int#
               -> Int#
               -> (# Double, Int# #)

and without,

$wfold :: Double#
               -> Int#
               -> Int#
               -> (# Double#, Int# #)

yielding to boxed returns and heap checks.

-}

-- $references
--
-- * Chan, T. F.; Golub, G.H.; LeVeque, R.J. (1979) Updating formulae
--   and a pairwise algorithm for computing sample
--   variances. Technical Report STAN-CS-79-773, Department of
--   Computer Science, Stanford
--   University. <ftp://reports.stanford.edu/pub/cstr/reports/cs/tr/79/773/CS-TR-79-773.pdf>
--
-- * Knuth, D.E. (1998) The art of computer programming, volume 2:
--   seminumerical algorithms, 3rd ed., p. 232.
--
-- * Welford, B.P. (1962) Note on a method for calculating corrected
--   sums of squares and products. /Technometrics/
--   4(3):419–420. <http://www.jstor.org/stable/1266577>
--
-- * West, D.H.D. (1979) Updating mean and variance estimates: an
--   improved method. /Communications of the ACM/
--   22(9):532–535. <http://doi.acm.org/10.1145/359146.359153>


{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Statistics
-- Copyright   :  (c) A. V. H. McPhail 2010, 2012, 2014
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Useful statistical functions
--
-----------------------------------------------------------------------------

module Numeric.Statistics (
                           Sample,Samples
                          , covarianceMatrix, correlationCoefficientMatrix
                          , meanList, meanArray, meanMatrix
                          , varianceList, varianceArray, varianceMatrix
                          --
                          , centre, cloglog, corcoeff, cut
                          , ranks, kendall, logit
                          , mahalanobis
                          , mode, moment
                          , ols, percentile, range
                          , run_count
                          , spearman, studentize
                          ) where

-----------------------------------------------------------------------------

import Numeric.LinearAlgebra hiding(range)
--import Numeric.LinearAlgebra.Data hiding(range)
--import Numeric.LinearAlgebra.Devel
import Prelude hiding ((<>))
import qualified Data.Array.IArray as I
import qualified Data.List as DL
import qualified Data.Vector.Generic as GV
--import qualified Data.Vector.Storable as SV
import Foreign.Storable
import Numeric.GSL.Statistics
import Numeric.GSL.Sort(sort)

-----------------------------------------------------------------------------

type Sample a = Vector a
type Samples a = I.Array Int (Vector a)

-----------------------------------------------------------------------------

-- | the covariance matrix
covarianceMatrix :: Samples Double              -- ^ the dimensions of data (each vector being one dimension)
                 -> Matrix Double               -- ^ the symmetric covariance matrix
covarianceMatrix d = let (s,f) = I.bounds d
                      in fromArray2D $ I.array ((s,s),(f,f)) $ concat $ map (\(x,y) -> let c = covariance (d I.! x) (d I.! y) in if x == y then [((x,y),c)] else [((x,y),c),((y,x),c)]) $ filter (\(x,y) -> x <= y) $ I.range ((s,s),(f,f))

-- | the correlation coefficient: (cov x y) / (std x) (std y)
correlationCoefficientMatrix :: Samples Double -> Matrix Double
correlationCoefficientMatrix d = let (s,f) = I.bounds d
                           in fromArray2D $ I.array ((s,s),(f,f)) $ concat $ map (\(x,y) -> let { x' = d I.! x ; y' = d I.! y ; c = (covariance x' y') / ((stddev x') * (stddev y')) } in if x == y then [((x,y),c)] else [((x,y),c),((y,x),c)]) $ filter (\(x,y) -> x <= y) $ I.range ((s,s),(f,f))

-----------------------------------------------------------------------------
-----------------------------------------------------------------------------
-----------------------------------------------------------------------------

-- | the mean of a list of vectors
meanList :: (Container Vector a, Num (Vector a), Fractional a) => [Sample a] -> Sample a
meanList []     = error "meanVectors: empty list"
meanList [s]    = s
meanList (s:ss) = let ln = fromIntegral $ length ss + 1
                  in scale (recip ln) $ foldl (+) s ss

-- | the mean of an array of vectors
meanArray :: (Container Vector a, Num (Vector a), Fractional a) => Samples a -> Sample a
meanArray a = meanList $ I.elems a

-- | the mean of a matrix with data series in rows
meanMatrix :: (Container Vector a, Num (Vector a), Fractional a) => Matrix a -> Sample a
meanMatrix a = meanList $ toRows a

-----------------------------------------------------------------------------

-- | the variance of a list of vectors
varianceList :: (Container Vector a, Floating (Vector a), Num a, Fractional a) => [Sample a] -> Sample a
varianceList []  = error "varianceList: empty list"
varianceList [s] = konst 0 (size s)
varianceList l   = let mxs = meanList (map (** 2) l)
                       msx = (meanList l) ** 2
                   in mxs - msx

-- | the variance of an array of vectors
varianceArray :: (Container Vector a, Floating (Vector a), Fractional a) => Samples a -> Sample a
varianceArray a = varianceList $ I.elems a

-- | the variance of a matrix with data series in rows
varianceMatrix :: (Container Vector a, Floating (Vector a), Fractional a) => Matrix a -> Sample a
varianceMatrix a = varianceList $ toRows a

-----------------------------------------------------------------------------
-----------------------------------------------------------------------------
-----------------------------------------------------------------------------

-- | centre the data to 0: (x - (mean x))
centre :: Vector Double -> Vector Double
centre v = v - (realToFrac (mean v))

-----------------------------------------------------------------------------

-- | complementary log-log function
--cloglog :: Vector Double -> Vector Double
cloglog :: Floating a => a -> a
cloglog v = - log (- (log v))

-----------------------------------------------------------------------------

-- | corcoeff = covariance x / (std dev x * std dev y)
corcoeff :: Vector Double -> Vector Double -> Double
corcoeff x y = (covariance x y)/((stddev x)*(stddev y))

-----------------------------------------------------------------------------

-- | cut numerical data into intervals, data must fall inside the bounds
cut :: Vector Double
    -> Vector Double -- ^ intervals
    -> Vector Int    -- ^ data indexed by bin
cut v c  = let c' = sort c
           in GV.map (\x -> cut_helper 0 x c') v
    where
      cut_helper j x d
          | j >= size d                       = error "Numeric.Statistics: cut: data point not within interval"
          | x >= (d `atIndex` j) && x <= (d `atIndex` (j+1)) = j
          | otherwise                       = cut_helper (j + 1) x d

-----------------------------------------------------------------------------

-- | return the rank of each element of the vector
--     multiple identical entries result in the average rank of those entries
--ranks :: Vector Double -> Vector Double
ranks :: (Fractional b, Storable b) => Vector Double -> Vector b
ranks v = let v' = sort v
          in GV.map (\x -> 1 + rank_helper x v') v
              where rank_helper x v' = let is = GV.elemIndices x v'
                                       in (realToFrac (GV.foldl (+) 0 is)) / (fromIntegral $ GV.length is)

-----------------------------------------------------------------------------

-- | kendall's rank correlation τ
kendall :: Vector Double -> Vector Double -> Matrix Double
kendall x y = let ln = size x
                  rx = ranks x
                  ry = ranks y
                  r = fromColumns [rx,ry]
                  m = signum $ (kronecker r (asColumn $ konst 1.0 ln)) - (kronecker (asRow $ konst 1.0 ln) r)
                  c = rows m - 1
              in correlationCoefficientMatrix $ I.listArray (0,c) (toColumns m)

-----------------------------------------------------------------------------

-- | (logit p) = log(p/(1-p))
--logit :: Vector Double -> Vector Double
logit :: (Floating b, Storable b)
        => Vector b -> Vector b
logit v =  GV.map (\x -> - (log ((1 / x) - 1))) v

-----------------------------------------------------------------------------

-- | the Mahalanobis D-square distance between samples
--     columns are components and rows are observations (uses pseudoinverse)
mahalanobis :: Samples Double        -- ^ the data set
            -> Maybe (Sample Double) -- ^ (Just sample) to be measured or use mean when Nothing
            -> Double                -- ^ D^2
mahalanobis x u = let (_,xr) = I.bounds x
                      xl     = I.elems x
                      s'     = pinv $ covarianceMatrix x
                      xu     = case u of
                                 Nothing -> fromList $ map mean xl
                                 Just m  -> m
                      xm     = fromRows $ map ((-) xu) $ toRows $ fromColumns xl
                      --um     = asColumn xu
                      --w      = ((tr' xm) <> xm + (tr' um) <> um)/(fromIntegral $ xr - 1)
                      --w'     = inv w
                  in ((xm <> s' <> (tr' xm)) `atIndex` (0,0))

-----------------------------------------------------------------------------

-- | a list of element frequencies
mode :: Vector Double -> [(Double,Integer)]
mode v = let w = sort v
         in DL.sortBy (\(_,n) (_,n') -> compare n' n) $ GV.foldr freqs [] w
            where freqs x []          = [(x,1)]
                  freqs x ((f,n):fns)
                      | f == x         = ((f,n+1):fns)
                      | otherwise     = ((x,1):(f,n):fns)

-----------------------------------------------------------------------------

-- | the p'th moment of a vector
moment :: Integral a
       => a             -- ^ moment
       -> Bool          -- ^ calculate central moment
       -> Bool          -- ^ calculate absolute moment
       -> Vector Double -- ^ data
       -> Double
moment p c a v
    | p <= 0     = error "Numeric.Statistics.moment: negative moment requested"
--    | p == 1     = mean v
--    | p == 2     = variance v -- gives sample variance
    | otherwise = let u = if c then centre v else v
                      w = if a then abs u else u
                      x = GV.map (** (fromIntegral p)) w
                  in mean x

-----------------------------------------------------------------------------

-- | ordinary least squares estimation for the multivariate model
--   Y = X B + e        rows are observations, columns are elements
--   mean e = 0, cov e = kronecker s I
ols :: (Num (Vector t), Field t)
      => Matrix t         -- ^ X
    -> Matrix t           -- ^ Y
    -> (Matrix t, Matrix t, Matrix t) -- ^ (OLS estimator for B, OLS estimator for s, OLS residuals)
ols x y
    | rows x /= rows y = error "Numeric.Statistics: ols: incorrect matrix dimensions"
    | otherwise       = let (xr,xc) = (rows x,cols x)
                            (yr,yc) = (rows y,cols y)
                            z = (tr' x) <> x
                            r = rank z
                            beta = if r == xc
                                      then (inv z) <> (tr' x) <> y
                                      else (pinv x) <> y
                            rr = y - x <> beta
                            sigma = ((tr' rr) <> rr) / (fromIntegral $ xr - r)
                        in (beta,rr,sigma)

-----------------------------------------------------------------------------

-- | compute quantiles in percent
percentile :: Double        -- ^ percentile (0 - 100)
           -> Vector Double -- ^ data
           -> Double        -- ^ result
percentile p d = quantile (0.01*p) d

-----------------------------------------------------------------------------

-- | the difference between the maximum and minimum of the input
range :: (Container c e, Num e) => c e -> e
range v = maxElement v - minElement v

-----------------------------------------------------------------------------

-- | count the number of runs greater than or equal to @n@ in the data
run_count :: (Num a, Num t, Ord b, Ord a, Container Vector b)
            => a             -- ^ longest run to count
          -> Vector b        -- ^ data
          -> [(a, t)]        -- ^ [(run length,count)]
run_count n v = let w = subVector 1 (size v - 1) v
                    x = GV.foldr run_count' [(1,v `atIndex` 0)] w
                    y = map fst x
                    z = takeWhile (<= n) $  DL.sort y
                in foldr count [] z
    where run_count' m ((c,g):cs)
              | m < g             = ((c+1,m):cs)
              | otherwise         = ((1,m):(c,g):cs)
          run_count' _ [] = error "Lib.Numeric.Statistics.run_count.run_count'"
          count x []           = [(x,1)]
          count x ((yv,yc):ys)
              | x == yv         = ((yv,yc+1):ys)
              | otherwise      = ((x,1):(yv,yc):ys)

-----------------------------------------------------------------------------

-- | Spearman's rank correlation coefficient
spearman :: Vector Double -> Vector Double -> Double
spearman x y = corcoeff (ranks x) (ranks y)

-----------------------------------------------------------------------------

-- | centre and normalise a vector
studentize :: Vector Double -> Vector Double
studentize x = (centre x)/(fromList $ [stddev x])

-----------------------------------------------------------------------------

--table

-----------------------------------------------------------------------------

-----------------------------------------------------------------------------


import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.LAPACK

m :: Matrix Double
m = (3><3)
  [7.589183,1.703609,-4.477162,
    -4.597851,9.434889,-6.543450,
    0.4588202,-6.115153,1.331191]

v :: Matrix Double
v = (3><1)
  [1.745005,-4.448092,-4.160842]


module EnvelopeTest where

import           Control.Monad                  as M
import           Control.Monad.Parallel         as MP
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.Repa                as R
import           Data.Binary
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           Filter.Utils
import           FokkerPlanck.FourierSeries
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2D
import           Image.IO
import           Pinwheel.FourierSeries2D
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.List
import           Utils.Time

main = do
  args@(deviceIDsStr:numPointsStr:deltaStr:numPointsReconStr:deltaReconStr:numR2FreqStr:periodR2Str:phiFreqsStr:rhoFreqsStr:stdR2Str:stdStr:numBatchR2Str:numBatchR2FreqsStr:sStr:numThreadStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoints = read numPointsStr :: Int
      delta = read deltaStr :: Double
      numPointsRecon = read numPointsReconStr :: Int
      deltaRecon = read deltaReconStr :: Double
      numR2Freq = read numR2FreqStr :: Int
      periodR2 = read periodR2Str :: Double
      phiFreq = read phiFreqsStr :: Int
      rhoFreq = read rhoFreqsStr :: Int
      numThread = read numThreadStr :: Int
      stdR2 = read stdR2Str :: Double
      std = read stdStr :: Double
      numBatchR2 = read numBatchR2Str :: Int
      numBatchR2Freqs = read numBatchR2FreqsStr :: Int
      s = read sStr :: Double
      folderPath = "output/test/EnvelopeTest"
      func a =
        R.fromFunction (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) $ \(Z :. _ :. i :. j) ->
          let r =
                sqrt
                  (fromIntegral $
                   (i - div numPointsRecon 2) ^ 2 +
                   (j - div numPointsRecon 2) ^ 2) :: Double
          in if r == 0
               then 0
               else r ** a :+ 0
  createDirectoryIfMissing True folderPath
  plotImageRepaComplex (folderPath </> "Reference(-0.5).png") .
    ImageRepa 8 . computeS . func $
    (-0.5)
  plotImageRepaComplex (folderPath </> "Reference(-1).png") .
    ImageRepa 8 . computeS . func $
    (-1)
  plotImageRepaComplex (folderPath </> "Reference(-1.5).png") .
    ImageRepa 8 . computeS . func $
    (-1.5)
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  let envelope' =
        VU.convert . toUnboxed . computeS $
        analyticalFourierCoefficients2 numR2Freq 1 phiFreq rhoFreq (-s) periodR2
      envelope =
        VS.map
          (\x -> x - VS.sum envelope' / (fromIntegral $ numR2Freq ^ 2))
          envelope'
      envelopeMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        envelope
  plotImageRepaComplex (folderPath </> "EnvelopeFreq.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    envelope
  envelopeR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPointsRecon
      1
      periodR2
      deltaRecon
      numBatchR2
      envelopeMat
  plotImageRepaComplex (folderPath </> "Envelope.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    envelopeR2
      -- envelope2 = VS.map (^ 2) envelope
  let envelope2' =
        VS.map (^ 2) . VU.convert . toUnboxed . computeS $
        analyticalFourierCoefficients2 numR2Freq 1 phiFreq rhoFreq (-s) periodR2
      envelope2 =
        VS.map
          (\x -> x - VS.sum envelope2' / (fromIntegral $ numR2Freq ^ 2))
          envelope2'
      envelope2Mat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        envelope2
  plotImageRepaComplex (folderPath </> "Envelope2Freq.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    envelope2
  envelope2R2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPointsRecon
      1
      periodR2
      deltaRecon
      numBatchR2
      envelope2Mat
  plotImageRepaComplex (folderPath </> "Envelope2.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    envelope2R2
  plan <-
    makePlan
      folderPath
      emptyPlan
      numPointsRecon
      numPointsRecon
      numR2Freq
      (2 * phiFreq + 1)
      (2 * rhoFreq + 1)
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      -- bias =
      --   VU.convert . toUnboxed . computeUnboxedS $
      --   analyticalFourierCoefficients2 numR2Freq 1 0 0 stdR2 periodR2
      bias' =
        computeUnboxedS $
        analyticalFourierCoefficients2 numR2Freq 1 0 0 stdR2 periodR2
      bias =
        VU.convert .
        toUnboxed .
        computeUnboxedS .
        R.map (\x -> x - sumAllS bias' / (fromIntegral $ numR2Freq ^ 2)) $
        bias' :: VS.Vector (Complex Double)
      biasMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        bias
  plotImageRepaComplex (folderPath </> "BiasFreq.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    bias
  biasR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPointsRecon
      1
      periodR2
      deltaRecon
      numBatchR2
      biasMat
  plotImageRepaComplex (folderPath </> "Bias.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    biasR2
  let forwardPlanID = DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]
      backwardPlanID = DFTPlanID IDFT1DG [numR2Freq, numR2Freq] [0, 1]
      forwardPlan = getDFTPlan plan forwardPlanID
      backwardPlan = getDFTPlan plan backwardPlanID
      biasVec' =
        VU.convert .
        toUnboxed .
        computeUnboxedS .
        makeFilter2D . fromUnboxed (Z :. numR2Freq :. numR2Freq) . VS.convert $
        bias
      biasVec = VS.map (\x -> x - VS.sum biasVec' / (fromIntegral $ numR2Freq^2)) biasVec'
  biasF <- dftExecuteWithPlan forwardPlanID forwardPlan biasVec
  envelope2F <- dftExecuteWithPlan forwardPlanID forwardPlan envelope2
  biased' <-
    dftExecuteWithPlan backwardPlanID backwardPlan $
    VS.zipWith (*) biasF envelope2F
  let biased =
        VS.map
          (\x -> x - VS.sum biased' / (fromIntegral $ numR2Freq ^ 2))
          biased'
      biasedMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        biased
  plotImageRepaComplex (folderPath </> "BiasedFreq.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    biased
  biasedR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPointsRecon
      1
      periodR2
      deltaRecon
      numBatchR2
      biasedMat
  plotImageRepaComplex (folderPath </> "Biased.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    biasedR2