text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      :  Spiral.Exp
-- Copyright   :  (c) 2016-2017 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Spiral.Exp (
    Var(..),
    Const(..),
    Exp(..),
    Unop(..),
    Binop(..),
    BBinop(..),

    Type(..),
    Typed(..),

    ToConst(..),
    RealFloatConst,
    fromConst,
    fromComplex,
    lower,

    LiftNum(..),
    LiftFloating(..),

    IfThenElse(..),
    IsEq(..),
    IsOrd(..),

    intE,
    complexE,
    unComplexE,
    forceUnComplexE,

    isConstE,

    IsZeroOne(..),
    isI,
    isNegI,

    true,
    false
  ) where

import Data.Complex
import Data.Complex.Cyclotomic
import Data.Modular
import Data.Modular.Instances ()
import Data.Monoid ((<>))
import Data.Proxy (Proxy(..))
import Data.Ratio
import Data.String
import Data.Symbol
import GHC.TypeLits (KnownNat,
                     natVal)
import Language.C.Quote (ToIdent(..))
import Test.QuickCheck (Arbitrary(..))
import Text.PrettyPrint.Mainland
import Text.PrettyPrint.Mainland.Class

import Data.Heterogeneous
import Spiral.Globals
import Spiral.RootOfUnity
import Spiral.Util.Name
import Spiral.Util.Pretty
import Spiral.Util.Uniq

newtype Var = Var Name
  deriving (Eq, Ord, Show, IsString, ToIdent, Pretty, Named)

instance Gensym Var where
    gensymAt s _ = do
        u <- newUnique
        return $ Var $ Name (intern s) (Just u)

    uniquify (Var (Name sym _)) = do
        u <- newUnique
        return $ Var $ Name sym (Just u)

-- | Representation of scalar expressions.
data Const a where
    BoolC     :: Bool -> Const Bool
    IntC      :: Int -> Const Int
    IntegerC  :: Integer -> Const Integer
    FloatC    :: Float -> Const Float
    DoubleC   :: Double -> Const Double
    RationalC :: FractionalConst a => Rational -> Const a
    ComplexC  :: RealFloatConst a => Const a -> Const a -> Const (Complex a)

    -- | Constant multiple of root of unity tagged with n and k.
    W :: RootOfUnity (Const a) => !Int -> !Int -> Const a -> Const a

    -- | Cyclotomic numbers
    CycC :: RealFloatConst a => Cyclotomic -> Const (Complex a)

    -- | Multiple of $\pi$
    PiC :: (Floating a, ToConst a) => Rational -> Const a

    -- | Integers modulo a prime
    ModularC :: KnownNat p => ℤ/p -> Const (ℤ/p)

deriving instance Show (Const a)

instance Eq (Const a) where
    (==) = heq
    (/=) = hne

instance Ord (Const a) where
    compare = hcompare

class (Eq a, Pretty a) => ToConst a where
    toConst :: a -> Const a

instance ToConst Bool where
    toConst = BoolC

instance ToConst Int where
    toConst = IntC

instance ToConst Integer where
    toConst = IntegerC

instance ToConst Float where
    toConst x | isIntegral x = RationalC (fromIntegral (truncate x :: Int))
              | otherwise    = FloatC x

instance ToConst Double where
    toConst x | isIntegral x = RationalC (fromIntegral (truncate x :: Int))
              | otherwise    = DoubleC x

instance ToConst Rational where
    toConst = RationalC

instance RealFloatConst a => ToConst (Complex a) where
    toConst (r :+ i) = ComplexC (toConst r) (toConst i)

instance KnownNat p => ToConst (ℤ/p) where
    toConst = ModularC

class ( Fractional a
      , Typed a
      , ToConst a
      ) => FractionalConst a where

instance FractionalConst Rational where

instance FractionalConst Float where

instance FractionalConst Double where

instance FractionalConst (Complex Float) where

instance FractionalConst (Complex Double) where

class ( Eq a
      , RealFloat a
      , Pretty a
      , FractionalConst a
      , Floating (Const a)
      , Floating (Exp a)
      , Floating (Exp (Complex a))
      ) => RealFloatConst a where

instance RealFloatConst Float where

instance RealFloatConst Double where

-- | Convert a 'Const a' to the value of type 'a' that it represents.
fromConst :: Const a -> a
fromConst (BoolC x)      = x
fromConst (IntC x)       = x
fromConst (IntegerC x)   = x
fromConst (FloatC x)     = x
fromConst (DoubleC x)    = x
fromConst (RationalC x)  = fromRational x
fromConst (ComplexC r i) = fromConst r :+ fromConst i
fromConst (W _ _ x)      = fromConst x
fromConst (CycC x)       = toComplex x
fromConst (PiC x)        = pi*fromRational x
fromConst (ModularC x)   = x

lift :: ToConst a => (a -> a) -> Const a -> Const a
lift f (IntC x)      = IntC (f x)
lift f (IntegerC x)  = IntegerC (f x)
lift f (FloatC x)    = FloatC (f x)
lift f (DoubleC x)   = DoubleC (f x)
lift f x@RationalC{} = toConst (f (fromConst x))
lift f x@ComplexC{}  = toConst (f (fromConst x))
lift f x             = lift f (lower x)

lift2 :: ToConst a => (a -> a -> a) -> Const a -> Const a -> Const a
lift2 f (IntC x)      (IntC y)      = IntC (f x y)
lift2 f (IntegerC x)  (IntegerC y)  = IntegerC (f x y)
lift2 f (FloatC x)    (FloatC y)    = FloatC (f x y)
lift2 f (DoubleC x)   (DoubleC y)   = DoubleC (f x y)
lift2 f x@RationalC{} y@RationalC{} = toConst (f (fromConst x) (fromConst y))
lift2 f x@ComplexC{}  y@ComplexC{}  = toConst (f (fromConst x) (fromConst y))
lift2 f x             y             = lift2 f (lower x) (lower y)

instance Arbitrary (Const Int) where
    arbitrary = IntC <$> arbitrary

instance Arbitrary (Const Integer) where
    arbitrary = IntegerC <$> arbitrary

instance Arbitrary (Const Double) where
    arbitrary = DoubleC <$> arbitrary

instance Arbitrary (Const Float) where
    arbitrary = FloatC <$> arbitrary

instance Arbitrary (Const Rational) where
    arbitrary = RationalC <$> arbitrary

instance Arbitrary (Const (Complex Double)) where
    arbitrary = ComplexC <$> arbitrary <*> arbitrary

instance KnownNat p => Arbitrary (Const (ℤ/p)) where
    arbitrary = ModularC . toMod <$> arbitrary

instance Pretty (Const a) where
    pprPrec _ (BoolC x)     = if x then text "true" else text "false"
    pprPrec _ (IntC x)      = ppr x
    pprPrec _ (IntegerC x)  = ppr x
    pprPrec _ (FloatC x)    = ppr x <> char 'f'
    pprPrec _ (DoubleC x)   = ppr x
    pprPrec p (RationalC x)
      | denominator x == 1  = ppr (numerator x)
      | otherwise           = infixop p FDiv (numerator x) (denominator x)

    pprPrec p x@ComplexC{} = pprComplex p (fromConst x)

    pprPrec _ (W _ 0 _) = text "1"
    pprPrec _ (W n 1 _) = text "ω_" <> ppr n
    pprPrec _ (W n k _) = text "ω_" <> ppr n <> char '^' <> ppr k

    pprPrec p (CycC x) = text (showsPrec p x "")

    pprPrec p (PiC x) = parensIf (p > mulPrec) $
                        pprPrec mulPrec1 x <> char '*' <> text "pi"

    pprPrec p (ModularC x) = pprPrec p x

instance RootOfUnity (Const (Complex Float)) where
    rootOfUnity n k = mkW (k%n) (CycC (rootOfUnity n k))

instance RootOfUnity (Const (Complex Double)) where
    rootOfUnity n k = mkW (k%n) (CycC (rootOfUnity n k))

instance KnownNat p => RootOfUnity (Const (ℤ/p)) where
    rootOfUnity n k = ModularC (rootOfUnity n k)

pprComplex :: (Eq a, Num a, Pretty a) => Int -> Complex a -> Doc
pprComplex p (r :+ 0)    = pprPrec p r
pprComplex _ (0 :+ 1)    = char 'i'
pprComplex _ (0 :+ (-1)) = text "-i"
pprComplex _ (0 :+ i)    = pprPrec mulPrec i <> char 'i'
pprComplex p (r :+ i)    = parensIf (p > addPrec) $
                           pprPrec addPrec r <+> text "+" <+> pprPrec mulPrec i <> char 'i'

-- | Convert a 'Complex Double' to a 'Const (Complex Double)'
fromComplex :: RealFloatConst a => Complex a -> Const (Complex a)
fromComplex (r :+ i) = ComplexC (toConst r) (toConst i)

-- | Convert a constant into the most general value that can be used for exact
-- comparison.
exact :: forall m a . Monad m => Const a -> m (Const a)
exact x@BoolC{}     = return x
exact x@IntC{}      = return x
exact x@IntegerC{}  = return x
exact x@FloatC{}    = return x
exact x@DoubleC{}   = return x

exact x@(RationalC r) =
    case tau of
      FloatT     -> return $ FloatC $ fromConst x
      DoubleT    -> return $ DoubleC $ fromConst x
      RationalT  -> return x
      ComplexT{} -> return $ CycC $ fromRational r
      _          -> fail "Cannot represent exactly"
  where
    tau :: Type a
    tau = typeOf (undefined :: a)

exact (ComplexC cre cim) | isIntegral re && isIntegral im =
    return $ CycC $ gaussianRat (fromIntegral (round re :: Integer)) (fromIntegral (round im :: Integer))
  where
    re, im :: a ~ Complex b => b
    re = fromConst cre
    im = fromConst cim

exact (W _ _ x)    = exact x
exact x@CycC{}     = return x
exact x@ModularC{} = return x
exact _            = fail "Cannot represent exactly"

isIntegral :: RealFrac a => a -> Bool
isIntegral x = x == fromInteger (round x)

-- | Lower a constant's representation from a specialized representation to a
-- standard representation. This ensures the constant is not constructed with
-- the 'W' or 'CycC' data constructors.
lower :: forall a . Const a -> Const a
lower x@(RationalC r) =
    case tau of
      FloatT     -> FloatC $ fromConst x
      DoubleT    -> DoubleC $ fromConst x
      ComplexT{} -> lower $ ComplexC (RationalC r) 0
      _          -> x
  where
    tau :: Type a
    tau = typeOf (undefined :: a)

lower (ComplexC r i) = ComplexC (lower r) (lower i)
lower (W _ _ x)      = lower x
lower x@CycC{}       = toConst (fromConst x)
lower x@PiC{}        = toConst (fromConst x)
lower e              = e

-- | Representation of scalar constants.
data Exp a where
    ConstE :: Const a -> Exp a
    VarE   :: Var -> Exp a
    UnopE  :: Unop -> Exp a -> Exp a
    BinopE :: Num (Exp a) => Binop -> Exp a -> Exp a -> Exp a
    IdxE   :: Var -> [Exp Int] -> Exp a

    ComplexE :: (Typed a, Num (Exp a)) => Exp a -> Exp a -> Exp (Complex a)
    ReE :: RealFloatConst a => Exp (Complex a) -> Exp a
    ImE :: RealFloatConst a => Exp (Complex a) -> Exp a

    BBinopE :: (Typed a, ToConst a) => BBinop -> Exp a -> Exp a -> Exp Bool
    IfE     :: Exp Bool -> Exp a -> Exp a -> Exp a

deriving instance Show (Exp a)

instance Eq (Exp a) where
    (==) = heq
    (/=) = hne

instance Ord (Exp a) where
    compare = hcompare

instance IsString (Exp a) where
    fromString s = VarE (fromString s)

instance RootOfUnity (Exp (Complex Float)) where
    rootOfUnity n k = ConstE (rootOfUnity n k)

instance RootOfUnity (Exp (Complex Double)) where
    rootOfUnity n k = ConstE (rootOfUnity n k)

instance KnownNat p => RootOfUnity (Exp (ℤ/p)) where
    rootOfUnity n k = ConstE (rootOfUnity n k)

--
-- Heterogeneous quality and comparison
--

modPair :: forall p . KnownNat p => ℤ/p -> (Integer, Integer)
modPair x = (unMod x, natVal (Proxy :: Proxy p))

