{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DerivingVia               #-}
{-# OPTIONS_GHC -fno-warn-deprecations #-}

-- |
-- Module      : Data.Functor.Invariant.Inplicative
-- Copyright   : (c) Justin Le 2021
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Contains the classes 'Inply' and 'Inplicative', the invariant
-- counterparts to 'Apply'/'Divise' and 'Applicative'/'Divisible'.
--
-- @since 0.4.0.0
module Data.Functor.Invariant.Inplicative (
  -- * Typeclass
    Inply(..)
  , Inplicative(..)
  -- * Deriving
  , WrappedApplicativeOnly(..)
  , WrappedDivisibleOnly(..)
  -- * Invariant 'Day'
  , runDay
  , dather
  , runDayApply
  , runDayDivise
  -- * Assembling Helpers
  , gatheredN
  , gatheredNMap
  , gatheredN1
  , gatheredN1Map
  , gatheredNRec
  , gatheredNMapRec
  , gatheredN1Rec
  , gatheredN1MapRec
  , gatherN
  , gatherN1
  ) where

import           Control.Applicative
import           Control.Applicative.Backwards               (Backwards(..))
import           Control.Applicative.Lift                    (Lift(Pure, Other))
import           Control.Arrow                               (Arrow)
import           Control.Monad.Trans.Cont                    (ContT)
import           Control.Monad.Trans.Error                   (ErrorT(..))
import           Control.Monad.Trans.Except                  (ExceptT(..))
import           Control.Monad.Trans.Identity                (IdentityT(..))
import           Control.Monad.Trans.List                    (ListT(..))
import           Control.Monad.Trans.Maybe                   (MaybeT(..))
import           Control.Monad.Trans.RWS                     (RWST(..))
import           Control.Monad.Trans.Reader                  (ReaderT(..))
import           Control.Monad.Trans.State                   (StateT)
import           Control.Monad.Trans.Writer                  (WriterT(..))
import           Control.Natural
import           Data.Complex                                (Complex)
import           Data.Deriving
import           Data.Functor.Apply
import           Data.Functor.Bind.Class                     (Bind)
import           Data.Functor.Constant                       (Constant)
import           Data.Functor.Contravariant
import           Data.Functor.Contravariant.Divise
import           Data.Functor.Contravariant.Divisible
import           Data.Functor.Identity
import           Data.Functor.Invariant
import           Data.Functor.Invariant.Day
import           Data.Functor.Product                        (Product(..))
import           Data.Functor.Reverse                        (Reverse(..))
import           Data.Hashable                               (Hashable)
import           Data.Kind
import           Data.List.NonEmpty                          (NonEmpty)
import           Data.SOP hiding                             (hmap)
import           Data.Sequence                               (Seq)
import           Data.StateVar                               (SettableStateVar)
import           Data.Tagged                                 (Tagged)
import           Data.Tree                                   (Tree)
import           GHC.Generics                                (Generic)
import qualified Control.Monad.Trans.RWS.Strict as Strict    (RWST(..))
import qualified Control.Monad.Trans.State.Strict as Strict  (StateT)
import qualified Control.Monad.Trans.Writer.Strict as Strict (WriterT(..))
import qualified Data.HashMap.Lazy                           as HM
import qualified Data.IntMap                                 as IM
import qualified Data.Map                                    as M
import qualified Data.Monoid                                 as Monoid
import qualified Data.Semigroup                              as Semigroup
import qualified Data.Sequence.NonEmpty                      as NESeq
import qualified Data.Vinyl                                  as V
import qualified Data.Vinyl.Curry                            as V
import qualified Data.Vinyl.Functor                          as V
import qualified GHC.Generics                                as Generics

-- | The invariant counterpart of 'Apply' and 'Divise'.
--
-- Conceptually you can think of 'Apply' as, given a way to "combine" @a@ and
-- @b@ to @c@, lets you merge @f a@ (producer of @a@) and @f b@ (producer
-- of @b@) into a @f c@ (producer of @c@).  'Divise' can be thought of as,
-- given a way to "split" a @c@ into an @a@ and a @b@, lets you merge @f
-- a@ (consumer of @a@) and @f b@ (consumder of @b@) into a @f c@ (consumer
-- of @c@).
--
-- 'Inply', for 'gather', requires both a combining function and
-- a splitting function in order to merge @f b@ (producer and consumer of
-- @b@) and @f c@ (producer and consumer of @c@) into a @f a@.  You can
-- think of it as, for the @f a@, it "splits" the a into @b@ and @c@ with
-- the @a -> (b, c)@, feeds it to the original @f b@ and @f c@, and then
-- re-combines the output back into a @a@ with the @b -> c -> a@.
--
-- @since 0.4.0.0
class Invariant f => Inply f where
    -- | Like '<.>', '<*>', 'divise', or 'divide', but requires both
    -- a splitting and a recombining function.  '<.>' and '<*>' require
    -- only a combining function, and 'divise' and 'divide' require only
    -- a splitting function.
    --
    -- It is used to merge @f b@ (producer and consumer of @b@) and @f c@
    -- (producer and consumer of @c@) into a @f a@.  You can think of it
    -- as, for the @f a@, it "splits" the a into @b@ and @c@ with the @a ->
    -- (b, c)@, feeds it to the original @f b@ and @f c@, and then
    -- re-combines the output back into a @a@ with the @b -> c -> a@.
    --
    -- An important property is that it will always use @both@ of the
    -- ccomponents given in order to fulfil its job.  If you gather an @f
    -- a@ and an @f b@ into an @f c@, in order to consume/produdce the @c@,
    -- it will always use both the @f a@ or the @f b@ -- exactly one of
    -- them.
    --
    -- @since 0.4.0.0
    gather
        :: (b -> c -> a)
        -> (a -> (b, c))
        -> f b
        -> f c
        -> f a
    gather f g x y = invmap (uncurry f) g (gathered x y)
    -- | A simplified version of 'gather' that combines into a tuple.  You
    -- can then use 'invmap' to reshape it into the proper shape.
    --
    -- @since 0.4.0.0
    gathered
        :: f a
        -> f b
        -> f (a, b)
    gathered = gather (,) id
    {-# MINIMAL gather | gathered #-}

-- | The invariant counterpart of 'Applicative' and 'Divisible'.
--
-- The main important action is described in 'Inply', but this adds 'knot',
-- which is the counterpart to 'pure' and 'conquer'.  It's the identity to
-- 'gather'; if combine two @f a@s with 'gather', and one of them is
-- 'knot', it will leave the structure unchanged.
--
-- Conceptually, if you think of 'gather' as "splitting and re-combining"
-- along multiple forks, then 'knot' introduces a fork that is never taken.
--
-- @since 0.4.0.0
class Inply f => Inplicative f where
    knot :: a -> f a

-- | Interpret out of a contravariant 'Day' into any instance of 'Inply' by
-- providing two interpreting functions.
--
-- This should go in "Data.Functor.Invariant.Day", but that module is in
-- a different package.
--
-- @since 0.4.0.0
runDay
    :: Inply h
    => (f ~> h)
    -> (g ~> h)
    -> Day f g ~> h
runDay f g (Day x y a b) = gather a b (f x) (g y)

-- | Squash the two items in a 'Day' using their natural 'Inply'
-- instances.
--
-- This should go in "Data.Functor.Invariant.Day", but that module is in
-- a different package.
--
-- @since 0.4.0.0
dather
    :: Inply f
    => Day f f ~> f
dather (Day x y a b) = gather a b x y

-- | Ignores the contravariant part of 'gather'
instance Apply f => Inply (WrappedFunctor f) where
    gather f _ (WrapFunctor x) (WrapFunctor y) = WrapFunctor (liftF2 f x y)
    gathered (WrapFunctor x) (WrapFunctor y) = WrapFunctor (liftF2 (,) x y)
-- | @'knot' = 'pure'@
instance (Applicative f, Apply f) => Inplicative (WrappedFunctor f) where
    knot = pure

-- | Ignores the covariant part of 'gather'
instance Divise f => Inply (WrappedContravariant f) where
    gather _ g (WrapContravariant x) (WrapContravariant y) = WrapContravariant (divise g x y)
    gathered (WrapContravariant x) (WrapContravariant y) = WrapContravariant (divised x y)
-- | @'knot' _ = 'conquer'@
instance (Divisible f, Divise f) => Inplicative (WrappedContravariant f) where
    knot _ = conquer

-- | Ignores the covariant part of 'gather'
instance Divise f => Inply (WrappedDivisible f) where
    gather _ g (WrapDivisible x) (WrapDivisible y) = WrapDivisible (divise g x y)
    gathered (WrapDivisible x) (WrapDivisible y) = WrapDivisible (divised x y)
-- | @'knot' _ = 'conquer'@
instance (Divisible f, Divise f) => Inplicative (WrappedDivisible f) where
    knot _ = conquer

-- | Wrap an 'Applicative' that is not necessarily an 'Apply'.
newtype WrappedApplicativeOnly f a =
    WrapApplicativeOnly { unwrapApplicativeOnly :: f a }
  deriving (Generic, Eq, Show, Ord, Read, Functor, Foldable, Traversable)
  deriving newtype (Applicative, Monad)

deriveShow1 ''WrappedApplicativeOnly
deriveRead1 ''WrappedApplicativeOnly
deriveEq1 ''WrappedApplicativeOnly
deriveOrd1 ''WrappedApplicativeOnly

instance Invariant f => Invariant (WrappedApplicativeOnly f) where
    invmap f g (WrapApplicativeOnly x) = WrapApplicativeOnly (invmap f g x)
instance (Applicative f, Invariant f) => Apply (WrappedApplicativeOnly f) where
    x <.> y = x <*> y
-- | Ignores the contravariant part of 'gather'
instance (Applicative f, Invariant f) => Inply (WrappedApplicativeOnly f) where
    gather f _ (WrapApplicativeOnly x) (WrapApplicativeOnly y) = WrapApplicativeOnly (liftA2 f x y)
    gathered (WrapApplicativeOnly x) (WrapApplicativeOnly y) = WrapApplicativeOnly (liftA2 (,) x y)
-- | @'knot' = 'pure'@
instance (Applicative f, Invariant f) => Inplicative (WrappedApplicativeOnly f) where
    knot = pure

-- | Wrap an 'Divisible' that is not necessarily a 'Divise'.
newtype WrappedDivisibleOnly f a =
    WrapDivisibleOnly { unwrapDivisibleOnly :: f a }
  deriving (Generic, Eq, Show, Ord, Read, Functor, Foldable, Traversable)
  deriving newtype (Divisible, Contravariant)

deriveShow1 ''WrappedDivisibleOnly
deriveRead1 ''WrappedDivisibleOnly
deriveEq1 ''WrappedDivisibleOnly
deriveOrd1 ''WrappedDivisibleOnly

instance Invariant f => Invariant (WrappedDivisibleOnly f) where
    invmap f g (WrapDivisibleOnly x) = WrapDivisibleOnly (invmap f g x)
instance (Divisible f, Invariant f) => Divise (WrappedDivisibleOnly f) where
    divise g (WrapDivisibleOnly x) (WrapDivisibleOnly y) = WrapDivisibleOnly (divide g x y)
-- | Ignores the covariant part of 'gather'
instance (Divisible f, Invariant f) => Inply (WrappedDivisibleOnly f) where
    gather _ g (WrapDivisibleOnly x) (WrapDivisibleOnly y) = WrapDivisibleOnly (divide g x y)
    gathered (WrapDivisibleOnly x) (WrapDivisibleOnly y) = WrapDivisibleOnly (divided x y)
-- | @'knot' _ = 'conquer'@
instance (Divisible f, Invariant f) => Inplicative (WrappedDivisibleOnly f) where
    knot _ = conquer

funzip :: Functor f => f (a, b) -> (f a, f b)
funzip x = (fmap fst x, fmap snd x)

-- | @since 0.4.1.0
instance Inply f => Inply (MaybeT f) where
    gather f g (MaybeT x) (MaybeT y) = MaybeT $
      gather (liftA2 f) (funzip . fmap g) x y
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (MaybeT f) where
    knot x = MaybeT (knot (Just x))

-- | @since 0.4.1.0
instance (Inply f, Semigroup w) => Inply (WriterT w f) where
    gather f g (WriterT x) (WriterT y) = WriterT $
      gather (\case (a, q) -> \case (b, r) -> (f a b, q <> r))
             (\case (a, s) -> case g a of (b, c) -> ((b, s), (c, s)))
             x y
-- | @since 0.4.1.0
instance (Inplicative f, Monoid w) => Inplicative (WriterT w f) where
    knot x = WriterT (knot (x, mempty))

-- | @since 0.4.1.0
instance (Inply f, Semigroup w) => Inply (Strict.WriterT w f) where
    gather f g (Strict.WriterT x) (Strict.WriterT y) = Strict.WriterT $
      gather (\(~(a, q)) (~(b, r)) -> (f a b, q <> r))
             (\(~(a, s)) -> let ~(b, c) = g a in ((b, s), (c, s)))
             x y
-- | @since 0.4.1.0
instance (Inplicative f, Monoid w) => Inplicative (Strict.WriterT w f) where
    knot x = Strict.WriterT (knot (x, mempty))

-- | @since 0.4.1.0
instance Inply f => Inply (ReaderT r f) where
    gather f g (ReaderT x) (ReaderT y) = ReaderT $ \r ->
      gather f g (x r) (y r)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (ReaderT r f) where
    knot x = ReaderT (\_ -> knot x)

-- | @since 0.4.1.0
instance Inply f => Inply (ExceptT e f) where
    gather f g (ExceptT x) (ExceptT y) = ExceptT $
      gather (liftA2 f) (funzip . fmap g) x y
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (ExceptT e f) where
    knot x = ExceptT (knot (Right x))

-- | @since 0.4.1.0
instance Inply f => Inply (ErrorT e f) where
    gather f g (ErrorT x) (ErrorT y) = ErrorT $
      gather (liftA2 f) (funzip . fmap g) x y
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (ErrorT e f) where
    knot x = ErrorT (knot (Right x))

-- | @since 0.4.1.0
instance Inply f => Inply (ListT f) where
    gather f g (ListT x) (ListT y) = ListT $
      gather (liftA2 f) (funzip . fmap g) x y
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (ListT f) where
    knot x = ListT (knot [x])

-- | @since 0.4.1.0
deriving via WrappedFunctor (RWST r w s m) instance (Bind m, Invariant m, Semigroup w) => Inply (RWST r w s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (RWST r w s m) instance (Monad m, Bind m, Invariant m, Monoid w) => Inplicative (RWST r w s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Strict.RWST r w s m) instance (Bind m, Invariant m, Semigroup w) => Inply (Strict.RWST r w s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Strict.RWST r w s m) instance (Monad m, Bind m, Invariant m, Monoid w) => Inplicative (Strict.RWST r w s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (StateT s m) instance (Bind m, Invariant m) => Inply (StateT s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (StateT s m) instance (Monad m, Bind m, Invariant m) => Inplicative (StateT s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Strict.StateT s m) instance (Bind m, Invariant m) => Inply (Strict.StateT s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Strict.StateT s m) instance (Monad m, Bind m, Invariant m) => Inplicative (Strict.StateT s m)

-- | @since 0.4.1.0
instance Inply f => Inply (Generics.M1 i t f :: Type -> Type) where
    gather f g (Generics.M1 x) (Generics.M1 y) = Generics.M1 (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Generics.M1 i t f :: Type -> Type) where
    knot = Generics.M1 . knot
-- | @since 0.4.1.0
instance (Inply f, Inply g) => Inply (f Generics.:*: g) where
    gather f g (x1 Generics.:*: y1) (x2 Generics.:*: y2) =
        gather f g x1 x2 Generics.:*: gather f g y1 y2
-- | @since 0.4.1.0
instance (Inplicative f, Inplicative g) => Inplicative (f Generics.:*: g) where
    knot x = knot x Generics.:*: knot x
-- | @since 0.4.1.0
instance (Inply f, Inply g) => Inply (Product f g) where
    gather f g (Pair x1 y1) (Pair x2 y2) =
      gather f g x1 x2 `Pair` gather f g y1 y2
-- | @since 0.4.1.0
instance (Inplicative f, Inplicative g) => Inplicative (Product f g) where
    knot x = knot x `Pair` knot x
-- | @since 0.4.1.0
instance Inply f => Inply (Generics.Rec1 f :: Type -> Type) where
    gather f g (Generics.Rec1 x) (Generics.Rec1 y) = Generics.Rec1 (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Generics.Rec1 f :: Type -> Type) where
    knot = Generics.Rec1 . knot
-- | @since 0.4.1.0
instance Inply f => Inply (Monoid.Alt f) where
    gather f g (Monoid.Alt x) (Monoid.Alt y) = Monoid.Alt (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Monoid.Alt f) where
    knot = Monoid.Alt . knot
-- | @since 0.4.1.0
instance Inply f => Inply (IdentityT f :: Type -> Type) where
    gather f g (IdentityT x) (IdentityT y) = IdentityT (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (IdentityT f :: Type -> Type) where
    knot = IdentityT . knot
-- | @since 0.4.1.0
instance Inply f => Inply (Reverse f :: Type -> Type) where
    gather f g (Reverse x) (Reverse y) = Reverse (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Reverse f :: Type -> Type) where
    knot = Reverse . knot
-- | @since 0.4.1.0
instance Inply f => Inply (Backwards f :: Type -> Type) where
    gather f g (Backwards x) (Backwards y) = Backwards (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Backwards f :: Type -> Type) where
    knot = Backwards . knot
-- | @since 0.4.1.0
instance Inply f => Inply (Lift f) where
    gather f g = \case
      Pure  x -> \case
        Pure  y -> Pure (f x y)
        Other y -> Other (invmap (f x) (snd . g) y)
      Other x -> \case
        Pure  y -> Other (invmap (`f` y) (fst . g) x)
        Other y -> Other (gather f g x y)
-- | @since 0.4.1.0
instance Inply f => Inplicative (Lift f) where
    knot = Pure

-- | @since 0.4.1.0
deriving via WrappedApplicativeOnly (Tagged a) instance Inply (Tagged a)
-- | @since 0.4.1.0
deriving via WrappedApplicativeOnly (Tagged a) instance Inplicative (Tagged a)

-- | @since 0.4.1.0
deriving via WrappedFunctor Identity instance Inply Identity
-- | @since 0.4.1.0
deriving via WrappedFunctor Identity instance Inplicative Identity
-- | @since 0.4.1.0
deriving via WrappedFunctor (Proxy :: Type -> Type) instance Inply Proxy
-- | @since 0.4.1.0
deriving via WrappedFunctor (Proxy :: Type -> Type) instance Inplicative Proxy
-- | @since 0.4.1.0
deriving via WrappedFunctor [] instance Inply []
-- | @since 0.4.1.0
deriving via WrappedFunctor [] instance Inplicative []
-- | @since 0.4.1.0
deriving via WrappedFunctor ((->) r) instance Inply ((->) r)
-- | @since 0.4.1.0
deriving via WrappedFunctor ((->) r) instance Inplicative ((->) r)
-- | @since 0.4.1.0
deriving via WrappedFunctor Maybe instance Inply Maybe
-- | @since 0.4.1.0
deriving via WrappedFunctor Maybe instance Inplicative Maybe
-- | @since 0.4.1.0
deriving via WrappedFunctor (Either e) instance Inply (Either e)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Either e) instance Inplicative (Either e)
-- | @since 0.4.1.0
deriving via WrappedFunctor IO instance Inply IO
-- | @since 0.4.1.0
deriving via WrappedFunctor IO instance Inplicative IO
-- | @since 0.4.1.0
deriving via WrappedFunctor Generics.Par1 instance Inply Generics.Par1
-- | @since 0.4.1.0
deriving via WrappedFunctor Generics.Par1 instance Inplicative Generics.Par1
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.U1 :: Type -> Type) instance Inply Generics.U1
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.U1 :: Type -> Type) instance Inplicative Generics.U1
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.K1 i c :: Type -> Type) instance Semigroup c => Inply (Generics.K1 i c)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.K1 i c :: Type -> Type) instance Monoid c => Inplicative (Generics.K1 i c)
-- | @since 0.4.1.0
deriving via WrappedFunctor Complex instance Inply Complex
-- | @since 0.4.1.0
deriving via WrappedFunctor Complex instance Inplicative Complex
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Min instance Inply Semigroup.Min
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Min instance Inplicative Semigroup.Min
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Max instance Inply Semigroup.Max
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Max instance Inplicative Semigroup.Max
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.First instance Inply Semigroup.First
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.First instance Inplicative Semigroup.First
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Last instance Inply Semigroup.Last
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Last instance Inplicative Semigroup.Last
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Option instance Inply Semigroup.Option
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Option instance Inplicative Semigroup.Option
-- | @since 0.4.1.0
deriving via WrappedFunctor ZipList instance Inply ZipList
-- | @since 0.4.1.0
deriving via WrappedFunctor ZipList instance Inplicative ZipList
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.First instance Inply Monoid.First
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.First instance Inplicative Monoid.First
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Last instance Inply Monoid.Last
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Last instance Inplicative Monoid.Last
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Dual instance Inply Monoid.Dual
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Dual instance Inplicative Monoid.Dual
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Sum instance Inply Monoid.Sum
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Sum instance Inplicative Monoid.Sum
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Product instance Inply Monoid.Product
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Product instance Inplicative Monoid.Product
-- | @since 0.4.1.0
deriving via WrappedFunctor NonEmpty instance Inply NonEmpty
-- | @since 0.4.1.0
deriving via WrappedFunctor NonEmpty instance Inplicative NonEmpty
-- | @since 0.4.1.0
deriving via WrappedFunctor Tree instance Inply Tree
-- | @since 0.4.1.0
deriving via WrappedFunctor Tree instance Inplicative Tree
-- | @since 0.4.1.0
deriving via WrappedFunctor Seq instance Inply Seq
-- | @since 0.4.1.0
deriving via WrappedFunctor Seq instance Inplicative Seq
-- | @since 0.4.1.0
deriving via WrappedFunctor NESeq.NESeq instance Inply NESeq.NESeq
-- | @since 0.4.1.0
deriving via WrappedFunctor (WrappedArrow a b) instance Arrow a => Inply (WrappedArrow a b)
-- | @since 0.4.1.0
deriving via WrappedFunctor (WrappedArrow a b) instance Arrow a => Inplicative (WrappedArrow a b)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.V1 :: Type -> Type) instance Inply Generics.V1
-- | @since 0.4.1.0
deriving via WrappedFunctor IM.IntMap instance Inply IM.IntMap
-- | @since 0.4.1.0
deriving via WrappedFunctor (M.Map k) instance Ord k => Inply (M.Map k)
-- | @since 0.4.1.0
deriving via WrappedFunctor (HM.HashMap k) instance (Hashable k, Eq k) => Inply (HM.HashMap k)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Const w :: Type -> Type) instance Semigroup w => Inply (Const w)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Const w :: Type -> Type) instance Monoid w => Inplicative (Const w)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Constant w :: Type -> Type) instance Semigroup w => Inply (Constant w)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Constant w :: Type -> Type) instance Monoid w => Inplicative (Constant w)
-- | @since 0.4.1.0
deriving via WrappedFunctor (ContT r (m :: Type -> Type)) instance Inply (ContT r m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (ContT r (m :: Type -> Type)) instance Inplicative (ContT r m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (WrappedMonad m) instance Monad m => Inply (WrappedMonad m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (WrappedMonad m) instance Monad m => Inplicative (WrappedMonad m)
-- | @since 0.4.1.0
deriving via WrappedFunctor ((,) w :: Type -> Type) instance Semigroup w => Inply ((,) w)
-- | @since 0.4.1.0
deriving via WrappedFunctor ((,) w :: Type -> Type) instance Monoid w => Inplicative ((,) w)

-- | @since 0.4.1.0
deriving via WrappedDivisible SettableStateVar instance Inply SettableStateVar
-- | @since 0.4.1.0
deriving via WrappedDivisible SettableStateVar instance Inplicative SettableStateVar
-- | @since 0.4.1.0
deriving via WrappedDivisible Predicate instance Inply Predicate
-- | @since 0.4.1.0
deriving via WrappedDivisible Predicate instance Inplicative Predicate
-- | @since 0.4.1.0
deriving via WrappedDivisible Comparison instance Inply Comparison
-- | @since 0.4.1.0
deriving via WrappedDivisible Comparison instance Inplicative Comparison
-- | @since 0.4.1.0
deriving via WrappedDivisible Equivalence instance Inply Equivalence
-- | @since 0.4.1.0
deriving via WrappedDivisible Equivalence instance Inplicative Equivalence
-- | @since 0.4.1.0
deriving via WrappedDivisible (Op r) instance Semigroup r => Inply (Op r)
-- | @since 0.4.1.0
deriving via WrappedDivisible (Op r) instance Monoid r => Inplicative (Op r)

-- | Convenient wrapper to build up an 'Inplicative' instance by providing
-- each component of it.  This makes it much easier to build up longer
-- chains because you would only need to write the splitting/joining
-- functions in one place.
--
-- For example, if you had a data type
--
-- @
-- data MyType = MT Int Bool String
-- @
--
-- and an invariant functor and 'Inplicative' instance @Prim@
-- (representing, say, a bidirectional parser, where @Prim Int@ is
-- a bidirectional parser for an 'Int'@), then you could assemble
-- a bidirectional parser for a @MyType@ using:
--
-- @
-- invmap (\(MyType x y z) -> I x :* I y :* I z :* Nil)
--        (\(I x :* I y :* I z :* Nil) -> MyType x y z) $
--   gatheredN $ intPrim
--                    :* boolPrim
--                    :* stringPrim
--                    :* Nil
-- @
--
-- Some notes on usefulness depending on how many components you have:
--
-- *    If you have 0 components, use 'knot' directly.
-- *    If you have 1 component, you don't need anything.
-- *    If you have 2 components, use 'gather' directly.
-- *    If you have 3 or more components, these combinators may be useful;
--      otherwise you'd need to manually peel off tuples one-by-one.
--
-- @since 0.4.1.0
gatheredN
    :: Inplicative f
    => NP f as
    -> f (NP I as)
gatheredN = \case
    Nil     -> knot Nil
    x :* xs -> gather
      (\y ys -> I y :* ys)
      (\case I y :* ys -> (y, ys))
      x
      (gatheredN xs)

-- | Given a function to "break out" a data type into a 'NP' (tuple) and one to
-- put it back together from the tuple, 'gather' all of the components
-- together.
--
-- For example, if you had a data type
--
-- @
-- data MyType = MT Int Bool String
-- @
--
-- and an invariant functor and 'Inplicative' instance @Prim@
-- (representing, say, a bidirectional parser, where @Prim Int@ is
-- a bidirectional parser for an 'Int'@), then you could assemble
-- a bidirectional parser for a @MyType@ using:
--
-- @
--   concaMapInplicative
--      (\(MyType x y z) -> I x :* I y :* I z :* Nil)
--      (\(I x :* I y :* I z :* Nil) -> MyType x y z)
--      $ intPrim
--     :* boolPrim
--     :* stringPrim
--     :* Nil
-- @
--
-- See notes on 'gatheredNMap' for more details and caveats.
--
-- @since 0.4.1.0
gatheredNMap
    :: Inplicative f
    => (NP I as -> b)
    -> (b -> NP I as)
    -> NP f as
    -> f b
gatheredNMap f g = invmap f g . gatheredN

-- | A version of 'gatheredN' for non-empty 'NP', but only
-- requiring an 'Inply' instance.
--
-- @since 0.4.1.0
gatheredN1
    :: Inply f
    => NP f (a ': as)
    -> f (NP I (a ': as))
gatheredN1 (x :* xs) = case xs of
    Nil    -> invmap ((:* Nil) . I) (\case I y :* _ -> y) x
    _ :* _ -> gather
      (\y ys -> I y :* ys)
      (\case I y :* ys -> (y, ys))
      x
      (gatheredN1 xs)

-- | A version of 'gatheredNMap' for non-empty 'NP', but only
-- requiring an 'Inply' instance.
--
-- @since 0.4.1.0
gatheredN1Map
    :: Inplicative f
    => (NP I (a ': as) -> b)
    -> (b -> NP I (a ': as))
    -> NP f (a ': as)
    -> f b
gatheredN1Map f g = invmap f g . gatheredN1

-- | A version of 'gatheredN' using 'V.XRec' from /vinyl/ instead of
-- 'NP' from /sop-core/.  This can be more convenient because it doesn't
-- require manual unwrapping/wrapping of tuple components.
--
-- @since 0.4.1.0
gatheredNRec
    :: Inplicative f
    => V.Rec f as
    -> f (V.XRec V.Identity as)
gatheredNRec = \case
    V.RNil    -> knot V.RNil
    x V.:& xs -> gather
      (V.::&)
      (\case y V.::& ys -> (y, ys))
      x
      (gatheredNRec xs)

-- | A version of 'gatheredNMap' using 'V.XRec' from /vinyl/ instead of
-- 'NP' from /sop-core/.  This can be more convenient because it doesn't
-- require manual unwrapping/wrapping of tuple components.
--
-- @since 0.4.1.0
gatheredNMapRec
    :: Inplicative f
    => (V.XRec V.Identity as -> b)
    -> (b -> V.XRec V.Identity as)
    -> V.Rec f as
    -> f b
gatheredNMapRec f g = invmap f g . gatheredNRec

-- | Convenient wrapper to 'gather' over multiple arguments using tine
-- vinyl library's multi-arity uncurrying facilities.  Makes it a lot more
-- convenient than using 'gather' multiple times and needing to accumulate
-- intermediate types.
--
-- For example, if you had a data type
--
-- @
-- data MyType = MT Int Bool String
-- @
--
-- and an invariant functor and 'Inplicative' instance @Prim@
-- (representing, say, a bidirectional parser, where @Prim Int@ is
-- a bidirectional parser for an 'Int'@), then you could assemble
-- a bidirectional parser for a @MyType@ using:
--
-- @
-- 'gatherN'
--   MT                                         -- ^ curried assembling function
--   (\(MT x y z) -> x ::& y ::& z ::& XRNil)   -- ^ disassembling function
--   (intPrim :: Prim Int)
--   (boolPrim :: Prim Bool)
--   (stringPrim :: Prim String)
-- @
--
-- Really only useful with 3 or more arguments, since with two arguments
-- this is just 'gather' (and with zero arguments, you can just use
-- 'knot').
--
-- The generic type is a bit tricky to understand, but it's easier to
-- understand what's going on if you instantiate with concrete types:
--
-- @
-- ghci> :t gatherN @MyInplicative @'[Int, Bool, String]
--      (Int -> Bool -> String -> b)
--   -> (b -> XRec Identity '[Int, Bool, String])
--   -> MyInplicative Int
--   -> MyInplicative Bool
--   -> MyInplicative String
--   -> MyInplicative b
-- @
--
-- @since 0.4.1.0
gatherN
    :: forall f as b. (Inplicative f, V.IsoXRec V.Identity as, V.RecordCurry as)
    => V.Curried as b
    -> (b -> V.XRec V.Identity as)
    -> V.CurriedF f as (f b)
gatherN f g = V.rcurry @as @f $
    invmap (V.runcurry' f . V.fromXRec) g
  . gatheredNRec

-- | A version of 'gatheredN1' using 'V.XRec' from /vinyl/ instead of
-- 'NP' from /sop-core/.  This can be more convenient because it doesn't
-- require manual unwrapping/wrapping of components.
--
-- @since 0.4.1.0
gatheredN1Rec
    :: Inply f
    => V.Rec f (a ': as)
    -> f (V.XRec V.Identity (a ': as))
gatheredN1Rec (x V.:& xs) = case xs of
    V.RNil   -> invmap (V.::& V.RNil) (\case z V.::& _ -> z) x
    _ V.:& _ -> gather
      (V.::&)
      (\case y V.::& ys -> (y, ys))
      x
      (gatheredN1Rec xs)

-- | A version of 'gatheredNMap' using 'V.XRec' from /vinyl/ instead of
-- 'NP' from /sop-core/.  This can be more convenient because it doesn't
-- require manual unwrapping/wrapping of tuple components.
--
-- @since 0.4.1.0
gatheredN1MapRec
    :: Inplicative f
    => (V.XRec V.Identity (a ': as) -> b)
    -> (b -> V.XRec V.Identity (a ': as))
    -> V.Rec f (a ': as)
    -> f b
gatheredN1MapRec f g = invmap f g . gatheredN1Rec

-- | 'gatherN' but with at least one argument, so can be used with any
-- 'Inply'.
--
-- @since 0.4.1.0
gatherN1
    :: forall f a as b. (Inply f, V.IsoXRec V.Identity as, V.RecordCurry as)
    => V.Curried (a ': as) b
    -> (b -> V.XRec V.Identity (a ': as))
    -> V.CurriedF f (a ': as) (f b)
gatherN1 f g = V.rcurry @(a ': as) @f $
    invmap (V.runcurry' f . V.fromXRec) g
  . gatheredN1Rec

-- | Interpret out of a contravariant 'Day' into any instance of 'Apply' by
-- providing two interpreting functions.
--
-- In theory, this should not need to exist, since you should always be
-- able to use 'runDay' because every instance of 'Apply' is also an
-- instance of 'Inply'.  However, this can be handy if you are using an
-- instance of 'Apply' that has no 'Inply' instance.  Consider also
-- 'unsafeInplyCo' if you are using a specific, concrete type for @h@.
runDayApply
    :: forall f g h. Apply h
    => f ~> h
    -> g ~> h
    -> Day f g ~> h
runDayApply f g (Day x y j _) = liftF2 j (f x) (g y)

-- | Interpret out of a contravariant 'Day' into any instance of 'Divise'
-- by providing two interpreting functions.
--
-- In theory, this should not need to exist, since you should always be
-- able to use 'runDay' because every instance of 'Divise' is also an
-- instance of 'Inply'.  However, this can be handy if you are using an
-- instance of 'Divise' that has no 'Inply' instance.  Consider also
-- 'unsafeInplyContra' if you are using a specific, concrete type for @h@.
runDayDivise
    :: forall f g h. Divise h
    => f ~> h
    -> g ~> h
    -> Day f g ~> h
runDayDivise f g (Day x y _ h) = divise h (f x) (g y)


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeOperators       #-}
module Test.Grenade.Network where

import           Control.Monad                   (guard)
import           Control.Monad.ST                (runST)
import           Data.Constraint
import           Data.Proxy
import           Data.Singletons.Prelude.List
import           Data.Singletons.TypeLits
import           Data.Typeable
import qualified Data.Vector.Storable            as VS
import qualified Data.Vector.Storable.Mutable    as VS (write)
import           Hedgehog
import qualified Hedgehog.Gen                    as Gen
import           Hedgehog.Internal.Property      (failWith)
import           Hedgehog.Internal.Source

import           Grenade

#if MIN_VERSION_base(4,11,0)
import           GHC.TypeLits                    hiding (natVal)
#else
import           GHC.TypeLits
#endif
#if MIN_VERSION_base(4,9,0)
import           Data.Kind                       (Type)
#endif
import           Data.Singletons
import           Data.Singletons.Prelude.Num     ((%*))
import           GHC.TypeLits                    (someNatVal)

import           Test.Grenade.Layers.Convolution
import           Test.Hedgehog.Compat
import           Test.Hedgehog.Hmatrix
import           Test.Hedgehog.TypeLits

import           Numeric.LinearAlgebra           (flatten)
import           Numeric.LinearAlgebra.Static    (extract, norm_Inf)
import           Unsafe.Coerce

data SomeNetwork :: Type where
    SomeNetwork :: ( SingI shapes, SingI (Head shapes), SingI (Last shapes), Show (Network layers shapes) ) => Network layers shapes -> SomeNetwork

instance Show SomeNetwork where
  show (SomeNetwork net) = show net

-- | Generate a random network of a random type
--
-- This is slightly insane for a few reasons. Everything must be wrapped up
-- in a SomeNetwork.
genNetwork :: Gen SomeNetwork
genNetwork =
  Gen.recursive Gen.choice [
      do SomeSing ( r :: Sing final ) <- genShape
         withSingI r $
           pure (SomeNetwork (NNil :: Network '[] '[ final ] ))
    ] [
      do SomeNetwork ( rest :: Network layers shapes ) <- genNetwork
         case ( sing :: Sing shapes ) of
           SNil -> Gen.discard -- Can't occur
           SCons ( h :: Sing h ) ( _ :: Sing hs ) ->
             withSingI h $
              case h of
                 D1Sing l@SNat ->
                   Gen.choice [
                     pure (SomeNetwork (Tanh    :~> rest :: Network ( Tanh    ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Logit   :~> rest :: Network ( Logit   ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Relu    :~> rest :: Network ( Relu    ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Elu     :~> rest :: Network ( Elu     ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Softmax :~> rest :: Network ( Softmax ': layers ) ( h ': h ': hs )))
                   , do -- Reshape to two dimensions
                        let divisors :: Integer -> [Integer]
                            divisors n = 1 : [x | x <- [2..(n-1)], n `rem` x == 0]
                        let len = fromIntegral $ natVal l
                        rs  <- Gen.element $ divisors len
                        let cs  = len `quot` rs
                        case ( someNatVal rs, someNatVal cs, someNatVal len ) of
                          ( Just (SomeNat (rs' :: Proxy inRows)), Just (SomeNat (cs' :: Proxy inCols)), Just (SomeNat (_ :: Proxy outLen ) )) ->
                            let p1 = singByProxy rs'
                                p2 = singByProxy cs'
                            in  case ( p1 %* p2, unsafeCoerce (Dict :: Dict ()) :: Dict ((inRows * inCols) ~ outLen), unsafeCoerce (Dict :: Dict ()) :: Dict (( 'D1 outLen ) ~ h )) of
                                  ( SNat, Dict, Dict ) ->
                                    pure (SomeNetwork (Reshape :~> rest :: Network ( Reshape ': layers ) ( ('D2 inRows inCols) ': h ': hs )))
                          _ -> Gen.discard -- Doesn't occur
                   ]
                 D2Sing r@SNat c@SNat ->
                   Gen.choice [
                     pure (SomeNetwork (Tanh    :~> rest :: Network ( Tanh    ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Logit   :~> rest :: Network ( Logit   ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Relu    :~> rest :: Network ( Relu    ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Elu     :~> rest :: Network ( Elu     ': layers ) ( h ': h ': hs )))
                   , do -- Build a convolution layer with one filter output
                        -- Figure out some kernel sizes which work for this layer
                        -- There must be a better way than this...
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c

                        let ok extent kernel = [stride | stride <- [ 1 .. extent ], (extent - kernel) `mod` stride == 0]

                        -- Get some kernels which will fit
                        kernel_r <- choose 1 output_r
                        kernel_c <- choose 1 output_c

                        -- Build up some strides which also fit
                        stride_r <- Gen.element $ ok output_r kernel_r
                        stride_c <- Gen.element $ ok output_c kernel_c

                        -- Determine the input size
                        let input_r = (output_r - 1) * stride_r + kernel_r
                        let input_c = (output_c - 1) * stride_c + kernel_c

                        guard (input_r < 100)
                        guard (input_c < 100)

                        -- Remake types for input
                        case ( someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal kernel_r, someNatVal kernel_c, someNatVal stride_r, someNatVal stride_c ) of
                          ( Just (SomeNat (_    :: Proxy inRows)),     Just (SomeNat  (_    :: Proxy inCols)),
                            Just (SomeNat (_    :: Proxy outRows)),    Just (SomeNat  (_    :: Proxy outCols)),
                            Just (SomeNat (pkr  :: Proxy kernelRows)), Just (SomeNat  (pkc  :: Proxy kernelCols)),
                            Just (SomeNat (_    :: Proxy strideRows)), Just (SomeNat  (_    :: Proxy strideCols))) ->
                              let p1 = singByProxy pkr
                                  p2 = singByProxy pkc
                              in  case ( p1 %* p2
                                        -- Fake it till you make it.
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D2 outRows outCols ) ~ h ))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outRows - 1) * strideRows) ~ (inRows - kernelRows)))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outCols - 1) * strideCols) ~ (inCols - kernelCols)))) of
                                    (SNat, Dict, Dict, Dict) -> do
                                        conv <- genConvolution
                                        pure (SomeNetwork (conv :~> rest :: Network ( Convolution 1 1 kernelRows kernelCols strideRows strideCols ': layers ) ( ('D2 inRows inCols) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                   , do -- Build a convolution layer with one filter output
                        -- Figure out some kernel sizes which work for this layer
                        -- There must be a better way than this...
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c

                        let ok extent kernel = [stride | stride <- [ 1 .. extent ], (extent - kernel) `mod` stride == 0]

                        -- Get some kernels which will fit
                        kernel_r <- choose 1 output_r
                        kernel_c <- choose 1 output_c
                        channels <- choose 1 10

                        -- Build up some strides which also fit
                        stride_r <- Gen.element $ ok output_r kernel_r
                        stride_c <- Gen.element $ ok output_c kernel_c

                        -- Determine the input size
                        let input_r = (output_r - 1) * stride_r + kernel_r
                        let input_c = (output_c - 1) * stride_c + kernel_c

                        guard (input_r < 100)
                        guard (input_c < 100)

                        -- Remake types for input
                        case ( someNatVal channels, someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal kernel_r, someNatVal kernel_c, someNatVal stride_r, someNatVal stride_c ) of
                          ( Just (SomeNat (chan :: Proxy channels)),
                            Just (SomeNat (pinr :: Proxy inRows)),     Just (SomeNat  (_    :: Proxy inCols)),
                            Just (SomeNat (_    :: Proxy outRows)),    Just (SomeNat  (_    :: Proxy outCols)),
                            Just (SomeNat (pkr  :: Proxy kernelRows)), Just (SomeNat  (pkc  :: Proxy kernelCols)),
                            Just (SomeNat (_    :: Proxy strideRows)), Just (SomeNat  (_    :: Proxy strideCols))) ->
                              let p1 = singByProxy pkr
                                  p2 = singByProxy pkc
                                  p3 = singByProxy chan
                              in  case ( p1 %* p2 %* p3
                                        , singByProxy pinr %* singByProxy chan
                                        -- Fake it till you make it.
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D2 outRows outCols ) ~ h ))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outRows - 1) * strideRows) ~ (inRows - kernelRows)))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outCols - 1) * strideCols) ~ (inCols - kernelCols)))) of
                                    (SNat, SNat, Dict, Dict, Dict) -> do
                                        conv <- genConvolution
                                        pure (SomeNetwork (conv :~> rest :: Network ( Convolution channels 1 kernelRows kernelCols strideRows strideCols ': layers ) ( ('D3 inRows inCols channels) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                   , do -- Build a Pooling layer
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c

                        let ok extent kernel = [stride | stride <- [ 1 .. extent ], (extent - kernel) `mod` stride == 0]

                        -- Get some kernels which will fit
                        kernel_r <- choose 1 output_r
                        kernel_c <- choose 1 output_c

                        -- Build up some strides which also fit
                        stride_r <- Gen.element $ ok output_r kernel_r
                        stride_c <- Gen.element $ ok output_c kernel_c

                        -- Determine the input size
                        let input_r = (output_r - 1) * stride_r + kernel_r
                        let input_c = (output_c - 1) * stride_c + kernel_c

                        guard (input_r < 100)
                        guard (input_c < 100)

                        -- Remake types for input
                        case ( someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal kernel_r, someNatVal kernel_c, someNatVal stride_r, someNatVal stride_c ) of
                          ( Just (SomeNat (_    :: Proxy inRows)),     Just (SomeNat  (_    :: Proxy inCols)),
                            Just (SomeNat (_    :: Proxy outRows)),    Just (SomeNat  (_    :: Proxy outCols)),
                            Just (SomeNat (pkr  :: Proxy kernelRows)), Just (SomeNat  (pkc  :: Proxy kernelCols)),
                            Just (SomeNat (_    :: Proxy strideRows)), Just (SomeNat  (_    :: Proxy strideCols))) ->
                              let p1 = singByProxy pkr
                                  p2 = singByProxy pkc
                              in  case ( p1 %* p2
                                        -- Fake it till you make it.
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D2 outRows outCols ) ~ h ))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outRows - 1) * strideRows) ~ (inRows - kernelRows)))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outCols - 1) * strideCols) ~ (inCols - kernelCols)))) of
                                    (SNat, Dict, Dict, Dict) ->
                                        pure (SomeNetwork (Pooling :~> rest :: Network ( Pooling kernelRows kernelCols strideRows strideCols ': layers ) ( ('D2 inRows inCols) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                   , do -- Build a Pad layer
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c

                        pad_left   <- choose 0 (output_r - 1)
                        pad_right  <- choose 0 (output_r - 1 - pad_left)
                        pad_top    <- choose 0 (output_c - 1)
                        pad_bottom <- choose 0 (output_c - 1 - pad_top)

                        let input_r = output_r - pad_top - pad_bottom
                        let input_c = output_c - pad_left - pad_right

                        guard (input_r > 1)
                        guard (input_c > 1)

                        -- Remake types for input
                        case ( someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal pad_left, someNatVal pad_right, someNatVal pad_top, someNatVal pad_bottom ) of
                          ( Just (SomeNat (_    :: Proxy inputRows)),  Just (SomeNat  (_    :: Proxy inputColumns)),
                            Just (SomeNat (_    :: Proxy outputRows)), Just (SomeNat  (_    :: Proxy outputColumns)),
                            Just (SomeNat (_    :: Proxy padLeft)),    Just (SomeNat  (_    :: Proxy padRight)),
                            Just (SomeNat (_    :: Proxy padTop)),     Just (SomeNat  (_    :: Proxy padBottom))) ->
                              case (  (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D2 outputRows outputColumns) ~ h ))
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (inputRows + padTop + padBottom) ~ outputRows) )
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (inputColumns + padLeft + padRight) ~ outputColumns ))) of
                                    ( Dict, Dict, Dict ) ->
                                        pure (SomeNetwork (Pad :~> rest :: Network ( Pad padLeft padTop padRight padBottom ': layers ) ( ('D2 inputRows inputColumns) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                   , do -- Build a Crop layer
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c

                        crop_left   <- choose 0 10
                        crop_right  <- choose 0 10
                        crop_top    <- choose 0 10
                        crop_bottom <- choose 0 10

                        let input_r = output_r + crop_top + crop_bottom
                        let input_c = output_c + crop_left + crop_right

                        -- Remake types for input
                        case ( someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal crop_left, someNatVal crop_right, someNatVal crop_top, someNatVal crop_bottom ) of
                          ( Just (SomeNat (_    :: Proxy inputRows)),  Just (SomeNat  (_    :: Proxy inputColumns)),
                            Just (SomeNat (_    :: Proxy outputRows)), Just (SomeNat  (_    :: Proxy outputColumns)),
                            Just (SomeNat (_    :: Proxy cropLeft)),   Just (SomeNat  (_    :: Proxy cropRight)),
                            Just (SomeNat (_    :: Proxy cropTop)),    Just (SomeNat  (_    :: Proxy cropBottom))) ->
                              case (  (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D2 outputRows outputColumns) ~ h ))
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (outputRows + cropTop + cropBottom) ~ inputRows ) )
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (outputColumns + cropLeft + cropRight) ~ inputColumns ))) of
                                    (Dict, Dict, Dict) ->
                                        pure (SomeNetwork (Crop :~> rest :: Network ( Crop cropLeft cropTop cropRight cropBottom ': layers ) ( ('D2 inputRows inputColumns) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                   ]
                 D3Sing r@SNat c@SNat f@SNat ->
                   Gen.choice [
                     pure (SomeNetwork (Tanh    :~> rest :: Network ( Tanh    ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Logit   :~> rest :: Network ( Logit   ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Relu    :~> rest :: Network ( Relu    ': layers ) ( h ': h ': hs )))
                   , pure (SomeNetwork (Elu     :~> rest :: Network ( Elu     ': layers ) ( h ': h ': hs )))
                   , do -- Build a convolution layer with one filter output
                        -- Figure out some kernel sizes which work for this layer
                        -- There must be a better way than this...
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c
                        let output_f = fromIntegral $ natVal f

                        let ok extent kernel = [stride | stride <- [ 1 .. extent ], (extent - kernel) `mod` stride == 0]

                        -- Get some kernels which will fit
                        kernel_r <- choose 1 output_r
                        kernel_c <- choose 1 output_c
                        channels <- choose 1 10

                        -- Build up some strides which also fit
                        stride_r <- Gen.element $ ok output_r kernel_r
                        stride_c <- Gen.element $ ok output_c kernel_c

                        -- Determine the input size
                        let input_r = (output_r - 1) * stride_r + kernel_r
                        let input_c = (output_c - 1) * stride_c + kernel_c

                        guard (input_r < 100)
                        guard (input_c < 100)

                        -- Remake types for input
                        case ( someNatVal channels, someNatVal output_f, someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal kernel_r, someNatVal kernel_c, someNatVal stride_r, someNatVal stride_c ) of
                          ( Just (SomeNat (chan :: Proxy channels)),   Just (SomeNat  (_    :: Proxy filters)),
                            Just (SomeNat (pinr :: Proxy inRows)),     Just (SomeNat  (_    :: Proxy inCols)),
                            Just (SomeNat (_    :: Proxy outRows)),    Just (SomeNat  (_    :: Proxy outCols)),
                            Just (SomeNat (pkr  :: Proxy kernelRows)), Just (SomeNat  (pkc  :: Proxy kernelCols)),
                            Just (SomeNat (_    :: Proxy strideRows)), Just (SomeNat  (_    :: Proxy strideCols))) ->
                              let p1 = singByProxy pkr
                                  p2 = singByProxy pkc
                                  p3 = singByProxy chan
                              in  case ( p1 %* p2 %* p3
                                        , singByProxy pinr %* singByProxy chan
                                        -- Fake it till you make it.
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D3 outRows outCols filters) ~ h ))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outRows - 1) * strideRows) ~ (inRows - kernelRows)))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outCols - 1) * strideCols) ~ (inCols - kernelCols)))) of
                                    (SNat, SNat, Dict, Dict, Dict) -> do
                                        conv <- genConvolution
                                        pure (SomeNetwork (conv :~> rest :: Network ( Convolution channels filters kernelRows kernelCols strideRows strideCols ': layers ) ( ('D3 inRows inCols channels) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                   , do -- Build a Pooling layer
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c
                        let output_f = fromIntegral $ natVal f

                        let ok extent kernel = [stride | stride <- [ 1 .. extent ], (extent - kernel) `mod` stride == 0]

                        -- Get some kernels which will fit
                        kernel_r <- choose 1 output_r
                        kernel_c <- choose 1 output_c

                        -- Build up some strides which also fit
                        stride_r <- Gen.element $ ok output_r kernel_r
                        stride_c <- Gen.element $ ok output_c kernel_c

                        -- Determine the input size
                        let input_r = (output_r - 1) * stride_r + kernel_r
                        let input_c = (output_c - 1) * stride_c + kernel_c

                        guard (input_r < 100)
                        guard (input_c < 100)

                        -- Remake types for input
                        case ( someNatVal output_f, someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal kernel_r, someNatVal kernel_c, someNatVal stride_r, someNatVal stride_c ) of
                          ( Just (SomeNat (chan :: Proxy filters)),
                            Just (SomeNat (pinr :: Proxy inRows)),     Just (SomeNat  (_    :: Proxy inCols)),
                            Just (SomeNat (_    :: Proxy outRows)),    Just (SomeNat  (_    :: Proxy outCols)),
                            Just (SomeNat (pkr  :: Proxy kernelRows)), Just (SomeNat  (pkc  :: Proxy kernelCols)),
                            Just (SomeNat (_    :: Proxy strideRows)), Just (SomeNat  (_    :: Proxy strideCols))) ->
                              let p1 = singByProxy pkr
                                  p2 = singByProxy pkc
                                  p3 = singByProxy chan
                              in  case ( p1 %* p2 %* p3
                                        , singByProxy pinr %* singByProxy chan
                                        -- Fake it till you make it.
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D3 outRows outCols filters) ~ h ))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outRows - 1) * strideRows) ~ (inRows - kernelRows)))
                                        , (unsafeCoerce (Dict :: Dict ()) :: Dict (((outCols - 1) * strideCols) ~ (inCols - kernelCols)))) of
                                    (SNat, SNat, Dict, Dict, Dict) ->
                                        pure (SomeNetwork (Pooling :~> rest :: Network ( Pooling kernelRows kernelCols strideRows strideCols ': layers ) ( ('D3 inRows inCols filters) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                   , do -- Build a Pad layer
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c
                        let output_f = fromIntegral $ natVal f

                        pad_left   <- choose 0 (output_r - 1)
                        pad_right  <- choose 0 (output_r - 1 - pad_left)
                        pad_top    <- choose 0 (output_c - 1)
                        pad_bottom <- choose 0 (output_c - 1 - pad_top)

                        let input_r = output_r - pad_top - pad_bottom
                        let input_c = output_c - pad_left - pad_right

                        guard (input_r > 1)
                        guard (input_c > 1)

                        -- Remake types for input
                        case ( someNatVal output_f, someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal pad_left, someNatVal pad_right, someNatVal pad_top, someNatVal pad_bottom ) of
                          ( Just (SomeNat (chan :: Proxy filters)),
                            Just (SomeNat (pinr :: Proxy inputRows)),  Just (SomeNat  (_    :: Proxy inputColumns)),
                            Just (SomeNat (_    :: Proxy outputRows)), Just (SomeNat  (_    :: Proxy outputColumns)),
                            Just (SomeNat (_    :: Proxy padLeft)),    Just (SomeNat  (_    :: Proxy padRight)),
                            Just (SomeNat (_    :: Proxy padTop)),     Just (SomeNat  (_    :: Proxy padBottom))) ->
                              case (   singByProxy pinr %* singByProxy chan
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D3 outputRows outputColumns filters) ~ h ))
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (inputRows + padTop + padBottom) ~ outputRows) )
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (inputColumns + padLeft + padRight) ~ outputColumns ))) of
                                    (SNat, Dict, Dict, Dict) ->
                                        pure (SomeNetwork (Pad :~> rest :: Network ( Pad padLeft padTop padRight padBottom ': layers ) ( ('D3 inputRows inputColumns filters) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                   , do -- Build a Crop layer
                        let output_r = fromIntegral $ natVal r
                        let output_c = fromIntegral $ natVal c
                        let output_f = fromIntegral $ natVal f

                        crop_left   <- choose 0 10
                        crop_right  <- choose 0 10
                        crop_top    <- choose 0 10
                        crop_bottom <- choose 0 10

                        let input_r = output_r + crop_top + crop_bottom
                        let input_c = output_c + crop_left + crop_right

                        -- Remake types for input
                        case ( someNatVal output_f, someNatVal input_r, someNatVal input_c, someNatVal output_r, someNatVal output_c, someNatVal crop_left, someNatVal crop_right, someNatVal crop_top, someNatVal crop_bottom ) of
                          ( Just (SomeNat (chan :: Proxy filters)),
                            Just (SomeNat (pinr :: Proxy inputRows)),  Just (SomeNat  (_    :: Proxy inputColumns)),
                            Just (SomeNat (_    :: Proxy outputRows)), Just (SomeNat  (_    :: Proxy outputColumns)),
                            Just (SomeNat (_    :: Proxy cropLeft)),   Just (SomeNat  (_    :: Proxy cropRight)),
                            Just (SomeNat (_    :: Proxy cropTop)),    Just (SomeNat  (_    :: Proxy cropBottom))) ->
                              case ( singByProxy pinr %* singByProxy chan
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( ( 'D3 outputRows outputColumns filters) ~ h ))
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (outputRows + cropTop + cropBottom) ~ inputRows ) )
                                    , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (outputColumns + cropLeft + cropRight) ~ inputColumns ))) of
                                    (SNat, Dict, Dict, Dict) ->
                                        pure (SomeNetwork (Crop :~> rest :: Network ( Crop cropLeft cropTop cropRight cropBottom ': layers ) ( ('D3 inputRows inputColumns filters) ': h ': hs )))
                          _ -> Gen.discard -- Can't occur
                     ]
    ]

