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

module Numeric.Trainee.Gradee (
	Gradee(..), gradee, ad,
	Unary, Binary,
	unary, binary,
	dup, vdup, vdupWith,
	conjoin, plus, vconjoin, vsum, vfold,
	swap,

	matMat, matVec, odot,
	corrVec, corrMat,
	flattenMat, reshapeVec, concatVecs,
	transposeMat,
	vecRow, vecCol,
	biasVec, biasMat,
	maxPoolVec, maxPoolMat,
	padVec, padMat
	) where

import Prelude hiding (id, (.))
import Prelude.Unicode

import Control.Category
import Control.Lens (lens)
import Data.Reflection (Reifies, reify)
import qualified Data.Vector as V
import Numeric.AD (grad, auto)
import Numeric.AD.Internal.Reverse (Reverse, primal, Tape)
import Numeric.LinearAlgebra

import Numeric.Trainee.Types

-- | Make Gradee like lens
gradee ∷ (a → b) → (a → b → a) → Gradee a b
gradee g s = Gradee $ lens g s

-- | Make like lens from binary op
gradee2 ∷ (a → b → c) → (a → b → c → (a, b)) → Gradee (a, b) c
gradee2 g s = gradee (uncurry g) (uncurry s)

-- | Make Gradee from any function
ad ∷ (Traversable f, Num a) ⇒ (forall s . Reifies s Tape ⇒ f (Reverse s a) → Reverse s a) → Gradee (f a) a
ad f = gradee f' (\x dx → fmap (* dx) (grad f x)) where
	f' = reify undefined (\p → primal ∘ spec p f ∘ fmap auto)
	spec ∷ Reifies t Tape ⇒ proxy t → (forall s . Reifies s Tape ⇒ g (Reverse s a) → Reverse s a) → g (Reverse t a) → Reverse t a
	spec _ h = h

type Unary a = forall s . Reifies s Tape ⇒ Reverse s a → Reverse s a
type Binary a = forall s . Reifies s Tape ⇒ Reverse s a → Reverse s a → Reverse s a

-- | Make Gradee from unary function
unary ∷ Num a ⇒ Unary a → Gradee a a
unary f = ad (\[x] → f x) . gradee return (const head)

-- | Make @Gradee@ from binary function
binary ∷ Num a ⇒ Binary a → Gradee (a, a) a
binary f = ad (\[x, y] → f x y) . gradee g s where
	g (x, y) = [x, y]
	s _ [x, y] = (x, y)
	s _ _ = error "binary"

