{-# LANGUAGE CPP, DerivingVia, GADTs, InstanceSigs, KindSignatures    #-}
{-# LANGUAGE PatternSynonyms, RankNTypes, RoleAnnotations             #-}
{-# LANGUAGE ScopedTypeVariables, StandaloneDeriving, TemplateHaskell #-}
{-# LANGUAGE TypeApplications, TypeFamilies, TypeOperators            #-}
{-# LANGUAGE UndecidableSuperClasses                                  #-}
{-# OPTIONS_GHC -Wno-orphans #-}
module Control.Subcategory.Functor
  ( Constrained(..), Dom(), CFunctor (..),
    (<$:>),
    defaultCmapConst,
    WrapFunctor (..),
    WrapMono (WrapMono, unwrapMono),
    coerceToMono, withMonoCoercible,
  )
where
import qualified Control.Applicative                  as App
import           Control.Arrow                        (Arrow, ArrowMonad)
import           Control.Exception                    (Handler)
import qualified Control.Monad.ST.Lazy                as LST
import qualified Control.Monad.ST.Strict              as SST
import           Control.Subcategory.Wrapper.Internal
import           Data.Coerce
import           Data.Complex                         (Complex)
import qualified Data.Functor.Compose                 as SOP
import           Data.Functor.Const                   (Const)
import           Data.Functor.Identity                (Identity)
import qualified Data.Functor.Product                 as SOP
import qualified Data.Functor.Sum                     as SOP
import           Data.Hashable                        (Hashable)
import qualified Data.HashMap.Strict                  as HM
import qualified Data.HashSet                         as HS
import qualified Data.IntMap                          as IM
import           Data.Kind                            (Constraint, Type)
import           Data.List.NonEmpty                   (NonEmpty)
import qualified Data.Map                             as Map
import qualified Data.Monoid                          as Mon
import           Data.MonoTraversable                 (Element,
                                                       MonoFunctor (..))
#if MIN_VERSION_mono_traversable(1,0,14)
import Data.MonoTraversable (WrappedMono)
#endif

import qualified Data.IntSet                     as IS
import           Data.Ord                        (Down (..))
import qualified Data.Primitive.Array            as A
import qualified Data.Primitive.PrimArray        as PA
import qualified Data.Primitive.SmallArray       as SA
import           Data.Proxy                      (Proxy)
import qualified Data.Semigroup                  as Sem
import qualified Data.Sequence                   as Seq
import qualified Data.Set                        as Set
import qualified Data.Tree                       as Tree
import qualified Data.Vector                     as V
import qualified Data.Vector.Primitive           as P
import qualified Data.Vector.Storable            as S
import qualified Data.Vector.Unboxed             as U
import           Foreign.Ptr                     (Ptr)
import           GHC.Conc                        (STM)
import           GHC.Generics                    ((:*:) (..), (:+:) (..),
                                                  (:.:) (..), K1, M1, Par1,
                                                  Rec1, U1, URec, V1)
import qualified System.Console.GetOpt           as GetOpt
import           Text.ParserCombinators.ReadP    (ReadP)
import           Text.ParserCombinators.ReadPrec (ReadPrec)

infixl 4 <$:

class Constrained (f :: Type -> Type) where
  type Dom f (a :: Type) :: Constraint
  type Dom f a = ()

class Constrained f => CFunctor f where
  cmap :: (Dom f a, Dom f b) => (a -> b) -> f a -> f b
  default cmap :: Functor f => (a -> b) -> f a -> f b
  cmap = fmap
  {-# INLINE cmap #-}
  (<$:) :: (Dom f a, Dom f b) => a -> f b -> f a
  (<$:) = cmap . const
  {-# INLINE (<$:) #-}

defaultCmapConst :: (CFunctor f, Dom f a, Dom f b) => a -> f b -> f a
defaultCmapConst = cmap . const
{-# INLINE defaultCmapConst #-}

instance Constrained (WrapFunctor f) where
  type Dom (WrapFunctor f) a = ()

instance Functor f => CFunctor (WrapFunctor f) where
  cmap :: (a -> b) -> WrapFunctor f a -> WrapFunctor f b
  cmap = fmap
  {-# INLINE cmap #-}
  (<$:) :: a -> WrapFunctor f b -> WrapFunctor f a
  (<$:) = (<$)
  {-# INLINE (<$:) #-}

instance Constrained []
instance CFunctor []
instance Constrained Maybe
instance CFunctor Maybe
instance Constrained IO
instance CFunctor IO
instance Constrained Par1
instance CFunctor Par1
instance Constrained NonEmpty
instance CFunctor NonEmpty
instance Constrained ReadP
instance CFunctor ReadP
instance Constrained ReadPrec
instance CFunctor ReadPrec

instance Constrained Down
instance CFunctor Down
instance Constrained Mon.Product
instance CFunctor Mon.Product

instance Constrained Mon.Sum
instance CFunctor Mon.Sum
instance Constrained Mon.Dual
instance CFunctor Mon.Dual

instance Constrained Mon.Last
instance CFunctor Mon.Last
instance Constrained Mon.First
instance CFunctor Mon.First

instance Constrained STM
instance CFunctor STM
instance Constrained Handler
instance CFunctor Handler

instance Constrained Identity
instance CFunctor Identity
instance Constrained App.ZipList
instance CFunctor App.ZipList
instance Constrained GetOpt.ArgDescr
instance CFunctor GetOpt.ArgDescr
instance Constrained GetOpt.OptDescr
instance CFunctor GetOpt.OptDescr
instance Constrained GetOpt.ArgOrder
instance CFunctor GetOpt.ArgOrder
instance Constrained Sem.Option
instance CFunctor Sem.Option

instance Constrained Sem.Last
instance CFunctor Sem.Last
instance Constrained Sem.First
instance CFunctor Sem.First

instance Constrained Sem.Max
instance CFunctor Sem.Max
instance Constrained Sem.Min
instance CFunctor Sem.Min

instance Constrained Complex
instance CFunctor Complex
instance Constrained (Either a)
instance CFunctor (Either a)

instance Constrained V1
instance CFunctor V1
instance Constrained U1
instance CFunctor U1

instance Constrained ((,) a)
instance CFunctor ((,) a)
instance Constrained (SST.ST s)
instance CFunctor (SST.ST s)

instance Constrained (LST.ST s)
instance CFunctor (LST.ST s)
instance Constrained Proxy
instance CFunctor Proxy

instance Constrained (ArrowMonad a)
instance Arrow a => CFunctor (ArrowMonad a)
instance Constrained (App.WrappedMonad m)
instance Monad m => CFunctor (App.WrappedMonad m)

instance Constrained (Sem.Arg a)
instance CFunctor (Sem.Arg a)
instance Constrained (Rec1 f)
instance Functor f => CFunctor (Rec1 f)

instance Constrained (URec Char)
instance CFunctor (URec Char)
instance Constrained (URec Double)
instance CFunctor (URec Double)

instance Constrained (URec Float)
instance CFunctor (URec Float)
instance Constrained (URec Int)
instance CFunctor (URec Int)

instance Constrained (URec Word)
instance CFunctor (URec Word)
instance Constrained (URec (Ptr ()))
instance CFunctor (URec (Ptr ()))

instance Constrained f => Constrained (Mon.Ap f) where
  type Dom (Mon.Ap f) a = Dom f a

deriving newtype instance CFunctor f => CFunctor (Mon.Ap f)

instance Constrained (Mon.Alt f) where
  type Dom (Mon.Alt f) a = Dom f a
deriving newtype instance CFunctor f => CFunctor (Mon.Alt f)

instance Constrained (Const m)
instance CFunctor (Const m)
instance Constrained (App.WrappedArrow a b)
instance Arrow a => CFunctor (App.WrappedArrow a b)

instance Constrained ((->) r)
instance CFunctor ((->) r)
instance Constrained (K1 i c)
instance CFunctor (K1 i c)

instance Constrained (f :+: g) where
  type Dom (f :+: g) a = (Dom f a, Dom g a)
instance (CFunctor f, CFunctor g) => CFunctor (f :+: g) where
  cmap f (L1 xs) = L1 $ cmap f xs
  cmap f (R1 xs) = R1 $ cmap f xs
  {-# INLINE [1] cmap #-}
instance Constrained (f :*: g) where
  type Dom (f :*: g) a = (Dom f a, Dom g a)
instance (CFunctor f, CFunctor g) => CFunctor (f :*: g) where
  cmap f (l :*: r) = cmap f l :*: cmap f r
  {-# INLINE cmap #-}

instance Constrained (f :.: (g :: Type -> Type)) where
  type Dom (f :.: g) a = (Dom f (g a), Dom g a)
instance (CFunctor f, CFunctor g) => CFunctor (f :.: g) where
  cmap f gfa = Comp1 $ cmap (cmap f) $ unComp1 gfa
  {-# INLINE cmap #-}
instance (Constrained f, Constrained g) => Constrained (SOP.Sum f g) where
  type Dom (SOP.Sum f g) a = (Dom f a, Dom g a)

instance (CFunctor f, CFunctor g) => CFunctor (SOP.Sum f g) where
  cmap f (SOP.InL a) = SOP.InL $ cmap f a
  cmap f (SOP.InR b) = SOP.InR $ cmap f b
  {-# INLINE cmap #-}

  (<$:) = defaultCmapConst
  {-# INLINE (<$:) #-}

instance (Constrained f, Constrained g) => Constrained (SOP.Product f g) where
  type Dom (SOP.Product f g) a = (Dom f a, Dom g a)

instance (CFunctor f, CFunctor g) => CFunctor (SOP.Product f g) where
  cmap f (SOP.Pair a b) = SOP.Pair (cmap f a) (cmap f b)
  {-# INLINE cmap #-}

  (<$:) = defaultCmapConst
  {-# INLINE (<$:) #-}

instance (Constrained (f ::Type -> Type), Constrained (g :: Type -> Type))
  => Constrained (SOP.Compose f g) where
  type Dom (SOP.Compose f g) a = (Dom g a, Dom f (g a))

instance (CFunctor f, CFunctor g) => CFunctor (SOP.Compose f g) where
  cmap f (SOP.Compose a) = SOP.Compose $ cmap (cmap f) a
  (<$:) = defaultCmapConst

  {-# INLINE (<$:) #-}

instance Constrained (M1 i c f)
instance Functor f => CFunctor (M1 i c f)

instance Constrained Seq.Seq
instance CFunctor Seq.Seq

#if MIN_VERSION_mono_traversable(1,0,14)
instance Constrained (WrappedMono mono) where
  type Dom (WrappedMono mono) a = a ~ Element mono

instance MonoFunctor IS.IntSet where
  omap = IS.map

instance MonoFunctor mono => CFunctor (WrappedMono mono) where
  cmap = omap
  (<$:) = omap . const
#endif

instance Constrained (WrapMono mono) where
  type Dom (WrapMono mono) b = b ~ Element mono

instance {-# OVERLAPPABLE #-} MonoFunctor a
      => CFunctor (WrapMono a) where
  cmap = coerce @((Element a -> Element a) -> a -> a) omap
  {-# INLINE [1] cmap #-}

  (<$:) = defaultCmapConst
  {-# INLINE [1] (<$:) #-}

instance Constrained IM.IntMap
instance CFunctor IM.IntMap

instance Constrained (Map.Map k)
instance Ord k => CFunctor (Map.Map k)

instance Constrained Set.Set where
  type Dom Set.Set a = Ord a

instance CFunctor Set.Set where
  cmap = Set.map
  {-# INLINE [1] cmap #-}
  (<$:) = flip $ \s ->
    if Set.null s
    then const Set.empty else Set.singleton
  {-# INLINE [1] (<$:) #-}

instance Constrained HS.HashSet where
  type Dom HS.HashSet a = (Hashable a, Eq a)

instance CFunctor HS.HashSet where
  cmap :: (Hashable b, Eq b) => (a -> b) -> HS.HashSet a -> HS.HashSet b
  cmap = HS.map
  {-# INLINE [1] cmap #-}
  (<$:) = flip $ \s -> if HS.null s
    then const HS.empty else HS.singleton
  {-# INLINE (<$:) #-}

instance Constrained (HM.HashMap k)
instance CFunctor (HM.HashMap k)
instance Constrained Tree.Tree
instance CFunctor Tree.Tree

infixl 4 <$:>
(<$:>) :: (CFunctor f, Dom f a, Dom f b) => (a -> b) -> f a -> f b
(<$:>) = cmap
{-# INLINE [1] (<$:>) #-}

instance Constrained V.Vector where
  type Dom V.Vector a = ()

instance CFunctor V.Vector where
  cmap = V.map
  {-# INLINE [1] cmap #-}

instance Constrained U.Vector where
  type Dom U.Vector a = U.Unbox a
instance CFunctor U.Vector where
  cmap = U.map
  {-# INLINE [1] cmap #-}
instance Constrained S.Vector where
  type Dom S.Vector a = S.Storable a
instance CFunctor S.Vector where
  cmap = S.map
  {-# INLINE [1] cmap #-}

instance Constrained P.Vector where
  type Dom P.Vector a = P.Prim a
instance CFunctor P.Vector where
  cmap = P.map
  {-# INLINE [1] cmap #-}

instance Constrained PA.PrimArray where
  type Dom PA.PrimArray a = P.Prim a

instance CFunctor PA.PrimArray where
  cmap = PA.mapPrimArray
  {-# INLINE [1] cmap #-}

deriving via WrapFunctor SA.SmallArray
  instance Constrained SA.SmallArray
deriving via WrapFunctor SA.SmallArray
  instance CFunctor SA.SmallArray

deriving via WrapFunctor A.Array
  instance Constrained A.Array
deriving via WrapFunctor A.Array
  instance CFunctor A.Array


{- |
Copyright: (c) 2020 Thomas Tuegel
SPDX-License-Identifier: BSD-3-Clause
Maintainer: Thomas Tuegel <ttuegel@mailbox.org>

-}

module Injection
    ( Injection (..)
    , Retraction (..)
    ) where

import Data.Complex (Complex ((:+)))
import Data.Dynamic (Dynamic, Typeable, fromDynamic, toDyn)
import Data.Fixed (Fixed, HasResolution)
import Data.Functor.Const (Const (..))
import Data.Functor.Identity (Identity (..))
import Data.List.NonEmpty (NonEmpty (..))
import Data.Maybe (maybeToList)
import Data.Monoid (Dual (..))
import Data.Monoid (Product (..))
import Data.Monoid (Sum (..))
import Data.Monoid (Any (..))
import Data.Monoid (All (..))
import qualified Data.Monoid as Monoid (First (..), Last (..))
import Data.Ord (Down (..))
import Data.Ratio (Ratio)
import qualified Data.Ratio as Ratio
import Data.Semigroup (Max (..), Min (..))
import qualified Data.Semigroup as Semigroup (First (..), Last (..))
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Lazy as Lazy (Text)
import qualified Data.Text.Lazy as Text.Lazy
import Data.Void (Void)
import Numeric.Natural (Natural)

{- | @Injection@ describes a lossless conversion that includes one type in another.

The sole method of this class,

> inject :: from -> into

takes a value @input :: from@ and returns a value @output :: into@ which preserves all the information contained in the input.
Specifically, each @input@ is mapped to a /unique/ @output@.
In mathematical terminology, @inject@ is /injective/:

> inject a ≡ inject b → a ≡ b

The name of the class is derived from the mathematical term.

@Injection@ models the "is-a" relationship used in languages with subtypes (such as in object-oriented programming),
but an explicit cast with @inject@ is required in Haskell.

Although it is often possible to infer the type parameters of this class,
it is advisable to specify one or both of the parameters to @inject@
using a type signature or the @TypeApplications@ language extension.
Specifying the type parameters will give clearer error messages from the type checker in any case.

-}
class Injection from into where
    inject :: from -> into

{- | @Retraction@ undoes an 'Injection'.

Because 'Injection' is a lossless conversion, we can define a @Retraction@ which undoes it.
The method

> retract :: into -> Maybe from

is the (left) inverse of 'inject':

> retract (inject x) = Just x

'retract' is partial (returns 'Maybe') because the type @into@ may be larger than the type @from@;
that is, there may be values in @into@ which are not 'inject'-ed from @from@,
and in that case @retract@ may return 'Nothing'.

Although it is often possible to infer the type parameters of this class,
it is advisable to specify one or both of the parameters to @retract@
using a type signature or the @TypeApplications@ language extension.
Specifying the type parameters will give clearer error messages from the type checker in any case.

-}
class Injection from into => Retraction from into where
    retract :: into -> Maybe from

instance Injection a a where
    inject = id
    {-# INLINE inject #-}

instance Retraction a a where
    retract = Just
    {-# INLINE retract #-}

instance Typeable a => Injection a Dynamic where
    inject = toDyn
    {-# INLINE inject #-}

instance Typeable a => Retraction a Dynamic where
    retract = fromDynamic
    {-# INLINE retract #-}

instance Injection a b => Injection a (Maybe b) where
    inject = Just . inject
    {-# INLINE inject #-}

instance Retraction a b => Retraction a (Maybe b) where
    retract = \x -> x >>= retract @a @b
    {-# INLINE retract #-}

instance Injection a b => Injection (Maybe a) [b] where
    inject = maybeToList . fmap (inject @a @b)
    {-# INLINE inject #-}

instance Retraction a b => Retraction (Maybe a) [b] where
    retract [] = Just Nothing
    retract [b] = Just <$> retract @a @b b
    retract _ = Nothing

instance Injection Natural Integer where
    inject = toInteger
    {-# INLINE inject #-}

instance Retraction Natural Integer where
    retract x
        | x < 0 = Nothing
        | otherwise = Just (fromInteger x)
    {-# INLINE retract #-}

instance Injection Void any where
    inject = \case {}
    {-# INLINE inject #-}

-- | 'Text.unpack' is the canonical injection @'Text' -> 'String'@.
-- There is no injection in the other direction because 'String' can represent
-- invalid surrogate code points, but 'Text' cannot; for details, see
-- "Data.Text".
instance Injection Text String where
    inject = Text.unpack
    {-# INLINE inject #-}

-- | 'Text.Lazy.unpack' is the canonical injection @'Lazy.Text' -> 'String'@.
-- There is no injection in the other direction because 'String' can represent
-- invalid surrogate code points, but 'Lazy.Text' cannot; for details, see
-- "Data.Text.Lazy".
instance Injection Lazy.Text String where
    inject = Text.Lazy.unpack
    {-# INLINE inject #-}

instance Injection Text Lazy.Text where
    inject = Text.Lazy.fromStrict
    {-# INLINE inject #-}

instance Injection Lazy.Text Text where
    inject = Text.Lazy.toStrict
    {-# INLINE inject #-}

instance HasResolution a => Injection Integer (Fixed a) where
    inject = fromInteger
    {-# INLINE inject #-}

instance HasResolution a => Retraction Integer (Fixed a) where
    retract x = retract @Integer (toRational x)
    {-# INLINE retract #-}

instance Injection a (Const a b) where
    inject = Const
    {-# INLINE inject #-}

instance Injection (Const a b) a where
    inject = getConst
    {-# INLINE inject #-}

instance Injection Integer (Ratio Integer) where
    inject = fromInteger
    {-# INLINE inject #-}

instance Retraction Integer (Ratio Integer) where
    retract x
        | Ratio.denominator x == 1 = Just (Ratio.numerator x)
        | otherwise = Nothing
    {-# INLINE retract #-}

instance Num a => Injection a (Complex a) where
    inject = (:+ 0)
    {-# INLINE inject #-}

instance (Eq a, Num a) => Retraction a (Complex a) where
    retract (x :+ y)
      | y == 0 = Just x
      | otherwise = Nothing
    {-# INLINE retract #-}

instance Injection a (Identity a) where
    inject = Identity
    {-# INLINE inject #-}

instance Injection (Identity a) a where
    inject = runIdentity
    {-# INLINE inject #-}

instance Injection (NonEmpty a) [a] where
    inject (x :| xs) = x : xs
    {-# INLINE inject #-}

instance Retraction (NonEmpty a) [a] where
    retract (x : xs) = Just (x :| xs)
    retract [] = Nothing
    {-# INLINE retract #-}

instance Injection a (Down a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Down a) a where
    inject = \(Down a) -> a
    {-# INLINE inject #-}

instance Injection a (Product a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Product a) a where
    inject = getProduct
    {-# INLINE inject #-}

instance Injection a (Sum a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Sum a) a where
    inject = getSum
    {-# INLINE inject #-}

instance Injection a (Dual a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Dual a) a where
    inject = getDual
    {-# INLINE inject #-}

instance Injection a (Monoid.Last a) where
    inject = pure
    {-# INLINE inject #-}

instance Retraction a (Monoid.Last a) where
    retract = Monoid.getLast
    {-# INLINE retract #-}

instance Injection a (Monoid.First a) where
    inject = pure
    {-# INLINE inject #-}

instance Retraction a (Monoid.First a) where
    retract = Monoid.getFirst
    {-# INLINE retract #-}

instance Injection a (Semigroup.First a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Semigroup.First a) a where
    inject = Semigroup.getFirst
    {-# INLINE inject #-}

instance Injection a (Semigroup.Last a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Semigroup.Last a) a where
    inject = Semigroup.getLast
    {-# INLINE inject #-}

instance Injection a (Max a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Max a) a where
    inject = getMax
    {-# INLINE inject #-}

instance Injection a (Min a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Min a) a where
    inject = getMin
    {-# INLINE inject #-}

instance Injection a (r -> a) where
    inject = const
    {-# INLINE inject #-}

instance Injection Bool Any where
    inject = Any
    {-# INLINE inject #-}

instance Injection Any Bool where
    inject = getAny
    {-# INLINE inject #-}

instance Injection Bool All where
    inject = All
    {-# INLINE inject #-}

instance Injection All Bool where
    inject = getAll
    {-# INLINE inject #-}


{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, ScopedTypeVariables, RecordWildCards #-}

import           Control.Monad.Primitive
import           Control.Monad
import           Foreign.C.Types
import           Foreign.Ptr
import           Unsafe.Coerce
import           Data.Complex
import           Foreign.Marshal.Array

import qualified Data.Vector.Generic               as VG
import qualified Data.Vector.Generic.Mutable       as VGM
import qualified Data.Vector.Storable              as VS
import qualified Data.Vector.Storable.Mutable      as VSM
import qualified Data.Vector.Fusion.Stream         as VFS
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM

import           Foreign.Storable.Complex
import           Criterion.Main
import           Test.QuickCheck
import           Test.QuickCheck.Monadic

-- | A class for things that can be multiplied by a scalar.
class Mult a b where
    mult :: a -> b -> a

instance (Num a) => Mult a a where
    mult = (*)

instance (Num a) => Mult (Complex a) a where
    mult (x :+ y) z = (x * z) :+ (y * z)

-- | Fill a mutable vector from a monadic stream
{-# INLINE fill #-}
fill :: (PrimMonad m, Functor m, VGM.MVector vm a) => VFS.MStream m a -> vm (PrimState m) a -> m ()
fill str outBuf = void $ VFSM.foldM' put 0 str
    where
    put i x = do
        VGM.unsafeWrite outBuf i x
        return $ i + 1

{-# INLINE stride #-}
stride :: VG.Vector v a => Int -> v a -> v a
stride str inv = VG.unstream $ VFS.unfoldr func 0
    where
    len = VG.length inv
    func i | i >= len  = Nothing
           | otherwise = Just (VG.unsafeIndex inv i, i + str)

-- | The functions to be benchmarked

-- | Filters

filterHighLevel :: (PrimMonad m, Functor m, Num a, Mult a b, VG.Vector v a, VG.Vector v b, VGM.MVector vm a) => Int -> v b -> v a -> vm (PrimState m) a -> m ()
filterHighLevel num coeffs inBuf outBuf = fill (VFSM.generate num dotProd) outBuf
    where
    dotProd offset = VG.sum $ VG.zipWith mult (VG.unsafeDrop offset inBuf) coeffs

type FilterCRR = CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()
type FilterRR  = Int -> VS.Vector Float -> VS.Vector Float -> VS.MVector RealWorld Float -> IO ()
type FilterRC  = Int -> VS.Vector Float -> VS.Vector (Complex Float) -> VS.MVector RealWorld (Complex Float) -> IO ()

filterFFIR :: FilterCRR -> FilterRR
filterFFIR func num coeffs inBuf outBuf =
    VS.unsafeWith (unsafeCoerce coeffs) $ \cPtr ->
        VS.unsafeWith (unsafeCoerce inBuf) $ \iPtr ->
            VSM.unsafeWith (unsafeCoerce outBuf) $ \oPtr ->
                func (fromIntegral num) (fromIntegral $ VG.length coeffs) cPtr iPtr oPtr

filterFFIC :: FilterCRR -> FilterRC
filterFFIC func num coeffs inBuf outBuf =
    VS.unsafeWith (unsafeCoerce coeffs) $ \cPtr ->
        VS.unsafeWith (unsafeCoerce inBuf) $ \iPtr ->
            VSM.unsafeWith (unsafeCoerce outBuf) $ \oPtr ->
                func (fromIntegral num) (fromIntegral $ VG.length coeffs) cPtr iPtr oPtr

-- | Decimation

decimateHighLevel :: (PrimMonad m, Functor m, Num a, Mult a b, VG.Vector v a, VG.Vector v b, VGM.MVector vm a) => Int -> Int -> v b -> v a -> vm (PrimState m) a -> m ()
decimateHighLevel num factor coeffs inBuf outBuf = fill x outBuf
    where
    x = VFSM.map dotProd (VFSM.iterateN num (+ factor) 0)
    dotProd offset = VG.sum $ VG.zipWith mult (VG.unsafeDrop offset inBuf) coeffs

type DecimateCRR = CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()
type DecimateRR  = Int -> Int -> VS.Vector Float -> VS.Vector Float -> VS.MVector RealWorld Float -> IO ()
type DecimateRC  = Int -> Int -> VS.Vector Float -> VS.Vector (Complex Float) -> VS.MVector RealWorld (Complex Float) -> IO ()

decimateFFIR :: DecimateCRR -> DecimateRR
decimateFFIR func num factor coeffs inBuf outBuf =
    VS.unsafeWith (unsafeCoerce coeffs) $ \cPtr ->
        VS.unsafeWith (unsafeCoerce inBuf) $ \iPtr ->
            VSM.unsafeWith (unsafeCoerce outBuf) $ \oPtr ->
                func (fromIntegral num) (fromIntegral factor) (fromIntegral $ VG.length coeffs) cPtr iPtr oPtr

decimateFFIC :: DecimateCRR -> DecimateRC
decimateFFIC func num factor coeffs inBuf outBuf =
    VS.unsafeWith (unsafeCoerce coeffs) $ \cPtr ->
        VS.unsafeWith (unsafeCoerce inBuf) $ \iPtr ->
            VSM.unsafeWith (unsafeCoerce outBuf) $ \oPtr ->
                func (fromIntegral num) (fromIntegral factor) (fromIntegral $ VG.length coeffs) cPtr iPtr oPtr

-- | Rational downsampling
resampleHighLevel :: (PrimMonad m, Num a, Mult a b, VG.Vector v a, VG.Vector v b, VGM.MVector vm a) => Int -> Int -> Int -> Int -> v b -> v a -> vm (PrimState m) a -> m Int
resampleHighLevel count interpolation decimation filterOffset coeffs inBuf outBuf = fill 0 filterOffset 0
    where
    fill i filterOffset inputOffset
        | i < count = do
            let dp = dotProd filterOffset inputOffset
            VGM.unsafeWrite outBuf i dp
            let (q, r)        = quotRem (decimation - filterOffset - 1) interpolation
                inputOffset'  = inputOffset + q + 1
                filterOffset' = interpolation - 1 - r
            filterOffset' `seq` inputOffset' `seq` fill (i + 1) filterOffset' inputOffset'
        | otherwise = return filterOffset
    dotProd filterOffset offset = VG.sum $ VG.zipWith mult (VG.unsafeDrop offset inBuf) (stride interpolation (VG.unsafeDrop filterOffset coeffs))

pad :: a -> Int -> [a] -> [a]
pad with num list = list ++ replicate (num - length list) with

strideList :: Int -> [a] -> [a]
strideList s xs = go 0 xs
    where
    go _ []     = []
    go 0 (x:xs) = x : go (s-1) xs
    go n (x:xs) = go (n - 1) xs

roundUp :: Int -> Int -> Int
roundUp num div = ((num + div - 1) `quot` div) * div

data Coeffs = Coeffs {
    numCoeffs  :: Int,
    numGroups  :: Int,
    increments :: [Int],
    groups     :: [[Float]]
}

prepareCoeffs :: Int -> Int -> Int -> [Float] -> Coeffs
prepareCoeffs n interpolation decimation coeffs = Coeffs {..}
    where
    numCoeffs   = maximum $ map (length . snd) dats
    numGroups   = length groups
    increments  = map fst dats

    groups      :: [[Float]]
    groups      = map (pad 0 (roundUp numCoeffs n)) $ map snd dats

    dats :: [(Int, [Float])]
    dats = func 0
        where

        func' 0      = []
        func' x      = func x

        func :: Int -> [(Int, [Float])]
        func offset = (increment, strideList interpolation $ drop offset coeffs) : func' offset'
            where
            (q, r)    = quotRem (decimation - offset - 1) interpolation
            increment = q + 1
            offset'   = interpolation - 1 - r

resampleFFIR :: (Ptr CFloat -> Ptr CFloat -> IO ()) -> VS.Vector Float -> VSM.MVector RealWorld Float -> IO ()
resampleFFIR func inBuf outBuf =
    VS.unsafeWith (unsafeCoerce inBuf) $ \iPtr ->
        VS.unsafeWith (unsafeCoerce outBuf) $ \oPtr ->
            func iPtr oPtr

type ResampleR = CInt -> CInt -> CInt -> CInt -> Ptr CInt -> Ptr (Ptr CFloat) -> Ptr CFloat -> Ptr CFloat -> IO ()

mkResampler :: ResampleR -> Int -> Int -> Int -> [Float] -> IO (Int -> Int -> VS.Vector Float -> VS.MVector RealWorld Float -> IO ())
mkResampler func n interpolation decimation coeffs = do
    groupsP     <- mapM newArray $ map (map realToFrac) groups
    groupsPP    <- newArray groupsP
    incrementsP <- newArray $ map fromIntegral increments
    return $ \num offset -> resampleFFIR $ func (fromIntegral num) (fromIntegral numCoeffs) (fromIntegral offset) (fromIntegral numGroups) incrementsP groupsPP
    where
    Coeffs {..} = prepareCoeffs n interpolation decimation coeffs

theBench :: IO ()
theBench = do
    --Setup
    let size       =  16384
        numCoeffs  =  128
        num        =  size - numCoeffs + 1
        decimation =  4
        interpolation = 3
        numCoeffsDiv2  =  64

        coeffsList :: [Float]
        coeffsList = take numCoeffs [0 ..]
        coeffs    :: VS.Vector Float
        coeffs    =  VG.fromList $ take numCoeffs [0 ..]
        coeffsSym    :: VS.Vector Float
        coeffsSym    =  VG.fromList $ take numCoeffsDiv2 [0 ..]
        inBuf     :: VS.Vector Float
        inBuf     =  VG.fromList $ take size [0 ..]
        inBufComplex :: VS.Vector (Complex Float)
        inBufComplex =  VG.fromList $ take size $ do
            i <- [0..]
            return $ i :+ i

        numConv   = 16386
        inBufConv :: VS.Vector CUChar
        inBufConv = VG.fromList $ take size $ concat $ repeat [0 .. 255]

        duplicate :: [a] -> [a]
        duplicate = concatMap func
            where func x = [x, x]

        coeffs2 :: VS.Vector Float
        coeffs2 =  VG.fromList $ duplicate $ take numCoeffs [0 ..]

    outBuf        :: VS.MVector RealWorld Float <- VGM.new size
    outBufComplex :: VS.MVector RealWorld (Complex Float) <- VGM.new size

    --Benchmarks
    defaultMain [
            bgroup "filter" [
                bgroup "real" [
                    bench "highLevel"   $ nfIO $ filterHighLevel        num coeffs inBuf outBuf
                ],
                bgroup "complex" [
                    bench "highLevel"   $ nfIO $ filterHighLevel        num coeffs  inBufComplex outBufComplex
                ]
            ],
            bgroup "decimate" [
                bgroup "real" [
                    bench "highLevel"   $ nfIO $ decimateHighLevel        (num `quot` decimation) decimation coeffs inBuf outBuf
                ],
                bgroup "complex" [
                    bench "highLevel"   $ nfIO $ decimateHighLevel      (num `quot` decimation) decimation coeffs  inBufComplex outBufComplex
                ]
            ],
            bgroup "resample" [
                bgroup "real" [
                    bench "highLevel"   $ nfIO $ resampleHighLevel      (num `quot` decimation) interpolation decimation 0 coeffs inBuf outBuf
                ],
                bgroup "complex" [
                    bench "highLevel"   $ nfIO $ resampleHighLevel      (num `quot` decimation) interpolation decimation 0 coeffs inBufComplex outBufComplex
                ]
            ]
        ]

theTest = quickCheck $ conjoin [propFiltersComplex]
    where
    sizes           = elements [1024, 2048, 4096, 8192, 16384, 32768, 65536]
    numCoeffs       = elements [32, 64, 128, 256, 512]
    factors         = elements [1, 2, 3, 4, 7, 9, 12, 15, 21]
    factors'        = [1, 2, 3, 4, 7, 9, 12, 15, 21]
    propFiltersReal = forAll sizes $ \size ->
                          forAll (vectorOf size (choose (-10, 10))) $ \inBuf ->
                              forAll numCoeffs $ \numCoeffs ->
                                  forAll (vectorOf numCoeffs (choose (-10, 10))) $ \coeffs ->
                                      testFiltersReal size numCoeffs coeffs inBuf
    testFiltersReal :: Int -> Int -> [Float] -> [Float] -> Property
    testFiltersReal size numCoeffs coeffs inBuf = monadicIO $ do
        let vCoeffsHalf = VS.fromList coeffs
            vCoeffs     = VS.fromList $ coeffs ++ reverse coeffs
            vInput      = VS.fromList inBuf
            num         = size - numCoeffs*2 + 1

        r1 <- run $ getResult num $ filterHighLevel       num vCoeffs     vInput

    propFiltersComplex = forAll sizes $ \size ->
                             forAll (vectorOf size (choose (-10, 10))) $ \inBufR ->
                                 forAll (vectorOf size (choose (-10, 10))) $ \inBufI ->
                                     forAll numCoeffs $ \numCoeffs ->
                                         forAll (vectorOf numCoeffs (choose (-10, 10))) $ \coeffs ->
        assert $ all (r1 `eqDelta`) [r1]
                                             testFiltersComplex size numCoeffs coeffs $ zipWith (:+) inBufR inBufI
    testFiltersComplex :: Int -> Int -> [Float] -> [Complex Float] -> Property
    testFiltersComplex size numCoeffs coeffs inBuf = monadicIO $ do
        let vCoeffsHalf = VS.fromList coeffs
            vCoeffs     = VS.fromList $ coeffs ++ reverse coeffs
            vInput      = VS.fromList inBuf
            num         = size - numCoeffs*2 + 1
            --vCoeffs2    = VG.fromList $ duplicate $ coeffs ++ reverse coeffs

        r1 <- run $ getResult num $ filterHighLevel       num vCoeffs     vInput

    propDecimationReal = forAll sizes $ \size ->
                             forAll (vectorOf size (choose (-10, 10))) $ \inBuf ->
                                 forAll numCoeffs $ \numCoeffs ->
                                     forAll (vectorOf numCoeffs (choose (-10, 10))) $ \coeffs ->
                                        forAll factors $ \factor ->
        assert $ all (r1 `eqDeltaC`) [r1]
                                             testDecimationReal size numCoeffs factor coeffs inBuf
    testDecimationReal :: Int -> Int -> Int -> [Float] -> [Float] -> Property
    testDecimationReal size numCoeffs factor coeffs inBuf = monadicIO $ do
        let vCoeffsHalf = VS.fromList coeffs
            vCoeffs     = VS.fromList $ coeffs ++ reverse coeffs
            vInput      = VS.fromList inBuf
            num         = (size - numCoeffs*2 + 1) `quot` factor

        r1 <- run $ getResult num $ decimateHighLevel       num factor vCoeffs     vInput

    propDecimationComplex = forAll sizes $ \size ->
                                forAll (vectorOf size (choose (-10, 10))) $ \inBufR ->
                                    forAll (vectorOf size (choose (-10, 10))) $ \inBufI ->
                                        forAll numCoeffs $ \numCoeffs ->
                                            forAll (vectorOf numCoeffs (choose (-10, 10))) $ \coeffs ->
                                                forAll factors $ \factor ->
        assert $ all (r1 `eqDelta`) [r1]
                                                    testDecimationComplex size numCoeffs factor coeffs $ zipWith (:+) inBufR inBufI
    testDecimationComplex :: Int -> Int -> Int -> [Float] -> [Complex Float] -> Property
    testDecimationComplex size numCoeffs factor coeffs inBuf = monadicIO $ do
        let vCoeffsHalf = VS.fromList coeffs
            vCoeffs     = VS.fromList $ coeffs ++ reverse coeffs
            vInput      = VS.fromList inBuf
            num         = (size - numCoeffs*2 + 1) `quot` factor
            --vCoeffs2    = VG.fromList $ duplicate $ coeffs ++ reverse coeffs

        r1 <- run $ getResult num $ decimateHighLevel       num factor vCoeffs     vInput

    propResamplingReal = forAll sizes $ \size ->
                             forAll (vectorOf size (choose (-10, 10))) $ \inBuf ->
                                 forAll numCoeffs $ \numCoeffs ->
                                     forAll (vectorOf numCoeffs (choose (-10, 10))) $ \coeffs ->
                                        forAll (elements $ tail factors') $ \decimation ->
                                            forAll (elements $ filter (< decimation) factors') $ \interpolation ->
        assert $ all (r1 `eqDeltaC`) [r1]
                                                 testResamplingReal size numCoeffs interpolation decimation coeffs inBuf
    testResamplingReal :: Int -> Int -> Int -> Int -> [Float] -> [Float] -> Property
    testResamplingReal size numCoeffs interpolation decimation coeffs inBuf = monadicIO $ do
        let vCoeffsHalf = VS.fromList coeffs
            vCoeffs     = VS.fromList $ coeffs ++ reverse coeffs
            vInput      = VS.fromList inBuf
            num         = (size - numCoeffs*2 + 1) `quot` decimation

        r1 <- run $ getResult num $ resampleHighLevel       num interpolation decimation 0 vCoeffs vInput

        assert $ all (r1 `eqDelta`) [r1]
    getResult :: (VSM.Storable a) => Int -> (VS.MVector RealWorld a -> IO b) -> IO [a]
    getResult size func = do
        outBuf <- VGM.new size
        func outBuf
        out :: VS.Vector a <- VG.freeze outBuf
        return $ VG.toList out
    eqDelta x y = all (uncurry eqDelta') $ zip x y
        where
        eqDelta' x y = abs (x - y) < 0.01
    eqDeltaC x y = all (uncurry eqDelta') $ zip x y
        where
        eqDelta' x y = magnitude (x - y) < 0.01
    duplicate :: [a] -> [a]
    duplicate = concatMap func
        where func x = [x, x]

main = theBench


{-# LANGUAGE TypeFamilies #-}
-- | Tools for implementing (and debugging the use of) gradient descent schemes.
module HordeAd.Core.OptimizerTools
  ( updateWithGradient
  , gradientIsNil, minimumGradient, maximumGradient
  , ArgsAdam(..), defaultArgsAdam
  , StateAdam(..), initialStateAdam
  , liftMatrix43, updateWithGradientAdam
  ) where

import Prelude

import           Control.Monad.ST.Strict (runST)
import qualified Data.Array.DynamicS as OT
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as VM
import           HordeAd.Internal.OrthotopeOrphanInstances (liftVT2)
import           Numeric.LinearAlgebra (Element, Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM
import           Numeric.LinearAlgebra.Data (flatten)
import           Numeric.LinearAlgebra.Devel
  (MatrixOrder (..), liftMatrix, liftMatrix2, matrixFromVector, orderOf)

import HordeAd.Internal.Delta (Domains, isTensorDummy)

{-
60% of heap allocation in matrix- and vector-based MNIST
with simple gradient descent (no mini-batches) is performed
by @updateWithGradient.updateVector@ below

  let updateVector i r = i - HM.scale gamma r

due to allocating once in @scale@ and again in @-@
(and there seems to be one more allocation judging by the numbers).
Something like the following code would be needed to eliminate
one allocation, but it would requre the hmatrix maintainer to expose
internal modules.

import           Internal.Vectorized
  ( FunCodeSV (Scale), FunCodeVV (Sub), applyRaw, c_vectorMapValR
  , c_vectorZipR, createVector )

minusTimesGamma :: Storable r => r -> Vector r -> Vector r -> Vector r
minusTimesGamma gamma u v = unsafePerformIO $ do
  r <- createVector (dim0 u)
  pval <- newArray [gamma]
  (v `applyRaw` (r `applyRaw` id))
    (c_vectorMapValR (fromei Scale) pval)
    #| "minusTimesGamma1"
  free pval
  (u `applyRaw` (v `applyRaw` (r `applyRaw` id)))
    (c_vectorZipR (fromei Sub))
    #| "minusTimesGamma2"
  return r

BTW, a version with HM.Devel.zipVectorWith makes
the test twice slower and allocate twice more

  let updateVector = zipVectorWith (\i r -> i - gamma * r)

and a version with Vector.Storable makes the test thrice slower
and allocate thrice more

  let updateVector = V.zipWith (\i r -> i - gamma * r)

which is probably a bug in stream fusion which, additionally in this case,
can't fuse with anything and so can't pay for its overhead.

-}

updateWithGradient :: (Numeric r, Floating (Vector r))
                   => r -> Domains r -> Domains r -> Domains r
updateWithGradient gamma (params0, params1, params2, paramsX)
                         (gradient0, gradient1, gradient2, gradientX) =
  let updateVector i r = i - HM.scale gamma r
      !params0New = updateVector params0 gradient0
      update1 i r = if V.null r  -- eval didn't update it, would crash
                    then i
                    else updateVector i r
      !params1New = V.zipWith update1 params1 gradient1
      update2 i r = if HM.rows r <= 0  -- eval didn't update it, would crash
                    then i
                    else liftMatrix2 updateVector i r
      !params2New = V.zipWith update2 params2 gradient2
      updateX i r = if isTensorDummy r  -- eval didn't update it, would crash
                    then i
                    else liftVT2 updateVector i r
                      -- TODO: this is slow; add @liftArray2@ and use HM,
                      -- unless we move away from HM; similarly other OT calls
      !paramsXNew = V.zipWith updateX paramsX gradientX
  in (params0New, params1New, params2New, paramsXNew)

gradientIsNil :: (Eq r, Numeric r)
              => Domains r -> Bool
gradientIsNil (gradient0, gradient1, gradient2, gradientX) =
  V.all (== 0) gradient0
  && V.all V.null gradient1
  && V.all (\r -> HM.rows r <= 0) gradient2
  && V.all isTensorDummy gradientX

minimumGradient :: (Ord r, Numeric r)
                => Domains r -> r
minimumGradient (gradient0, gradient1, gradient2, gradientX) =
  min (if V.null gradient0 then 0 else HM.minElement gradient0)
      (min (if V.null gradient1 then 0
            else V.minimum (V.map HM.minElement gradient1))
           (min (if V.null gradient2 then 0
                 else V.minimum (V.map HM.minElement gradient2))
                (if V.null gradientX then 0
                 else V.minimum (V.map OT.minimumA gradientX))))

maximumGradient :: (Ord r, Numeric r)
                => Domains r -> r
maximumGradient (gradient0, gradient1, gradient2, gradientX) =
  max (if V.null gradient0 then 0 else HM.maxElement gradient0)
      (max (if V.null gradient1 then 0
            else V.maximum (V.map HM.maxElement gradient1))
           (max (if V.null gradient2 then 0
                 else V.maximum (V.map HM.maxElement gradient2))
                (if V.null gradientX then 0
                 else V.maximum (V.map OT.maximumA gradientX))))

data ArgsAdam r = ArgsAdam
  { alpha   :: r
  , betaOne :: r
  , betaTwo :: r
  , epsilon :: r
  }

-- The defaults taken from
-- https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
defaultArgsAdam :: Fractional r => ArgsAdam r
defaultArgsAdam = ArgsAdam
  { alpha = 0.001
  , betaOne = 0.9
  , betaTwo = 0.999
  , epsilon = 1e-7
  }

data StateAdam r = StateAdam
  { tAdam :: Int  -- iteration count
  , mAdam :: Domains r
  , vAdam :: Domains r
  }

-- The arguments are just sample params0, for dimensions.
zeroParameters :: Numeric r
               => Domains r -> Domains r
zeroParameters (params0, params1, params2, paramsX) =
  let zeroVector v = runST $ do
        vThawed <- V.thaw v
        VM.set vThawed 0
        V.unsafeFreeze vThawed
  in ( zeroVector params0
     , V.map zeroVector params1
     , V.map (liftMatrix zeroVector) params2
     , V.map (\a -> OT.constant (OT.shapeL a) 0) paramsX )  -- fast allright

initialStateAdam :: Numeric r
                 => Domains r -> StateAdam r
initialStateAdam parameters0 =
  let zeroP = zeroParameters parameters0
  in StateAdam
       { tAdam = 0
       , mAdam = zeroP
       , vAdam = zeroP
       }

-- | Application of a vector function on the flattened matrices elements.
liftMatrix43 :: ( Numeric a, Numeric b, Numeric c, Numeric d
                , Element x, Element y, Element z )
             => (Vector a -> Vector b -> Vector c -> Vector d
                 -> (Vector x, Vector y, Vector z))
             -> Matrix a -> Matrix b -> Matrix c -> Matrix d
             -> (Matrix x, Matrix y, Matrix z)
liftMatrix43 f m1 m2 m3 m4 =
  let sz@(r, c) = HM.size m1
      rowOrder = orderOf m1
        -- checking @m4@ (gradient) makes RNN LL test much faster and BB slower
        -- so this needs much more benchmarking and understading to tweak
  in if sz == HM.size m2 && sz == HM.size m3 && sz == HM.size m4
     then
       let (vx, vy, vz) = case rowOrder of
             RowMajor -> f (flatten m1) (flatten m2) (flatten m3) (flatten m4)
             ColumnMajor -> f (flatten (HM.tr' m1)) (flatten (HM.tr' m2))
                              (flatten (HM.tr' m3)) (flatten (HM.tr' m4))
       in ( matrixFromVector rowOrder r c vx
          , matrixFromVector rowOrder r c vy
          , matrixFromVector rowOrder r c vz
          )
     else error $ "nonconformant matrices in liftMatrix43: "
                  ++ show (HM.size m1, HM.size m2, HM.size m3, HM.size m4)

-- TOOD: make sure this is not worse that OT.zipWith3A when transposing
-- between each application or that we never encounter such situations
--
-- | Application of a vector function on the flattened arrays elements.
liftArray43 :: ( Numeric a, Numeric b, Numeric c, Numeric d
               , Numeric x, Numeric y, Numeric z )
            => (Vector a -> Vector b -> Vector c -> Vector d
                -> (Vector x, Vector y, Vector z))
            -> OT.Array a -> OT.Array b -> OT.Array c -> OT.Array d
            -> (OT.Array x, OT.Array y, OT.Array z)
liftArray43 f m1 m2 m3 m4 =
  let sz = OT.shapeL m1
  in if sz == OT.shapeL m2 && sz == OT.shapeL m3 && sz == OT.shapeL m4
     then
       let (vx, vy, vz) = f (OT.toVector m1) (OT.toVector m2)
                            (OT.toVector m3) (OT.toVector m4)
       in ( OT.fromVector sz vx
          , OT.fromVector sz vy
          , OT.fromVector sz vz
          )
     else error
          $ "nonconformant arrays in liftArray43: "
            ++ show (OT.shapeL m1, OT.shapeL m2, OT.shapeL m3, OT.shapeL m4)

updateWithGradientAdam
  :: forall r. (Numeric r, Floating r, Floating (Vector r))
  => ArgsAdam r -> StateAdam r -> Domains r -> Domains r
  -> (Domains r, StateAdam r)
updateWithGradientAdam ArgsAdam{..}
                       StateAdam{ tAdam
                                , mAdam = (mAdam0, mAdam1, mAdam2, mAdamX)
                                , vAdam = (vAdam0, vAdam1, vAdam2, vAdamX)
                                }
                       (params0, params1, params2, paramsX)
                       (gradient0, gradient1, gradient2, gradientX) =
  let tAdamNew = tAdam + 1
      oneMinusBeta1 = 1 - betaOne
      oneMinusBeta2 = 1 - betaTwo
      updateVector :: Vector r -> Vector r
                   -> Vector r -> Vector r
                   -> (Vector r, Vector r, Vector r)
      updateVector mA vA p g =
        let mANew = HM.scale betaOne mA + HM.scale oneMinusBeta1 g
            vANew = HM.scale betaTwo vA + HM.scale oneMinusBeta2 (g * g)
            alphat = alpha * sqrt (1 - betaTwo ^ tAdamNew)
                             / (1 - betaOne ^ tAdamNew)
        in ( mANew
           , vANew
           , p - HM.scale alphat mANew
                 / (sqrt vANew + HM.scalar epsilon) )  -- the @scalar@ is safe
                      -- @addConstant@ would be better, but it's not exposed
      (!mAdam0New, !vAdam0New, !params0New) =
        updateVector mAdam0 vAdam0 params0 gradient0
      update1 mA vA p g = if V.null g  -- eval didn't update it, would crash
                          then (mA, vA, p)
                          else updateVector mA vA p g
      (!mAdam1New, !vAdam1New, !params1New) =
        V.unzip3 $ V.zipWith4 update1 mAdam1 vAdam1 params1 gradient1
      update2 mA vA p g = if HM.rows g <= 0  -- eval didn't update it; crash
                          then (mA, vA, p)
                          else liftMatrix43 updateVector mA vA p g
      (!mAdam2New, !vAdam2New, !params2New) =
        V.unzip3 $ V.zipWith4 update2 mAdam2 vAdam2 params2 gradient2
      updateX mA vA p g = if isTensorDummy g  -- eval didn't update it
                          then (mA, vA, p)
                          else liftArray43 updateVector mA vA p g
      (!mAdamXNew, !vAdamXNew, !paramsXNew) =
        V.unzip3 $ V.zipWith4 updateX mAdamX vAdamX paramsX gradientX
  in ( (params0New, params1New, params2New, paramsXNew)
     , StateAdam { tAdam = tAdamNew
                 , mAdam = (mAdam0New, mAdam1New, mAdam2New, mAdamXNew)
                 , vAdam = (vAdam0New, vAdam1New, vAdam2New, vAdamXNew)
                 }
     )


{-# LANGUAGE OverloadedLists #-}

module OBJReader where

import qualified Data.HashMap as HM
import qualified Data.Vector.Storable as VS
import qualified Data.Vector as V
import           Data.Vector.Storable ((!))
import           Control.Arrow
import           Control.Lens
import           Data.List (nub, nubBy)
import qualified Numeric.LinearAlgebra as L
import qualified Graphics.Rendering.OpenGL.GL as GL
import           Data.Function (on)
import Debug.Trace

import Physics.Collision

splitOn x =
  foldr (\c -> if c == x then ([]:) else over (ix 0) (c:)) [[]]

addP (a,b) (c,d) = (a+b,c+d)

getFace :: [String] -> [(Int, Int)]
getFace =
  map ((\[v,_,n] -> (read v - 1,read n - 1)) . splitOn '/') . tail

readVectors :: String -> [[String]] -> [VS.Vector Float]
readVectors vtype =
  map (VS.fromList.map read.tail) . filter ((==vtype) . head)

readPoints = readVectors "v"
readNormals = readVectors "vn"

readBoxFromObj :: FilePath -> IO Primitive
readBoxFromObj fn = do
  text <- map words . lines <$> readFile fn
  let p = V.fromList . readPoints $ text
      q = V.fromList . readNormals $  text
      f = map (V.unsafeIndex p *** V.unsafeIndex q >>> uncurry (flip Squad)) .
          nubBy ((==) `on` snd) .
          map (head . getFace) .
          filter ((=="f").head) $ text
  return $ Box p (V.fromList f)

readOBJ :: FilePath -> IO (VS.Vector Float, VS.Vector GL.BaseInstance,[VS.Vector Float])
readOBJ fn = do
  text <- map words . lines <$> readFile fn
  let
    verts    = readPoints text
    normals :: [VS.Vector Float]
    normals  = readNormals text
    faces    = map getFace . filter ((=="f").head) $ text
    vertData = nub . concat $ faces
    vertIxMap= HM.fromList . flip zip [0..] $ vertData
    --                 u v color
    addColor = (VS.++ [1,1,0.5,0.5,0.5])
    computeT :: [(Int,Int)] -> VS.Vector Float
    computeT [v0, v1, v2] =
      let dpos1 = verts !! fst v1 - verts !! fst v0
          -- dpos2 = verts !! fst v2 - verts !! fst v0
          tang  = dpos1
          -- bitan = dpos2
      in tang
    nullTans = HM.fromList . flip zip (repeat $ VS.replicate 3 0) $ vertData
    insertT l h = let p = computeT l
                  in foldr (\k -> HM.insertWith (+) k p) h l
    tangents = foldr insertT nullTans faces
    normalize :: (Int, Int) -> VS.Vector Float -> VS.Vector Float
    normalize (_,k) t =
      let n                 = normals !! k
          normalizedTangent = normalizeVec $ t - L.scale (L.dot n t) n
          normalizedBitangent = L.cross normalizedTangent n
      in -- n VS.++ t
        normalizedTangent VS.++ normalizedBitangent
    normalizedTangents = HM.fromList . map (fst &&& uncurry normalize) . HM.toList $ tangents
    faceProc p = (verts!!) *** (normals!!) >>> uncurry (VS.++) >>> addColor >>> flip (VS.++) (normalizedTangents HM.! p) $ p
    vBuffData = VS.concat $ map faceProc vertData
    vertices  = map ((verts!!).fst) vertData
    indices  = VS.fromList . map (vertIxMap HM.!) . concat $ faces
--      result   = VS.concat $ map faceProc faces
  return (vBuffData, indices, vertices)

normalizeVec :: VS.Vector Float -> VS.Vector Float
normalizeVec v = L.scale (realToFrac . (1/) $ L.norm_2 v) v

-- compute


import System.Environment
import Data.Complex
import Paraiso

main = do
  args <- getArgs
  let arch = if "--cuda" `elem` args then
               CUDA 128 128
             else
               X86
  putStrLn $ compile arch code
    where
      code = do
         parallel 16384 $ do
           c <- allocate
           z <- allocate
           c =$ (Rand (-2.0) 2.0) :+ (Rand (-2.0) 2.0)
           z =$ (0 :+ 0 :: Complex (Expr Double))
           cuda $ do
             sequential (65536) $ do
               z =$ z * z + c
           output [realPart c,imagPart c, realPart z, imagPart z]


{-# LANGUAGE NamedFieldPuns #-}

module Zambini where

import Data.Complex
import Data.Fixed (mod')

import Graphics.Gloss.Interface.Pure.Simulate

import ViewUtil

type R = Double
type C = Complex Double

data Zambini = Zambini { m :: R -- physical attributes
                       , i :: R
                       , r :: R
                       , fMax :: R
                       -- input
                       , f1 :: R
                       , f2 :: R
                       -- state
                       , x :: C
                       , v :: C
                       , theta :: R
                       , omega :: R
                       }

step :: R -> Zambini -> Zambini
step dt c@Zambini{m, i, r, fMax, f1, f2, x, v, theta, omega} = c{x=x', v=v', theta=theta', omega=omega'}
  where
    f1' = min 1 f1
    f2' = min 1 f2
    a = mkPolar ((f1' + f2') / m) theta
    v' = v + a * realToFrac dt
    x' = x + v' * realToFrac dt
    torque = r * (f1' - f2')
    omega' = omega + torque / i * dt
    theta' = theta + omega' * dt `mod'` 2*pi

render :: Zambini -> Picture
render Zambini{f1, f2, x, theta} = move x $ rotateRad (realToFrac theta) $ pictures [zambiniPic, control1 f1, control2 f2]
  where
    zambiniPic = color (greyN 0.5) $ pictures [ polygon [(-5,0), (-10, 25), (10, 25), (5, 0), (10, -25), (-10, -25)]
                                              , polygon [(0, 10), (20, 0), (0, -10)]
                                              ]

    control = rotate 180 . arrow (5, 20) (15, 100) (low, high)
      where
        low  = makeColorI 0 64 192 255
        high = makeColorI 0 160 192 255

    control1 = translate (-20) (-20) . control
    control2 = translate (-20) (20) . control

    move (x :+ y) = translate (realToFrac x) (realToFrac y)


{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# HLINT ignore "Use camelCase"      #-}
{-|
Module      : Grenade.Layers.Internal.Add
Description : Fast functions for performing convolutions and helper functions
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Internal.Convolution (
  -- * im2col functions
    im2col
  , col2im
  , col2vid
  , vid2col

  -- * convolution functions
  , forwardConv2d
  , forwardBiasConv2d
  , backwardConv2d
  , backwardBiasConv2d
  ) where

import qualified Data.Vector.Storable        as U (unsafeFromForeignPtr0,
                                                   unsafeToForeignPtr0)

import           Foreign                     (mallocForeignPtrArray,
                                              withForeignPtr)
import           Foreign.Ptr                 (Ptr)
import           Grenade.Types
import           Numeric.LinearAlgebra       (Matrix, Vector, cols, flatten,
                                              rows)
import qualified Numeric.LinearAlgebra       as LA
import qualified Numeric.LinearAlgebra.Devel as U
import           System.IO.Unsafe            (unsafePerformIO)

import Grenade.Layers.Internal.Hmatrix

-- | Rearrange the matrix columns into blocks, calculates the number of channels automatically.
col2vid :: Int            -- ^ kernel rows
        -> Int            -- ^ kernel columns
        -> Int            -- ^ stride rows
        -> Int            -- ^ stride columns
        -> Int            -- ^ input height
        -> Int            -- ^ input width
        -> Matrix RealNum -- ^ input matrix
        -> Matrix RealNum -- ^ output matrix
col2vid kernelRows kernelColumns strideRows strideColumns height width dataCol =
  let channels = cols dataCol `div` (kernelRows * kernelColumns)
      outRows  = (height - kernelRows) `div` strideRows + 1
      outCols  = (width - kernelColumns) `div` strideColumns + 1
  in  col2im_c channels height width kernelRows kernelColumns strideRows strideColumns 0 0 outRows outCols dataCol

-- | Rearrange the matrix columns into blocks, assumes there is only one channel
col2im :: Int            -- ^ kernel rows
       -> Int            -- ^ kernel columns
       -> Int            -- ^ stride rows
       -> Int            -- ^ stride columns
       -> Int            -- ^ input height
       -> Int            -- ^ input width
       -> Matrix RealNum -- ^ input matrix
       -> Matrix RealNum -- ^ output matrix
col2im kernelRows kernelColumns strideRows strideColumns height width dataCol =
  let channels = 1
      outRows  = (height - kernelRows) `div` strideRows + 1
      outCols  = (width - kernelColumns) `div` strideColumns + 1
  in  col2im_c channels height width kernelRows kernelColumns strideRows strideColumns 0 0 outRows outCols dataCol

col2im_c :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2im_c channels height width kernelRows kernelColumns strideRows strideColumns padl padt outRows outCols dataCol =
  let vec = flatten dataCol
  in unsafePerformIO $ do
    outPtr <- mallocForeignPtrArray (height * width * channels)
    let (inPtr, _) = U.unsafeToForeignPtr0 vec

    withForeignPtr inPtr $ \inPtr' ->
      withForeignPtr outPtr $ \outPtr' ->
        col2im_cpu inPtr' channels height width kernelRows kernelColumns strideRows strideColumns padt padl outRows outCols outPtr'

    let matVec = U.unsafeFromForeignPtr0 outPtr (height * width * channels)
    return $ U.matrixFromVector U.RowMajor (height * channels) width matVec
{-# INLINE col2im_c #-}

foreign import ccall unsafe
    col2im_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr RealNum -> IO ()

-- | Rearrange the matrix blocks into columns, calculates the number of channels automatically.
vid2col :: Int            -- ^ kernel rows
        -> Int            -- ^ kernel columns
        -> Int            -- ^ stride rows
        -> Int            -- ^ stride columns
        -> Int            -- ^ input height
        -> Int            -- ^ input width
        -> Matrix RealNum -- ^ input matrix
        -> Matrix RealNum -- ^ output matrix
vid2col kernelRows kernelColumns strideRows strideColumns height width dataVid =
  let channels = rows dataVid `div` height
      rowOut          = (height - kernelRows) `div` strideRows + 1
      colOut          = (width - kernelColumns) `div` strideColumns + 1
  in  im2col_c channels height width kernelRows kernelColumns strideRows strideColumns 0 0 dataVid rowOut colOut

-- | Rearrange the matrix blocks into columns, assumes there is only one channel
im2col :: Int            -- ^ kernel rows
       -> Int            -- ^ kernel columns
       -> Int            -- ^ stride rows
       -> Int            -- ^ stride columns
       -> Matrix RealNum -- ^ input matrix
       -> Matrix RealNum -- ^ output matrix
im2col kernelRows kernelColumns strideRows strideColumns dataIm =
  let channels = 1
      height = rows dataIm
      width  = cols dataIm
      rowOut          = (height - kernelRows) `div` strideRows + 1
      colOut          = (width - kernelColumns) `div` strideColumns + 1
  in  im2col_c channels height width kernelRows kernelColumns strideRows strideColumns 0 0 dataIm rowOut colOut

im2col_c :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Int -> Int -> Matrix RealNum
im2col_c channels height width kernelRows kernelColumns strideRows strideColumns padl padt dataIm outRows outCols =
  let vec             = flatten dataIm
      kernelSize      = kernelRows * kernelColumns
      numberOfPatches = outRows * outCols
  in unsafePerformIO $ do
    outPtr <- mallocForeignPtrArray (numberOfPatches * kernelSize * channels)
    let (inPtr, _) = U.unsafeToForeignPtr0 vec

    withForeignPtr inPtr $ \inPtr' ->
      withForeignPtr outPtr $ \outPtr' ->
        im2col_cpu inPtr' channels height width kernelRows kernelColumns strideRows strideColumns padt padl outRows outCols outPtr'

    let matVec = U.unsafeFromForeignPtr0 outPtr (numberOfPatches * kernelSize * channels)
    return $ U.matrixFromVector U.RowMajor (kernelSize * channels) numberOfPatches matVec
{-# INLINE im2col_c #-}

foreign import ccall unsafe
    im2col_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr RealNum -> IO ()

foreign import ccall unsafe
    in_place_add_per_channel_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Ptr RealNum -> IO ()

-- | Efficient implementation of convolutions using the im2col trick as popularised by Caffe
forwardConv2d :: Matrix RealNum -- ^ input
              -> Int            -- ^ input channels
              -> Int            -- ^ input rows
              -> Int            -- ^ input columns
              -> Matrix RealNum -- ^ kernel
              -> Int            -- ^ kernel filters
              -> Int            -- ^ kernel rows
              -> Int            -- ^ kernel columns
              -> Int            -- ^ stride height
              -> Int            -- ^ stride width
              -> Int            -- ^ output rows
              -> Int            -- ^ output columns
              -> Int            -- ^ pad left
              -> Int            -- ^ pad top
              -> Matrix RealNum -- ^ output of convolution
forwardConv2d input channels rows cols kernel filters kernelRows kernelCols strideRows strideCols outRows outCols padLeft padTop
  = let dataCol  = im2col_c channels rows cols kernelRows kernelCols strideRows strideCols padLeft padTop input outRows outCols
        gemmM    = LA.tr kernel LA.<> dataCol
        dataVid  = reshapeMatrix (filters * outRows) outCols gemmM
    in  dataVid

-- | Efficient implementation of convolutions using the im2col trick as popularised by Caffe
--   also adds some bias to the functions, note that it does it in place to prevent unnecessary allocation.
forwardBiasConv2d :: Matrix RealNum -- ^ input
                  -> Int            -- ^ input channels
                  -> Int            -- ^ input rows
                  -> Int            -- ^ input columns
                  -> Vector RealNum -- ^ bias
                  -> Matrix RealNum -- ^ kernel
                  -> Int            -- ^ kernel filters
                  -> Int            -- ^ kernel rows
                  -> Int            -- ^ kernel columns
                  -> Int            -- ^ stride height
                  -> Int            -- ^ stride width
                  -> Int            -- ^ output rows
                  -> Int            -- ^ output columns
                  -> Int            -- ^ pad left
                  -> Int            -- ^ pad top
                  -> Matrix RealNum -- ^ output of convolution
forwardBiasConv2d input channels rows cols bias kernel filters kernelRows kernelCols strideRows strideCols outRows outCols padLeft padTop =
  let outSize         = outRows * outCols * filters
  in  unsafePerformIO $ do
    let dataCol  = im2col_c channels rows cols kernelRows kernelCols strideRows strideCols padLeft padTop input outRows outCols
        gemmM    = LA.tr kernel LA.<> dataCol
        dataVid  = flatten gemmM
        (xPtr, _) = U.unsafeToForeignPtr0 dataVid
        (bPtr, _) = U.unsafeToForeignPtr0 bias

    withForeignPtr xPtr $ \xPtr' ->
      withForeignPtr bPtr $ \bPtr' ->
        in_place_add_per_channel_cpu xPtr' filters outRows outCols bPtr'

    let matVec = U.unsafeFromForeignPtr0 xPtr outSize
    return $ U.matrixFromVector U.RowMajor (outRows * filters) outCols matVec

-- | Efficient implementation of the backward pass for convolutions using the im2col trick as popularised by Caffe
backwardConv2d :: Matrix RealNum -- ^ input
               -> Int            -- ^ input channels
               -> Int            -- ^ input rows
               -> Int            -- ^ input columns
               -> Matrix RealNum -- ^ kernel
               -> Int            -- ^ kernel filters
               -> Int            -- ^ kernel rows
               -> Int            -- ^ kernel columns
               -> Int            -- ^ stride height
               -> Int            -- ^ stride width
               -> Int            -- ^ pad left
               -> Int            -- ^ pad top
               -> Int            -- ^ pad right
               -> Int            -- ^ pad bottom
               -> Matrix RealNum -- ^ derivative wrt out of layer
               -> Int            -- ^ output rows
               -> Int            -- ^ output columns
               -> (Matrix RealNum, Matrix RealNum)    -- ^ (derivated wrt input, derivative wrt kernel)
backwardConv2d input channels rows cols kernel filters kernelRows kernelCols strideRows strideCols padl padt padr padb dout outRows outCols =
  let dout'  = LA.reshape (outRows * outCols) $ LA.flatten dout
      dX_col = kernel LA.<> dout'
      rowCol = (div (rows + padt + padb - kernelRows) strideRows) + 1
      colCol = (div (cols + padl + padr - kernelCols) strideCols) + 1
      dX     = col2im_c channels rows cols kernelRows kernelCols strideRows strideCols padl padt rowCol colCol dX_col

      x_col  = im2col_c channels rows cols kernelRows kernelCols strideRows strideCols padl padt input outRows outCols
      dw_col = dout' LA.<> (LA.tr x_col)
      dw     = LA.tr $ reshapeMatrix filters (kernelRows * kernelCols * channels) dw_col
  in (dX, dw)

-- | Efficient implementation of the backward pass for convolutions using the im2col trick as popularised by Caffe
backwardBiasConv2d :: Matrix RealNum    -- ^ input
                   -> Int -- ^ input channels
                   -> Int -- ^ input rows
                   -> Int -- ^ input columns
                   -> Matrix RealNum -- ^ kernel
                   -> Int -- ^ kernel filters
                   -> Int -- ^ kernel rows
                   -> Int -- ^ kernel columns
                   -> Int -- ^ stride height
                   -> Int -- ^ stride width
                   -> Int -- ^ pad left
                   -> Int -- ^ pad top
                   -> Int -- ^ pad right
                   -> Int -- ^ pad bottom
                   -> Matrix RealNum -- ^ derivative wrt out of layer
                   -> Int -- ^ output rows
                   -> Int -- ^ output columns
                   -> (Matrix RealNum, Matrix RealNum, Vector RealNum)  -- ^ (derivated wrt input, derivative wrt kernel, derivative wrt bias)
backwardBiasConv2d input channels rows cols kernel filters kernelRows kernelCols strideRows strideCols padl padt padr padb dout outRows outCols =
  let (dX, dw) = backwardConv2d input channels rows cols kernel filters kernelRows kernelCols strideRows strideCols padl padt padr padb dout outRows outCols
      dB       = sum_over_channels_c dout filters outRows outCols
  in (dX, dw, dB)

sum_over_channels_c :: Matrix RealNum -> Int -> Int -> Int -> Vector RealNum
sum_over_channels_c mat channels rows cols = unsafePerformIO $ do
  let (xPtr, _) = U.unsafeToForeignPtr0 . flatten $ mat
  outPtr <- mallocForeignPtrArray channels

  withForeignPtr xPtr $ \xPtr' ->
    withForeignPtr outPtr $ \outPtr' ->
      sum_over_channels_cpu xPtr' channels rows cols outPtr'

  return $ U.unsafeFromForeignPtr0 outPtr channels

foreign import ccall unsafe
    sum_over_channels_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Ptr RealNum -> IO ()


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE RankNTypes #-}

module Main(main) where

import Graphics.Gloss (Display, Color, Picture(..),
    circleSolid, red, color, translate, dark,
    pictures, black, white, rotate)
import Graphics.Gloss.Data.ViewPort
import Graphics.Gloss.Interface.Pure.Game (
      Event(EventKey)
    , MouseButton(..)
    , Key(Char, MouseButton)
    , KeyState(Up, Down)
    , play
    , Display(InWindow)
    )

import qualified Data.Map.Strict as Map
import Data.List (elemIndex, mapAccumL, partition)
import Data.Tuple (swap)
import Data.Maybe (catMaybes)
import System.Random
import Control.Monad.State

import Control.Lens

{-} ###########
    # PHYSICS #
    ########### -}

import Data.Complex

type Vector = Complex Float
j :: Vector
j = 0:+1

normalize :: RealFloat a => Complex a -> Complex a
normalize (0:+0) = 0:+0
normalize a = a / (sqrt $ a * conjugate a)

toxy :: Complex a -> (a, a)
toxy (x :+ y) = (x, y)

toVector :: Float -> Vector
toVector x = x:+0

{-} ##############
    # STRUCTURES #
    ############## -}

data Ball    = Ball Vector Vector Float
data Timer   = Timer
    { _total  :: Float
    , _actual :: Float
    }
makeLenses ''Timer

data Body    = Body Ball Int

data Entity d = Entity {
      _pos    :: Vector
    , _vel    :: Vector
    , _rad    :: Float
    , _lifes  :: Int
    , _damage :: Int
    , _body :: d
    }
makeLenses ''Entity

type Bullet  = Entity ()
type Barrel  = Entity ()
type Player  = Entity ()
type Smoke   = Entity Timer

data Satelite = Satelite Float Float Float Float Float Float -- r w h rot v t

data MonsterData = MonsterData {
    _hull  :: [Satelite]
  , _shift :: Int
  }
makeLenses ''MonsterData

type Monster = Entity MonsterData

data Game = Game {
      _player   :: Player
    , _keyboard :: Map.Map Key KeyState
    , _bullets  :: [Bullet]
    , _barrels  :: [Barrel]
    , _clicks   :: [Vector]
    , _smoke    :: [Smoke]
    , _seed     :: StdGen
    , _ticks    :: Int
    , _monsters :: [Monster]
    }
makeLenses ''Game

{-} ##########
    # ENTITY #
    ########## -}

renderBody :: Entity e -> Picture
renderBody ent = translate x y $
                 color white $
                 circleSolid (ent^.rad)
                 where (x:+y) = ent^.pos

updateBody :: Float -> Entity e -> Entity e
updateBody time ent = ent & pos +~ (toVector time)*(ent^.vel)

collide :: Entity a -> Entity b -> Bool
collide a b = (realPart $ abs $ (a^.pos - b^.pos)) < (a^.rad + b^.rad)

isAlive :: Entity e -> Bool
isAlive ent = (ent^.lifes) > 0

{-} ##########
    #  GAME  #
    ########## -}

window :: Display
window = InWindow "JampaZen" (500, 500) (5, 5) -- w, h, offset

background :: Color
background = black

initialState :: Game
initialState = Game
    { _player = Entity
        { _pos = 0
        , _vel = 0
        , _rad = 10
        , _lifes = 10
        , _damage = 1
        , _body = ()
        }
    , _keyboard = Map.fromList $ zip (map Char "wasd") (repeat Up)
    , _bullets  = []
    , _barrels  = []
    , _monsters = []
    , _clicks   = []
    , _smoke    = []
    , _seed     = mkStdGen 1
    , _ticks    = 0
    }

fps :: Int
fps = 60

handleKeys :: Event -> Game -> Game
handleKeys (EventKey (Char key) keystate _ _) =
    keyboard %~ Map.adjust (const keystate) (Char key)
handleKeys (EventKey (MouseButton LeftButton) Down _ (mx, my)) =
    clicks %~ (:) (mx:+my)
handleKeys _ = id

renderGame :: Game -> Picture
renderGame game =
    let (mx:+my) = game^.player.pos
        canmap x = map renderBody (game^.x)
    in  translate (-mx) (-my) $
        pictures $ concat $ [
          [renderBody (game^.player)]
        , map renderBody (game^.bullets)
        , map renderSmoke (game^.smoke)
        , map renderMonster (game^.monsters)
        ]

updateGame :: Float -> Game -> Game
updateGame time game =
    game & ifPlayerAlive shot
         & removeSmoke
         & onceUpon 4 bulletSmoke
         & ifPlayerAlive (onceUpon 50 spawnMonster)
         & playerMove
         & collisions
         & player   %~ updateBody time
         & bullets  %~ map (updateBody time)
         & smoke    %~ map (updateSmoke time)
         & monsters %~ map (updateMonster time game)
         & ticks    +~ 1

main :: IO ()
main = play window background fps initialState renderGame handleKeys updateGame

{-} ########
    # SHOT #
    ######## -}

newBullet :: Vector -> Vector -> Bullet
newBullet p v = Entity
    { _pos=p, _vel=v, _rad=5, _lifes=1, _damage=1, _body=()}

shot :: Game -> Game
shot game = game & (clicks .~ [])
                 & (bullets %~ (++) newBullets)
      where newBullets = (newBullet (game^.player.pos) . (500*) . normalize) <$> game^.clicks

{-} ##############
    # EXPLOSIONS #
    ############## -}

randomSt :: (Float, Float) -> State StdGen Float
randomSt = state . randomR

generateExplosionSmoke :: Float -> Vector -> State StdGen Smoke
generateExplosionSmoke size pos = do
    pa <- randomSt (0, 2*pi)
    pr <- randomSt (0, 5*size)
    let p = mkPolar pr pa
    va <- randomSt (0, 2*pi)
    vr <- randomSt (10, 10*size)
    let v = mkPolar vr va
    r <- randomSt (10, 50)
    ttl <- randomSt (0.25, (sqrt size)/3)
    return Entity {_pos=p+pos, _vel=v, _rad=r,
                   _lifes=1, _damage=0, _body=newTimer ttl}

generateExplosion :: Float -> Vector -> Game -> Game
generateExplosion size pos game =
    game & seed  .~ seed'
         & smoke %~ (++) newsmoke
    where
        (newsmoke, seed') = flip runState (game^.seed) $
            sequence $ take 30 $ repeat $ generateExplosionSmoke size pos

{-} ##############
    # COLLISIONS #
    ############## -}

-- two entities collide
(<..>) :: Entity a -> Entity b -> (Entity a, Entity b)
a <..> b | collide a b && isAlive a && isAlive b = (,)
              (a & lifes %~ flip (-) (b^.damage))
              (b & lifes %~ flip (-) (a^.damage))
         | otherwise = (a, b)

(<.+>) :: Entity a -> [Entity b] -> (Entity a, [Entity b])
(<.+>) = mapAccumL (<..>)

(<+.>) :: [Entity a] -> Entity b -> ([Entity a], Entity b)
(<+.>) as b = swap $ b <.+> as

-- two list of entities collide
(<++>) :: [Entity a] -> [Entity b] -> ([Entity a], [Entity b])
(<++>) = mapAccumL (<+.>)

smash :: Lens' Game [Entity a]
      -> Lens' Game [Entity b]
      -> (Entity a -> Game -> Game)
      -> (Entity b -> Game -> Game)
      -> Game -> Game

--smash as bs destroyA destroyB b game =
smash as bs destroyA destroyB game =
    game & as .~ aliveAs
         & bs .~ aliveBs
         & foldr (.) id destructions
    where
        (as', bs') = (game^.as) <++> (game^.bs)
        (aliveAs, deadAs) = partition isAlive as'
        (aliveBs, deadBs) = partition isAlive bs'
        destructions = map destroyA deadAs ++ map destroyB deadBs

collisions :: Game -> Game
collisions =
    smash bullets monsters destroyBullet destroyMonster

{-} #########
    # TIMER #
    ######### -}

newTimer :: Float -> Timer
newTimer time = Timer {_total=time, _actual=time}

onceUpon :: Int -> (Game -> Game) -> Game -> Game
onceUpon n f game | game^.ticks `mod` n == 0 = f game
                  | otherwise                = game

{-} ##########
    # PLAYER #
    ########## -}

playerMove game = game & player.vel .~ 100*normalize (keyboardToVector $ game^.keyboard)

ifPlayerAlive :: (Game -> Game) -> Game -> Game
ifPlayerAlive f game = if isAlive (game^.player) then f game else game

stateToScalar :: KeyState -> Vector
stateToScalar Up = 1
stateToScalar Down = 0

keyToVector :: Key -> Vector
keyToVector (Char c) = maybe 0 (j^) (elemIndex c "asdw")

keyboardToVector :: Map.Map Key KeyState -> Vector
keyboardToVector = Map.foldrWithKey f (0:+0) where
    f k v a = a + (stateToScalar v)*(keyToVector k)

destroyPlayer :: Player -> Game -> Game
destroyPlayer player = generateExplosion 100 (player^.pos)

{-} ###########
    #  SMOKE  #
    ########### -}

bulletSmoke :: Game -> Game
bulletSmoke game = game & smoke %~ (++) (map bulletToSmoke $ game^.bullets)
      where bulletToSmoke bullet = Entity{
            _pos=bullet^.pos
          , _vel=0
          , _rad=bullet^.rad
          , _lifes=1
          , _damage= 0
          , _body=newTimer 0.25
          }

updateSmoke :: Float -> Smoke -> Smoke
updateSmoke time = (updateBody time) . (body.actual -~ time)

renderSmoke smoke = translate x y $
               color white $
               circleSolid r
    where (x:+y) = smoke^.pos
          r = smoke^.rad * smoke^.body.actual / smoke^.body.total

removeSmoke :: Game -> Game
removeSmoke = smoke %~ filter ((>0) . (^.body.actual))

{-} ##########
    # BULLET #
    ########## -}

destroyBullet :: Bullet -> Game -> Game
destroyBullet = const id

{-} ###########
    # MONSTER #
    ########### -}

generateSatelite :: State StdGen Satelite
generateSatelite = do
    r <- randomSt (10, 20)
    w <- randomSt (10, 50)
    h <- randomSt (10, 50)
    v <- randomSt (1, 5)
    t <- randomSt (0, 2*pi)
    rot <- randomSt (0, 2*pi)
    o <- signum <$> randomSt (-1, 1) -- orientation
    return $ Satelite r w h rot (o*v) t

generateMonster :: Vector -> Game -> Game
generateMonster p game =
    game & seed     .~ seed'
         & monsters %~ (:) Entity {
                _pos=p
              , _vel=0
              , _rad=50
              , _lifes=3
              , _damage=1
              , _body=MonsterData {
                    _hull=hull'
                  , _shift=game^.ticks
                  }
              }
    where (hull', seed') = flip runState (game^.seed) $
              sequence $ take 20 $ repeat $ generateSatelite

spawnMonster :: Game -> Game
spawnMonster game =
    game & generateMonster (mkPolar d a)
         & seed     .~ seed'
    where ((a, d), seed') = runState randomMonster (game^.seed)
          randomMonster = do
              a <- randomSt (0, 2*pi)
              d <- randomSt (1000, 2000)
              return (a, d)

renderSatelite :: Satelite -> Picture
renderSatelite (Satelite r w h rot v t) =
    rotate rot $
    translate (w*cos t) (h*sin t)$
    color white $
    circleSolid r

updateSatelite :: Float -> Satelite -> Satelite
updateSatelite time (Satelite r w h rot v t) = (Satelite r w h rot v (t + v*time))

updateMonster :: Float -> Game -> Monster -> Monster
updateMonster time game monster =
        monster & updateBody time
                & followPlayer
                & body.hull %~ map (updateSatelite time)
        where
            pp = game^.player.pos
            mp = monster^.pos
            followPlayer
              | ((game^.ticks) + (monster^.body.shift)) `mod` 120 /= 0 = id
              | otherwise = vel .~ 100*normalize (pp - mp)

renderMonster :: Monster -> Picture
renderMonster monster =
    translate x y $ pictures $ renderSatelite <$> (monster^.body.hull)
    where (x:+y) = monster^.pos

destroyMonster :: Monster -> Game -> Game
destroyMonster monster = generateExplosion 15 (monster^.pos)


{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE BangPatterns,
             ScopedTypeVariables,
             RecordWildCards,
             FlexibleContexts,
             TypeFamilies #-}
import Data.IDX (IDXData,  decodeIDXFile, idxIntContent )
{------------------------}
import Numeric.LinearAlgebra as NL (cmap, (#>), vjoin, Numeric,  Vector, fromList )
{------------------------}
import Data.Vector.Unboxed as DTU ( toList )
import qualified Data.Vector.Storable as DT
import Data.List (sort)
import Prelude hiding (readFile)
import Neuro
    ( ActivationFunction, Network(..), loadNetwork )

kol_vo :: Int
kol_vo = 1000

main :: IO ()
main = do
        res <- loadNetwork "smartNet5.nn"
        Just idxTrain  <- decodeIDXFile "test/t10k-images.idx3-ubyte" -- image
        Just idxResult <- decodeIDXFile "test/t10k-labels.idx1-ubyte" -- result
        --writeFile "resNeuro.txt" (concat $ take (28*kol_vo) $ drowImage idxTrain)
        --writeFile "resLabel.txt" (show $ take kol_vo $ DTU.toList $  idxIntContent idxResult)
        putStrLn "Точность алгоритма :"
        print $ (\x -> (fromIntegral x) / fromIntegral kol_vo) $ length $ filter (\((_,x),y) -> x == y) $ flip zip (filterCorrect idxResult) $  printRes kol_vo res idxTrain $ DTU.toList $ idxIntContent idxResult
        where filterCorrect idxResult = reverse $ take kol_vo $ DTU.toList $ idxIntContent idxResult

printRes :: Int -> Network Double -> IDXData -> t -> [(Double,Int)]
printRes 0 _   _           _      = []
printRes x res idxTrain idxResult = last (sort $ flip zip [0..] $ softmax $ DT.toList $ Main.output res tanh (last $ take x $ image idxTrain)) : printRes (x-1) res idxTrain idxResult

output :: (Floating (Vector a), Numeric a, DT.Storable a, Num (Vector a)) => Network a -> ActivationFunction a -> Vector a -> Vector a
output (Network{..}) act input = foldl f (vjoin [input, 1]) matrices
  where f (!inp) m = cmap act $ m #> inp

softmax :: [Double] -> [Double]
softmax xs = let xs' = exp <$> xs
                 s   = sum xs'
             in map (/ s) xs'

drowImage :: IDXData -> [String]
drowImage = _0_255to0or1 . DTU.toList . idxIntContent

_0_255to0or1 :: [Int] -> [String]
_0_255to0or1 [] = []
_0_255to0or1 xs = (show (concat $ fmap show $ convert $ take 28 xs) ++ " \n") : _0_255to0or1 (drop 28 xs)

convert [] = []
convert (x:xs)
              | x > 240 = 1 : convert xs
              | otherwise = 0 : convert xs

image :: IDXData -> [Vector Double]
image  = matrix2x2 . _0_1to0_255 . DTU.toList . idxIntContent
result :: IDXData -> [Vector Double]
result = unitar . DTU.toList . idxIntContent

unitar :: [Int] -> [NL.Vector Double]
unitar []     = []
unitar (0:xs) = NL.fromList [1,-1,-1,-1,-1,-1,-1,-1,-1,-1] : unitar xs
unitar (1:xs) = NL.fromList [-1,1,-1,-1,-1,-1,-1,-1,-1,-1] : unitar xs
unitar (2:xs) = NL.fromList [-1,-1,1,-1,-1,-1,-1,-1,-1,-1] : unitar xs
unitar (3:xs) = NL.fromList [-1,-1,-1,1,-1,-1,-1,-1,-1,-1] : unitar xs
unitar (4:xs) = NL.fromList [-1,-1,-1,-1,1,-1,-1,-1,-1,-1] : unitar xs
unitar (5:xs) = NL.fromList [-1,-1,-1,-1,-1,1,-1,-1,-1,-1] : unitar xs
unitar (6:xs) = NL.fromList [-1,-1,-1,-1,-1,-1,1,-1,-1,-1] : unitar xs
unitar (7:xs) = NL.fromList [-1,-1,-1,-1,-1,-1,-1,1,-1,-1] : unitar xs
unitar (8:xs) = NL.fromList [-1,-1,-1,-1,-1,-1,-1,-1,1,-1] : unitar xs
unitar (9:xs) = NL.fromList [-1,-1,-1,-1,-1,-1,-1,-1,-1,1] : unitar xs

_0_1to0_255 :: [Int] -> [Double]
_0_1to0_255 [] = []
_0_1to0_255 (x:xs) = (fromIntegral x / 255) : _0_1to0_255 xs

matrix2x2 :: [Double] -> [NL.Vector Double]
matrix2x2 [] = []
matrix2x2 xs = NL.fromList (Prelude.take 784 xs) : matrix2x2 (Prelude.drop 784 xs)


{-# LANGUAGE CPP, OverloadedStrings, TypeOperators #-}

module Main where

import Control.Exception
import qualified Data.Array.Accelerate as A
import qualified Data.Array.Accelerate.Interpreter as ALI
import Data.Array.Accelerate.Math.Hilbert
import Data.Array.Accelerate.Math.Qtfd
import Data.Array.Accelerate.Math.WindowFunc
import Data.Attoparsec.Text
import ParseArgs
import ParseFile
import System.Directory
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
import qualified Data.Array.Accelerate.LLVM.Native as ALN
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
import qualified Data.Array.Accelerate.LLVM.PTX as ALP
#endif
import Data.Array.Accelerate.Array.Sugar as S
import qualified Data.ByteString.Builder as BB
import Data.Complex
import qualified Data.Double.Conversion.Convertable as DC
import Data.List
import Data.Monoid
import qualified Data.Text as T
import qualified Data.Text.IO as TI
import qualified Data.Vector as V
import GHC.Float
import Math.Gamma
import System.IO

main :: IO ()
main = do
  args <- opts
  makeQTFDAll args

-- | Transform data to text(using Float conversion. It is much faster) and write it to file.

writeData ::
  Handle            -- ^ output file
  -> [Float]       -- ^ Data to write
  -> Int            -- ^ number of element
  -> Int            -- ^ elements per row
  -> IO ()
writeData _ [] _ _  = return ()
writeData file (x:xs) w n = do
  TI.hPutStr file $ DC.toPrecision 5 x <> " "
  if (n `mod` w) == 0
  then do
    TI.hPutStr file "\n"
    writeData file xs w (n + 1)
  else writeData file xs w (n + 1)

makeString ::
     Handle
  -> [Float]       -- ^ Data to write
  -> Int            -- ^ number of element
  -> Int            -- ^ elements per row
  -> BB.Builder
  -> IO ()
makeString _ [] _ _ _  = return ()
makeString file (x:xs) w n acc =
  if (n `mod` w) == 0
  then BB.hPutBuilder file (acc <> DC.toPrecision 5 x <> " " <> "\n") >> makeString file xs w (n+1) ""
  else makeString file xs w (n+1) (acc <> DC.toPrecision 5 x <> " ")

-- | make Pseudo Wigner-Ville transform to all files in a directory.

makeQTFDAll :: Opts -> IO ()
makeQTFDAll opts = do
  let path = getPath opts
      nosAVGflag = subAVGflag opts
  files <- listDirectory path
  setCurrentDirectory path
  let fFiles = filter (isSuffixOf ".txt") files
      sAVGflag = A.constant $ not nosAVGflag
  case opts of
    (OptsPWV _ dev wlen wfun _) ->
      do
        if (even wlen)
        then error "Window length maust be odd !"
        else
          do
            let window = makeWindow wfun (A.unit $ A.constant wlen) :: A.Acc (A.Array A.DIM1 Float)
            mapM_ (makePWV dev window sAVGflag) fFiles
    (OptsWV path dev nosAVGflag) -> mapM_ (makeWV dev sAVGflag) fFiles
    (OptsCW path dev sigma uWindow uWindowFunc nWindow nWindowFunc normalise nosAVGflag) ->
      do
        if (even uWindow || even nWindow)
        then error "Window length maust be odd !"
        else
          let uWindowArray = makeWindow uWindowFunc (A.unit $ A.constant uWindow) :: A.Acc (A.Array A.DIM1 Float)
              nWindowArray = makeWindow nWindowFunc (A.unit $ A.constant nWindow) :: A.Acc (A.Array A.DIM1 Float)
              bothRectangular = uWindowFunc == Rect && nWindowFunc == Rect
          in mapM_ (makeCW dev sAVGflag sigma uWindowArray nWindowArray normalise bothRectangular) fFiles
    (OptsBJ path dev alpha uWindow uWindowFunc nWindow nWindowFunc nosAVGflag) ->
      do
        if (even uWindow || even nWindow)
        then error "Window length maust be odd !"
        else
          let uWindowArray = makeWindow uWindowFunc (A.unit $ A.constant uWindow) :: A.Acc (A.Array A.DIM1 Float)
              nWindowArray = makeWindow nWindowFunc (A.unit $ A.constant nWindow) :: A.Acc (A.Array A.DIM1 Float)
              bothRectangular = uWindowFunc == Rect && nWindowFunc == Rect
          in mapM_ (makeBJ dev sAVGflag alpha uWindowArray nWindowArray bothRectangular) fFiles
    (OptsEMBD path dev alpha beta uWindow uWindowFunc nWindow nWindowFunc normalise nosAVGflag) ->
      do
        if (even uWindow || even nWindow)
        then error "Window length maust be odd !"
        else let gammas = makeGammaArray beta uWindow
                 uWindowArray = makeWindow uWindowFunc (A.unit $ A.constant uWindow) :: A.Acc (A.Array A.DIM1 Float)
                 nWindowArray = makeWindow nWindowFunc (A.unit $ A.constant nWindow) :: A.Acc (A.Array A.DIM1 Float)
                 bothRectangular = uWindowFunc == Rect && nWindowFunc == Rect
             in mapM_ (makeEMBD dev sAVGflag alpha gammas uWindowArray nWindowArray bothRectangular normalise) fFiles
  return ()

makeCW :: CalcDev                   -- ^ Calculation device - CPU or GPU
  -> A.Exp Bool                     -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float                          -- ^ sigma
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if the core array should be normalised
  -> Bool                           -- ^ if Both windows are rectangular
  -> FilePath                       -- ^ Name of file with data.
  -> IO ()
makeCW dev sAVGflag sigma uWindowArray nWindowArray normalise bothRect file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startCW dev file sAVGflag sigma uWindowArray nWindowArray normalise bothRect) dataF

makeBJ :: CalcDev                   -- ^ Calculation device - CPU or GPU
  -> A.Exp Bool                     -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float                          -- ^ sigma
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if Both windows are rectangular
  -> FilePath                       -- ^ Name of file with data.
  -> IO ()
makeBJ dev sAVGflag alpha uWindowArray nWindowArray bothRect file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startBJ dev file sAVGflag alpha uWindowArray nWindowArray bothRect) dataF

-- | Make Pseudo Wigner-Ville transform to all columns in the given file

makePWV ::
  CalcDev                            -- ^ Calculation device - CPU or GPU
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> A.Exp Bool                      -- ^ Apply subtraction of the mean value from all elements in a column.
  -> FilePath                        -- ^ Name of file with data.
  -> IO ()
makePWV dev window sAVGflag file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startPWV dev window file sAVGflag) dataF

--makeEMBD dev alpha gammas uWindowArray nWindowArray bothRectangular normalise

makeEMBD :: CalcDev                   -- ^ Calculation device - CPU or GPU
  -> A.Exp Bool                     -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float                          -- ^ alpha
  -> [Float]                        -- ^ gammas
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if the core array should be normalised
  -> Bool                           -- ^ if Both windows are rectangular
  -> FilePath                       -- ^ Name of file with data.
  -> IO ()
makeEMBD dev sAVGflag alpha gammas uWindowArray nWindowArray normalise bothRect file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startEMBD dev file sAVGflag alpha gammas uWindowArray nWindowArray normalise bothRect) dataF

-- | Make Pseudo Wigner-Ville transform to the given column.
-- At first it makes subtraction of average value (if supAVGflag) and applies hilbert transform to make analytic signal. After that it applies Pseudo-Wigner transftorm
-- and save data to file with

startPWV ::
  CalcDev                          -- ^ Calculation device - CPU or GPU
  -> A.Acc (A.Array A.DIM1 Float) -- ^ Smoothing window in time-domain. Length must be odd.
  -> FilePath                      -- ^ Name of file with data.
  -> A.Exp Bool                    -- ^ Apply subtraction of the mean value from all elements in a column.
  -> (T.Text,[Float])             -- ^ Name of column and parsed data from column
  -> IO ()
startPWV dev window oldName sAVGflag (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      pData = A.fromList (A.Z A.:. leng) dataF
      appPWV = (pWignerVille window) . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appPWV pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appPWV pData
#endif
                    CPU -> ALI.run1 appPWV pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $! processed
  file <- openFile newFName WriteMode
  onException (writeData file pList leng 1) (removeFile newFName)
  hClose file
  putStrLn "Done !"

-- | Make Wigner-Ville transform to all columns in the given file

makeWV ::
  CalcDev        -- ^ Calculation device - CPU or GPU
  -> A.Exp Bool  -- ^ Apply subtraction of the mean value from all elements in a column.
  -> FilePath    -- ^ Name of file with data.
  -> IO ()
makeWV dev sAVGflag file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startWV dev file sAVGflag) dataF

-- | Make Wigner-Ville transform to the given column.
-- At first it makes subtraction of average value (if supAVGflag) and applies hilbert transform to make analytic signal. After that it applies Pseudo-Wigner transftorm
-- and save data to file with

startWV ::
  CalcDev                 -- ^ Calculation device - CPU or GPU
  -> FilePath             -- ^ Name of file with data.
  -> A.Exp Bool           -- ^ Apply subtraction of the mean value from all elements in a column.
  -> (T.Text,[Float])    -- ^ Name of column and parsed data from column
  -> IO ()
startWV dev oldName sAVGflag (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      pData = A.fromList (A.Z A.:. leng) dataF
      appWV = wignerVilleNew . hilbert . supAVG sAVGflag
      appWVGPU = wignerVille . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appWV pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appWV pData
#endif
                    CPU -> ALI.run1 appWV pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $  processed
  file <- openFile newFName WriteMode
  hSetBinaryMode file True >> hSetBuffering file LineBuffering
  -- let resultBsBuilder = makeString pList leng 1 ""
  onException (makeString file pList leng 1 "") (removeFile newFName) >> hFlush file
  hClose file
  putStrLn "Done !"

startCW ::
  CalcDev                 -- ^ Calculation device - CPU or GPU
  -> FilePath             -- ^ Name of file with data.
  -> A.Exp Bool           -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float               -- ^ Sigma
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if the core array should be normalised
  -> Bool                           -- ^ if Both windows are rectangular
  -> (T.Text,[Float])    -- ^ Name of column and parsed data from column
  -> IO ()
startCW dev oldName sAVGflag sigma uWindowArray nWindowArray normalise bothRect (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      uWindow = A.length uWindowArray
      nWindow = A.length nWindowArray
      mWindowArrays = if bothRect then Nothing else Just (uWindowArray,nWindowArray)
      pData = A.fromList (A.Z A.:. leng) dataF
      appCW = (choiWilliams (A.constant sigma) mWindowArrays uWindow nWindow normalise) . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appCW pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appCW pData
#endif
                    CPU -> ALI.run1 appCW pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $  processed
  file <- openFile newFName WriteMode
  onException (writeData file pList leng 1) (removeFile newFName)
  hClose file
  putStrLn "Done !"

startBJ ::
  CalcDev                 -- ^ Calculation device - CPU or GPU
  -> FilePath             -- ^ Name of file with data.
  -> A.Exp Bool           -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float               -- ^ Sigma
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if Both windows are rectangular
  -> (T.Text,[Float])    -- ^ Name of column and parsed data from column
  -> IO ()
startBJ dev oldName sAVGflag alpha uWindowArray nWindowArray bothRect (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      uWindow = A.length uWindowArray
      nWindow = A.length nWindowArray
      mWindowArrays = if bothRect then Nothing else Just (uWindowArray,nWindowArray)
      pData = A.fromList (A.Z A.:. leng) dataF
      appCW = (bornJordan (A.constant alpha) mWindowArrays uWindow nWindow) . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appCW pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appCW pData
#endif
                    CPU -> ALI.run1 appCW pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $  processed
  file <- openFile newFName WriteMode
  onException (writeData file pList leng 1) (removeFile newFName)
  hClose file
  putStrLn "Done !"

startEMBD ::
  CalcDev                 -- ^ Calculation device - CPU or GPU
  -> FilePath             -- ^ Name of file with data.
  -> A.Exp Bool           -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float                -- ^ Alpha
  -> [Float]              -- ^ Gammas
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if the core array should be normalised
  -> Bool                           -- ^ if Both windows are rectangular
  -> (T.Text,[Float])    -- ^ Name of column and parsed data from column
  -> IO ()
startEMBD dev oldName sAVGflag alpha gammas uWindowArray nWindowArray normalise bothRect (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      uWindow = A.length uWindowArray
      nWindow = A.length nWindowArray
      gammasArr = A.use $ A.fromList (A.Z A.:. (length gammas)) gammas
      mWindowArrays = if bothRect then Nothing else Just (uWindowArray,nWindowArray)
      pData = A.fromList (A.Z A.:. leng) dataF
      appEMBD = (eModifiedB (A.constant alpha) gammasArr mWindowArrays uWindow nWindow normalise) . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appEMBD pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appEMBD pData
#endif
                    CPU -> ALI.run1 appEMBD pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $  processed
  file <- openFile newFName WriteMode
  onException (writeData file pList leng 1) (removeFile newFName)
  hClose file
  putStrLn "Done !"

-- | Subtraction of average value from all elements of column

supAVG :: A.Exp Bool -> A.Acc (Array DIM1 Float) -> A.Acc (Array DIM1 Float)
supAVG flag arr =
  A.acond flag (A.map (\x -> x - avg) arr) arr
  where leng = A.length arr
        avg = (A.the $ A.sum arr)/(A.fromIntegral leng)

makeGammaArray :: Float -> Int -> [Float]
makeGammaArray beta wLength =
  let dwLength = fromIntegral wLength :: Float
      v = [(-0.5), ((-0.5) + 1.0/dwLength)..0.5]
      s = map (\x -> beta :+ pi*x) v
      f = map (\x -> abs $ ((magnitude $ gamma x) ^ 2)/((gamma beta) ^ 2)) s
      n = floor $ (fromIntegral $ length f)/2.0
      frst = Data.List.take (n-1) f
      secnd  = Data.List.drop n f
  in secnd ++ frst