instance HEq Const where
    heq x y =
      case (exact x, exact y) of
        (Just x', Just y') -> go x' y'
        _                  -> go (lower x) (lower y)
      where
        go :: Const a -> Const b -> Bool
        go (BoolC x)        (BoolC y)        = x == y
        go (IntC x)         (IntC y)         = x == y
        go (IntegerC x)     (IntegerC y)     = x == y
        go (FloatC x)       (FloatC y)       = x == y
        go (DoubleC x)      (DoubleC y)      = x == y
        go (RationalC x)    (RationalC y)    = x == y
        go (ComplexC r1 i1) (ComplexC r2 i2) = (Some r1, Some i1) == (Some r2, Some i2)
        go (CycC x)         (CycC y)         = x == y
        go (ModularC x)     (ModularC y)     = modPair x == modPair y
        go _                _                = False

instance HEq Exp where
    heq (ConstE c1)           (ConstE c2)           = heq c1 c2
    heq (VarE x)              (VarE y)              = x == y
    heq (UnopE op1 e1)        (UnopE op2 e2)        = (op1, Some e1) == (op2, Some e2)
    heq (BinopE op1 e1a e1b)  (BinopE op2 e2a e2b)  = (op1, Some e1a, Some e1b) == (op2, Some e2a, Some e2b)
    heq (IdxE v1 es1)         (IdxE v2 es2)         = (v1, es1) == (v2, es2)
    heq (ComplexE r1 i1)      (ComplexE r2 i2)      = (Some r1, Some i1) == (Some r2, Some i2)
    heq (ReE e1)              (ReE e2)              = heq e1 e2
    heq (ImE e1)              (ImE e2)              = heq e1 e2
    heq (BBinopE op1 e1a e1b) (BBinopE op2 e2a e2b) = (op1, Some e1a, Some e1b) == (op2, Some e2a, Some e2b)
    heq (IfE e1a e1b e1c)     (IfE e2a e2b e2c)     = (Some e1a, Some e1b, Some e1c) == (Some e2a, Some e2b, Some e2c)
    heq _                     _                     = False

instance HOrd Const where
    hcompare x y =
      case (exact x, exact y) of
        (Just x', Just y') -> go x' y'
        _                  -> go (lower x) (lower y)
      where
        go :: Const a -> Const b -> Ordering
        go (BoolC x)        (BoolC y)        = compare x y
        go (IntC x)         (IntC y)         = compare x y
        go (IntegerC x)     (IntegerC y)     = compare x y
        go (FloatC x)       (FloatC y)       = compare x y
        go (DoubleC x)      (DoubleC y)      = compare x y
        go (RationalC x)    (RationalC y)    = compare x y
        go (ComplexC r1 i1) (ComplexC r2 i2) = compare (Some r1, Some i1) (Some r2, Some i2)
        go (CycC x)         (CycC y)         = compare (xr, xi) (yr, yi)
          where
            xr :+ xi = fromConst (CycC x :: Const (Complex Double))
            yr :+ yi = fromConst (CycC y :: Const (Complex Double))
        go (ModularC x)     (ModularC y)     = compare (modPair x) (modPair y)

        go x y = compare (tag x) (tag y)
          where
            tag :: Const a -> Int
            tag BoolC{}     = 0
            tag IntC{}      = 1
            tag IntegerC{}  = 2
            tag FloatC{}    = 3
            tag DoubleC{}   = 4
            tag RationalC{} = 5
            tag ComplexC{}  = 6
            tag W{}         = 7
            tag CycC{}      = 8
            tag PiC{}       = 9
            tag ModularC{}  = 10

instance HOrd Exp where
    hcompare (ConstE c1)           (ConstE c2)           = hcompare c1 c2
    hcompare (VarE x)              (VarE y)              = compare x y
    hcompare (UnopE op1 e1)        (UnopE op2 e2)        = compare (op1, Some e1)(op2, Some e2)
    hcompare (BinopE op1 e1a e1b)  (BinopE op2 e2a e2b)  = compare (op1, Some e1a, Some e1b) (op2, Some e2a, Some e2b)
    hcompare (IdxE v1 es1)         (IdxE v2 es2)         = compare (v1, es1) (v2, es2)
    hcompare (ComplexE r1 i1)      (ComplexE r2 i2)      = compare (Some r1, Some i1) (Some r2, Some i2)
    hcompare (ReE e1)              (ReE e2)              = hcompare e1 e2
    hcompare (ImE e1)              (ImE e2)              = hcompare e1 e2
    hcompare (BBinopE op1 e1a e1b) (BBinopE op2 e2a e2b) = compare (op1, Some e1a, Some e1b) (op2, Some e2a, Some e2b)
    hcompare (IfE e1a e1b e1c)     (IfE e2a e2b e2c)     = compare (Some e1a, Some e1b, Some e1c) (Some e2a, Some e2b, Some e2c)
    hcompare e1                    e2                    = compare (tag e1) (tag e2)
      where
        tag :: Exp a -> Int
        tag ConstE{}   = 0
        tag VarE{}     = 1
        tag UnopE{}    = 2
        tag BinopE{}   = 3
        tag IdxE{}     = 4
        tag ComplexE{} = 5
        tag ReE{}      = 6
        tag ImE{}      = 7
        tag BBinopE{}  = 8
        tag IfE{}      = 9

-- | Unary operators
data Unop = Neg
          | Abs
          | Signum
          | Exp
          | Log
          | Sqrt
          | Sin
          | Cos
          | Asin
          | Acos
          | Atan
          | Sinh
          | Cosh
          | Asinh
          | Acosh
          | Atanh
  deriving (Eq, Ord, Show)

-- | Binary operators
data Binop = Add
           | Sub
           | Mul
           | Quot
           | Rem
           | Div
           | Mod
           | FDiv
  deriving (Eq, Ord, Show, Enum)

-- | Boolean binary operators
data BBinop = Eq
            | Ne
            | Lt
            | Le
            | Ge
            | Gt
  deriving (Eq, Ord, Show, Enum)

instance HasFixity Unop where
    fixity Neg    = infixr_ 10
    fixity Abs    = infixr_ 10
    fixity Signum = infixr_ 10
    fixity Exp    = infixr_ 10
    fixity Log    = infixr_ 10
    fixity Sqrt   = infixr_ 10
    fixity Sin    = infixr_ 10
    fixity Cos    = infixr_ 10
    fixity Asin   = infixr_ 10
    fixity Acos   = infixr_ 10
    fixity Atan   = infixr_ 10
    fixity Sinh   = infixr_ 10
    fixity Cosh   = infixr_ 10
    fixity Asinh  = infixr_ 10
    fixity Acosh  = infixr_ 10
    fixity Atanh  = infixr_ 10

instance HasFixity Binop where
    fixity Add  = infixl_ 8
    fixity Sub  = infixl_ 8
    fixity Mul  = infixl_ 9
    fixity Quot = infixl_ 9
    fixity Rem  = infixl_ 9
    fixity Div  = infixl_ 9
    fixity Mod  = infixl_ 9
    fixity FDiv = infixl_ 9

instance HasFixity BBinop where
    fixity Eq = infix_ 4
    fixity Ne = infix_ 4
    fixity Lt = infix_ 4
    fixity Le = infix_ 4
    fixity Ge = infix_ 4
    fixity Gt = infix_ 4

instance Pretty Unop where
    ppr Neg     = text "-"
    ppr Abs     = text "abs" <> space
    ppr Signum  = text "signum" <> space
    ppr Exp     = text "exp" <> space
    ppr Log     = text "log" <> space
    ppr Sqrt    = text "sqrt" <> space
    ppr Sin     = text "sin" <> space
    ppr Cos     = text "cos" <> space
    ppr Asin    = text "asin" <> space
    ppr Acos    = text "acos" <> space
    ppr Atan    = text "atan" <> space
    ppr Sinh    = text "sinh" <> space
    ppr Cosh    = text "cosh" <> space
    ppr Asinh   = text "asinh" <> space
    ppr Acosh   = text "acosh" <> space
    ppr Atanh   = text "atanh" <> space

instance Pretty Binop where
    ppr Add  = text "+"
    ppr Sub  = text "-"
    ppr Mul  = text "*"
    ppr Quot = text "`quot`"
    ppr Rem  = text "`rem`"
    ppr Div  = text "`div`"
    ppr Mod  = text "`mod`"
    ppr FDiv = text "/"

instance Pretty BBinop where
    ppr Eq = text "=="
    ppr Ne = text "/="
    ppr Lt = text "<"
    ppr Le = text "<="
    ppr Ge = text ">="
    ppr Gt = text ">"

instance Pretty (Exp a) where
    pprPrec p (ConstE c) = pprPrec p c
    pprPrec p (VarE v)   = pprPrec p v

    pprPrec p (UnopE op e) =
        unop p op e

    pprPrec p (BinopE op e1 e2) =
        infixop p op e1 e2

    pprPrec _ (IdxE ev eis) =
        ppr ev <> mconcat [brackets (ppr ei) | ei <- eis]

    pprPrec p (ComplexE er ei) =
        pprComplex p (er :+ ei)

    pprPrec _ (ReE e) =
        text "re" <> parens (ppr e)

    pprPrec _ (ImE e) =
        text "im" <> parens (ppr e)

    pprPrec p (BBinopE op e1 e2) =
        infixop p op e1 e2

    pprPrec _ (IfE e1 e2 e3) =
        text "if" <+> ppr e1 <+>
        text "then" <+> ppr e2 <+>
        text "else" <+> ppr e3

-- | Representation of types.
data Type a where
    BoolT     :: Type Bool
    IntT      :: Type Int
    IntegerT  :: Type Integer
    FloatT    :: Type Float
    DoubleT   :: Type Double
    RationalT :: Type Rational
    ComplexT  :: RealFloatConst a => Type a -> Type (Complex a)
    ModPT     :: KnownNat p => Integer -> Type (ℤ/p)

deriving instance Eq (Type a)
deriving instance Show (Type a)

instance Pretty (Type a) where
    ppr BoolT          = text "bool"
    ppr IntT           = text "int"
    ppr IntegerT       = text "integer"
    ppr FloatT         = text "float"
    ppr DoubleT        = text "double"
    ppr RationalT      = text "rational"
    ppr (ComplexT tau) = text "complex" <+> ppr tau
    ppr (ModPT p)      = text "ℤ/" <> ppr p

class Typed a where
    typeOf :: a -> Type a

instance Typed Bool where
    typeOf _ = BoolT

instance Typed Int where
    typeOf _ = IntT

instance Typed Integer where
    typeOf _ = IntegerT

instance Typed Float where
    typeOf _ = FloatT

instance Typed Double where
    typeOf _ = DoubleT

instance Typed Rational where
    typeOf _ = RationalT

instance RealFloatConst a => Typed (Complex a) where
    typeOf _ = ComplexT (typeOf (undefined :: a))

instance KnownNat p => Typed (ℤ/p) where
    typeOf _ = ModPT (natVal (Proxy :: Proxy p))

--------------------------------------------------------------------------------
--
-- Num instances
--
--------------------------------------------------------------------------------

-- | Class to lift 'Num' operators.
class LiftNum b where
    -- | Lift a unary operation on 'Num' to the type 'b'
    liftNum :: Unop -> (forall a . Num a => a -> a) -> b -> b

    -- | Lift a binary operation on 'Num' to the type 'b
    liftNum2 :: Binop -> (forall a . Num a => a -> a -> a) -> b -> b -> b

instance (Typed a, Num a, ToConst a, Num (Const a)) => LiftNum (Const a) where
    liftNum Neg _ e
      | isZero e = 0

    liftNum Neg _ (ComplexC a b) = ComplexC (-a) (-b)
    liftNum Neg _ (CycC c)       = CycC (-c)

    liftNum Neg _ (W n k c) = mkW (k % n + 1 % 2) (-c)

    liftNum _op f (RationalC x) = RationalC (f x)
    liftNum op  f (W _ _ c)     = liftNum op f c
    liftNum _op f (ModularC x)  = ModularC (f x)

    liftNum _op f c = lift f c

    liftNum2 Add _ e1 e2
      | isZero e1 = e2
      | isZero e2 = e1

    liftNum2 Sub _ e1 e2
      | isZero e1 = -e2
      | isZero e2 = e1

    liftNum2 Mul _ e1 e2
      | isZero   e1 = 0
      | isZero   e2 = 0
      | isOne    e1 = e2
      | isOne    e2 = e1
      | isNegOne e1 = -e2
      | isNegOne e2 = -e1

    liftNum2 Add _ (ComplexC a b) (ComplexC c d) = ComplexC (a + c) (b + d)
    liftNum2 Sub _ (ComplexC a b) (ComplexC c d) = ComplexC (a - c) (b - d)
    liftNum2 Mul _ (ComplexC a b) (ComplexC c d) = ComplexC (a*c - b*d) (b*c + a*d)

    liftNum2 Mul _ (W n k x) (W n' k' y) = mkW (k % n + k' % n') (x*y)

    liftNum2 Add _ (PiC x) (PiC y) = PiC (x + y)
    liftNum2 Sub _ (PiC x) (PiC y) = PiC (x - y)

    liftNum2 Mul _ (PiC x) (RationalC y) = PiC (x * y)
    liftNum2 Mul _ (RationalC x) (PiC y) = PiC (x * y)

    liftNum2 _op f (RationalC x) (RationalC y) = RationalC (f x y)
    liftNum2 op  f (W _ _ c)     y             = liftNum2 op f c y
    liftNum2 op  f x             (W _ _ c)     = liftNum2 op f x c
    liftNum2 _op f (CycC x)      (CycC y)      = CycC (f x y)
    liftNum2 _op f (ModularC x)  (ModularC y)  = ModularC (f x y)

    liftNum2 op f x y
      | Just x'@CycC{} <- exact x, Just y'@CycC{} <- exact y = liftNum2 op f x' y'
      | otherwise                                            = lift2 f x y