-- | Test a partial derivative numerically for a random network and input
--
-- This is the most important test.
prop_auto_diff :: Property
prop_auto_diff = withDiscards 1000 . withTests 10000 . property $ do
  SomeNetwork (network :: Network layers shapes) <- forAll genNetwork
  (input  :: S (Head shapes))     <- forAllWith nice genOfShape
  (target :: S (Last shapes))     <- forAllWith nice oneUp
  (tested :: S (Head shapes))     <- forAllWith nice oneUp

  let (!tapes, !output) = runNetwork network input
  let (_, !backgrad)    = runGradient network tapes target
  let inputDiff         = input + tested * 0.00001
  let expected          = maxVal ( backgrad * tested )
  let (_, !outputDiff)  = runNetwork network inputDiff
  let result            = maxVal ( outputDiff * target - output * target ) / 0.00001

  result ~~~ expected

-- Make a shape where all are 0 except for 1 value, which is 1.
oneUp :: forall shape. ( SingI shape ) => Gen (S shape)
oneUp =
  case ( sing :: Sing shape ) of
    D1Sing SNat ->
      let x = 0 :: S ( shape )
      in  case x of
            ( S1D x' ) -> do
              let ex = extract x'
              let len = VS.length ex
              ix <- choose 0 (len - 1)
              let nx = runST $ do ex' <- VS.thaw ex
                                  VS.write ex' ix 1.0
                                  VS.freeze ex'
              maybe Gen.discard pure . fromStorable $ nx

    D2Sing SNat SNat ->
      let x = 0 :: S ( shape )
      in  case x of
            ( S2D x' ) -> do
              let ex = flatten ( extract x' )
              let len = VS.length ex
              ix <- choose 0 (len - 1)
              let nx = runST $ do ex' <- VS.thaw ex
                                  VS.write ex' ix 1
                                  VS.freeze ex'
              maybe Gen.discard pure . fromStorable $ nx

    D3Sing SNat SNat SNat ->
      let x = 0 :: S ( shape )
      in  case x of
            ( S3D x' ) -> do
              let ex = flatten ( extract x' )
              let len = VS.length ex
              ix <- choose 0 (len - 1)
              let nx = runST $ do ex' <- VS.thaw ex
                                  VS.write ex' ix 1
                                  VS.freeze ex'
              maybe Gen.discard pure . fromStorable $ nx

maxVal :: S shape -> RealNum
maxVal ( S1D x ) = norm_Inf x
maxVal ( S2D x ) = norm_Inf x
maxVal ( S3D x ) = norm_Inf x

(~~~) :: (Monad m, HasCallStack) => RealNum -> RealNum -> PropertyT m ()
(~~~) x y =
  if abs (x - y) < precision then
    success
  else
    withFrozenCallStack $
    failWith Nothing $ unlines [
        "━━━ Not Simliar ━━━"
        , show x
        , show y
        ]
  where precision | nameF == show (typeRep (Proxy :: Proxy Float)) = 2e-2
                  | otherwise = 2e-4
infix 4 ~~~

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


{-# LANGUAGE
  Rank2Types,
  TypeFamilies,
  TupleSections
 #-}

-- Import all models under maintenance.
-- Models not imported here will not be compiled
-- when invoking `stack bench`.
import qualified BetaBin
import qualified Dice
import qualified DPmixture
import qualified Gamma
import qualified HMM
import Plotting

import Numeric.LogDomain
import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Simple
import Control.Monad.Bayes.Weighted
import Control.Monad.Bayes.Sampler
import Control.Monad.Bayes.Enumerator
import Control.Monad.Bayes.Inference
import Control.Monad.Bayes.Population

import System.IO
import Control.Arrow (first, second)
import Control.Monad (unless, when)
import Control.Monad.IO.Class (MonadIO, liftIO)

import Statistics.Sample
import qualified Data.Vector as Vector
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo
import Options.Applicative
import System.Directory

opts :: ParserInfo Bool
opts = flip info fullDesc $ switch
  ( long "trial"
 <> help "Run a quick version of benchmarks to check that all is working correctly")

tryCache :: (MonadIO m, Read a, Show a) => FilePath -> m a -> m a
tryCache filepath fresh = do
  exists <- liftIO $ doesFileExist filepath
  if exists then
    liftIO $ fmap read (readFile filepath)
  else do
    value <- fresh
    liftIO $ writeFile filepath (show value)
    return value

main = do
  -- make sure `putStrLn` prints to console immediately
  hSetBuffering stdout LineBuffering

  trial <- execParser opts
  when trial $ putStrLn "Trial run"

  sampleIO hmmBenchmark

meanVar :: (MonadDist m, CustomReal m ~ Double) => Int -> m Double -> m (Double, Double)
meanVar n d = do
  xs <- Vector.replicateM n d
  return $ meanVariance xs

smcParams :: [Int]
smcParams = [10,20,50,100,200,500,1000]

ns :: [Int]
--ns = [10,20,50,100,200,500,1000]
ns = [10,20..1000]

smcParamsDouble :: [Double]
smcParamsDouble = map fromIntegral smcParams

smcResults :: (MonadDist m, CustomReal m ~ Double) => [m (Vector.Vector Double)]
smcResults = map (\p -> Vector.replicateM 10 $ fmap HMM.hmmKL $ explicitPopulation $ smcMultinomial (length HMM.values) p HMM.hmm) smcParams

hmmBenchmark :: SamplerIO ()
hmmBenchmark = do
  liftIO $ putStrLn "running HMM benchmark"

  isSamples <- fmap (drop 5000) $ explicitPopulation $ importance 10000 HMM.hmm
  let isRes = map (\n -> HMM.hmmKL $ take n isSamples) ns
  -- mhSamples <- fmap (drop 5000) $ traceMH 10000 HMM.hmm
  -- let mhRes = map (\n -> HMM.hmmKL $ take n $ map (,1) mhSamples) ns
  mhPriorSamples <- fmap (drop 5000) $ mhPrior HMM.hmm 10000
  let mhPriorRes = map (\n -> HMM.hmmKL $ take n $ map (,1) mhPriorSamples) ns
  pimhSamples <- pimh (length HMM.values) 100 1000 HMM.hmm
  let pimhRes = map (\n -> HMM.hmmKL $ take n $ map (,1) pimhSamples) ns
  liftIO $ toFile (fo_format .~ PDF $ def) "anytime.pdf" $ do
    layout_title .= "HMM"
    anytimePlot "#samples" "KL" ns [
      ("IS", isRes),
      -- ("MHtrace", mhRes),
      ("MHprior", mhPriorRes),
      ("PIMH", pimhRes)]

  smcRes <- sequence smcResults
  liftIO $ toFile (fo_format .~ PDF $ def) "smc.pdf" $ do
    layout_title .= "HMM"
    oneShotPlot "#particles" "KL" [
      ("SMC", zip smcParamsDouble (map (errBars 2) smcRes))]


{-# LANGUAGE Haskell2010
    , TypeFamilies
    , FlexibleContexts
    , Trustworthy
    , StandaloneDeriving
    , DeriveDataTypeable
    , CPP
 #-}
{-# OPTIONS -Wall -fno-warn-name-shadowing #-}
{-# OPTIONS -cpp  -pgmPcpphs  -optP--cpp #-}

-- | Efficient matrix operations in 100% pure Haskell.
--
-- This package uses miscellaneous implementations,
-- depending on the type of its components. Typically unboxed
-- arrays will perform best, while unboxed arrays give you
-- certain features such as 'Rational' or 'Complex' components.
--
-- The following component types are supported by 'Matrix':
--
-- [@Int@] Uses unboxed arrays internally. 'inv' will always
--      return 'Nothing'.
--
-- [@Integer@] Uses boxed arrays internally. 'inv' will always
--      return 'Nothing'.
--
-- [@Double@ and @Float@] Uses unboxed arrays internally.
--      All matrix operations will work as expected.
--      @Matrix Double@ will probably yield the best peformance.
--
-- [@Rational@] Best choice if precision is what you aim for.
--      Uses boxed arrays internally. All matrix operations will
--      work as expected.
--
-- [@Complex@] Experimental. Uses boxed arrays internally.
--      All matrix operations will work as expected, though
--      finding the inverse of a matrix isa tad less numerically
--      stable than with a @Double@ matrix.
module Numeric.Matrix (

    Matrix,

    MatrixElement (..),

    -- * Matrix property and utility functions.
    (<|>),
    (<->),
    scale,

    -- ** Matrix properties
    isUnit,
    isZero,
    isDiagonal,
    isEmpty,
    isSquare,

    -- ** Conversions
    toDoubleMatrix,
    toComplexMatrix,
    toRationalMatrix

) where

import Control.Applicative ((<$>))
import Control.DeepSeq
import Control.Monad
import Control.Monad.ST

import Data.Function (on)
import Data.Ratio
import Data.Complex
import Data.Maybe
import Data.Int
import Data.Word

import qualified Data.List as L
import Data.Array.IArray
import Data.Array.MArray
import Data.Array.Unboxed
import Data.Array.ST
import Data.STRef
import Data.Binary

import qualified Data.Array.Unsafe as U

import Data.Typeable

import Prelude hiding (any, all, read, map)
import qualified Prelude as P

-- | Matrices are represented by a type which fits best the component type.
-- For example a @Matrix Double@ is represented by unboxed arrays,
-- @Matrix Integer@ by boxed arrays.
--
-- Data instances exist for 'Int', 'Float', 'Double', 'Integer', 'Ratio',
-- and 'Complex'. Certain types do have certain disadvantages, like for
-- example you can not compute the inverse matrix of a @Matrix Int@.
--
-- Every matrix (regardless of the component type) has instances for
-- 'Show', 'Read', 'Num', 'Fractional', 'Eq', 'Typeable', 'Binary',
-- and 'NFData'. This means that you can use arithmetic operations like
-- '+', '*', and '/', as well as functions like 'show', 'read', or 'typeOf'.
--
-- [@Show (Matrix e)@]
-- Note that a Show instance for the component type @e@ must exist.
--
-- [@Read (Matrix e)@]
-- You can read a matrix like so:
--
-- > read "1 0\n0 1\n" :: Matrix Double
--
-- [@Num (Matrix e)@]
-- '+', '-', '*', 'negate', 'abs', 'signum', and 'fromInteger'.
--
-- 'signum' will compute the determinant and return the signum
-- of it.
--
-- 'abs' applies @map abs@ on the matrix (that is, it applies
-- @abs@ on every component in the matrix and returns a new
-- matrix without negative components).
--
-- @fromInteger@ yields a 1-x-1-matrix.
--
-- [@Fractional (Matrix e)@]
-- Only available if there exists an instance @Fractional e@
-- (the component type needs to have a @Fractional@ instance, too).
-- Note that while the 'Num' operations are safe, 'recip' and
-- '/' will fail (with an 'error') if the involved matrix is
-- not invertible or not a square matrix.
--
-- [@NFData (Matrix e)@]
-- Matrices have instances for NFData so that you can use a
-- matrix in parallel computations using the @Control.Monad.Par@
-- monad (see the @monad-par@ package for details).
--
-- [@Typeable (Matrix e)@]
-- Allows you to use matrices as 'Dynamic' values.
--
-- [@Binary (Matrix e)@]
-- Serialize and unserialize matrices using the @binary@ package.
-- See @encode@ and @decode@.
data family Matrix e

#if defined(__GLASGOW_HASKELL__) && (__GLASGOW_HASKELL__ >= 707)
deriving instance Typeable Matrix
#else
deriving instance Typeable1 Matrix
#endif

data instance Matrix Int
    = IntMatrix !Int !Int (Array Int (UArray Int Int))

data instance Matrix Float
    = FloatMatrix !Int !Int (Array Int (UArray Int Float))

data instance Matrix Double
    = DoubleMatrix !Int !Int (Array Int (UArray Int Double))

data instance Matrix Integer
    = IntegerMatrix !Int !Int (Array Int (Array Int Integer))

data instance Matrix (Ratio a)
    = RatioMatrix !Int !Int (Array Int (Array Int (Ratio a)))

data instance Matrix (Complex a)
    = ComplexMatrix !Int !Int (Array Int (Array Int (Complex a)))

instance (MatrixElement e, Show e) => Show (Matrix e) where
    show = unlines . P.map showRow . toList
      where
        showRow = unwords . P.map ((' ':) . show)

instance (Read e, MatrixElement e) => Read (Matrix e) where
    readsPrec _ = (\x -> [(x, "")]) . fromList . P.map (P.map P.read . words) . lines

instance (MatrixElement e) => Num (Matrix e) where
    (+) = plus
    (-) = minus
    (*) = times
    abs         = map abs
    signum      = matrix (1,1) . const . signum . det
    fromInteger = matrix (1,1) . const . fromInteger

instance (MatrixElement e, Fractional e) => Fractional (Matrix e) where
    recip        = fromJust . inv
    fromRational = matrix (1,1) . const . fromRational

instance (MatrixElement e) => Eq (Matrix e) where
    m == n
        | dimensions m == dimensions n
            = allWithIndex (\ix e -> m `at` ix == e) n
        | otherwise = False

instance (MatrixElement e) => NFData (Matrix e) where
    rnf matrix = matrix `deepseq` ()

instance (MatrixElement e, Binary e) => Binary (Matrix e) where

    put m = do
        let (rows, cols) = dimensions m

        put rows >> put cols

        forM_ [1..rows] $ \i -> do
            forM_ [1..cols] $ \j -> do
                put (m `at` (i,j))

    get = do
        rows <- get :: Get Int
        cols <- get :: Get Int

        forM [1..rows] (const (forM [1..cols] (const get)))
            >>= return . fromList

(<|>) :: MatrixElement e => Matrix e -> Matrix e -> Matrix e
-- ^ Joins two matrices horizontally.
--
-- > 1 2 3     1 0 0      1 2 3 1 0 0
-- > 3 4 5 <|> 2 1 0  ->  3 4 5 2 1 0
-- > 5 6 7     3 2 1      5 6 7 3 2 1
m1 <|> m2 = let m = numCols m1
                n1 = numRows m1
                n2 = numRows m2
            in matrix (max n1 n2, m + numCols m2)
              $ \(i,j) -> if j > m
                    then (if i > n2 then 0 else m2 `at` (i,j-m))
                    else (if i > n1 then 0 else m1 `at` (i,j))

(<->) :: MatrixElement e => Matrix e -> Matrix e -> Matrix e
-- ^ Joins two matrices vertically.
--
-- > 1 2 3     1 0 0      1 2 3
-- > 3 4 5 <-> 2 1 0  ->  3 4 5
-- > 5 6 7     3 2 1      5 6 7
-- >                      1 0 0
-- >                      2 1 0
-- >                      3 2 1
m1 <-> m2 = let m = numRows m1
                n1 = numCols m1
                n2 = numCols m2
            in matrix (m + numRows m2, max n1 n2)
              $ \(i,j) -> if i > m
                    then (if j > n2 then 0 else m2 `at` (i-m,j))
                    else (if j > n1 then 0 else m1 `at` (i,j))

scale :: MatrixElement e => Matrix e -> e -> Matrix e
-- ^ Scales a matrix by the given factor.
--
-- > scale s == map (*s)
scale m s = map (*s) m

isUnit, isDiagonal, isZero, isEmpty, isSquare :: MatrixElement e => Matrix e -> Bool

-- | Check whether the matrix consists of all zeros.
--
-- > isZero == all (== 0)
isZero = all (== 0)

-- | Check whether the matrix is an identity matrix.
--
-- > 1 0 0
-- > 0 1 0
-- > 0 0 1 (True)
isUnit m = isSquare m && allWithIndex (uncurry check) m
    where check = \i j e -> if i == j then e == 1 else e == 0

-- | Checks whether the matrix is empty.
--
-- > isEmpty m = numCols == 0 || numRows == 0
isEmpty m = numRows m == 0 || numCols m == 0

-- | Checks whether the matrix is a diagonal matrix.
--
-- > 4 0 0 0
-- > 0 7 0 0
-- > 0 0 3 0
-- > 0 0 0 9 (True)
isDiagonal m = isSquare m && allWithIndex (uncurry check) m
    where check = \i j e -> if i /= j then e == 0 else True

-- | Checks whether the matrix is a square matrix.
--
-- > isSquare == uncurry (==) . dimensions
isSquare m = let (a, b) = dimensions m in a == b

toDoubleMatrix :: (MatrixElement a, Integral a) => Matrix a -> Matrix Double
toDoubleMatrix = map fromIntegral

toRationalMatrix :: (MatrixElement a, Real a) => Matrix a -> Matrix Rational
toRationalMatrix = map toRational

toComplexMatrix :: (MatrixElement a, RealFloat a, Show a) => Matrix a -> Matrix (Complex a)
toComplexMatrix = map (:+ 0)

class Division e where
    divide :: e -> e -> e

instance Division Int     where divide = quot
instance Division Int8    where divide = quot
instance Division Int16   where divide = quot
instance Division Int32   where divide = quot
instance Division Int64   where divide = quot
instance Division Word8   where divide = quot
instance Division Word16  where divide = quot
instance Division Word32  where divide = quot
instance Division Word64  where divide = quot
instance Division Integer where divide = quot
instance Division Float   where divide = (/)
instance Division Double  where divide = (/)
instance Integral a => Division (Ratio a) where divide = (/)
instance RealFloat a => Division (Complex a) where divide = (/)

class (Eq e, Num e) => MatrixElement e where

    -- | Creates a matrix of the given size using a generator
    -- function for the value of each component.
    matrix :: (Int, Int) -> ((Int, Int) -> e) -> Matrix e

    -- | Builds a list from a matrix for the indices for which
    -- the given predicate matches.
    --
    -- > trace == select (uncurry (==))
    select :: ((Int, Int) -> Bool) -> Matrix e -> [e]

    -- | Returns the component at the given position in the matrix.
    -- Note that indices start at one, not at zero.
    at :: Matrix e -> (Int, Int) -> e

    -- | Returns the row at the given index in the matrix.
    -- Note that indices start at one, not at zero.
    row :: Int -> Matrix e -> [e]

    -- | Returns the row at the given index in the matrix.
    -- Note that indices start at one, not at zero.
    col :: Int -> Matrix e -> [e]

    -- | The dimensions of a given matrix.
    dimensions :: Matrix e -> (Int, Int)

    -- | The number of rows in the given matrix.
    numRows :: Matrix e -> Int

    -- | The number of columns in the given matrix.
    numCols :: Matrix e -> Int

    -- | Builds a matrix from a list of lists.
    --
    -- The innermost lists represent the rows. This function will create a m-n-matrix,
    -- where m is the number of rows, which is the minimum length of the row lists
    -- and n is the number of columns, i.e. the length of the outer list.
    --
    -- > fromList [[1,2,3],[2,1,3],[3,2,1]] :: Matrix Rational
    fromList :: [[e]] -> Matrix e

    -- | Turns a matrix into a list of lists.
    --
    -- > (toList . fromList) xs == xs
    --
    -- > (fromList . toList) mat == mat
    toList   :: Matrix e -> [[e]]

    -- | An identity square matrix of the given size.
    --
    -- >>> unit 4
    -- 1 0 0 0
    -- 0 1 0 0
    -- 0 0 1 0
    -- 0 0 0 1
    unit  :: Int -> Matrix e

    -- | A square matrix of the given size consisting of all zeros.
    --
    -- >>> zero 2
    -- 0 0
    -- 0 0
    zero  :: Int -> Matrix e

    -- | A square matrix which trace is the given list, all other components
    -- set to zero.
    --
    -- >>> diag [1,4,7,9]
    -- 1 0 0 0
    -- 0 4 0 0
    -- 0 0 7 0
    -- 0 0 0 9
    diag  :: [e] -> Matrix e

    -- | Check whether the matrix is the empty matrix.
    --
    -- > dimensions empty == (0, 0)
    empty :: Matrix e

    -- | Subtract a matrix from another.
    minus :: Matrix e -> Matrix e -> Matrix e

    -- | Add two matrices.
    --
    -- You may also use the 'Num' instance for matrices,
    -- i.e. @matrix1 + matrix2@ will work, too.
    plus  :: Matrix e -> Matrix e -> Matrix e

    -- | Multiply two matrices /O(n^3)/.
    --
    -- You may also use the 'Num' instance for matrices,
    -- i.e. @matrix1 * matrix2@ will work, too.
    times :: Matrix e -> Matrix e -> Matrix e

    -- | Compute the inverse of a matrix. Returns @Nothing@
    -- if the matrix is not invertible.
    inv   :: Matrix e -> Maybe (Matrix e)

    -- | Applies Bareiss multistep integer-preserving
    -- algorithm for finding the determinant of a matrix.
    -- Returns 0 if the matrix is not a square matrix.
    det       :: Matrix e -> e

    -- | Flip rows and columns.
    --
    -- > 1 8 9                1 2 3
    -- > 2 1 8  --transpose-> 8 1 2
    -- > 3 2 1                9 8 1
    transpose :: Matrix e -> Matrix e

    -- | Compute the rank of a matrix.
    rank      :: Matrix e -> e

    -- | Select the diagonal elements of a matrix as a list.
    --
    -- > 1 8 3
    -- > 3 6 5 --trace-> [1, 6, 2]
    -- > 7 4 2
    trace     :: Matrix e -> [e]

    -- | Select the minor of a matrix, that is the determinant
    -- of the 'minorMatrix'.
    --
    -- > minor = det . minorMatrix
    minor :: MatrixElement e => (Int, Int) -> Matrix e -> e

    -- | Select the minor matrix of a matrix, a matrix that is obtained
    -- by deleting the i-th row and j-th column.
    --
    -- > 10  9 95 45
    -- >  8  7  3 27                        8  3 27
    -- > 13 17 19 23 --minorMatrix (1,2)-> 13 19 23
    -- >  1  2  5  8                        1  5  8
    minorMatrix :: MatrixElement e => (Int, Int) -> Matrix e -> Matrix e

    cofactors :: MatrixElement e => Matrix e -> Matrix e
    adjugate :: MatrixElement e => Matrix e -> Matrix e

    -- | Apply a function on every component in the matrix.
    map :: MatrixElement f => (e -> f) -> Matrix e -> Matrix f

    -- | Apply a predicate on every component in the matrix
    -- and returns True iff all components satisfy it.
    all :: (e -> Bool) -> Matrix e -> Bool

    -- | Apply a predicate on every component in the matrix
    -- and return True if one or more components satisfy it.
    any :: (e -> Bool) -> Matrix e -> Bool

    mapWithIndex :: MatrixElement f => ((Int, Int) -> e -> f) -> Matrix e -> Matrix f
    allWithIndex :: ((Int, Int) -> e -> Bool) -> Matrix e -> Bool
    anyWithIndex :: ((Int, Int) -> e -> Bool) -> Matrix e -> Bool

    unit n  = fromList [[ if i == j then 1 else 0 | j <- [1..n]] | i <- [1..n] ]
    zero n  = matrix (n,n) (const 0)
    empty   = fromList []
    diag xs = matrix (n,n) (\(i,j) -> if i == j then xs !! (i-1) else 0)
      where n = length xs

    select p m = [ at m (i,j) | i <- [1..numRows m]
                              , j <- [1..numCols m]
                              , p (i,j) ]

    at mat (i, j) = ((!! j) . (!! i) . toList) mat

    row i = (!! (i-1)) . toList
    col i = row i . transpose

    numRows = fst . dimensions
    numCols = snd . dimensions
    dimensions m = case toList m of [] -> (0, 0)
                                    (x:xs) -> (length xs + 1, length x)

    adjugate = transpose . cofactors
    transpose mat = matrix (n, m) (\(i,j) -> mat `at` (j,i))
      where (m, n) = dimensions mat
    trace = select (uncurry (==))
    inv _ = Nothing

    minorMatrix (i,j) mat = matrix (numRows mat - 1, numCols mat - 1) $
                \(i',j') -> mat `at` (if i' >= i then i' + 1 else i',
                                      if j' >= j then j' + 1 else j')

    minor ix = det . minorMatrix ix

    cofactors mat = matrix (dimensions mat) $
       \(i,j) -> fromIntegral ((-1 :: Int)^(i+j)) * minor (i,j) mat

    map f = mapWithIndex (const f)
    all f = allWithIndex (const f)
    any f = anyWithIndex (const f)

    mapWithIndex f m = matrix (dimensions m) (\x -> f x (m `at` x))
    allWithIndex f m = P.all id [ f (i, j) (m `at` (i,j))
                                | i <- [1..numRows m], j <- [1..numCols m]]
    anyWithIndex f m = P.any id [ f (i, j) (m `at` (i,j))
                                | i <- [1..numRows m], j <- [1..numCols m]]

    a `plus` b
        | dimensions a /= dimensions b = error "Matrix.plus: dimensions don't match."
        | otherwise = matrix (dimensions a) (\x -> a `at` x + b `at` x)
    a `minus` b
        | dimensions a /= dimensions b = error "Matrix.minus: dimensions don't match."
        | otherwise = matrix (dimensions a) (\x -> a `at` x - b `at` x)
    a `times` b
        | numCols a /= numRows b = error "Matrix.times: `numRows a' and `numCols b' don't match."
        | otherwise = _mult a b

instance MatrixElement Int where
    matrix d g = runST (_matrix IntMatrix arrayST arraySTU d g)
    fromList = _fromList IntMatrix

    at         (IntMatrix _ _ arr) = _at arr
    dimensions (IntMatrix m n _)   = (m, n)
    row i      (IntMatrix _ _ arr) = _row i arr
    col j      (IntMatrix _ _ arr) = _col j arr
    toList     (IntMatrix _ _ arr) = _toList arr
    det        (IntMatrix m n arr) = if m /= n then 0 else runST (_det thawsUnboxed arr)
    rank       (IntMatrix _ _ arr) = runST (_rank thawsBoxed arr)

instance MatrixElement Integer where
    matrix d g = runST (_matrix IntegerMatrix arrayST arrayST d g)
    fromList   = _fromList IntegerMatrix

    at         (IntegerMatrix _ _ arr) = _at arr
    dimensions (IntegerMatrix m n _)   = (m, n)
    row i      (IntegerMatrix _ _ arr) = _row i arr
    col j      (IntegerMatrix _ _ arr) = _col j arr
    toList     (IntegerMatrix _ _ arr) = _toList arr
    det        (IntegerMatrix m n arr) = if m /= n then 0 else runST (_det thawsBoxed arr)
    rank       (IntegerMatrix _ _ arr) = runST (_rank thawsBoxed arr)

instance MatrixElement Float where
    matrix d g = runST (_matrix FloatMatrix arrayST arraySTU d g)
    fromList   = _fromList FloatMatrix

    at         (FloatMatrix _ _ arr) = _at arr
    dimensions (FloatMatrix m n _  ) = (m, n)
    row i      (FloatMatrix _ _ arr) = _row i arr
    col j      (FloatMatrix _ _ arr) = _col j arr
    toList     (FloatMatrix _ _ arr) = _toList arr
    det        (FloatMatrix m n arr) = if m /= n then 0 else runST (_det thawsUnboxed arr)
    rank       (FloatMatrix _ _ arr) = runST (_rank thawsBoxed arr)
    inv        (FloatMatrix m n arr) = if m /= n then Nothing else
                                         let x = runST (_inv unboxedST pivotMax arr)
                                         in maybe Nothing (Just . FloatMatrix m n) x

instance MatrixElement Double where
    matrix d g = runST (_matrix DoubleMatrix arrayST arraySTU d g)
    fromList   = _fromList DoubleMatrix

    at         (DoubleMatrix _ _ arr) = _at arr
    dimensions (DoubleMatrix m n _  ) = (m, n)
    row i      (DoubleMatrix _ _ arr) = _row i arr
    col j      (DoubleMatrix _ _ arr) = _col j arr
    toList     (DoubleMatrix _ _ arr) = _toList arr
    det        (DoubleMatrix m n arr) = if m /= n then 0 else runST (_det thawsUnboxed arr)
    rank       (DoubleMatrix _ _ arr) = runST (_rank thawsBoxed arr)
    inv        (DoubleMatrix m n arr) = if m /= n then Nothing else
                                         let x = runST (_inv unboxedST pivotMax arr)
                                         in maybe Nothing (Just . DoubleMatrix m n) x

instance (Show a, Integral a) => MatrixElement (Ratio a) where
    matrix d g = runST (_matrix RatioMatrix arrayST arrayST d g)
    fromList   = _fromList RatioMatrix

    at         (RatioMatrix _ _ arr) = _at arr
    dimensions (RatioMatrix m n _  ) = (m, n)
    row i      (RatioMatrix _ _ arr) = _row i arr
    col j      (RatioMatrix _ _ arr) = _col j arr
    toList     (RatioMatrix _ _ arr) = _toList arr
    det        (RatioMatrix m n arr) = if m /= n then 0 else  runST (_det thawsBoxed arr)
    rank       (RatioMatrix _ _ arr) = runST (_rank thawsBoxed arr)
    inv        (RatioMatrix m n arr) = if m /= n then Nothing else
                                        let x = runST (_inv boxedST pivotMax arr)
                                        in maybe Nothing (Just . RatioMatrix m n) x

instance (Show a, RealFloat a) => MatrixElement (Complex a) where
    matrix d g = runST (_matrix ComplexMatrix arrayST arrayST d g)
    fromList   = _fromList ComplexMatrix

    at         (ComplexMatrix _ _ arr) = _at arr
    dimensions (ComplexMatrix m n _  ) = (m, n)
    row i      (ComplexMatrix _ _ arr) = _row i arr
    col j      (ComplexMatrix _ _ arr) = _col j arr
    toList     (ComplexMatrix _ _ arr) = _toList arr
    det        (ComplexMatrix m n arr) = if m /= n then 0 else runST (_det thawsBoxed arr)
    rank       (ComplexMatrix _ _ arr) = runST (_rank thawsBoxed arr)
    inv        (ComplexMatrix m n arr) = if m /= n then Nothing else
                                          let x = runST (_inv boxedST pivotNonZero arr)
                                          in maybe Nothing (Just . ComplexMatrix m n) x

_at :: (IArray a (u Int e), IArray u e)
    => a Int (u Int e) -> (Int, Int) -> e
_at arr (i,j) = arr ! i ! j

_row, _col :: (IArray a (u Int e), IArray u e) => Int -> a Int (u Int e) -> [e]
_row i arr = let row = arr ! i in [ row ! j | j <- [1..(snd (bounds row))] ]
_col j arr = [ arr ! i ! j | i <- [1..(snd (bounds arr))] ]

_toList :: (IArray a e) => Array Int (a Int e) -> [[e]]
_toList = P.map elems . elems

_fromList :: (IArray a (u Int e), IArray u e)
          => (Int -> Int -> a Int (u Int e) -> matrix e) -> [[e]] -> matrix e
_fromList c xs =
    let lengths = P.map length xs
        numCols = if null lengths then 0 else foldl1 min lengths
        numRows = length lengths

    in  c numRows numCols
          $ array (1, numRows)
          $ zip [1..numRows]
          $ P.map (array (1, numCols) . zip [1..numCols]) xs

thawsBoxed :: (IArray a e, MArray (STArray s) e (ST s))
           => Array Int (a Int e)
           -> ST s [STArray s Int e]
thawsBoxed = mapM thaw . elems

thawsUnboxed :: (IArray a e, MArray (STUArray s) e (ST s))
             => Array Int (a Int e)
             -> ST s [STUArray s Int e]
thawsUnboxed = mapM thaw . elems

arrays :: [(u s) Int e]
       -> ST s ((STArray s) Int ((u s) Int e))
arrays list = newListArray (1, length list) list

augment :: (IArray a e, MArray (u s) e (ST s), Num e)
        => ((Int, Int) -> [e] -> ST s ((u s) Int e))
        -> Array Int (a Int e)
        -> ST s (STArray s Int (u s Int e))
augment _ arr = do
    let (_, n) = bounds arr
        row (a,i) = newListArray (1, 2*n)
                                 [ if j > n then (if j == i + n then 1 else 0)
                                            else a ! j
                                 | j <- [1..2*n] ]

    mapM row (zip (elems arr) [1..]) >>= newListArray (1, n)

boxedST :: MArray (STArray s) e (ST s)
        => (Int, Int) -> [e] -> ST s ((STArray s) Int e)
boxedST = newListArray

unboxedST :: MArray (STUArray s) e (ST s)
          => (Int, Int) -> [e] -> ST s ((STUArray s) Int e)
unboxedST = newListArray

arrayST :: MArray (STArray s) e (ST s)
        => (Int, Int) -> e -> ST s ((STArray s) Int e)
arrayST = newArray

arraySTU :: MArray (STUArray s) e (ST s)
         => (Int, Int) -> e -> ST s ((STUArray s) Int e)
arraySTU = newArray

tee :: Monad m => (b -> m a) -> b -> m b
tee f x = f x >> return x

read :: (MArray a1 b m, MArray a (a1 Int b) m) =>
                       a Int (a1 Int b) -> Int -> Int -> m b
read a i j = readArray a i >>= flip readArray j

pivotMax :: Ord v => [(i, v)] -> i
pivotMax = fst . L.maximumBy (compare `on` snd)

pivotNonZero :: (Num v, Eq v) => [(i, v)] -> i
pivotNonZero xs = maybe (fst $ head xs) fst $ L.find ((/= 0) . snd) xs

_inv :: (IArray a e, MArray (u s) e (ST s), Fractional e, Show e, Eq e)
     => ((Int, Int) -> [e] -> ST s ((u s) Int e))
        -- ^ A function for building a new array
     -> ([(Int, e)] -> Int)
        -- ^ A function for selecting pivot elements
     -> Array Int (a Int e)
        -- ^ A matrix as arrays or arrays
     -> ST s (Maybe (Array Int (a Int e)))
_inv mkArrayST selectPivot mat = do
    let m = snd $ bounds mat
        n = 2 * m

        swap a i j = do
            tmp <- readArray a i
            readArray a j >>= writeArray a i
            writeArray a j tmp

    okay <- newSTRef True

    a <- augment mkArrayST mat

    forM_ [1..m] $ \k -> do
        iPivot <- selectPivot <$> zip [k..m]
                              <$> mapM (\i -> abs <$> read a i k) [k..m]

        p <- read a iPivot k
        if p == 0 then writeSTRef okay False else do

            swap a iPivot k

            forM_ [k+1..m] $ \i -> do
                a_i <- readArray a i
                a_k <- readArray a k
                forM_ [k+1..n] $ \j -> do
                    a_ij <- readArray a_i j
                    a_kj <- readArray a_k j
                    a_ik <- readArray a_i k
                    writeArray a_i j (a_ij - a_kj * (a_ik / p))
                writeArray a_i k 0

    invertible <- readSTRef okay

    if invertible then
      do
        forM_ [ m - v | v <- [0..m-1] ] $ \i -> do
            a_i <- readArray a i
            p   <- readArray a_i i
            writeArray a_i i 1
            forM_ [i+1..n] $ \j -> do
                readArray a_i j >>= writeArray a_i j . (/ p)

            unless (i == m) $ do
                forM_ [i+1..m] $ \k -> do
                    a_k <- readArray a k
                    p   <- readArray a_i k

                    forM_ [k..n] $ \j -> do
                        a_ij <- readArray a_i j
                        a_kj <- readArray a_k j
                        writeArray a_i j (a_ij - p * a_kj)

        mapM (\i -> readArray a i >>= getElems
                        >>= return . listArray (1, m) . drop m) [1..m]
            >>= return . Just . listArray (1, m)

      else return Nothing

_rank :: (IArray a e, MArray (u s) e (ST s), Num e, Division e, Eq e)
      => (Array Int (a Int e) -> ST s [(u s) Int e])
      -- ^ A function for thawing a boxed array
      -> Array Int (a Int e)
      -- ^ A matrix given as array of arrays
      -> ST s e
_rank thaws mat = do
    let m = snd $ bounds mat
        n = snd $ bounds (mat ! 1)

        swap a i j = do
            tmp <- readArray a i
            readArray a j >>= writeArray a i
            writeArray a j tmp

    a <- thaws mat >>= arrays

    ixPivot <- newSTRef 1
    prevR   <- newSTRef 1

    forM_ [1..n] $ \k -> do
        pivotRow <- readSTRef ixPivot

        switchRow <- mapM (\i -> read a i k) [pivotRow .. m]
            >>= return . L.findIndex (/= 0)

        when (isJust switchRow) $ do
            let ix = fromJust switchRow + pivotRow
            when (pivotRow /= ix) (swap a pivotRow ix)

            a_p   <- readArray a k
            pivot <- readArray a_p k
            prev  <- readSTRef prevR

            forM_ [pivotRow+1..m] $ \i -> do
                a_i <- readArray a i
                forM_ [k+1..n] $ \j -> do
                    a_ij <- readArray a_i j
                    a_ik <- readArray a_i k
                    a_pj <- readArray a_p j
                    writeArray a_i j ((pivot * a_ij - a_ik * a_pj)
                                        `divide` prev)

            writeSTRef ixPivot (pivotRow + 1)
            writeSTRef prevR pivot

    readSTRef ixPivot >>= return . (+ negate 1) . fromIntegral

_det :: (IArray a e, MArray (u s) e (ST s),
         Num e, Eq e, Division e)
     => (Array Int (a Int e) -> ST s [(u s) Int e])
     -> Array Int (a Int e) -> ST s e
_det thaws mat = do

    let size = snd $ bounds mat

    a <- thaws mat >>= arrays

    signR  <- newSTRef 1
    pivotR <- newSTRef 1

    forM_ [1..size] $ \k -> do
        sign <- readSTRef signR
        unless (sign == 0) $ do

            prev  <- readSTRef pivotR
            pivot <- read a k k >>= tee (writeSTRef pivotR)

            when (pivot == 0) $ do
                s <- forM [(k+1)..size] $ \r -> do
                    a_rk <- read a r k
                    if a_rk == 0 then return 0 else return r
                let sf = filter (>0) s

                when (not $ null sf) $ do
                    let sw = head sf

                    row <- readArray a sw
                    readArray a k >>= writeArray a sw
                    writeArray a k row

                    read a k k >>= writeSTRef pivotR
                    readSTRef signR >>= writeSTRef signR . negate

                when (null sf) (writeSTRef signR 0)

            sign' <- readSTRef signR
            unless (sign' == 0) $ do
                pivot' <- readSTRef pivotR
                forM_ [(k+1)..size] $ \i -> do
                    a_i <- readArray a i
                    forM [(k+1)..size] $ \j -> do
                        a_ij <- readArray a_i j
                        a_ik <- readArray a_i k
                        a_kj <- read a k j
                        writeArray a_i j ((pivot' * a_ij - a_ik * a_kj) `divide` prev)

    liftM2 (*) (readSTRef pivotR) (readSTRef signR)

-- TODO: More efficient implementation (decrease the constant factors
-- a little bit by working in the ST monad)
-- [ remark: not sure if this will be faster than lists -> benchmark! ]
_mult :: MatrixElement e => Matrix e -> Matrix e -> Matrix e
_mult a b = let rowsA = numRows a
                rowsB = numRows b
                colsB = numCols b
            in  matrix (rowsA, colsB) (\(i,j) -> L.foldl' (+) 0 [a `at` (i, k) * b `at` (k, j) | k <- [1..rowsB]])

_matrix :: (IArray a1 (u Int e), IArray u e,
            MArray a2 (u Int e) (ST s), MArray a3 e (ST s),
            Num e)
        => (Int -> Int -> a1 Int (u Int e) -> matrix)
        -> ((Int, Int) -> a -> ST s (a2 Int (u Int e)))
        -> ((Int, Int) -> e -> ST s (a3 Int e))
        -> (Int, Int)
        -> ((Int, Int) -> e)
        -> ST s matrix
_matrix c newArray newArrayU (m, n) g = do
    rows <- newArray (1, m) undefined
    forM_ [1..m] $ \i -> do
        cols <- newArrayU (1, n) 0
        forM_ [1..n] $ \j -> do
            writeArray cols j (g (i,j))
        U.unsafeFreeze cols >>= writeArray rows i
    U.unsafeFreeze rows >>= return . c m n

{-# RULES

"det/pow"
    forall a k. det (a ^ k) = (det a) ^ k

 #-}