{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NoImplicitPrelude          #-}
{-# LANGUAGE Rank2Types                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}

-----------------------------------------------------------------------------
-- |
-- A class for semirings (types with two binary operations, one commutative and one associative, and two respective identities), with various general-purpose instances.
--
-----------------------------------------------------------------------------

module Data.Semiring
  ( -- * Semiring typeclass
    Semiring(..)
  , (+)
  , (*)
  , (^)
  , foldMapP
  , foldMapT
  , sum
  , product
  , sum'
  , product'
  , isZero
  , isOne

    -- * Types
  , Add(..)
  , Mul(..)
  , WrappedNum(..)
  , Mod2(..)
#if defined(VERSION_containers)
  , IntSetOf(..)
  , IntMapOf(..)
#endif

    -- * Ring typeclass
  , Ring(..)
  , fromInteger
  , fromIntegral
  , minus
  , (-)
  ) where

import           Control.Applicative (Applicative(..), Const(..), liftA2)
import           Data.Bits (Bits)
import           Data.Bool (Bool(..), (||), (&&), otherwise)
import           Data.Coerce (Coercible, coerce)
import           Data.Complex (Complex(..))
import           Data.Eq (Eq(..))
import           Data.Fixed (Fixed, HasResolution)
import           Data.Foldable (Foldable(foldMap))
import qualified Data.Foldable as Foldable
import           Data.Function ((.), const, id)
#if defined(VERSION_unordered_containers) || defined(VERSION_containers)
import           Data.Function (flip)
#endif
import           Data.Functor (Functor(..))
#if MIN_VERSION_base(4,12,0)
import           Data.Functor.Contravariant (Predicate(..), Equivalence(..), Op(..))
#endif
import           Data.Functor.Identity (Identity(..))
#if defined(VERSION_unordered_containers)
import           Data.Hashable (Hashable)
import           Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import           Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
#endif
import           Data.Int (Int, Int8, Int16, Int32, Int64)
import           Data.Maybe (Maybe(..))
#if MIN_VERSION_base(4,12,0)
import           Data.Monoid (Ap(..))
#endif
#if defined(VERSION_containers)
#if MIN_VERSION_base(4,7,0)
import           Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import           Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
#endif
import           Data.Map (Map)
import qualified Data.Map as Map
#endif
import           Data.Monoid (Monoid(..), Dual(..))
import           Data.Ord (Ord((<)), (>=))
import           Data.Ord (Down(..))
import           Data.Proxy (Proxy(..))
import           Data.Ratio (Ratio, Rational, (%))
import           Data.Semigroup.Compat (Semigroup(..))
#if defined(VERSION_containers)
import           Data.Set (Set)
import qualified Data.Set as Set
#endif
import           Data.Traversable (Traversable)
import           Data.Typeable (Typeable)
import           Data.Word (Word, Word8, Word16, Word32, Word64)
import           Foreign.C.Types
  (CChar, CClock, CDouble, CFloat, CInt,
   CIntMax, CIntPtr, CLLong, CLong,
   CPtrdiff, CSChar, CSUSeconds, CShort,
   CSigAtomic, CSize, CTime, CUChar, CUInt,
   CUIntMax, CUIntPtr, CULLong, CULong,
   CUSeconds, CUShort, CWchar)
import           Foreign.Ptr (IntPtr, WordPtr)
import           Foreign.Storable (Storable)
import           GHC.Enum (Enum, Bounded)
import           GHC.Err (error)
import           GHC.Float (Float, Double)
import           GHC.Generics (Generic,Generic1)
import           GHC.IO (IO)
import           GHC.Integer (Integer)
import qualified GHC.Num as Num
import           GHC.Read (Read)
import           GHC.Real (Integral, Fractional, Real, RealFrac)
import qualified GHC.Real as Real
import           GHC.Show (Show)
import           Numeric.Natural (Natural)

#ifdef mingw32_HOST_OS
#define HOST_OS_WINDOWS 1
#else
#define HOST_OS_WINDOWS 0
#endif

#if !HOST_OS_WINDOWS
import           System.Posix.Types
  (CCc, CDev, CGid, CIno, CMode, CNlink,
   COff, CPid, CRLim, CSpeed, CSsize,
   CTcflag, CUid, Fd)
#endif

infixl 7 *, `times`
infixl 6 +, `plus`, -, `minus`
infixr 8 ^

{--------------------------------------------------------------------
  Helpers
--------------------------------------------------------------------}

-- | Raise a number to a non-negative integral power.
-- If the power is negative, this will call 'error'.
{-# SPECIALISE [1] (^) ::
        Integer -> Integer -> Integer,
        Integer -> Int -> Integer,
        Int -> Int -> Int #-}
{-# INLINABLE [1] (^) #-} -- See note [Inlining (^)]
(^) :: (Semiring a, Integral b) => a -> b -> a
x ^ y
  | y < 0 = error "Data.Semiring.^: negative power"
  | y == 0 = one
  | otherwise = getMul (stimes y (Mul x))

{- Note [Inlining (^)]
   ~~~~~~~~~~~~~~~~~~~
   The INLINABLE pragma allows (^) to be specialised at its call sites.
   If it is called repeatedly at the same type, that can make a huge
   difference, because of those constants which can be repeatedly
   calculated.

   Currently the fromInteger calls are not floated because we get
             \d1 d2 x y -> blah
   after the gentle round of simplification.
-}

{- Rules for powers with known small exponent
    see Trac #5237
    For small exponents, (^) is inefficient compared to manually
    expanding the multiplication tree.
    Here, rules for the most common exponent types are given.
    The range of exponents for which rules are given is quite
    arbitrary and kept small to not unduly increase the number of rules.
    It might be desirable to have corresponding rules also for
    exponents of other types (e.g., Word), but it's doubtful they
    would fire, since the exponents of other types tend to get
    floated out before the rule has a chance to fire. (Why?)

    Note: Trying to save multiplication by sharing the square for
    exponents 4 and 5 does not save time, indeed, for Double, it is
    up to twice slower, so the rules contain flat sequences of
    multiplications.
-}

{-# RULES
"^0/Int" forall x. x ^ (0 :: Int) = one
"^1/Int" forall x. x ^ (1 :: Int) = let u = x in u
"^2/Int" forall x. x ^ (2 :: Int) = let u = x in u*u
"^3/Int" forall x. x ^ (3 :: Int) = let u = x in u*u*u
"^4/Int" forall x. x ^ (4 :: Int) = let u = x in u*u*u*u
"^5/Int" forall x. x ^ (5 :: Int) = let u = x in u*u*u*u*u
"^0/Integer" forall x. x ^ (0 :: Integer) = one
"^1/Integer" forall x. x ^ (1 :: Integer) = let u = x in u
"^2/Integer" forall x. x ^ (2 :: Integer) = let u = x in u*u
"^3/Integer" forall x. x ^ (3 :: Integer) = let u = x in u*u*u
"^4/Integer" forall x. x ^ (4 :: Integer) = let u = x in u*u*u*u
"^5/Integer" forall x. x ^ (5 :: Integer) = let u = x in u*u*u*u*u
  #-}

-- | Infix shorthand for 'plus'.
(+) :: Semiring a => a -> a -> a
(+) = plus
{-# INLINE (+) #-}

-- | Infix shorthand for 'times'.
(*) :: Semiring a => a -> a -> a
(*) = times
{-# INLINE (*) #-}

-- | Infix shorthand for 'minus'.
(-) :: Ring a => a -> a -> a
(-) = minus
{-# INLINE (-) #-}

-- | Map each element of the structure to a semiring, and combine the results
--   using 'plus'.
foldMapP :: (Foldable t, Semiring s) => (a -> s) -> t a -> s
foldMapP f = Foldable.foldr (plus  . f) zero
{-# INLINE foldMapP #-}

-- | Map each element of the structure to a semiring, and combine the results
--   using 'times'.
foldMapT :: (Foldable t, Semiring s) => (a -> s) -> t a -> s
foldMapT f = Foldable.foldr (times . f) one
{-# INLINE foldMapT #-}

infixr 9 #.
(#.) :: Coercible b c => (b -> c) -> (a -> b) -> a -> c
(#.) _ = coerce

-- | The 'sum' function computes the additive sum of the elements in a structure.
--   This function is lazy. For a strict version, see 'sum''.
sum :: (Foldable t, Semiring a) => t a -> a
sum = getAdd #. foldMap Add
{-# INLINE sum #-}

-- | The 'product' function computes the product of the elements in a structure.
--   This function is lazy. for a strict version, see 'product''.
product :: (Foldable t, Semiring a) => t a -> a
product = getMul #. foldMap Mul
{-# INLINE product #-}

-- | The 'sum'' function computes the additive sum of the elements in a structure.
--   This function is strict. For a lazy version, see 'sum'.
sum'  :: (Foldable t, Semiring a) => t a -> a
sum'  = Foldable.foldl' plus zero
{-# INLINE sum' #-}

-- | The 'product'' function computes the additive sum of the elements in a structure.
--   This function is strict. For a lazy version, see 'product'.
product' :: (Foldable t, Semiring a) => t a -> a
product' = Foldable.foldl' times one
{-# INLINE product' #-}

-- | Monoid under 'plus'. Analogous to 'Data.Monoid.Sum', but
--   uses the 'Semiring' constraint rather than 'Num.Num'.
newtype Add a = Add { getAdd :: a }
  deriving
    ( Bounded
    , Enum
    , Eq
    , Foldable
    , Fractional
    , Functor
    , Generic
    , Generic1
    , Num.Num
    , Ord
    , Read
    , Real
    , RealFrac
    , Show
    , Storable
    , Traversable
    , Typeable
    )

instance Semiring a => Semigroup (Add a) where
  Add a <> Add b = Add (a + b)
  stimes n (Add a) = Add (fromNatural (Real.fromIntegral n) * a)
  {-# INLINE (<>) #-}

instance Semiring a => Monoid (Add a) where
  mempty = Add zero
  mappend = (<>)
  {-# INLINE mempty #-}
  {-# INLINE mappend #-}

-- | This is an internal type, solely for purposes
-- of default implementation of 'fromNatural'.
newtype Add' a = Add' { getAdd' :: a }

instance Semiring a => Semigroup (Add' a) where
  Add' a <> Add' b = Add' (a + b)

-- | Monoid under 'times'. Analogous to 'Data.Monoid.Product', but
--   uses the 'Semiring' constraint rather than 'Num.Num'.
newtype Mul a = Mul { getMul :: a }
  deriving
    ( Bounded
    , Enum
    , Eq
    , Foldable
    , Fractional
    , Functor
    , Generic
    , Generic1
    , Num.Num
    , Ord
    , Read
    , Real
    , RealFrac
    , Show
    , Storable
    , Traversable
    , Typeable
    )

instance Semiring a => Semigroup (Mul a) where
  Mul a <> Mul b = Mul (a * b)
  {-# INLINE (<>) #-}

instance Semiring a => Monoid (Mul a) where
  mempty = Mul one
  mappend = (<>)
  {-# INLINE mempty #-}
  {-# INLINE mappend #-}

-- | Provide Semiring and Ring for an arbitrary 'Num.Num'. It is useful with GHC 8.6+'s DerivingVia extension.
newtype WrappedNum a = WrapNum { unwrapNum :: a }
  deriving
    ( Bounded
    , Enum
    , Eq
    , Foldable
    , Fractional
    , Functor
    , Generic
    , Generic1
    , Num.Num
    , Ord
    , Read
    , Real
    , RealFrac
    , Show
    , Storable
    , Traversable
    , Typeable
    , Bits
    )

instance Num.Num a => Semiring (WrappedNum a) where
  plus  = (Num.+)
  zero  = 0
  times = (Num.*)
  one   = 1
  fromNatural = Real.fromIntegral

instance Num.Num a => Ring (WrappedNum a) where
  negate = Num.negate

-- | 'Mod2' represents the integers mod 2.
--
--   It is useful in the computing of <https://en.wikipedia.org/wiki/Zhegalkin_polynomial Zhegalkin polynomials>.
newtype Mod2 = Mod2 { getMod2 :: Bool }
  deriving
    ( Bounded
    , Enum
    , Eq
    , Ord
    , Read
    , Show
    , Generic
    )

instance Semiring Mod2 where
  -- we inline the definition of 'xor'
  -- on Bools, since the instance did not exist until
  -- base-4.7.0.
  plus (Mod2 x) (Mod2 y) = Mod2 (x /= y)
  times (Mod2 x) (Mod2 y) = Mod2 (x && y)
  zero = Mod2 False
  one = Mod2 True

instance Ring Mod2 where
  negate = id
  {-# INLINE negate #-}

{--------------------------------------------------------------------
  Classes
--------------------------------------------------------------------}

-- | The class of semirings (types with two binary
-- operations and two respective identities). One
-- can think of a semiring as two monoids of the same
-- underlying type, with the first being commutative.
-- In the documentation, you will often see the first
-- monoid being referred to as @additive@, and the second
-- monoid being referred to as @multiplicative@, a typical
-- convention when talking about semirings.
--
-- For any type R with a 'Num.Num'
-- instance, the additive monoid is (R, 'Prelude.+', 0)
-- and the multiplicative monoid is (R, 'Prelude.*', 1).
--
-- For 'Prelude.Bool', the additive monoid is ('Prelude.Bool', 'Prelude.||', 'Prelude.False')
-- and the multiplicative monoid is ('Prelude.Bool', 'Prelude.&&', 'Prelude.True').
--
-- Instances should satisfy the following laws:
--
-- [/additive left identity/]
--     @'zero' '+' x = x@
-- [/additive right identity/]
--     @x '+' 'zero' = x@
-- [/additive associativity/]
--     @x '+' (y '+' z) = (x '+' y) '+' z@
-- [/additive commutativity/]
--     @x '+' y = y '+' x@
-- [/multiplicative left identity/]
--     @'one' '*' x = x@
-- [/multiplicative right identity/]
--     @x '*' 'one' = x@
-- [/multiplicative associativity/]
--     @x '*' (y '*' z) = (x '*' y) '*' z@
-- [/left-distributivity of '*' over '+'/]
--     @x '*' (y '+' z) = (x '*' y) '+' (x '*' z)@
-- [/right-distributivity of '*' over '+'/]
--     @(x '+' y) '*' z = (x '*' z) '+' (y '*' z)@
-- [/annihilation/]
--     @'zero' '*' x = x '*' 'zero' = 'zero'@

class Semiring a where
#if __GLASGOW_HASKELL__ >= 708
  {-# MINIMAL plus, times, (zero, one | fromNatural) #-}
#endif
  plus  :: a -> a -> a -- ^ Commutative Operation
  zero  :: a           -- ^ Commutative Unit
  zero = fromNatural 0
  times :: a -> a -> a -- ^ Associative Operation
  one   :: a           -- ^ Associative Unit
  one = fromNatural 1
  fromNatural :: Natural -> a -- ^ Homomorphism of additive semigroups
  fromNatural 0 = zero
  fromNatural n = getAdd' (stimes n (Add' one))

-- | The class of semirings with an additive inverse.
--
--     @'negate' a '+' a = 'zero'@

class Semiring a => Ring a where
#if __GLASGOW_HASKELL__ >= 708
  {-# MINIMAL negate #-}
#endif
  negate :: a -> a

-- | Subtract two 'Ring' values. For any type @R@ with
-- a 'Num.Num' instance, this is the same as '(Prelude.-)'.
--
--     @x `minus` y = x '+' 'negate' y@
minus :: Ring a => a -> a -> a
minus x y = x + negate y
{-# INLINE minus #-}

-- | Convert from integer to ring.
--
-- When @{-#@ @LANGUAGE RebindableSyntax #-}@ is enabled,
-- this function is used for desugaring integer literals.
-- This may be used to facilitate transition from 'Num.Num' to 'Ring':
-- no need to replace 0 and 1 with 'one' and 'zero'
-- or to cast numeric literals.
fromInteger :: Ring a => Integer -> a
fromInteger x
  | x >= 0    = fromNatural (Num.fromInteger x)
  | otherwise = negate (fromNatural (Num.fromInteger (Num.negate x)))
{-# INLINE fromInteger #-}

-- | Convert from integral to ring.
fromIntegral :: (Integral a, Ring b) => a -> b
fromIntegral x
  | x >= 0    = fromNatural (Real.fromIntegral x)
  | otherwise = negate (fromNatural (Real.fromIntegral (Num.negate x)))
{-# INLINE fromIntegral #-}

{--------------------------------------------------------------------
  Instances (base)
--------------------------------------------------------------------}

instance Semiring b => Semiring (a -> b) where
  plus f g    = \x -> f x `plus` g x
  zero        = const zero
  times f g   = \x -> f x `times` g x
  one         = const one
  fromNatural = const . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring b => Ring (a -> b) where
  negate f x = negate (f x)
  {-# INLINE negate #-}

instance Semiring () where
  plus _ _  = ()
  zero      = ()
  times _ _ = ()
  one       = ()
  fromNatural _ = ()
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring () where
  negate _ = ()
  {-# INLINE negate #-}

instance Semiring (Proxy a) where
  plus _ _  = Proxy
  zero      = Proxy
  times _ _ = Proxy
  one       = Proxy
  fromNatural _ = Proxy
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Semiring Bool where
  plus  = (||)
  zero  = False
  times = (&&)
  one   = True
  fromNatural 0 = False
  fromNatural _ = True
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Semiring a => Semiring (Maybe a) where
  zero  = Nothing
  one   = Just one

  plus Nothing y = y
  plus x Nothing = x
  plus (Just x) (Just y) = Just (plus x y)

  times Nothing _ = Nothing
  times _ Nothing = Nothing
  times (Just x) (Just y) = Just (times x y)

  fromNatural 0 = Nothing
  fromNatural n = Just (fromNatural n)
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Semiring a => Semiring (IO a) where
  zero  = pure zero
  one   = pure one
  plus  = liftA2 plus
  times = liftA2 times
  fromNatural = pure . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring a => Ring (IO a) where
  negate = fmap negate
  {-# INLINE negate #-}

instance Semiring a => Semiring (Dual a) where
  zero = Dual zero
  Dual x `plus` Dual y = Dual (y `plus` x)
  one = Dual one
  Dual x `times` Dual y = Dual (y `times` x)
  fromNatural = Dual . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring a => Ring (Dual a) where
  negate (Dual x) = Dual (negate x)
  {-# INLINE negate #-}

instance Semiring a => Semiring (Const a b) where
  zero = Const zero
  one  = Const one
  plus  (Const x) (Const y) = Const (x `plus`  y)
  times (Const x) (Const y) = Const (x `times` y)
  fromNatural = Const . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring a => Ring (Const a b) where
  negate (Const x) = Const (negate x)
  {-# INLINE negate #-}

-- | This instance can suffer due to floating point arithmetic.
instance Ring a => Semiring (Complex a) where
  zero = zero :+ zero
  one  = one  :+ zero
  plus  (x :+ y) (x' :+ y') = plus x x' :+ plus y y'
  times (x :+ y) (x' :+ y')
    = (x * x' - (y * y')) :+ (x * y' + y * x')
  fromNatural n = fromNatural n :+ zero
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring a => Ring (Complex a) where
  negate (x :+ y) = negate x :+ negate y
  {-# INLINE negate #-}

#if MIN_VERSION_base(4,12,0)
instance (Semiring a, Applicative f) => Semiring (Ap f a) where
  zero  = pure zero
  one   = pure one
  plus  = liftA2 plus
  times = liftA2 times
  fromNatural = pure . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance (Ring a, Applicative f) => Ring (Ap f a) where
  negate = fmap negate
  {-# INLINE negate #-}
#endif

#if MIN_VERSION_base(4,12,0)
deriving instance Semiring (Predicate a)

deriving instance Semiring a => Semiring (Equivalence a)

deriving instance Semiring a => Semiring (Op a b)
deriving instance Ring a => Ring (Op a b)
#endif

#define deriveSemiring(ty)         \
instance Semiring (ty) where {     \
   zero  = 0                       \
;  one   = 1                       \
;  plus  x y = (Num.+) x y         \
;  times x y = (Num.*) x y         \
;  fromNatural = Real.fromIntegral \
;  {-# INLINE zero #-}             \
;  {-# INLINE one  #-}             \
;  {-# INLINE plus #-}             \
;  {-# INLINE times #-}            \
;  {-# INLINE fromNatural #-}      \
}

deriveSemiring(Int)
deriveSemiring(Int8)
deriveSemiring(Int16)
deriveSemiring(Int32)
deriveSemiring(Int64)
deriveSemiring(Integer)
deriveSemiring(Word)
deriveSemiring(Word8)
deriveSemiring(Word16)
deriveSemiring(Word32)
deriveSemiring(Word64)
deriveSemiring(Float)
deriveSemiring(Double)
deriveSemiring(CUIntMax)
deriveSemiring(CIntMax)
deriveSemiring(CUIntPtr)
deriveSemiring(CIntPtr)
deriveSemiring(CSUSeconds)
deriveSemiring(CUSeconds)
deriveSemiring(CTime)
deriveSemiring(CClock)
deriveSemiring(CSigAtomic)
deriveSemiring(CWchar)
deriveSemiring(CSize)
deriveSemiring(CPtrdiff)
deriveSemiring(CDouble)
deriveSemiring(CFloat)
deriveSemiring(CULLong)
deriveSemiring(CLLong)
deriveSemiring(CULong)
deriveSemiring(CLong)
deriveSemiring(CUInt)
deriveSemiring(CInt)
deriveSemiring(CUShort)
deriveSemiring(CShort)
deriveSemiring(CUChar)
deriveSemiring(CSChar)
deriveSemiring(CChar)
deriveSemiring(IntPtr)
deriveSemiring(WordPtr)

#if !HOST_OS_WINDOWS
deriveSemiring(CCc)
deriveSemiring(CDev)
deriveSemiring(CGid)
deriveSemiring(CIno)
deriveSemiring(CMode)
deriveSemiring(CNlink)
deriveSemiring(COff)
deriveSemiring(CPid)
deriveSemiring(CRLim)
deriveSemiring(CSpeed)
deriveSemiring(CSsize)
deriveSemiring(CTcflag)
deriveSemiring(CUid)
deriveSemiring(Fd)
#endif

deriveSemiring(Natural)

instance Integral a => Semiring (Ratio a) where
  {-# SPECIALIZE instance Semiring Rational #-}
  zero  = 0 % 1
  one   = 1 % 1
  plus  = (Num.+)
  times = (Num.*)
  fromNatural n = Real.fromIntegral n % 1
  {-# INLINE zero  #-}
  {-# INLINE one   #-}
  {-# INLINE plus  #-}
  {-# INLINE times #-}
  {-# INLINE fromNatural #-}
deriving instance Semiring a => Semiring (Identity a)
deriving instance Semiring a => Semiring (Down a)
instance HasResolution a => Semiring (Fixed a) where
  zero  = 0
  one   = 1
  plus  = (Num.+)
  times = (Num.*)
  fromNatural = Real.fromIntegral
  {-# INLINE zero  #-}
  {-# INLINE one   #-}
  {-# INLINE plus  #-}
  {-# INLINE times #-}
  {-# INLINE fromNatural #-}

#define deriveRing(ty)          \
instance Ring (ty) where {      \
  negate = Num.negate           \
; {-# INLINE negate #-}         \
}

deriveRing(Int)
deriveRing(Int8)
deriveRing(Int16)
deriveRing(Int32)
deriveRing(Int64)
deriveRing(Integer)
deriveRing(Word)
deriveRing(Word8)
deriveRing(Word16)
deriveRing(Word32)
deriveRing(Word64)
deriveRing(Float)
deriveRing(Double)
deriveRing(CUIntMax)
deriveRing(CIntMax)
deriveRing(CUIntPtr)
deriveRing(CIntPtr)
deriveRing(CSUSeconds)
deriveRing(CUSeconds)
deriveRing(CTime)
deriveRing(CClock)
deriveRing(CSigAtomic)
deriveRing(CWchar)
deriveRing(CSize)
deriveRing(CPtrdiff)
deriveRing(CDouble)
deriveRing(CFloat)
deriveRing(CULLong)
deriveRing(CLLong)
deriveRing(CULong)
deriveRing(CLong)
deriveRing(CUInt)
deriveRing(CInt)
deriveRing(CUShort)
deriveRing(CShort)
deriveRing(CUChar)
deriveRing(CSChar)
deriveRing(CChar)
deriveRing(IntPtr)
deriveRing(WordPtr)

#if !HOST_OS_WINDOWS
deriveRing(CCc)
deriveRing(CDev)
deriveRing(CGid)
deriveRing(CIno)
deriveRing(CMode)
deriveRing(CNlink)
deriveRing(COff)
deriveRing(CPid)
deriveRing(CRLim)
deriveRing(CSpeed)
deriveRing(CSsize)
deriveRing(CTcflag)
deriveRing(CUid)
deriveRing(Fd)
#endif

instance Integral a => Ring (Ratio a) where
  negate = Num.negate
  {-# INLINE negate #-}

deriving instance Ring a => Ring (Down a)
deriving instance Ring a => Ring (Identity a)
instance HasResolution a => Ring (Fixed a) where
  negate = Num.negate
  {-# INLINE negate #-}

{--------------------------------------------------------------------
  Instances (containers)
--------------------------------------------------------------------}

#if defined(VERSION_containers)

-- | The multiplication laws are satisfied for
--   any underlying 'Monoid', so we require a
--   'Monoid' constraint instead of a 'Semiring'
--   constraint since 'times' can use
--   the context of either.
instance (Ord a, Monoid a) => Semiring (Set a) where
  zero  = Set.empty
  one   = Set.singleton mempty
  plus  = Set.union
  times xs ys = Foldable.foldMap (flip Set.map ys . mappend) xs
  fromNatural 0 = zero
  fromNatural _ = one
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

-- | Wrapper to mimic 'Set' ('Data.Semigroup.Sum' 'Int'),
-- 'Set' ('Data.Semigroup.Product' 'Int'), etc.,
-- while having a more efficient underlying representation.
newtype IntSetOf a = IntSetOf { getIntSet :: IntSet }
  deriving
    ( Eq
    , Generic
    , Generic1
    , Ord
    , Read
    , Show
    , Typeable
    , Semigroup
    , Monoid
    )

instance (Coercible Int a, Monoid a) => Semiring (IntSetOf a) where
  zero  = coerce IntSet.empty
  one   = coerce IntSet.singleton (mempty :: a)
  plus  = coerce IntSet.union
  xs `times` ys
    = coerce IntSet.fromList
        [ mappend k l
        | k :: a <- coerce IntSet.toList xs
        , l :: a <- coerce IntSet.toList ys
        ]
  fromNatural 0 = zero
  fromNatural _ = one
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

-- | The multiplication laws are satisfied for
--   any underlying 'Monoid' as the key type,
--   so we require a 'Monoid' constraint instead of
--   a 'Semiring' constraint since 'times' can use
--   the context of either.
instance (Ord k, Monoid k, Semiring v) => Semiring (Map k v) where
  zero = Map.empty
  one  = Map.singleton mempty one
  plus = Map.unionWith (+)
  xs `times` ys
    = Map.fromListWith (+)
        [ (mappend k l, v * u)
        | (k,v) <- Map.toList xs
        , (l,u) <- Map.toList ys
        ]
  fromNatural 0 = zero
  fromNatural n = Map.singleton mempty (fromNatural n)
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

-- | Wrapper to mimic 'Map' ('Data.Semigroup.Sum' 'Int') v,
-- 'Map' ('Data.Semigroup.Product' 'Int') v, etc.,
-- while having a more efficient underlying representation.
newtype IntMapOf k v = IntMapOf { getIntMap :: IntMap v }
  deriving
    ( Eq
    , Generic
    , Generic1
    , Ord
    , Read
    , Show
    , Typeable
    , Semigroup
    , Monoid
    )

instance (Coercible Int k, Monoid k, Semiring v) => Semiring (IntMapOf k v) where
  zero = coerce (IntMap.empty :: IntMap v)
  one  = coerce (IntMap.singleton :: Int -> v -> IntMap v) (mempty :: k) (one :: v)
  plus = coerce (IntMap.unionWith (+) :: IntMap v -> IntMap v -> IntMap v)
  xs `times` ys
    = coerce (IntMap.fromListWith (+) :: [(Int, v)] -> IntMap v)
        [ (mappend k l, v * u)
        | (k :: k, v :: v) <- coerce (IntMap.toList :: IntMap v -> [(Int, v)]) xs
        , (l :: k, u :: v) <- coerce (IntMap.toList :: IntMap v -> [(Int, v)]) ys
        ]
  fromNatural 0 = zero
  fromNatural n = coerce (IntMap.singleton :: Int -> v -> IntMap v) (mempty :: k) (fromNatural n :: v)
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

#endif

{--------------------------------------------------------------------
  Instances (unordered-containers)
--------------------------------------------------------------------}

#if defined(VERSION_unordered_containers)

-- | The multiplication laws are satisfied for
--   any underlying 'Monoid', so we require a
--   'Monoid' constraint instead of a 'Semiring'
--   constraint since 'times' can use
--   the context of either.
instance (Eq a, Hashable a, Monoid a) => Semiring (HashSet a) where
  zero = HashSet.empty
  one  = HashSet.singleton mempty
  plus = HashSet.union
  times xs ys = Foldable.foldMap (flip HashSet.map ys . mappend) xs
  fromNatural 0 = zero
  fromNatural _ = one
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

-- | The multiplication laws are satisfied for
--   any underlying 'Monoid' as the key type,
--   so we require a 'Monoid' constraint instead of
--   a 'Semiring' constraint since 'times' can use
--   the context of either.
instance (Eq k, Hashable k, Monoid k, Semiring v) => Semiring (HashMap k v) where
  zero = HashMap.empty
  one  = HashMap.singleton mempty one
  plus = HashMap.unionWith (+)
  xs `times` ys
    = HashMap.fromListWith (+)
        [ (mappend k l, v * u)
        | (k,v) <- HashMap.toList xs
        , (l,u) <- HashMap.toList ys
        ]
  fromNatural 0 = zero
  fromNatural n = HashMap.singleton mempty (fromNatural n)
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}
#endif

-- | Is the value 'zero'?
isZero :: (Eq a, Semiring a) => a -> Bool
isZero x = x == zero
{-# INLINEABLE isZero #-}

-- | Is the value 'one'?
isOne :: (Eq a, Semiring a) => a -> Bool
isOne x = x == one
{-# INLINEABLE isOne #-}


module Main where

import Data.Ord (comparing)
import qualified Data.Vector.Algorithms.Intro as Intro
import qualified Data.Vector.Storable as VS
import Data.Vector.Storable.Mutable (IOVector)
import Numeric.LinearAlgebra hiding (eig, vector)
import Numeric.LinearAlgebra.Arnoldi
import Test.Hspec
import Test.QuickCheck

main :: IO ()
main = hspec $ do
  describe "Numeric.LinearAlgebra.Arnoldi.eig" $ do

    it "returns the diagonal of an arbitrary diagonal matrix" $ property $ do
      dim <- arbitrary `suchThat` (> 3)
      diagonal <- VS.fromList <$> vector dim
      let
        _matrix = diag diagonal :: Matrix Double
        nev = dim - 3
        options = Options { which = SR
                          , number = nev
                          , maxIterations = Nothing
                          }
        sort = VS.modify Intro.sort
        actual = (sort . fst) (eig options dim (multiply _matrix))
        expected = (VS.take nev . sort) diagonal
        relative = VS.maximum (VS.zipWith (%) expected actual)
      pure (counterexample (show (expected, actual)) (relative < 1E-4))

    it "returns the diagonal of an arbitrary diagonal matrix" $ property $ do
      dim <- arbitrary `suchThat` (> 3)
      diagonal <- VS.fromList <$> vector dim
      let
        _matrix = diag diagonal :: Matrix (Complex Double)
        nev = dim - 3
        options = Options { which = SR
                          , number = nev
                          , maxIterations = Nothing
                          }
        sort = VS.modify (Intro.sortBy (comparing realPart))
        actual = (sort . fst) (eig options dim (multiply _matrix))
        expected = (VS.take nev . sort) diagonal
        relative = VS.maximum (VS.zipWith (%%) expected actual)
      pure (counterexample (show (expected, actual)) (relative < 1E-4))

multiply :: Numeric a => Matrix a -> IOVector a -> IOVector a -> IO ()
multiply _matrix dst src = do
  x <- VS.freeze src
  let y = _matrix #> x
  VS.copy dst y

(%) :: Double -> Double -> Double
(%) a b =
  let a_plus_b = a + b
  in if a_plus_b /= 0
     then abs ((a - b) / a_plus_b)
     else a - b

(%%) :: Complex Double -> Complex Double -> Double
(%%) a b =
  let a_plus_b = a + b
  in if a_plus_b /= 0
     then magnitude ((a - b) / a_plus_b)
     else magnitude (a - b)


{-# LANGUAGE FlexibleContexts, FlexibleInstances, TypeFamilies #-}

module Tests.ApproxEq
    (
      ApproxEq(..)
    ) where

import Data.Complex (Complex(..), realPart)
import Data.List (intersperse)
import Data.Maybe (catMaybes)
import Numeric.MathFunctions.Constants (m_epsilon)
import Statistics.Matrix hiding (map, toList)
import Test.QuickCheck
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Statistics.Matrix as M

class (Eq a, Show a) => ApproxEq a where
    type Bounds a

    eq   :: Bounds a -> a -> a -> Bool
    eql  :: Bounds a -> a -> a -> Property
    eql eps a b = counterexample (show a ++ " /=~ " ++ show b) (eq eps a b)

    (=~)  :: a -> a -> Bool

    (==~) :: a -> a -> Property
    a ==~ b = counterexample (show a ++ " /=~ " ++ show b) (a =~ b)

instance ApproxEq Double where
    type Bounds Double = Double

    eq eps a b
      | a == 0 && b == 0 = True
      | otherwise        = abs (a - b) <= eps * max (abs a) (abs b)
    (=~)  = eq m_epsilon

instance ApproxEq (Complex Double) where
    type Bounds (Complex Double) = Double

    eq eps a@(ar :+ ai) b@(br :+ bi)
      | a == 0 && b == 0 = True
      | otherwise        = abs (ar - br) <= eps * d
                        && abs (ai - bi) <= eps * d
      where
        d = max (realPart $ abs a) (realPart $ abs b)

    (=~)  = eq m_epsilon

instance ApproxEq [Double] where
    type Bounds [Double] = Double

    eq  eps (x:xs) (y:ys) = eq eps x y && eq eps xs ys
    eq  _   []     []     = True
    eq  _   _      _      = False

    eql   = eqll length id id
    (=~)  = eq m_epsilon
    (==~) = eql m_epsilon

instance ApproxEq (U.Vector Double) where
    type Bounds (U.Vector Double) = Double

    eq    = eqv
    (=~)  = eq m_epsilon
    eql   = eqlv
    (==~) = eqlv m_epsilon

instance ApproxEq (V.Vector Double) where
    type Bounds (V.Vector Double) = Double

    eq    = eqv
    (=~)  = eq m_epsilon
    eql   = eqlv
    (==~) = eqlv m_epsilon

instance ApproxEq Matrix where
    type Bounds Matrix = Double

    eq eps (Matrix r1 c1 e1 v1) (Matrix r2 c2 e2 v2) =
      (r1,c1,e1) == (r2,c2,e2) && eq eps v1 v2
    (=~)  = eq m_epsilon
    eql eps a b = eqll dimension M.toList (`quotRem` cols a) eps a b
    (==~) = eql m_epsilon

eqv :: (ApproxEq a, G.Vector v Bool, G.Vector v a) =>
       Bounds a -> v a -> v a -> Bool
eqv eps a b = G.length a == G.length b && G.and (G.zipWith (eq eps) a b)

eqlv :: (ApproxEq [a], G.Vector v a) => Bounds [a] -> v a -> v a -> Property
eqlv eps a b = eql eps (G.toList a) (G.toList b)

eqll :: (ApproxEq l, ApproxEq a, Show c, Show d, Eq d, Bounds l ~ Bounds a) =>
        (l -> d) -> (l -> [a]) -> (Int -> c) -> Bounds l -> l -> l -> Property
eqll dim toList coord eps a b = counterexample fancy $ eq eps a b
  where
    fancy
      | la /= lb  = "size mismatch: " ++ show la ++ " /= " ++ show lb
      | length summary < length full = summary
      | otherwise = full
    summary = concat . intersperse ", " . catMaybes $
              zipWith3 whee (map coord [(0::Int)..]) xs ys
    full | '\n' `elem` sa = sa ++ "  /=~\n" ++ sb
         | otherwise      = sa ++ " /=~" ++ sb
    (sa, sb) = (show a, show b)
    (xs, ys) = (toList a, toList b)
    (la, lb) = (dim a, dim b)
    whee i x y | eq eps x y = Nothing
               | otherwise  = Just $ show i ++ ": " ++ show x ++ " /=~ " ++ show y


module Brain
(
  NetworkWithSamples (..)
, constructNeuralNetwork
, trainN
, activate
)

where

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

import AI.HNN.FF.Network          ( Network
                                  , Samples
                                  , createNetwork
                                  , trainNTimes
                                  , tanh
                                  , tanh'
                                  , output )
import Numeric.LinearAlgebra.Data ( Vector )
import Debug.Trace

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

type Network' = Network Double
type Samples' = Samples Double

data NetworkWithSamples = NetworkWithSamples { network :: Network'
                                             , samples :: Samples' }
                          deriving (Show)

constructNeuralNetwork :: Int -> [Int] -> Int -> Samples' -> IO (NetworkWithSamples)
constructNeuralNetwork i h o s = do
    net <- createNetwork i h o
    return NetworkWithSamples { network = net, samples = s }

trainN :: NetworkWithSamples -> Int -> NetworkWithSamples
trainN nn n = nn { network = nn' }
  where nn' = trainNTimes n 0.05 tanh tanh' (network nn) (samples nn)

activate :: NetworkWithSamples -> Vector Double -> Vector Double
activate n s = output (network n) tanh s


{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
-- | Minimil linear algebra lib.
module Math.Regression.Simple.LinAlg (
    -- * Operations
    Add (..),
    Eye (..),
    Mult (..),
    Det (..),
    Inv (..),
    -- * Zeros
    zerosLin,
    zerosQuad,
    optimaQuad,
    -- * Two dimensions
    V2 (..),
    M22 (..),
    SM22 (..),
    -- * Three dimensions
    V3 (..),
    M33 (..),
    SM33 (..),
) where

import Control.DeepSeq (NFData (..))
import Data.Complex    (Complex (..))

-------------------------------------------------------------------------------
-- Classes
-------------------------------------------------------------------------------

-- | Addition
class Add a where
    zero :: a
    add  :: a -> a -> a

-- | Identity
class Eye a where
    eye :: a

-- | Multiplication of different things.
class Eye a => Mult a b c | a b -> c where
    mult :: a -> b -> c

-- | Determinant
class Eye a => Det a where
    det :: a -> Double

-- | Inverse
class Det a => Inv a where
    inv :: a -> a

infixl 6 `add`
infixl 7 `mult`

instance Eye Double where
    eye = 1

instance Add Double where
    zero = 0
    add = (+)

instance Mult Double Double Double where
    mult = (*)

instance Det Double where
    det = id

instance Inv Double where
    inv = recip

-------------------------------------------------------------------------------
-- Zeros
-------------------------------------------------------------------------------

-- | Solve linear equation.
--
-- >>> zerosLin (V2 1 2)
-- -2.0
--
zerosLin :: V2 -> Double
zerosLin (V2 a b) = negate (b / a)

-- | Solve quadratic equation.
--
-- >>> zerosQuad (V3 2 0 (-1))
-- Right (-0.7071067811865476,0.7071067811865476)
--
-- >>> zerosQuad (V3 2 0 1)
-- Left ((-0.0) :+ (-0.7071067811865476),(-0.0) :+ 0.7071067811865476)
--
-- Double root is not treated separately:
--
-- >>> zerosQuad (V3 1 0 0)
-- Right (-0.0,0.0)
--
-- >>> zerosQuad (V3 1 (-2) 1)
-- Right (1.0,1.0)
--
zerosQuad :: V3 -> Either (Complex Double, Complex Double) (Double, Double)
zerosQuad (V3 a b c)
    | delta < 0 = Left ((-b/da) :+ (-sqrtNDelta/da), (-b/da) :+ (sqrtNDelta/da))
    | otherwise = Right ((- b - sqrtDelta) / da, (-b + sqrtDelta) / da)
  where
    delta = b*b - 4 * a * c
    sqrtDelta = sqrt delta
    sqrtNDelta = sqrt (- delta)
    da = 2 * a

-- | Find an optima point.
--
-- >>> optimaQuad (V3 1 (-2) 0)
-- 1.0
--
-- compare to
--
-- >>> zerosQuad (V3 1 (-2) 0)
-- Right (0.0,2.0)
--
optimaQuad :: V3 -> Double
optimaQuad (V3 a b _) = zerosLin (V2 (2 * a) b)

-------------------------------------------------------------------------------
-- 2 dimensions
-------------------------------------------------------------------------------

-- | 2d vector. Strict pair of 'Double's.
--
-- Also used to represent linear polynomial: @V2 a b@  \(= a x + b\).
--
data V2 = V2 !Double !Double
  deriving (Eq, Show)

instance NFData V2 where
    rnf V2 {} = ()

instance Add V2 where
    zero = V2 0 0
    add (V2 x y) (V2 x' y') = V2 (x + x') (y + y')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Mult Double V2 V2 where
    mult k (V2 x y) = V2 (k * x) (k * y)
    {-# INLINE mult #-}

-- | 2×2 matrix.
data M22 = M22 !Double !Double !Double !Double
  deriving (Eq, Show)

instance NFData M22 where
    rnf M22 {} = ()

instance Add M22 where
    zero = M22 0 0 0 0
    add (M22 a b c d) (M22 a' b' c' d') = M22 (a + a') (b + b') (c + c') (d + d')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Eye M22 where
    eye = M22 1 0 0 1
    {-# INLINE eye #-}

instance Det M22 where det = det2
instance Inv M22 where inv = inv2

instance Mult Double M22 M22 where
    mult k (M22 a b c d) = M22 (k * a) (k * b) (k * c) (k * d)
    {-# INLINE mult #-}

instance Mult M22 V2 V2 where
    mult (M22 a b c d) (V2 u v) = V2 (a * u + b * v) (c * u + d * v)
    {-# INLINE mult #-}

-- | >>> M22 1 2 3 4 `mult` eye @M22
-- M22 1.0 2.0 3.0 4.0
instance Mult M22 M22 M22 where
    mult (M22 a b c d) (M22 x y z w) = M22
        (a * x + b * z) (a * y + b * w)
        (c * x + d * z) (c * y + d * w)
    {-# INLINE mult #-}

det2 :: M22 -> Double
det2 (M22 a b c d) = a * d - b * c
{-# INLINE det2 #-}

inv2 :: M22 -> M22
inv2 m@(M22 a b c d) = M22
    (  d / detm) (- b / detm)
    (- c / detm) (  a / detm)
  where
    detm = det2 m
{-# INLINE inv2 #-}

-- | Symmetric 2x2 matrix.
data SM22 = SM22 !Double !Double !Double
  deriving (Eq, Show)

instance NFData SM22 where
    rnf SM22 {} = ()

instance Add SM22 where
    zero = SM22 0 0 0
    add (SM22 a b d) (SM22 a' b' d') = SM22 (a + a') (b + b') (d + d')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Eye SM22 where
    eye = SM22 1 0 1
    {-# INLINE eye #-}

instance Det SM22 where det = detS2
instance Inv SM22 where inv = invS2

instance Mult Double SM22 SM22 where
    mult k (SM22 a b d) = SM22 (k * a) (k * b) (k * d)
    {-# INLINE mult #-}

instance Mult SM22 V2 V2 where
    mult (SM22 a b d) (V2 u v) = V2 (a * u + b * v) (b * u + d * v)
    {-# INLINE mult #-}

detS2 :: SM22 -> Double
detS2 (SM22 a b d) = a * d - b * b
{-# INLINE detS2 #-}

invS2 :: SM22 -> SM22
invS2 m@(SM22 a b d) = SM22
    (  d / detm)
    (- b / detm) (  a / detm)
  where
    detm = detS2 m
{-# INLINE invS2 #-}

-------------------------------------------------------------------------------
-- 3 dimensions
-------------------------------------------------------------------------------

-- | 3d vector. Strict triple of 'Double's.
--
-- Also used to represent quadratic polynomial: @V3 a b c@  \(= a x^2 + b x + c\).
data V3 = V3 !Double !Double !Double
  deriving (Eq, Show)

instance NFData V3 where
    rnf V3 {} = ()

instance Add V3 where
    zero = V3 0 0 0
    add (V3 x y z) (V3 x' y' z') = V3 (x + x') (y + y') (z + z')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Mult Double V3 V3 where
    mult k (V3 x y z) = V3 (k * x) (k * y) (k * z)
    {-# INLINE mult #-}

-- | 3×3 matrix.
data M33 = M33
    !Double !Double !Double
    !Double !Double !Double
    !Double !Double !Double
  deriving (Eq, Show)

instance NFData M33 where
    rnf M33 {} = ()

instance Add M33 where
    zero = M33 0 0 0 0 0 0 0 0 0

    add (M33 a b c d e f g h i) (M33 a' b' c' d' e' f' g' h' i') = M33
        (a + a') (b + b') (c + c')
        (d + d') (e + e') (f + f')
        (g + g') (h + h') (i + i')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Eye M33 where
    eye = M33 1 0 0
              0 1 0
              0 0 1
    {-# INLINE eye #-}

instance Det M33 where det = det3
instance Inv M33 where inv = inv3

instance Mult Double M33 M33 where
    mult k (M33 a b c d e f g h i) = M33
        (k * a) (k * b) (k * c)
        (k * d) (k * e) (k * f)
        (k * g) (k * h) (k * i)
    {-# INLINE mult #-}

instance Mult M33 V3 V3 where
    mult (M33 a b c
           d e f
           g h i) (V3 u v w) = V3
        (a * u + b * v + c * w)
        (d * u + e * v + f * w)
        (g * u + h * v + i * w)
    {-# INLINE mult #-}

-- TODO: instance Mult M33 M33 M33 where

det3 :: M33 -> Double
det3 (M33 a b c
          d e f
          g h i)
    = a * (e*i-f*h) - d * (b*i-c*h) + g * (b*f-c*e)
{-# INLINE det3 #-}

inv3 :: M33 -> M33
inv3 m@(M33 a b c
            d e f
            g h i)
    = M33 a' b' c'
          d' e' f'
          g' h' i'
  where
    a' = cofactor e f h i / detm
    b' = cofactor c b i h / detm
    c' = cofactor b c e f / detm
    d' = cofactor f d i g / detm
    e' = cofactor a c g i / detm
    f' = cofactor c a f d / detm
    g' = cofactor d e g h / detm
    h' = cofactor b a h g / detm
    i' = cofactor a b d e / detm
    cofactor q r s t = det2 (M22 q r s t)
    detm = det3 m
{-# INLINE inv3 #-}

-- | Symmetric 3×3 matrix.
data SM33 = SM33
    !Double
    !Double !Double
    !Double !Double !Double
  deriving (Eq, Show)

instance NFData SM33 where
    rnf SM33 {} = ()

instance Add SM33 where
    zero = SM33 0 0 0 0 0 0

    add (SM33 a d e g h i) (SM33 a' d' e' g' h' i') = SM33
        (a + a')
        (d + d') (e + e')
        (g + g') (h + h') (i + i')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Eye SM33 where
    eye = SM33 1
               0 1
               0 0 1
    {-# INLINE eye #-}

instance Det SM33 where det = detS3
instance Inv SM33 where inv = invS3

instance Mult Double SM33 SM33 where
    mult k (SM33 a d e g h i) = SM33
        (k * a)
        (k * d) (k * e)
        (k * g) (k * h) (k * i)
    {-# INLINE mult #-}

instance Mult SM33 V3 V3 where
    mult (SM33 a
               d e
               g h i) (V3 u v w) = V3
        (a * u + d * v + g * w)
        (d * u + e * v + h * w)
        (g * u + h * v + i * w)
    {-# INLINE mult #-}

detS3 :: SM33 -> Double
detS3 (SM33 a
            d e
            g h i)
    = a * (e*i-h*h) - d * (d*i-g*h) + g * (d*h-g*e)
{-# INLINE detS3 #-}

invS3 :: SM33 -> SM33
invS3 m@(SM33 a
              d e
              g h i)
     = SM33 a'
            d' e'
            g' h' i'
  where
    a' = cofactor e h h i / detm
    d' = cofactor h d i g / detm
    e' = cofactor a g g i / detm
    g' = cofactor d e g h / detm
    h' = cofactor d a h g / detm
    i' = cofactor a d d e / detm
    cofactor q r s t = det2 (M22 q r s t)
    detm = detS3 m
{-# INLINE invS3 #-}


module NeuralNetwork
  () where

import           Control.Monad         (mapM)
import qualified Data.Sequence         as Seq
import           Debug.Trace
import           Numeric.LinearAlgebra ((<>), (><))
import qualified Numeric.LinearAlgebra as LA

type MatrixNum = LA.Matrix Double
type Cache = ((MatrixNum, MatrixNum, MatrixNum), MatrixNum)

data Parameter = Parameter
  { weightsParam :: MatrixNum
  , biasesParam  :: MatrixNum
  }

-- Utils
emptyx :: MatrixNum
emptyx = (0><0) []

emptyCache :: Cache
emptyCache = ((emptyx, emptyx, emptyx), emptyx)

-- Activation functions
relu :: MatrixNum -> MatrixNum
relu = LA.cmap (max 0)

sigmoid :: MatrixNum -> MatrixNum
sigmoid = LA.cmap (\x -> (1 / (1 + exp (-1 * x))))

relu' :: MatrixNum -> MatrixNum
relu' = LA.cmap f
  where
    f x | x < 0 = 0
        | otherwise = 1

sigmoid' :: MatrixNum -> MatrixNum
sigmoid' xw = (sigmoid xw) * (1 - sigmoid xw)

-- | Initialize parameters. Randomize weights and zeroes biases
initParams :: [Int] -> IO [Parameter]
initParams layerDims =
  (fmap . fmap) (\(bx, wx) -> Parameter wx bx) ((zip biases) <$> weights)
    where
      weights = mapM genParams layerDims'
      layerDims' = (zip layerDims (tail layerDims))
      genParams (prevDim, currentDim) = LA.randn currentDim prevDim
      biases =
        fmap
        (\(_, currentDim) -> (currentDim><1) (replicate currentDim 0))
        layerDims'

-- | Propagate forward to next layer. Calulates next a
linearForward axPrev wx bx activationFun =
  let zx = linearForward' axPrev wx bx
      ax = activationFun zx
      cache = ((axPrev, wx, bx), zx)
      linearForward' axPrev wx bx = (wx <> axPrev) + bx :: MatrixNum
  in  (ax, cache)

-- | Complete forward propagation. Calculates final ax
propagateForward :: MatrixNum
                 -> [Parameter]
                 -> (MatrixNum, [Cache])
propagateForward xx params  =
  (outputAx, hiddenCache ++ [outputCache])
    where
      params' = take ((length params) - 1) params
      f (aPrev, _) (Parameter wx bx) = linearForward aPrev wx bx relu
      hidden = scanl f (xx, emptyCache) params'
      hiddenCache = drop 1 (fmap snd hidden)
      (lastHiddenAx, _) = last hidden
      Parameter wxLast bxLast = last params
      (outputAx, outputCache) = linearForward lastHiddenAx wxLast bxLast sigmoid

-- | Propagate backward to previous layer
linearBackward :: MatrixNum -> Cache -> (MatrixNum -> MatrixNum)
               -> (MatrixNum, MatrixNum, MatrixNum)
linearBackward dax ((axPrev, wx, bx), zx) af' =
  (daxPrev, dwx, dbx)
    where
      dzx = dax * af' zx
      m = fromIntegral $ (snd . LA.size) axPrev
      dwx = (dzx <> (LA.tr axPrev)) / m
      sums = fmap sum (LA.toLists dzx)
      layerSize = (fst . LA.size) wx
      dbx = ((layerSize><1) sums) / m
      daxPrev = (LA.tr wx) <> dzx

-- | Complete backpropagation. Computes gradients (dax, dwx, dbx)
-- of each layer
propagateBackward :: MatrixNum -> (LA.Vector Double) -> [Cache]
                  -> [(MatrixNum, MatrixNum, MatrixNum)]
propagateBackward axOut yx caches =
  let l = length caches
      yx' = LA.reshape 1 yx
      daxOut = - (yx' / axOut) - ((1 - yx') / (1 - axOut))
      gradOut = linearBackward daxOut (last caches) sigmoid'
      hiddenCaches = reverse $ take (length caches - 1) caches
      grads = scanl f gradOut hiddenCaches
      f (dax, _, _) cache = linearBackward dax cache relu'
  in  reverse grads

-- | Perform gradient descent. Computes new decreased gradients (dax, dwx, dbx)
updateParameters :: [Parameter]
                 -> [(MatrixNum, MatrixNum, MatrixNum)]
                 -> Double
                 -> [Parameter]
updateParameters params grads eta = fmap fun (zip params grads)
   where fun ((Parameter wx bx), (_, dwx, dbx)) =
           Parameter
             (wx - (LA.scalar eta) * dwx)
             (bx - (LA.scalar eta) * dbx)

train :: MatrixNum
      -> LA.Vector Double
      -> [Int]
      -> Double
      -> Int
      -> IO [Parameter]
train xx yx layerDims eta iterations =
  do
    params <- initParams layerDims
    return $ train' params iterations
  where
    train' params 0 = params
    train' params times =
      let (ax, caches) = propagateForward xx params
          grads = propagateBackward ax yx caches
          newParams = updateParameters params grads eta
      in  train' newParams (times - 1)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeApplications      #-}

import           Control.Monad
import           Data.Bifunctor
import           Data.Maybe
import           Hedgehog
import           Lens.Micro
import           Nudge
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static          (L, R)
import           System.Exit
import           System.IO
import qualified Numeric.LinearAlgebra.Static.Backprop as B

prop_vec2 :: Property
prop_vec2 = nudgeProp2 B.vec2

prop_vec3 :: Property
prop_vec3 = nudgeProp @(Double, Double, Double)
                (\t -> B.vec3 (t ^^. _1) (t ^^. _2) (t ^^. _3))

prop_vec4 :: Property
prop_vec4 = nudgeProp2 @(Double, Double) @(Double, Double)
                (\x y -> B.vec4 (x ^^. _1) (x ^^. _2) (y ^^. _1) (y ^^. _2))

prop_snoc :: Property
prop_snoc = nudgeProp2 @(R 3) (B.&)

prop_append :: Property
prop_append = nudgeProp2 @(R 3) @(R 2) (B.#)

prop_split1 :: Property
prop_split1 = nudgeProp @(R 3) (fst . B.split @2)

prop_split2 :: Property
prop_split2 = nudgeProp @(R 3) (snd . B.split @2)

prop_headTail1 :: Property
prop_headTail1 = nudgeProp @(R 3) (fst . B.headTail)

prop_headTail2 :: Property
prop_headTail2 = nudgeProp @(R 3) (snd . B.headTail)

prop_vector :: Property
prop_vector = nudgeProp (B.vector @3 . sequenceVar)

prop_linspace :: Property
prop_linspace = nudgeProp2 (B.linspace @3)

prop_row :: Property
prop_row = nudgeProp @(R 3) B.row

prop_col :: Property
prop_col = nudgeProp @(R 3) B.col

prop_horzcat :: Property
prop_horzcat = nudgeProp2 @(L 3 2) @(L 3 1) (B.|||)

prop_vertcat :: Property
prop_vertcat = nudgeProp2 @(L 2 3) @(L 1 3) (B.===)

prop_splitRows1 :: Property
prop_splitRows1 = nudgeProp @(L 2 3) (fst . B.splitRows @1)

prop_splitRows2 :: Property
prop_splitRows2 = nudgeProp @(L 2 3) (snd . B.splitRows @1)

prop_splitCols1 :: Property
prop_splitCols1 = nudgeProp @(L 3 2) (fst . B.splitCols @1)

prop_splitCols2 :: Property
prop_splitCols2 = nudgeProp @(L 3 2) (snd . B.splitCols @1)

prop_unrow :: Property
prop_unrow = nudgeProp @(L 1 3) B.unrow

prop_uncol :: Property
prop_uncol = nudgeProp @(L 3 1) B.uncol

prop_tr :: Property
prop_tr = nudgeProp @(L 3 2) B.tr

prop_diag :: Property
prop_diag = nudgeProp @(R 3) B.diag

prop_svd :: Property
prop_svd = nudgeProp @(L 3 2) B.svd

prop_svd_ :: Property
prop_svd_ = nudgeProp @(L 3 2) ((\(_,x,_) -> x) . B.svd_)

prop_eigensystem1 :: Property
prop_eigensystem1 = nudgeProp @(L 3 2) (fst . B.eigensystem . B.mTm)

prop_eigensystem2 :: Property
prop_eigensystem2 = nudgeProp @(L 3 2) (snd . B.eigensystem . B.mTm)

prop_eigenvalues :: Property
prop_eigenvalues = nudgeProp @(L 3 2) (B.eigenvalues . B.mTm)

prop_chol :: Property
prop_chol = nudgeProp @(L 3 2) (B.chol . B.mTm)

prop_norm_0V :: Property
prop_norm_0V = nudgeProp @(R 3) B.norm_0

prop_norm_0M :: Property
prop_norm_0M = nudgeProp @(L 3 2) B.norm_0

prop_norm_1V :: Property
prop_norm_1V = nudgeProp @(R 3) B.norm_1V

prop_norm_1M :: Property
prop_norm_1M = nudgeProp @(L 3 2) B.norm_1M

prop_norm_2V :: Property
prop_norm_2V = nudgeProp @(R 3) B.norm_2V

prop_norm_2M :: Property
prop_norm_2M = nudgeProp @(L 3 2) B.norm_2M

prop_norm_InfV :: Property
prop_norm_InfV = nudgeProp @(R 3) B.norm_InfV

prop_norm_InfM :: Property
prop_norm_InfM = nudgeProp @(L 3 2) B.norm_InfM

prop_mean :: Property
prop_mean = nudgeProp @(R 3) B.mean

prop_meanCov1 :: Property
prop_meanCov1 = nudgeProp @(L 3 2) (fst . B.meanCov)

prop_meanCov2 :: Property
prop_meanCov2 = nudgeProp @(L 3 2) (B.unSym . snd . B.meanCov)

prop_meanL :: Property
prop_meanL = nudgeProp @(L 3 2) B.meanL

prop_cov :: Property
prop_cov = nudgeProp @(L 3 2) (B.unSym . B.cov)

prop_mul :: Property
prop_mul = nudgeProp2 @(L 3 2) @(L 2 3) B.mul

prop_app :: Property
prop_app = nudgeProp2 @(L 3 2) @(R 2) B.app

prop_dot :: Property
prop_dot = nudgeProp2 @(R 3) @(R 3) B.dot

prop_cross :: Property
prop_cross = nudgeProp2 @(R 3) B.cross

-- TODO: bug in diagR?
-- prop_diagR :: Property
-- prop_diagR = nudgeProp2 genDouble (genVec @3) (B.diagR @5 @4)

-- TODO: Mappers
-- , dvmap
-- , dvmap'
-- , dmmap
-- , dmmap'

prop_outer :: Property
prop_outer = nudgeProp2 @(R 3) @(R 2) B.outer

-- TODO: Zippers
-- , zipWithVector
-- , zipWithVector'

prop_det :: Property
prop_det = nudgeProp @(L 3 3) B.det

prop_invlndet1 :: Property
prop_invlndet1 = nudgeProp @(L 3 3) (fst . B.invlndet)

prop_invlndet2 :: Property
prop_invlndet2 = nudgeProp @(L 3 3) (fst . snd . B.invlndet)

prop_invlndet3 :: Property
prop_invlndet3 = nudgeProp @(L 3 3) (snd . snd . B.invlndet)

prop_lndet :: Property
prop_lndet = nudgeProp @(L 3 3) B.lndet

-- TODO: more general invertible matrix
prop_inv :: Property
prop_inv = nudgeProp @(L 3 2) (B.inv . B.unSym . B.mTm)

prop_toRows :: Property
prop_toRows = nudgeProp @(L 3 2) (collectVar . B.toRows)

prop_toColumns :: Property
prop_toColumns = nudgeProp @(L 2 3) (collectVar . B.toColumns)

prop_fromRows :: Property
prop_fromRows = nudgeProp (B.fromRows @3 @2 . sequenceVar)

prop_fromColumns :: Property
prop_fromColumns = nudgeProp (B.fromColumns @2 @3 . sequenceVar)

prop_konstV :: Property
prop_konstV = nudgeProp (B.konst @_ @(B.R 3))

prop_konstM :: Property
prop_konstM = nudgeProp (B.konst @_ @(B.L 3 2))

prop_sumElementsV :: Property
prop_sumElementsV = nudgeProp @(R 3) B.sumElements

prop_sumElementsM :: Property
prop_sumElementsM = nudgeProp @(L 3 2) B.sumElements

prop_extractV :: Property
prop_extractV = nudgeProp (B.extractV @_ @(R 3))

prop_extractM :: Property
prop_extractM = nudgeProp (B.extractM @_ @(L 3 2))

prop_createV :: Property
prop_createV = nudgeProp (fromMaybe 0 . B.create @_ @(R 3))

prop_createM :: Property
prop_createM = nudgeProp (fromMaybe 0 . B.create @_ @(L 3 2))

prop_takeDiag :: Property
prop_takeDiag = nudgeProp @(L 3 3) B.takeDiag

prop_sym :: Property
prop_sym = nudgeProp @(L 3 3) (B.unSym . B.sym)

prop_mTm :: Property
prop_mTm = nudgeProp @(L 3 2) (B.unSym . B.mTm)

prop_unSym :: Property
prop_unSym = nudgeProp @(L 3 3) (B.unSym . B.sym)

tryGroup :: (forall a. Num a => a) -> Group -> Group
tryGroup n Group{..} =
    Group groupName
          ((map . second) (withDiscards n . withTests n)
                          groupProperties
          )

main :: IO ()
main = do
  hSetBuffering stdout LineBuffering
  hSetBuffering stderr LineBuffering

  results <- checkParallel (tryGroup 100 $$(discover))

  unless results exitFailure


module Mbrot where
import           Codec.Picture.Types
import           Coloring
import           Control.Arrow       ((***))
import           Data.Complex
import           ImageOutput

-- |Size parameters
size :: (Int,Int)
size = (1400,800)

-- |Floating point size, calculated
sizef :: (Float,Float)
sizef = (fromIntegral *** fromIntegral) size

-- |Maximum iterations
maxIter :: Int
maxIter = 1000

-- |Color step for integer esacpe coloring
mult :: Float
mult = 255 / fromIntegral maxIter

-- |Color step for continuously colored
multg :: Float
multg = 400 / fromIntegral maxIter

-- |Function for rendering pixel in greyscale
drawGreyscale :: Bool    -- ^Continuously colored
              -> Int
              -> Int
              -> Pixel8
drawGreyscale c x y
    |c =  fromIntegral (round(float * mult))
    |otherwise = fromIntegral(int * round mult)
    where int = mbrot maxIter (((x1/x2)-0.5) :+ (y1/y2))
          float = mbrotContinuous maxIter (((x1/x2)-0.5) :+ (y1/y2))
          x1 = 2 * (fromIntegral x - x2)
          y1 = 2 * (fromIntegral y - y2)
          x2 = fst sizef /2
          y2 = snd sizef /2

-- |Deprecated. Use 'drawColorZoom' instead
drawColor :: Int -> Int -> PixelRGB8
drawColor x y =  colorToPixel $ color (fromIntegral (int*400)/1000)
    where int = mbrot maxIter ((x1/x2) :+ (y1/y2))
          x1 = 1.75 * (fromIntegral x - x2)
          y1 = fromIntegral y - y2
          x2 = fst sizef /2
          y2 = snd sizef /2

-- |Function for rendering pixel in color, with zoom level and focus
drawColorZoom :: Float         -- ^ Zoom level
              -> (Float,Float) -- ^ Focus point (x,y)
              -> Int
              -> Int
              -> PixelRGB8
drawColorZoom zoom (fx, fy) x y =  colorToPixel $ color (float*multg)
    where float = mbrotContinuous maxIter (((x1/x2)+fx) :+ ((y1/y2)+fy))
          x1 = 1.75 * (fromIntegral x - x2) / zoom
          y1 = (fromIntegral y - y2) / zoom
          x2 = fst sizef /2
          y2 = snd sizef /2

-- |Function for drawing color gradient. Tests for correct linear interpolation of color
drawTestGradient :: Int -> Int -> PixelRGB8
drawTestGradient x _ = colorToPixel $ color (fromIntegral x *400/1000)

-- |Function for integer escape time of mbrot set of Re(x) + Im(y)
mbrot :: RealFloat x
     => Int           -- ^ Maximum iterations
     -> Complex x     -- ^ Re(x) + Im(y)
     -> Int
mbrot n = mbrot' maxIter (0 :+ 0)
    where
        mbrot' :: RealFloat x => Int -> Complex x -> Complex x -> Int
        mbrot' n a c
            |magnitude newIter < 65535 && n > 0 = mbrot' (n-1) newIter c
            |otherwise = maxIter - n
                where
                    newIter = a*a + c

-- |Function for real floating escape time of mbrot set of Re(x) + Im(y)
mbrotContinuous :: RealFloat x
                => Int           -- ^ Maximum iterations
                -> Complex x     -- ^ Re(x) + Im(y)
                -> x
mbrotContinuous n = mbrot' maxIter (0 :+ 0)
    where
        mbrot' :: RealFloat x => Int -> Complex x -> Complex x -> x
        mbrot' n a c
            |magnitude newIter < 65535 && n > 0 = mbrot' (n-1) newIter c
            |otherwise = normalized
                where
                    newIter = a*a + c
                    normalized =  if n > 0
                                  then
                                    fromIntegral(maxIter - n +1) - logBase 2 ( logBase 2 (magnitude newIter ^ 2) / 2 )
                                    --fromIntegral (maxIter - n +1) - (log (log (magnitude newIter)))/ log (2.0)
                                  else fromIntegral maxIter


module MachineLearning.Classification.BinaryTest
(
  tests
  , testOptPath
)

where

import Test.Framework (testGroup)
import Test.Framework.Providers.HUnit
import Test.HUnit
import Test.HUnit.Approx
import Test.HUnit.Plus

import MachineLearning.Types
import qualified Data.Vector as V
import qualified Numeric.LinearAlgebra as LA

import MachineLearning.DataSets (dataset2)

import qualified MachineLearning as ML
import MachineLearning.Classification.Binary

(x, y) = ML.splitToXY dataset2

processX muSigma x = ML.addBiasDimension $ ML.featureNormalization muSigma $ ML.mapFeatures 6 x

muSigma = ML.meanStddev (ML.mapFeatures 6 x)
x1 = processX muSigma x
zeroTheta = LA.konst 0 (LA.cols x1)

xPredict = LA.matrix 2 [ -0.5, 0.5
                       , 0.2, -0.2
                       , 1, 1
                       , 1, 0
                       , 0, 0
                       , 0, 1]
xPredict1 = processX muSigma xPredict
yExpected = LA.vector [1, 1, 0, 0, 1, 0]

eps = 0.0001

-- Binary

(thetaCGFR, optPathCGFR) = learn (ConjugateGradientFR 0.1 0.1) eps 50 (L2 0.5) x1 y zeroTheta
(thetaCGPR, optPathCGPR) = learn (ConjugateGradientPR 0.1 0.1) eps 50 (L2 0.5) x1 y zeroTheta
(thetaBFGS, optPathBFGS) = learn (BFGS2 0.1 0.1) eps 50 (L2 0.5) x1 y zeroTheta

isInDescendingOrder :: V.Vector Double -> Bool
isInDescendingOrder lst = V.and . snd . V.unzip $ V.scanl (\(prev, _) current -> (current, prev-current > (-0.001))) (1/0, True) lst

testOptPath optPath = do
  let js = V.convert $ (LA.toColumns optPath) !! 1
  assertBool ("non-increasing errors: " ++ show js) $ isInDescendingOrder js

testAccuracyBinary theta eps = do
  let yPredicted = predict x1 theta
      accuracy = calcAccuracy y yPredicted
  assertApproxEqual "" eps 1 accuracy

tests = [
  testGroup "learn" [
      testCase "Conjugate Gradient FR" $ assertVector "" 0.01 yExpected (predict xPredict1 thetaCGFR)
      , testCase "Conjugate Gradient PR" $ assertVector "" 0.01 yExpected (predict xPredict1 thetaCGPR)
      , testCase "BFGS" $ assertVector "" 0.01 yExpected (predict xPredict1 thetaBFGS)
      ]
  , testGroup "optPath" [
      testCase "Conjugate Gradient FR" $ testOptPath optPathCGFR
      , testCase "Conjugate Gradient PR" $ testOptPath optPathCGPR
      , testCase "BFGS" $ testOptPath optPathBFGS
      ]
    , testGroup "accuracy" [
        testCase "Conjugate Gradient FR" $ testAccuracyBinary thetaCGFR 0.2
        , testCase "Conjugate Gradient PR" $ testAccuracyBinary thetaCGPR 0.2
        , testCase "BFGS" $ testAccuracyBinary thetaBFGS 0.2
        ]
  ]


{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RecordWildCards   #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

module Internal.CG(
    cgSolve, cgSolve',
    CGState(..), R, V
) where

import           Control.Arrow       ((***))
import           Internal.Algorithms (NormType (..), linearSolve, linearSolveLS, pnorm,
                                      relativeError)
import           Internal.Container
import           Internal.Element
import           Internal.IO
import           Internal.Matrix
import           Internal.Numeric
import           Internal.Sparse
import           Internal.Vector
import           Numeric.Vector      ()

{-
import Util.Misc(debug, debugMat)

(//) :: Show a => a -> String -> a
infix 0 // -- , ///
a // b = debug b id a

(///) :: V -> String -> V
infix 0 ///
v /// b = debugMat b 2 asRow v
-}

type V = Vector R

data CGState = CGState
    { cgp  :: Vector R  -- ^ conjugate gradient
    , cgr  :: Vector R  -- ^ residual
    , cgr2 :: R         -- ^ squared norm of residual
    , cgx  :: Vector R  -- ^ current solution
    , cgdx :: R         -- ^ normalized size of correction
    }

cg :: Bool -> (V -> V) -> (V -> V) -> CGState -> CGState
cg sym at a (CGState p r r2 x _) = CGState p' r' r'2 x' rdx
  where
    ap1 = a p
    ap  | sym       = ap1
        | otherwise = at ap1
    pap | sym       = p <.> ap1
        | otherwise = norm2 ap1 ** 2
    alpha = r2 / pap
    dx = scale alpha p
    x' = x + dx
    r' = r - scale alpha ap
    r'2 = r' <.> r'
    beta = r'2 / r2
    p' = r' + scale beta p

    rdx = norm2 dx / max 1 (norm2 x)

conjugrad
  :: Bool -> GMatrix -> V -> V -> R -> R -> [CGState]
conjugrad sym a b = solveG sym (tr a !#>) (a !#>) (cg sym) b

solveG
    :: Bool
    -> (V -> V) -> (V -> V)
    -> ((V -> V) -> (V -> V) -> CGState -> CGState)
    -> V
    -> V
    -> R -> R
    -> [CGState]
solveG sym mat ma meth rawb x0' ϵb ϵx
    = takeUntil ok . iterate (meth mat ma) $ CGState p0 r0 r20 x0 1
  where
    a = if sym then ma else mat . ma
    b = if sym then rawb else mat rawb
    x0  = if x0' == 0 then konst 0 (dim b) else x0'
    r0  = b - a x0
    r20 = r0 <.> r0
    p0  = r0
    nb2 = b <.> b
    ok CGState {..}
        =  cgr2 <nb2*ϵb**2
        || cgdx < ϵx

takeUntil :: (a -> Bool) -> [a] -> [a]
takeUntil q xs = a++ take 1 b
  where
    (a,b) = break q xs

-- | Solve a sparse linear system using the conjugate gradient method with default parameters.
cgSolve
  :: Bool          -- ^ is symmetric
  -> GMatrix       -- ^ coefficient matrix
  -> Vector R      -- ^ right-hand side
  -> Vector R      -- ^ solution
cgSolve sym a b  = cgx $ last $ cgSolve' sym 1E-4 1E-3 n a b 0
  where
    n = max 10 (round $ sqrt (fromIntegral (dim b) :: Float))

-- | Solve a sparse linear system using the conjugate gradient method with default parameters.
cgSolve'
  :: Bool      -- ^ symmetric
  -> R         -- ^ relative tolerance for the residual (e.g. 1E-4)
  -> R         -- ^ relative tolerance for δx (e.g. 1E-3)
  -> Int       -- ^ maximum number of iterations
  -> GMatrix   -- ^ coefficient matrix
  -> Vector R  -- ^ initial solution
  -> Vector R  -- ^ right-hand side
  -> [CGState] -- ^ solution
cgSolve' sym er es n a b x = take n $ conjugrad sym a b x er es

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

instance Testable GMatrix
  where
    checkT _ = (ok,info)
      where
        sma = convo2 20 3
        x1 = vect [1..20]
        x2 = vect [1..40]
        sm = mkSparse sma
        dm = toDense sma

        s1 = sm !#> x1
        d1 = dm #> x1

        s2 = tr sm !#> x2
        d2 = tr dm #> x2

        sdia = mkDiagR 40 20 (vect [1..10])
        s3 =    sdia !#> x1
        s4 = tr sdia !#> x2
        ddia = diagRect 0 (vect [1..10])  40 20
        d3 = ddia #> x1
        d4 = tr ddia #> x2

        v = testb 40
        s5 = cgSolve False sm v
        d5 = denseSolve dm v

        symassoc = [((0,0),1.0),((1,1),2.0),((0,1),0.5),((1,0),0.5)]
        b = vect [3,4]
        d6 = flatten $ linearSolve (toDense symassoc) (asColumn b)
        s6 = cgSolve True (mkSparse symassoc) b

        info = do
            print sm
            disp (toDense sma)
            print s1; print d1
            print s2; print d2
            print s3; print d3
            print s4; print d4
            print s5; print d5
            print $ relativeError (pnorm Infinity) s5 d5
            print s6; print d6
            print $ relativeError (pnorm Infinity) s6 d6

        ok = s1==d1
          && s2==d2
          && s3==d3
          && s4==d4
          && relativeError (pnorm Infinity) s5 d5 < 1E-10
          && relativeError (pnorm Infinity) s6 d6 < 1E-10

        disp = putStr . dispf 2

        vect = fromList :: [Float] -> Vector Float

        convomat :: Int -> Int -> AssocMatrix
        convomat n k = [ ((i,j `mod` n),1) | i<-[0..n-1], j <- [i..i+k-1]]

        convo2 :: Int -> Int -> AssocMatrix
        convo2 n k = m1 ++ m2
          where
            m1 = convomat n k
            m2 = map (((+n) *** id) *** id) m1

        testb n = vect $ take n $ cycle ([0..10]++[9,8..1])

        denseSolve a = flatten . linearSolveLS a . asColumn

        -- mkDiag v = mkDiagR (dim v) (dim v) v


import Data.Bool
import Data.Complex (Complex((:+)), magnitude)

mandelbrot
  :: RealFloat a
  => Complex a -> Complex a
mandelbrot a = iterate ((a +) . (^ 2)) 0 !! 50

main :: IO ()
main =
  mapM_
    putStrLn
    [ [ bool ' ' '*' (2 > magnitude (mandelbrot (x :+ y)))
      | x <- [-2,-1.9685 .. 0.5] ]
    | y <- [1,0.95 .. -1] ]


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE BangPatterns #-}

{-# OPTIONS_GHC -fno-warn-orphans #-} -- do not warn on local instance of Random for Complex

module BBrotCompute(compute, selectCells, inSet) where

import Control.Monad.Par
import Data.Aeson(encode)
import Data.Array
import qualified Data.ByteString.Lazy as BS
import Data.Complex
import Data.List
import Data.Maybe
import Data.Time
import System.Environment
import System.Console.CmdArgs(whenNormal, whenLoud)
import System.Random
import Text.Printf

import BBrotConf
import BBrotSelection

instance Random (Complex Double) where
  randomR (!(x0 :+ y0), !(x1 :+ y1)) g = (r :+ i, g2)
    where (!r, !g1) = randomR (x0, x1) g
          (!i, !g2) = randomR (y0, y1) g1
  random = randomR (0.0 :+ 0.0, 1.0 :+ 1.0)

data GridPoint = GridPoint { xbase :: !Double
                           , ybase :: !Double
                           , inMandelbrotSet :: !Bool
                           }
                 deriving Show

makeGrid :: Double -> Int -> Array (Int, Int) GridPoint
makeGrid step bailout = array ((0, 0), (xsteps, ysteps)) points
  where (xmin :+ ymin) = loCorner
        (xmax :+ ymax) = hiCorner
        xsteps = floor $ (xmax - xmin) / step :: Int
        ysteps = floor $ (ymax - ymin) / step :: Int
        points = [ ((i,j), g i j) | i <- [0..xsteps], j <- [0..ysteps] ]
        g i j = GridPoint x y p
          where x = xmin + fromIntegral i * step
                y = ymin + fromIntegral j * step
                p = not $ inSet 0 bailout (x :+ y)

selectCells :: Double -> Int -> [(Double, Double)]
selectCells step bailout = [ (xbase g, ybase g) | g <- cells ]
  where !grid = makeGrid step bailout
        ((x0, y0), (xsteps, ysteps)) = bounds grid
        !cells = [ grid!(i,j) | i <- [x0 + 1 .. xsteps - 1],
                                j <- [y0 + 1 .. ysteps - 1],
                                isNearBorder i j ]
        isNearBorder i j = someIn && someOut
            where neighbours = [ grid!(i+a,j+b) | (a, b) <- range ((-1,-1), (1,1)) ]
                  inMandel = map inMandelbrotSet neighbours
                  someIn = or inMandel
                  someOut = not $ and inMandel

iterations :: Int -> Double -> Double -> Double -> Double -> Int -> Int
iterations !maxK !x !y !x0 !y0 !k =
  let x2 = x * x
      y2 = y * y
  in
   if k == maxK || x2 + y2 > 4
   then k
   else iterations maxK (x2 - y2 + x0) (2 * x * y + y0) x0 y0 (k + 1)

inCardioBulb :: Double -> Double -> Bool
inCardioBulb !x !y = inCardio || inBulb
  where sqr a = a * a
        q = sqr (x - 1/4) + sqr y
        inCardio = q * (q + (x - 1/4)) < sqr y / 4
        inBulb = sqr (x + 1) + sqr y < 1 / 16

toMaybeBBPoint :: Int -> Int -> Complex Double -> Maybe BBPoint
toMaybeBBPoint !minK !maxK !z =
  let x = realPart z
      y = imagPart z
      k = iterations maxK x y x y 0
  in if not (inCardioBulb x y) && k >= minK && k < maxK
     then Just (BBPoint x y k)
     else Nothing

inSet :: Int -> Int -> Complex Double -> Bool
inSet !minK !maxK !z = p $ toMaybeBBPoint minK maxK z
  where p (Just _) = False
        p Nothing  = True

toUnit :: Int -> String
toUnit n = show (n `div` d) ++ unit
  where pow :: Int -> Int -> Int
        pow x y = x ^ y
        l = [ (10 `pow` 9, "G"), (10 `pow` 6, "M"), (10 `pow` 3, "K") ]
        (d, unit) = fromMaybe (1, "") $ find ((<= n) . fst) l

compute :: BBrotConf -> IO ()
compute conf = do
  -- Pick random points in the complex plane, test their orbits,
  -- save the interesting ones.

  whenNormal $ putStrLn $ printf "Sampling %s points..." $ toUnit $ samples conf

  initRandGen <- case seed conf of
    Just s -> return $ mkStdGen s
    Nothing -> getStdGen

  let step = gridStep conf
      cells = selectCells step 255
      pointsPerCell = samples conf `div` length cells
      generators = iterate (fst . split) initRandGen
      genPoints (cell, gen) = take pointsPerCell $ randomRs (corners cell) gen
      corners (x, y) = ((x :+ y) - cellDiag / 2, (x :+ y) + cellDiag / 2)
      cellDiag = step :+ step

  let pointLists = map genPoints $ zip cells generators

  let selected = concat $ concat $ runPar $ parMap (map f) pointLists
        where f = maybeToList . toMaybeBBPoint (minIters conf) (maxIters conf)

  let cachefile = fromMaybe defPath (ocachepath conf)
        where defPath = printf "/tmp/buddhabrot-%s-%s_%s.bbc"
                        (toUnit $ samples conf)
                        (toUnit $ minIters conf)
                        (toUnit $ maxIters conf)

  whenNormal $ do
    putStrLn $ "Selected cells: " ++ show (length cells)
    putStrLn $ "Writing cache " ++ cachefile ++ " ..."
  args <- getArgs
  currentTime <- getCurrentTime
  BS.writeFile cachefile $ encode $
    PointSelection selected (Just args) (Just $ show initRandGen) (Just $ show currentTime)
  whenLoud $ putStrLn $ "selected points: " ++ show (length selected)


{-|
Module      : Univariate
Description : Univariate polynomials (notably cyclotomics)
Copyright   : (c) Matthew Amy, 2020
Maintainer  : matt.e.amy@gmail.com
Stability   : experimental
Portability : portable
-}

module Feynman.Algebra.Polynomial.Univariate(
  Univariate,
  Cyclotomic,
  var,
  constant,
  (*|),
  primitiveRoot,
  evaluate
  )
where

import Data.List
import Data.Map(Map)
import qualified Data.Map as Map
import Data.Complex
import Data.Maybe(maybe)

import qualified Feynman.Util.Unicode as Unicode
import Feynman.Algebra.Base
import Feynman.Algebra.Polynomial

{-------------------------------
 Univariate polynomials
 -------------------------------}

-- | Univariate polynomials over the ring 'r'
data Univariate r = Univariate { getCoeffs :: !(Map Integer r) } deriving (Eq, Ord)

instance (Eq r, Num r, Show r) => Show (Univariate r) where
  show = showWithName "x"

instance Degree (Univariate r) where
  degree = maybe 0 (fromIntegral . fst) . Map.lookupMax . getCoeffs

-- | Print a univariate polynomial with a particular variable name
showWithName :: (Eq r, Num r, Show r) => String -> Univariate r -> String
showWithName x p
    | p == 0    = "0"
    | otherwise = intercalate " + " $ map showTerm (Map.assocs $ getCoeffs p)
    where showTerm (expt, a)
            | expt == 0  = show a
            | a    == 1  = showExpt expt
            | a    == -1 = "-" ++ showExpt expt
            | otherwise  = show a ++ showExpt expt
          showExpt expt
            | expt == 1  = x
            | otherwise  = Unicode.sup x expt

instance (Eq r, Num r) => Num (Univariate r) where
  (+)           = add
  (*)           = mult
  negate        = Univariate . Map.map negate . getCoeffs
  abs           = id
  signum        = id
  fromInteger 0 = Univariate Map.empty
  fromInteger i = Univariate $ Map.singleton 0 (fromInteger i)

-- | Normalize a univariate polynomial
normalize :: (Eq r, Num r) => Univariate r -> Univariate r
normalize = Univariate . Map.filter (/=0) . getCoeffs

-- | The unique univariate variable
var :: Num r => Univariate r
var = Univariate $ Map.singleton 1 1

-- | Constant polynomial
constant :: (Eq r, Num r) => r -> Univariate r
constant a
  | a == 0    = Univariate $ Map.empty
  | otherwise = Univariate $ Map.singleton 0 a

-- | Multiply by a scalar
(*|) :: (Eq r, Num r) => r -> Univariate r -> Univariate r
(*|) 0 = \_p -> zero
(*|) a = Univariate . Map.map (a*) . getCoeffs

-- | Add two univariate polynomials
add :: (Eq r, Num r) => Univariate r -> Univariate r -> Univariate r
add p = normalize . Univariate . Map.unionWith (+) (getCoeffs p) . getCoeffs

-- | Multiply two univariate polynomials
mult :: (Eq r, Num r) => Univariate r -> Univariate r -> Univariate r
mult p = normalize . Map.foldrWithKey (\expt a -> add (mulTerm expt a)) 0 . getCoeffs
  where mulTerm expt a = Univariate . Map.mapKeysMonotonic (+ expt) . Map.map (* a) $ getCoeffs p

{-------------------------------
 Cyclotomics
 -------------------------------}
-- | Cyclotomic polynomials over the ring 'r'
data Cyclotomic r = Cyc { getOrder :: !Integer, getPoly :: !(Univariate r) } deriving (Eq, Ord)

instance (Eq r, Num r, Show r) => Show (Cyclotomic r) where
  show p = showWithName (Unicode.sub Unicode.zeta (getOrder p)) $ getPoly p

instance Degree (Cyclotomic r) where
  degree = degree . getPoly

instance (Eq r, Num r) => Num (Cyclotomic r) where
  p + q = reduceOrder $ Cyc m (p' + q')
    where (Cyc m p', Cyc _ q') = unifyOrder p q
  p * q = reduceOrder $ Cyc m (p' * q')
    where (Cyc m p', Cyc _ q') = unifyOrder p q
  negate (Cyc m p) = Cyc m $ negate p
  abs    p = p
  signum p = p
  fromInteger i = Cyc 1 (fromInteger i)

-- | Unify the order of two cyclotomics
unifyOrder :: Cyclotomic r -> Cyclotomic r -> (Cyclotomic r, Cyclotomic r)
unifyOrder (Cyc n p) (Cyc m q)
  | n == m    = (Cyc n p, Cyc m q)
  | otherwise = (Cyc r p', Cyc r q')
  where r  = lcm n m
        p' = Univariate . Map.mapKeysMonotonic ((r `div` n) *) $ getCoeffs p
        q' = Univariate . Map.mapKeysMonotonic ((r `div` m) *) $ getCoeffs q

-- | Rewrite the cyclotomic in lowest order
reduceOrder :: (Eq r, Num r) => Cyclotomic r -> Cyclotomic r
reduceOrder (Cyc m c) = Cyc m' c'
  where m' = m `div` d
        c' = Univariate . Map.mapKeysMonotonic (`div` d) $ getCoeffs c
        d  = foldr gcd m . Map.keys $ getCoeffs c

-- | Construct the cyclotomic polynomial \(\zeta_i\)
primitiveRoot :: Num r => Integer -> Cyclotomic r
primitiveRoot i = Cyc i var

-- | Convert to a complex number
evaluate :: (Real r, RealFloat f) => Cyclotomic r -> Complex f
evaluate (Cyc m p) = Map.foldrWithKey f (0.0 :+ 0.0) $ getCoeffs p
  where f root coeff = (mkPolar (realToFrac coeff) (expnt root) +)
        expnt root   = 2.0*pi*(fromInteger root)/(fromInteger m)


import Graphics.Gloss
import Graphics.Gloss.Juicy
import Codec.Picture
import Codec.Picture.RGBA8
import Data.Complex

width = 512
height = width `div` 2

mandelbrotness :: Complex Float -> Complex Float -> Float
mandelbrotness c s = (fromIntegral $ length $ takeWhile (\c -> magnitude c <= 2) $ take maxIter $ z c) / (fromIntegral maxIter) where
  z c = iterate (\x -> x^2 + c) s
  maxIter = 32

image :: Float -> Image PixelRGB8
image time = generateImage (\x y -> let [x',y',w,h] = map fromIntegral [x,y,width,height] in mandelbrotColor ((4*x'/w - 2.5) :+ (2*y'/h - 1)) (0.1*time :+ 0) ) width height where
  mandelbrotColor c s = PixelRGB8 (floor $ 255 * mandelbrotness c s) 120 120

main :: IO ()
main = animate (InWindow "Mandelbrot" (width,height) (0,0)) black (fromImageRGB8 . image)


module Mobius where

import Data.Complex

moebius :: (Floating a, Fractional a, Eq a) => a -> a -> a -> a -> a -> a
moebius a b c d z
    | a*d - b*c == 0 = undefined
    | otherwise = (a*z + b) / (c*z + d)

lambdamoebius :: (RealFloat a, Fractional a, Eq a) =>
    Complex a -> Complex a -> Complex a
lambdamoebius a z = (z - a) / (1 - ((conjugate a) * z))

circular :: (RealFloat a, Fractional a) => Complex a -> Complex a
circular z = möbius 0 (-1) 1 0 z

elliptic :: (Floating a, Fractional a, Eq a) => a -> a -> a
elliptic theta z = möbius (cos a) (-sin a) (sin a) (cos a) z --(cis theta) * z
    where a = theta

parabolic :: (Floating a, Fractional a) => a -> a -> a
parabolic a z = z + a -- translation

hyperbolic :: (Floating a, Fractional a) => a -> a -> a
hyperbolic theta z = (exp theta) * z

loxodromic :: (Floating a, Fractional a) => a -> a -> a
loxodromic a z = a * z -- S-shaped path

projection :: (Floating a, Fractional a) => a -> a
projection z = tanh


-- |
-- Module      : BattleHack.Utilities.Vector
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created September 14 2015

-- TODO | - Is it silly to use Complex Double for everything because of Cairo (yes it is; fixed) (✓)
--        - Many of these functions should probably be moved to Southpaw

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module BattleHack.Utilities.Vector where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex

import BattleHack.Types
import BattleHack.Lenses

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------
-- |
dotwise :: (a -> a -> a) -> Complex a -> Complex a -> Complex a
dotwise f (x:+y) (x':+y') = f x x':+f y y'

-- |
dotmap :: (a -> a) -> Complex a -> Complex a
dotmap f (x:+y) = f x:+f y

-- |
toscalar :: (a -> a -> a) -> Complex a -> a
toscalar f (x:+y) = f x y

--  | Converts a 2-tuple to a vector
tovector :: (a, a) -> Complex a
tovector =  uncurry (:+)

-- |
vectorise :: RealFloat f => (f -> f -> a) -> Complex f -> a
vectorise f (x:+y) = f x y


{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds, PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module TensorHMatrix where

import Nats
import Numeric.LinearAlgebra hiding ((!))
import Data.Vector.Storable
import Data.Default
import Data.Singletons.Prelude

import Data.Vinyl

import Prelude hiding (zipWith)

type Usable a = (Element a, Num a, Numeric a, Num (Vector a), Container Vector a)

type IntegralK (p :: KProxy k) = (SingKind p, Integral (DemoteRep p))
type IntegralN (n :: k) = IntegralK ('KProxy :: KProxy k)

natVal' :: (Num a, IntegralN n) => Sing n -> a
natVal' = fromIntegral . fromSing

data Tensor (dims :: [k]) a where
  Scalar :: a -> Tensor '[] a
  Vector :: Vector a -> Tensor '[n] a
  Matrix :: Matrix a -> Tensor '[n, m] a

deriving instance (Show a, Element a) => Show (Tensor dims a)

instance (Default a) => Default (Tensor '[] a) where
  def = Scalar def

fill1 :: forall a n. (SingI n, IntegralN n, Usable a) => a -> Tensor '[n] a
fill1 a = Vector $ konst a (natVal' (sing::Sing n))

instance (SingI n, IntegralN n, Default a, Usable a) => Default (Tensor '[n] a) where
  def = fill1 def

instance (SingI n, IntegralN n, SingI m, Default a, Usable a) => Default (Tensor '[n, m] a) where
  def = Matrix $ konst def (natVal' (sing::Sing n), natVal' (sing::Sing m))

--type instance IndexOf (Tensor l) = Rec LT l
--type instance IndexOf (Tensor l) = Rec (Const Int) l
--type instance IndexOf (Tensor l) = Vec (Length l) Int

{- Fails backwards fundep :(
instance Indexable (Tensor (n ': l) a) (Tensor l a) where
  Matrix m ! i = Vector (m ! i)
  Vector v ! i = Scalar (v ! i)
-}

select :: Storable a => Int -> Tensor '[n] a -> a
select i (Vector v) = v ! i

-- pretty much all of these instances should be automatically derivable

instance Num a => Num (Tensor '[] a) where
  Scalar a + Scalar b = Scalar (a + b)
  Scalar a - Scalar b = Scalar (a - b)
  Scalar a * Scalar b = Scalar (a * b)
  negate (Scalar a) = Scalar (negate a)
  abs (Scalar a) = Scalar (abs a)
  signum (Scalar a) = Scalar (signum a)
  fromInteger n = Scalar (fromInteger n)

instance (SingI n, IntegralN n, Usable a) => Num (Tensor '[n] a) where
  Vector a + Vector b = Vector (a + b)
  Vector a - Vector b = Vector (a - b)
  Vector a * Vector b = Vector (a * b)
  negate (Vector a) = Vector (negate a)
  abs (Vector a) = Vector (abs a)
  signum (Vector a) = Vector (signum a)
  fromInteger n = Vector $ konst (fromInteger n) (natVal' (sing::Sing n))

instance (SingI n, IntegralN n, SingI m, Usable a) => Num (Tensor '[n, m] a) where
  Matrix a + Matrix b = Matrix (a + b)
  Matrix a - Matrix b = Matrix (a - b)
  Matrix a * Matrix b = Matrix (a * b)
  negate (Matrix a) = Matrix (negate a)
  abs (Matrix a) = Matrix (abs a)
  signum (Matrix a) = Matrix (signum a)
  fromInteger n = Matrix $ konst (fromInteger n) (natVal' (sing::Sing n), natVal' (sing::Sing m))

instance Fractional a => Fractional (Tensor '[] a) where
  Scalar a / Scalar b = Scalar (a / b)
  recip (Scalar a) = Scalar (recip a)
  fromRational r = Scalar (fromRational r)

instance (SingI n, IntegralN n, Numeric a, Fractional a, Fractional (Vector a)) => Fractional (Tensor '[n] a) where
  Vector a / Vector b = Vector (a / b)
  recip (Vector a) = Vector (recip a)
  fromRational r = Vector $ konst (fromRational r) (natVal' (sing::Sing n))

instance Floating a => Floating (Tensor '[] a) where
  pi = Scalar pi
  exp (Scalar a) = Scalar (exp a)
  log (Scalar a) = Scalar (log a)
  sqrt (Scalar a) = Scalar (sqrt a)
  Scalar a ** Scalar b = Scalar (a ** b)
  logBase (Scalar a) (Scalar b) = Scalar (logBase a b)
  sin (Scalar a) = Scalar (sin a)
  cos (Scalar a) = Scalar (cos a)
  tan (Scalar a) = Scalar (tan a)
  asin (Scalar a) = Scalar (asin a)
  acos (Scalar a) = Scalar (acos a)
  atan (Scalar a) = Scalar (atan a)
  sinh (Scalar a) = Scalar (sinh a)
  cosh (Scalar a) = Scalar (cosh a)
  tanh (Scalar a) = Scalar (tanh a)
  asinh (Scalar a) = Scalar (asinh a)
  acosh (Scalar a) = Scalar (acosh a)
  atanh (Scalar a) = Scalar (atanh a)

instance (SingI n, IntegralN n, Usable a, Floating a, Floating (Vector a)) => Floating (Tensor '[n] a) where
  pi = fill1 pi
  exp (Vector a) = Vector (exp a)
  log (Vector a) = Vector (log a)
  sqrt (Vector a) = Vector (sqrt a)
  Vector a ** Vector b = Vector (a ** b)
  logBase (Vector a) (Vector b) = Vector (logBase a b)
  sin (Vector a) = Vector (sin a)
  cos (Vector a) = Vector (cos a)
  tan (Vector a) = Vector (tan a)
  asin (Vector a) = Vector (asin a)
  acos (Vector a) = Vector (acos a)
  atan (Vector a) = Vector (atan a)
  sinh (Vector a) = Vector (sinh a)
  cosh (Vector a) = Vector (cosh a)
  tanh (Vector a) = Vector (tanh a)
  asinh (Vector a) = Vector (asinh a)
  acosh (Vector a) = Vector (acosh a)
  atanh (Vector a) = Vector (atanh a)

tmap :: (Container Vector a, Num a, Element b) => (a -> b) -> Tensor dims a -> Tensor dims b
tmap f (Scalar a) = Scalar $ f a
tmap f (Vector v) = Vector $ cmap f v
tmap f (Matrix m) = Matrix $ cmap f m

tZipWith :: (Storable a, Storable b, Storable c) => (a -> b -> c) -> Tensor '[n] a -> Tensor '[n] b -> Tensor '[n] c
tZipWith f (Vector a) (Vector b) = Vector (zipWith f a b)

transpose :: (Numeric a) => Tensor '[n, m] a -> Tensor '[m, n] a
transpose (Matrix m) = Matrix (tr' m)

asCol :: Storable a => Tensor '[n] a -> Tensor '[n, FromInteger 1] a
asCol (Vector v) = Matrix $ asColumn v

asRow' :: Storable a => Tensor '[n] a -> Tensor '[FromInteger 1, n] a
asRow' (Vector v) = Matrix $ asRow v

dot :: (Storable a, Numeric a) => Tensor '[n] a -> Tensor '[n] a -> a
dot (Vector a) (Vector b) = a <.> b

mv :: Numeric a => Tensor '[n, m] a -> Tensor '[m] a -> Tensor '[n] a
mv (Matrix m) (Vector v) = Vector $ m #> v

mm :: Numeric a => Tensor '[n, m] a -> Tensor '[m, k] a -> Tensor '[n, k] a
mm (Matrix m1) (Matrix m2) = Matrix $ m1 <> m2

gradDot :: Numeric a => Tensor '[n] a -> Tensor '[n] a -> a -> (Tensor '[n] a, Tensor '[n] a)
gradDot (Vector a) (Vector b) g = (Vector $ scale g b, Vector $ scale g a)

gradMV :: Numeric a => Tensor '[n, m] a -> Tensor '[m] a -> Tensor '[n] a -> (Tensor '[n, m] a, Tensor '[m] a)
gradMV m v g = (mm (asCol g) (asRow' v), mv (transpose m) g)

gradMM :: Numeric a => Tensor '[n, m] a -> Tensor '[m, k] a -> Tensor '[n, k] a -> (Tensor '[n, m] a, Tensor '[m, k] a)
gradMM a b g = (mm g (transpose b), mm (transpose a) g)

oneHot :: forall a n. (SingI n, IntegralN n, Usable a) => Int -> Tensor '[n] a
oneHot m = Vector $ (konst 0 (natVal' (sing::Sing n))) // [(m, 1)]

gradSelect :: forall a n. (SingI n, IntegralN n, Usable a) => Int -> Tensor '[n] a -> a -> Tensor '[n] a
gradSelect i _ a = Vector $ (konst 0 (natVal' (sing::Sing n))) // [(i, a)]


module School.Train.GradientDescent
( gradientDescent ) where

import Conduit ((.|), ConduitM, mapMC, runConduit, sinkNull, takeWhileC)
import Control.Monad.State.Lazy (get)
import Data.Maybe (isJust)
import Numeric.LinearAlgebra (Matrix)
import School.Train.AppTrain (AppTrain, runAppTrain)
import School.Train.GradientDescentPass (gradientDescentPass)
import School.Train.IterationHandler (IterationHandler(..))
import School.Train.UpdateParams (UpdateParams)
import School.Train.StoppingCondition (StoppingCondition(..))
import School.Train.TrainState (TrainState)
import School.Unit.CostFunction (CostFunction(..))
import School.Unit.Unit (Unit)
import School.Unit.UnitBackward (BackwardStack)

stopping :: StoppingCondition a
         -> ConduitM (BackwardStack a)
                     (Maybe ())
                     (AppTrain a)
                     ()
stopping condition = mapMC $ \_ -> do
  state <- get
  if runCondition condition state
    then return Nothing
    else return (Just ())

gradientDescent :: ConduitM () (Matrix a) (AppTrain a) ()
                -> [Unit a]
                -> CostFunction a (AppTrain a)
                -> UpdateParams a
                -> StoppingCondition a
                -> IterationHandler a
                -> TrainState a
                -> IO (Either String
                              (TrainState a))
gradientDescent source
                units
                cost
                update
                condition
                handler
                initState = do
  let pass = setupCost cost source
          .| gradientDescentPass units cost update
          .| runHandler handler
          .| stopping condition
          .| takeWhileC isJust
          .| sinkNull
  result <- runAppTrain initState . runConduit $ pass
  return $ snd <$> result


module HLearn.Metrics.Mahalanobis.Normal
    where

import Control.DeepSeq
import qualified Data.Vector.Generic as VG

import Foreign.Storable
import Numeric.LinearAlgebra hiding ((<>))
import qualified Numeric.LinearAlgebra as LA

import HLearn.Algebra
import HLearn.Metrics.Mahalanobis
import HLearn.Models.Distributions.Multivariate.MultiNormalFast

-------------------------------------------------------------------------------
-- data types

newtype MahalanobisParams dp = MahalanobisParams
    { normal :: MultiNormal dp
    }

deriving instance Show (MultiNormal dp) => Show (MahalanobisParams dp)
deriving instance Read (MultiNormal dp) => Read (MahalanobisParams dp)
deriving instance Eq (MultiNormal dp) => Eq (MahalanobisParams dp)
deriving instance Ord (MultiNormal dp) => Ord (MahalanobisParams dp)
deriving instance NFData (MultiNormal dp) => NFData (MahalanobisParams dp)

-------------------------------------------------------------------------------
-- algebra

deriving instance Monoid (MultiNormal dp) => Monoid (MahalanobisParams dp)

instance HasRing dp => HasRing (MahalanobisParams dp) where
    type Ring (MahalanobisParams dp) = Ring dp

-------------------------------------------------------------------------------
-- training

instance HomTrainer (MultiNormal dp) => HomTrainer (MahalanobisParams dp) where
    type Datapoint (MahalanobisParams dp) = Datapoint (MultiNormal dp)

    train1dp dp = MahalanobisParams $ train1dp dp
    train dps = MahalanobisParams $ train dps

-------------------------------------------------------------------------------
-- metric

instance
    ( VG.Vector dp r
    , Ring (dp r) ~ r
    , LA.Product r
    , Field r
    ) =>  MkMahalanobis (MahalanobisParams (dp r))
        where
    type MetricDatapoint (MahalanobisParams (dp r)) = dp r
    mkMahalanobis (MahalanobisParams m) dp = Mahalanobis
        { rawdp = dp
        , moddp = VG.fromList $ toList $ flatten $ inv (covar m) LA.<> asColumn v
        }
        where
            v = fromList $ VG.toList dp

-- mkMahalanobis ::
--     ( VG.Vector dp r
--     , Ring (dp r) ~ r
--     , LA.Product r
--     ) => MahalanobisParams (dp r) -> dp r -> Mahalanobis (dp r)
-- mkMahalanobis (MahalanobisParams m) dp = Mahalanobis
--     { rawdp = dp
--     , moddp = VG.fromList $ toList $ flatten $ covar m LA.<> asColumn v
--     }
--     where
--         v = fromList $ VG.toList dp


{-# LANGUAGE FlexibleContexts #-}
module SentimentAnalysis where

import Numeric.LinearAlgebra as LA
import Data.Bifunctor (Bifunctor(bimap))
import Data.List.Split (splitOn)

hingeLossSingle :: (Ord a, Numeric a) => Vector a -> a -> Vector a -> a -> a
hingeLossSingle x y theta theta0 = max 0 (1 - y * (theta LA.<.> x + theta0))

hingeLossFull :: (Numeric a, Ord a, Num (Vector a), Fractional a) => Matrix a -> Vector a -> Vector a -> a -> a
hingeLossFull x y theta theta0 = loss
    where
        theta0s = size y |> repeat theta0
        ones = size y |> repeat 1.0
        agreement = y * (theta LA.<# LA.tr' x + theta0s)
        losses = LA.cmap (max 0) (ones - agreement)
        loss = sumElements losses / fromIntegral (size losses)

perceptronSingleStepUpdate :: (Ord b, Numeric b, Num (Vector b)) => Vector b -> b -> Vector b -> b -> (Vector b, b)
perceptronSingleStepUpdate x y theta theta0 =
    if y * (theta LA.<.> x + theta0) <= 0
        then (theta + (scalar y * x), theta0 + y)
        else (theta, theta0)

perceptronUpdate :: (Ord t, Numeric t, Num (Vector t)) => [Vector t] -> [t] -> Vector t -> t -> [Int] -> (Vector t, t)
perceptronUpdate _ _ theta theta0 [] = (theta, theta0)
perceptronUpdate xs ys theta theta0 (z:zs) =
    let (theta', theta0') = perceptronSingleStepUpdate (xs !! z) (ys !! z) theta theta0
    in perceptronUpdate xs ys theta' theta0' zs

perceptron :: (Ord t, Numeric t, Num (Vector t)) => Matrix t -> Vector t -> Int -> [Int] -> (Vector t, t)
perceptron x y t indices = iterate step (cols x |> repeat 0, 0) !! t
    where step (theta, theta0) = perceptronUpdate (toRows x) (toList y) theta theta0 indices

averagePerceptronUpdate :: (Ord t, Numeric t, Num (Vector t))
    => [Vector t] -> [t] -> Vector t -> t -> [Int] -> [(Vector t, t)]
averagePerceptronUpdate _ _ theta theta0 [] = []
averagePerceptronUpdate xs ys theta theta0 (z:zs) =
    let (theta', theta0') = perceptronSingleStepUpdate (xs !! z) (ys !! z) theta theta0
    in (theta', theta0') : averagePerceptronUpdate xs ys theta' theta0' zs

averagePerceptronList :: (Ord t, Numeric t, Num (Vector t)) => Matrix t -> Vector t -> [Int] -> [[(Vector t, t)]]
averagePerceptronList x y indices = iterate step [(cols x |> repeat 0, 0)]
    where
        step xs = averagePerceptronUpdate (toRows x) (toList y) theta theta0 indices
            where (theta, theta0) = last xs

averagePerceptron :: (Fractional t, Ord t, Numeric t, Num (Vector t))
                        => Matrix t -> Vector t -> Int -> [Int] -> (Vector t, t)
averagePerceptron x y t indices = (bimap ((/ scalar n) . sum) ((/ n) . sum) . unzip . concat) xss
    where
        n = fromIntegral $ t * rows x
        xss = drop 1 $ take (t + 1) $ averagePerceptronList x y indices

pegasosSingleStepUpdate :: (Ord e, Numeric e, Num (Vector e)) => Vector e -> e -> e -> e -> Vector e -> e -> (Vector e, e)
pegasosSingleStepUpdate x y lambda eta theta theta0 =
    if y * (theta LA.<.> x + theta0) <= 1
        then
            (scalar (1 - eta * lambda) * theta + scalar (eta * y) * x, theta0 + eta * y)
        else
            (scalar (1 - eta * lambda) * theta, theta0)

pegasosUpdate :: (Ord t, Numeric t, Num (Vector t), Floating t)
                    => [Vector t] -> [t] -> t -> t -> Vector t -> t -> [Int] -> (Vector t, t, t)
pegasosUpdate _ _ lambda i theta theta0 [] = (theta, theta0, i)
pegasosUpdate xs ys lambda i theta theta0 (z:zs) =
    let (theta', theta0') = pegasosSingleStepUpdate (xs !! z) (ys !! z) lambda (1 / sqrt i) theta theta0
    in pegasosUpdate xs ys lambda (i + 1) theta' theta0' zs

pegasos :: (Ord t, Numeric t, Floating t, Num (Vector t)) => Matrix t -> Vector t -> Int -> t -> [Int] -> (Vector t, t)
pegasos x y t lambda indices = (theta', theta0')
    where
        (theta', theta0', i') = iterate step (cols x |> repeat 0, 0, 1.0) !! t
        step (theta, theta0, i) = pegasosUpdate (toRows x) (toList y) lambda i theta theta0 indices

main :: IO ()
main = do
    dta <- loadMatrix "toy_data.tsv"
    orderDta <- readFile "200.txt"

    let t = 10
    let lambda = 0.2
    let x = dta ¿ [1, 2]
    let y = flatten $ dta ¿ [0]
    let indices = map read $ splitOn "," orderDta

    let (theta, theta0) = perceptron x y t indices
    putStrLn "Perceptron:"
    print theta
    print theta0

    let (theta, theta0) = averagePerceptron x y t indices
    putStrLn "Average Perceptron:"
    print theta
    print theta0

    let (theta, theta0) = pegasos x y t lambda indices
    putStrLn "Pegasos:"
    print theta
    print theta0


module Mnist
    ( Image
    , Label
    , DataSet
    , loadMnist
    ) where

import Control.Monad
import Numeric.LinearAlgebra
import Network.HTTP.Simple (parseRequest, httpLBS, getResponseBody)
import System.Directory (doesFileExist)
import qualified Data.ByteString.Lazy as BL
import qualified Codec.Compression.GZip as GZ (compress, decompress)
import Data.Binary (encode, decode)

type Image = Matrix R
type Label = Vector R
type DataSet = (Image, Label)

baseUrl = "http://yann.lecun.com/exdb/mnist"
keyFiles = [
            ("train_img", "train-images-idx3-ubyte.gz"),
            ("train_label", "train-labels-idx1-ubyte.gz"),
            ("test_img", "t10k-images-idx3-ubyte.gz"),
            ("test_label", "t10k-labels-idx1-ubyte.gz")
          ]

assetsDir = "assets"
pickleFile = "mnist.dat"
imgSize = 784

generatePath :: String -> String
generatePath p = assetsDir ++ "/" ++ p

download :: String -> IO ()
download fn = do
    let p = generatePath fn

    e <- doesFileExist p
    unless e $ do
      putStrLn $ "Downloading " ++ fn ++ " ..."
      res <- httpLBS =<< parseRequest (baseUrl ++ "/" ++ fn)
      BL.writeFile p (getResponseBody res)
      putStrLn "Done"

downloadMnist :: [(String, String)] -> IO ()
downloadMnist [] = return ()
downloadMnist (x:xs) = do
    download $ snd x
    downloadMnist xs

toDoubleList :: BL.ByteString -> [Double]
toDoubleList = fmap (read . show . fromEnum) . BL.unpack

loadLabel :: String -> IO Label
loadLabel fn = do
    c <- fmap GZ.decompress (BL.readFile $ generatePath fn)
    return . vector . toDoubleList $ BL.drop 8 c

loadImg :: String -> IO Image
loadImg fn = do
    c <- fmap GZ.decompress (BL.readFile $ generatePath fn)
    return . matrix imgSize . toDoubleList $ BL.drop 16 c

convertDataset :: IO [DataSet]
convertDataset = do
    tri <- loadImg . snd . head $ keyFiles
    trl <- loadLabel . snd . (!!1) $ keyFiles
    ti <- loadImg . snd . (!!2) $ keyFiles
    tl <- loadLabel . snd . (!!3) $ keyFiles

    return [(tri, trl), (ti, tl)]

createPickle :: String -> [DataSet] -> IO ()
createPickle p ds = BL.writeFile p $ (GZ.compress . encode) ds

loadPickle :: String -> IO [DataSet]
loadPickle p = do
    eds <- BL.readFile p
    return $ (decode . GZ.decompress) eds

initMnist :: IO ()
initMnist = do
    downloadMnist keyFiles
    putStrLn "Creating binary DataSet file ..."
    createPickle (generatePath pickleFile) =<< convertDataset
    putStrLn "Done"

normalizeImg :: Bool -> [DataSet] -> IO [DataSet]
normalizeImg f ds@[tr, t]
    | f         = return [ ((/255) $ fst tr, snd tr), ((/255) $ fst t, snd t) ]
    | otherwise = return ds

loadMnist :: Bool -> IO [DataSet]
loadMnist nml = do
    let p = generatePath pickleFile
    e <- doesFileExist p
    unless e initMnist

    normalizeImg nml =<< loadPickle p