--- XXX Need UndecidableInstances for this...but we *must* call
--- liftNum/liftNum2 on constants to ensure they are properly simplified.
--- The other option would be to forgo the Const instance and push it into the
--- Exp instance.
instance (Typed a, Num (Const a), LiftNum (Const a), Num (Exp a)) => LiftNum (Exp a) where
    liftNum op f (ConstE c) = ConstE $ liftNum op f c

    liftNum Neg _ e
      | isZero e = 0

    liftNum Neg _ (UnopE Neg x)      = x
    liftNum Neg _ (BinopE Sub e1 e2) = e2 - e1
    liftNum Neg _ (ComplexE a b)     = ComplexE (-a) (-b)

    liftNum op _ e = UnopE op e

    liftNum2 op f (ConstE c1) (ConstE c2) = ConstE $ liftNum2 op f c1 c2

    liftNum2 Add _ e1 e2
      | isZero e1 = e2
      | isZero e2 = e1

    liftNum2 Sub _ e1 e2
      | isZero e1 = -e2
      | isZero e2 = e1

    liftNum2 Mul _ e1 e2
      | isZero   e1 = 0
      | isZero   e2 = 0
      | isOne    e1 = e2
      | isOne    e2 = e1
      | isNegOne e1 = -e2
      | isNegOne e2 = -e1

    -- Simplify multiplication by i and -i
    liftNum2 Mul _ (ConstE k) (ComplexE a b)
      | isI k    = ComplexE (-b) a
      | isNegI k = ComplexE b (-a)

    liftNum2 Mul _ (ComplexE a b) (ConstE k)
      | isI k    = ComplexE (-b) a
      | isNegI k = ComplexE b (-a)

    -- Push negation out
    liftNum2 Add _ e1 (UnopE Neg e2) =
        e1 - e2

    liftNum2 Add _ (UnopE Neg e1) e2 =
        e2 - e1

    liftNum2 Sub _ e1 (UnopE Neg e2) =
        e1 + e2

    liftNum2 Sub _ (UnopE Neg e1) e2 =
        -(e1 + e2)

    liftNum2 Mul _ e1 (UnopE Neg e2) =
        -(e1 * e2)

    liftNum2 Mul _ (UnopE Neg e1) e2 =
        -(e1 * e2)

    -- Standard component-wise add/subtract
    liftNum2 Add _ (ComplexE a b) (ComplexE c d) = ComplexE (a + c) (b + d)
    liftNum2 Sub _ (ComplexE a b) (ComplexE c d) = ComplexE (a - c) (b - d)

    -- 3-multiply/5-add complex multiplication
    liftNum2 Mul _ (ComplexE a b) (ComplexE c d) | threeMults =
          ComplexE (t1 - t2) (t1 + t3)
      where
        t1 = a*(c+d)
        t2 = d*(b+a)
        t3 = c*(b-a)

    -- Usual 4-multiply/4-add complex multiplication
    liftNum2 Mul _ (ComplexE a b) (ComplexE c d) =
        ComplexE (a*c - b*d) (b*c + a*d)

    liftNum2 op _ e1 e2 = BinopE op e1 e2

instance Num (Const Int) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = IntC . fromInteger

instance Num (Const Integer) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = IntegerC

instance Num (Const Rational) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = RationalC . fromInteger

instance Num (Const Float) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = RationalC . fromInteger

instance Num (Const Double) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = RationalC . fromInteger

instance Num (Const (Complex Float)) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = RationalC . fromInteger

instance Num (Const (Complex Double)) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = RationalC . fromInteger

instance KnownNat p => Num (Const (ℤ/p)) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ModularC . fromInteger

instance Num (Exp Int) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ConstE . fromInteger

instance Num (Exp Integer) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ConstE . fromInteger

instance Num (Exp Float) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ConstE . fromInteger

instance Num (Exp Double) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ConstE . fromInteger

instance Num (Exp Rational) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ConstE . fromInteger

instance Num (Exp (Complex Float)) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ConstE . fromInteger

instance Num (Exp (Complex Double)) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ConstE . fromInteger

instance KnownNat p => Num (Exp (ℤ/p)) where
    (+) = liftNum2 Add (+)
    (-) = liftNum2 Sub (-)
    (*) = liftNum2 Mul (*)

    abs = liftNum Abs abs

    negate = liftNum Neg negate

    signum  = liftNum Signum signum

    fromInteger = ConstE . fromInteger

--------------------------------------------------------------------------------
--
-- Enum/Real/Integral instances
--
--------------------------------------------------------------------------------

-- | Class to lift 'Integral' operators.
class LiftIntegral b where
    -- | Lift a binary operation on 'LiftIntegral' to the type 'b
    liftIntegral2 :: Binop -> (forall a . Integral a => a -> a -> a) -> b -> b -> b

instance LiftIntegral (Const Int) where
    liftIntegral2 _ f (IntC x) (IntC y) = IntC (f x y)

    liftIntegral2 _op f c1 c2 = lift2 f c1 c2

instance Enum (Const Int) where
    toEnum n = IntC (fromIntegral n)

    fromEnum (IntC i)    = fromIntegral i
    fromEnum RationalC{} = error "can't happen"
    fromEnum PiC{}       = error "can't happen"
    fromEnum W{}         = error "can't happen"

instance Real (Const Int) where
    toRational (IntC i)    = toRational i
    toRational RationalC{} = error "can't happen"
    toRational PiC{}       = error "can't happen"
    toRational W{}         = error "can't happen"

instance Integral (Const Int) where
    quot = liftIntegral2 Quot quot
    rem  = liftIntegral2 Rem rem
    div  = liftIntegral2 Div div
    mod  = liftIntegral2 Mod mod

    x `quotRem` y = (x `quot` y, x `rem` y)

    x `divMod` y = (x `div` y, x `mod` y)

    toInteger (IntC i)    = fromIntegral i
    toInteger RationalC{} = error "can't happen"
    toInteger W{}         = error "can't happen"
    toInteger PiC{}       = error "can't happen"

instance (LiftIntegral (Const a), Integral (Exp a)) => LiftIntegral (Exp a) where
    liftIntegral2 op f (ConstE x) (ConstE y) = ConstE $ liftIntegral2 op f x y

    liftIntegral2 Quot _ (BinopE Mul  e1 e2) e2' | e2' == e2 = e1
    liftIntegral2 Rem  _ (BinopE Mul _e1 e2) e2' | e2' == e2 = 0

    liftIntegral2 Quot _ (BinopE Add (BinopE Mul  e1 e2) e3) e2' | e2' == e2 = e1 + e3 `quot` e2
    liftIntegral2 Rem  _ (BinopE Add (BinopE Mul _e1 e2) e3) e2' | e2' == e2 = e3 `rem` e2

    liftIntegral2 op _ e1 e2 = BinopE op e1 e2

instance Enum (Exp Int) where
    toEnum n = intE (fromIntegral n)

    fromEnum (ConstE (IntC i)) = fromIntegral i
    fromEnum _                 = error "Enum Exp: fromEnum not implemented"

instance Real (Exp Int) where
    toRational (ConstE (IntC i)) = toRational i
    toRational _                 = error "Real Exp: toRational not implemented"

instance Integral (Exp Int) where
    quot = liftIntegral2 Quot quot
    rem  = liftIntegral2 Rem rem
    div  = liftIntegral2 Div div
    mod  = liftIntegral2 Mod mod

    x `quotRem` y = (x `quot` y, x `rem` y)

    x `divMod` y = (x `div` y, x `mod` y)

    toInteger (ConstE (IntC i)) = fromIntegral i
    toInteger _                 = error "Integral Exp: toInteger not implemented"

--------------------------------------------------------------------------------
--
-- Fractional instances
--
--------------------------------------------------------------------------------

-- | Class to lift 'Fractional' operators.
class LiftFrac b where
    -- | Lift a binary operation on 'Num' to the type 'b
    liftFrac2 :: Binop -> (forall a . Fractional a => a -> a -> a) -> b -> b -> b

instance (Fractional a, ToConst a, Fractional (Const a)) => LiftFrac (Const a) where
    liftFrac2 FDiv _ c1 (W n k x) | isOne c1 = mkW (-k % n) (1/x)

    liftFrac2 FDiv _ (W n k x) (W n' k' y) = mkW (k % n - k' % n') (x/y)

    liftFrac2 FDiv _ (PiC x) (RationalC y) = PiC (x / y)

    liftFrac2 FDiv _ c1 c2 | isOne c2 =
        c1

    liftFrac2 _op f (RationalC x) (RationalC y) = RationalC (f x y)
    liftFrac2 op  f (W _ _ c)     y             = liftFrac2 op f c y
    liftFrac2 op  f x             (W _ _ c)     = liftFrac2 op f x c
    liftFrac2 _op f (CycC x)      (CycC y)      = CycC (f x y)
    liftFrac2 _op f (ModularC x)  (ModularC y)  = ModularC (f x y)

    liftFrac2 op f x y
      | Just x'@CycC{} <- exact x, Just y'@CycC{} <- exact y = liftFrac2 op f x' y'
      | otherwise                                            = lift2 f x y

instance (LiftFrac (Const a), Num (Exp a)) => LiftFrac (Exp a) where
    liftFrac2 op f (ConstE c1) (ConstE c2) = ConstE $ liftFrac2 op f c1 c2

    liftFrac2 FDiv _ e1 e2 | isOne e2 =
        e1

    liftFrac2 op _ e1 e2 = BinopE op e1 e2

instance Fractional (Const Rational) where
    (/) = liftFrac2 FDiv (/)

    fromRational = RationalC

instance Fractional (Const Float) where
    (/) = liftFrac2 FDiv (/)

    fromRational = RationalC

instance Fractional (Const Double) where
    (/) = liftFrac2 FDiv (/)

    fromRational = RationalC

instance Fractional (Const (Complex Float)) where
    (/) = liftFrac2 FDiv (/)

    fromRational = RationalC

instance Fractional (Const (Complex Double)) where
    (/) = liftFrac2 FDiv (/)

    fromRational = RationalC

instance KnownNat p => Fractional (Const (ℤ/p)) where
    (/) = liftFrac2 FDiv (/)

    fromRational = ModularC . fromRational

instance Fractional (Exp Float) where
    (/) = liftFrac2 FDiv (/)

    fromRational = ConstE . fromRational

instance Fractional (Exp Double) where
    (/) = liftFrac2 FDiv (/)

    fromRational = ConstE . fromRational

instance Fractional (Exp (Complex Float)) where
    (/) = liftFrac2 FDiv (/)

    fromRational = ConstE . fromRational

instance Fractional (Exp (Complex Double)) where
    (/) = liftFrac2 FDiv (/)

    fromRational = ConstE . fromRational

instance KnownNat p => Fractional (Exp (ℤ/p)) where
    (/) = liftFrac2 FDiv (/)

    fromRational = ConstE . fromRational

--------------------------------------------------------------------------------
--
-- Floating instances
--
--------------------------------------------------------------------------------

-- | Class to lift 'Floating' operators.
class LiftFloating b where
    -- | Lift a binary operation on 'Floating' to the type 'b
    liftFloating :: Unop -> (forall a . Floating a => a -> a) -> b -> b

