{-# LANGUAGE ViewPatterns #-}
--------------------------------------------------------------------------------
-- |
-- Module      : ArrayFire.Signal
-- Copyright   : David Johnson (c) 2019-2020
-- License     : BSD 3
-- Maintainer  : David Johnson <djohnson.m@gmail.com>
-- Stability   : Experimental
-- Portability : GHC
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft.htm)
--
-- Functions performing signal processing on 'Array'
--
-- >>> input = vector 3 [10,20,30]
-- >>> positions = vector 5 [0.0, 0.5, 1.0, 1.5, 2.0]
-- >>> approx1 @Double input positions Cubic 0.0
-- ArrayFire Array
-- [5 1 1 1]
--    10.0000
--    13.7500
--    20.0000
--    26.2500
--    30.0000
--
--------------------------------------------------------------------------------
module ArrayFire.Signal where

import Data.Complex

import ArrayFire.FFI
import ArrayFire.Internal.Signal
import ArrayFire.Internal.Types

-- | 'approx1' interpolates data along the first dimensions
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__approx1.htm)
--
-- Interpolation is performed assuming input data is equally spaced with indices in the range [0, n). The positions are sampled with respect to data at these locations.
--
-- >>> input = vector 3 [10,20,30]
-- >>> positions = vector 5 [0.0, 0.5, 1.0, 1.5, 2.0]
-- >>> approx1 @Double input positions Cubic 0.0
-- ArrayFire Array
-- [5 1 1 1]
--    10.0000
--    13.7500
--    20.0000
--    26.2500
--    30.0000
approx1
  :: AFType a
  => Array a
  -- ^ the input array
  -> Array a
  -- ^ array contains the interpolation locations
  -> InterpType
  -- ^ is the interpolation type, it can take one of the values defined by 'InterpType'
  -> Float
  -- ^ is the value that will set in the output array when certain index is out of bounds
  -> Array a
  -- ^ is the array with interpolated values
approx1 arr1 arr2 (fromInterpType -> i1) f =
  op2 arr1 arr2 (\p x y -> af_approx1 p x y i1 f)

-- | approx2 performs interpolation on data along the first and second dimensions.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__approx2.htm)
--
-- Interpolation is performed assuming input data is equally spaced with indices in the range [0, n) along each dimension. The positions are sampled with respect to data at these locations.
--
-- >>> input = matrix @Double (3,3) [ [ 1.0,1.0,1.0 ], [ 2.0, 2.0, 2.0 ], [ 3.0,3.0,3.0 ] ]
-- >>> positions1 = matrix @Double (2,2) [ [ 0.5,1.5 ],[ 0.5,1.5 ] ]
-- >>> positions2 = matrix @Double (2,2) [ [ 0.5,0.5 ],[ 1.5,1.5 ] ]
-- >>> approx2 @Double input positions1 positions2 Cubic 0.0
-- ArrayFire Array
-- [2 2 1 1]
--     1.3750     2.6250
--     1.3750     2.6250
--
approx2
  :: AFType a
  => Array a
  -- ^ is the input array
  -> Array a
  -- ^ array contains the interpolation locations for first dimension
  -> Array a
  -- ^ array contains the interpolation locations for second dimension
  -> InterpType
  -- ^ is the interpolation type, it can take one of the values defined by 'InterpType'
  -> Float
  -- ^ is the value that will set in the output array when certain index is out of bounds
  -> Array a
  -- ^	is the array with interpolated values
approx2 arr1 arr2 arr3 (fromInterpType -> i1) f =
  op3 arr1 arr2 arr3 (\p x y z -> af_approx2 p x y z i1 f)

-- DMJ: Where did these functions go? Were they removed?
-- http://arrayfire.org/docs/group__approx__mat.htm
-- approx1Uniform
--   :: AFType a
--   => Array a
--   -> Array a
--   -> Int
--   -> Double
--   -> Double
--   -> InterpType
--   -> Float
--   -> Array a
-- approx1Uniform arr1 arr2 (fromIntegral -> i1) d1 d2 (fromInterpType -> interp) f =
--   op2 arr1 arr2 (\p x y -> af_approx1_uniform p x y i1 d1 d2 interp f)

-- approx2Uniform
--   :: AFType a
--   => Array a
--   -> Array a
--   -> Int
--   -> Double
--   -> Double
--   -> Array a
--   -> Int
--   -> Double
--   -> Double
--   -> InterpType
--   -> Float
--   -> Array a
-- approx2Uniform arr1 arr2 (fromIntegral -> i1) d1 d2 arr3 (fromIntegral -> i2) d3 d4 (fromInterpType -> interp) f =
--   op3 arr1 arr2 arr3 (\p x y z -> af_approx2_uniform p x y i1 d1 d2 z i2 d3 d4 interp f)

-- | Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft.htm#ga64d0db9e59c9410ba738591ad146a884)
--
-- The Fast Fourier Transform (FFT) is an efficient algorithm to compute the discrete Fourier transform (DFT) of a signal or array. This is most commonly used to convert data in the time (or space) domain to the frequency domain, Then, the inverse FFT (iFFT) is used to return the data to the original domain.
--
-- >>> fft (vector @Double 10 [1..]) 2.0 10
-- ArrayFire Array
-- [10 1 1 1]
--          (110.0000,0.0000)
--          (-10.0000,30.7768)
--          (-10.0000,13.7638)
--          (-10.0000,7.2654)
--          (-10.0000,3.2492)
--          (-10.0000,0.0000)
--          (-10.0000,-3.2492)
--          (-10.0000,-7.2654)
--          (-10.0000,-13.7638)
--          (-10.0000,-30.7768)
fft
  :: (AFType a, Fractional a)
  => Array a
  -- ^ input 'Array'
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals - used to either truncate or pad the input signals.
  -> Array a
  -- ^ is the transformed array
fft a d (fromIntegral -> x) =
  op1 a (\j k -> af_fft j k d x)