dup ∷ Num a ⇒ Gradee a (a, a)
dup = gradee (\x → (x, x)) (\_ (dx', dx'') → dx' + dx'')

vdup ∷ Num a ⇒ Int → Gradee a (V.Vector a)
vdup n = gradee (V.replicate n) (const V.sum)

vdupWith ∷ (a → a → a) → Int → Gradee a (V.Vector a)
vdupWith fn n
	| n ≤ 0 = error "vdupWith: negative argument"
	| otherwise = gradee (V.replicate n) (const $ V.foldr1 fn)

conjoin ∷ Num a ⇒ Gradee (a, a) a
conjoin = binary (+)

plus ∷ Num a ⇒ Gradee (a, a) a
plus = conjoin

vconjoin ∷ Num a ⇒ Gradee (V.Vector a) a
vconjoin = ad V.sum

vsum ∷ Num a ⇒ Gradee (V.Vector a) a
vsum = vconjoin

vfold ∷ (a → a → a) → Gradee (V.Vector a) a
vfold fn = gradee (V.foldr1 fn) (V.replicate ∘ V.length)

swap ∷ Gradee (a, b) (b, a)
swap = gradee (\(x, y) → (y, x)) (\_ (dy, dx) → (dx, dy))

matMat ∷ Numeric a ⇒ Gradee (Matrix a, Matrix a) (Matrix a)
matMat = gradee2 (<>) backprop where
	backprop a b dc = (dc <> tr b, tr a <> dc)

matVec ∷ Numeric a ⇒ Gradee (Matrix a, Vector a) (Vector a)
matVec = gradee2 (#>) backprop where
	backprop a b dc = (outer dc b, tr a #> dc)

odot ∷ Num (Vector a) ⇒ Gradee (Vector a, Vector a) (Vector a)
odot = gradee2 (+) backprop where
	backprop _ _ dc = (dc, dc)

corrVec ∷ Numeric a ⇒ Gradee (Vector a, Vector a) (Vector a)
corrVec = gradee2 corr backprop where
	backprop a b dc = (corr dc b, conv a dc)

corrMat ∷ (Numeric a, Num (Vector a)) ⇒ Gradee (Matrix a, Matrix a) (Matrix a)
corrMat = gradee2 corr2 backprop where
	backprop a b dc = (corr2 dc b, conv2 a dc)

flattenMat ∷ Numeric a ⇒ Gradee (Matrix a) (Vector a)
flattenMat = gradee flatten backprop where
	backprop a = reshape (cols a)

reshapeVec ∷ Numeric a ⇒ Int → Gradee (Vector a) (Matrix a)
reshapeVec cols' = gradee (reshape cols') (const flatten)

concatVecs ∷ Numeric a ⇒ Gradee (V.Vector (Vector a)) (Vector a)
concatVecs = gradee
	(vjoin ∘ V.toList)
	(\vs → V.fromList ∘ takesV (V.toList (V.map size vs)))

transposeMat ∷ Numeric a ⇒ Gradee (Matrix a) (Matrix a)
transposeMat = gradee tr (const tr)

vecRow ∷ Numeric a ⇒ Gradee (Vector a) (Matrix a)
vecRow = gradee asRow (const flatten)

vecCol ∷ Numeric a ⇒ Gradee (Vector a) (Matrix a)
vecCol = gradee asColumn (const flatten)

biasVec ∷ Numeric a ⇒ Gradee (a, Vector a) (Vector a)
biasVec = gradee2 (\b v → cmap (+ b) v) backprop where
	backprop _ _ dv = (sumElements dv, dv)

biasMat ∷ Numeric a ⇒ Gradee (a, Matrix a) (Matrix a)
biasMat = gradee2 (\b m → cmap (+ b) m) backprop where
	backprop _ _ dm = (sumElements dm, dm)

maxPoolVec ∷ (RealFrac a, Numeric a) ⇒ Int → Gradee (Vector a) (Vector a)
maxPoolVec n = gradee pool' unpool' where
	pool' v = build (size v `div` n) $ \i → maxElement (subVector (round i * n) n v)
	unpool' v dv = accum zero (+) (zip maxIndices (toList dv)) where
		maxIndices = map (\i → maxIndex (subVector (i * n) n v) + i * n) [0 .. size dv - 1]
		zero = build (size v) (const 0)

maxPoolMat ∷ (RealFrac a, Numeric a) ⇒ Int → Int → Gradee (Matrix a) (Matrix a)
maxPoolMat w h = gradee pool' unpool' where
	pool' m = build (rows m `div` h, cols m `div` w) $ \i j → maxElement (subMatrix (round i * h, round j * w) (h, w) m)
	unpool' m dm = accum zero (+) (zip maxIndices (toList $ flatten dm)) where
		maxIndices = map (\(i, j) → maxIndex (subMatrix (i * h, j * w) (h, w) m) `biplus` (i * h, j * w)) [(i, j) | i ← [0 .. rows dm - 1], j ← [0 .. cols dm - 1]]
		zero = build (size m) (const $ const 0)
		biplus (lx, ly) (rx, ry) = (lx + rx, ly + ry)

padVec ∷ Numeric a ⇒ Int → Gradee (Vector a) (Vector a)
padVec n = gradee pad' unpad' where
	pad' v = vjoin [zeroes', v, zeroes'] where
		zeroes' = build n (const 0)
	unpad' = subVector n ∘ size

padMat ∷ Numeric a ⇒ Int → Int → Gradee (Matrix a) (Matrix a)
padMat w h = gradee pad' unpad' where
	pad' m = diagBlock [zeroes', m, zeroes'] where
		zeroes' = build (h, w) (const $ const 0)
	unpad' = subMatrix (h, w) ∘ size


-- Copyright (c) 2015 James M. Lawrence
--
-- Permission is hereby granted, free of charge, to any person obtaining
-- a copy of this software and associated documentation files (the
-- "Software"), to deal in the Software without restriction, including
-- without limitation the rights to use, copy, modify, merge, publish,
-- distribute, sublicense, and/or sell copies of the Software, and to
-- permit persons to whom the Software is furnished to do so, subject to
-- the following conditions:
--
-- The above copyright notice and this permission notice shall be included
-- in all copies or substantial portions of the Software.
--
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
-- EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
-- MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
-- IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
-- CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
-- TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
-- SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
--
-----------------------------------------------------------------------------
-- |
-- Module      :  Math.CayleyDickson
-- Copyright   :  (c) James M. Lawrence
-- License     :  MIT
--
-- Maintainer  :  James M. Lawrence <llmjjmll@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Cayley-Dickson constructions (complex numbers, quaternions,
-- octonions, sedenions, etc.) over general scalars without limit to
-- the number of dimensions.
--
-- An element of this structure is composed of an m-dimensional
-- /scalar part/ and an m*(2^n - 1)-dimensional /pure part/ (unrelated
-- to Haskell's uses of "pure"). An element whose scalar part is zero
-- is called a /pure/. Construction with real scalars yields the
-- Cayley-Dickson algebras, in which case the scalar part is also
-- called the /real part/. Other structures may be obtained by
-- considering general scalars, for instance the quaternions over
-- complex scalars.
-----------------------------------------------------------------------------

module Math.CayleyDickson (
    -- * Types
    Nion, Complex, Quaternion, Octonion, Sedenion,

    -- * Construction
    nion, fromScalar, complex, quaternion, octonion, sedenion,

    -- * Operations
    dot, cross, sqnorm, norm, polar,
    -- ** Operations with scalars
    --
    -- | The mnemonic is that the period (".") is on the side of the
    -- scalar.
    (**.),
    (.+), (+.), (.-), (-.), (.*), (*.), (/.),

    -- * Accessors
    coord, coords, setCoord, scalarPart, purePart,

    -- * Constants
    basisElement,

    -- * Classes
    Conjugable(..),

    -- ** Tags
    Tag(tagVal),
    Tag0,  Tag1,  Tag2,  Tag3,  Tag4,  Tag5,  Tag6,  Tag7,  Tag8,  Tag9,
    Tag10, Tag11, Tag12, Tag13, Tag14, Tag15, Tag16, Tag17, Tag18, Tag19,
    Tag20, Tag21, Tag22, Tag23, Tag24, Tag25, Tag26, Tag27, Tag28, Tag29,
    Tag30,

  ) where

----------------------------------------------------------
-- import

import Data.List (genericSplitAt, genericTake, genericReplicate, genericLength)
import Data.Bits (Bits, testBit)
import Data.Proxy (Proxy(Proxy))
import qualified Data.Int as Z
import qualified Data.Ratio as Q
import qualified Data.Complex as C
import qualified Data.Fixed as F
import qualified Data.Word as W

----------------------------------------------------------
-- infix

infix 6 :@

infix 6 .+
infix 6 +.
infix 6 .-
infix 6 -.

infix 7 .*
infix 7 *.
infix 7 /.

infixr 8 **.

----------------------------------------------------------
-- Nion

-- | General Cayley-Dickson construction producing \"N-ions\". The
-- first parameter is a 'Tag' instance that determines the dimension,
-- which is 2 raised to 'tagVal'. The second parameter is the scalar
-- type.
data Nion n a = Scalar a | Nion n a :@ Nion n a

----------------------------------------------------------
-- basic operations

-- | Equivalent to @'coord' x 0@.
scalarPart :: Nion n a -> a
scalarPart (Scalar x) = x
scalarPart (x :@ _) = scalarPart x

-- | Equivalent to @'setCoord' x 0 0@.
purePart :: Num a => Nion n a -> Nion n a
purePart (Scalar _) = Scalar 0
purePart (x :@ y) = purePart x :@ y

-- | Dot product. For real scalars, @(1 \`dot\`)@ is equivalent to
-- @'scalarPart'@.
dot :: Conjugable a => Nion n a -> Nion n a -> a
-- (y * conj x + x * conj y) / 2
Scalar x `dot` Scalar y = scalarPart' $ y * conj x
x@(Scalar _) `dot` (y1 :@ _) = x `dot` y1
(x1 :@ _) `dot` y@(Scalar _) = x1 `dot` y
(x1 :@ x2) `dot` (y1 :@ y2) = (x1 `dot` y1) + (x2 `dot` y2)

-- | Cross product. The cross product of two pures yields a pure that
-- is orthogonal to both operands. For real scalars, @(1 \`cross\`)@
-- is equivalent to @'purePart'@.
cross :: Conjugable a => Nion n a -> Nion n a -> Nion n a
-- (y * conj x - x * conj y) / 2
x `cross` y = y * conj x -. x `dot` y

-- | Squared norm: the dot product of an element with itself.
sqnorm :: Conjugable a => Nion n a -> a
sqnorm x = x `dot` x

-- | Square root of @sqnorm@.
norm :: (Conjugable a, Floating a) => Nion n a -> a
norm = sqrt . sqnorm

-- | Promote a scalar, returning an element whose scalar part is the
-- argument and whose pure part is zero. The element behaves as if it
-- were padded with zeros, but no actual padding is done.
fromScalar :: a -> Nion n a
fromScalar = Scalar

----------------------------------------------------------
-- operations with scalars

leftScalarOp :: (Nion n a -> Nion n a -> Nion n a) -> a -> Nion n a -> Nion n a
leftScalarOp f x y = f (Scalar x) y

rightScalarOp :: (Nion n a -> Nion n a -> Nion n a) -> Nion n a -> a -> Nion n a
rightScalarOp f x y = f x (Scalar y)

-- | Equivalent to @'fromScalar' x + y@.
(.+) :: Conjugable a => a -> Nion n a -> Nion n a
(.+) = leftScalarOp (+)

-- | Equivalent to @'fromScalar' x - y@.
(.-) :: Conjugable a => a -> Nion n a -> Nion n a
(.-) = leftScalarOp (-)

-- | Equivalent to @'fromScalar' x * y@.
(.*) :: Conjugable a => a -> Nion n a -> Nion n a
(.*) = leftScalarOp (*)

-- | Equivalent to @x + 'fromScalar' y@.
(+.) :: Conjugable a => Nion n a -> a -> Nion n a
(+.) = rightScalarOp (+)

-- | Equivalent to @x - 'fromScalar' y@.
(-.) :: Conjugable a => Nion n a -> a -> Nion n a
(-.) = rightScalarOp (-)

-- | Equivalent to @x * 'fromScalar' y@.
(*.) :: Conjugable a => Nion n a -> a -> Nion n a
(*.) = rightScalarOp (*)

-- | Equivalent to @x / 'fromScalar' y@.
(/.) :: (Conjugable a, Fractional a) => Nion n a -> a -> Nion n a
(/.) = rightScalarOp (/)

-- | Raise to a scalar power.
(**.) :: (Tag n, Conjugable a, RealFloat a) => Nion n a -> a -> Nion n a
Scalar x **. y = Scalar $ x ** y
x **. y = exp (y .* log x)

----------------------------------------------------------
-- polar form and complex function application

realPolar :: RealFloat a => Nion n a -> a -> (a, a, Nion n a)
realPolar sqrtMinus1 r
  | r >= 0 = (r, 0, sqrtMinus1)
  | otherwise = (-r, pi, sqrtMinus1)

polarUsing :: (Conjugable a, RealFloat a) =>
              Nion n a -> Nion n a -> (a, a, Nion n a)
polarUsing sqrtMinus1 (Scalar r) = realPolar sqrtMinus1 r
polarUsing sqrtMinus1 x
  | sqnormp == 0 = realPolar sqrtMinus1 r
  | otherwise = (normx, acos (r / normx), u)
  where
    p = purePart x
    sqnormp = sqnorm p
    u = p /. sqrt sqnormp
    r = scalarPart x
    normx = sqrt $ sqnormp + r * conj r

polar' :: (Tag n, Conjugable a, RealFloat a) =>
          Proxy n -> Nion n a -> (a, a, Nion n a)
polar' n x
  | tagVal n == 0 = error "polar: no polar form in 1 dimension"
  | otherwise = polarUsing basisElement1 x

-- | Return @(s, t, u)@ such that (approximately)
--
--     @x == s .* 'exp' (t .* u)@
--
-- where @s@ and @t@ are scalars, @s >= 0@, and @u@ is a unit pure.
--
-- If @x@ has no pure part then @u@ is arbitrarily chosen to be the
-- first pure basis element.
polar :: (Tag n, Conjugable a, RealFloat a) => Nion n a -> (a, a, Nion n a)
polar = polar' Proxy

applyUsing' :: (Tag n, Conjugable a, RealFloat a) =>
               Proxy n ->
               Nion n a -> (a -> a) -> (C.Complex a -> C.Complex a) ->
               Nion n a -> Nion n a
applyUsing' n sqrtMinus1 fr f z
  | tagVal n == 0 = Scalar $ fr $ scalarPart z
  | otherwise = x .+ u *. y
  where (s, t, u) = polarUsing sqrtMinus1 z
        -- handle special cases for a little more accuracy
        x C.:+ y | t == 0 = f $ c s 0
                 | t == pi = f $ c (-s) 0
                 | otherwise = f $ c s 0 * exp (c t 0 * c 0 1)
                 where c = (C.:+)

applyUsing :: (Tag n, Conjugable a, RealFloat a) =>
              Nion n a -> (a -> a) -> (C.Complex a -> C.Complex a) ->
              Nion n a -> Nion n a
applyUsing = applyUsing' Proxy

----------------------------------------------------------
-- constants

fill' :: Tag n => Proxy n -> a -> Nion n a
fill' n s = f $ tagVal n where
  f 0 = Scalar s
  f k = f k' :@ f k' where k' = k - 1

fill :: Tag n => a -> Nion n a
fill = fill' Proxy

paddedZero :: (Tag n, Num a) => Nion n a
paddedZero = fill 0

validIndex :: (Tag n, Num b, Ord b) => Proxy n -> b -> Bool
validIndex n i = i >= 0 && i < 2 ^ tagVal n

basisElement' :: (Tag n, Conjugable a, Bits i, Integral i) =>
                 Proxy n -> i -> Nion n a
basisElement' _ 0 = Scalar 1
basisElement' n index
  | validIndex n index = setCoord paddedZero index 1
  | otherwise = error "basisElement: out of range"

-- | The nth basis element.
basisElement :: (Tag n, Conjugable a, Bits i, Integral i) => i -> Nion n a
basisElement = basisElement' Proxy

basisElement1 :: (Tag n, Conjugable a) => Nion n a
basisElement1 = basisElement (1 :: Integer)

----------------------------------------------------------
-- util

-- Proper Functor and Foldable instances would have to translate
-- top-level scalars to their equivalent representation with padded
-- zeros. The machinations needed for this indirection are rather
-- cumbersome relative to the benefits of having the instances, whose
-- use would seem uncommon.

smap :: (a -> a) -> Nion n a -> Nion n a
smap f (Scalar s) = Scalar $ f s
smap f (x :@ y) = smap f x :@ smap f y

sfoldr :: (a -> b -> b) -> b -> Nion n a -> b
sfoldr f acc (Scalar x) = f x acc
sfoldr f acc (x :@ y) = sfoldr f (sfoldr f acc y) x

----------------------------------------------------------
-- accessors

coords' :: (Tag n, Num a) => Proxy n -> Nion n a -> [a]
coords' n (Scalar x) = x : genericReplicate k 0 where
                         k = 2 ^ tagVal n - 1 :: Integer
coords' _ x = sfoldr (:) [] x

-- | List of coordinates for this element.
coords :: (Tag n, Num a) => Nion n a -> [a]
coords = coords' Proxy

coord' :: (Tag n, Num a, Integral b, Bits b) => Proxy n -> Nion n a -> b -> a
coord' _ (Scalar x) 0 = x
coord' n elt index
  | validIndex n index = case elt of
                           Scalar _ -> 0
                           _ -> f elt $ fromInteger $ tagVal n - 1
  | otherwise = error "coord: out of range"
  where
    f (Scalar x) _ = x
    f (x :@ y) k = case testBit index k of
                     False -> f x k'
                     True  -> f y k'
                   where k' = k - 1

-- | Get the nth coordinate.
coord :: (Tag n, Num a, Integral b, Bits b) => Nion n a -> b -> a
coord = coord' Proxy

setCoord' :: (Tag n, Conjugable a, Integral b, Bits b) =>
             Proxy n -> Nion n a -> b -> a -> Nion n a
setCoord' _ (Scalar _) 0 value = Scalar value
setCoord' n elt index value
  | validIndex n index = case elt of
                           Scalar x -> setCoord (x .+ paddedZero) index value
                           _ -> f elt $ fromInteger $ tagVal n - 1
  | otherwise = error "setCoord: out of range"
  where
    f (Scalar _) _ = Scalar value
    f (x :@ y) k = case testBit index k of
                     False -> f x k' :@ y
                     True  -> x :@ f y k'
                   where k' = k - 1

-- | Set the nth coordinate, returning a new element.
setCoord :: (Tag n, Conjugable a, Integral b, Bits b) =>
            Nion n a -> b -> a -> Nion n a
setCoord = setCoord' Proxy

----------------------------------------------------------
-- construction

fromList :: Integer -> [a] -> Nion n a
fromList _ (x:[]) = Scalar x
fromList k xs = fromList k' l :@ fromList k' r where
                  k' = k `div` 2
                  (l, r) = genericSplitAt k' xs

nion' :: (Tag n, Num a) => Proxy n -> [a] -> Nion n a
nion' n elems = fromList d $ taken ++ padding where
                  d = 2 ^ tagVal n
                  taken = genericTake d elems
                  padding = genericReplicate (d - genericLength taken) 0

-- | Construct an element from a list of coordinates. If the list is
-- too small then the remaining coordinates are padded with zeros. If
-- the list is too large then the extra values are ignored.
nion :: (Tag n, Num a) => [a] -> Nion n a
nion = nion' Proxy

----------------------------------------------------------
-- instances

instance (Tag n, Show a, Num a) => Show (Nion n a) where
  show x = "nion " ++ show (coords x)

instance (Conjugable a, Eq a) => Eq (Nion n a) where
  Scalar x == Scalar y = x == y
  x@(Scalar _) == y1 :@ y2 = x == y1 && y2 == 0
  x1 :@ x2 == y@(Scalar _) = x1 == y && x2 == 0
  x1 :@ x2 == y1 :@ y2 = x1 == y1 && x2 == y2

instance Conjugable a => Num (Nion n a) where
  Scalar x + Scalar y = Scalar $ x + y
  x@(Scalar _) + (y1 :@ y2) = (x + y1) :@ y2
  (x1 :@ x2) + y@(Scalar _) = (x1 + y) :@ x2
  (x1 :@ y1) + (x2 :@ y2) = (x1 + x2) :@ (y1 + y2)

  Scalar x - Scalar y = Scalar $ x - y
  x@(Scalar _) - (y1 :@ y2) = (x - y1) :@ negate y2
  (x1 :@ x2) - y@(Scalar _) = (x1 - y) :@ x2
  (x1 :@ y1) - (x2 :@ y2) = (x1 - x2) :@ (y1 - y2)

  Scalar x * Scalar y = Scalar $ x * y
  x@(Scalar _) * (y1 :@ y2) = (x * y1) :@ (x * y2)
  (x1 :@ x2) * y@(Scalar _) = (x1 * y) :@ (x2 * y)
  (x1 :@ x2) * (y1 :@ y2) = (x1 * y1 - conj y2 * x2) :@ (y2 * x1 + x2 * conj y1)

  negate = smap negate
  fromInteger = fromScalar . fromInteger
  abs = doNotUse
  signum = doNotUse

instance (Conjugable a, Fractional a) => Fractional (Nion n a) where
  Scalar x / Scalar y = Scalar $ x / y
  x@(_ :@ _) / Scalar y = smap (/ y) x
  x / y = (x * conj y) /. sqnorm y

  recip x = conj x /. sqnorm x
  fromRational = fromScalar . fromRational

-- | The first pure basis element is arbitrarily chosen as sqrt (-1).
instance (Tag n, Conjugable a, RealFloat a) => Floating (Nion n a) where
  pi    = Scalar pi
  exp   = applyUsing basisElement1 exp exp
  log   = applyUsing basisElement1 log log
  sqrt  = applyUsing basisElement1 sqrt sqrt
  sin   = applyUsing basisElement1 sin sin
  cos   = applyUsing basisElement1 cos cos
  tan   = applyUsing basisElement1 tan tan
  asin  = applyUsing basisElement1 asin asin
  acos  = applyUsing basisElement1 acos acos
  atan  = applyUsing basisElement1 atan atan
  sinh  = applyUsing basisElement1 sinh sinh
  cosh  = applyUsing basisElement1 cosh cosh
  tanh  = applyUsing basisElement1 tanh tanh
  asinh = applyUsing basisElement1 asinh asinh
  acosh = applyUsing basisElement1 acosh acosh
  atanh = applyUsing basisElement1 atanh atanh

----------------------------------------------------------
-- convenience types

-- | Complex numbers, the 2^1-dimensional construction.
type Complex a = Nion Tag1 a

-- | Quaternions, the 2^2-dimensional construction.
type Quaternion a = Nion Tag2 a

-- | Octonions, the 2^3-dimensional construction.
type Octonion a = Nion Tag3 a

-- | Sedenions, the 2^4-dimensional construction.
type Sedenion a = Nion Tag4 a

-- | Construct a complex number.
complex :: a -> a -> Complex a
complex x y = (:@) (Scalar x) (Scalar y)

-- | Construct a quaternion.
quaternion :: a -> a -> a -> a -> Quaternion a
quaternion w x y z = (:@) ((:@) (Scalar w) (Scalar x))
                          ((:@) (Scalar y) (Scalar z))

-- | Construct an octonion.
octonion :: a -> a -> a -> a ->
            a -> a -> a -> a -> Octonion a
octonion s t u v
         w x y z = (:@) ((:@) ((:@) (Scalar s) (Scalar t))
                              ((:@) (Scalar u) (Scalar v)))
                        ((:@) ((:@) (Scalar w) (Scalar x))
                              ((:@) (Scalar y) (Scalar z)))

-- | Construct a sedenion.
sedenion :: a -> a -> a -> a ->
            a -> a -> a -> a ->
            a -> a -> a -> a ->
            a -> a -> a -> a -> Sedenion a
sedenion k l m n
         o p q r
         s t u v
         w x y z = (:@) ((:@) ((:@) ((:@) (Scalar k) (Scalar l))
                                    ((:@) (Scalar m) (Scalar n)))
                              ((:@) ((:@) (Scalar o) (Scalar p))
                                    ((:@) (Scalar q) (Scalar r))))
                        ((:@) ((:@) ((:@) (Scalar s) (Scalar t))
                                    ((:@) (Scalar u) (Scalar v)))
                              ((:@) ((:@) (Scalar w) (Scalar x))
                                    ((:@) (Scalar y) (Scalar z))))

----------------------------------------------------------
-- Conjugable

-- | The /conjugate/ of an element is obtained by negating the pure
-- part and conjugating the scalar part. The conjugate of a real
-- number is the identity ('id'), which is the default implementation.
class Num a => Conjugable a where
  -- | Conjugate.
  conj :: a -> a
  conj = id

  -- | Equivalent to @(x + conj x) / 2@.
  scalarPart' :: a -> a
  scalarPart' = id

instance Conjugable a => Conjugable (Nion n a) where
  conj (Scalar x) = Scalar $ conj x
  conj (x :@ y) = conj x :@ negate y

  scalarPart' (Scalar x) = Scalar $ scalarPart' x
  scalarPart' (x :@ _) = scalarPart' x

instance (Conjugable a, RealFloat a) => Conjugable (C.Complex a) where
  conj = C.conjugate
  scalarPart' (x C.:+ _) = scalarPart' x C.:+ 0

instance Conjugable Int
instance Conjugable Integer
instance Conjugable Float
instance Conjugable Double
instance Conjugable Z.Int8
instance Conjugable Z.Int16
instance Conjugable Z.Int32
instance Conjugable Z.Int64
instance Conjugable W.Word8
instance Conjugable W.Word16
instance Conjugable W.Word32
instance Conjugable W.Word64
instance Integral a => Conjugable (Q.Ratio a)
instance F.HasResolution a => Conjugable (F.Fixed a)

-----------------------------------------------------------------------------
-- doNotUse

rant :: String
rant = unlines $
  ["",
   "The Num class is a bit messed up, having tied (+), (-), and (*) to abs",
   "and signum. Number systems that have no appropriate definition for abs",
   "or signum must either invent their own operators for addition,",
   "subtraction, and multiplication, else break the contract with Num by",
   "raising an error such as this one when someone uses abs or signum.",
   "",
   "For some time I resisted hijacking Num, but eventually the replacement",
   "operators became too cumbersome and, coupled with the lack of numeric",
   "promotion, significantly detracted from the usability of the package.",
   "So here we are. Good luck, and stay away from abs and signum, which",
   "officially have cooties."]

doNotUse :: a -> a
doNotUse _ = error rant

----------------------------------------------------------
-- Tag

-- | Tags serve to determine a type's dimension, which is 2 raised to
-- `tagVal`. Tag instances are included for convenience only, as you
-- may create your own tag.
class Tag n where
  tagVal :: Proxy n -> Integer

data     Tag0
data     Tag1
data     Tag2
data     Tag3
data     Tag4
data     Tag5
data     Tag6
data     Tag7
data     Tag8
data     Tag9
data    Tag10
data    Tag11
data    Tag12
data    Tag13
data    Tag14
data    Tag15
data    Tag16
data    Tag17
data    Tag18
data    Tag19
data    Tag20
data    Tag21
data    Tag22
data    Tag23
data    Tag24
data    Tag25
data    Tag26
data    Tag27
data    Tag28
data    Tag29
data    Tag30

instance Tag    Tag0 where tagVal _ =    0
instance Tag    Tag1 where tagVal _ =    1
instance Tag    Tag2 where tagVal _ =    2
instance Tag    Tag3 where tagVal _ =    3
instance Tag    Tag4 where tagVal _ =    4
instance Tag    Tag5 where tagVal _ =    5
instance Tag    Tag6 where tagVal _ =    6
instance Tag    Tag7 where tagVal _ =    7
instance Tag    Tag8 where tagVal _ =    8
instance Tag    Tag9 where tagVal _ =    9
instance Tag   Tag10 where tagVal _ =   10
instance Tag   Tag11 where tagVal _ =   11
instance Tag   Tag12 where tagVal _ =   12
instance Tag   Tag13 where tagVal _ =   13
instance Tag   Tag14 where tagVal _ =   14
instance Tag   Tag15 where tagVal _ =   15
instance Tag   Tag16 where tagVal _ =   16
instance Tag   Tag17 where tagVal _ =   17
instance Tag   Tag18 where tagVal _ =   18
instance Tag   Tag19 where tagVal _ =   19
instance Tag   Tag20 where tagVal _ =   20
instance Tag   Tag21 where tagVal _ =   21
instance Tag   Tag22 where tagVal _ =   22
instance Tag   Tag23 where tagVal _ =   23
instance Tag   Tag24 where tagVal _ =   24
instance Tag   Tag25 where tagVal _ =   25
instance Tag   Tag26 where tagVal _ =   26
instance Tag   Tag27 where tagVal _ =   27
instance Tag   Tag28 where tagVal _ =   28
instance Tag   Tag29 where tagVal _ =   29
instance Tag   Tag30 where tagVal _ =   30


-- |
-- Copyright  : (c) 2012 Daniël de Kok
-- License    : BSD3
--
-- Maintainer : Daniël de Kok <me@danieldk.eu>
-- Stability  : experimental
--
-- This module provides functionality to perform approximate randomization
-- tests (Noreen, 1989).

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DoAndIfThenElse #-}
{-# LANGUAGE FlexibleContexts #-}

module Statistics.Test.ApproxRand (
  -- * Description
  -- $description

  -- * Examples
  -- $examples

  -- * Data types
  TestOptions(..),
  TestResult(..),
  Significance(..),
  RandWithError,

  -- * Approximate randomization tests
  approxRandTest,
  approxRandStats,

  approxRandPairTest,
  approxRandPairStats,

  -- * Test statistics
  TestStatistic,
  differenceMean,
  meanDifference,
  varianceRatio
) where

import           Prelude hiding ((++))
import           Control.Monad (liftM, replicateM, when)
import           Control.Monad.Error (ErrorT)
import           Control.Monad.Error.Class (throwError)
import           Control.Monad.Mersenne.Random (R(..), Rand(..), getBool)
import           Control.Monad.ST (runST)
import           Control.Monad.Trans.Class (lift)
import           Data.Vector.Generic ((++))
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as GM
import           Data.Word (Word)
import           Statistics.Sample (variance)
import           Statistics.Test.Types (TestType(..))
import           Statistics.Types
import           System.Random.Mersenne.Pure64 (PureMT, randomInt, randomWord)

-- $description
--
-- Approximate randomization tests rely on a simple premise: given a test
-- statistic, if the null-hypothesis (the samples do not differ) is true,
-- we can randomly swap values between samples without an (extreme) impact
-- on the test statistic. Otherwise, the null-hypothesis must be rejected.
--
-- The test works by generating a given number of sample shuffles and computing
-- the test statistic for each shuffle. If /r/ is the number of shuffled
-- samples where the test statistic is at least as high as the test statistic
-- applied on the original samples; and /N/ the number of shuffles, then
-- the null-hypothesis is rejected iff /(r + 1):(N + 1) < p-value/ (for
-- one-sided tests).
--
-- Two kinds of test are supported:
--
-- * /Paired sample/ ('approxRandPairTest'): values from samples are shuffled
--   pair-wise. This requires the samples to have an equal length.
--
-- * /Unpaired sample/ ('approxRandTest'): values from samples are shuffled
--   among both samples. Consequently the i-th element of one sample does not
--   bear a relationship with the i-th element of the other sample. The
--   shuffled samples retain the sizes of the original samples.
--
-- Both tests can be performed as a one-tailed or two-tailed test.

-- $examples
-- Both unpaired and paired sample tests use the 'Rand' monad to obtain
-- random numbers. We can obtain a pseudo-random number generator that
-- is seeded using the system clock using the
-- 'System.Random.Mersenne.Pure64.newPureMT' function (please refer to
-- the documentation of 'System.Random.Mersenne.Pure64' for more
-- information):
--
-- > prng <- newPureMT
--
-- Suppose that we have the samples 's1' and 's2'. We could now perform
-- a Two-Tailed randomization test with 10,000 shuffles and the mean
-- difference as the test statistic, by running 'approxRandTest' in the 'Rand'
-- monad (at the /p = 0.01/ level):
--
-- > evalRandom (approxRandTest TwoTailed meanDifference 10000 0.01 s1 s2) prng
--
-- It is also possible to obtain the test statistics of the shuffled samples
-- directly (e.g. to inspect the distribution of test statistics) using the
-- 'approxRandStats'/'approxRandPiarStats' functions:
--
-- > evalRandom (approxRandStats meanDifference 10000 0.01 s1 s2) prng

-- | Computations with random numbers that can fail.
type RandWithError a = ErrorT String Rand a

-- | Options for randomization tests
--
data TestOptions = TestOptions {
    toTestType      :: TestType,      -- ^ Type of test ('OneTailed' or 'TwoTailed')
    toTestStatistic :: TestStatistic, -- ^ Test statistic
    toIterations    :: Int,           -- ^ Number of shuffled samples to create
    toPValue        :: Double         -- ^ he p-value at which to test (e.g. 0.05)
}

-- |
-- The result of hypothesis testing.
data TestResult = TestResult {
    trSignificance     :: Significance, -- ^ Significance
    trStat             :: Double,       -- ^ Test statistic for the samples
    trRandomizedStats  :: Sample        -- ^ Test statistics for the
                                        --   randomized samples
  } deriving (Eq, Ord, Show)

-- |
-- Significance.
data Significance =
    Significant    Double -- ^ The null hypothesis should be rejected
  | NotSignificant Double -- ^ Data is compatible with the null hypothesis
  deriving (Eq, Ord, Show)

-- |
-- Apply a pair-wise approximate randomization test.
--
-- In pair-wise approximate randomization tests the data points at a given
-- index are swapped between samples with a probability of 0.5. Since
-- swapping is pairwise, the samples should have the same length.
approxRandPairTest ::
     TestOptions              -- ^ Options for the test
  -> Sample                   -- ^ First sample
  -> Sample                   -- ^ Second sample
  -> RandWithError TestResult -- ^ The test result
approxRandPairTest (TestOptions testType stat n pTest) s1 s2 = do
  stats <- approxRandPairStats stat n s1 s2
  let tOrig = stat s1 s2
  let sig = significance testType pTest n $ countExtremes tOrig stats
  return $ TestResult sig tOrig stats

-- |
-- Apply an approximate randomization test.
--
-- In approximate randomization tests, the values of two samples are
-- shuffled among those samples. A test statistic is calculated for
-- the original samples and the shuffled samples, to detect whether the
-- difference of the samples is extreme or not.
approxRandTest ::
     TestOptions     -- ^ Options for the test
  -> Sample          -- ^ First sample
  -> Sample          -- ^ Second sample
  -> Rand TestResult -- ^ The test result
approxRandTest (TestOptions testType stat n pTest) s1 s2 = do
  stats <- approxRandStats stat n s1 s2
  let tOrig = stat s1 s2
  let sig = significance testType pTest n $ countExtremes tOrig stats
  return $ TestResult sig tOrig stats

-- | Determine the significance.
significance ::
     TestType     -- ^ Type of test ('OneTailed' or 'TwoTailed')
  -> Double       -- ^ The p-value at which to test (e.g. 0.05)
  -> Int          -- ^ Number of sample shuffles
  -> (Int, Int)   -- ^ Extreme statistic counts
  -> Significance -- ^ The test result
significance TwoTailed pTest n =
  significant (pTest / 2) . pValue n . uncurry min
significance OneTailed pTest n =
  significant pTest . pValue n . snd

-- | Wrap a p-value in a 'TestResult'.
significant ::
     Double       -- ^ The p-value at which to test
  -> Double       -- ^ The p-value
  -> Significance -- ^ The test result
significant pTail p =
  if p < pTail then Significant p else NotSignificant p

-- | Calculate a p-value
pValue ::
     Int    -- ^ Number of extreme outcomes
  -> Int    -- ^ Number of shuffles
  -> Double -- ^ The p-value
pValue n r = (fromIntegral r + 1) / (fromIntegral n + 1)

-- |
-- Count extreme test statistic values. If the test statistic value of the
-- original sample is in the right tail, we want to count values equal to
-- or larger than that value. If the value is in the left tail, we want to
-- count value smaller than or equal to that value. Since we do not know
-- the tail (yet), we count both.
--
-- Note: we can determine the tail by (1) averaging the test statistics of
-- the randomized samples, or (2) taking the smaller of the two counts.
countExtremes ::
     Double     -- ^ Test statistic value of the original samples
  -> Sample     -- ^ Test statistic values of the randomized samples.
  -> (Int, Int) -- ^ Count of left- and right-tail extremes.
countExtremes tOrig =
  VG.foldl' count (0, 0)
  where
    count (left, right) tPerm =
      let !newLeft = if tPerm <= tOrig then succ left else left in
      let !newRight = if tPerm >= tOrig then succ right else right in
        (newLeft, newRight)

-- |
-- Generate a given number of pairwise shuffled samples, and calculate
-- the test statistic for each shuffle.
--
-- Since the data points at a given index are swapped (with a probability of
-- 0.5), the samples should have the same length.
approxRandPairStats ::
     TestStatistic          -- ^ Test statistic
  -> Int                    -- ^ Number of shuffled samples to create
  -> Sample                 -- ^ First sample
  -> Sample                 -- ^ Second sample
  -> RandWithError Sample   -- ^ The statistics of the shuffles
approxRandPairStats stat n s1 s2 = do
  when (VG.length s1 /= VG.length s2) $
    throwError "Cannot calculate pairwise statistic: samples have different sizes"
  lift $ liftM VG.fromList $ replicateM n $
    uncurry stat `liftM` shuffleVectorsPairwise s1 s2

-- |
-- Generate a given number of shuffled samples, and calculate the test
-- statistic for each shuffle.
--
-- This function does not require the samples to have an equal length.
approxRandStats ::
     TestStatistic -- ^ Test statistic
  -> Int           -- ^ Number of shuffled samples to create
  -> Sample        -- ^ First sample
  -> Sample        -- ^ Second sample
  -> Rand Sample   -- ^ The statistics of the shuffles
approxRandStats stat n s1 s2 =
  liftM VG.fromList $ replicateM n $ uncurry stat `liftM` shuffleVectors s1 s2

-- | Pair-wise shuffle of two vectors.
shuffleVectorsPairwise :: (VG.Vector v a, VG.Vector v Bool) =>
  v a -> v a -> Rand (v a, v a)
shuffleVectorsPairwise vec1 vec2 = do
  randomVec <- randomVector (VG.length vec1)
  let pv1 = VG.zipWith3 permute vec1 vec2 randomVec
  let pv2 = VG.zipWith3 permute vec2 vec1 randomVec
  return (pv1, pv2)
  where
    permute val1 val2 coin =
      if coin then val1 else val2

randomVector :: VG.Vector v Bool => Int -> Rand (v Bool)
randomVector len =
  VG.replicateM len getBool

-- Shuffle values amongst two vectors, keeping the original vector lengths.
shuffleVectors :: VG.Vector v a => v a -> v a -> Rand (v a, v a)
shuffleVectors v1 v2 = do
  shuffledVectors <- shuffleVector $ v1 ++ v2
  return (VG.slice 0 (VG.length v1) shuffledVectors,
    VG.slice (VG.length v1) (VG.length v2) shuffledVectors)

-- Fisher-Yates shuffle in the Rand monad
shuffleVector :: VG.Vector v a => v a -> Rand (v a)
shuffleVector v =
  Rand $ \s -> case shuffleVector' s v of (sv, s') -> R sv s'

-- Fisher-Yates shuffle
shuffleVector' :: VG.Vector v a => PureMT -> v a -> (v a, PureMT)
shuffleVector' gen v = runST $ do
  let maxIdx = VG.length v - 1
  vm   <- VG.thaw v
  gen' <- swaps vm 0 maxIdx gen
  vmf  <- VG.unsafeFreeze vm
  return (vmf, gen')
  where
    swaps vm idx maxIdx gen'
      | idx < maxIdx = do
          let (newIdx, gen'') = randomIntR gen' (idx, maxIdx)
          GM.unsafeSwap vm idx newIdx
          swaps vm (idx + 1) maxIdx gen''
      | otherwise = return gen'

-- |
-- A test stastic calculates the difference between two samples. See
-- 'meanDifference' and 'varianceRatio' for examples.
type TestStatistic = Sample -> Sample -> Double

-- |
-- Calculates the difference mean of two samples (/mean(s1 - s2)/). When the
-- two samples do not have an equal length, the trailing elements of the
-- longer vector are ignored.
differenceMean :: TestStatistic
differenceMean v1 v2 =
  VG.sum (subVector v1 v2) / fromIntegral (VG.length v1)

-- | Calculates the mean difference of two samples (/mean(s1) - mean(s2)/).
meanDifference :: TestStatistic
meanDifference s1 s2 =
  mean s1 - mean s2

-- | Calculate the mean of a sample.
mean :: Sample -> Double
mean = do
  t <- VG.sum
  l <- VG.length
  return $ t / fromIntegral l

-- | Calculate the ratio of sample variances (/var(s1) : var(s2)/).
varianceRatio :: TestStatistic
varianceRatio v1 v2 =
  variance v1 / variance v2

-- | Subtract two vectors.
subVector :: (VG.Vector v n, Num n) => v n -> v n -> v n
subVector = VG.zipWith (-)

subIIW :: Int -> Int -> Word
subIIW a b = fromIntegral a - fromIntegral b
{-# INLINE subIIW #-}

addIWI :: Int -> Word -> Int
addIWI a b = a + fromIntegral b
{-# INLINE addIWI #-}

-- | Generate Int numbers within a range
randomIntR :: PureMT -> (Int, Int) -> (Int, PureMT)
randomIntR gen (a, b)
  | n == 0    = randomInt gen
  | otherwise = loop gen
  where
    (a', b') = if a < b then (a, b) else (b, a)
    -- Number of different Ints that should be generated
    n = 1 + subIIW b' a'
    -- The total range of Word can hold x complete n ranges
    x = maxBound `div` n
    -- Pick from a range the is dividable by n without remainders
    s = x * n
    loop gen'
      | r >= s    = loop gen'' -- r is outside the range, discard it...
      | otherwise = (addIWI a' (r `div` x), gen'')
      where
        (!r, !gen'') = randomWord gen'
{-# INLINE randomIntR #-}


module Math.IRT.Model.ThreePLM
  ( ThreePLM (..)
  ) where

import Statistics.Distribution

import Math.IRT.Internal.Distribution
import Math.IRT.Internal.LogLikelihood
import Math.IRT.Model.FourPLM ( FourPLM(..) )
import Math.IRT.Model.Generic

data ThreePLM = ThreePLM { discrimination :: !Double
                         , difficulty     :: !Double
                         , pseudoGuessing :: !Double
                         } deriving (Show)

instance Distribution ThreePLM where
    cumulative = cumulative . toFourPLM

instance ContDistr ThreePLM where
    density    = density . toFourPLM
    quantile _ = error "This shouldn't be needed"

instance DensityDeriv ThreePLM where
    densityDeriv = densityDeriv . toFourPLM

instance GenericModel ThreePLM where
    fromRasch         b = ThreePLM 1.0 b 0.0
    fromOnePLM        b = ThreePLM 1.7 b 0.0
    fromTwoPLM      a b = ThreePLM   a b 0.0
    fromThreePLM        = ThreePLM
    fromFourPLM a b c _ = ThreePLM   a b   c

instance LogLikelihood ThreePLM where
    logLikelihood b = logLikelihood b . toFourPLM

toFourPLM :: ThreePLM -> FourPLM
toFourPLM (ThreePLM sa sb sc) = FourPLM sa sb sc 1.0


module Math.IRT.Model.Rasch
  ( RaschModel (..)
  ) where

import Statistics.Distribution

import Math.IRT.Internal.Distribution
import Math.IRT.Internal.LogLikelihood
import Math.IRT.Model.FourPLM ( FourPLM(..) )
import Math.IRT.Model.Generic

data RaschModel = RaschModel { difficulty :: !Double
                             } deriving (Show)

instance Distribution RaschModel where
    cumulative = cumulative . toFourPLM

instance ContDistr RaschModel where
    density    = density . toFourPLM
    quantile _ = error "This shouldn't be needed"

instance DensityDeriv RaschModel where
    densityDeriv = densityDeriv . toFourPLM

instance GenericModel RaschModel where
    fromRasch           = RaschModel
    fromOnePLM          = RaschModel
    fromTwoPLM      _ b = RaschModel b
    fromThreePLM  _ b _ = RaschModel b
    fromFourPLM _ b _ _ = RaschModel b

instance LogLikelihood RaschModel where
    logLikelihood b = logLikelihood b . toFourPLM

toFourPLM :: RaschModel -> FourPLM
toFourPLM (RaschModel sb) = FourPLM 1.0 sb 0.0 1.0


{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict #-}
module STC.Convolution where

import           Control.Arrow
import           Control.Monad              as M
import           Control.Monad.Parallel     as MP (bindM2, mapM)
import           Data.Array.IArray          as IA
import           Data.Array.Repa            as R
import           Data.Complex
import           Data.Ix
import           Data.List                  as L
import           Data.Vector.Storable       as VS
import           Data.Vector.Unboxed        as VU
import           Debug.Trace
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.FourierSeries
import           Image.IO
import           STC.DFTArray
import           Types
import           Utils.Parallel
import FokkerPlanck.GPUKernel
import qualified Data.Array.Accelerate              as A
import qualified Data.Array.Accelerate.Data.Complex as A
import           Data.Array.Accelerate.LLVM.PTX
import Debug.Trace
import Utils.List
import Utils.BLAS

data Field
  = Source
  | Sink
  deriving (Read, Show)

{-# INLINE dftHarmonicsArray #-}
dftHarmonicsArray ::
     DFTPlan
  -> Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> IO (IA.Array (Int, Int) (VS.Vector (Complex Double)))
dftHarmonicsArray !plan !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff = do
  let !arr =
        computeHarmonicsArray
          numRows
          deltaRow
          numCols
          deltaCol
          phiFreqs
          rhoFreqs
          thetaFreqs
          rFreqs
          halfLogPeriod
          cutoff
  fmap (listArray (bounds arr)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [numCols, numRows] [0, 1]) .
    L.map
      (VS.convert .
       toUnboxed .
       computeUnboxedS . makeFilter2D . fromUnboxed (Z :. numCols :. numRows)) .
    IA.elems $
    arr

{-# INLINE dftHarmonicsArrayG #-}
dftHarmonicsArrayG ::
     DFTPlan
  -> Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double -> VS.Vector (Complex Double)
  -> IO (IA.Array (Int, Int) (VS.Vector (Complex Double)))
dftHarmonicsArrayG !plan !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff !gaussian = do
  let !arr =
        computeHarmonicsArray
          numRows
          deltaRow
          numCols
          deltaCol
          phiFreqs
          rhoFreqs
          thetaFreqs
          rFreqs
          halfLogPeriod
          cutoff
  xs <-
    fmap (L.map (VS.zipWith (*) gaussian)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [numCols, numRows] [0, 1]) .
    IA.elems $
    arr
  ys <- dftExecuteBatchP plan (DFTPlanID IDFT1DG [numCols, numRows] [0, 1]) xs
  fmap (listArray (bounds arr)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [numCols, numRows] [0, 1]) .
    L.map
      (\vec ->
         let arr = fromUnboxed (Z :. numCols :. numRows) . VS.convert $ vec
         in VS.convert .
            toUnboxed . computeUnboxedS . makeFilter2D . R.traverse arr id $ \f idx@(Z :. i :. j) ->
              if (i == div numCols 2 && j == div numRows 2)
                 then 0
                 else f idx) $
    ys

{-# INLINE convolve #-}
convolve ::
     Field
  -> DFTPlan
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> IO DFTArray
convolve !field !plan !coefficients !harmonicsArray !arr@(DFTArray rows cols thetaFreqs rFreqs vecs) = do
  dftVecs <- dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs
  let !initVec = VS.replicate (VS.length . L.head $ vecs) 0
      idx = (,) <$> (L.zip [0 ..] rFreqs) <*> (L.zip [0 ..] thetaFreqs)
  fmap (DFTArray rows cols thetaFreqs rFreqs) .
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
    parMap
      rdeepseq
      (\((!r, !rFreq), (!theta, !thetaFreq)) ->
         L.foldl'
           (\(!vec) (((!rho, !rhoFreq), (!phi, !phiFreq)), inputVec) ->
              VS.zipWith
                (+)
                vec
                (VS.map (* (coefficients R.! (Z :. r :. theta :. rho :. phi))) .
                 VS.zipWith
                   (*)
                   (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq) $
                 inputVec))
           initVec .
         L.zip idx $
         dftVecs) $
    idx

{-# INLINE convolve' #-}
convolve' ::
     Field
  -> DFTPlan
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> IO DFTArray
convolve' !field !plan !coefficients !harmonicsArray !arr@(DFTArray rows cols thetaFreqs rFreqs vecs) = do
  dftVecs <- dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs
  let !initVec = VS.replicate (VS.length . L.head $ vecs) 0
      idxTheta = L.zip [0 ..] thetaFreqs
  fmap (DFTArray rows cols thetaFreqs rFreqs) .
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
    parMap
      rdeepseq
      (\(!theta, !thetaFreq) ->
         VS.map
           (\x ->
              case field of
                Source -> x
                Sink -> x * cis (pi * thetaFreq)) .
         L.foldl'
           (\vec1 (rho, rhoFreq) ->
              L.foldl'
                (\(!vec2) ((!phi, !phiFreq), inputVec) ->
                   VS.zipWith
                     (+)
                     vec2
                     (VS.map
                        (* (coefficients R.!
                            (Z :. (0 :: Int) :. theta :. rho :. phi))) .
                      VS.zipWith
                        (*)
                        (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq 0) $
                      inputVec))
                vec1 .
              L.zip idxTheta $
              dftVecs)
           initVec .
         L.zip [0 ..] $
         rFreqs) $
    idxTheta

{-# INLINE convolveG' #-}
convolveG' ::
     Field
  -> DFTPlan
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> DFTArray
  -> IO DFTArray
convolveG' !field !plan !coefficients !harmonicsArray !gaussianFilter !arr@(DFTArray rows cols thetaFreqs rFreqs vecs) = do
  dftVecs <-
    L.map (VS.zipWith (*) gaussianFilter) <$>
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs
  let !initVec = VS.replicate (VS.length . L.head $ vecs) 0
      idxTheta = L.zip [0 ..] thetaFreqs
  fmap (DFTArray rows cols thetaFreqs rFreqs) .
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
    parMap
      rdeepseq
      (\(!theta, !thetaFreq) ->
         L.foldl'
           (\vec1 (rho, rhoFreq) ->
              L.foldl'
                (\(!vec2) ((!phi, !phiFreq), inputVec) ->
                   VS.zipWith
                     (+)
                     vec2
                     (VS.map
                        (* (case field of
                              Source ->
                                coefficients R.!
                                (Z :. (0 :: Int) :. theta :. rho :. phi)
                              Sink ->
                                (coefficients R.!
                                 (Z :. (0 :: Int) :. theta :. rho :. phi)) *
                                (cis $ (-(phiFreq + thetaFreq)) * pi))) .
                      VS.zipWith
                        (*)
                        (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq 0) $
                      inputVec))
                vec1 .
              L.zip idxTheta $
              dftVecs)
           initVec .
         L.zip [0 ..] $
         rFreqs) $
    idxTheta

-- {-# INLINE convolve'' #-}
-- convolve'' ::
--      Field
--   -> DFTPlan
--   -> R.Array U DIM4 (Complex Double)
--   -> IA.Array (Int, Int) (VS.Vector (Complex Double))
--   -> DFTArray
--   -> IO DFTArray
-- convolve'' !field !plan !coefficients !harmonicsArray !arr@(DFTArray rows cols phiFreqs rhoFreqs vecs) = do
--   let !initVec = VS.replicate (VS.length . L.head $ vecs) 0
--       idx = L.zip [0 ..] phiFreqs
--   return .
--     (DFTArray rows cols phiFreqs rhoFreqs) .
--     parMap
--       rdeepseq
--       (\(theta, thetaFreq) ->
--          L.foldl'
--            (\vec1 (rho, rhoFreq) ->
--               L.foldl'
--                 (\vec2 (phi, phiFreq) ->
--                    VS.zipWith (+) vec2 .
--                    VS.map
--                      (* (case field of
--                            Source ->
--                              coefficients R.!
--                              (Z :. (0 :: Int) :. theta :. rho :. phi)
--                            Sink ->
--                              (coefficients R.!
--                               (Z :. (0 :: Int) :. theta :. rho :. phi)) *
--                              (cis $ (-(thetaFreq + phiFreq)) * pi))) $
--                    (getHarmonics harmonicsArray phiFreq rhoFreq 0 0))
--                 vec1
--                 idx)
--            initVec .
--          L.zip [0 ..] $
--          rhoFreqs) $
--     idx

-- {-# INLINE dftHarmonicsArrayGPU #-}
-- dftHarmonicsArrayGPU ::
--      DFTPlan
--   -> Int
--   -> Double
--   -> Int
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> IO (Acc (A.Array A.DIM4 (A.Complex Double)))
-- dftHarmonicsArrayGPU !plan !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff = do
--   let !arr =
--         computeHarmonicsArrayGPU
--           numRows
--           deltaRow
--           numCols
--           deltaCol
--           phiFreqs
--           rhoFreqs
--           thetaFreqs
--           rFreqs
--           halfLogPeriod
--           cutoff
--       rangeFunc2 xs ys =
--         (round (L.head xs - L.last ys), round (L.last xs - L.head ys))
--       (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
--       (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
--   fmap
--     (A.use .
--      A.fromList
--        (A.Z A.:. (rfUB - rfLB + 1) A.:. (tfUB - tfLB + 1) A.:. numCols A.:.
--         numRows) .
--      VS.toList . VS.concat) .
--     dftExecuteBatchP plan (DFTPlanID DFT1DG [numCols, numRows] [0, 1]) .
--     L.map
--       (VS.convert .
--        toUnboxed .
--        computeUnboxedS . makeFilter2D . fromUnboxed (Z :. numCols :. numRows)) .
--     IA.elems $
    -- arr

-- {-# INLINE convolveGPU #-}
-- convolveGPU ::
--      PTX
--   -> Field
--   -> DFTPlan
--   -> Acc (A.Array A.DIM4 (A.Complex Double))
--   -> Acc (A.Array A.DIM4 (A.Complex Double))
--   -> Acc (A.Array A.DIM4 (A.Complex Double))
--   -> DFTArray
--   -> IO DFTArray
-- convolveGPU !ptx !field !plan !coefficients !coefficientsSink !harmonics !arr@(DFTArray rows cols thetaFreqs rFreqs vecs) = do
--   dftVecs <- dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs
--   let !numRFreq = L.length rFreqs
--       !numThetaFreq = L.length thetaFreqs
--       idx =
--         [ (r, theta)
--         | r <- [0 .. (numRFreq - 1)]
--         , theta <- [0 .. (numThetaFreq - 1)]
--         ]
--   case field of
--     Source ->
--       fmap (DFTArray rows cols thetaFreqs rFreqs) .
--       dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
--       L.map
--         (\(r, theta) ->
--            VS.fromList . A.toList $
--            (runNWith ptx (convolveKernel coefficients harmonics))
--              (A.fromList A.Z [theta])
--              (A.fromList A.Z [r])
--              (A.fromList
--                 (A.Z A.:. numRFreq A.:. numThetaFreq A.:. cols A.:. rows) .
--               VS.toList . VS.concat $
--               dftVecs)) $
--       idx
--     -- Sink ->
--     --   fmap (DFTArray rows cols thetaFreqs rFreqs) .
--     --   dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
--     --   L.map
--     --     (\(r, theta) ->
--     --        VS.fromList .
--     --        A.toList .
--     --        (runNWith ptx (convolveKernel coefficientsSink harmonics))
--     --          (A.fromList (Z :. (1 :: Int)) theta)
--     --          (A.fromList (Z :. (1 :: Int)) r)
--     --          (A.fromList (Z :. numRFreq :. numThetaFreq :. cols :. rows) .
--     --           VS.toList . VS.concat $
--     --           dftVecs)) $
--     --   idx

{-# INLINE convolveSingle #-}
convolveSingle ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> Double
  -> R.Array U DIM2 (Complex Double)
  -> VU.Vector (Complex Double)
convolveSingle !field !coefficients !harmonicsArray !thetaFreqs !rFreqs !x !y !input =
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
      rowCenter = div rows 2
      colCenter = div cols 2
      (Z :. _ :. _ :. numRFreq :. numThetaFreq) = extent coefficients
      -- harmonics =
      --   R.traverse2 thetaFreqs rFreqs (\_ _ -> extent coefficients) $ \fThetaFreq fRFreq (Z :. r :. theta :. rho :. phi) ->
      --     (harmonicsArray IA.!
      --      ( round (fRFreq (Z :. rho) - fRFreq (Z :. r))
      --      , round $ (fThetaFreq (Z :. phi) - (fThetaFreq (Z :. theta))))) R.!
      --     (Z :. (x + colCenter) :. (y + rowCenter))
      harmonics =
        R.traverse2 thetaFreqs rFreqs (\_ _ -> extent coefficients) $ \fPhiFreq fRhoFreq (Z :. r :. theta :. rho :. phi) ->
          let !tf = fPhiFreq (Z :. phi) - fPhiFreq (Z :. theta)
              !rf = fRhoFreq (Z :. rho) - fRhoFreq (Z :. r)
          in (x :+ y) ** (tf :+ 0) *
             ((x ^ 2 + y ^ 2) :+ 0) **
             (((-tf - 1) :+ 2 * pi * rf / (log $ (sqrt 2) * 512)) / 2)
      product =
        R.traverse2 (R.zipWith (*) coefficients harmonics) input const $ \fCoefficient fInput idx@(Z :. _ :. _ :. rho :. phi) ->
          fCoefficient idx * fInput (Z :. rho :. phi)
      arr =
        case field of
          Source -> product
          Sink ->
            error
              "convolveSingle: You should not use it to compute a sink field."
  in if x == 0 && y == 0
       then VU.replicate (numRFreq * numThetaFreq) 0
       else toUnboxed . sumS . sumS $ arr

-- {-# INLINE convolveSparse #-}
convolveSparse ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> [R.Array U DIM2 (Complex Double)]
  -> [R.Array U DIM2 (Complex Double)]
convolveSparse !field !coefficients !harmonicsArray !thetaFreqs !rFreqs !cutoff !xs !inputs =
  L.map (fromUnboxed (extent . L.head $ inputs)) .
  parMap
    rdeepseq
    (\(x, y) ->
       L.foldl1' (VU.zipWith (+)) .
       L.zipWith
         (\(x', y') input ->
            let !x'' = x - x'
                !y'' = y - y'
                (!a, !b) =
                  if (x'' ^ 2 + y'' ^ 2) <= cutoff ^ 2
                    then (x'', y'')
                    else (0, 0)
            in convolveSingle
                 field
                 coefficients
                 harmonicsArray
                 thetaFreqs
                 rFreqs
                 a
                 b
                 input)
         xs $
       inputs) $
  xs

{-# INLINE convolveSingle' #-}
convolveSingle' ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> Double
  -> R.Array U DIM1 (Complex Double)
  -> VU.Vector (Complex Double)
convolveSingle' !field !coefficients !harmonicsArray !phiFreqs !rhoFreqs !x !y !input =
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
      rowCenter = div rows 2
      colCenter = div cols 2
      (Z :. _ :. _ :. numRhoFreq :. numPhiFreq) = extent coefficients
      -- harmonics =
      --   R.traverse2 phiFreqs rhoFreqs (\_ _ -> extent coefficients) $ \fPhiFreq fRhoFreq (Z :. _ :. theta :. rho :. phi) ->
      --     (harmonicsArray IA.!
      --      ( round (fRhoFreq (Z :. rho))
      --      , round (fPhiFreq (Z :. phi) - (fPhiFreq (Z :. theta))))) R.!
      --     (Z :. (x + colCenter) :. (y + rowCenter))
      harmonics =
        R.traverse2 phiFreqs rhoFreqs (\_ _ -> extent coefficients) $ \fPhiFreq fRhoFreq (Z :. _ :. theta :. rho :. phi) ->
          let !tf = fPhiFreq (Z :. phi) - fPhiFreq (Z :. theta)
              !rf = fRhoFreq (Z :. rho)
          in (x :+ y) ** (tf :+ 0) *
             ((x ^ 2 + y ^ 2) :+ 0) ** (((-tf - 1) :+ rf) / 2)
      product =
        R.traverse2 (R.zipWith (*) coefficients harmonics) input const $ \fCoefficient fInput idx@(Z :. _ :. _ :. _ :. phi) ->
          fCoefficient idx * fInput (Z :. phi)
      arr =
        case field of
          Source -> product
          Sink ->
            R.traverse2 product phiFreqs const $ \fProd fPhiFreq idx@(Z :. _ :. theta :. _ :. phi) ->
              fProd idx *
              (cis $ (-(fPhiFreq (Z :. phi) - fPhiFreq (Z :. theta))) * pi)
  in if x == 0 && y == 0
       then VU.replicate (numRhoFreq * numPhiFreq) 0
       else toUnboxed . sumS . sumS $ arr

convolveSparse' ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> [R.Array U DIM1 (Complex Double)]
  -> [R.Array U DIM1 (Complex Double)]
convolveSparse' !field !coefficients !harmonicsArray !phiFreqs !rhoFreqs !cutoff !xs !inputs =
  L.map (fromUnboxed (extent . L.head $ inputs)) .
  parMap
    rdeepseq
    (\(x, y) ->
       L.foldl1' (VU.zipWith (+)) .
       L.zipWith
         (\(x', y') input ->
            let !x'' = x - x'
                !y'' = y - y'
                (!a, !b) =
                  if (x'' ^ 2 + y'' ^ 2) <= cutoff ^ 2
                    then (x'', y'')
                    else (0, 0)
            in convolveSingle'
                 field
                 coefficients
                 harmonicsArray
                 phiFreqs
                 rhoFreqs
                 a
                 b
                 input)
         xs $
       inputs) $
  xs

{-# INLINE gaussianFilter2D #-}
gaussianFilter2D :: DFTPlan -> Int -> Double -> IO (VS.Vector (Complex Double))
gaussianFilter2D plan numPoint stdG =
  dftExecute plan (DFTPlanID DFT1DG [numPoint, numPoint] [0, 1]) .
  VU.convert .
  toUnboxed . computeS . makeFilter2D . fromFunction (Z :. numPoint :. numPoint) $ \(Z :. i :. j) ->
    let x = fromIntegral $ i - div numPoint 2
        y = fromIntegral $ j - div numPoint 2
    in (exp (-(x ^ 2 + y ^ 2) / (2 * stdG ^ 2))) :+ 0

{-# INLINE convoluvePinhweelBasis #-}
convoluvePinhweelBasis ::
     R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> DFTArray
convoluvePinhweelBasis coefficients harmonicsArray arr@(DFTArray xFreq yFreq thetaFreqs rFreqs vecs) =
 let initVec = VS.replicate (VS.length . L.head $ vecs) 0
     idxs = (,) <$> (L.zip [0 ..] rFreqs) <*> (L.zip [0 ..] thetaFreqs)
  in DFTArray xFreq yFreq thetaFreqs rFreqs .
     parMap
      rdeepseq
       (\((r, rFreq), (theta, thetaFreq)) ->
          L.foldl'
            (\vec (((rho, rhoFreq), (phi, phiFreq)), inputVec) ->
               VS.zipWith (+) vec .
               VS.map (* (coefficients R.! (Z :. r :. theta :. rho :. phi))) .
               VS.zipWith
                 (*)
                 (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq) $
               inputVec)
            initVec .
          L.zip idxs $
          vecs) $
     idxs

{-# INLINE convoluvePinhweelBasis' #-}
convoluvePinhweelBasis' ::
     R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> DFTArray
convoluvePinhweelBasis' coefficients harmonicsArray arr@(DFTArray xFreq yFreq thetaFreqs rFreqs vecs) =
  let idxTheta = L.zip [0 ..] thetaFreqs
      initVec = VS.replicate (VS.length . L.head $ vecs) 0
  in DFTArray xFreq yFreq thetaFreqs rFreqs .
     parMap
       rdeepseq
       (\(!theta, !thetaFreq) ->
          L.foldl'
            (\vec1 (rho, rhoFreq) ->
               L.foldl'
                 (\(!vec2) ((!phi, !phiFreq), inputVec) ->
                    VS.zipWith
                      (+)
                      vec2
                      (VS.map
                         (* (coefficients R.!
                             (Z :. (0 :: Int) :. theta :. rho :. phi))) .
                       VS.zipWith
                         (*)
                         (getHarmonics
                            harmonicsArray
                            phiFreq
                            rhoFreq
                            thetaFreq
                            0) $
                       inputVec))
                 vec1 .
               L.zip idxTheta $
               vecs)
            initVec .
          L.zip [0 ..] $
          rFreqs) $
     idxTheta

{-# INLINE convolveFull'' #-}
convolveFull'' ::
     R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> DFTArray
convolveFull'' !coefficients !harmonicsArray !arr@(DFTArray r2Freq _ thetaFreqs rFreqs vecs) =
  let idxTheta = L.zip [0 ..] thetaFreqs
      !initVec = VS.replicate (VS.length . L.head $ vecs) 0
  in DFTArray r2Freq r2Freq thetaFreqs rFreqs .
     parMap
       rdeepseq
       (\(!theta, !thetaFreq) ->
          L.foldl'
            (\vec1 (rho, rhoFreq) ->
               L.foldl'
                 (\(!vec2) ((!phi, !phiFreq), inputVec) ->
                    VS.zipWith
                      (+)
                      vec2
                      (VS.map
                         (* (coefficients R.!
                             (Z :. (0 :: Int) :. theta :. rho :. phi)))
                         (getHarmonics
                            harmonicsArray
                            phiFreq
                            rhoFreq
                            thetaFreq
                            0)))
                 vec1 .
               L.zip idxTheta $
               vecs)
            initVec .
          L.zip [0 ..] $
          rFreqs) $
     idxTheta

{-# INLINE convoluvePinhweelBasisTest #-}
convoluvePinhweelBasisTest ::
     VS.Vector (Complex Double)
  -> [VS.Vector (Complex Double)]
  -> [VS.Vector (Complex Double)]
  -> DFTArray
  -> IO DFTArray
convoluvePinhweelBasisTest coefficients harmonicsArray1 harmonicsArray2 arr@(DFTArray xFreq yFreq thetaFreqs rFreqs vecs) = do
  let vecs1 = parZipWith rdeepseq (VS.zipWith (*)) vecs harmonicsArray1
      m = L.length thetaFreqs * L.length rFreqs
      n = xFreq * yFreq
  vec2 <- gemmBLAS m n m coefficients . VS.concat $ vecs1
  let arr =
        fromUnboxed
          (Z :. (L.length rFreqs) :. (L.length thetaFreqs) :. xFreq :. yFreq) .
        VS.convert $
        vec2
      vecs2 =
        parMap
          rdeepseq
          (\(i, j) ->
             VS.convert . toUnboxed . computeS . R.slice arr $
             (Z :. i :. j :. All :. All))
          [ (i, j)
          | i <- [0 .. L.length rFreqs - 1]
          , j <- [0 .. L.length thetaFreqs - 1]
          ] :: [VS.Vector (Complex Double)]
      vecs3 = parZipWith rdeepseq (VS.zipWith (*)) vecs2 harmonicsArray2
  return . DFTArray xFreq yFreq thetaFreqs rFreqs $ vecs3

{-# INLINE convoluvePinhweelBasisTest' #-}
convoluvePinhweelBasisTest' ::
     R.Array U DIM4 (Complex Double)
  -> [VS.Vector (Complex Double)]
  -> [VS.Vector (Complex Double)]
  -> DFTArray
  -> IO DFTArray
convoluvePinhweelBasisTest' coefficients harmonicsArray1 harmonicsArray2 arr@(DFTArray xFreq yFreq thetaFreqs rFreqs vecs) = do
  let (Z :. numRFreq :. numThetaFreq :. numRhoFreq :. numPhiFreq) =
        extent coefficients
      coefVec = VU.convert . toUnboxed $ coefficients
      vecs1 =
        parZipWith rdeepseq (VS.zipWith (*)) harmonicsArray1 .
        L.concat . L.replicate numRhoFreq $
        vecs
      m = numRFreq * numThetaFreq
      n = xFreq * yFreq
      k = numRhoFreq * numPhiFreq
  print (m, n, k)
  vec2 <- gemmBLAS m n k coefVec . VS.concat $ vecs1
  let arr =
        fromUnboxed (Z :. numRFreq :. numThetaFreq :. xFreq :. yFreq) .
        VS.convert $
        vec2
      vecs2 =
        parMap
          rdeepseq
          (\(i, j) ->
             VS.convert . toUnboxed . computeS . R.slice arr $
             (Z :. i :. j :. All :. All))
          [(i, j) | i <- [0 .. numRFreq - 1], j <- [0 .. numThetaFreq - 1]] :: [VS.Vector (Complex Double)]
      vecs3 = parZipWith rdeepseq (VS.zipWith (*)) vecs2 harmonicsArray2
  return . DFTArray xFreq yFreq thetaFreqs rFreqs $ vecs3


module DFT (dft, idft) where
import FFT (fft)
import Data.Complex (Complex)
import Data.List (reverse, map, take)

dft :: [Complex Float] -> [Complex Float]
dft [] = []
dft xs = fft $ preprocess xs

idft :: [Complex Float] -> [Complex Float]
idft [] = []
idft xs = (invertFFT . dft) xs

invertFFT :: [Complex Float] -> [Complex Float]
invertFFT [] = []
invertFFT coeffs@(x:xs) = [ x' / n | x' <- (x : reverse xs) ]
  where n = (fromIntegral . length) coeffs

preprocess :: [Complex Float] -> [Complex Float]
preprocess [] = []
preprocess xs = take targetLength (xs ++ (repeat 0))
  where targetLength = 2 ^ ceiling (log ((fromIntegral . length) xs) / log 2)


{-# LANGUAGE AllowAmbiguousTypes, DataKinds, RankNTypes, TypeFamilies,
             TypeOperators #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | Shaped tensor-based implementation of Recurrent Neural Network
-- for classification of MNIST digits. Sports 2 hidden layers.
module HordeAd.Tool.MnistRnnShaped where

import Prelude

import           Control.Monad (foldM)
import qualified Data.Array.DynamicS as OT
import           Data.Array.Internal (valueOf)
import qualified Data.Array.Shape
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import           Data.Proxy (Proxy)
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat)
import           Numeric.LinearAlgebra (Vector)
import qualified Numeric.LinearAlgebra as HM

-- until stylish-haskell accepts NoStarIsType
import qualified GHC.TypeLits

import HordeAd.Core.DualNumber
import HordeAd.Core.Engine
import HordeAd.Core.PairOfVectors (DualNumberVariables, varS)
import HordeAd.Tool.MnistData

zeroStateS
  :: (IsScalar d r, OS.Shape sh)
  => (DualNumber d (OS.Array sh r)  -- state
      -> a)
  -> a
zeroStateS f = f (konstS 0)

unrollLastS :: forall state c w m d r k rest.
               (DualMonad d r m, KnownNat k, OS.Shape rest)
            => (state -> OS.Array rest r -> w -> m (c, state))
            -> (state -> OS.Array (k : rest) r -> w -> m (c, state))
unrollLastS f s0 xs w = do
  let g :: (c, state) -> OS.Array rest r -> m (c, state)
      g (_, s) x = f s x w
  foldM g (undefined, s0) $ OSB.toList $ OS.unravel xs

type LayerWeigthsRNN in_width out_width d r =
  ( DualNumber d (OS.Array '[out_width, in_width] r)   -- input weight
  , DualNumber d (OS.Array '[out_width, out_width] r)  -- state weight
  , DualNumber d (OS.Array '[out_width] r) )           -- bias

rnnMnistLayerS
  :: forall in_width out_width batch_size d r m.
     ( DualMonad d r m
     , KnownNat in_width, KnownNat out_width, KnownNat batch_size )
  => DualNumber d (OS.Array '[out_width, batch_size] r)  -- in state
  -> DualNumber d (OS.Array '[in_width, batch_size] r)  -- in
  -> LayerWeigthsRNN in_width out_width d r
  -> m ( DualNumber d (OS.Array '[out_width, batch_size] r)  -- out
       , DualNumber d (OS.Array '[out_width, batch_size] r) )  -- out state
rnnMnistLayerS s x (wX, wS, b) = do
  let y = wX <>$ x + wS <>$ s + asColumnS b
  yTanh <- returnLet $ tanh y
  return (yTanh, yTanh)

rnnMnistTwoS
  :: forall out_width batch_size d r m.
     (DualMonad d r m, KnownNat out_width, KnownNat batch_size)
  => DualNumber d (OS.Array '[2 GHC.TypeLits.* out_width, batch_size] r)
       -- initial state
  -> OS.Array '[SizeMnistWidth, batch_size] r
  -> ( LayerWeigthsRNN SizeMnistWidth out_width d r
     , LayerWeigthsRNN out_width out_width d r )
  -> m ( DualNumber d (OS.Array '[out_width, batch_size] r)
       , DualNumber d (OS.Array '[2 GHC.TypeLits.* out_width, batch_size] r) )
           -- final state
rnnMnistTwoS s x ((wX, wS, b), (wX2, wS2, b2)) = do
  let s1 = sliceS @0 @out_width s
      s2 = sliceS @out_width @out_width s
  (vec1, s1') <- rnnMnistLayerS s1 (constant x) (wX, wS, b)
  (vec2, s2') <- rnnMnistLayerS s2 vec1 (wX2, wS2, b2)
  s3 <- returnLet $ appendS s1' s2'
  return (vec2, s3)

rnnMnistZeroS
  :: forall out_width batch_size d r m.
     (DualMonad d r m, KnownNat out_width, KnownNat batch_size)
  => OS.Array '[SizeMnistHeight, SizeMnistWidth, batch_size] r
  -- All below is the type of all paramters of this nn. The same is reflected
  -- in the length function below and read from variables further down.
  -> ( LayerWeigthsRNN SizeMnistWidth out_width d r
     , LayerWeigthsRNN out_width out_width d r )
  -> DualNumber d (OS.Array '[SizeMnistLabel, out_width] r)
  -> DualNumber d (OS.Array '[SizeMnistLabel] r)
  -> m (DualNumber d (OS.Array '[SizeMnistLabel, batch_size] r))
rnnMnistZeroS xs ((wX, wS, b), (wX2, wS2, b2)) w3 b3 = do
  (out, _s) <- zeroStateS (unrollLastS rnnMnistTwoS) xs
                          ((wX, wS, b), (wX2, wS2, b2))
  returnLet $ w3 <>$ out + asColumnS b3

rnnMnistLenS
  :: forall out_width. KnownNat out_width
  => Proxy out_width -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
rnnMnistLenS _ =
  ( 0
  , []
  , []
  , [ Data.Array.Shape.shapeT @'[out_width, SizeMnistWidth]
    , Data.Array.Shape.shapeT @'[out_width, out_width]
    , Data.Array.Shape.shapeT @'[out_width]
    , Data.Array.Shape.shapeT @'[out_width, out_width]
    , Data.Array.Shape.shapeT @'[out_width, out_width]
    , Data.Array.Shape.shapeT @'[out_width]
    , Data.Array.Shape.shapeT @'[SizeMnistLabel, out_width]
    , Data.Array.Shape.shapeT @'[SizeMnistLabel]
    ]
  )

rnnMnistS
  :: forall out_width batch_size d r m.
     (DualMonad d r m, KnownNat out_width, KnownNat batch_size)
  => OS.Array '[SizeMnistHeight, SizeMnistWidth, batch_size] r
  -> DualNumberVariables d r
  -> m (DualNumber d (OS.Array '[SizeMnistLabel, batch_size] r))
rnnMnistS xs variables = do
  let wX = varS variables 0
      wS = varS variables 1
      b = varS variables 2
      wX2 = varS variables 3
      wS2 = varS variables 4
      b2 = varS variables 5
      w3 = varS variables 6
      b3 = varS variables 7
  rnnMnistZeroS @out_width xs ((wX, wS, b), (wX2, wS2, b2)) w3 b3

rnnMnistLossFusedS
  :: forall out_width batch_size d r m.
     (DualMonad d r m, KnownNat out_width, KnownNat batch_size)
  => Proxy out_width
  -> MnistDataBatchS batch_size r
  -> DualNumberVariables d r
  -> m (DualNumber d r)
rnnMnistLossFusedS _ (glyphS, labelS) variables = do
  let xs = OS.transpose @'[2, 1, 0] glyphS
  result <- rnnMnistS @out_width xs variables
  let targets2 = HM.tr $ HM.reshape (valueOf @SizeMnistLabel)
                       $ OS.toVector labelS
  vec <- lossSoftMaxCrossEntropyL targets2 (fromS2 result)
  returnLet $ scale (recip $ fromIntegral (valueOf @batch_size :: Int))
            $ sumElements0 vec

rnnMnistTestS
  :: forall out_width batch_size r.
     (IsScalar 'DModeGradient r, KnownNat out_width, KnownNat batch_size)
  => Proxy r -> Proxy out_width
  -> MnistDataBatchS batch_size r
  -> Domains r
  -> r
rnnMnistTestS _ _ (glyphS, labelS) parameters =
  let xs = OS.transpose @'[2, 1, 0] glyphS
      outputS = primalValue (rnnMnistS @out_width xs) parameters
      outputs = map OS.toVector $ OSB.toList $ OS.unravel
                $ OS.transpose @'[1, 0] $ outputS
      labels = map OS.toVector $ OSB.toList $ OS.unravel labelS
      matchesLabels :: Vector r -> Vector r -> Int
      matchesLabels output label | V.maxIndex output == V.maxIndex label = 1
                                 | otherwise = 0
  in fromIntegral (sum (zipWith matchesLabels outputs labels))
     / fromIntegral (valueOf @batch_size :: Int)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Lib (benchmark) where

import Test.QuickCheck
import Criterion.Main
import Data.Proxy
import GHC.TypeLits
import Control.DeepSeq
import qualified Numeric.LinearAlgebra.Data as HM
import qualified Numeric.LinearAlgebra.HMatrix as HM
import qualified Data.Matrix as DM
import qualified LAoP.Matrix.Type as TM
import qualified Linear as LM
import qualified Linear.V as LM
import qualified Linear.Matrix as LM
import qualified Data.Vector as V
import LAoP.Utils hiding ((.))

randomMatrixHM :: Int -> Int -> Gen (HM.Matrix Double)
randomMatrixHM cols rows = do
  l <- vectorOf (cols * rows) arbitrary
  let lr = buildList l cols
      m  = HM.fromLists lr
  return m

randomMatrixDM :: Int -> Int -> Gen (DM.Matrix Double)
randomMatrixDM cols rows = do
  l <- vectorOf (cols * rows) arbitrary
  let lr = buildList l cols
      m  = DM.fromLists lr
  return m

randomMatrixTM :: forall a b . (TM.CountableDimensions a b, TM.FromListsN Double a b) => Gen (TM.Matrix Double a b)
randomMatrixTM = do
  let cols = fromInteger (natVal (Proxy :: Proxy (TM.Count a)))
      rows = fromInteger (natVal (Proxy :: Proxy (TM.Count b)))
  l <- vectorOf (cols * rows) arbitrary
  let lr = buildList l cols
      m  = TM.fromLists lr
  return m

randomMatrixLM :: forall a b . (KnownNat a, KnownNat b, LM.Dim a, LM.Dim b) => Gen (LM.V a (LM.V b Double))
randomMatrixLM = do
  let cols = fromInteger (natVal (Proxy :: Proxy a))
      rows = fromInteger (natVal (Proxy :: Proxy b))
  l <- vectorOf (cols * rows) (arbitrary :: Gen Double)
  let lr = buildList l cols
      v  = map (fromJust . LM.fromVector @b . V.fromList) lr
      (Just m)  = LM.fromVector @a (V.fromList v)
  return m

fromJust (Just a) = a

buildList [] _ = []
buildList l r  = take r l : buildList (drop r l) r

--
setupEnv1 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 10 10))
  hm12 <- generate (resize 1 (randomMatrixHM 10 10))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 10 10))
  dm12 <- generate (resize 1 (randomMatrixDM 10 10))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @10 @10))
  lm12 <- generate (resize 1 (randomMatrixLM @10 @10))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 10) @(Natural 0 10)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 10) @(Natural 0 10)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

setupEnv2 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 25 25))
  hm12 <- generate (resize 1 (randomMatrixHM 25 25))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 25 25))
  dm12 <- generate (resize 1 (randomMatrixDM 25 25))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @25 @25))
  lm12 <- generate (resize 1 (randomMatrixLM @25 @25))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 25) @(Natural 0 25)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 25) @(Natural 0 25)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

setupEnv3 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 50 50))
  hm12 <- generate (resize 1 (randomMatrixHM 50 50))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 50 50))
  dm12 <- generate (resize 1 (randomMatrixDM 50 50))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @50 @50))
  lm12 <- generate (resize 1 (randomMatrixLM @50 @50))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 50) @(Natural 0 50)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 50) @(Natural 0 50)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

setupEnv4 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 75 75))
  hm12 <- generate (resize 1 (randomMatrixHM 75 75))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 75 75))
  dm12 <- generate (resize 1 (randomMatrixDM 75 75))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @75 @75))
  lm12 <- generate (resize 1 (randomMatrixLM @75 @75))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 75) @(Natural 0 75)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 75) @(Natural 0 75)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

setupEnv5 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 100 100))
  hm12 <- generate (resize 1 (randomMatrixHM 100 100))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 100 100))
  dm12 <- generate (resize 1 (randomMatrixDM 100 100))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @100 @100))
  lm12 <- generate (resize 1 (randomMatrixLM @100 @100))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 100) @(Natural 0 100)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 100) @(Natural 0 100)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

setupEnv6 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 200 200))
  hm12 <- generate (resize 1 (randomMatrixHM 200 200))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 200 200))
  dm12 <- generate (resize 1 (randomMatrixDM 200 200))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @200 @200))
  lm12 <- generate (resize 1 (randomMatrixLM @200 @200))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 200) @(Natural 0 200)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 200) @(Natural 0 200)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

setupEnv7 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 400 400))
  hm12 <- generate (resize 1 (randomMatrixHM 400 400))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 400 400))
  dm12 <- generate (resize 1 (randomMatrixDM 400 400))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @400 @400))
  lm12 <- generate (resize 1 (randomMatrixLM @400 @400))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 400) @(Natural 0 400)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 400) @(Natural 0 400)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

setupEnv8 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 800 800))
  hm12 <- generate (resize 1 (randomMatrixHM 800 800))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 800 800))
  dm12 <- generate (resize 1 (randomMatrixDM 800 800))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @800 @800))
  lm12 <- generate (resize 1 (randomMatrixLM @800 @800))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 800) @(Natural 0 800)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 800) @(Natural 0 800)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

setupEnv9 = do
  -- HMatrix
  hm11 <- generate (resize 1 (randomMatrixHM 1600 1600))
  hm12 <- generate (resize 1 (randomMatrixHM 1600 1600))
  -- Matrix
  dm11 <- generate (resize 1 (randomMatrixDM 1600 1600))
  dm12 <- generate (resize 1 (randomMatrixDM 1600 1600))
  -- Linear
  lm11 <- generate (resize 1 (randomMatrixLM @1600 @1600))
  lm12 <- generate (resize 1 (randomMatrixLM @1600 @1600))
  -- Laop
  tm11 <- generate (resize 1 (randomMatrixTM @(Natural 0 1600) @(Natural 0 1600)))
  tm12 <- generate (resize 1 (randomMatrixTM @(Natural 0 1600) @(Natural 0 1600)))
  return (
   (hm11, hm12)
   ,(dm11, dm12)
   ,(lm11, lm12)
   ,(tm11, tm12)
         )

benchmark :: IO ()
benchmark = do
  print "Starting benchmarks..."
  defaultMain [
   env setupEnv1 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 1" [
   bgroup "10x10" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ],
   env setupEnv2 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 2" [
   bgroup "25x25" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ],
   env setupEnv3 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 3" [
   bgroup "50x50" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ],
   env setupEnv4 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 4" [
   bgroup "75x75" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ],
   env setupEnv5 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 5" [
   bgroup "100x100" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ],
   env setupEnv6 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 6" [
   bgroup "200x200" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ],
   env setupEnv7 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 7" [
   bgroup "400x400" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ],
   env setupEnv8 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 8" [
   bgroup "800x800" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ],
   env setupEnv9 $ \ ~((hm11, hm12), (dm11, dm12), (lm11, lm12), (tm11, tm12)) -> bgroup "Matrix Composition 9" [
   bgroup "1600x1600" [
     bench "hmatrix" $ nf (HM.mul hm11) hm12
   , bench "matrix" $ nf (DM.multStd2 dm11) dm12
   , bench "linear" $ nf (LM.!*! lm11) lm12
   , bench "laop" $ nf (TM.comp tm11) tm12
   ] ]
              ]


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
-- | Main entry point for executable, mainly for tests
module Main (main) where

import AI.Network.RNN
import AI.GeneticAlgorithm.Simple
import Control.Monad
import Control.Monad.Random hiding (fromList)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.IORef
import System.Directory
--import Debug.Trace
import System.Environment

import Data.Binary

import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL

import Numeric.LinearAlgebra.HMatrix

--main::IO()
--main = do
--    let m :: [Expr Double] = map (\n->Var ("m" ++ show n)) [1..400]
--        v :: [Expr Double] = map (\n->Var ("v" ++ show n)) [1..4]
--        mycost v m = sum $ listMProd m v
--        mygrad = grad (mycost (map auto v)) m
--        --mygrads = gradientDescent (mycost (map auto v))  m
--    print $ mygrad

main :: IO()
main = do
    --txt <- liftM (T.take 100) $ T.readFile "data/tinyshakespeare.txt"
    let txt= "hello world!"
        trainData = train txt
        fn = "output/lstm.one.helloworld"
        generateLength = min 50 (T.length txt)
    print "Text:"
    print txt
    args <- getArgs
    when (length args<2) $ error "rnn gradient <maxgen> | rmsprop <maxgen> | generic <maxgen>"
    let fn2=fn++"."++head args
    ex <- return False -- doesFileExist fn2
    rnn<-case args of
        ("gradient":mg:_) -> do
            let maxGen = read mg
            rd <- readEx fn2 ex trainData
            gradient fn2 rd generateLength maxGen
        ("rmsprop":mg:_) -> do
            let maxGen = read mg
            rd <- readEx fn2 ex trainData
            rmsprop fn2 rd generateLength maxGen
        ("genetic":mg:_) -> do
            let maxGen = read mg
            b <- if ex
               then do
                 rnn1::LSTMNetwork <- fromVector (tdRecSize trainData) <$> decodeFile fn2
                 return $ buildExisting rnn1 trainData
               else return $ buildTD trainData
--            let b = buildTDs trainData
            genetic fn2 b generateLength maxGen
        _ -> error "rnn gradient <maxgen> | rmsprop <maxgen> | generic <maxgen>"
    encodeFile fn2 $ toVector rnn
    print rnn
    where
      readEx fn ex trainData =
          if ex
                then do
                 r::LSTMNetwork <- fromVector (tdRecSize trainData) <$> decodeFile fn
                 return $ RNNData r trainData (fromIntegral $ length $ tdInputs trainData)
                else
                  evalRandIO $ buildTD trainData
      (train, gener) = (textToTrainDataS,generateS)
      genStep maxGen = max 1 $ maxGen `div` 10
      gradient fn (RNNData r td _) generateLength maxGen = learnGradientDescent r td $ test fn generateLength maxGen
      rmsprop fn (RNNData r td _) generateLength maxGen = learnRMSProp r td $ test fn generateLength maxGen
      genetic fn r generateLength maxGen = do
        fitnessList <- newIORef []
        (RNNData rnn _ _) <- runGAIO 64 0.2 r $ stopf2 fn generateLength maxGen fitnessList
        return rnn
      test fn generateLength maxGen rnn td gen= do
        when (mod gen (genStep maxGen) == 0) $ do
            let c= cost rnn td
            print (show gen ++ ":" ++ show c)
            t <- evalRandIO $ gener (tdData td) generateLength (size $ head $ tdInputs td) rnn
            print t
            let fn2 = fn ++ "." ++ (show gen)
            let bs= BSL.toStrict $ encode $ toVector rnn
            BS.writeFile fn2 bs
        return $ gen < maxGen
      stopf2 fn generateLength maxGen fs rd@((RNNData rnn td _),_) gen = do
        _ <- test fn generateLength maxGen rnn td gen
        stopf fs maxGen rd gen

-- | Build LSTM network
buildTD ::(Monad m,RandomGen g) =>  TrainData a -> RandT g m (RNNData LSTMNetwork a)
buildTD td = buildNetworkData td (tdRecSize td) lstmFullSize

buildTDIO ::(Monad m,RandomGen g) =>  TrainData a -> RandT g m (RNNData LSTMIO a)
buildTDIO td =
    let layers = 2
        sz   = 40
     --   tds    = TrainData (tdInputs td) (tdOutputs td) (tdRecSize td,sz,layers,tdRecSize td) (tdData td)
     in buildNetworkData td (tdRecSize td,sz,layers,tdRecSize td) lstmioFullSize

-- | Build LSTM network
buildTDs ::(Monad m,RandomGen g) =>  TrainData a -> RandT g m (RNNData [LSTMNetwork] a)
buildTDs td =
    let layers = 2
     --   tds = TrainData (tdInputs td) (tdOutputs td) (replicate layers $ tdRecSize td) (tdData td)
    in  buildNetworkData td (replicate layers $ tdRecSize td)(sum . map lstmFullSize)

-- | Build from existing data for genetic algorithm, to restart from where we were
buildExisting :: (Monad m,RandomGen g) =>  LSTMNetwork -> TrainData a -> RandT g m (RNNData LSTMNetwork a)
buildExisting rnn1 td@(TrainData is _ _ _) = do
  -- g <- getSplit
  --let rnn2 = evalRand (mutateNetwork rnn1) g
  return $ RNNData rnn1 td (fromIntegral $ length is)

-- | Build from existing data for genetic algorithm, to restart from where we were
buildExistings :: (Monad m,RandomGen g) =>  [LSTMNetwork] -> TrainData a -> RandT g m (RNNData [LSTMNetwork] a)
buildExistings rnns td@(TrainData is _ _ _) = do
  -- g <- getSplit
  --let tds = TrainData (tdInputs td) (tdOutputs td) (replicate (length rnns) $ tdRecSize td) (tdData td)
  --let rnn2 = evalRand (mutateNetwork rnn1) g
  return $ RNNData rnns td (fromIntegral $ length is)

-- | Build from existing data for genetic algorithm, to restart from where we were
buildExistingIO :: (Monad m,RandomGen g) =>  LSTMIO -> TrainData a -> RandT g m (RNNData LSTMIO a)
buildExistingIO rnn1 td@(TrainData is _ _ _) = do
  -- g <- getSplit
  --let tds = TrainData (tdInputs td) (tdOutputs td) (rnnsize rnn1) (tdData td)
  --let rnn2 = evalRand (mutateNetwork rnn1) g
  return $ RNNData rnn1 td (fromIntegral $ length is)

-- DRAGONS
--
--main :: IO()
--main = do
--   -- txt <- T.readFile "data/tinyshakespeare.txt"
--    let txt= "hello world!"
--        (is,os,m) = textToTrainData txt
--        sz = DM.size m
--        fn = "lstm.learn1.helloworld"
--    print sz
--    print $ length is
--    print $ length os
--    ex <- doesFileExist fn
--    (RNNData r _ _ _) <- if ex
--        then do
--          r <- read <$> readFile fn
--          return $ RNNData r [] [] 0
--        else
--          evalRandIO $ build sz is os
--    let ls = toList $ toVector r
--    --print ls
--    --print $ toList $ chead ocs
--    --print ls
--    --print is
--    --print $ lstmList sz ls (toList $ head is)
--    --print $ mapAccumL (lstmList sz) ls (map toList is)
--    --headJet $ tailJet $ jet $ grads
----    let c= cost sz (fromIntegral $ length is) (map (toList) is) (map (toList) os) ls
----    print c
----    let gs= grad (cost sz (auto $ fromIntegral $ length is) (map (map auto . toList) is) (map (map auto . toList) os)) ls
----    --let gs=jacobian (lstmList (lstmFullSize sz) (map auto ls)) (toList $ head os)
----    -- print gs
----    let ls2 = zipWith (\l g->l-g*0.1) ls gs
----    let c2= cost sz (fromIntegral $ length is) (map (toList) is) (map (toList) os) ls2
----    print c2
----    where
----        gFun :: [AD s (Sparse Double)]
----                   -> AD s (Sparse Double)
----        gFun =  map sigmoid
--    let (l,lis,los)=((fromIntegral $ length is),(map toList is),(map toList os))
--    -- lsf <- learn sz ls l lis los 0
--    let myl= length ls
--    print myl
--    -- let adCost = cost sz (auto l) (map (map auto) lis) (map (map auto) los)
--    -- lsf <- learn2 sz ls (replicate myl 0) (replicate myl 0) (replicate myl 0) l lis los 0 (m,(min 50 (T.length txt)))
--    lsf <- learn sz ls l lis los 0 (m,(min 50 (T.length txt)))
--    --print lsf
--    let rnn::LSTMNetwork = fromVector sz (fromList lsf)
--    writeFile fn $ show rnn
--    --print rnn
--    -- print $ length is
--    let s = size $ head is
----    -- print s
----    let tl= T.length txt
----    -- print tl
----    -- print m
----    let c= cost sz l lis los ls
----    print c
----    let c2= cost sz l lis los lsf
----    print c2
--    t <- evalRandIO $ generate m (min 50 (T.length txt)) s rnn
--    print t
--    return ()
--    where
--          test sz ls l is os gen (m,gl) = when ((mod gen 500) == 0) $ do
--                let c= cost sz l is os ls
--                print (show gen ++ ":" ++ show c)
--                let rnn::LSTMNetwork = fromVector sz (fromList ls)
--                t <- evalRandIO $ generate m gl (length $ head is) rnn
--                print t
--          learn sz ls l is os gen testInfo = do
--            test sz ls l is os gen testInfo
--            if gen < 10000
--                then do
--                    let ls2 = learnGradientDescent (cost sz (auto l) (map (map auto) is) (map (map auto) os)) ls
--                    --let gs= grad (cost sz (auto l) (map (map auto) is) (map (map auto) os)) ls
--                    --    ls2 = zipWith (\o g->o-g*0.1) ls gs
--                    learn sz ls2 l is os (gen + 1) testInfo
--                else return ls
--          learn2 sz ls rgs rgs2 ugs l is os gen testInfo= do
--            test sz ls l is os gen testInfo
--            if gen < 10000
--                then do
--                    let gs= force $ grad (cost sz (auto l) (map (map auto) is) (map (map auto) os)) ls
--                        rgup = force $ zipWith (\rg g-> 0.95 * rg + 0.05 * g) rgs ls
--                        rg2up = force $ zipWith (\rg2 g-> 0.95 * rg2 + 0.05 * (g ** 2)) rgs2 ls
--                        ugup = force $ zipWith4 (\ud zg rg rg2 -> 0.9 * ud - 1e-4 * zg / sqrt(rg2 - rg ** 2 + 1e-4)) ugs gs rgup rg2up
--                        -- ls2 = zipWith (\o g->o-g*0.1) ls gs
--                        ls2 = force $ zipWith (+) ls ugup
--                    learn2 sz ls2 rgup rg2up ugup l is os (gen + 1) testInfo
--                else return ls
--          learn3 _ ls _ _ _ _ [] [] _ = return ls
--          learn3 sz ls rgs rgs2 ugs l is os gen = do
--            let batchSize = 100
--                (is1,is2) = splitAt batchSize is
--                (os1,os2) = splitAt batchSize os
--            when ((mod gen 100) == 0) $ do
--                let c= cost sz l is1 os1 ls
--                print (show gen ++ ":" ++ show c)
--            --if gen < 1
--            --    then do
--            let gs= force $ grad (cost sz (auto l) (map (map auto) is1) (map (map auto) os1)) ls
--                rgup = force $ zipWith (\rg g-> 0.95 * rg + 0.05 * g) rgs ls
--                rg2up = force $ zipWith (\rg2 g-> 0.95 * rg2 + 0.05 * (g ** 2)) rgs2 ls
--                ugup = force $ zipWith4 (\ud zg rg rg2 -> 0.9 * ud - 1e-4 * zg / sqrt(rg2 - rg ** 2 + 1e-4)) ugs gs rgup rg2up
--                -- ls2 = zipWith (\o g->o-g*0.1) ls gs
--                ls2 = force $ zipWith (+) ls ugup
--            learn3 sz ls2 rgup rg2up ugup l is2 os2 (gen + 1)
--            --    else return ls

--main1 :: IO ()
--main1 = do
----    n<-createRandomNetwork (RNNDimensions 1 2 3 True)
----    let (n2,out)=evalStep n [1]
----        (n3,out1)=evalStep n2 [3]
----        (n4,out2)=evalSteps n [[1],[3]]
----    print $ n3==n4
----    print $ out1 == (last out2)
----  txt <- liftM (T.take 100) $ T.readFile "data/tinyshakespeare.txt"
--  let -- dim =  (RNNDimensions 4 6 4 True)
--      --is  = [[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,1,0],[0,0,0,1]]
--      --os  = [[0,1,0,0],[0,0,1,0],[0,0,1,0],[0,0,0,1],[0,0,0,0]]
--      txt= "hello world!"
--      (is,os,m) = textToTrainData txt
--      sz = DM.size m
--      fn = "out.lstm.h"
--     --  dim = RNNDimensions sz (sz) sz True
--  print sz
--  fitnessList <- newIORef []
----  rnn1::LSTMNetwork <- read <$> readFile fn
--  ex <- doesFileExist fn
--  b <- if ex
--    then do
--        rnn1::LSTMNetwork <- read <$> readFile fn
--        return $ buildExisting rnn1 sz is os
--    else return $ build sz is os
--  r@(RNNData rnn _ _ _) <- runGAIO 64 0.1 b $ stopf fitnessList 100
----  -- (RNNData rnn _ _) <- evalRandIO $ buildNetworkData dim is os
----  print $ fitness r
--  -- print rnn
--  --let real = snd $ evalSteps rnn is
--  --print $ dataToText m real
--  t <- evalRandIO $ generate m (min 50 (T.length txt)) (size $ head is) rnn
--  print t
--  -- writeFile fn $ show rnn

-- build ::(Monad m,RandomGen g) =>  Int -> [Vector Double] -> [Vector Double] -> RandT g m (RNNData RNNetwork RNNDimensions)
-- build sz is os = buildNetworkData (RNNDimensions sz (sz) sz True) totalDataLength is os

--build ::(Monad m,RandomGen g) =>  TrainData a Int  -> RandT g m (RNNData LSTMNetwork Int)
--build sz td = buildNetworkData sz lstmFullSize is os


module Cube where

import Prelude hiding (Left, Right)
import Data.Function (on)
import Numeric.LinearAlgebra (Z, Vector)
import qualified Data.List as List
import qualified Numeric.LinearAlgebra as LA

type SideField = (Int, Int)

fieldsOnSide :: [SideField]
fieldsOnSide = [(x,y)|x <- [0..2], y <- [0..2]]

data Side = Top | Front | Right | Back | Left | Bottom
  deriving (Show, Read, Eq, Enum, Ord, Bounded)

data Color = White | Green | Red | Blue | Orange | Yellow
  deriving (Show, Read, Eq, Enum, Ord, Bounded)

type Cube = [(Side, Int, Int, Color)]

solvedCube :: Cube
solvedCube = do
  (side, color) <- zip [Top ..] [White ..]
  (x, y) <- fieldsOnSide
  [(side, x, y, color)]

computeIndex :: (Side, Int, Int, Color) -> Int
computeIndex (_, x, y, c) = x * 3 + y + fromEnum c * 9

cubeToVector :: Cube -> Vector Z
cubeToVector = LA.fromList . fmap go . List.sortBy (compare `on` project)
  where
    project (a,b,c,_) = (a,b,c)

    go :: (Side, Int, Int, Color) -> Z
    go (_, _, _, c) = fromIntegral $ fromEnum c

vectorToCube :: Vector Z -> Cube
vectorToCube = zipWith go solvedCube . fmap toColor . LA.toList
  where
    toColor :: Z -> Color
    toColor = toEnum . fromIntegral

    go :: (Side, Int, Int, Color) -> Color -> (Side, Int, Int, Color)
    go (s, x, y, _) c = (s, x, y, c)

oppositeSides :: Side -> Side -> Bool
oppositeSides Top Bottom = True
oppositeSides Front Back = True
oppositeSides Right Left = True
oppositeSides x y
  | fromEnum x > fromEnum y = oppositeSides y x
  | otherwise = False

adjacentSides :: Side -> Side -> Bool
adjacentSides x y
  | x == y = False
  | otherwise = not $ oppositeSides x y