instance (Floating a, Typed a, ToConst a) => LiftFloating (Const a) where
    liftFloating Sqrt f c@(RationalC r) =
        case tau of
          ComplexT{} -> CycC (sqrtRat r)
          _          -> lift f c
      where
        tau :: Type a
        tau = typeOf (undefined :: a)

    liftFloating Sin f c@(PiC r) =
        case tau of
          ComplexT{} -> CycC (sinRev (r / 2))
          _          -> lift f c
      where
        tau :: Type a
        tau = typeOf (undefined :: a)

    liftFloating Cos f c@(PiC r) =
        case tau of
          ComplexT{} -> CycC (cosRev (r / 2))
          _          -> lift f c
      where
        tau :: Type a
        tau = typeOf (undefined :: a)

    liftFloating _op f c = lift f c

instance LiftFloating (Const a) => LiftFloating (Exp a) where
    liftFloating op f (ConstE c) = ConstE $ liftFloating op f c

    liftFloating op _ e = UnopE op e

instance Floating (Const Float) where
    pi = PiC 1

    exp   = liftFloating Exp exp
    log   = liftFloating Log log
    sqrt  = liftFloating Sqrt sqrt
    asin  = liftFloating Asin asin
    acos  = liftFloating Acos acos
    atan  = liftFloating Atan atan
    sinh  = liftFloating Sinh sinh
    cosh  = liftFloating Cosh cosh
    asinh = liftFloating Asinh asinh
    acosh = liftFloating Acosh acosh
    atanh = liftFloating Atanh atanh
    sin   = liftFloating Sin sin
    cos   = liftFloating Cos cos

instance Floating (Const Double) where
    pi = PiC 1

    exp   = liftFloating Exp exp
    log   = liftFloating Log log
    sqrt  = liftFloating Sqrt sqrt
    asin  = liftFloating Asin asin
    acos  = liftFloating Acos acos
    atan  = liftFloating Atan atan
    sinh  = liftFloating Sinh sinh
    cosh  = liftFloating Cosh cosh
    asinh = liftFloating Asinh asinh
    acosh = liftFloating Acosh acosh
    atanh = liftFloating Atanh atanh
    sin   = liftFloating Sin sin
    cos   = liftFloating Cos cos

instance Floating (Const (Complex Float)) where
    pi = PiC 1

    exp   = liftFloating Exp exp
    log   = liftFloating Log log
    sqrt  = liftFloating Sqrt sqrt
    asin  = liftFloating Asin asin
    acos  = liftFloating Acos acos
    atan  = liftFloating Atan atan
    sinh  = liftFloating Sinh sinh
    cosh  = liftFloating Cosh cosh
    asinh = liftFloating Asinh asinh
    acosh = liftFloating Acosh acosh
    atanh = liftFloating Atanh atanh
    sin   = liftFloating Sin sin
    cos   = liftFloating Cos cos

instance Floating (Const (Complex Double)) where
    pi = PiC 1

    exp   = liftFloating Exp exp
    log   = liftFloating Log log
    sqrt  = liftFloating Sqrt sqrt
    asin  = liftFloating Asin asin
    acos  = liftFloating Acos acos
    atan  = liftFloating Atan atan
    sinh  = liftFloating Sinh sinh
    cosh  = liftFloating Cosh cosh
    asinh = liftFloating Asinh asinh
    acosh = liftFloating Acosh acosh
    atanh = liftFloating Atanh atanh
    sin   = liftFloating Sin sin
    cos   = liftFloating Cos cos

instance Floating (Exp Float) where
    pi = ConstE pi

    exp   = liftFloating Exp exp
    log   = liftFloating Log log
    sqrt  = liftFloating Sqrt sqrt
    asin  = liftFloating Asin asin
    acos  = liftFloating Acos acos
    atan  = liftFloating Atan atan
    sinh  = liftFloating Sinh sinh
    cosh  = liftFloating Cosh cosh
    asinh = liftFloating Asinh asinh
    acosh = liftFloating Acosh acosh
    atanh = liftFloating Atanh atanh
    sin   = liftFloating Sin sin
    cos   = liftFloating Cos cos

instance Floating (Exp Double) where
    pi = ConstE pi

    exp   = liftFloating Exp exp
    log   = liftFloating Log log
    sqrt  = liftFloating Sqrt sqrt
    asin  = liftFloating Asin asin
    acos  = liftFloating Acos acos
    atan  = liftFloating Atan atan
    sinh  = liftFloating Sinh sinh
    cosh  = liftFloating Cosh cosh
    asinh = liftFloating Asinh asinh
    acosh = liftFloating Acosh acosh
    atanh = liftFloating Atanh atanh
    sin   = liftFloating Sin sin
    cos   = liftFloating Cos cos

instance Floating (Exp (Complex Float)) where
    pi = ConstE pi

    exp   = liftFloating Exp exp
    log   = liftFloating Log log
    sqrt  = liftFloating Sqrt sqrt
    asin  = liftFloating Asin asin
    acos  = liftFloating Acos acos
    atan  = liftFloating Atan atan
    sinh  = liftFloating Sinh sinh
    cosh  = liftFloating Cosh cosh
    asinh = liftFloating Asinh asinh
    acosh = liftFloating Acosh acosh
    atanh = liftFloating Atanh atanh
    sin   = liftFloating Sin sin
    cos   = liftFloating Cos cos

instance Floating (Exp (Complex Double)) where
    pi = ConstE pi

    exp   = liftFloating Exp exp
    log   = liftFloating Log log
    sqrt  = liftFloating Sqrt sqrt
    asin  = liftFloating Asin asin
    acos  = liftFloating Acos acos
    atan  = liftFloating Atan atan
    sinh  = liftFloating Sinh sinh
    cosh  = liftFloating Cosh cosh
    asinh = liftFloating Asinh asinh
    acosh = liftFloating Acosh acosh
    atanh = liftFloating Atanh atanh
    sin   = liftFloating Sin sin
    cos   = liftFloating Cos cos

--------------------------------------------------------------------------------
--
-- Staged operations
--
--------------------------------------------------------------------------------

class IfThenElse a b where
    ifThenElse :: a -> b -> b -> b

instance IfThenElse Bool a where
    ifThenElse c t e = if c then t else e

infix 4 .==., ./=.

class IsEq a b | a -> b where
    (.==.) :: a -> a -> b
    (./=.) :: a -> a -> b

infix 4 .<., .<=., .>=., .>.

class IsEq a b => IsOrd a b where
    (.<.)  :: a -> a -> b
    (.<=.) :: a -> a -> b
    (.>=.) :: a -> a -> b
    (.>.)  :: a -> a -> b

instance (Typed a, ToConst a, Eq a) => IsEq (Exp a) (Exp Bool) where
    e1 .==. e2 = BBinopE Eq e1 e2
    e1 ./=. e2 = BBinopE Ne e1 e2

instance (Typed a, ToConst a, Eq a) => IsOrd (Exp a) (Exp Bool) where
    e1 .<.  e2 = BBinopE Lt e1 e2
    e1 .<=. e2 = BBinopE Le e1 e2
    e1 .>.  e2 = BBinopE Gt e1 e2
    e1 .>=. e2 = BBinopE Ge e1 e2

instance IfThenElse (Exp Bool) (Exp a) where
    ifThenElse = IfE

--------------------------------------------------------------------------------
--
-- Smart constructors
--
--------------------------------------------------------------------------------

class IsZeroOne a where
    isZero, isOne, isNegOne :: a -> Bool

isI :: Const (Complex a) -> Bool
isI RationalC{}    = False
isI (ComplexC r i) = isZero r && isOne i
isI (W _ _ c)      = isI c
isI (CycC c)       = c == i
isI PiC{}          = False

isNegI :: Const (Complex a) -> Bool
isNegI RationalC{}    = False
isNegI (ComplexC r i) = isZero r && isNegOne i
isNegI (W _ _ c)      = isNegI c
isNegI (CycC c)       = c == -i
isNegI PiC{}          = False

instance IsZeroOne (Const a) where
    isZero (IntC 0)       = True
    isZero (IntegerC 0)   = True
    isZero (FloatC 0)     = True
    isZero (DoubleC 0)    = True
    isZero (RationalC 0)  = True
    isZero (ComplexC 0 0) = True
    isZero (CycC 0)       = True
    isZero _              = False

    isOne (IntC 1)       = True
    isOne (IntegerC 1)   = True
    isOne (FloatC 1)     = True
    isOne (DoubleC 1)    = True
    isOne (RationalC 1)  = True
    isOne (ComplexC 1 i) = isZero i
    isOne (W n k _)      = k % n == 0 || k % n == 1
    isOne (CycC 1)       = True
    isOne _              = False

    isNegOne (IntC (-1))       = True
    isNegOne (IntegerC (-1))   = True
    isNegOne (FloatC (-1))    = True
    isNegOne (DoubleC (-1))    = True
    isNegOne (RationalC (-1))  = True
    isNegOne (ComplexC (-1) i) = isZero i
    isNegOne (W _ _ c)         = isNegOne c
    isNegOne (CycC (-1))       = True
    isNegOne _                 = False

instance IsZeroOne (Exp a) where
    isZero (ConstE c) = isZero c
    isZero _          = False

    isOne (ConstE c) = isOne c
    isOne _          = False

    isNegOne (ConstE c)             = isNegOne c
    isNegOne (UnopE Neg (ConstE c)) = isOne c
    isNegOne _                      = False

mkW :: RootOfUnity (Const a) => Ratio Int -> Const a -> Const a
mkW r c = W n (k `mod` n) c
  where
    k = numerator r
    n = denominator r

isConstE :: Exp a -> Bool
isConstE ConstE{} = True
isConstE _        = False

-- | Create an 'Exp Int' from an 'Int'.
intE :: Int -> Exp Int
intE = ConstE . IntC

-- | Create an 'Exp (Complex Double)' from a 'Complex Double'.
complexE :: RealFloatConst a => Complex a -> Exp (Complex a)
complexE (r :+ i) = ConstE $ ComplexC (toConst r) (toConst i)

-- | Extract the complex and real portions of a 'Const (Complex a)'.
unComplexC :: RealFloatConst a => Const (Complex a) -> (Const a, Const a)
unComplexC (RationalC r)  = (RationalC r, 0)
unComplexC (ComplexC r i) = (r, i)
unComplexC (W _ _ c)      = unComplexC c
unComplexC (PiC k)        = (PiC k, 0)
unComplexC x@CycC{}       = (toConst r, toConst i)
  where
    r :+ i = fromConst x

-- | Extract the complex and real portions of an 'Exp (Complex a)'.
unComplexE :: RealFloatConst a => Exp (Complex a) -> (Exp a, Exp a)
unComplexE (ConstE c) =
    (ConstE cr, ConstE ci)
  where
    (cr, ci) = unComplexC c

unComplexE (ComplexE er ei) =
    (er, ei)

unComplexE e =
    (ReE e, ImE e)

-- | Force extraction of the complex and real portions of an 'Exp (Complex a)'.
forceUnComplexE :: RealFloatConst a => Exp (Complex a) -> (Exp a, Exp a)
forceUnComplexE (ConstE c) =
    (ConstE cr, ConstE ci)
  where
    (cr, ci) = unComplexC c

forceUnComplexE (ComplexE er ei) =
    (er, ei)

forceUnComplexE (UnopE Neg e) =
    forceUnComplexE $ - (ensureComplexE e)

forceUnComplexE (BinopE Add e1 e2) =
    forceUnComplexE $ ensureComplexE e1 + ensureComplexE e2

forceUnComplexE (BinopE Sub e1 e2) =
    forceUnComplexE $ ensureComplexE e1 - ensureComplexE e2

forceUnComplexE (BinopE Mul e1 e2) =
    forceUnComplexE $ ensureComplexE e1 * ensureComplexE e2

forceUnComplexE e =
    (ReE e, ImE e)

-- | Ensure that a complex expression is represented using its constituent real
-- and imaginary parts.
ensureComplexE :: RealFloatConst a => Exp (Complex a) -> Exp (Complex a)
ensureComplexE e = ComplexE er ei
  where
    (er, ei) = forceUnComplexE e

true, false :: Exp Bool
true  = ConstE (BoolC True)
false = ConstE (BoolC False)

max_stars_repo_name: mainland/hspiral

max_issues_repo_name: mainland/hspiral

max_forks_repo_name: mainland/hspiral


text:
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

module Test.Grenade.Layers.Trivial where

import           Grenade.Layers.Trivial
import           Grenade.Core

import           Hedgehog

import           Test.Hedgehog.Hmatrix

import           Numeric.LinearAlgebra.Static as H hiding ((===), create)
import           Data.Serialize
import           Data.Either
import           System.Random.MWC             (create)