-- | Fast Fourier Transform (in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft.htm#gaa2f03c9ee1cb80dc184c0b0a13176da1)
--
-- C Interface for fast fourier transform on one dimensional signals.
--
-- *Note* The input in must be a complex array
--
fftInPlace
  :: (AFType a, Fractional a)
  => Array (Complex a)
  -- ^ is the input array on entry and the output of 1D forward fourier transform at exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
fftInPlace a d = a `inPlace` (flip af_fft_inplace d)

-- | Fast Fourier Transform (2-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft2.htm#gaab3fb1ed398e208a615036b4496da611)
--
-- C Interface for fast fourier transform on two dimensional signals.
--
fft2
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Array a
  -- ^ the transformed array
fft2 a d x y =
  op1 a (\j k -> af_fft2 j k d (fromIntegral x) (fromIntegral y))

-- | Fast Fourier Transform (2-dimensional, in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft2.htm#gacdeebb3f221ae698833dc4900a172b8c)
--
-- C Interface for fast fourier transform on two dimensional signals.
--
-- *Note* The input in must be a complex array
--
fft2_inplace
  :: (Fractional a, AFType a)
  => Array (Complex a)
  -- ^  input array on entry and the output of 2D forward fourier transform on exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
fft2_inplace a d = a `inPlace` (flip af_fft2_inplace d)

-- | Fast Fourier Transform (3-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft3.htm#ga5138ef1740ece0fde2c796904d733c12)
--
-- C Interface for fast fourier transform on three dimensional signals.
--
fft3
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along third dimension - used to either truncate/pad the input
  -> Array a
  -- ^ the transformed array
fft3 a d x y z =
  op1 a (\j k -> af_fft3 j k d (fromIntegral x) (fromIntegral y) (fromIntegral z))

-- | Fast Fourier Transform (3-dimensional, in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft2.htm#gacdeebb3f221ae698833dc4900a172b8c)
--
-- C Interface for fast fourier transform on three dimensional signals.
--
-- *Note* The input in must be a complex array
--
fft3_inplace
  :: (Fractional a, AFType a)
  => Array (Complex a)
  -- ^  input array on entry and the output of 3D forward fourier transform on exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
fft3_inplace a d = a `inPlace` (flip af_fft3_inplace d)

-- | Inverse Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft.htm#ga2d62c120b474b3b937b0425c994645fe)
--
-- C Interface for inverse fast fourier transform on one dimensional signals.
--
ifft
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^  is the length of output signals - used to either truncate or pad the input signals
  -> Array a
  -- ^ the transformed array
ifft a d x =
  op1 a (\j k -> af_ifft j k d (fromIntegral x))

-- | Inverse Fast Fourier Transform (in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft.htm#ga827379bef0e2cadb382c1b6301c91429)
--
-- C Interface for fast fourier transform on one dimensional signals.
--
ifft_inplace
  :: (AFType a, Fractional a)
  => Array (Complex a)
  -- ^ is the input array on entry and the output of 1D forward fourier transform at exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
ifft_inplace a d = a `inPlace` (flip af_ifft_inplace d)

-- | Inverse Fast Fourier Transform (2-dimensional signals)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft2.htm#ga7cd29c6a35c19240635b62cc5c30dc4f)
--
-- C Interface for inverse fast fourier transform on two dimensional signals.
--
ifft2
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Array a
  -- ^ the transformed array
ifft2 a d x y =
  op1 a (\j k -> af_ifft2 j k d (fromIntegral x) (fromIntegral y))

-- | Inverse Fast Fourier Transform (2-dimensional, in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft2.htm#ga9e6a165d44306db4552a56d421ce5d05)
--
-- C Interface for fast fourier transform on two dimensional signals.
--
ifft2_inplace
  :: (AFType a, Fractional a)
  => Array (Complex a)
  -- ^ is the input array on entry and the output of 1D forward fourier transform at exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
ifft2_inplace a d = a `inPlace` (flip af_ifft2_inplace d)

-- | Inverse Fast Fourier Transform (3-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft3.htm)
--
-- C Interface for inverse fast fourier transform on three dimensional signals.
--
ifft3
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along third dimension - used to either truncate/pad the input
  -> Array a
  -- ^ the transformed array
ifft3 a d x y z =
  op1 a (\j k -> af_ifft3 j k d (fromIntegral x) (fromIntegral y) (fromIntegral z))

-- | Inverse Fast Fourier Transform (3-dimensional, in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft3.htm#ga439a7a49723bc6cf77cf4fe7f8dfe334)
--
-- C Interface for fast fourier transform on two dimensional signals.
--
ifft3_inplace
  :: (AFType a, Fractional a)
  => Array (Complex a)
  -- ^ is the input array on entry and the output of 1D forward fourier transform at exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
ifft3_inplace a d = a `inPlace` (flip af_ifft3_inplace d)

-- | Real to Complex Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__r2c.htm#ga7486f342182a18e773f14cc2ab4cb551)
--
-- C Interface for real to complex fast fourier transform for one dimensional signals.
--
-- The first dimension of the output will be of size (pad0 / 2) + 1. The second dimension of the output will be pad1. The third dimension of the output will be pad 2.
--
fftr2c
  :: (Fractional a, AFType a)
  => Array a
  -- ^ is a real array
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Array a
  -- ^ is a complex array containing the non redundant parts of in.
fftr2c a d x =
  op1 a (\j k -> af_fft_r2c j k d (fromIntegral x))

-- | Real to Complex Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__r2c.htm#ga7486f342182a18e773f14cc2ab4cb551)
--
-- C Interface for real to complex fast fourier transform for two dimensional signals.
--
-- The first dimension of the output will be of size (pad0 / 2) + 1. The second dimension of the output will be pad1. The third dimension of the output will be pad 2.
--
fft2r2c
  :: (Fractional a, AFType a)
  => Array a
  -- ^ is a real array
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Array a
  -- ^ is a complex array containing the non redundant parts of in.
fft2r2c a d x y =
  op1 a (\j k -> af_fft2_r2c j k d (fromIntegral x) (fromIntegral y))

-- | Real to Complex Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__r2c.htm#gab4ca074b54218b74d8cfbda63d38be51)
--
-- C Interface for real to complex fast fourier transform for three dimensional signals.
--
-- The first dimension of the output will be of size (pad0 / 2) + 1. The second dimension of the output will be pad1. The third dimension of the output will be pad 2.
--
fft3r2c
  :: (Fractional a, AFType a)
  => Array a
  -- ^ is a real array
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along third dimension - used to either truncate/pad the input
  -> Array a
  -- ^ is a complex array containing the non redundant parts of in.
fft3r2c a d x y z =
  op1 a (\j k -> af_fft3_r2c j k d (fromIntegral x) (fromIntegral y) (fromIntegral z))

-- | Complex to Real Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__c2r.htm#gaa5efdfd84213a4a07d81a5d534cde5ac)
--
-- C Interface for complex to real fast fourier transform for one dimensional signals.
--
-- The first dimension of the output will be 2 * dim0 - 1 if is_odd is true else 2 * dim0 - 2 where dim0 is the first dimension of the input. The remaining dimensions are unchanged.
--
fftc2r
  :: AFType a
  => Array a
  -- ^ is a complex array containing only the non redundant parts of the signals.
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Bool
  -- ^ is a flag signifying if the output should be even or odd size
  -> Array a
  -- ^ is a real array containing the output of the transform.
fftc2r a cm (fromIntegral . fromEnum -> cd) = op1 a (\x y -> af_fft_c2r x y cm cd)

-- | Complex to Real Fast Fourier Transform (2-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__c2r.htm#gaaa7da16f226cacaffced631e08da4493)
--
-- C Interface for complex to real fast fourier transform for two dimensional signals.
--
-- The first dimension of the output will be 2 * dim0 - 1 if is_odd is true else 2 * dim0 - 2 where dim0 is the first dimension of the input. The remaining dimensions are unchanged.
--
fft2C2r
  :: AFType a
  => Array a
  -- ^ is a complex array containing only the non redundant parts of the signals.
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Bool
  -- ^ is a flag signifying if the output should be even or odd size
  -> Array a
  -- ^ is a real array containing the output of the transform.
fft2C2r a cm (fromIntegral . fromEnum -> cd) = op1 a (\x y -> af_fft2_c2r x y cm cd)

-- | Complex to Real Fast Fourier Transform (3-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__c2r.htm#gaa9b3322d9ffab15268919e1f114bed24)
--
-- C Interface for complex to real fast fourier transform for three dimensional signals.
--
-- The first dimension of the output will be 2 * dim0 - 1 if is_odd is true else 2 * dim0 - 2 where dim0 is the first dimension of the input. The remaining dimensions are unchanged.
--
fft3C2r
  :: AFType a
  => Array a
  -- ^ is a complex array containing only the non redundant parts of the signals.
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Bool
  -- ^ is a flag signifying if the output should be even or odd size
  -> Array a
  -- ^ is a real array containing the output of the transform.
fft3C2r a cm (fromIntegral . fromEnum -> cd) = op1 a (\x y -> af_fft3_c2r x y cm cd)

-- | Convolution Integral for one dimensional data
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__convolve1.htm#ga25d77b794463b5cd72cd0b7f4af140d7)
--
-- C Interface for convolution on one dimensional signals.
--
-- *Note* The default parameter of domain, AF_CONV_AUTO, heuristically switches between frequency and spatial domain.
--
convolve1
  :: AFType a
  => Array a
  -- ^ the input signal
  -> Array a
  -- ^ the signal that shall be flipped for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> ConvDomain
  -- ^ specifies if the convolution should be performed in frequency os spatial domain
  -> Array a
  -- ^ convolved array
convolve1 a b (toConvMode -> cm) (fromConvDomain -> cd) = op2 a b (\x y z -> af_convolve1 x y z cm cd)

-- | Convolution Integral for two dimensional data
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__convolve2.htm#ga25d77b794463b5cd72cd0b7f4af140d7)
--
-- C Interface for convolution on two dimensional signals.
--
-- *Note* The default parameter of domain, AF_CONV_AUTO, heuristically switches between frequency and spatial domain.
--
convolve2
  :: AFType a
  => Array a
  -- ^ the input signal
  -> Array a
  -- ^ the signal that shall be flipped for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> ConvDomain
  -- ^ specifies if the convolution should be performed in frequency os spatial domain
  -> Array a
  -- ^ convolved array
convolve2 a b (toConvMode -> cm) (fromConvDomain -> cd) = op2 a b (\x y z -> af_convolve2 x y z cm cd)

-- | Convolution Integral for three dimensional data
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__convolve3.htm#ga25d77b794463b5cd72cd0b7f4af140d7)
--
-- C Interface for convolution on three dimensional signals.
--
-- *Note* The default parameter of domain, AF_CONV_AUTO, heuristically switches between frequency and spatial domain.
--
convolve3
  :: AFType a
  => Array a
  -- ^ the input signal
  -> Array a
  -- ^ the signal that shall be flipped for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> ConvDomain
  -- ^ specifies if the convolution should be performed in frequency os spatial domain
  -> Array a
  -- ^ convolved array
convolve3 a b (toConvMode -> cm) (fromConvDomain -> cd) =
  op2 a b (\x y z -> af_convolve3 x y z cm cd)

-- | C Interface for separable convolution on two dimensional signals.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__convolve.htm#gaeb6ba88155cf3ef29d93f97b147e372f)
--
-- C Interface for separable convolution on two dimensional signals.
--
-- *Note* The default parameter of domain, AF_CONV_AUTO, heuristically switches between frequency and spatial domain.
--
convolve2Sep
  :: AFType a
  => Array a
  -- ^ filter that has to be applied along the coloumns
  -> Array a
  -- ^ filter that has to be applied along the rows
  -> Array a
  -- ^ the input array
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> Array a
  -- ^ convolved array
convolve2Sep a b c (toConvMode -> d) = op3 a b c (\x y z j -> af_convolve2_sep x y z j d)

-- DMJ: did this get removed? Can't find in latest docs
-- fftConvolve1
--   :: AFType a
--   => Array a
--   -> Array a
--   -> ConvMode
--   -> Array a
-- fftConvolve1 a b (toConvMode -> c) = op2 a b (\x y z -> af_fft_convolve1 x y z c)

-- | 2D Convolution using Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__convolve2.htm#gab52ebe631d8358cdef1b5c8a95550556)
--
-- A convolution is a common operation between a source array, a, and a filter (or kernel) array b. The answer to the convolution is the same as computing the coefficients in polynomial multiplication, if a and b are the coefficients.
--
-- C Interface for FFT-based convolution on two dimensional signals.
--
fftConvolve2
  :: AFType a
  => Array a
  -- ^ is the input signal
  -> Array a
  --  ^ is the signal that shall be used for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> Array a
  -- ^  is convolved array
fftConvolve2 a b (toConvMode -> c) = op2 a b (\x y z -> af_fft_convolve2 x y z c)

-- | 3D Convolution using Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__convolve3.htm)
--
-- A convolution is a common operation between a source array, a, and a filter (or kernel) array b. The answer to the convolution is the same as computing the coefficients in polynomial multiplication, if a and b are the coefficients.
--
-- C Interface for FFT-based convolution on three dimensional signals.
--
fftConvolve3
  :: AFType a
  => Array a
  -- ^ is the input signal
  -> Array a
  --  ^ is the signal that shall be used for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> Array a
  -- ^  is convolved array
fftConvolve3 a b (toConvMode -> c) = op2 a b (\x y z -> af_fft_convolve3 x y z c)

-- | Finite Impulse Filter.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fir.htm#ga2a850e69775eede4709e0d607bba240b)
--
-- Finite impulse filters take an input x and a co-efficient array b to generate an output y such that:
--
-- C Interface for finite impulse response filter.
--
fir
  :: AFType a
  => Array a
  -- ^ is the input signal to the filter
  -> Array a
  -- ^ is the array containing the coefficients of the filter
  -> Array a
  -- ^ is the output signal from the filter
fir a b = op2 a b af_fir

-- | Infinite Impulse Filter.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__iir.htm#ga7adcc364da0a66cdfd2bb351215456c4)
--
-- Infinite impulse filters take an input x and a feedforward array b, feedback array a to generate an output y such that:
--
-- C Interface for infinite impulse response filter.
--
-- *Note* The feedforward coefficients are currently limited to a length of 512
--
iir
  :: AFType a
  => Array a
  -- ^ the array containing the feedforward coefficient
  -> Array a
  -- ^ is the array containing the feedback coefficients
  -> Array a
  -- ^ is the input signal to the filter
  -> Array a
  -- ^ the output signal from the filter
iir a b c = op3 a b c af_iir

-- | Median Filter
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__image__func__medfilt.htm#gaaf3f62f2de0f4dc315b831e494e1b2c0)
--
-- A median filter is similar to the arbitrary filter except that instead of a weighted sum, the median value of the pixels covered by the kernel is returned.
--
-- C Interface for median filter.
--
medFilt
  :: AFType a
  => Array a
  -- ^ 'Array' is the input image
  -> Int
  -> Int
  -> BorderType
  -> Array a
  -- ^ 'Array' is the processed image
medFilt a l w (fromBorderType -> b) =
 a `op1` (\x y -> af_medfilt x y (fromIntegral l) (fromIntegral w) b)

-- | 1D Median Filter
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__image__func__medfilt.htm#gad108ea62cbbb5371bd14a17d06384359)
--
-- A median filter is similar to the arbitrary filter except that instead of a weighted sum, the median value of the pixels covered by the kernel is returned.
--
-- C Interface for 1D median filter.
--
medFilt1
  :: AFType a
  => Array a
  -- ^ 'Array' is the input signal
  -> Int
  -- ^ Is the kernel width
  -> BorderType
  -- ^ value will decide what happens to border when running filter in their neighborhood. It takes one of the values [AF_PAD_ZERO | AF_PAD_SYM]
  -> Array a
  -- ^ 'Array' is the processed signal
medFilt1 a w (fromBorderType -> b) =
 a `op1` (\x y -> af_medfilt1 x y (fromIntegral w) b)

-- | 2D Median Filter
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__image__func__medfilt.htm#ga2cb99dca5842f74f6b9cd28eb187a9cd)
--
-- A median filter is similar to the arbitrary filter except that instead of a weighted sum, the median value of the pixels covered by the kernel is returned.
--
-- C Interface for 2D median filter.
--
medFilt2
  :: AFType a
  => Array a
  -- ^ 'Array' is the input image
  -> Int
  -- ^ the kernel height
  -> Int
  -- ^ the kernel width
  -> BorderType
  -- ^ value will decide what happens to border when running filter in their neighborhood. It takes one of the values [AF_PAD_ZERO | AF_PAD_SYM]
  -> Array a
  -- ^ 'Array' is the processed image
medFilt2 a l w (fromBorderType -> b) =
 a `op1` (\x y -> af_medfilt2 x y (fromIntegral l) (fromIntegral w) b)

-- | C Interface for setting plan cache size.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft.htm#ga4ddef19b43d9a50c97b1a835df60279a)
--
-- This function doesn't do anything if called when CPU backend is active. The plans associated with the most recently used array sizes are cached.
--
-- >>> setFFTPlanCacheSize 2
-- ()
--
setFFTPlanCacheSize
  :: Int
  -- ^ is the number of plans that shall be cached
  -> IO ()
setFFTPlanCacheSize =
 afCall . af_set_fft_plan_cache_size . fromIntegral


module Numeric.DataFrame
  ( module Numeric.DataFrame.Type
    -- * Simplified type aliases
  , module Numeric.Scalar
  , module Numeric.Vector
  , module Numeric.Matrix
    -- * Functionality
  , module Numeric.DataFrame.SubSpace
  , module Numeric.DataFrame.Contraction
  , module Numeric.Basics
  , module Numeric.Subroutine.Sort
  ) where

import Numeric.DataFrame.Internal.Backend ()

import Numeric.DataFrame.Contraction
import Numeric.DataFrame.SubSpace
import Numeric.DataFrame.Type

import Numeric.Matrix
import Numeric.Scalar
import Numeric.Vector

import Numeric.Basics

import Numeric.Subroutine.Sort (sort, sortBy)


{-# LANGUAGE BangPatterns #-}

-- | 2D FFT functions
module SatOpt.FFT where

import           Control.Parallel.Strategies   (parMap, rdeepseq)
import           Numeric.GSL.Fourier           (fft, ifft)
import           Numeric.LinearAlgebra.HMatrix

half [] = []
half (x:y:xs) = x : half xs

clone [] = []
clone (x:xs) = x : x : clone xs

-- | Take the FFT of a matrix of complex numbers
fft2dM :: Matrix ℂ -> Matrix ℂ
fft2dM m = fromColumns $ clone $ clone $ inter0
  where
    !inter0 = mapFFT $ inter1
    !inter1 = half $ half $ toColumns $ fromRows $ clone $ clone $ inter2
    !inter2 = mapFFT $ inter3
    !inter3 = half $ half $ toRows m
    mapFFT = parMap rdeepseq fft

-- | Take the inverse FFT of a matrix of complex numbers
ifft2dM :: Matrix ℂ -> Matrix ℂ
ifft2dM m = fromColumns $ mapIFFT $ toColumns inter
  where
    !inter = fromRows $ mapIFFT $ toRows m
    mapIFFT = parMap rdeepseq ifft


module RandomUtil.RandomMatrixTest
  (
    testsRandomMatrix
  ) where

import Test.HUnit

import Numeric.LinearAlgebra.HMatrix hiding (corr)
import System.Random

import RandomUtil.RandomMatrix
import Matrix.State

mat :: Matrix Double
mat = reshape 3 $ fromList $ take 9 [1..]

data Two = Two

instance RandomGen Two where
  next g  = (2, g)
  split g = (g, g)

testsRandomMatrix :: Test
testsRandomMatrix = TestList [ TestLabel "selectRandomMatrixElement" testsSelectRandomMatrixElement,
                               TestLabel "randomMatrixIndex" testsRandomMatrixIndex
                             ]

testsSelectRandomMatrixElement :: Test
testsSelectRandomMatrixElement = TestList [TestLabel "test1" testSelects]
testsRandomMatrixIndex :: Test
testsRandomMatrixIndex = TestList [TestLabel "test1" testInBounds]

-- testsSelectRandomMatrixElement
testSelects :: Test
testSelects = TestCase (assertEqual "selectRandomMatrixElement - Select"
                         5
                         (fst $ selectRandomMatrixElement Two mat))

-- testsRandomMatrixIndex
testInBounds :: Test
testInBounds = TestCase (do gen <- getStdGen
                            let listI  = getListRandomIndex gen 100 mat
                            let result = checkAllTrue listI (isValid mat)
                            assertEqual "randomMatrixIndex - Indices in bounds" True result)

-- testsRandomMatrixIndex helper function
isValid :: Matrix a -> MatrixIndex -> Bool
isValid m (i,j) = i >= 0 && j >= 0 && i < rows m && j < cols m

getListRandomIndex :: (Element a) => StdGen -> Int -> Matrix a -> [MatrixIndex]
getListRandomIndex _ 0 _ = []
getListRandomIndex g n m = rIndex : getListRandomIndex g' (n-1) m
  where
    (rIndex,g') = randomMatrixIndex g m

checkAllTrue :: [a] -> (a -> Bool) -> Bool
checkAllTrue xs f = foldr (&&) True $ map f xs


{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators    #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
module Main where

import qualified Control.Foldl                 as FL
import qualified Data.List                     as List
import           Data.Maybe                     ( fromMaybe )
import qualified Data.Text                     as T
import qualified Data.Vector.Storable          as V
import qualified Data.Vinyl                    as V
import qualified Frames                        as F

import           Numeric.LinearAlgebra          ( (#>)
                                                , (<#)
                                                , (<.>)
                                                , (<\>)
                                                )
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data     ( Matrix
                                                , R
                                                , Vector
                                                )
import qualified Numeric.LinearAlgebra.Data    as LA
import           Statistics.Types               ( cl95 )

import           Control.Applicative
import           Control.Monad
import           EasyTest
import           Control.Monad.IO.Class         ( MonadIO
                                                , liftIO
                                                )

import qualified Math.Regression.LeastSquares  as LS
import qualified Math.Regression.Regression    as RE

import qualified Knit.Effect.Logger            as KL
import qualified Polysemy.IO                   as P
import qualified Polysemy                      as P

main = runOnly "all" suite

suite :: Test
suite = scope "all.Regression" $ tests [testOLS, testWOLS, testTLS, testWTLS]

logged
  :: forall m a
   . MonadIO m
  => P.Sem '[KL.Logger KL.LogEntry, KL.PrefixLog, P.Lift IO, P.Lift m] a
  -> m a
logged = P.runM . P.runIO @m . KL.filteredLogEntriesToIO KL.logAll
-- regression tests

-- build some data for testing
-- uniformly distributed measurements, normally distributed noise.  Separate noise amplitudes for xs and ys
-- also allow building heteroscedastic data
buildRegressable
  :: [Double]
  -> Maybe (LA.Vector R)
  -> Double
  -> LA.Vector R
  -> Double
  -> IO (LA.Vector R, LA.Matrix R)
buildRegressable variances offsetsM noiseObs coeffs noiseMeas = do
  -- generate random measurements
  let d    = LA.size coeffs
      nObs = List.length variances
      xsO  = LA.asColumn (LA.fromList (List.replicate nObs 1))
        LA.<> LA.asRow (fromMaybe (LA.fromList $ List.replicate d 0) offsetsM)
  xs0    <- LA.rand nObs d
  xNoise <- LA.randn nObs d
  let xs  = xs0 + xsO + LA.scale noiseMeas xNoise
  let ys0 = (xs0 + xsO) LA.<> LA.asColumn coeffs
  yNoise <- fmap (List.head . LA.toColumns) (LA.randn nObs 1)
  let ys = ys0 + LA.asColumn
        (LA.scale
          noiseObs
          (V.zipWith (*) yNoise (LA.cmap sqrt (LA.fromList variances)))
        )
  return (List.head (LA.toColumns ys), xs)

unweighted :: Int -> [Double]
unweighted n = List.replicate n (1.0)

coneWeighted :: Int -> Double -> [Double]
coneWeighted n increment =
  let w0 = [ 1 + (realToFrac i) * increment | i <- [0 .. n] ]
      s  = realToFrac n / FL.fold FL.sum w0
  in  fmap (* s) w0

showText :: Show a => a -> T.Text
showText = T.pack . show

errR2 res = let r2 = RE.rSquared res in (r2 <= 1 && r2 > 0.0)
coeffs :: LA.Vector R = LA.fromList [1.0, 2.2, 0.3]
offsets :: LA.Vector R = LA.fromList [0, 1, 0]
vars = coneWeighted 100 0.1
varListToWeights = LA.cmap (\x -> 1 / sqrt x) . LA.fromList
wgts = varListToWeights vars

testRegression
  :: ( r
         ~
         '[KL.Logger KL.LogEntry, KL.PrefixLog, P.Lift IO, P.Lift (PropertyT IO)]
     ) --KL.LogWithPrefixesLE r
  => Double
  -> Double
  -> Bool
  -> Bool
  -> (  Matrix R
     -> Vector R
     -> Vector R
     -> P.Sem r (RE.RegressionResult R)
     )
  -> Test
testRegression yNoise xNoise weighted offset f = do
  let scopeT =
        "Vy="
          <> showText yNoise
          <> "; Vx="
          <> showText xNoise
          <> (if weighted then "; cone weights" else "")
          <> (if offset then "; w/offsets" else "")
      vars    = if weighted then coneWeighted 100 0.1 else unweighted 100
      wgts    = varListToWeights vars
      offsetM = if offset then (Just offsets) else Nothing
  property $ do
    (ys, xs) <- liftIO $ buildRegressable vars offsetM yNoise coeffs xNoise
    result   <- logged $ f xs ys wgts
    footnote
      ( T.unpack
      $ RE.prettyPrintRegressionResult "y" ["x1", "x2", "x3"] result cl95
      )
    success -- expect $ errR2 result

const3 f x y z = f x y

type RegressF = (Matrix R
                -> Vector R
                -> Vector R
                -> P.Sem
                   '[KL.Logger KL.LogEntry, KL.PrefixLog, P.Lift IO, P.Lift
                                                                     (PropertyT IO)]
                   (RE.RegressionResult R))

testOLS :: Test
testOLS =
  let regress :: RegressF
      regress xs ys ws = LS.ordinaryLS False xs ys
  in  scope "OLS" $ tests
        [ testRegression 0   0   False False regress
        , testRegression 0.1 0   False False regress
        , testRegression 0.5 0   False False regress
        , testRegression 0.1 0   False True  regress
        , testRegression 0.3 0.3 False False regress
        , testRegression 0.3 0   True  False regress
        , testRegression 0.3 0.3 True  False regress
        ]

testWOLS :: Test
testWOLS =
  let regress :: RegressF = LS.weightedLS False
  in  scope "WOLS" $ tests
        [ testRegression 0   0   True  False regress
        , testRegression 0.1 0   True  False regress
        , testRegression 0.1 0   False False regress
        , testRegression 0.3 0   True  False regress
        , testRegression 0.1 0.1 True  False regress
        , testRegression 0.1 0.1 True  True  regress
        , testRegression 0.3 0.3 True  False regress
        ]

testTLS :: Test
testTLS =
  let regress :: RegressF
      regress xs ys ws = LS.totalLS False xs ys
  in  scope "TLS" $ tests
        [ testRegression 0   0   False False regress
        , testRegression 0.1 0   False False regress
        , testRegression 0.5 0   False False regress
        , testRegression 0.1 0   False True  regress
        , testRegression 0.3 0.3 False False regress
        , testRegression 0.3 0   True  False regress
        , testRegression 0.3 0.3 True  False regress
        ]

testWTLS :: Test
testWTLS =
  let regress :: RegressF
      regress = LS.weightedTLS False
  in  scope "WTLS" $ tests
        [ testRegression 0   0   True  False regress
        , testRegression 0.1 0   True  False regress
        , testRegression 0.1 0   False False regress
        , testRegression 0.3 0   True  False regress
        , testRegression 0.1 0.1 True  False regress
        , testRegression 0.1 0.1 True  True  regress
        , testRegression 0.3 0.3 True  False regress
        ]


{-# LANGUAGE BangPatterns #-}

{- | Some extra facilities to work with 'Rand' monad and 'PureMT'
     random number generator.
-}

module Moo.GeneticAlgorithm.Random
    (
    -- * Random numbers from given range
      getRandomR
    , getRandom
    -- * Probability distributions
    , getNormal2
    , getNormal
    -- * Random samples and shuffles
    , randomSample
    , randomSampleIndices
    , shuffle
    -- * Building blocks
    , withProbability
    -- * Re-exports from random number generator packages
    , getBool, getInt, getWord, getInt64, getWord64, getDouble
    , runRandom, evalRandom, newPureMT
    , Rand, Random, PureMT
    ) where

import Control.Monad (liftM)
import Control.Monad.State.Strict (MonadState(..), State, runState, evalState)
import Data.Complex (Complex (..))
import System.Random (RandomGen, Random(..))
import System.Random.Mersenne.Pure64
import qualified System.Random.Shuffle as S
import qualified Data.Set as Set
import Data.Word (Word64)
import Data.Int (Int64)

-- | Monad for handling randomized computation. This uses the mersenne twister RNG.
type Rand = State PureMT

-- | Run a random computation using the generator @g@, returning the result
-- and the updated generator.
runRandom  :: Rand a -> PureMT -> (a, PureMT)
runRandom = runState

-- | Evaluate a random computation using the mersenne generator @g@.  Note that the
-- generator @g@ is not returned, so there's no way to recover the
-- updated version of @g@.
evalRandom :: Rand a -> PureMT -> a
evalRandom = evalState

-- | Yield a new boolean value from the generator.
getBool     :: Rand Bool
getBool     = fmap (<0) getInt

-- | Yield a new 'Int' value from the generator.
getInt      :: Rand Int
getInt      = state randomInt

-- | Yield a new 'Word' value from the generator.
getWord     :: Rand Word
getWord     = state randomWord

-- | Yield a new 'Int64' value from the generator.
getInt64    :: Rand Int64
getInt64    = state randomInt64

-- | Yield a new 53-bit precise 'Double' value from the generator.
getDouble :: Rand Double
getDouble = state randomDouble

-- | Yield a new 'Word64' value from the generator.
getWord64   :: Rand Word64
getWord64   = state randomWord64

-- | Yield a new randomly selected value of type @a@ in the range @(lo, hi)@.
-- See 'System.Random.randomR' for details.
getRandomR :: Random a => (a, a) -> Rand a
getRandomR range = state $ randomR range

-- | Yield a new randomly selected value of type @a@.
-- See 'System.Random.random' for details.
getRandom :: Random a => Rand a
getRandom = state random

-- | Yield two randomly selected values which follow standard
-- normal distribution.
getNormal2 :: Rand (Double, Double)
getNormal2 = do
  -- Box-Muller method
  u <- getDouble
  v <- getDouble
  let (c :+ s) = exp (0 :+ (2*pi*v))
  let r = sqrt $ (-2) * log u
  return (r*c, r*s)

-- | Yield one randomly selected value from standard normal distribution.
getNormal :: Rand Double
getNormal = fst `liftM` getNormal2

-- | Take at most n random elements from the list. Preserve order.
randomSample :: Int -> [a] -> Rand [a]
randomSample n xs =
  state $ \g -> select g n (length xs) xs []
  where
    select rng _ _ [] acc = (reverse acc, rng)
    select rng n m xs acc
        | n <= 0     = (reverse acc, rng)
        | otherwise  =
            let (k, rng') = randomR (0, m - n) rng
                (x:rest) = drop k xs
            in  select rng' (n-1) (m-k-1) rest (x:acc)

-- | Select @sampleSize@ numbers in the range from @0@ to @(populationSize-1)@.
-- The function works best when @sampleSize@ is much smaller than @populationSize@.
randomSampleIndices :: Int -> Int -> Rand [Int]
randomSampleIndices sampleSize populationSize =
    state $ \g ->
        let (sampleSet, g') = buildSampleSet g sampleSize Set.empty
        in  (Set.toList sampleSet, g')
  where
    buildSampleSet g n s
        | n <= 0 = (s, g)
        | otherwise =
            let (i, g') = randomR (0, populationSize-1) g
            in  if (i `Set.member` s)
                then buildSampleSet g' n s
                else buildSampleSet g' (n-1) (Set.insert i s)

-- | Randomly reorder the list.
shuffle :: [a] -> Rand [a]
shuffle xs = state $ randomShuffle xs (length xs)

-- | Given a sequence (e1,...en) to shuffle, its length, and a random
-- generator, compute the corresponding permutation of the input
-- sequence, return the permutation and the new state of the
-- random generator.
randomShuffle :: RandomGen gen => [a] -> Int -> gen -> ([a], gen)
randomShuffle elements len g =
    let (rs, g') = rseq len g
    in  (S.shuffle elements rs, g')
  where
  -- | The sequence (r1,...r[n-1]) of numbers such that r[i] is an
  -- independent sample from a uniform random distribution
  -- [0..n-i]
  rseq :: RandomGen gen => Int -> gen -> ([Int], gen)
  rseq n g = second lastGen . unzip $ rseq' (n - 1) g
      where
        rseq' :: RandomGen gen => Int -> gen -> [(Int, gen)]
        rseq' i gen
          | i <= 0    = []
          | otherwise = let (j, gen') = randomR (0, i) gen
                        in  (j, gen') : rseq' (i - 1) gen'
        -- apply a function on the second element of a pair
        second :: (b -> c) -> (a, b) -> (a, c)
        second f (x,y) = (x, f y)
        -- the last returned random number generator
        lastGen [] = g   -- didn't use the generator yet
        lastGen (lst:[]) = lst
        lastGen gens = lastGen (drop 1 gens)

-- |Modify value with probability @p@. Return the unchanged value with probability @1-p@.
withProbability :: Double -> (a -> Rand a) -> (a -> Rand a)
withProbability p modify x = do
  t <- getDouble
  if t < p
     then modify x
     else return x


module NeuralNet.TrainSpec (trainSpec) where

import Test.Hspec
import Numeric.LinearAlgebra
import NeuralNet.Net
import NeuralNet.Activation
import NeuralNet.Example
import NeuralNet.Layer
import NeuralNet.Train

trainSpec :: SpecWith ()
trainSpec =
  describe "NeuralNet.Train" $ do
    describe "nnBackward" $ do
      it "calculates correctly for simple example" $
        let
          layer1W = [-1.31386475, 0.88462238, 0.88131804, 1.70957306
                    , 0.05003364, -0.40467741, -0.54535995, -1.54647732
                    , 0.98236743, -1.10106763, -1.18504653, -0.2056499]
          layer1B = [1.48614836, 0.23671627, -1.02378514]
          layer2W = [-1.02387576, 1.12397796, -0.13191423]
          layer2B = [-1.62328545]
          nn = buildNNFromList (4, [LayerDefinition ReLU 3, LayerDefinition Sigmoid 1]) (layer1W ++ layer1B ++ layer2W ++ layer2B)
          examples = createExampleSet [ ([0.09649747, -0.2773882, -0.08274148, -0.04381817], 1)
                                      , ([-1.8634927, -0.35475898, -0.62700068, -0.47721803], 0)]

          inputs = exampleSetX examples
          layer1As = fromLists [[1.97611078, -1.24412333], [-0.62641691, -0.80376609], [-2.41908317, -0.92379202]]
          layer1Zs = fromLists [[-0.7129932, 0.62524497], [-0.16051336, -0.76883635], [-0.23003072, 0.74505627]]
          layer2As = fromLists [[1.78862847, 0.43650985]]
          layer2Zs = fromLists [[0.64667545, -0.35627076]]
          steps = [createInputForwardPropStep inputs, createForwardPropStep layer1As layer1Zs, createForwardPropStep layer2As layer2Zs]

          dW1 = fromLists [ [0.4101000190122487, 7.807203346853248e-2, 0.13798443685274048, 0.10502167417988162]
                          , [0, 0, 0, 0]
                          , [5.2836516249770524e-2, 1.0058654166727896e-2, 1.77776556985907e-2, 1.3530795262454464e-2]]
          db1 = fromLists [[-0.22007063350033446], [0], [-2.835348711039787e-2]]
          dW2 = fromLists [[-0.39202432174003965, -0.13325854898009654, -4.601088848081533e-2]]
          db2 = fromLists [[0.15187860742650333]]
        in
          nnBackward nn steps examples `shouldBe` [(dW1, db1), (dW2, db2)]

    describe "updateNNParams" $ do
      it "calculates correctly for logreg case" $
        let
          nn = buildNNFromList (2, [LayerDefinition ReLU 1]) [1, 2, 3]
          grads = [(fromLists [[10, 100]], fromLists [[1000]])]
          learningRate = 0.1
          expected = buildNNFromList (2, [LayerDefinition ReLU 1]) [0, -8, -97]
        in
          updateNNParams nn grads learningRate `shouldBe` expected


{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveAnyClass                           #-}
{-# LANGUAGE DeriveDataTypeable                       #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE DerivingVia                              #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TemplateHaskell                          #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Model.Neural.LSTM (
    -- * LSTM
    lstm
  , LSTMp(..), lstmForget, lstmInput, lstmUpdate, lstmOutput
  , reshapeLSTMpInput
  , reshapeLSTMpOutput
  , lstm'
    -- * GRU
  , gru
  , GRUp(..), gruMemory, gruUpdate, gruOutput
  , gru'
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Function
import           Backprop.Learn.Model.Neural
import           Backprop.Learn.Model.Regression
import           Backprop.Learn.Model.State
import           Backprop.Learn.Model.Types
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Primitive
import           Data.Type.Tuple
import           Data.Typeable
import           GHC.Generics                          (Generic)
import           GHC.TypeNats
import           Lens.Micro
import           Lens.Micro.TH
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.OneLiner
import           Numeric.Opto.Ref
import           Numeric.Opto.Update
import           Statistics.Distribution
import qualified Data.Binary                           as Bi
import qualified Numeric.LinearAlgebra.Static          as H
import qualified System.Random.MWC                     as MWC

-- TODO: allow parameterize internal activation function?
-- TODO: Peepholes

-- | 'LSTM' layer parmateters
data LSTMp (i :: Nat) (o :: Nat) =
    LSTMp { _lstmForget :: !(FCp (i + o) o)
          , _lstmInput  :: !(FCp (i + o) o)
          , _lstmUpdate :: !(FCp (i + o) o)
          , _lstmOutput :: !(FCp (i + o) o)
          }
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Bi.Binary, Regularize, Backprop)

deriving via (GNum (LSTMp i o)) instance (KnownNat i, KnownNat o) => Num (LSTMp i o)
deriving via (GNum (LSTMp i o)) instance (KnownNat i, KnownNat o) => Fractional (LSTMp i o)
deriving via (GNum (LSTMp i o)) instance (KnownNat i, KnownNat o) => Floating (LSTMp i o)

makeLenses ''LSTMp

instance (PrimMonad m, KnownNat i, KnownNat o) => Mutable m (LSTMp i o) where
    type Ref m (LSTMp i o) = GRef m (LSTMp i o)
    thawRef   = gThawRef
    freezeRef = gFreezeRef
    copyRef   = gCopyRef
instance (PrimMonad m, KnownNat i, KnownNat o) => LinearInPlace m Double (LSTMp i o)

instance (PrimMonad m, KnownNat i, KnownNat o) => Learnable m (LSTMp i o)

-- | Stateless version of 'lstm' that takes the "previous input" as a part
-- of the input vector.
lstm'
    :: (KnownNat i, KnownNat o)
    => Model ('Just (LSTMp i o)) ('Just (R o)) (R (i + o)) (R o)
lstm' = modelD $ \(PJust p) x (PJust s) ->
    let forget = logistic $ runLRp (p ^^. lstmForget) x
        input  = logistic $ runLRp (p ^^. lstmInput ) x
        update = tanh     $ runLRp (p ^^. lstmUpdate) x
        s'     = forget * s + input * update
        o      = logistic $ runLRp (p ^^. lstmOutput) x
        h      = o * tanh s'
    in  (h, PJust s')

-- | Long-term short-term memory layer
--
-- <http://colah.github.io/posts/2015-08-Understanding-LSTMs/>
--
lstm
    :: (KnownNat i, KnownNat o)
    => Model ('Just (LSTMp i o)) ('Just (R o :# R o)) (R i) (R o)
lstm = recurrent H.split (H.#) id lstm'

reshapeLSTMpInput
    :: (ContGen d, PrimMonad m, KnownNat i, KnownNat i', KnownNat o)
    => d
    -> MWC.Gen (PrimState m)
    -> LSTMp i o
    -> m (LSTMp i' o)
reshapeLSTMpInput d g (LSTMp forget input update output) =
    LSTMp <$> reshaper forget
          <*> reshaper input
          <*> reshaper update
          <*> reshaper output
  where
    reshaper = reshapeLRpInput d g

reshapeLSTMpOutput
    :: (ContGen d, PrimMonad m, KnownNat i, KnownNat o, KnownNat o')
    => d
    -> MWC.Gen (PrimState m)
    -> LSTMp i o
    -> m (LSTMp i o')
reshapeLSTMpOutput d g (LSTMp forget input update output) =
    LSTMp <$> reshaper forget
          <*> reshaper input
          <*> reshaper update
          <*> reshaper output
  where
    reshaper = reshapeLRpInput  d g
           <=< reshapeLRpOutput d g

-- | Forget biases initialized to 1
instance (KnownNat i, KnownNat o) => Initialize (LSTMp i o) where
    initialize d g = LSTMp <$> set (mapped . fcBias) 1 (initialize d g)
                           <*> initialize d g
                           <*> initialize d g
                           <*> initialize d g

-- | 'GRU' layer parmateters
data GRUp (i :: Nat) (o :: Nat) =
    GRUp { _gruMemory :: !(FCp (i + o) o)
         , _gruUpdate :: !(FCp (i + o) o)
         , _gruOutput :: !(FCp (i + o) o)
         }
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Bi.Binary, Initialize, Regularize, Backprop)

deriving via (GNum (GRUp i o)) instance (KnownNat i, KnownNat o) => Num (GRUp i o)
deriving via (GNum (GRUp i o)) instance (KnownNat i, KnownNat o) => Fractional (GRUp i o)
deriving via (GNum (GRUp i o)) instance (KnownNat i, KnownNat o) => Floating (GRUp i o)

makeLenses ''GRUp

instance (PrimMonad m, KnownNat i, KnownNat o) => Mutable m (GRUp i o) where
    type Ref m (GRUp i o) = GRef m (GRUp i o)
    thawRef   = gThawRef
    freezeRef = gFreezeRef
    copyRef   = gCopyRef
instance (KnownNat i, KnownNat o, Mutable m (GRUp i o)) => LinearInPlace m Double (GRUp i o)

instance (KnownNat i, KnownNat o, PrimMonad m) => Learnable m (GRUp i o)

-- | Stateless version of 'gru' that takes the "previous input" as a part
-- of the input vector.
gru'
    :: forall i o. (KnownNat i, KnownNat o)
    => Model ('Just (GRUp i o)) 'Nothing (R (i + o)) (R o)
gru' = modelStatelessD $ \(PJust p) x ->
    let z      = logistic $ runLRp (p ^^. gruMemory) x
        r      = logistic $ runLRp (p ^^. gruUpdate) x
        r'     = 1 # r
        h'     = tanh     $ runLRp (p ^^. gruOutput) (r' * x)
    in  (1 - z) * snd (split @i x) + z * h'

-- | Gated Recurrent Unit
--
-- <http://colah.github.io/posts/2015-08-Understanding-LSTMs/>
--
gru :: (KnownNat i, KnownNat o)
    => Model ('Just (GRUp i o)) ('Just (R o)) (R i) (R o)
gru = recurrent H.split (H.#) id gru'


module Main where

import           Control.Monad.Error (runErrorT)
import           Control.Monad.Mersenne.Random (evalRandom)
import           Numeric.IEEE (IEEE, epsilon)
import           System.Random.Mersenne.Pure64 (pureMT)
import qualified Data.Vector.Unboxed as V
import           Statistics.Test.ApproxRand
import           Test.HUnit (Assertion, assertBool, assertEqual)
import           Test.Framework
import           Test.Framework.Providers.HUnit

tests :: Test
tests = testGroup "Paired approximate randomization tests" $
  concat [statTests, randomizationTests]

main :: IO ()
main = defaultMain [ tests ]

-- Statistics tests

statTests :: [Test]
statTests = [meanDifferenceTest]

meanDifferenceTest :: Test
meanDifferenceTest =
  testIEEEEquality "mean difference robot competition"
    1.8 $ meanDifference cohenRobotsAlpha cohenRobotsBeta

-- Approximate andomization tests

randomizationTests :: [Test]
randomizationTests = [pairApproxExactTestScores]

pairApproxExactTestScores :: Test
pairApproxExactTestScores =
  testEquality  "number of extreme values robot competition"
    (Right 21) $ V.length `fmap` V.filter (>= 1.8) `fmap` scores
  where
    test = runErrorT $
      approxRandPairStats differenceMean 1024 cohenRobotsAlpha cohenRobotsBeta
    scores = evalRandom test $ pureMT 42

-- Helper functions

testEquality :: (Show a, Eq a) => String -> a -> a -> Test
testEquality msg a b = testCase msg $ assertEqual msg a b

testIEEEEquality :: IEEE a => String -> a -> a -> Test
testIEEEEquality msg a b = testCase msg $ assertEqualIEEE msg a b

assertEqualIEEE :: IEEE a => String -> a -> a -> Assertion
assertEqualIEEE msg a b = assertBool msg $ fracEq epsilon a b

-- Suggested in The Floating Point Guide: http://floating-point-gui.de/
fracEq :: (Fractional a, Ord a) => a -> a -> a -> Bool
fracEq eps a b
  | a == b     = True
  | a * b == 0 = diff < (eps * eps)
  | otherwise  = diff / (abs a + abs b) < eps
  where diff = abs (a - b)

-- Example from Cohen, 1995
cohenRobotsAlpha :: V.Vector Double
cohenRobotsAlpha = V.fromList [8,3,9,6,5,8,7,8,9,9]

cohenRobotsBeta :: V.Vector Double
cohenRobotsBeta  = V.fromList [7,0,9,4,5,9,8,3,4,5]


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
module Numeric.Layer where

import Dependent.Size
import Numeric.Vector.Sized
import Data.Kind(Type)

class Layer layer where
  type Inputs layer :: Size
  type Outputs layer :: Size
  type Tape layer :: Type
  type Gradient layer :: Type

  forward :: layer
          -> SizedArray (Inputs layer)
          -> (SizedArray (Outputs layer), Tape layer)

  backward :: layer
           -> Tape layer
           -> SizedArray (Outputs layer)
           -> (Gradient layer, SizedArray (Inputs layer))

  applyGradient :: layer -> Gradient layer -> layer


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE OverloadedLists     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

import Data.Maybe
import Data.Number.Symbolic
import qualified Data.Number.Symbolic as Sym
import Data.Proxy

import qualified Naperian as N
import qualified Data.Foldable as F
import           Control.Applicative ( liftA2 )
import qualified GHC.TypeLits as M
import           Data.Functor
import           Data.List.Split

import           Numeric.Sundials.ARKode.ODE
import           Numeric.LinearAlgebra

kx, ky :: Floating a => a
kx = 0.5
ky = 0.75

-- spatial mesh size
nx, ny :: Int
nx = 30
ny = 60

-- x mesh spacing
-- y mesh spacing
dx :: Floating a => a
dx = 1 / (fromIntegral nx - 1)

dy :: Floating a => a
dy = 1 / (fromIntegral ny - 1)

c1, c2 :: Floating a => a
c1 = kx/dx/dx
c2 = ky/dy/dy

cc4' :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
cc4' = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f m n i j k l | i == 0               = 0
                        | j == 0               = 0
                        | i == n - 1           = 0
                        | j == m - 1           = 0
                        | k == i - 1 && l == j = 1
                        | k == i     && l == j = -2
                        | k == i + 1 && l == j = 1
                        | otherwise            = 0

cc5' :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
cc5' = N.binary (*) (N.Scalar c2) cc4'

cc5Sym' :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
          N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
cc5Sym' = N.binary (*) (N.Scalar $ var "c2") cc4'

yy4' :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
yy4' = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f :: Int -> Int -> Int -> Int -> Int -> Int -> b
          f m n i j k l | i == 0                   = 0
                        | j == 0                   = 0
                        | i == n - 1               = 0
                        | j == m - 1               = 0
                        | k == i     && l == j - 1 = 1
                        | k == i     && l == j     = -2
                        | k == i     && l == j + 1 = 1
                        | otherwise                = 0

yy5' :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
yy5' = N.binary (*) (N.Scalar c1) yy4'

yy5Sym' :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
           N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
yy5Sym' = N.binary (*) (N.Scalar $ var "c1") yy4'

ccSym5 = cc5Sym' @Double @4 @5
yy5Sym = yy5Sym' @Double @4 @5

ccSym5

yy5Sym

fmap (N.elements . N.Prism . N.Prism . N.Scalar) $ N.elements $ N.crystal $ N.crystal $ N.binary (+) cc5Sym yy5Sym

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedLists       #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}

import Data.Maybe
import Data.Number.Symbolic
import Data.Proxy

import qualified Naperian as N
import qualified Data.Foldable as F
import           Control.Applicative ( liftA2 )
import qualified GHC.TypeLits as M

import           Numeric.Sundials.ARKode.ODE
import           Numeric.LinearAlgebra

x1, a, x2 :: Double
x1 = 0
a = 1.0
x2 = a

y1, y2 :: Double
y1 = 0.0
y2 = 1.0

bigT :: Double
bigT = 1000.0

n :: Int
n = 2

dx :: Double
dx = a / (fromIntegral n + 1)

dy :: Double
dy = a / (fromIntegral n + 1)

beta, s :: Double
beta = 1.0e-5
s = beta / (dx * dx)

kx, ky :: Double
kx = 0.5
ky = 0.75

c1, c2 :: Double
c1 = kx/dx/dx
c2 = ky/dy/dy

bigAA1 :: Matrix Double
bigAA1 = assoc (n * n, n * n) 0.0 [((i, j), f (i, j)) | i <- [0 .. n * n - 1]
                                                      , j <- [i - n, i,  i + n]
                                                      , j `elem` [0 .. n * n -1]]
  where
    f (i, j) | i     == j = (-2.0) * c1
             | i - n == j = 1.0    * c1
             | i + n == j = 1.0    * c1
             | otherwise = error $ show (i, j)

bigAA2 :: Matrix Double
bigAA2 = diagBlock (replicate n bigA)
  where
    bigA :: Matrix Double
    bigA = assoc (n, n) 0.0 [((i, j), f (i, j)) | i <- [0 .. n - 1]
                                                , j <- [i-1..i+1]
                                                , j `elem` [0..n-1]]
      where
        f (i, j) | i     == j = (-2.0) * c2
                 | i - 1 == j = 1.0    * c2
                 | i + 1 == j = 1.0    * c2

bigAA :: Matrix Double
bigAA = bigAA1 + bigAA2

bigAA1

bigAA2

n

bigZZ1 :: Matrix Double
bigZZ1 = assoc (m * m, m * m) 0.0 [((i, j), f (i, j)) | i <- [0 .. m * m - 1]
                                                      , j <- [0 .. m * m - 1]]
  where
    m = n + 2
    f (i, j) | i     == 0     = 0.0
             | j     == 0     = 0.0
             | i     == j     = (-2.0) * c1
             | i - n == j     = 1.0    * c1
             | i + n == j     = 1.0    * c1
             | i     == n + 1 = 0.0
             | j     == n + 1 = 0.0
             | otherwise      = 0.0

bigZZ1

x :: forall m n . (M.KnownNat m, M.KnownNat n) => N.Vector n (N.Vector m (Sym Int))
x = (fromJust . N.fromList) $
    map (fromJust . N.fromList) ([[var $ (\(x,y) -> "A" ++ show x ++ "," ++ show y) (x,y) | y <- [1..m]] | x <- [1..n]] :: [[Sym Int]])
    where
      m = M.natVal (undefined :: Proxy m)
      n = M.natVal (undefined :: Proxy n)

u1 :: N.Hyper '[N.Vector 3, N.Vector 2] (Sym Int)
u1 = N.Prism $ N.Prism (N.Scalar x)

u1

y :: forall n . M.KnownNat n => N.Vector n (Sym Int)
y = (fromJust . N.fromList) $
    (map (var . ("v" ++) . show) [1..n ] :: [Sym Int])
    where
    n = M.natVal (undefined :: Proxy n)

u2 :: N.Hyper '[N.Vector 3] (Sym Int)
u2 = N.Prism (N.Scalar y)

N.innerH u1 u2


{-# LANGUAGE Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Herm
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Hermitian views of matrices.
--

module Numeric.LinearAlgebra.Matrix.Herm (
    -- * Immutable interface

    -- ** Vector multiplication
    hermMulVector,
    hermMulVectorWithScale,
    addHermMulVectorWithScales,

    -- ** Matrix  multiplication
    hermMulMatrix,
    hermMulMatrixWithScale,
    addHermMulMatrixWithScales,

    -- ** Updates
    hermRank1Update,
    hermRank2Update,
    hermRankKUpdate,
    hermRank2KUpdate,

    -- * Mutable interface
    hermCreate,

    -- ** Vector multiplication
    hermMulVectorTo,
    hermMulVectorWithScaleTo,
    addHermMulVectorWithScalesM_,

    -- ** Matrix multiplication
    hermMulMatrixTo,
    hermMulMatrixWithScaleTo,
    addHermMulMatrixWithScalesM_,

    -- ** Updates
    hermRank1UpdateM_,
    hermRank2UpdateM_,
    hermRankKUpdateM_,
    hermRank2KUpdateM_,

    ) where

import Control.Monad( when )
import Control.Monad.ST( ST, runST, unsafeIOToST )
import Text.Printf( printf )

import Numeric.LinearAlgebra.Vector( Vector, RVector, STVector )
import qualified Numeric.LinearAlgebra.Vector as V
import Numeric.LinearAlgebra.Matrix.Base( Matrix )
import Numeric.LinearAlgebra.Matrix.STBase( STMatrix, RMatrix )
import qualified Numeric.LinearAlgebra.Matrix.STBase as M
import Numeric.LinearAlgebra.Types
import qualified Foreign.BLAS as BLAS

-- | A safe way to create and work with a mutable Herm Matrix before returning
-- an immutable one for later perusal.
hermCreate :: (Storable e)
           => (forall s. ST s (Herm (STMatrix s) e))
           -> Herm Matrix e
hermCreate mh = runST $ do
    (Herm u ma) <- mh
    a <- M.unsafeFreeze ma
    return $ Herm u a

-- | @hermRank1Update alpha x a@ returns
-- @alpha * x * x^H + a@.
hermRank1Update :: (BLAS2 e)
                => Double -> Vector e -> Herm Matrix e -> Herm Matrix e
hermRank1Update alpha x (Herm uplo a) = runST $ do
    ma' <- M.newCopy a
    hermRank1UpdateM_ alpha x (Herm uplo ma')
    a' <- M.unsafeFreeze ma'
    return $ Herm uplo a'

-- | @hermRank2Update alpha x y a@ returns
-- @alpha * x * y^H + conj(alpha) * y * x^H + a@.
hermRank2Update :: (BLAS2 e)
                => e -> Vector e -> Vector e -> Herm Matrix e
                -> Herm Matrix e
hermRank2Update alpha x y (Herm uplo a) = runST $ do
    ma' <- M.newCopy a
    hermRank2UpdateM_ alpha x y (Herm uplo ma')
    a' <- M.unsafeFreeze ma'
    return $ Herm uplo a'

-- | @hermRankKUpdate alpha trans a beta c@ returns
-- @c := alpha * a * a^H + beta * c@ when @trans@ is @NoTrans@ and
-- @c := alpha * a^H * a + beta * c@ when @trans@ is @ConjTrans@.  The
-- function signals an error when @trans@ is @Trans@.
hermRankKUpdate :: (BLAS3 e)
                => e -> Trans -> Matrix e -> e -> Herm Matrix e
                -> Herm Matrix e
hermRankKUpdate alpha trans a beta (Herm uplo c) = runST $ do
    mc' <- M.newCopy c
    hermRankKUpdateM_ alpha trans a beta (Herm uplo mc')
    c' <- M.unsafeFreeze mc'
    return $ Herm uplo c'

-- | @hermRank2KUpdate alpha trans a b beta c@ returns
-- @c := alpha * a * b^H + conj(alpha) * b * a^H + beta * c@ when @trans@ is
-- @NoTrans@ and @c := alpha * b^H * a + conj(alpha) * a^H * b + beta * c@
-- when @trans@ is @ConjTrans@.  The function signals an error when @trans@
-- is @Trans@.
hermRank2KUpdate :: (BLAS3 e)
                 => e -> Trans -> Matrix e -> Matrix e -> e -> Herm Matrix e
                 -> Herm Matrix e
hermRank2KUpdate alpha trans a b beta (Herm uplo c) = runST $ do
    mc' <- M.newCopy c
    hermRank2KUpdateM_ alpha trans a b beta (Herm uplo mc')
    c' <- M.unsafeFreeze mc'
    return $ Herm uplo c'

-- | @hermRank1UpdateM_ alpha x a@ sets
-- @a := alpha * x * x^H + a@.
hermRank1UpdateM_ :: (RVector v, BLAS2 e)
                  => Double -> v e -> Herm (STMatrix s) e -> ST s ()
hermRank1UpdateM_ alpha x (Herm uplo a) = do
    nx <- V.getDim x
    (ma,na) <- M.getDim a
    let n = nx

    when ((not . and) [ nx == n, (ma,na) == (n,n) ]) $ error $
        printf ("hermRank1UpdateM_ _ <vector with dim %d>"
                 ++ " (Herm _ <matrix with dim (%d,%d)>):"
                 ++ " invalid dimensions") nx ma na

    unsafeIOToST $
        V.unsafeWith x $ \px ->
        M.unsafeWith a $ \pa lda ->
            BLAS.her uplo n alpha px 1 pa lda

-- | @hermRank2UpdateM_ alpha x y a@ sets
-- @a := alpha * x * y^H + conj(alpha) * y * x^H + a@.
hermRank2UpdateM_ :: (RVector v1, RVector v2, BLAS2 e)
                  => e -> v1 e -> v2 e -> Herm (STMatrix s) e -> ST s ()
hermRank2UpdateM_ alpha x y (Herm uplo a) = do
    nx <- V.getDim x
    ny <- V.getDim y
    (ma,na) <- M.getDim a
    let n = nx

    when ((not . and) [ nx == n, ny == n, (ma,na) == (n,n) ]) $ error $
        printf ("hermRank2UpdateM_ _ <vector with dim %d>"
                 ++ " <vector with dim %d>"
                 ++ " (Herm _ <matrix with dim (%d,%d)>):"
                 ++ " invalid dimensions") nx ny ma na

    unsafeIOToST $
        V.unsafeWith x $ \px ->
        V.unsafeWith y $ \py ->
        M.unsafeWith a $ \pa lda ->
            BLAS.her2 uplo n alpha px 1 py 1 pa lda

-- | @hermRankKUpdateM_ alpha trans a beta c@ sets
-- @c := alpha * a * a^H + beta * c@ when @trans@ is @NoTrans@ and
-- @c := alpha * a^H * a + beta * c@ when @trans@ is @ConjTrans@.  The
-- function signals an error when @trans@ is @Trans@.
hermRankKUpdateM_ :: (RMatrix m, BLAS3 e)
                  => e -> Trans -> m e -> e -> Herm (STMatrix s) e
                  -> ST s ()
hermRankKUpdateM_ alpha trans a beta (Herm uplo c) = do
    (ma,na) <- M.getDim a
    (mc,nc) <- M.getDim c
    let (n,k) = if trans == NoTrans then (ma,na) else (na,ma)

    when (trans == Trans) $ error $
        printf ("hermRankKUpdateM_ _ %s:"
                 ++ " trans argument must be NoTrans or ConjTrans")
               (show trans)

    when ((not . and) [ (mc,nc) == (n,n)
                      , case trans of NoTrans -> (ma,na) == (n,k)
                                      _       -> (ma,na) == (k,n)
                      ]) $ error $
            printf ("hermRankKUpdateM_ _ %s <matrix with dim (%d,%d)> _"
                    ++ " (Herm _ <matrix with dim (%d,%d)>):"
                    ++ " invalid dimensions") (show trans) ma na mc nc

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith c $ \pc ldc ->
            BLAS.herk uplo trans n k alpha pa lda beta pc ldc

-- | @hermRank2KUpdateM_ alpha trans a b beta c@ sets
-- @c := alpha * a * b^H + conj(alpha) * b * a^H + beta * c@ when @trans@ is
-- @NoTrans@ and @c := alpha * b^H * a + conj(alpha) * a^H * b + beta * c@
-- when @trans@ is @ConjTrans@.  The function signals an error when @trans@
-- is @Trans@.
hermRank2KUpdateM_ :: (RMatrix m1, RMatrix m2, BLAS3 e)
                   => e -> Trans -> m1 e -> m2 e -> e -> Herm (STMatrix s) e
                   -> ST s ()
hermRank2KUpdateM_ alpha trans a b beta (Herm uplo c) = do
    (ma,na) <- M.getDim a
    (mb,nb) <- M.getDim b
    (mc,nc) <- M.getDim c
    let (n,k) = if trans == NoTrans then (ma,na) else (na,ma)

    when (trans == Trans) $ error $
        printf ("hermRank2KUpdateM_ _ %s:"
                 ++ " trans argument must be NoTrans or ConjTrans")
               (show trans)

    when ((not . and) [ (mc,nc) == (n,n)
                      , (mb,nb) == (ma,na)
                      , case trans of NoTrans -> (ma,na) == (n,k)
                                      _       -> (ma,na) == (k,n)
                      ]) $ error $
            printf ("hermRank2KUpdateM_ _ %s <matrix with dim (%d,%d)>"
                    ++ " <matrix with dim (%d,%d)> _"
                    ++ " (Herm _ <matrix with dim (%d,%d)>):"
                    ++ " invalid dimensions") (show trans) ma na mb nb mc nc

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith b $ \pb ldb ->
        M.unsafeWith c $ \pc ldc ->
            BLAS.her2k uplo trans n k alpha pa lda pb ldb beta pc ldc

-- | @hermMulVector a x@ returns @a * x@.
hermMulVector :: (BLAS2 e)
              => Herm Matrix e
              -> Vector e
              -> Vector e
hermMulVector a x =
    V.create $ do
        n <- V.getDim x
        y <- V.new_ n
        hermMulVectorTo y a x
        return y

-- | @hermMulVectorWithScale alpha a x@ retunrs @alpha * a * x@.
hermMulVectorWithScale :: (BLAS2 e)
                       => e
                       -> Herm Matrix e
                       -> Vector e
                       -> Vector e
hermMulVectorWithScale alpha a x =
    V.create $ do
        n <- V.getDim x
        y <- V.new_ n
        hermMulVectorWithScaleTo y alpha a x
        return y

-- | @addHermMulVectorWithScales alpha a x y@
-- returns @alpha * a * x + beta * y@.
addHermMulVectorWithScales :: (BLAS2 e)
                           => e
                           -> Herm Matrix e
                           -> Vector e
                           -> e
                           -> Vector e
                           -> Vector e
addHermMulVectorWithScales alpha a x beta y =
    V.create $ do
        y' <- V.newCopy y
        addHermMulVectorWithScalesM_ alpha a x beta y'
        return y'

-- | @hermMulMatrix side a b@
-- returns @alpha * a * b@ when @side@ is @LeftSide@ and
-- @alpha * b * a@ when @side@ is @RightSide@.
hermMulMatrix :: (BLAS3 e)
              => Side -> Herm Matrix e
              -> Matrix e
              -> Matrix e
hermMulMatrix side a b =
    M.create $ do
        mn <- M.getDim b
        c <- M.new_ mn
        hermMulMatrixTo c side a b
        return c

-- | @hermMulMatrixWithScale alpha side a b@
-- returns @alpha * a * b@ when @side@ is @LeftSide@ and
-- @alpha * b * a@ when @side@ is @RightSide@.
hermMulMatrixWithScale :: (BLAS3 e)
                       => e
                       -> Side -> Herm Matrix e
                       -> Matrix e
                       -> Matrix e
hermMulMatrixWithScale alpha side a b =
    M.create $ do
        mn <- M.getDim b
        c <- M.new_ mn
        hermMulMatrixWithScaleTo c alpha side a b
        return c

-- | @addHermMulMatrixWithScales alpha side a b beta c@
-- returns @alpha * a * b + beta * c@ when @side@ is @LeftSide@ and
-- @alpha * b * a + beta * c@ when @side@ is @RightSide@.
addHermMulMatrixWithScales :: (BLAS3 e)
                           => e
                           -> Side -> Herm Matrix e
                           -> Matrix e
                           -> e
                           -> Matrix e
                           -> Matrix e
addHermMulMatrixWithScales alpha side a b beta c =
    M.create $ do
        c' <- M.newCopy c
        addHermMulMatrixWithScalesM_ alpha side a b beta c'
        return c'

-- | @hermMulVectorTo dst a x@ sets @dst := a * x@.
hermMulVectorTo :: (RMatrix m, RVector v, BLAS2 e)
                => STVector s e
                -> Herm m e
                -> v e
                -> ST s ()
hermMulVectorTo dst = hermMulVectorWithScaleTo dst 1

-- | @hermMulVectorWithScaleTo dst alpha a x@
-- sets @dst := alpha * a * x@.
hermMulVectorWithScaleTo :: (RMatrix m, RVector v, BLAS2 e)
                         => STVector s e
                         -> e
                         -> Herm m e
                         -> v e
                         -> ST s ()
hermMulVectorWithScaleTo dst alpha a x =
    addHermMulVectorWithScalesM_ alpha a x 0 dst

-- | @addHermMulVectorWithScalesM_ alpha a x beta y@
-- sets @y := alpha * a * x + beta * y@.
addHermMulVectorWithScalesM_ :: (RMatrix m, RVector v, BLAS2 e)
                             => e
                             -> Herm m e
                             -> v e
                             -> e
                             -> STVector s e
                             -> ST s ()
addHermMulVectorWithScalesM_ alpha (Herm uplo a) x beta y = do
    (ma,na) <- M.getDim a
    nx <- V.getDim x
    ny <- V.getDim y
    let n = ny

    when (ma /= na) $ error $
        printf ("addHermMulVectorWithScalesM_ _"
                ++ " (Herm %s <matrix with dim (%d,%d)>)"
                ++ " %s <vector with dim %d>"
                ++ " _"
                ++ " <vector with dim %d>: Herm matrix is not square")
               (show uplo) ma na
               nx ny

    when ((not . and) [ (ma,na) == (n,n)
                      , nx == n
                      , ny == n
                      ]) $ error $
        printf ("addHermMulVectorWithScalesM_ _"
                ++ " (Herm %s <matrix with dim (%d,%d)>)"
                ++ " %s <vector with dim %d>"
                ++ " _"
                ++ " <vector with dim %d>: dimension mismatch")
               (show uplo) ma na
               nx ny

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        V.unsafeWith x $ \px ->
        V.unsafeWith y $ \py ->
            BLAS.hemv uplo n alpha pa lda px 1 beta py 1

-- | @hermMulMatrixTo dst side a b@
-- sets @dst := a * b@ when @side@ is @LeftSide@ and
-- @dst := b * a@ when @side@ is @RightSide@.
hermMulMatrixTo :: (RMatrix m1, RMatrix m2, BLAS3 e)
                => STMatrix s e
                -> Side -> Herm m1 e
                -> m2 e
                -> ST s ()
hermMulMatrixTo dst = hermMulMatrixWithScaleTo dst 1

-- | @hermMulMatrixWithScaleTo dst alpha side a b@
-- sets @dst := alpha * a * b@ when @side@ is @LeftSide@ and
-- @dst := alpha * b * a@ when @side@ is @RightSide@.
hermMulMatrixWithScaleTo :: (RMatrix m1, RMatrix m2, BLAS3 e)
                         => STMatrix s e
                         -> e
                         -> Side -> Herm m1 e
                         -> m2 e
                         -> ST s ()
hermMulMatrixWithScaleTo dst alpha side a b =
    addHermMulMatrixWithScalesM_ alpha side a b 0 dst

-- | @addHermMulMatrixWithScalesM_ alpha side a b beta c@
-- sets @c := alpha * a * b + beta * c@ when @side@ is @LeftSide@ and
-- @c := alpha * b * a + beta * c@ when @side@ is @RightSide@.
addHermMulMatrixWithScalesM_ :: (RMatrix m1, RMatrix m2, BLAS3 e)
                             => e
                             -> Side -> Herm m1 e
                             -> m2 e
                             -> e
                             -> STMatrix s e
                             -> ST s ()
addHermMulMatrixWithScalesM_ alpha side (Herm uplo a) b beta c = do
    (ma,na) <- M.getDim a
    (mb,nb) <- M.getDim b
    (mc,nc) <- M.getDim c
    let (m,n) = (mc,nc)

    when (ma /= na) $ error $
        printf ("addHermMulMatrixWithScalesM_ _"
                ++ " %s (Herm %s <matrix with dim (%d,%d)>)"
                ++ " <matrix with dim (%d,%d)>"
                ++ " _"
                ++ " <matrix with dim (%d,%d)>: Herm matrix is not square")
               (show side) (show uplo) ma na
               mb nb
               mc nc
    when ((not . and) [ case side of LeftSide  -> (ma,na) == (m,m)
                                     RightSide -> (ma,na) == (n,n)
                      , (mb, nb ) == (m,n)
                      , (mc, nc ) == (m,n)
                      ]) $ error $
        printf ("addHermMulMatrixWithScalesM_ _"
                ++ " %s (Herm %s <matrix with dim (%d,%d)>)"
                ++ " <matrix with dim (%d,%d)>"
                ++ " _"
                ++ " <matrix with dim (%d,%d)>: dimension mismatch")
               (show side) (show uplo) ma na
               mb nb
               mc nc

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith b $ \pb ldb ->
        M.unsafeWith c $ \pc ldc ->
            BLAS.hemm side uplo m n alpha pa lda pb ldb beta pc ldc


module Test.Grenade.Layers.Internal.Reference where

import           Numeric.LinearAlgebra

im2col :: Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
im2col nrows ncols srows scols m =
  let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
  in  im2colFit starts nrows ncols m

vid2col :: Int -> Int -> Int -> Int -> Int -> Int -> [Matrix Double] -> Matrix Double
vid2col nrows ncols srows scols inputrows inputcols ms =
  let starts = fittingStarts inputrows nrows srows inputcols ncols scols
      subs   = fmap (im2colFit starts nrows ncols) ms
  in  foldl1 (|||) subs

im2colFit :: [(Int,Int)] -> Int -> Int -> Matrix Double -> Matrix Double
im2colFit starts nrows ncols m =
  let imRows = fmap (\start -> flatten $ subMatrix start (nrows, ncols) m) starts
  in  fromRows imRows

col2vid :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> [Matrix Double]
col2vid nrows ncols srows scols drows dcols m =
  let starts = fittingStart (cols m) (nrows * ncols) (nrows * ncols)
      r      = rows m
      mats   = fmap (\s -> subMatrix (0,s) (r, nrows * ncols) m) starts
      colSts = fittingStarts drows nrows srows dcols ncols scols
  in  fmap (col2imfit colSts nrows ncols drows dcols) mats

col2im :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
col2im krows kcols srows scols drows dcols m =
  let starts     = fittingStarts drows krows srows dcols kcols scols
  in  col2imfit starts krows kcols drows dcols m

col2imfit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
col2imfit starts krows kcols drows dcols m =
  let indicies   = (\[a,b] -> (a,b)) <$> sequence [[0..(krows-1)], [0..(kcols-1)]]
      convs      = fmap (zip indicies . toList) . toRows $ m
      pairs      = zip convs starts
      accums     = concatMap (\(conv',(stx',sty')) -> fmap (\((ix,iy), val) -> ((ix + stx', iy + sty'), val)) conv') pairs
  in  accum (konst 0 (drows, dcols)) (+) accums

poolForward :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
poolForward nrows ncols srows scols outputRows outputCols m =
  let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
  in  poolForwardFit starts nrows ncols outputRows outputCols m

poolForwardList :: Functor f => Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> f (Matrix Double) -> f (Matrix Double)
poolForwardList nrows ncols srows scols inRows inCols outputRows outputCols ms =
  let starts = fittingStarts inRows nrows srows inCols ncols scols
  in  poolForwardFit starts nrows ncols outputRows outputCols <$> ms

poolForwardFit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
poolForwardFit starts nrows ncols _ outputCols m =
  let els    = fmap (\start -> maxElement $ subMatrix start (nrows, ncols) m) starts
  in  matrix outputCols els

poolBackward :: Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double -> Matrix Double
poolBackward krows kcols srows scols inputMatrix gradientMatrix =
  let inRows     = rows inputMatrix
      inCols     = cols inputMatrix
      starts     = fittingStarts inRows krows srows inCols kcols scols
  in  poolBackwardFit starts krows kcols inputMatrix gradientMatrix

poolBackwardList :: Functor f => Int -> Int -> Int -> Int -> Int -> Int -> f (Matrix Double, Matrix Double) -> f (Matrix Double)
poolBackwardList krows kcols srows scols inRows inCols inputMatrices =
  let starts     = fittingStarts inRows krows srows inCols kcols scols
  in  uncurry (poolBackwardFit starts krows kcols) <$> inputMatrices

poolBackwardFit :: [(Int,Int)] -> Int -> Int -> Matrix Double -> Matrix Double -> Matrix Double
poolBackwardFit starts krows kcols inputMatrix gradientMatrix =
  let inRows     = rows inputMatrix
      inCols     = cols inputMatrix
      inds       = fmap (\start -> maxIndex $ subMatrix start (krows, kcols) inputMatrix) starts
      grads      = toList $ flatten gradientMatrix
      grads'     = zip3 starts grads inds
      accums     = fmap (\((stx',sty'),grad,(inx, iny)) -> ((stx' + inx, sty' + iny), grad)) grads'
  in  accum (konst 0 (inRows, inCols)) (+) accums

-- | These functions are not even remotely safe, but it's only called from the statically typed
--   commands, so we should be good ?!?!?
--   Returns the starting sub matrix locations which fit inside the larger matrix for the
--   convolution. Takes into account the stride and kernel size.
fittingStarts :: Int -> Int -> Int -> Int -> Int -> Int -> [(Int,Int)]
fittingStarts nrows kernelrows steprows ncols kernelcols stepcolsh =
  let rs = fittingStart nrows kernelrows steprows
      cs = fittingStart ncols kernelcols stepcolsh
      ls = sequence [rs, cs]
  in  fmap (\[a,b] -> (a,b)) ls

-- | Returns the starting sub vector which fit inside the larger vector for the
--   convolution. Takes into account the stride and kernel size.
fittingStart :: Int -> Int -> Int -> [Int]
fittingStart width kernel steps =
  let go left | left + kernel < width
              = left : go (left + steps)
              | left + kernel == width
              = [left]
              | otherwise
              = []
  in go 0


{-# LANGUAGE FlexibleContexts #-}

module STC.OrientationScaleAnalysis where

import           Control.Monad.Parallel    as MP
import           Data.Array.Repa           as R
import           Data.Complex
import           Data.List                 as L
import           Data.Vector.Unboxed       as VU
import           DFT.Plan
import           FokkerPlanck.DomainChange (r2z1Tor2s1)
import           Graphics.Gnuplot.Simple
import           Image.IO                  (ImageRepa (..), plotImageRepa,
                                            plotImageRepaComplex)
import           STC.Convolution
import           System.FilePath
import           Types
import           Utils.Array
import           Utils.Parallel
import Text.Printf

-- {-# INLINE analyzeOrientation #-}
-- analyzeOrientation :: Int -> [Double] -> R2T0Array -> (R2T0Array, R2T0Array)
-- analyzeOrientation numOrientation thetaFreqs arr =
--   let arrR2S1 = R.map magnitude . r2z1Tor2s1 numOrientation thetaFreqs $ arr
--       deltaTheta = 2 * pi / fromIntegral numOrientation :: Double
--       (Z :. _ :. cols :. rows) = extent arrR2S1
--       orientationArr =
--         fromListUnboxed (Z :. cols :. rows) .
--         parMap
--           rdeepseq
--           (\(i, j) ->
--              deltaTheta *
--              (fromIntegral .
--               VU.maxIndex . toUnboxed . computeS . R.slice arrR2S1 $
--               (Z :. All :. i :. j))) $
--         [(i, j) | i <- [0 .. cols - 1], j <- [0 .. rows - 1]]
--       freqArr = fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs
--       magArr =
--         R.map sqrt . sumS . R.map (\x -> (magnitude x) ^ 2) . rotate3D $ arr
--       func theta =
--         computeS .
--         R.traverse3 orientationArr freqArr magArr (\_ _ _ -> extent arr) $ \fOri fFreq fMag (Z :. k :. i :. j) ->
--           (fMag (Z :. i :. j) :+ 0) *
--           (exp $ 0 :+ (1) * fFreq (Z :. k) * (fOri (Z :. i :. j) + theta))
--    in (func 0, func 0)

-- {-# INLINE analyzeOrientationR2Z1T0 #-}
-- analyzeOrientationR2Z1T0 ::
--      Int -> [Double] -> [Double] -> R2T0Array -> (R2Z1T0Array, R2Z1T0Array)
-- analyzeOrientationR2Z1T0 numOrientation thetaFreqs theta0Freqs arr =
--   let arrR2S1 = R.map magnitude . r2z1Tor2s1 numOrientation thetaFreqs $ arr
--       deltaTheta = 2 * pi / fromIntegral numOrientation :: Double
--       (Z :. _ :. cols :. rows) = extent arrR2S1
--       orientationArr =
--         fromListUnboxed (Z :. cols :. rows) .
--         parMap
--           rdeepseq
--           (\(i, j) ->
--              deltaTheta *
--              (fromIntegral .
--               VU.maxIndex . toUnboxed . computeS . R.slice arrR2S1 $
--               (Z :. All :. i :. j))) $
--         [(i, j) | i <- [0 .. cols - 1], j <- [0 .. rows - 1]]
--       freqArr = fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs
--       freq0Arr = fromListUnboxed (Z :. (L.length theta0Freqs)) theta0Freqs
--       magArr =
--         R.map sqrt . sumS . R.map (\x -> (magnitude x) ^ 2) . rotate3D $ arr
--       func theta =
--         computeS .
--         R.traverse4
--           orientationArr
--           freqArr
--           freq0Arr
--           magArr
--           (\_ _ _ _ ->
--              (Z :. (L.length thetaFreqs) :. (L.length theta0Freqs) :. cols :.
--               rows)) $ \fOri fFreq fFreq0 fMag (Z :. k :. l :. i :. j) ->
--           (fMag (Z :. i :. j) :+ 0) *
--           (exp $ 0 :+ (-1) * (fFreq (Z :. k) + fFreq0 (Z :. l)) * (fOri (Z :. i :. j) + theta))
--    in (func (pi / 2),func 0)

-- {-# INLINE normalizeList #-}
-- normalizeList :: (Ord e, Fractional e) => [e] -> [e]
-- normalizeList xs = L.map (/ L.maximum xs) xs

-- plotMagnitudeOrientation ::
--      (R.Source s (Complex Double))
--   => FilePath
--   -> Int
--   -> [Double]
--   -> R.Array s DIM3 (Complex Double)
--   -> (Int, Int)
--   -> IO (Int,Int)
-- plotMagnitudeOrientation folderPath numOrientationSample thetaFreqs arr (i', j') = do
--   let orientationSampleRad =
--         [ 2 * pi / fromIntegral numOrientationSample * fromIntegral i
--         | i <- [0 .. numOrientationSample - 1]
--         ]
--       orientationSampleDeg =
--         [ 360 * fromIntegral i / fromIntegral numOrientationSample
--         | i <- [0 .. numOrientationSample - 1]
--         ]
--       freqArr = fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs
--       xsFreqDomain
--         -- normalizeList $
--        =
--         parMap
--           rdeepseq
--           (\theta ->
--              magnitude .
--              R.sumAllS .
--              R.zipWith (\freq x -> x * exp (0 :+ theta * freq)) freqArr .
--              R.slice arr $
--              (Z :. All :. i :. j))
--           orientationSampleRad
--       xs =
--         normalizeList $
--         R.toList .
--         R.map magnitude .
--         r2z1Tor2s1 numOrientationSample thetaFreqs .
--         extend (Z :. All :. (1 :: Int) :. (1 :: Int)) . R.slice arr $
--         (Z :. All :. i :. j)
--       (Z :. _ :. cols :. rows) = extent arr
--       magVec =
--         VU.concat $
--         parMap
--           rdeepseq
--           (\theta ->
--              toUnboxed .
--              computeS .
--              R.map magnitude . R.sumS . rotate3D . R.traverse2 arr freqArr const $ \f1 f2 idx@(Z :. k :. _ :. _) ->
--                f1 idx * exp (0 :+ theta * f2 (Z :. k)))
--           orientationSampleRad
--       (Z :. c :. a :. b) =
--         fromIndex (Z :. (L.length thetaFreqs) :. cols :. rows) . VU.maxIndex $
--         magVec
--       maxMag = VU.maximum magVec
--       (i,j) = (a,b)
--   printf
--     "Max magnitude: %0.5f at (%d,%d) %f degree.\n"
--     maxMag
--     a
--     b
--     ((fromIntegral c :: Double) / (fromIntegral numOrientationSample) * 360)
--   plotPathsStyle
--     [ PNG (folderPath </> "Magnitude.png")
--     , Title ("Magnitude at " L.++ show (i, j))
--     ] $
--     L.zip
--         -- defaultStyle
--       --     { plotType = LinesPoints
--       --     , lineSpec = CustomStyle [LineTitle "Spatial Domain", PointType 1]
--       --     }
--       -- ,
--       [ defaultStyle
--           { plotType = Lines
--           , lineSpec =
--               CustomStyle
--                 [ LineTitle "Frequency Domain" -- , PointType 0
--                 ]
--           }
--       ]
--        -- L.zip orientationSampleDeg xs,
--       [L.zip orientationSampleDeg xsFreqDomain]
--   return (a,b)

-- plotMagnitudeOrientationSource ::
--      (R.Source s (Complex Double))
--   => DFTPlan
--   -> FilePath
--   -> Int
--   -> Int
--   -> [Double]
--   -> R2Z1T0Array
--   -> R.Array s DIM3 (Complex Double)
--   -> (Int, Int)
--   -> IO ()
-- plotMagnitudeOrientationSource plan folderPath numOrientationSample numOrientation thetaFreqs filter input (i, j) = do
--   let orientationSampleRad =
--         [ 2 * pi / fromIntegral numOrientationSample * fromIntegral i
--         | i <- [0 .. numOrientationSample - 1]
--         ]
--       orientationSampleDeg =
--         [ 360 * fromIntegral i / fromIntegral numOrientationSample
--         | i <- [0 .. numOrientationSample - 1]
--         ]
--       freqArr = fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs
--       (Z :. numThetaFreq :. cols :. rows) = extent input
--   xs <-
--     MP.mapM
--       (\theta -> do
--          initialDistF <-
--            fmap (fromUnboxed (Z :. numThetaFreq :. cols :. rows) . VU.convert) .
--            dftExecute plan (DFTPlanID DFT1DG [numThetaFreq, cols, rows] [1, 2]) .
--            VU.convert . toUnboxed . computeS . R.traverse2 input freqArr const $
--            (\f1 f2 idx@(Z :. k :. i :. j) ->
--               f1 idx * exp (0 :+ theta * (1) * f2 (Z :. k)))
--          sourceArr <- convolveR2T0 plan filter initialDistF
--          let sourceR2Z1 = R.sumS . rotateR2Z1T0Array $ sourceArr
--              mag =
--                magnitude .
--                R.sumAllS .
--                r2z1Tor2s1 numOrientation thetaFreqs .
--                extend (Z :. All :. (1 :: Int) :. (1 :: Int)) .
--                R.slice sourceR2Z1 $
--                (Z :. All :. i :. j)
--          return mag)
--       orientationSampleRad
--   plotPath
--     [ PNG (folderPath </> "SourceMagnitude.png")
--     , Title ("Source Magnitude at " L.++ show (i, j))
--     ] .
--     L.zip orientationSampleDeg . normalizeList $
--     xs