{-# 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)


{-
 -  <DeclList> ::= { (<FunDecl> | <ValDecl>) }
 -
 -  <FunDecl> ::= 'fun' <Ident> <Ident> '=' <Exp>
 -
 -  <ValDecl> ::= 'val' <Ident> '=' <Exp>
 -
 -  <Expr> ::= <Comp>
 -           | 'if' <Expr> 'then' <Expr> 'else' <Expr>
 -           | 'let' <DeclList> 'in' <Expr> 'end'
 -           | 'fn' <Ident> '=>' <Expr>
 -
 -  <Comp> ::= <Plus> { ('>' | '=' | '<') <Plus> }
 -
 -  <Plus> ::= <Mult> { ('+' | '-') <Mult> }
 -
 -  <Mult> ::= <App>  { ('*' | '/') <App> }
 -
 -  <App>  ::= <Fact> { <Fact> }
 -
 -  <Fact> ::= '(' <Expr> ')'
 -          |  <Integer>
 -          |  <Identifier>
 -}

 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE RankNTypes #-}

 {-# LANGUAGE FlexibleContexts #-}

 module Hw8 () where

 -- import Hw1
 -- import Hw2
 -- import Hw3
 -- import Hw4
 -- import Hw5
 -- import Hw6
 -- import Hw7
 import Data.Time.Clock
 import Data.Complex

 import Text.Parsec
 import Text.Parsec.Expr
 import Text.Parsec.Prim (unexpected)
 import Text.Parsec.Language (emptyDef)
 import Data.Either
 import qualified Text.Parsec.Token as Token

 -- Abstract Syntax Tree ----------------------------------------------------------------------------------

 -- function/variable declaration
 data Decl = Fun String String Exp -- fun f x = e;
           | Val String Exp        -- val x = e;

 -- expressions
 data Exp = Lt Exp Exp     -- e1 < e2
         | Gt Exp Exp     -- e1 > e2
         | Eq Exp Exp     -- e1 = e2
         | Plus Exp Exp   -- e1 + e2
         | Minus Exp Exp  -- e1 - e2
         | Times Exp Exp  -- e1 * e2
         | Div Exp Exp    -- e1 div e2
         | Var String     -- x
         | If Exp Exp Exp -- if e0 then e1 else e2
         | Fn String Exp  -- fn x => e
         | Let [Decl] Exp -- let val x = e0; fun f = e1; in e2 end
         | App Exp Exp    -- e1 e2
         | Const Integer  -- n

 instance Show Decl where
   show (Fun f x e) = "fun " ++ f ++ " " ++ x ++ " = " ++ show e
   show (Val x e) = "val " ++ x ++ " = " ++ show e

 instance Show Exp where
   show (Const x) = show x
   show (Plus t1 t2) = showBop t1 "+" t2
   show (Times t1 t2) = showBop t1 "*" t2
   show (Minus t1 t2) = showBop t1 "-" t2
   show (Div t1 t2) = showBop t1 "/" t2
   show (Lt t1 t2) = showBop t1 "<" t2
   show (Gt t1 t2) = showBop t1 ">" t2
   show (Eq t1 t2) = showBop t1 "=" t2
   show (If t0 t1 t2) = "if " ++ show t0 ++ " then " ++ show t1 ++ " else " ++ show t2
   show (Var s) = s
   show (App t1 t2) = "(" ++ show t1 ++ " " ++ show t2 ++ ")"
   show (Fn x t) = "(fn " ++ x ++ " => " ++ show t ++ ")"
   show (Let decls e) = "(let " ++ show decls ++ " in " ++ show e ++ " end)"

 showBop t1 op t2 = "(" ++ show t1 ++ " " ++ op ++ " " ++ show t2 ++ ")"

 -- Parser definitions ----------------------------------------------------------------------------------

 type Parser = Parsec [Char] ()

 -- Runner

 -- print AST or error
 run :: Show a => Parser a -> String -> IO()
 run p input = putStrLn $ case parse p "" input of
       Left error -> show error
       Right x    -> show x

 -- return Either AST or parse error
 runParser :: Parser a -> String -> Either ParseError a
 runParser p input = parse p "" input

 -- Lexer

 lexeme p = do { x <- p; skipMany space <?> ""; return x }  -- get rid of trailing spaces

 symbol name = lexeme $ string name                         -- parse string and then remove spaces

 lexer :: Token.TokenParser ()
 lexer = Token.makeTokenParser style
   where style = emptyDef {
           Token.reservedOpNames = ["+","*","-","=", ">", "<", "in"],
           Token.reservedNames = ["let", "in", "end"]
           }

 parens :: Parser a -> Parser a
 parens = Token.parens lexer

 ident :: Parser String                                     -- raw identifier
 ident = (lexeme $ many1 letter) <?> "identifier"

 integer :: Parser Integer                                  -- parse integer and return numeric value
 integer = (read <$> (lexeme $ many1 digit)) <?> "integer"

 keyword :: String -> Parser ()                             -- make keyword parser with backtrack
 keyword name = (try $ symbol name) >> return ()

 if_ = keyword "if"                                         -- keyword parsers
 then_ = keyword "then"
 else_ = keyword "else"
 fn_ = keyword "fn"
 fun_ = keyword "fun"
 let_ = keyword "let"
 in_ = keyword "in"
 end_ = keyword "end"
 val_ = keyword "val"
 assign_ = keyword "="
 gt_ = keyword ">"
 lt_ = keyword "<"
 eq_ = keyword "="
 pl_ = keyword "+"
 mn_ = keyword "-"
 tm_ = keyword "*"
 dv_ = keyword "/"
 fatArrow_ = keyword "=>"
 sCol_ = keyword ";"

 identifier :: Parser String                                -- identifier that won't collide with keywords
 identifier = try $ do
     name <- ident
     if name `elem` reserved
       then unexpected ("reserved word " ++ name) -- throw parse error
       else return name
   where
     reserved = ["if", "then", "else", "fn", "fun", "let", "in", "end", "val"]

 -- Parser for the (phrase) grammar

 prog :: Parser [Decl]
 prog = skipMany space >> declList                              -- skip leading spaces

 declList :: Parser [Decl]                                      -- parse a list of function or variable declarations
 declList = many $ valDecl <|> funDecl

 funDecl :: Parser Decl
 funDecl = do
   fun_
   name <- identifier
   param <- identifier
   assign_
   t1 <- parseExpression
   return (Fun name param t1)

 valDecl :: Parser Decl
 valDecl = do
   val_
   name <- identifier
   assign_
   t1 <- parseExpression
   return (Val name t1)

 parseExpression :: Parser Exp
 parseExpression =  parseIf <|> parseLet <|> expr

 parseApp :: Parser Exp
 parseApp = ((factor `chainl1` (return App)) <?> "parseApp expr")

 expr = expr' `chainl1` compOp
 expr' = term `chainl1` addOp
 term    = parseApp `chainl1` mulOp
 factor  =  parseVar <|> parseInt <|> parseFn <|> parens expr

 compOp = do {gt_ ; return gtExpr}
         <|> do {lt_ ; return ltExpr}
         <|> do {eq_ ; return eqExpr}

 mulOp   =   do{ tm_ ; return tmExpr   }
         <|> do{ dv_ ; return dvExpr }

 addOp   =   do{ pl_ ; return plExpr }
         <|> do{ mn_ ; return mlExpr }

 gtExpr :: Exp -> Exp -> Exp
 gtExpr = Gt

 ltExpr :: Exp -> Exp -> Exp
 ltExpr = Lt

 eqExpr :: Exp -> Exp -> Exp
 eqExpr = Eq

 tmExpr :: Exp -> Exp -> Exp
 tmExpr = Times

 dvExpr :: Exp -> Exp -> Exp
 dvExpr = Div

 plExpr :: Exp -> Exp -> Exp
 plExpr = Plus

 mlExpr :: Exp -> Exp -> Exp
 mlExpr = Minus

 parseIf :: Parser Exp
 parseIf = do
   if_
   cond <- parseExpression
   then_
   true <- parseExpression
   else_
   false <- parseExpression
   return (If cond true false)

 parseFn :: Parser Exp
 parseFn = do
   fn_
   name <- identifier
   fatArrow_
   exp <- parseExpression
   return (Fn name exp)

 -- parseLet :: Parser Exp
 -- parseLet = do
 --   let_
 --   decl <- parseExpression
 --   in_
 --   exp <- parseExpression
 --   end_
 --   return (Let decl exp)

 parseVar :: Parser Exp
 parseVar = Var `fmap` identifier

 parseInt :: Parser Exp
 parseInt = Const `fmap` integer

 parseLet :: Parser Exp
 parseLet = do
   let_
   decls <- declList
   in_
   exp <- parseExpression
   end_
   return (Let decls exp)
 -- Test code ----------------------------------------------------------------------------------

--  main :: IO ()
--  main = do

--          let fact = "fun fact x = if x < 1 then 1 else x * fact (x-1)\n"
--          let it = "val it = let val y = 10 val z = 15 in fact y end"
--          let fnFn = "fun f x = fn z => y"

--          let d = (fact ++ it)
--          let test =  "val x = x * 5 / 5 - 5\n"
--          let test2 =  "fun fact x = 5\n"
--          let test3 = "if x > 1 then 1 else -1"
--          let test4 = "fun fact x = y"
--          let bla = (test ++ test2 ++ test3)
--          -- run prog bla
--          run prog d

 -- main = do
     -- Hw2
     -- let n = 2^5
     -- let s1 = map (\x -> sin(20*pi*x) + sin(40*pi*x)/2) $ range 0 1 n
     -- -- print $ split [1..64]
     -- -- print(rd 3 s1)
     -- start <- getCurrentTime
     -- let dft1 = map (\x -> x/n) $ absolute $ dft s1
     -- print(rd 2 dft1)
     -- end <- getCurrentTime
     -- print (diffUTCTime end start)
     -- start2 <- getCurrentTime
     -- let fft1 = map (\x -> x/n) $ absolute $ fft s1
     -- print(rd 2 fft1)
     -- end2 <- getCurrentTime
     -- print (diffUTCTime end2 start2)
     -- let v1 = Vec [1.0,2.0,3.0]
     -- let v2 = Vec [2,3,4]
     -- let v3 = Vec [-10,0,10]

     -- Hw3
     -- print $ v1 == v2
     -- print $ show v3
     -- print $ v1 + v2
     -- print $ v1 - v2
     -- print $ v1 * v2
     -- print $ v1 / v3
     -- print $ negate v1
     -- print $ signum v3
     -- print $ abs v3
     -- print $ foldr (*) 1 v1
     -- print $ sin $ v1 * (pi / 2)
     -- print $ v1 + (pure' $ sqrt 2)
     -- print $ realV v1
     -- print $ pure' 'a'

     -- let v1 = Vec [1.123,2.213,3.321]
     -- let v2 = Vec [2,3,4]
     -- let v3 = Vec [-10,0,10]
     -- print $ v1 + v2
     -- print $ v1 - v2
     -- print $ v1 * v2
     -- print $ v1 / v2
     -- print $ negate v1
     -- print $ signum v3
     -- print $ abs v3
     -- print $ v1 + 10
     -- print $ v2 + 1.2
     -- print $ v1 + (pure' $ sqrt 2)
     -- print $ realV v1
     -- print $ imagV v1
     -- print $ realV v1 + imagV v2
     -- print $ sin $ v1 * (pi / 2)
     -- print $ sum v1

     -- Hw4
     -- let n = fromIntegral 2^8
     -- let s1 = realV $ Vec $ fmap (\x -> sin(20*pi*x) + sin(40*pi*x)/2) $ range 0 1 n
     -- print $ fmap (\z -> conjugate z) s1
     -- print s1
     -- start <- getCurrentTime
     -- let dft1 = fmap (/n) $ absolute $ dft s1
     -- print(rd 2 dft1)
     -- let n = fromIntegral 2^8
     -- let s1 = fmap (\x -> sin(20*pi*x) + sin(40*pi*x)/2) $ Vec $ range 0 1 n
     -- print(rd 3 s1)
     -- print s1
     -- let dft1 = fmap (/n) $ absolute $ dft s1
     -- print(rd 2 dft1)
     -- print(rd 2 $ fmap (/n) $ absolute $ idft $ dft s1)
     -- print(rd 2 $ fmap (/n) $ absolute $ ifft $fft s1)

     -- end <- getCurrentTime
     -- print (diffUTCTime end start)
     -- start2 <- getCurrentTime
     -- let fft1 = fmap (/n) $ absolute $ fft s1
     -- print(rd 2 fft1)
     -- end2 <- getCurrentTime
     -- print (diffUTCTime end2 start2)

     -- let v1 = Vec [1,2,3]
     -- print $ v1 + v2
     -- print $ rd 2 v1
     -- let n = fromIntegral 2^8
     -- let s1 = fmap (\x -> sin(20*pi*x) + sin(40*pi*x)/2) $ Vec $ range 0 1 n
     -- print(rd 3 s1)
     -- print(rd 3 $ fmap (\(r:+_) -> r) $ low_pass' 15 $ realV s1)
     -- print(rd 3 $ fmap (\(r:+_) -> r) $ low_pass 15 $ realV s1)
     -- let a = [2,4,1,11,9]
     -- let b = [3,1,5,0,2]
     -- print $ maxList a
     -- print $ maxList' (<) a
     -- print $ rd 1 $ makeNoise 12345 10 0.0 1.0
     -- let n' = 2^8
     -- let noise = makeNoise 0 n' (-1.0) 1.0
     -- let n = fromIntegral n'
     -- let s = fmap (\x -> sin(20*pi*x) + sin(40*pi*x)/2) $ Vec $ range 0 1 n
     -- let s1 = s + noise
     -- -- noisy signal
     -- print(rd 2 s1)
     -- -- noisy frequency
     -- print (rd 2 $ fmap (/n) $ absolute $ dft $ realV s1)
     -- -- filtered signal (dft)
     -- print (rd 3 $ fmap (\(r:+_) -> r) $ low_pass' 15 $ realV s1)
     -- -- filtered signal (fft)
     -- print (rd 2 $ fmap (\(r:+_) -> r) $ low_pass 15 $ realV s1)


module PropertiesTests
 ( properties
 ) where

import GradientDescent
import Test.Tasty
import qualified Test.Tasty.QuickCheck as QC
import qualified Test.Tasty.SmallCheck as SC
import Numeric.LinearAlgebra

properties :: TestTree
properties = testGroup "Properties" [scProps, qcProps]

scProps =
  testGroup
    "(checked by SmallCheck)"
    [checkTranspositionSC, checkMatrixSigmoidSC, checkMatrixSigmoidPrimeSC]

checkTranspositionSC = SC.testProperty "matrix == tr . tr matrix" $
                           \list -> ((length list >< 1) (list :: [Double]) :: Matrix Double)
                              == (tr $ tr $ (length list><1) (list :: [Double]) :: Matrix Double)

checkMatrixSigmoidSC = SC.testProperty "check matrixSigmoid" $
                         \list ->
                            (length (list :: [Double])) >=
                              1 SC.==> ((1 >< length list) (sigmoid <$> list) :: Matrix Double)
                                == matrixSigmoid ((1 >< length list) list :: Matrix Double)

checkMatrixSigmoidPrimeSC = SC.testProperty "check matrixSigmoidPrime" $
                         \list ->
                            (length (list :: [Double])) >=
                              1 SC.==> ((1 >< length list) (sigmoidPrime <$> list) :: Matrix Double)
                                == matrixSigmoidPrime ((1 >< length list) list :: Matrix Double)

qcProps =
  testGroup
    "(checked by QuickCheck)"
    [checkTranspositionQC, checkMatrixSigmoidQC, checkMatrixSigmoidPrimeQC]

checkTranspositionQC = QC.testProperty "matrix == tr . tr matrix" $
                           \list -> ((length list >< 1) (list :: [Double]) :: Matrix Double)
                              == (tr $ tr $ (length list><1) (list :: [Double]) :: Matrix Double)

checkMatrixSigmoidQC = QC.testProperty "check matrixSigmoid" $
                         \list ->
                            (length (list :: [Double])) >=
                              1 QC.==> ((1 >< length list) (sigmoid <$> list) :: Matrix Double)
                                == matrixSigmoid ((1 >< length list) list :: Matrix Double)

checkMatrixSigmoidPrimeQC = QC.testProperty "check matrixSigmoidPrime" $
                         \list ->
                            (length (list :: [Double])) >=
                              1 QC.==> ((1 >< length list) (sigmoidPrime <$> list) :: Matrix Double)
                                == matrixSigmoidPrime ((1 >< length list) list :: Matrix Double)


module P18b where
import Data.Complex (magnitude)
import Data.Char (isNumber)
import Data.Maybe (fromMaybe)
import Data.List (tails)
import Data.Tuple (swap)

data SnailNumber = RegularNumber Int | NumberPair SnailNumber SnailNumber deriving (Eq, Show)

run :: String -> Int
run = maximum . map (snailMagnitude . snailSum) . pairs . parse

parse :: String -> [SnailNumber]
parse = map (fst . parseNumber) . lines

parseNumber :: String -> (SnailNumber, String)
parseNumber ('[':xs) =
    let (left, ',':xs') = parseNumber xs
        (right, xs'') = parseNumber xs'
    in (NumberPair left right, tail xs'')
parseNumber xs =
    let (n, xs') = span isNumber xs
    in (RegularNumber (read n), xs')

snailSum :: (SnailNumber, SnailNumber) -> SnailNumber
snailSum (a, b) = snailAdd a b

snailAdd :: SnailNumber -> SnailNumber -> SnailNumber
snailAdd a b = snailReduce (snailAdd' a b)

snailAdd' :: SnailNumber -> SnailNumber -> SnailNumber
snailAdd' = NumberPair

snailReduce :: SnailNumber -> SnailNumber
snailReduce x =
    let x' = snailReduce' x
    in if x == x' then x' else snailReduce x'

snailReduce' :: SnailNumber -> SnailNumber
snailReduce' x =
    let x' = explode x
    in fromMaybe x' (split x')

explode :: SnailNumber -> SnailNumber
explode x =
    let (x', _, _) = explode' 0 x
    in x'

explode' :: Int -> SnailNumber -> (SnailNumber, Int, Int)
explode' 4 (NumberPair (RegularNumber left) (RegularNumber right)) = (RegularNumber 0, left, right)
explode' depth (NumberPair left right) =
    let (left', expLeft1, expRight1) = explode' (depth+1) left
        (right', expLeft2, expRight2) = explode' (depth+1) (incrementLeft right expRight1)
        left'' = incrementRight left' expLeft2
    in  (NumberPair left'' right', expLeft1, expRight2)
explode' _ x = (x, 0, 0)

incrementLeft :: SnailNumber -> Int -> SnailNumber
incrementLeft (RegularNumber x) a = RegularNumber (x+a)
incrementLeft (NumberPair left right) x = NumberPair (incrementLeft left x) right

incrementRight :: SnailNumber -> Int -> SnailNumber
incrementRight (RegularNumber x) a = RegularNumber (x+a)
incrementRight (NumberPair left right) x = NumberPair left (incrementRight right x)

split :: SnailNumber -> Maybe SnailNumber
split n@(RegularNumber x) = if x > 9 then Just (NumberPair (RegularNumber x') (RegularNumber (x - x'))) else Nothing
    where x' = x `div` 2
split (NumberPair left right) =
    case split left of
        Just left' -> Just (NumberPair left' right)
        Nothing -> case split right of
            Just right' -> Just (NumberPair left right')
            Nothing -> Nothing

snailMagnitude :: SnailNumber -> Int
snailMagnitude (RegularNumber x) = x
snailMagnitude (NumberPair left right) = 3*snailMagnitude left + 2*snailMagnitude right

pairs :: [a] -> [(a, a)]
pairs xs =
    let ps = concatMap (\xs' -> map (\x -> (head xs', x)) (tail xs')) $ filter (not . null) (tails xs)
    in ps ++ map swap ps


module Trainer where

import Numeric.LinearAlgebra
import Models
import Layer
import Optimization
import Data.Foldable
import Control.Monad.IO.Class
import Utils

data TrainConfig = TrainConfig {
  trainConfigRowNum :: Int,
  trainConfigColNum :: Int,
  trainConfigLearningRate :: Double,
  trainConfigMomentum :: Double,
  trainConfigDecay :: Double,
  trainConfigEpsilon :: Double,
  trainConfigBatchSize :: Double,
  trainConfigOptimization:: String
                               }

trainSingleData :: Model VLayer -> TrainConfig -> Matrix R -> IO ()
trainSingleData (Model []) _ _ = print "No models !" >> return ()
trainSingleData neuModel trainConfig inputData =
  let r = trainConfigRowNum trainConfig
      c = trainConfigColNum trainConfig
      paramsList = genParamsList neuModel
      gparamsList = genGParamsList neuModel
      accGradList = accMatrixInit paramsList
      accDeltaList = accMatrixInit paramsList
      optimize = Adadelta {
                adaDecay = trainConfigDecay trainConfig,
                adaEpsilon = trainConfigEpsilon trainConfig,
                adaBatchSize = (trainConfigBatchSize trainConfig),
                adaParams = paramsList,
                adaGparams = gparamsList,
                adaAccGrad = accGradList,
                adaAccDelta = accDeltaList
                            } in
  -- begin to train
  forM_ [1..10000] $ \idx -> do
    --print idx
    let output = forward inputData neuModel in
      --print "backpropagate the model" >>
      let neuModel = backward "mse" output inputData neuModel in
      --print "neumodel" >>
      let gparamsList = updateGParamsList neuModel gparamsList in
      --print "continue" >>
      if mod idx 100 == 0 && idx >= 100 then
        -- update the params in optimize
        -- update the params in model
        let optimize = paramUpdate optimize
            paramsList = adaParams optimize
            neuModel = updateModelParams neuModel paramsList in
        print (getAverageMat output) >>
        return()
      else return()


{-# LANGUAGE Strict, StrictData #-}
{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
{-# LANGUAGE TypeFamilies, TypeFamilyDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE DefaultSignatures #-}

------------------------------------------------
-- |
-- Module    :  Numeric.Algebra.Presentation
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Matrix presentations of morphisms
--
------------------------------------------------

module Numeric.Algebra.Presentation where

import GHC.Generics (Generic)
import Control.DeepSeq (NFData, force)

import Control.Monad
import Control.Monad.ST

import Data.Foldable

import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.LinearAlgebra.Devel as LA

import Numeric.Algebra.FreeModule
import Numeric.Matrix.Integral
import Numeric.F2 (F2)
import qualified Numeric.Matrix.F2 as F2
import qualified Numeric.Matrix.F2.Mutable as F2

---------------
-- * Classes
---------------
-- | Class which provides matrix presentations for homomorphisms with coefficients in a.
class (Eq a, Num a, NFData a, Semigroup (Mat a), NFData (Vec a), NFData (Mat a)) => Coefficient a where
  type Vec a = v | v -> a
  -- ^ The type of vectors with coefficients in a.
  type Mat a = m | m -> a
  -- ^ The type of matrices with coefficients in a.
  ident :: Int -> Mat a
  -- ^ Presentation of the identity of the given rank.
  present :: (Integral a', NFData a', Ord c, NFData c) => (b->FreeMod a' c) -> [b] -> [c] -> Mat a
  -- ^ Present a linear map in a matrix form with respect to a given bases
  vecToList :: Vec a -> [a]
  -- ^ Vec a is convertible to a list.
  rankVec :: Vec a -> Int
  -- ^ The rank of the module where the given vector lies.
  rankDom :: Mat a -> Int
  -- ^ The rank of the domain of the given matrix.
  rankCod :: Mat a -> Int
  -- ^ The rank of the codomain of the given matrix.
  fromVecs :: (Foldable t) => t (Vec a) -> Mat a
  -- ^ Give a matrix whose image is spanned by given vectors.
  -- Typically, it is the matrix with given vectors as columns.
  toVecs :: Mat a -> [Vec a]
  -- ^ Give a set of vectors which span the image of the matrix.
  -- Typically, it is the set of column vectors.
  -- ^ The scalar matrix of the given size.
  (@|>) :: Mat a -> Vec a -> Vec a
  -- ^ Multiplication of matrices and vectors.
  scale :: a -> Vec a -> Vec a
  -- ^ Scalar multiplication on vectors.
  scalar :: Int -> a -> Mat a
  -- ^ Scalar * identity matrix.
  -- | Default implementations
  ident n = scalar n 1
  rankVec = length . vecToList
  (@|>) x = head . toVecs . (x<>) . fromVecs . (:[])
  scale x v = scalar (rankVec v) x @|> v
  scalar n x = fromVecs (map (scale x) (toVecs (ident n)))

-- | Matrix presentations with normal forms.
class (Coefficient a, NFData (NormalForm a)) => Normalized a where
  data NormalForm a
  -- ^ The type of normal forms.
  fromNF :: NormalForm a -> Mat a
  -- ^ Convert into a matrix form.
  -- >> let xnf = toNF x
  -- Then
  -- >> x <> basisDom xnf == basisCod xnf <> fromNF xnf
  basisDom :: NormalForm a -> Mat a
  -- ^ The basis matrix on the domain.
  -- It should be invertible.
  basisCod :: NormalForm a -> Mat a
  -- ^ The basis matrix on the codomain.
  -- It should be invertible.
  toNF :: Mat a -> NormalForm a
  -- ^ Compute the normal form.
  invariants :: NormalForm a -> [a]
  -- ^ Each normal form should contain the data of "invariant factors;" e.g. those for Smith normal forms of matrices with coefficients in a PID.

------------------------------------------------
-- * Instances in terms of LA.Z using hmatrix
------------------------------------------------
instance Coefficient (LA.Z) where
  type Vec LA.Z = LA.Vector LA.Z
  type Mat LA.Z = LA.Matrix LA.Z
  ident = LA.ident
  present f dom cod
    = let domRk = length dom
          codRk = length cod
      in runST $ do
    stMat <- LA.newUndefinedMatrix LA.RowMajor codRk domRk
    forM_ [0..domRk-1] $ \j -> do
      let !vals = force $ f (dom !!j)
      forM_ [0..codRk-1] $ \i -> do
        let coeff = getCoeff (cod !! i) vals
        LA.unsafeWriteMatrix stMat i j (fromIntegral coeff)
    LA.unsafeFreezeMatrix stMat
  vecToList = LA.toList
  rankVec = LA.size
  rankDom = LA.cols
  rankCod = LA.rows
  fromVecs = LA.fromColumns . toList
  toVecs = LA.toColumns
  (@|>) = (LA.#>)
  scale = LA.scale

instance Normalized (LA.Z) where
  -- | Presentation in Smith normal forms.
  data NormalForm LA.Z = SmithNF {
    _invFactor :: LA.Vector LA.Z, -- Invariant factor
    _basisS :: LA.Matrix LA.Z,    -- basis for source
    _basisT :: LA.Matrix LA.Z     -- basis for target
    } deriving (Show, Generic, NFData)
  fromNF (SmithNF invs bss bts)
    = runST $ do
    stMat <- LA.newMatrix 0 (LA.cols bts) (LA.cols bss)
    forM_ [0..(LA.size invs)-1] $ \i ->
      LA.writeMatrix stMat i i  (invs `LA.atIndex` i)
    LA.unsafeFreezeMatrix stMat
  basisDom = _basisS
  basisCod = _basisT
  toNF x
    = let (u,s,v) = smithRep x
      in SmithNF (LA.takeDiag s) v u
  invariants = LA.toList . _invFactor

---------------------------------
-- * Instances in terms of F2
---------------------------------
instance Coefficient F2 where
  type Vec F2 = F2.VectorF2
  type Mat F2 = F2.MatrixF2
  ident = F2.ident
  present f dom cod
    = let domRk = length dom
          codRk = length cod
      in runST $ do
    stMat <- F2.unsafeNewMatrixF2 codRk domRk
    forM_ [0..domRk-1] $ \j -> do
      let !vals = force $ f (dom !!j)
      forM_ [0..codRk-1] $ \i -> do
        let coeff = getCoeff (cod !! i) vals
        F2.unsafeWriteMatrixF2 stMat i j (fromIntegral coeff :: F2)
    F2.unsafeFreezeMatrixF2 stMat
  vecToList = F2.toList
  rankVec = F2.sizeF2
  rankDom = F2.cols
  rankCod = F2.rows
  fromVecs = F2.fromColumns . toList
  toVecs = F2.toColumns
  scalar = F2.scalar

instance Normalized F2 where
  data NormalForm F2 =
    DiagNF {-# UNPACK #-} !Int {-# UNPACK #-} !F2.MatrixF2 {-# UNPACK #-} !F2.MatrixF2
    deriving (Show, Generic,NFData)
  invariants (DiagNF rk _ _) = replicate rk 1
  basisDom (DiagNF _ bss _) = bss
  basisCod (DiagNF _ _ bst) = bst
  fromNF (DiagNF rk bss bst) = F2.diagRectL (rankCod bst) (rankCod bss) (repeat 1 :: [Int])
  toNF x
    = let (u,rk,v) = F2.diagRep x
      in DiagNF rk v u


--
-- FlattenLayer: flatten/unflatten layer
--

module CNN.FlattenLayer (
  flatten
, unflatten
) where

import qualified Numeric.LinearAlgebra as A

import CNN.Image
import CNN.LayerType

{- |
flatten: flatten image

  IN: Image (=[Matrix R])

>>> let im = [A.fromLists [[2.0,3.0],[4.0,5.0]], A.fromLists [[1.0,2.0],[4.0,3.0]], A.fromLists [[2.0,1.0],[6.0,5.0]]]
>>> flatten im
[(1><12)
 [ 2.0, 3.0, 4.0, 5.0, 1.0, 2.0, 4.0, 3.0, 2.0, 1.0, 6.0, 5.0 ]]

-}

flatten :: Image -> Image
flatten im = [A.fromLists [concat $ map (A.toList . A.flatten) im]]

{- |
unflatten: unflatten image

  IN: [Matrix R]

>>> let ds = [A.fromLists [[2.0,3.0,4.0,5.0,1.0,2.0,4.0,3.0,2.0,1.0,6.0,5.0]]]
>>> unflatten 2 2 ds
[(2><2)
 [ 2.0, 3.0
 , 4.0, 5.0 ],(2><2)
 [ 1.0, 2.0
 , 4.0, 3.0 ],(2><2)
 [ 2.0, 1.0
 , 6.0, 5.0 ]]

-}

unflatten :: Int -> Int -> Image -> Image
unflatten x y [m] = map (A.reshape x . A.fromList) $ split (y*x) v
  where
    v = A.toList $ A.flatten m
    split :: Int -> [a] -> [[a]]
    split _ [] = []
    split n ds = d : split n ds'
      where
        (d, ds') = splitAt n ds


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE IncoherentInstances #-}

module Test.Utils where

import Test.Tasty.QuickCheck
import Data.Complex
import Data.Matrix.Static.LinearAlgebra
import Data.Singletons
import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Sparse as S
import Data.Vector.Storable (Storable)
import Data.List.Ordered
import qualified Data.Vector as V
import Data.Ord
import qualified Data.Vector.Generic as G
import Data.AEq
import Control.Monad

class Approx a where
    (~=) :: a -> a -> Bool
    infixl 4 ~=

instance Approx Double where
    a ~= b = abs (a - b) < 1e-10

instance Approx Float where
    a ~= b = abs (a - b) < 1e-3

instance Approx (Complex Double) where
    a ~= b = r1 ~= r2 && i1 ~= i2
      where
        (r1, i1) = polar a
        (r2, i2) = polar b

instance (SingI m, SingI n, G.Vector v Bool, G.Vector v a, Approx a) =>
    Approx (D.Matrix m n v a) where
        m1 ~= m2 = D.all id $ D.zipWith (~=) m1 m2

instance (G.Vector v a, Arbitrary a, SingI m, SingI n)
    => Arbitrary (D.Matrix m n v a) where
        arbitrary = D.fromList <$> vector (m*n)
          where
            m = fromIntegral $ fromSing (sing :: Sing m)
            n = fromIntegral $ fromSing (sing :: Sing n)
        shrink _v = []

instance (G.Vector v a, Arbitrary a, SingI m, SingI n, S.Zero a)
    => Arbitrary (S.SparseMatrix m n v a) where
        arbitrary = do
            vals <- filter (/=S.zero) <$> vector p
            xs <- fmap (nubSortBy (comparing fst)) $ forM vals $ \v -> do
                i <- choose (0, m-1)
                j <- choose (0, n-1)
                return ((i,j),v)
            return $ S.fromTriplet $ V.fromList $ map (\((a,b),c) -> (a,b,c)) xs
          where
            p = (m * n) `div` 10
            m = fromIntegral $ fromSing (sing :: Sing m)
            n = fromIntegral $ fromSing (sing :: Sing n)
        shrink _v = []


module Main where

import QuickSpec
import Test.QuickCheck
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Poly(OrdA(..))
import Data.List hiding ( insert )
import Data.Ord ( comparing )
import Data.Complex

{- Simple laws about addition -}
addSig :: [Sig]
addSig = [
  con "+" ((+) :: Int -> Int -> Int),
  con "0" (0 :: Int)
  ]

{- Simple laws about concat and map -}
concatMapSig :: [Sig]
concatMapSig = [
  con "map" (map :: (A -> B) -> [A] -> [B]),
  con "concat" (concat :: [[A]] -> [A]),
  con "concatMap" (concatMap :: (A -> [B]) -> [A] -> [B]),
  con "[]" ([] :: [A]),
  con ":" ((:) :: A -> [A] -> [A])
  ]

{- Laws about equivalence of two functions -}
rev :: [a] -> [a]
rev [] = []
rev (x:xs) = rev xs ++ [x]

revAcc :: [a] -> [a]
revAcc = aux []
  where
    aux acc [] = acc
    aux acc (x:xs) = aux (x:acc) xs

revSig :: [Sig]
revSig = [
  con "rev" (rev :: [A] -> [A]),
  con "revAcc" (revAcc :: [A] -> [A]),
  con "[]" ([] :: [A])
  ]

{- Laws about sortedness -}
sorted :: Ord a => [a] -> Bool
sorted [] = True
sorted [x] = True
sorted (x:y:xs) = x <= y && sorted (y:xs)

sortedSig :: [Sig]
sortedSig = [
  predicate "sorted" (sorted :: [Int] -> Bool),
  background [
    predicate "<" ((<) :: Int -> Int -> Bool),
    prelude
    ]
  ]

{- Many, many laws about foldr -}
foldrSig :: [Sig]
foldrSig = [
  con "foldr" (foldr :: (A -> B -> B) -> B -> [A] -> B),
  background [
    con "length" (length :: [A] -> Int),
    con "const" (const :: A -> B -> A),
    con "+" ((+) :: Int -> Int -> Int),
    con "1" (1 :: Int),
    con "0" (0 :: Int)
    ]
  ]

{- Custom datatypes similar to https://github.com/nick8325/quickspec/blob/master/examples/Heaps.hs -}
data Tree a = Leaf | Node (Tree a) a (Tree a) deriving Typeable

instance Eq a => Eq (Tree a) where
  t1 == t2 = inorder t1 == inorder t2

instance Ord a => Ord (Tree a) where
  compare = comparing inorder

instance (Ord a, Arbitrary a) => Arbitrary (Tree a) where
  arbitrary = fromList <$> arbitrary

inorder :: Tree a -> [a]
inorder Leaf = []
inorder (Node l v r) = inorder l ++ [v] ++ inorder r

fromList :: Ord a => [a] -> Tree a
fromList = foldr insert Leaf

insert :: Ord a => a -> Tree a -> Tree a
insert v Leaf = Node Leaf v Leaf
insert v (Node l vt r)
  | v <= vt = Node (insert v l) vt r
  | otherwise = Node l vt (insert v r)

mirror :: Tree a -> Tree a
mirror Leaf = Leaf
mirror (Node l v r) = Node (mirror r) v (mirror l)

{- Some laws on trees -}
treeSig1 :: [Sig]
treeSig1 = [
  monoType (Proxy :: Proxy OrdA),
  monoTypeWithVars ["t", "t1", "t2"] (Proxy :: Proxy (Tree OrdA)),
  monoTypeWithVars ["t", "t1", "t2"] (Proxy :: Proxy (Tree Int)),
  con "mirror" (mirror :: Tree OrdA -> Tree OrdA),
  con "rev" (rev :: [OrdA] -> [OrdA]),
  con "fromList" (fromList :: [OrdA] -> Tree OrdA),
  con "insert" (insert :: OrdA -> Tree OrdA -> Tree OrdA)
  ]

{- A real challenge... -}
treeSig2 :: [Sig]
treeSig2 = [
  monoTypeWithVars ["a", "b", "c"] (Proxy :: Proxy OrdA),
  monoTypeWithVars ["t", "t1", "t2"] (Proxy :: Proxy (Tree OrdA)),
  monoTypeWithVars ["t", "t1", "t2"] (Proxy :: Proxy (Tree Int)),
  con "fromList" (fromList :: [OrdA] -> Tree OrdA),
  con "inorder" (inorder :: Tree OrdA -> [OrdA]),
  predicate "sorted" (sorted :: [OrdA] -> Bool)
  ]

{- A "complex" example... observational equivalence -}
complexSig :: [Sig]
complexSig = [
  monoTypeObserve (Proxy :: Proxy Float),
  --monoType (Proxy :: Proxy (Complex Int)),
  monoTypeObserve (Proxy :: Proxy (Complex Float)),
  con "realPart" (realPart :: Complex Float -> Float),
  con "imagPart" (imagPart :: Complex Float -> Float),
  con ":+" ((:+) :: Float -> Float -> Complex Float),
  background [ prelude ]
  ]

main :: IO ()
main = return ()


{-# LANGUAGE NoMonomorphismRestriction #-}
import Numeric.LinearAlgebra
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.Quaternion
import Linear.Matrix
import Vectorization
import Space
import Kalman

-- Aided Inertial Navigation -- Example 5.4

pe0 = (V2 (V2 100 0) (V2 0 100))
xe0 = V2 0 0
se0 = (xe0,pe0)

he0 (V2 x _) = V1 x
he1 (V2 _ y) = V1 y
he2 (V2 x y) = V1 (x*0.7 + 0.3*y)
he3 (V2 x y) = V1 (x*0.5 + y*0.5)

m0 = (he0,re0,ye0)
m1 = (he1,re1,ye1)
m2 = (he2,re2,ye2)
m3 = (he3,re3,ye3)

mlist = [m0,m1,m2,m3]

re0 = V1 1
re1 = V1 1
re2 = V1 1
re3 = V1 1

ye0 = V1 10.24
ye1 = V1 21.20
ye2 = V1 13.91
ye3 = V1 14.84

measure' (h,r,y) = measure h r y
sef = foldr measure' se0 mlist

se1 = measure he0 re0 ye0 se0
se2 = measure he1 re1 ye1 se1
se3 = measure he2 re2 ye2 se2
se4 = measure he3 re3 ye3 se3


{-# LANGUAGE ScopedTypeVariables, TypeFamilies #-}

module Fuml.Supervised where

import qualified Fuml.Base.LinearRegression as LR
import Fuml.Core
import Numeric.LinearAlgebra
import Control.Monad.Identity

newtype Beta = Beta { unBeta :: Vector Double }

ols :: Supervisor Identity Double Beta Double
ols = Supervisor $ \_ theData ->
              let beta = LR.ols theData
              in return $ linPredict beta

type Weight = Double

olsWeighted :: Supervisor Identity (Double, Weight) Beta Double
olsWeighted
   = Supervisor $ \_ theDataW ->
              let wvs = fromList $ map (\(_,(_,w)) -> w) theDataW
                  theData = map (\(xs,(y,_)) -> (xs,y)) theDataW
                  beta = LR.wols wvs theData
              in return $ linPredict beta

ridge :: Double -> Supervisor Identity Double Beta Double
ridge gammac = Supervisor $ \_ xys ->
  let npars = size $ fst $ head xys
      gamma = diag $ konst gammac npars
      beta = LR.ridge gamma xys
  in return $ linPredict beta

linPredict :: Vector Double -> Predict Beta Double
linPredict beta = Predict (Beta beta) (`dot` beta)


{-|

Copyright:
  This file is part of the package zxcvbn-hs. It is subject to the
  license terms in the LICENSE file found in the top-level directory
  of this distribution and at:

    https://code.devalot.com/sthenauth/zxcvbn-hs

  No part of this package, including this file, may be copied,
  modified, propagated, or distributed except according to the terms
  contained in the LICENSE file.

License: MIT

-}
module Text.Password.Strength.Internal.Math
  ( variations
  , variations'
  , bruteForce
  , caps
  ) where

--------------------------------------------------------------------------------
-- Library Imports:
import Control.Lens ((^.))
import Data.Char (isUpper)
import qualified Data.Text as Text
import Numeric.SpecFunctions (choose)

--------------------------------------------------------------------------------
-- Project Imports:
import Text.Password.Strength.Internal.Token

--------------------------------------------------------------------------------
-- | Equation 2, section 4, page 163 (8/18)
--
-- NOTE: The other implementations don't seem to divide the sum by two
--       but the equation in the paper clearly does.
variations :: Int -> Int -> Integer
variations 0 _ = 1
variations _ 0 = 1
variations u l = max 1 (sum (floor . choose (u+l) <$> [1 .. min u l]) `div` 2)

--------------------------------------------------------------------------------
-- | Like equation 2, but modified for l33t and keyboard variations.
--
-- This equation does not appear in the 2016 paper although it is
-- mentioned.  Therefore it was adapted from the CoffeeScript and
-- Python implementations.
variations' :: Int -> Int -> Integer
variations' 0 _ = 2
variations' _ 0 = 2
variations' u l = variations u l

--------------------------------------------------------------------------------
-- | Calculate the brute force score for text with the given length.
bruteForce :: Int -> Integer
bruteForce = (10 ^)

--------------------------------------------------------------------------------
-- | Score the use of uppercase letters according to the paper.  This
-- is specified in the paragraph preceding equation 2.
caps :: Token -> Integer -> Integer
caps token n =
  let text       = token ^. tokenChars
      upper      = Text.length (Text.filter isUpper text)
      lower      = Text.length text - upper
      allLower   = lower == Text.length text
      allUpper   = lower == 0
      firstUpper = upper == 1 && Text.all isUpper (Text.take 1 text)
      lastUpper  = upper == 1 && Text.all isUpper (Text.takeEnd 1 text)
  in case () of
    () | allLower   -> n
       | firstUpper -> n * 2
       | lastUpper  -> n * 2
       | allUpper   -> n * 2
       | otherwise  -> n * variations upper lower

--------------------------------------------------------------------------------
-- NOTE: Equation 3 is in Keybaord.hs.


module Main where

--import Data.Complex
import Graphics.Gloss
import Data.Complex
import Data.Ratio
import Fractals

window :: Display
window = InWindow "Mandelbrot set" (1000, 600) (10, 10)

background :: Color
background = black

drawing :: Complex Double -> Picture
--drawing = bmp 1000 600 (((-0.1011):+0.9563)::Complex Double) 0.0000001
drawing z0 = bmp 1000 600 z0 0.0001

main :: IO ()
--main = print $ order (mandelbrot ((-1.6) :+ 0.0)) (2.0::Double) 256 (0.0 :+ 0.0)
main = print point >> (display window background $ drawing point)
    where point = (tangentPoint (9 % 17))::Complex Double


{-# LANGUAGE TypeFamilies #-}
module Data.Running
  ( Coord (Coord)
  , PointRecord (PointRecord)
  , pointDistance
  , pointCoord
  , pointLat
  , pointLng
  , pointTime
  , pointSpeed
) where
import Data.Time.Clock (UTCTime)
import Data.Time.Format ()
import Statistics.Sample (mean, stdDev)
import Data.AdditiveGroup
import Data.VectorSpace

newtype Coord = Coord (Double, Double) deriving (Eq, Ord, Show)

data PointRecord = PointRecord
  { pointTime :: UTCTime
  , pointCoord :: Coord
  , pointDistance :: Double
  , pointSpeed :: Double
  } deriving (Eq, Ord, Show)

pointLat (PointRecord _ (Coord (a, _)) _ _) = a
pointLng (PointRecord _ (Coord (_, b)) _ _) = b

instance AdditiveGroup Coord where
  zeroV = Coord (0, 0)
  Coord (a, b) ^+^ Coord (c, d) = Coord (a + c, b + d)
  negateV (Coord (a, b)) = Coord (-a, -b)

instance InnerSpace Coord where
  Coord (a, b) <.> Coord (c, d) = a * c + b * d

instance VectorSpace Coord where
  type Scalar Coord = Double
  s *^ Coord (a, b) = Coord (s * a, s * b)

--instance VectorSpace PointRecord where
--  type Scalar =


module Impulse.Impulse where

import           Control.Monad             as M
import           Control.Monad.Parallel    as MP
import           Data.Array.Repa           as R
import           Data.Binary
import           Data.Complex
import           Data.List                 as L
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC.CompletionField
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:initDistStr:histFilePath:alphaStr:pinwheelFlagStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      initDist = read initDistStr :: [R2S1RPPoint]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numThread = read numThreadStr :: Int
      sourceDist = L.take 1 initDist
      sinkDist = L.drop 1 initDist
      folderPath = "output/test/Impulse"
      (R2S1RPPoint (_, _, _, s)) = L.head sourceDist
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  createDirectoryIfMissing True folderPath
  let arr4d =
        R.slice
          arrR2Z2T0S0
          (Z :. All :. All :. (L.length theta0Freqs - 1) :.
           (L.length scale0Freqs - 1) :.
           All :.
           All)
  createDirectoryIfMissing True (folderPath </> "GreensR2Z2T0S0")
  -- MP.mapM_
  --   (\(i, j) ->
  --      plotImageRepaComplex
  --        (folderPath </> "GreensR2Z2T0S0" </> show (i + 1) L.++ "_" L.++
  --         show (j + 1) L.++
  --         ".png") .
  --      ImageRepa 8 .
  --      computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice arr4d $
  --      (Z :. i :. j :. All :. All))
  --   [ (i, j)
  --   | i <- [0 .. (L.length thetaFreqs) - 1]
  --   , j <- [0 .. (L.length scaleFreqs) - 1]
  --   ]
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  sourceDistArr <-
    computeInitialDistributionR2T0S0
      plan
      numPoint
      numPoint
      theta0Freqs
      scale0Freqs
      maxScale
      -- initDist
      sourceDist
  arrR2Z2T0S0F <- dftR2Z2T0S0 plan . computeS . makeFilter2D $ arrR2Z2T0S0
  -- Source field
  sourceArr <- convolveR2T0S0 plan arrR2Z2T0S0F sourceDistArr
  sourceR2Z2 <- R.sumP . R.sumS . rotateR2Z2T0S0Array $ sourceArr
  sourceField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP .
    R.sumS .
    rotate4D .
    rotate4D . r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs $
    sourceR2Z2
  plotImageRepaComplex (folderPath </> printf "Source%.0f.png" s) . ImageRepa 8 $
    sourceField


module CNC.IntegrationTests where

import CNC.Declarative
import CNC.HCode

import Data.Complex
import Control.Monad

flake_side = do
  [p1,p2,p3,p4,p5] <- declarePoints 5
  xSize p1 p2 1
  len p2 p3 1
  len p3 p4 1
  xAngle p2 p3 (pi / 3)
  xAngle p3 p4 (- pi / 3)
  xSize p4 p5 1

renderPolygon :: Int -> Path -> HCode ()
renderPolygon n path = render_side 0 path n
  where
    render_side p0 path 0 = return ()
    render_side p0 path k = do
      let d = deltaPath path
      renderPath (posToExpr p0) path
      render_side (p0 + d) (rotate (- 2*pi / fromIntegral n) path) (k-1)

test1 = do putStrLn "a snowflake"
           print $ (figure flake_side)
           putHCode $ renderPolygon 3 (figure flake_side)

integrationTests = do
  putStrLn "Integration tests"
  test1


{-# OPTIONS_GHC -Wall #-}
module Diagrams.Spirograph
  ( Spiro(Spiro)
  , drawSpiros
  , drawSpirosIntermediate
  ) where

import           Data.Complex                 (Complex(..), mkPolar)
import           Diagrams.Prelude             (Diagram, P2)
import qualified Diagrams.Backend.SVG.CmdLine as SVG
import qualified Diagrams.Prelude             as D
import           Data.Monoid                  (Sum(..))

data Spiro = Spiro
  { offset ∷ Double
  , speed  ∷ Double
  }

newtype Path w = Path (Double → w)

deriving newtype instance Semigroup w ⇒ Semigroup (Path w)
deriving newtype instance Monoid w    ⇒ Monoid    (Path w)

newtype P = P (Sum (P2 Double))

runP ∷ P → P2 Double
runP (P (Sum p)) = p

deriving newtype instance Semigroup P
deriving newtype instance Monoid P

-- | @'disk' d@ is the path that @d@ will travel.
disk ∷ Spiro → Path P
disk Spiro{..} = Path $ \t → P $ Sum $ p2 $ mkPolar offset (speed * t)

disks ∷ [Spiro] → Path P
disks = foldMap disk

allSpiros ∷ [Spiro] → [Path P]
allSpiros = \case
  []     → mempty
  (x:xs) → px : map f pxs
    where
    px  = disk x
    pxs = allSpiros xs
    f   ∷ Path P → Path P
    f p = px <> p

-- | Form a point in diagrams in terms of a complex number.
p2 ∷ Complex a → P2 a
p2 (x :+ y) = D.p2 (x, y)

-- | Convert a path into a path.
sample ∷ Path P → [P2 Double]
sample (Path f) = runP . f <$> enumFromThenTo 0 0.1 100

drawSpiros ∷ [Spiro] → Diagram SVG.B
drawSpiros xs = D.fromVertices $ sample (disks xs)

drawSpirosIntermediate ∷ [Spiro] → Diagram SVG.B
drawSpirosIntermediate xs = foldMap step (allSpiros xs)
  where
  step p = D.fromVertices $ sample p


{-# LANGUAGE LambdaCase #-}

module Statistics.Quantile.Approximate.Sampling (
    sampling
  , samplingJackknife
) where

import Control.Applicative
import Control.Monad.IO.Class

import qualified Data.List.Ordered as Bag
import Data.Monoid
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed as V

import Pipes
import qualified Pipes.Prelude as P

import qualified Statistics.Quantile as Q
import Statistics.Resampling
import Statistics.Types
import Statistics.Sample

import Statistics.Quantile.Types
import Statistics.Quantile.Util

import System.Random.MWC

data Jackknife = Jackknifed | NotJackknifed
  deriving (Eq, Show)

sampling :: Int
         -> Int
         -> Selector IO
sampling k n = Selector $ sampling' NotJackknifed k n

samplingJackknife :: Int
                  -> Int
                  -> Selector IO
samplingJackknife k n = Selector $ sampling' Jackknifed k n

-- FIXME: this breaks at 'select' with low probability.
-- FIXME: fold with something stricter than a tuple
sampling' :: Jackknife
          -> Int
          -> Int
          -> Quantile
          -> Stream IO
          -> IO Double
sampling' j k n q (Stream src) = withSystemRandom $ \gen -> do
  P.foldM (sample gen) (pure ([],0)) (finalize j) src
  where choose :: GenIO -> IO Bool
        choose g = choose' <$> uniformR (0.0, invP) g

        choose' :: Double -> Bool
        choose' x
          | x < 1.0   = True
          | otherwise = False

        -- 1/P(choose)
        invP :: Double
        invP = (fromIntegral n) / (fromIntegral k)

        sample :: GenIO -> ([Double], Int) -> Double -> IO ([Double], Int)
        sample g (vs,c) v = choose g >>= \case
          True -> let vs' = Bag.insertBag v vs in
                  pure (vs', c+1)
          False -> pure (vs,c)

        finalize NotJackknifed (vs, c) = pure $ vs !! (qidx q c)
        finalize Jackknifed (vs, _) = pure . estimateQ median $ jackknife (Function $ estimateQ q) $ V.fromList vs

estimateQ :: Quantile -> (Vector Double) -> Double
estimateQ (Quantile k n) = Q.weightedAvg k n


{-# LANGUAGE DataKinds #-}

module Matrix
  ( Matrix4
  , Vector3
  , vec3
  , perspective
  , identity
  , withFloatMatrix
  , translate
  , projection
  , normalize
  , xHat
  , yHat
  , zHat
  , cross
  , vector3ToFloatList
  , unwrap
  ) where

import           Data.Maybe
import           Foreign.C.Types
import           Foreign.Ptr
import qualified Numeric.LinearAlgebra        as L
import qualified Numeric.LinearAlgebra.Devel  as D
import           Numeric.LinearAlgebra.Static

type Matrix4 = L 4 4

type Vector3 = R 3

xHat = vec3 1 0 0

yHat = vec3 0 1 0

zHat = vec3 0 0 1

--type Matrix4 = Matrix4 Double
perspective :: Double -> Double -> Double -> Double -> Matrix4
perspective near far fov aspect = build gen
  where
    gen i j =
      case (i, j) of
        (0, 0) -> 1.0 / (aspect * tan (fov / 2.0))
        (1, 1) -> 1.0 / tan (fov / 2.0)
        (2, 2) -> -(far + near) / (far - near)
        (2, 3) -> -2.0 * far * near / (far - near)
        (3, 2) -> -1.0
        _      -> 0.0

identity :: Matrix4
identity = eye

translate :: Vector3 -> Matrix4
translate v = t ||| col (v & 1.0)
  where
    (t, _) = splitCols eye

projection :: Double -> Matrix4
projection = perspective 0.1 10000.0 (pi / 4.0)

withFloatMatrix ::
     Matrix4 -> (D.MatrixOrder -> CInt -> CInt -> Ptr Float -> IO r) -> IO r
withFloatMatrix m f = D.applyRaw singleM id (f (D.orderOf singleM))
  where
    singleM = L.single $ unwrap m

withFloatVector :: Vector3 -> (CInt -> Ptr Float -> IO r) -> IO r
withFloatVector m = D.applyRaw singleV id
  where
    singleV = L.single $ unwrap m

--normalize :: Vector3 -> Vector3
normalize v = realToFrac a * v
  where
    a = 1.0 / norm_2 v

vector3ToFloatList :: Vector3 -> [Float]
vector3ToFloatList = L.toList . L.single . unwrap


{-# LANGUAGE DataKinds
           , GADTs
           , KindSignatures
           , OverloadedLists
           , ScopedTypeVariables
           , TypeApplications #-}

module MorphGine.Data.Primitives where

import qualified Data.Vector.Storable  as VS
import qualified Numeric.LinearAlgebra as L
import           Numeric.LinearAlgebra ((><))
import qualified Numeric.LinearAlgebra.Data as DL --(flatten, reshape,cmap)
import           GHC.Ptr (Ptr(..))
import           GHC.TypeLits
import qualified Graphics.Rendering.OpenGL.GL as GL
import           Graphics.Rendering.OpenGL.GL (($=))
import qualified Graphics.GL as GLRaw
import           Graphics.GL (GLfloat)

data Matrix :: * -> Nat -> Nat -> * where
  M :: VS.Storable a => L.Matrix a -> Matrix a m n

withPtr :: (VS.Storable a, L.Element a) => Matrix a m n -> (Ptr a -> IO b) -> IO b
withPtr (M x) = VS.unsafeWith . L.flatten $ x

projection :: GLfloat -> GLfloat -> GLfloat -> GLfloat -> Matrix GLfloat 4 4
projection fov aspect near far = M . (4><4) $ [ t/aspect, 0, 0, 0
                                               , 0,        t, 0, 0
                                               , 0,        0, (far+near)/(near-far), (2*far*near)/(near-far)
                                               , 0,        0,-1, 0]
  where t = 1 / tan (fov/2)

degProjection :: GLfloat -> GLfloat -> GLfloat -> GLfloat -> Matrix  GLfloat 4 4
degProjection = projection . (/180) . (*pi)