prop_trivial_run_forward_passes_along_input :: Property
prop_trivial_run_forward_passes_along_input = property $ do
  input :: S ('D3 5 7 3) <- forAll genOfShape

  let layer = Trivial
      ((), out) = runForwards layer input
      ((), y)   = runBackwards layer () out
      S3D input' = input
      S3D out'   = out
      S3D y'     = y
  H.extract out' === H.extract input'
  H.extract out' === H.extract y'

prop_trivial_is_serializable :: Property
prop_trivial_is_serializable = withTests 1 $ property $ do
  let layer = Trivial
      bs = encode layer
      dec = decode bs :: Either String Trivial
  assert $ isRight dec

prop_can_randomise_trivial :: Property
prop_can_randomise_trivial = withTests 1 $ property $ do
  gen <- evalIO create
  _ :: Trivial <- evalIO $ createRandomWith UniformInit gen
  success

prop_can_show_trivial :: Property
prop_can_show_trivial = withTests 1 $ property $ do
  show Trivial `seq` success

prop_can_update_trivial_and_use_in_batches :: Property
prop_can_update_trivial_and_use_in_batches = property $ do
  let layer = Trivial
  runUpdate defSGD layer () `seq` success
  runUpdate defAdam layer () `seq` success
  reduceGradient @Trivial [()] `seq` success

tests :: IO Bool
tests = checkParallel $$(discover)

ext: hs

max_stars_repo_name: th-char/grenade

max_issues_repo_name: th-char/grenade

max_forks_repo_name: th-char/grenade


text:
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Normal
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The normal distribution.  This is a continuous probability
-- distribution that describes data that cluster around a mean.

module Statistics.Distribution.Normal
    (
      NormalDistribution
    -- * Constructors
    , normalDistr
    , normalDistrE
    , normalDistrErr
    , standard
    ) 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.MathFunctions.Constants (m_sqrt_2, m_sqrt_2_pi)
import Numeric.SpecFunctions (erfc, invErfc)
import qualified System.Random.MWC.Distributions as MWC
import qualified Data.Vector.Generic as G

import qualified Statistics.Distribution as D
import qualified Statistics.Sample as S
import Statistics.Internal

-- | The normal distribution.
data NormalDistribution = ND {
      mean       :: {-# UNPACK #-} !Double
    , stdDev     :: {-# UNPACK #-} !Double
    , ndPdfDenom :: {-# UNPACK #-} !Double
    , ndCdfDenom :: {-# UNPACK #-} !Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show NormalDistribution where
  showsPrec i (ND m s _ _) = defaultShow2 "normalDistr" m s i
instance Read NormalDistribution where
  readPrec = defaultReadPrecM2 "normalDistr" normalDistrE

instance ToJSON NormalDistribution
instance FromJSON NormalDistribution where
  parseJSON (Object v) = do
    m  <- v .: "mean"
    sd <- v .: "stdDev"
    either fail return $ normalDistrErr m sd
  parseJSON _ = empty

instance Binary NormalDistribution where
    put (ND m sd _ _) = put m >> put sd
    get = do
      m  <- get
      sd <- get
      either fail return $ normalDistrErr m sd

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

instance D.ContDistr NormalDistribution where
    logDensity    = logDensity
    quantile      = quantile
    complQuantile = complQuantile

instance D.MaybeMean NormalDistribution where
    maybeMean = Just . D.mean

instance D.Mean NormalDistribution where
    mean = mean

instance D.MaybeVariance NormalDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Variance NormalDistribution where
    stdDev = stdDev

instance D.Entropy NormalDistribution where
  entropy d = 0.5 * log (2 * pi * exp 1 * D.variance d)

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

instance D.ContGen NormalDistribution where
    genContVar d = MWC.normal (mean d) (stdDev d)

-- | Standard normal distribution with mean equal to 0 and variance equal to 1
standard :: NormalDistribution
standard = ND { mean       = 0.0
              , stdDev     = 1.0
              , ndPdfDenom = log m_sqrt_2_pi
              , ndCdfDenom = m_sqrt_2
              }

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistr :: Double            -- ^ Mean of distribution
            -> Double            -- ^ Standard deviation of distribution
            -> NormalDistribution
normalDistr m sd = either error id $ normalDistrErr m sd

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistrE :: Double            -- ^ Mean of distribution
             -> Double            -- ^ Standard deviation of distribution
             -> Maybe NormalDistribution
normalDistrE m sd = either (const Nothing) Just $ normalDistrErr m sd

-- | Create normal distribution from parameters.
--
normalDistrErr :: Double            -- ^ Mean of distribution
               -> Double            -- ^ Standard deviation of distribution
               -> Either String NormalDistribution
normalDistrErr m sd
  | sd > 0    = Right $ ND { mean       = m
                           , stdDev     = sd
                           , ndPdfDenom = log $ m_sqrt_2_pi * sd
                           , ndCdfDenom = m_sqrt_2 * sd
                           }
  | otherwise = Left $ errMsg m sd

errMsg :: Double -> Double -> String
errMsg _ sd = "Statistics.Distribution.Normal.normalDistr: standard deviation must be positive. Got " ++ show sd

-- | Variance is estimated using maximum likelihood method
--   (biased estimation).
--
--   Returns @Nothing@ if sample contains less than one element or
--   variance is zero (all elements are equal)
instance D.FromSample NormalDistribution Double where
  fromSample xs
    | G.length xs <= 1 = Nothing
    | v == 0           = Nothing
    | otherwise        = Just $! normalDistr m (sqrt v)
    where
      (m,v) = S.meanVariance xs

logDensity :: NormalDistribution -> Double -> Double
logDensity d x = (-xm * xm / (2 * sd * sd)) - ndPdfDenom d
    where xm = x - mean d
          sd = stdDev d

cumulative :: NormalDistribution -> Double -> Double
cumulative d x = erfc ((mean d - x) / ndCdfDenom d) / 2

complCumulative :: NormalDistribution -> Double -> Double
complCumulative d x = erfc ((x - mean d) / ndCdfDenom d) / 2

quantile :: NormalDistribution -> Double -> Double
quantile d p
  | p == 0         = -inf
  | p == 1         = inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.quantile: p must be in [0,1] range. Got: "++show p
  where x          = - invErfc (2 * p)
        inf        = 1/0

complQuantile :: NormalDistribution -> Double -> Double
complQuantile d p
  | p == 0         = inf
  | p == 1         = -inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.complQuantile: p must be in [0,1] range. Got: "++show p
  where x          = invErfc (2 * p)
        inf        = 1/0


text:
{-|
Module      : Numeric.MCMC.Diagnostics
Description : Diagnostics for Markov Chain Monte Carlo Sampling
Copyright   : (c) Adam Conner-Sax 2019
License     : BSD
Maintainer  : adam_conner_sax@yahoo.com
Stability   : experimental

Numeric.MCMC.Diagnostics contains functions to analyze the convergence of MCMC sampling and to estimate the distribution of expectations computed
using a set of MCMC chains.
-}
module Numeric.MCMC.Diagnostics
  ( summarize
  , mpsrf
  , ExpectationSummary(..)
  , mannWhitneyUTest
  )
where

import qualified Control.Foldl                 as FL
import qualified Statistics.Types              as S
import qualified Data.List                     as L
--import qualified Data.Vector.Generic           as V
import qualified Data.Vector.Unboxed           as VU
import qualified Statistics.Test.MannWhitneyU  as S
import qualified Numeric.LinearAlgebra         as LA

data ExpectationSummary a = ExpectationSummary { mean :: a, confidence :: (a, a), rHat :: a } deriving (Show)
-- all sub-chains should have same length here
-- do we need to throw away first half of all the chains??
{- |
Compute the exepctation, confidence interval and
PSRF statistic (from <http://www2.stat.duke.edu/~scs/Courses/Stat376/Papers/ConvergeDiagnostics/BrooksGelman.pdf>)
for a given confidence %, expectation function, and set of MCMC chains. In the version in the paper, each chain is
halved before the statistic is computed but it's not clear why.  Here we assume that is "burn-in" and the user may
choose to do that or not before calling this function.

All sub-chains should have the same length here, though that is not checked.
-}
summarize
  :: (RealFrac b, Ord b)
  => S.CL b -- ^ confidence interval to test
  -> (a -> b) -- ^ function to map a sample to a function under consideration
  -> [[a]] -- ^ list of chains
  -> Maybe (ExpectationSummary b) -- ^ @Maybe@ here in case chains aren't long enough
summarize confidence' expectationF chains = do
  let applied    = fmap (fmap expectationF) chains
      allApplied = concat applied
      chainInterval :: Ord b => [b] -> Maybe (b, b)
      chainInterval c =
        let n  = L.length c
            n' = round (realToFrac n * S.confidenceLevel confidence')
            d  = (n - n') `div` 2
            c' = L.take n' $ L.drop d $ L.sort c
        in if length c' >= 2 then (head &&& last) <$> nonEmpty c' else Nothing
--      intervalLength (a, b) = b - a
  intervals <- traverse chainInterval applied
  (lo, hi)  <- chainInterval allApplied
  let (mLo, mHi) = FL.fold
        ((,) <$> FL.premap fst FL.mean <*> FL.premap snd FL.mean)
        intervals
  let rHat' = (hi - lo) / (mHi - mLo)
      mean' = FL.fold FL.mean allApplied
  return $ ExpectationSummary mean' (lo, hi) rHat'

-- if the chain list is empty or chains have different sizes, etc. we return Nothing
mpsrf :: [a -> Double] -> [[a]] -> Maybe Double
mpsrf expectations chains = do
  let --mkVec :: a -> LA.Vector Double
      mkVec x = LA.fromList $ fmap ($ x) expectations
      m = length chains
  n <- length . head <$> nonEmpty chains --if m > 0 then Just (length $ head chains) else Nothing
  traverse_ (\c -> if length c == n then Just () else Nothing) chains -- check that all chains are length n
  let
    sumVecsF = FL.Fold (+) (LA.vector $ L.replicate (length expectations) 0) id
    sumMatsF = FL.Fold (+) (LA.scale 0 $ LA.ident (length expectations)) id
    meanVecsF =
      (\m' l -> LA.scale (1 / realToFrac l) m') <$> sumVecsF <*> FL.length
    avgVec     = FL.fold meanVecsF
    phiChains  = fmap (fmap mkVec) chains
    avgPhiEach = fmap avgVec phiChains
    diffProd :: LA.Vector Double -> LA.Vector Double -> LA.Matrix Double
    diffProd x y = (x - y) `LA.outer` (x - y)
    w1 = (\(ap, ps) -> (`diffProd` ap) <$> ps) <$> zip avgPhiEach phiChains
    w =
      LA.scale (1 / (realToFrac m * realToFrac (n - 1)))
        $ FL.fold sumMatsF
        $ fmap (FL.fold sumMatsF) w1
    avgPhiAll = avgVec avgPhiEach
    b1        = fmap (`diffProd` avgPhiAll) avgPhiEach
    b = LA.scale (realToFrac n / realToFrac (m - 1)) $ FL.fold sumMatsF b1
    mat       = LA.scale (1 / realToFrac n) $ LA.inv w LA.<> b
    eigen1    = LA.eigenvaluesSH (LA.trustSym mat) `LA.atIndex` 0
    rhat =
      (realToFrac (n - 1) / realToFrac n)
        + (realToFrac (m + 1) / realToFrac m)
        * eigen1
  return rhat

mannWhitneyUTest
  :: (Ord b, VU.Unbox b)
  => S.PValue Double
  -> (a -> b)
  -> [a]
  -> [a]
  -> Maybe S.TestResult
mannWhitneyUTest p f c1 c2 =
  let vOfF = VU.fromList . fmap f
  in  S.mannWhitneyUtest S.SamplesDiffer p (vOfF c1) (vOfF c2)

max_stars_repo_name: adamConnerSax/data-science-utils

max_issues_repo_name: adamConnerSax/data-science-utils

max_forks_repo_name: adamConnerSax/data-science-utils


text:
module Utils where

import Data.Matrix
import qualified Data.Complex as C

type ComplexMatrix a = Matrix (C.Complex a)

-- Complex conjugate of a matrix
conjugate :: (Num a) => ComplexMatrix a -> ComplexMatrix a
conjugate = fmap C.conjugate

-- Similar to numpy.vdot. Takes complex conjugate of the first matrix
-- and computes matrix product.
vdot :: (RealFloat a) => ComplexMatrix a -> ComplexMatrix a -> ComplexMatrix a
vdot m n = (conjugate m) * n

max_stars_repo_name: viswanathgs/haskell-fft

max_issues_repo_name: viswanathgs/haskell-fft

max_forks_repo_name: viswanathgs/haskell-fft


text:
module LinearRegression where

import Statistics.LinearRegression
import qualified Data.Vector.Unboxed as U
import Data.Decimal
import Graphics.EasyPlot

xs :: [Double]
xs = [25.47,67.80,20.65,22.50,30.40,55.30,22.74] ++
         [29.11,31.15,40.97,58.89,64.85,74.89,96.14] ++
         [65.37,80.23,95.80,92.71,98.16,89.77]

ys :: [Double]
ys = [390.50,632.70,449.36,501.81,621.72,684.62,685.53] ++
         [701.76,868.71,995.69,893.50,1240.48,1393.55] ++
         [1826.70,2145.32,2295.27,2522.79,2931.69] ++
         [2836.20,2456.30]

-- Values from a list
--
varIndep  = U.fromList xs
varDepend = U.fromList ys

-- computing theta0 and theta1 : Y = t0 + b1.X

(b0, b1) = linearRegression varIndep varDepend

b0n = realFracToDecimal 3 b0
b1n = realFracToDecimal 3 b1

points = zip xs ys

f x = b0 + b1*x
g x = ((*b0).sin.(/b1)) x

functionY = do
  plot X11 $ [
                    Function2D
                    [Title "Linear Regression", Color Red]
                    [Range 0 140] (g)
                    , Data2D
                    [Title "Data", Color Blue]
                    [] points
                    , Function2D
                    [Title "Linear Regression", Color Red]
                    [Range 0 140] (f)
                 ]

max_stars_repo_name: Levisoto/practiceAndLearning

max_issues_repo_name: Levisoto/practiceAndLearning

max_forks_repo_name: Levisoto/practiceAndLearning


text:
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
#if __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}
# if MIN_VERSION_array(0,4,0)
{-# LANGUAGE Safe #-}
# endif
#endif
-----------------------------------------------------------------------------
-- |
-- Module      :  Control.DeepSeq
-- Copyright   :  (c) The University of Glasgow 2001-2009
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  libraries@haskell.org
-- Stability   :  stable
-- Portability :  portable
--
-- This module provides an overloaded function, 'deepseq', for fully
-- evaluating data structures (that is, evaluating to \"Normal Form\").
--
-- A typical use is to prevent resource leaks in lazy IO programs, by
-- forcing all characters from a file to be read. For example:
--
-- > import System.IO
-- > import Control.DeepSeq
-- >
-- > main = do
-- >     h <- openFile "f" ReadMode
-- >     s <- hGetContents h
-- >     s `deepseq` hClose h
-- >     return s
--
-- 'deepseq' differs from 'seq' as it traverses data structures deeply,
-- for example, 'seq' will evaluate only to the first constructor in
-- the list:
--
-- > > [1,2,undefined] `seq` 3
-- > 3
--
-- While 'deepseq' will force evaluation of all the list elements:
--
-- > > [1,2,undefined] `deepseq` 3
-- > *** Exception: Prelude.undefined
--
-- Another common use is to ensure any exceptions hidden within lazy
-- fields of a data structure do not leak outside the scope of the
-- exception handler, or to force evaluation of a data structure in one
-- thread, before passing to another thread (preventing work moving to
-- the wrong threads).
--
-- /Since: 1.1.0.0/
module Control.DeepSeq (
     deepseq, ($!!), force,
     NFData(..),
  ) where

import Control.Applicative
import Control.Concurrent ( ThreadId )
import Data.Int
import Data.Word
import Data.Ratio
import Data.Complex
import Data.Array
import Data.Fixed
import Data.Version
import Data.Monoid
import Data.Unique ( Unique )
import Foreign.C.Types
import System.Mem.StableName ( StableName )

#if MIN_VERSION_base(4,6,0)
import Data.Ord ( Down(Down) )
#endif

#if MIN_VERSION_base(4,7,0)
import Data.Proxy ( Proxy(Proxy) )
#endif

#if MIN_VERSION_base(4,8,0)
import Data.Functor.Identity ( Identity(..) )
import Data.Typeable ( TypeRep, TyCon, rnfTypeRep, rnfTyCon )
import Data.Void ( Void, absurd )
import Numeric.Natural ( Natural )
#endif

#if __GLASGOW_HASKELL__ >= 702
import GHC.Fingerprint.Type ( Fingerprint(..) )
import GHC.Generics

-- | Hidden internal type-class
class GNFData f where
  grnf :: f a -> ()

instance GNFData V1 where
  grnf = error "Control.DeepSeq.rnf: uninhabited type"

instance GNFData U1 where
  grnf U1 = ()

instance NFData a => GNFData (K1 i a) where
  grnf = rnf . unK1
  {-# INLINEABLE grnf #-}

instance GNFData a => GNFData (M1 i c a) where
  grnf = grnf . unM1
  {-# INLINEABLE grnf #-}

instance (GNFData a, GNFData b) => GNFData (a :*: b) where
  grnf (x :*: y) = grnf x `seq` grnf y
  {-# INLINEABLE grnf #-}

instance (GNFData a, GNFData b) => GNFData (a :+: b) where
  grnf (L1 x) = grnf x
  grnf (R1 x) = grnf x
  {-# INLINEABLE grnf #-}
#endif

infixr 0 $!!

-- | 'deepseq': fully evaluates the first argument, before returning the
-- second.
--
-- The name 'deepseq' is used to illustrate the relationship to 'seq':
-- where 'seq' is shallow in the sense that it only evaluates the top
-- level of its argument, 'deepseq' traverses the entire data structure
-- evaluating it completely.
--
-- 'deepseq' can be useful for forcing pending exceptions,
-- eradicating space leaks, or forcing lazy I/O to happen.  It is
-- also useful in conjunction with parallel Strategies (see the
-- @parallel@ package).
--
-- There is no guarantee about the ordering of evaluation.  The
-- implementation may evaluate the components of the structure in
-- any order or in parallel.  To impose an actual order on
-- evaluation, use 'pseq' from "Control.Parallel" in the
-- @parallel@ package.
--
-- /Since: 1.1.0.0/
deepseq :: NFData a => a -> b -> b
deepseq a b = rnf a `seq` b

-- | the deep analogue of '$!'.  In the expression @f $!! x@, @x@ is
-- fully evaluated before the function @f@ is applied to it.
--
-- /Since: 1.2.0.0/
($!!) :: (NFData a) => (a -> b) -> a -> b
f $!! x = x `deepseq` f x

-- | a variant of 'deepseq' that is useful in some circumstances:
--
-- > force x = x `deepseq` x
--
-- @force x@ fully evaluates @x@, and then returns it.  Note that
-- @force x@ only performs evaluation when the value of @force x@
-- itself is demanded, so essentially it turns shallow evaluation into
-- deep evaluation.
--
-- 'force' can be conveniently used in combination with @ViewPatterns@:
--
-- > {-# LANGUAGE BangPatterns, ViewPatterns #-}
-- > import Control.DeepSeq
-- >
-- > someFun :: ComplexData -> SomeResult
-- > someFun (force -> !arg) = {- 'arg' will be fully evaluated -}
--
-- Another useful application is to combine 'force' with
-- 'Control.Exception.evaluate' in order to force deep evaluation
-- relative to other 'IO' operations:
--
-- > import Control.Exception (evaluate)
-- > import Control.DeepSeq
-- >
-- > main = do
-- >   result <- evaluate $ force $ pureComputation
-- >   {- 'result' will be fully evaluated at this point -}
-- >   return ()
--
-- /Since: 1.2.0.0/
force :: (NFData a) => a -> a
force x = x `deepseq` x

-- | A class of types that can be fully evaluated.
--
-- /Since: 1.1.0.0/
class NFData a where
    -- | 'rnf' should reduce its argument to normal form (that is, fully
    -- evaluate all sub-components), and then return '()'.
    --
    -- === 'Generic' 'NFData' deriving
    --
    -- Starting with GHC 7.2, you can automatically derive instances
    -- for types possessing a 'Generic' instance.
    --
    -- > {-# LANGUAGE DeriveGeneric #-}
    -- >
    -- > import GHC.Generics (Generic)
    -- > import Control.DeepSeq
    -- >
    -- > data Foo a = Foo a String
    -- >              deriving (Eq, Generic)
    -- >
    -- > instance NFData a => NFData (Foo a)
    -- >
    -- > data Colour = Red | Green | Blue
    -- >               deriving Generic
    -- >
    -- > instance NFData Colour
    --
    -- Starting with GHC 7.10, the example above can be written more
    -- concisely by enabling the new @DeriveAnyClass@ extension:
    --
    -- > {-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
    -- >
    -- > import GHC.Generics (Generic)
    -- > import Control.DeepSeq
    -- >
    -- > data Foo a = Foo a String
    -- >              deriving (Eq, Generic, NFData)
    -- >
    -- > data Colour = Red | Green | Blue
    -- >               deriving (Generic, NFData)
    -- >
    --
    -- === Compatibility with previous @deepseq@ versions
    --
    -- Prior to version 1.4.0.0, the default implementation of the 'rnf'
    -- method was defined as
    --
    -- @'rnf' a = 'seq' a ()@
    --
    -- However, starting with @deepseq-1.4.0.0@, the default
    -- implementation is based on @DefaultSignatures@ allowing for
    -- more accurate auto-derived 'NFData' instances. If you need the
    -- previously used exact default 'rnf' method implementation
    -- semantics, use
    --
    -- > instance NFData Colour where rnf x = seq x ()
    --
    -- or alternatively
    --
    -- > {-# LANGUAGE BangPatterns #-}
    -- > instance NFData Colour where rnf !_ = ()
    --
    rnf :: a -> ()

#if __GLASGOW_HASKELL__ >= 702
    default rnf :: (Generic a, GNFData (Rep a)) => a -> ()
    rnf = grnf . from
#endif

instance NFData Int      where rnf !_ = ()
instance NFData Word     where rnf !_ = ()
instance NFData Integer  where rnf !_ = ()
instance NFData Float    where rnf !_ = ()
instance NFData Double   where rnf !_ = ()

instance NFData Char     where rnf !_ = ()
instance NFData Bool     where rnf !_ = ()
instance NFData ()       where rnf !_ = ()

instance NFData Int8     where rnf !_ = ()
instance NFData Int16    where rnf !_ = ()
instance NFData Int32    where rnf !_ = ()
instance NFData Int64    where rnf !_ = ()

instance NFData Word8    where rnf !_ = ()
instance NFData Word16   where rnf !_ = ()
instance NFData Word32   where rnf !_ = ()
instance NFData Word64   where rnf !_ = ()

#if MIN_VERSION_base(4,7,0)
-- |/Since: 1.4.0.0/
instance NFData (Proxy a) where rnf Proxy = ()
#endif

#if MIN_VERSION_base(4,8,0)
-- |/Since: 1.4.0.0/
instance NFData a => NFData (Identity a) where
    rnf = rnf . runIdentity

-- | Defined as @'rnf' = 'absurd'@.
--
-- /Since: 1.4.0.0/
instance NFData Void where
    rnf = absurd

-- |/Since: 1.4.0.0/
instance NFData Natural  where rnf !_ = ()
#endif

-- |/Since: 1.3.0.0/
instance NFData (Fixed a) where rnf !_ = ()

-- |This instance is for convenience and consistency with 'seq'.
-- This assumes that WHNF is equivalent to NF for functions.
--
-- /Since: 1.3.0.0/
instance NFData (a -> b) where rnf !_ = ()

--Rational and complex numbers.

#if __GLASGOW_HASKELL__ >= 711
instance NFData a => NFData (Ratio a) where
#else
instance (Integral a, NFData a) => NFData (Ratio a) where
#endif
  rnf x = rnf (numerator x, denominator x)

#if MIN_VERSION_base(4,4,0)
instance (NFData a) => NFData (Complex a) where
#else
instance (RealFloat a, NFData a) => NFData (Complex a) where
#endif
  rnf (x:+y) = rnf x `seq`
               rnf y `seq`
               ()

instance NFData a => NFData (Maybe a) where
    rnf Nothing  = ()
    rnf (Just x) = rnf x

instance (NFData a, NFData b) => NFData (Either a b) where
    rnf (Left x)  = rnf x
    rnf (Right y) = rnf y

-- |/Since: 1.3.0.0/
instance NFData Data.Version.Version where
    rnf (Data.Version.Version branch tags) = rnf branch `seq` rnf tags

instance NFData a => NFData [a] where
    rnf [] = ()
    rnf (x:xs) = rnf x `seq` rnf xs

-- |/Since: 1.4.0.0/
instance NFData a => NFData (ZipList a) where
    rnf = rnf . getZipList

-- |/Since: 1.4.0.0/
instance NFData a => NFData (Const a b) where
    rnf = rnf . getConst

#if __GLASGOW_HASKELL__ >= 711
instance (NFData a, NFData b) => NFData (Array a b) where
#else
instance (Ix a, NFData a, NFData b) => NFData (Array a b) where
#endif
    rnf x = rnf (bounds x, Data.Array.elems x)

#if MIN_VERSION_base(4,6,0)
-- |/Since: 1.4.0.0/
instance NFData a => NFData (Down a) where
    rnf (Down x) = rnf x
#endif

-- |/Since: 1.4.0.0/
instance NFData a => NFData (Dual a) where
    rnf = rnf . getDual

-- |/Since: 1.4.0.0/
instance NFData a => NFData (First a) where
    rnf = rnf . getFirst

-- |/Since: 1.4.0.0/
instance NFData a => NFData (Last a) where
    rnf = rnf . getLast

-- |/Since: 1.4.0.0/
instance NFData Any where rnf = rnf . getAny

-- |/Since: 1.4.0.0/
instance NFData All where rnf = rnf . getAll

-- |/Since: 1.4.0.0/
instance NFData a => NFData (Sum a) where
    rnf = rnf . getSum

-- |/Since: 1.4.0.0/
instance NFData a => NFData (Product a) where
    rnf = rnf . getProduct

-- |/Since: 1.4.0.0/
instance NFData (StableName a) where
    rnf !_ = () -- assumes `data StableName a = StableName (StableName# a)`

-- |/Since: 1.4.0.0/
instance NFData ThreadId where
    rnf !_ = () -- assumes `data ThreadId = ThreadId ThreadId#`

-- |/Since: 1.4.0.0/
instance NFData Unique where
    rnf !_ = () -- assumes `newtype Unique = Unique Integer`

#if MIN_VERSION_base(4,8,0)
-- | __NOTE__: Only defined for @base-4.8.0.0@ and later
--
-- /Since: 1.4.0.0/
instance NFData TypeRep where
    rnf tyrep = rnfTypeRep tyrep

-- | __NOTE__: Only defined for @base-4.8.0.0@ and later
--
-- /Since: 1.4.0.0/
instance NFData TyCon where
    rnf tycon = rnfTyCon tycon
#endif

----------------------------------------------------------------------------
-- GHC Specifics

#if __GLASGOW_HASKELL__ >= 702
-- |/Since: 1.4.0.0/
instance NFData Fingerprint where
    rnf (Fingerprint _ _) = ()
#endif

----------------------------------------------------------------------------
-- Foreign.C.Types

-- |/Since: 1.4.0.0/
instance NFData CChar where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CSChar where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CUChar where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CShort where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CUShort where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CInt where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CUInt where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CLong where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CULong where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CPtrdiff where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CSize where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CWchar where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CSigAtomic where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CLLong where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CULLong where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CIntPtr where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CUIntPtr where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CIntMax where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CUIntMax where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CClock where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CTime where rnf !_ = ()

#if MIN_VERSION_base(4,4,0)
-- |/Since: 1.4.0.0/
instance NFData CUSeconds where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CSUSeconds where rnf !_ = ()
#endif

-- |/Since: 1.4.0.0/
instance NFData CFloat where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CDouble where rnf !_ = ()

-- NOTE: The types `CFile`, `CFPos`, and `CJmpBuf` below are not
-- newtype wrappers rather defined as field-less single-constructor
-- types.

-- |/Since: 1.4.0.0/
instance NFData CFile where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CFpos where rnf !_ = ()

-- |/Since: 1.4.0.0/
instance NFData CJmpBuf where rnf !_ = ()

----------------------------------------------------------------------------
-- Tuples

instance (NFData a, NFData b) => NFData (a,b) where
  rnf (x,y) = rnf x `seq` rnf y

instance (NFData a, NFData b, NFData c) => NFData (a,b,c) where
  rnf (x,y,z) = rnf x `seq` rnf y `seq` rnf z

instance (NFData a, NFData b, NFData c, NFData d) => NFData (a,b,c,d) where
  rnf (x1,x2,x3,x4) = rnf x1 `seq`
                      rnf x2 `seq`
                      rnf x3 `seq`
                      rnf x4

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5) =>
         NFData (a1, a2, a3, a4, a5) where
  rnf (x1, x2, x3, x4, x5) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5, NFData a6) =>
         NFData (a1, a2, a3, a4, a5, a6) where
  rnf (x1, x2, x3, x4, x5, x6) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5 `seq`
                  rnf x6

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5, NFData a6, NFData a7) =>
         NFData (a1, a2, a3, a4, a5, a6, a7) where
  rnf (x1, x2, x3, x4, x5, x6, x7) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5 `seq`
                  rnf x6 `seq`
                  rnf x7

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5, NFData a6, NFData a7, NFData a8) =>
         NFData (a1, a2, a3, a4, a5, a6, a7, a8) where
  rnf (x1, x2, x3, x4, x5, x6, x7, x8) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5 `seq`
                  rnf x6 `seq`
                  rnf x7 `seq`
                  rnf x8

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5, NFData a6, NFData a7, NFData a8, NFData a9) =>
         NFData (a1, a2, a3, a4, a5, a6, a7, a8, a9) where
  rnf (x1, x2, x3, x4, x5, x6, x7, x8, x9) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5 `seq`
                  rnf x6 `seq`
                  rnf x7 `seq`
                  rnf x8 `seq`
                  rnf x9

max_stars_repo_name: DavidAlphaFox/ghc

max_issues_repo_name: DavidAlphaFox/ghc

max_forks_repo_name: DavidAlphaFox/ghc


text:
{-# LANGUAGE FlexibleContexts #-}
module QModels where

import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit
import PCA.QModels
import PCA.Distribution
import Numeric.LinearAlgebra.Data
       (fromRows, toRows, size, (!), konst, Matrix, Vector, fromList)
import qualified Numeric.LinearAlgebra.HMatrix as H

tests :: TestTree
tests = testGroup "QModels module tests" [unitTests]

unitTests :: TestTree
unitTests =
    testCaseSteps "Tests for Q model calculation" $
    \step ->
         do step "Preparing..."
            let n = 10
                d = 3
            trainT <- H.rand n d
            let tau = initQtau
                mu = initQmu d
                alpha = initQalpha d
                w = initQW d
            step "Test calculation of Q"
            let sigmaX = calcSigmaX d tau w
            assertEqual "Size of Sigma_X must dxd" (size sigmaX) (d, d)
            let mX = calcMX trainT tau sigmaX w mu
            assertEqual
                "Fool impl and matrix impl of calcMX must give the same result"
                mX
                (foolCalcMX trainT tau sigmaX w mu)
            assertEqual "Size of mean X is Nxd" (size mX) (n, d)
            let x = MatrixMNormal mX sigmaX
            let sigmaMu = calcSigmaMu n tau mu
            assertEqual "Size of Sigma_Mu is dxd" (size sigmaMu) (d, d)
            let mMu = calcMMu trainT tau sigmaMu w x
            assertEqual "Size of m_Mu is 1xd" (size mMu) d
            assertEqual
                "Fool impl and matrix impl of calcMMu must give the same result"
                mMu
                (foolCalcMMu trainT tau sigmaMu w x)
            let sigmaW = calcSigmaW tau alpha x
            assertEqual "Size of sigma_W is dxd" (size sigmaW) (d, d)
            assertBool
                "Fool impl and matrix impl of calcSigmaW must give the same result"
                (H.maxElement (foolCalcSigmaW tau alpha x - sigmaW) < 0.001)
            let mW = calcMW tau sigmaW x trainT mu
            assertEqual "Size of m_W is dxd" (size mW) (d, d)
            assertBool
                "Fool impl and matrix impl of calcMW must give the same result"
                (H.maxElement (foolCalcMW tau sigmaW x trainT mu - mW) < 0.001)
            let bAlpha = calcBalpha (konst 1.0e-3 d) w
            assertEqual "Size of b_alpha is d" (size bAlpha) d
            assertEqual
                "Fool impl and matrix impl of calcBalpha must give the same result"
                (foolCalcBalpha (konst 1.0e-3 d) w)
                bAlpha
            let bTau = calcBtau 1.0e-3 trainT mu w x
            assertBool
                "Fool impl and matrix impl of calcBtau must give the same result"
                (abs (foolCalcBtau 1.0e-3 trainT mu w x) - bTau < 0.001)

foolCalcMX trainT tau sigmaX w mu =
    fromRows
        (map
             (\r ->
                   coef H.#> (r - mean mu))
             (toRows trainT))
  where
    coef = mean tau `H.scale` sigmaX H.<> H.tr (mean w)

foolCalcMMu trainT tau sigmaMu w x =
    (mean tau `H.scale` sigmaMu) H.#>
    foldr
        (\r resV ->
              resV + ((trainT ! r) - mW H.#> (mX ! r)))
        (konst 0 d)
        [0 .. (n - 1)]
  where
    mW = mean w
    mX = mean x
    (n,d) = size trainT

foolCalcSigmaW tau alpha x =
    H.inv
        (H.diag (mean alpha) +
         mean tau `H.scale`
         (foldr
              (\r resM ->
                    resM + variance x + r `H.outer` r)
              (H.diagl [0 | _ <- [1 .. d]])
              (toRows mX)))
  where
    mX = mean x
    (_,d) = size mX

foolCalcMW
    :: Distr Double Double
    -> Matrix Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Matrix Double
    -> Distr (Vector Double) (Matrix Double)
    -> Matrix Double
foolCalcMW tau sigmaW x trainT mu =
    fromRows
        (map
             (\k ->
                   (mean tau `H.scale` sigmaW) H.#>
                   (foldr
                        (\i resV ->
                              resV +
                              ((trainT ! i ! k) - (mMu ! k)) `H.scale` (mX ! i))
                        (konst 0 d)
                        [0 .. (n - 1)]))
                       [0 .. (d - 1)])
  where
    (n,d) = size trainT
    mMu = mean mu
    mX = mean x

foolCalcBalpha :: Vector Double
               -> Distr (Matrix Double) (Matrix Double)
               -> Vector Double
foolCalcBalpha b w =
    fromList
        (map
             (\i ->
                   b ! i + (mW ! i) `H.dot` (mW ! i) / 2)
             [0 .. (d - 1)])
  where
    mW = mean w
    (d,_) = size mW

foolCalcBtau
    :: Double
    -> Matrix Double
    -> Distr (Vector Double) (Matrix Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Double
foolCalcBtau b trainT mu w x =
    b +
    0.5 *
    sum
        (map
             (\i ->
                   normSq (trainT ! i) + normSq mMu +
                   2 * (mMu `H.dot` (mW H.#> (mX ! i))) -
                   2 * (trainT ! i `H.dot` (mW H.#> (mX ! i))) -
                   2 * (trainT ! i `H.dot` mMu))
             [0 .. (n - 1)])
  where
    normSq v = v `H.dot` v
    mMu = mean mu
    mX = mean x
    mW = mean w
    (n,_) = size trainT

max_stars_repo_name: DbIHbKA/vbpca

max_issues_repo_name: DbIHbKA/vbpca

max_forks_repo_name: DbIHbKA/vbpca


text:
{-# LANGUAGE ParallelListComp #-}
module Main where

import Numeric.LinearAlgebra.Data (fromList, toList, cmap)
import Numeric.GSL.Fitting.Linear (linear_w)
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo
import Text.Printf

x = fromList [10, 100, 1000, 10000, 100000, 1000000]
y = fromList [4772, 21148, 134194, 1343814, 13389821, 148509325]
sigma = fromList [2112, 8194, 61719, 804762, 4853113, 42071732]

(intercept, slope, _, _, _, chi_sq) = linear_w x (cmap (\s -> 1 / s^2) sigma) y

errbars = liftEC $ do
    let values = [symErrPoint x' y' 0 dy
                  | x' <- toList x
                  | y' <- toList y
                  | dy <- toList sigma
                 ]
    plot_errbars_values .= values

main = toFile def "benchmark-fit.png" $ do
    setColors (map opaque [blue,red,green])
    layout_title .=
        printf "t = %.2f (ns/element) n + %.2f (ns) [chi^2 = %.2f]" slope intercept chi_sq
    layout_x_axis . laxis_generate .= autoScaledLogAxis def
    layout_y_axis . laxis_generate .= autoScaledLogAxis def
    plot errbars
    plot $ points "Benchmark data"
        (zip (toList x) (toList y))
    plot $ line "Best fit" [[
            (t, slope*t + intercept)
            | t <- map (10**) [0.5, 0.6 .. 6.5 ]
        ]]

max_stars_repo_name: devonhollowood/radix-sort

max_issues_repo_name: devonhollowood/radix-sort

max_forks_repo_name: devonhollowood/radix-sort


text:
#!/usr/bin/env stack
-- stack --resolver lts-12.4 script
{-# LANGUAGE DeriveDataTypeable, ParallelListComp #-}
module Main ( main ) where

import Hledger

import Control.Exception ( bracket )
import Control.Monad
import System.Exit
import System.IO
import Data.Time.Calendar
import Text.Printf
import Data.Function (on)
import Data.List
import Data.Ord
import Statistics.Math.RootFinding
import Data.Decimal
import qualified Data.Text as T
import System.Console.CmdArgs
import Text.Tabular as Tbl
import qualified Text.Tabular.AsciiArt as Ascii

data Options = Options
  { cashFlow    :: Bool
  , debug       :: Bool
  , file        :: FilePath
  , investments :: String
  , pnl         :: String
  , begin       :: Maybe String
  , end         :: Maybe String
  , daily       :: Bool
  , weekly      :: Bool
  , monthly     :: Bool
  , yearly      :: Bool
  , period      :: Maybe String
  } deriving (Show, Data, Typeable)

queryFromOptions kind day accessor opts =
  -- We ignore QueryOpts with fst as they are irrelevant for this program
  case parseQuery day (T.pack $ accessor opts) of
    (Any, _) -> error $ "you need to specify " ++ kind ++ " query"
    (query, []) -> query
    (_,     opts) -> error $ "this program cannot process query options " ++ show opts

intervalAndSpanFromOptions :: Day -> Options -> (Maybe Interval, DateSpan)
intervalAndSpanFromOptions date opts =
  let begin' = fmap parsedate $ begin opts
      end' = fmap parsedate $ end opts
  in
   case (begin', end', daily opts, weekly opts, monthly opts, yearly opts, period opts) of
     (Nothing, Nothing, False, False, False, False, Nothing) -> (Nothing, nulldatespan)
     (b, e, True,  False, False, False, Nothing) -> (Just (Days 1),   DateSpan b e)
     (b, e, False, True,  False, False, Nothing) -> (Just (Weeks 1),  DateSpan b e)
     (b, e, False, False, True,  False, Nothing) -> (Just (Months 1), DateSpan b e)
     (b, e, False, False, False, True,  Nothing) -> (Just (Years 1),  DateSpan b e)
     (Nothing, Nothing, False, False, False, False, Just pexp) ->
       let (i, s) = parsePeriodExpr' date (T.pack pexp) in (Just i, s)
     _ -> error "Cannot work with this combination of date/period flags"

options =
  Options {
     cashFlow =
      False &= name "c"
      &= help "also show all relevant transactions"
    , debug =
      def
      &= help "print debugging info"
    , file =
      def &= name "f"
      &= typFile
      &= help "input ledger file (pass '-' for stdin)"
    , investments =
      def &= name "i"
      &= typ "QUERY"
      &= help "investments query"
    , pnl =
      def &= name "P"
      &= typ "QUERY"
      &= help "profit-and-loss query"
    , begin =
      Nothing &= name "b"
      &= typ "DATE"
      &= help "calculate interest from this date"
    , end =
      Nothing &= name "e"
      &= typ "DATE"
      &= help "calculate interest until this date"
    , daily =
      def &= name "D"
      &= help "calculate interest for each day"
    , weekly =
      def &= name "W"
      &= help "calculate interest for each week"
    , monthly =
      def &= name "M"
      &= help "calculate interest for each month"
    , yearly =
      def &= name "Y"
      &= help "calculate interest for each year"
    , period =
      Nothing &= name "p"
      &= typ "PERIODEXP"
      &= help "set start date, end date, and/or report interval all at once (overrides the flags above)"
    }
  &= program "hledger-roi"
  &= summary "compute return-on-investment for your portfolio using IRR and TWR"
  &= details [""]

main ::  IO ()
main = bracket (return ()) (\() -> hFlush stdout >> hFlush stderr) $ \() -> do
  opts <-cmdArgs options
  thisDay <- getCurrentDay

  let investmentsQuery = queryFromOptions "investments" thisDay investments opts
  let pnlQuery         = queryFromOptions "pnl" thisDay pnl opts

  jnl <- readJournalFile definputopts (file opts) >>= either fail return

  let trans = jtxns $ filterJournalTransactions investmentsQuery jnl

  when (null trans) $ do
    putStrLn "No relevant transactions found. Check your investments query"
    exitFailure

  when (debug opts) $ do
    putStrLn "DEBUG everything:"
    print $ trans
    putStrLn "DEBUG just non-PnL transactions:"
    print $ filter (matchesTransaction (And [investmentsQuery, Not pnlQuery])) trans
    putStrLn "DEBUG just PnL transactions:"
    print $ filter (matchesTransaction pnlQuery) trans

  let (requestedInterval, requestedSpan) = intervalAndSpanFromOptions thisDay opts
  let existingSpan =
        let dates = map transactionDate2 trans in
        DateSpan (Just $ minimum dates) (Just $ addDays 1 $ maximum dates)
  let wholeSpan = spanDefaultsFrom requestedSpan existingSpan

  let spans = case requestedInterval of
        Nothing -> [wholeSpan]
        Just interval ->
            splitSpan interval $
            spanIntersect existingSpan wholeSpan

  tableBody <- (flip mapM) spans $ \(DateSpan (Just ibegin) (Just iend)) -> do
    -- Spans are [begin,end), and iend here is 1 days after actual end date we are interested in
    let valueBeforeThisPeriod =
          total trans (And [ investmentsQuery
                           , Date (DateSpan Nothing (Just ibegin))])

    let valueAtTheEndOfPeriod  =
          total trans (And [investmentsQuery, Date (DateSpan Nothing (Just iend))])

    let cashFlowInThisPeriod =
          calculateCashFlow trans (And [ Not investmentsQuery
                                       , Not pnlQuery
                                       , Date (DateSpan (Just ibegin) (Just iend)) ] )

    irr <- internalRateOfReturn opts ibegin iend valueBeforeThisPeriod valueAtTheEndOfPeriod cashFlowInThisPeriod
    twr <- timeWeightedReturn opts investmentsQuery trans ibegin iend valueBeforeThisPeriod valueAtTheEndOfPeriod cashFlowInThisPeriod
    let cashFlowAmt = negate $ sum $ map snd cashFlowInThisPeriod
    return [ showDate ibegin
           , showDate (addDays (-1) iend)
           , show valueBeforeThisPeriod
           , show cashFlowAmt
           , show valueAtTheEndOfPeriod
           , show (valueAtTheEndOfPeriod - (valueBeforeThisPeriod + cashFlowAmt))
           , printf "%0.2f%%" irr
           , printf "%0.2f%%" twr ]

  let table = Tbl.Table
              (Group NoLine (map (Header . show) (take (length tableBody) [1..])))
              (Group DoubleLine
               [ Group SingleLine [Header "Begin", Header "End"]
               , Group SingleLine [Header "Value (begin)", Header "Cashflow", Header "Value (end)", Header "PnL"]
               , Group SingleLine [Header "IRR", Header "TWR"]])
              tableBody

  putStrLn $ Ascii.render id id id table

timeWeightedReturn opts investmentsQuery trans ibegin iend valueBefore valueAfter cashFlowInThisPeriod = do
  -- Span is [ibegin,iend), and iend is 1 days after actual end date we are interested in
  let initialUnitPrice = 100
  let initialUnits = valueBefore / initialUnitPrice
  let cashflow =
        -- Aggregate all entries for a single day, assuming that intraday interest is negligible
        map (\date_cash -> let (dates, cash) = unzip date_cash in (head dates, sum cash))
        $ groupBy ((==) `on` fst)
        $ sortBy (comparing fst)
        $ map (\(d,a) -> (d, negate a))
        $ filter ((/=0).snd) cashFlowInThisPeriod

  let units =
        tail $
        (flip scanl)
        (0,0,0,initialUnits)
        (\(_,_,_,unitBalance) (date, amt) ->
          let valueOnDate =
                total trans (And [investmentsQuery, Date (DateSpan Nothing (Just date))])
              unitPrice = if unitBalance == 0.0 then initialUnitPrice else valueOnDate / unitBalance
              unitsBoughtOrSold = amt / unitPrice
          in
           (valueOnDate, unitsBoughtOrSold, unitPrice, unitBalance + unitsBoughtOrSold)
        )
        cashflow

  let finalUnitBalance = if null units then initialUnits else let (_,_,_,u) = last units in u
      finalUnitPrice = valueAfter / finalUnitBalance
      totalTWR = roundTo 2 $ (finalUnitPrice - initialUnitPrice)
      years = (fromIntegral $ diffDays iend ibegin)/365 :: Double
      annualizedTWR = 100*((1+(realToFrac totalTWR/100))**(1/years)-1) :: Double

  let s d = show $ roundTo 2 d
  when (cashFlow opts) $ do
    printf "\nTWR cash flow for %s - %s\n" (showDate ibegin) (showDate (addDays (-1) iend))
    let (dates', amounts') = unzip cashflow
        (valuesOnDate',unitsBoughtOrSold', unitPrices', unitBalances') = unzip4 units
        add x lst = if valueBefore/=0 then x:lst else lst
        dates = add ibegin dates'
        amounts = add valueBefore amounts'
        unitsBoughtOrSold = add initialUnits unitsBoughtOrSold'
        unitPrices = add initialUnitPrice unitPrices'
        unitBalances = add initialUnits unitBalances'
        valuesOnDate = add 0 valuesOnDate'

    putStr $ Ascii.render id id id
      (Tbl.Table
       (Group NoLine (map (Header . showDate) dates))
       (Group DoubleLine [ Group SingleLine [Header "Portfolio value", Header "Unit balance"]
                         , Group SingleLine [Header "Cash", Header "Unit price", Header "Units"]
                         , Group SingleLine [Header "New Unit Balance"]])
       [ [value, oldBalance, amt, prc, udelta, balance]
       | value <- map s valuesOnDate
       | oldBalance <- map s (0:unitBalances)
       | balance <- map s unitBalances
       | amt <- map s amounts
       | prc <- map s unitPrices
       | udelta <- map s unitsBoughtOrSold ])

    printf "Final unit price: %s/%s=%s U.\nTotal TWR: %s%%.\nPeriod: %.2f years.\nAnnualized TWR: %.2f%%\n\n" (s valueAfter) (s finalUnitBalance) (s finalUnitPrice) (s totalTWR) years annualizedTWR

  return annualizedTWR

internalRateOfReturn opts ibegin iend valueBefore valueAfter cashFlowInPeriod = do
  -- Span is [ibegin,iend), and iend is 1 days after actual end date we are interested in
  let prefix = (ibegin, negate valueBefore)

      postfix = (iend, valueAfter)

      totalCF = filter ((/=0) . snd) $ prefix : (sortBy (comparing fst) cashFlowInPeriod) ++ [postfix]

  when (cashFlow opts) $ do
    printf "\nIRR cash flow for %s - %s\n" (showDate ibegin) (showDate (addDays (-1) iend))
    let (dates, amounts) = unzip totalCF
    putStrLn $ Ascii.render id id id
      (Tbl.Table
       (Group NoLine (map (Header . showDate) dates))
       (Group SingleLine [Header "Amount"])
       (map ((:[]) . show) amounts))

  -- 0% is always a solution, so require at least something here
  case ridders 0.00001 (0.000000000001,10000) (interestSum iend totalCF) of
    Root rate -> return ((rate-1)*100)
    NotBracketed -> error "Error: No solution -- not bracketed."
    SearchFailed -> error "Error: Failed to find solution."

type CashFlow = [(Day, Quantity)]

interestSum :: Day -> CashFlow -> Double -> Double
interestSum referenceDay cf rate = sum $ map go cf
    where go (t,m) = (fromRational $ toRational m) * (rate ** (fromIntegral (referenceDay `diffDays` t) / 365))

calculateCashFlow :: [Transaction] -> Query -> CashFlow
calculateCashFlow trans query = map go trans
    where
    go t = (transactionDate2 t, total [t] query)

total :: [Transaction] -> Query -> Quantity
total trans query = unMix $ sumPostings $ filter (matchesPosting query) $ concatMap realPostings trans

unMix :: MixedAmount -> Quantity
unMix a =
  case (normaliseMixedAmount $ costOfMixedAmount a) of
    (Mixed [a]) -> aquantity a
    _ -> error "MixedAmount failed to normalize"

max_stars_repo_name: adept/hledger-roi

max_issues_repo_name: adept/hledger-roi

max_forks_repo_name: adept/hledger-roi