{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Gelu
Description : Gaussian Error Linear Unit (GELU)
Copyright   : (c) Manuel Schneckenreither, 2020
License     : BSD2
Stability   : experimental

This module implements the Gaussian Error Linear Unit (GELU) activiation function. See

Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).

Available at: https://arxiv.org/pdf/1606.08415.pdf

As in the paper we simply let μ = 0 and σ = 1. Futher, we use the simplified and thus fast representation: x * sigmoid (1.702 * x)

-}
module Grenade.Layers.Gelu (
    Gelu (..)
  , SpecGelu (..)
  , specGelu1D
  , specGelu2D
  , specGelu3D
  , gelu
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Constraint                (Dict (..))
import           Data.Reflection                (reifyNat)
import           Data.Serialize
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS
import           Unsafe.Coerce                  (unsafeCoerce)

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build

-- | A Gaussion Error Linear Unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
--
--   Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).
data Gelu = Gelu
  deriving (Generic, NFData, Show)

instance UpdateLayer Gelu where
  type Gradient Gelu = ()
  runUpdate _ _ _ = Gelu

instance RandomLayer Gelu where
  createRandomWith _ _ = return Gelu

instance Serialize Gelu where
  put _ = return ()
  get = return Gelu

geluForwardFast :: Floating x => x -> x
geluForwardFast x = x / (e ** (-1.702 * x) + 1) -- = x * sigmoid (1.702 * x)
  where
    e = 2.71828

geluBackwardFast :: Floating x => x -> x
geluBackwardFast x = (e ** (1.702 * x) * (1 + e ** (1.702 * x) + 1.702 * x)) / (1 + e ** (1.702 * x)) ** 2
  where
    e = 2.71828

instance (KnownNat i) => Layer Gelu ('D1 i) ('D1 i) where
  type Tape Gelu ('D1 i) ('D1 i) = S ('D1 i)

  runForwards _ (S1D y) = (S1D y, S1D (gelu y))
    where
      gelu = LAS.dvmap geluForwardFast
  runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (gelu' y * dEdy))
    where
      gelu' = LAS.dvmap geluBackwardFast

instance (KnownNat i, KnownNat j) => Layer Gelu ('D2 i j) ('D2 i j) where
  type Tape Gelu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (gelu y))
    where
      gelu = LAS.dmmap geluForwardFast
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (gelu' y * dEdy))
    where
      gelu' = LAS.dmmap geluBackwardFast

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Gelu ('D3 i j k) ('D3 i j k) where

  type Tape Gelu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (gelu y))
    where
      gelu = LAS.dmmap geluForwardFast
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (gelu' y * dEdy))
    where
      gelu' = LAS.dmmap geluBackwardFast

-------------------- DynamicNetwork instance --------------------

instance FromDynamicLayer Gelu where
  fromDynamicLayer inp _ _ = SpecNetLayer $ SpecGelu (tripleFromSomeShape inp)

instance ToDynamicLayer SpecGelu where
  toDynamicLayer _ _ (SpecGelu (rows, cols, depth)) =
     reifyNat rows $ \(_ :: (KnownNat rows) => Proxy rows) ->
     reifyNat cols $ \(_ :: (KnownNat cols) => Proxy cols) ->
     reifyNat depth $ \(_ :: (KnownNat depth) => Proxy depth) ->
     case (rows, cols, depth) of
         (_, 1, 1)    -> return $ SpecLayer Gelu (sing :: Sing ('D1 rows)) (sing :: Sing ('D1 rows))
         (_, _, 1) -> return $ SpecLayer Gelu (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 rows cols))
         _    -> case (unsafeCoerce (Dict :: Dict()) :: Dict (KnownNat (rows GHC.TypeLits.* depth))) of
           Dict -> return $ SpecLayer Gelu (sing :: Sing ('D3 rows cols depth)) (sing :: Sing ('D3 rows cols depth))

-- | Create a specification for a elu layer.
specGelu1D :: Integer -> SpecNet
specGelu1D i = specGelu3D (i, 1, 1)

-- | Create a specification for a elu layer.
specGelu2D :: (Integer, Integer) -> SpecNet
specGelu2D (i, j) = specGelu3D (i, j, 1)

-- | Create a specification for a elu layer.
specGelu3D :: (Integer, Integer, Integer) -> SpecNet
specGelu3D = SpecNetLayer . SpecGelu

-- | Add a Gelu layer to your build.
gelu :: BuildM ()
gelu = buildGetLastLayerOut >>= buildAddSpec . SpecNetLayer . SpecGelu

-------------------- GNum instances --------------------

instance GNum Gelu where
  _ |* Gelu = Gelu
  _ |+ Gelu = Gelu


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.ChiSquared
-- Copyright : (c) 2010 Alexey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The chi-squared distribution. This is a continuous probability
-- distribution of sum of squares of k independent standard normal
-- distributions. It's commonly used in statistical tests
module Statistics.Distribution.ChiSquared (
          ChiSquared
        , chiSquaredNDF
        -- * Constructors
        , chiSquared
        , chiSquaredE
        ) where

import Control.Applicative
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.SpecFunctions ( incompleteGamma,invIncompleteGamma,logGamma,digamma)
import Numeric.MathFunctions.Constants (m_neg_inf)
import qualified System.Random.MWC.Distributions as MWC

import qualified Statistics.Distribution         as D
import Statistics.Internal

-- | Chi-squared distribution
newtype ChiSquared = ChiSquared
  { chiSquaredNDF :: Int
    -- ^ Get number of degrees of freedom
  }
  deriving (Eq, Typeable, Data, Generic)

instance Show ChiSquared where
  showsPrec i (ChiSquared n) = defaultShow1 "chiSquared" n i
instance Read ChiSquared where
  readPrec = defaultReadPrecM1 "chiSquared" chiSquaredE

-- | Construct chi-squared distribution. Number of degrees of freedom
--   must be positive.
chiSquared :: Int -> ChiSquared
chiSquared n = maybe (error $ errMsg n) id $ chiSquaredE n

-- | Construct chi-squared distribution. Number of degrees of freedom
--   must be positive.
chiSquaredE :: Int -> Maybe ChiSquared
chiSquaredE n
  | n <= 0    = Nothing
  | otherwise = Just (ChiSquared n)

errMsg :: Int -> String
errMsg n = "Statistics.Distribution.ChiSquared.chiSquared: N.D.F. must be positive. Got " ++ show n

instance D.Distribution ChiSquared where
  cumulative = cumulative

instance D.ContDistr ChiSquared where
  density chi x
    | x <= 0    = 0
    | otherwise = exp $ log x * (ndf2 - 1) - x2 - logGamma ndf2 - log 2 * ndf2
    where
      ndf  = fromIntegral $ chiSquaredNDF chi
      ndf2 = ndf/2
      x2   = x/2

  logDensity chi x
    | x <= 0    = m_neg_inf
    | otherwise = log x * (ndf2 - 1) - x2 - logGamma ndf2 - log 2 * ndf2
    where
      ndf  = fromIntegral $ chiSquaredNDF chi
      ndf2 = ndf/2
      x2   = x/2

  quantile = quantile

instance D.Mean ChiSquared where
    mean (ChiSquared ndf) = fromIntegral ndf

instance D.Variance ChiSquared where
    variance (ChiSquared ndf) = fromIntegral (2*ndf)

instance D.MaybeMean ChiSquared where
    maybeMean = Just . D.mean

instance D.MaybeVariance ChiSquared where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy ChiSquared where
  entropy (ChiSquared ndf) =
    let kHalf = 0.5 * fromIntegral ndf in
    kHalf
    + log 2
    + logGamma kHalf
    + (1-kHalf) * digamma kHalf

instance D.MaybeEntropy ChiSquared where
  maybeEntropy = Just . D.entropy

instance D.ContGen ChiSquared where
    genContVar (ChiSquared n) = MWC.chiSquare n

cumulative :: ChiSquared -> Double -> Double
cumulative chi x
  | x <= 0    = 0
  | otherwise = incompleteGamma (ndf/2) (x/2)
  where
    ndf = fromIntegral $ chiSquaredNDF chi

quantile :: ChiSquared -> Double -> Double
quantile (ChiSquared ndf) p
  | p == 0         = 0
  | p == 1         = 1/0
  | p > 0 && p < 1 = 2 * invIncompleteGamma (fromIntegral ndf / 2) p
  | otherwise      =
    error $ "Statistics.Distribution.ChiSquared.quantile: p must be in [0,1] range. Got: "++show p


{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}

#ifndef MIN_VERSION_hashable
#define MIN_VERSION_hashable(x,y,z) 1
#endif

-----------------------------------------------------------------------------
-- |
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Operations on affine spaces.
-----------------------------------------------------------------------------
module Linear.Affine where

import Control.Applicative
import Control.DeepSeq
import Control.Monad (liftM)
import Control.Lens
import Data.Binary as Binary
import Data.Bytes.Serial
import Data.Coerce
import Data.Complex (Complex)
import Data.Data
import Data.Distributive
import Data.Foldable as Foldable
import Data.Functor.Bind
import Data.Functor.Classes
import Data.Functor.Product
import Data.Functor.Rep as Rep
import Data.HashMap.Lazy (HashMap)
import Data.Hashable
import Data.Hashable.Lifted
import Data.IntMap (IntMap)
import Data.Ix
import Data.Kind
import Data.Map (Map)
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup (Semigroup)
#endif
import Data.Serialize as Cereal
import Data.Vector (Vector)
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed.Base as U
import Foreign.Storable
import GHC.Generics (Generic, Generic1)
import Linear.Epsilon
import Linear.Metric
import Linear.Plucker
import Linear.Quaternion
import Linear.V
import Linear.V0
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4
import Linear.Vector
import System.Random (Random(..))

-- | An affine space is roughly a vector space in which we have
-- forgotten or at least pretend to have forgotten the origin.
--
-- > a .+^ (b .-. a)  =  b@
-- > (a .+^ u) .+^ v  =  a .+^ (u ^+^ v)@
-- > (a .-. b) ^+^ v  =  (a .+^ v) .-. q@
class Additive (Diff p) => Affine p where
  type Diff p :: Type -> Type

  infixl 6 .-.
  -- | Get the difference between two points as a vector offset.
  (.-.) :: Num a => p a -> p a -> Diff p a

  infixl 6 .+^
  -- | Add a vector offset to a point.
  (.+^) :: Num a => p a -> Diff p a -> p a

  infixl 6 .-^
  -- | Subtract a vector offset from a point.
  (.-^) :: Num a => p a -> Diff p a -> p a
  p .-^ v = p .+^ negated v
  {-# INLINE (.-^) #-}

instance (Affine f, Affine g) => Affine (Product f g) where
  type Diff (Product f g) = Product (Diff f) (Diff g)
  Pair a b .-. Pair c d = Pair (a .-. c) (b .-. d)
  Pair a b .+^ Pair c d = Pair (a .+^ c) (b .+^ d)
  Pair a b .-^ Pair c d = Pair (a .+^ c) (b .+^ d)

-- | Compute the quadrance of the difference (the square of the distance)
qdA :: (Affine p, Foldable (Diff p), Num a) => p a -> p a -> a
qdA a b = Foldable.sum (fmap (join (*)) (a .-. b))
{-# INLINE qdA #-}

-- | Distance between two points in an affine space
distanceA :: (Floating a, Foldable (Diff p), Affine p) => p a -> p a -> a
distanceA a b = sqrt (qdA a b)
{-# INLINE distanceA #-}

#define ADDITIVEC(CTX,T) instance CTX => Affine T where type Diff T = T ; \
  (.-.) = (^-^) ; {-# INLINE (.-.) #-} ; (.+^) = (^+^) ; {-# INLINE (.+^) #-} ; \
  (.-^) = (^-^) ; {-# INLINE (.-^) #-}
#define ADDITIVE(T) ADDITIVEC((), T)

ADDITIVE([])
ADDITIVE(Complex)
ADDITIVE(ZipList)
ADDITIVE(Maybe)
ADDITIVE(IntMap)
ADDITIVE(Identity)
ADDITIVE(Vector)
ADDITIVE(V0)
ADDITIVE(V1)
ADDITIVE(V2)
ADDITIVE(V3)
ADDITIVE(V4)
ADDITIVE(Plucker)
ADDITIVE(Quaternion)
ADDITIVE(((->) b))
ADDITIVEC(Ord k, (Map k))
ADDITIVEC((Eq k, Hashable k), (HashMap k))
ADDITIVEC(Dim n, (V n))

-- | A handy wrapper to help distinguish points from vectors at the
-- type level
newtype Point f a = P (f a)
  deriving ( Eq, Ord, Show, Read, Monad, Functor, Applicative, Foldable
           , Eq1, Ord1, Show1, Read1
           , Traversable, Apply, Additive, Metric
           , Fractional , Num, Ix, Storable, Epsilon
           , Semigroup, Monoid
           , Random, Hashable
           , Generic, Generic1, Data
           )

instance Finite f => Finite (Point f) where
  type Size (Point f) = Size f
  toV (P v) = toV v
  fromV v = P (fromV v)

instance NFData (f a) => NFData (Point f a) where
  rnf (P x) = rnf x

instance Serial1 f => Serial1 (Point f) where
  serializeWith f (P p) = serializeWith f p
  deserializeWith m = P `liftM` deserializeWith m

instance Serial (f a) => Serial (Point f a) where
  serialize (P p) = serialize p
  deserialize = P `liftM` deserialize

instance Binary (f a) => Binary (Point f a) where
  put (P p) = Binary.put p
  get = P `liftM` Binary.get

instance Serialize (f a) => Serialize (Point f a) where
  put (P p) = Cereal.put p
  get = P `liftM` Cereal.get

instance Hashable1 f => Hashable1 (Point f) where
  liftHashWithSalt h s (P f) = liftHashWithSalt h s f
  {-# INLINE liftHashWithSalt #-}

lensP :: Lens' (Point g a) (g a)
lensP afb (P a) = P <$> afb a
{-# INLINE lensP #-}

_Point :: Iso' (Point f a) (f a)
_Point = iso (\(P a) -> a) P
{-# INLINE _Point #-}

instance (t ~ Point g b) => Rewrapped (Point f a) t
instance Wrapped (Point f a) where
  type Unwrapped (Point f a) = f a
  _Wrapped' = _Point
  {-# INLINE _Wrapped' #-}

-- These are stolen from Data.Profunctor.Unsafe
(.#) :: Coercible b a => (b -> c) -> (a -> b) -> a -> c
f .# _ = coerce f
{-# INLINE (.#) #-}

(#.) :: Coercible c b => (b -> c) -> (a -> b) -> a -> c
(#.) _ = coerce (\x -> x :: b) :: forall a b. Coercible b a => a -> b
{-# INLINE (#.) #-}

unP :: Point f a -> f a
unP (P x) = x
{-# INLINE unP #-}

-- We can't use GND to derive 'Bind' because 'join' causes
-- role troubles. However, GHC 7.8 and above let us use
-- explicit coercions for (>>-).
instance Bind f => Bind (Point f) where
  (>>-) = ((P .) . (. (unP .))) #. (>>-) .# unP
  join (P m) = P $ m >>- \(P m') -> m'

instance Distributive f => Distributive (Point f) where
  distribute = P . collect (\(P p) -> p)
  collect = (P .) #. collect .# (unP .)

instance Representable f => Representable (Point f) where
  type Rep (Point f) = Rep f
  tabulate = P #. tabulate
  {-# INLINE tabulate #-}
  index = Rep.index .# unP
  {-# INLINE index #-}

type instance Index (Point f a) = Index (f a)
type instance IxValue (Point f a) = IxValue (f a)

instance Ixed (f a) => Ixed (Point f a) where
  ix l = lensP . ix l
  {-# INLINE ix #-}

instance Traversable f => Each (Point f a) (Point f b) a b where
  each = traverse
  {-# INLINE each #-}

instance R1 f => R1 (Point f) where
  _x = lensP . _x
  {-# INLINE _x #-}

instance R2 f => R2 (Point f) where
  _y = lensP . _y
  {-# INLINE _y #-}
  _xy = lensP . _xy
  {-# INLINE _xy #-}

instance R3 f => R3 (Point f) where
  _z = lensP . _z
  {-# INLINE _z #-}
  _xyz = lensP . _xyz
  {-# INLINE _xyz #-}

instance R4 f => R4 (Point f) where
  _w = lensP . _w
  {-# INLINE _w #-}
  _xyzw = lensP . _xyzw
  {-# INLINE _xyzw #-}

instance Additive f => Affine (Point f) where
  type Diff (Point f) = f
  (.-.) = (. unP) #. (^-^) .# unP
  {-# INLINE (.-.) #-}
  (.+^) = (P .) #. (^+^) .# unP
  {-# INLINE (.+^) #-}
  (.-^) = (P .) #. (^-^) .# unP
  {-# INLINE (.-^) #-}

-- | Vector spaces have origins.
origin :: (Additive f, Num a) => Point f a
origin = P zero

-- | An isomorphism between points and vectors, given a reference
--   point.
relative :: (Additive f, Num a) => Point f a -> Iso' (Point f a) (f a)
relative p0 = iso (.-. p0) (p0 .+^)
{-# INLINE relative #-}

newtype instance U.Vector    (Point f a) =  V_P (U.Vector    (f a))
newtype instance U.MVector s (Point f a) = MV_P (U.MVector s (f a))
instance U.Unbox (f a) => U.Unbox (Point f a)

instance U.Unbox (f a) => M.MVector U.MVector (Point f a) where
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  basicLength (MV_P v) = M.basicLength v
  basicUnsafeSlice m n (MV_P v) = MV_P (M.basicUnsafeSlice m n v)
  basicOverlaps (MV_P v) (MV_P u) = M.basicOverlaps v u
  basicUnsafeNew n = MV_P `liftM` M.basicUnsafeNew n
  basicUnsafeRead (MV_P v) i = P `liftM` M.basicUnsafeRead v i
  basicUnsafeWrite (MV_P v) i (P x) = M.basicUnsafeWrite v i x
  basicInitialize (MV_P v) = M.basicInitialize v
  {-# INLINE basicInitialize #-}

instance U.Unbox (f a) => G.Vector U.Vector (Point f a) where
  {-# INLINE basicUnsafeFreeze #-}
  {-# INLINE basicUnsafeThaw   #-}
  {-# INLINE basicLength       #-}
  {-# INLINE basicUnsafeSlice  #-}
  {-# INLINE basicUnsafeIndexM #-}
  basicUnsafeFreeze (MV_P v) = V_P `liftM` G.basicUnsafeFreeze v
  basicUnsafeThaw   ( V_P v) = MV_P `liftM` G.basicUnsafeThaw   v
  basicLength       ( V_P v) = G.basicLength v
  basicUnsafeSlice m n (V_P v) = V_P (G.basicUnsafeSlice m n v)
  basicUnsafeIndexM (V_P v) i = P `liftM` G.basicUnsafeIndexM v i


{-# LANGUAGE Arrows, ExistentialQuantification, ScopedTypeVariables #-}

module SignalLevel.SpecAnalysis20 where

import Data.Complex
import Data.Maybe (catMaybes,listToMaybe)
import qualified Data.Map as M
import Control.Arrow.ArrowP
import Control.DeepSeq (NFData)
import Control.SF.SF (SF(..))
import FRP.UISF (SEvent,asyncVT,Time,DeltaT)
import FRP.UISF.Asynchrony (Automaton(..))
import FRP.UISF.Graphics.Color (rgbE)

import Euterpea hiding (SF)
import HSoM
import HSoM.MUI (asyncV)
import HSoM.Examples.FFT (fftA,quantize,presentFFT)

import SignalLevel.SigFun19

dft :: RealFloat a => [Complex a] -> [Complex a]
dft xs =
  let lenI = length xs
      lenR = fromIntegral lenI
      lenC = lenR :+ 0
  in [let i = -2*pi * fromIntegral k / lenR
      in (1/lenC) * sum [(xs !! n) * exp (0 :+ i * fromIntegral n)
                          | n <- [0,1..lenI-1]]
       | k <- [0,1..lenI-1]]

-- num = # samples
-- n   = term

mkTerm :: Int -> Double -> [Complex Double]
mkTerm num n = let f = 2*pi / fromIntegral num
  in [sin (n*f*fromIntegral i) / n :+ 0
       | i <- [0,1..num-1]]

mkxa,mkxb,mkxc :: Int -> [Complex Double]
mkxa num = mkTerm num 1
mkxb num = zipWith (+) (mkxa num) (mkTerm num 3)
mkxc num = zipWith (+) (mkxb num) (mkTerm num 5)

mkPulse :: Int -> [Complex Double]
mkPulse n = 100 : take (n-1) (repeat 0)

x1 num = let f = 2*pi*pi / fromIntegral num
  in map (:+ 0) [sin (f * fromIntegral i)
                  | i <- [0,1..num-1]]

mkPolars :: [Complex Double] -> [Complex Double]
mkPolars = map ((\(m,p) -> m :+ p) . polar)

---------------
-- Exercises --
---------------

{- Exercise 20.1©:Write a Haskell function idft that implements the inverse DFT
 - as captured in Equation 20.3. Test your code by applying idft to one of the
 - signals used earlier in this section. In other words, show empirically that,
 - up to round-off errors, idft (dft xs) == xs. -}

idft :: RealFloat a => [Complex a] -> [Complex a]
idft xs = let lenI = length xs
              lenR = fromIntegral lenI
          in [let i = 2*pi* fromIntegral n / lenR
              in sum [(xs !! k) * exp (0 :+ i * fromIntegral k)
                     | k <- [0..lenI-1]]
             | n <- [0..lenI-1]]

{- Exercise 20.2: Use dft to analyze some of the signals generated using signal
 - functions defined in Chapter 19. -}

fft1 :: forall p . Clock p => SigFun p () (SEvent (M.Map Double Double))
fft1 = s1 >>> (fftA 128 256) >>> arr (fmap $ presentFFT $ rate (undefined :: p))

{- Exercise 20.3:© Define a function mkSqWave :: Int -> Int -> [Complex Double]
 - such that mkSqWave num n is the sum of the first n terms of the Fourier
 - Series of a square wave, having num samples in the result. -}

mkSqWave :: Int -> Int -> [Complex Double]
mkSqWave num n =
  foldr (zipWith (+) . (\i -> mkTerm i))
        (replicate num 0)
        (filter odd [1..(2*n-1)])
  where mkTerm x = let f = 2 * pi / fromIntegral num
                       t = fromIntegral x
                   in [sin (t*f*fromIntegral i)/t:+0
                      | i <- [0,1..num-1]]

{- Exercise 20.4: -}

{- a. © Prove mathematically that x and x^ are inverses.

  Let X^ and X be the NxN matrices defined by

  X^(k,n) = (1/N) * w^(kn)
  X(n,k)    = (1/N) * w^(-(kn))

  where w = e^(-j2π/N)
  and (a,b) corresponds to the entry in the ath row and bth column of the matrix

  Prove X^ and X are inverses
  Proof:

    It suffices to prove (X^)X = I

    Thus I show
      if p=q then X^X(p,q) = 1
      otherwise   X^X(p,q) = 0

    X^X(p,q) = X^(p,*)·X(*,q) = (1/N) ∑w^(pn)*w^(-(n*q))

    where X^(p,*) is the pth row of X^ , X(*,q) is the qth column of X and
    ∑ is a sum from 0 to N-1

    if p = q then
      X^X(p,q) = (1/N) ∑ w^(pn)/w^(nq) = (1/N) ∑ 1 =  (1/N)*N = 1

    if p ≠ q then
      X^X(p,q) = (1/N) ∑ w^(n(p-q))
               = (1 - (e^(-j*2π(p-q)/N))^N)/(1 - e^(-j*2π(p-q)/N))
               = (1 - (e^(-j*2π(p-q)))/(1 - e^(-j*2π(p-q)/N))
               = (1 - 1/(cos(2π(p-q)) + jsin(2π(p-q))))/(1 - e^(-j*2π(p-q)/N))
               = (1 - 1)/(1 - e^(-j*2π(p-q)/N))
               = 0

    Thus X^X = I, proving X^ and X are inverses.
-}

{-
e1 =
  \i -> (1/lenC)*sum[(cxs !! n)*exp(0 :+ i*fromIntegral n) | n <- [0,1..lenI-1]]

f = \exp -> let lenI = length cxs
            in let lenR = fromIntegral lenI
               in let lenC = lenR :+ 0
                  in exp

-> fdft = (f (concatMap
               (\k -> let i = -2*pi*fromIntegral k/lenR in e1 i)
               [0,1..lenI-1]
             )
          )

fix exp = \x -> let x = exp x in x
-}

{- b. Also prove, using equational reasoning, that dft and idft are
 - inverses. (..., you may assume that Haskell numeric types
 - obey the standard axioms of real arithmetic.)

idft(dft(cxs))  = cxs

idft(dft(cxs))
-> idft (let lenI = length cxs
             lenR = fromIntegral lenI
             lenC = lenR :+ 0
         in [let i = -2*pi*fromIntegral k/lenR
             in (1/lenC)*sum[(cxs !! n)*exp (0 :+ i*fromIntegral n)
                             | n <-  [0,1..lenI-1]]
              |k <- [0,1..lenI-1]] )

-> idft (let lenI = length cxs
             lenR = fromIntegral lenI
             lenC = lenR :+ 0
         in [let i = -2*pi*fromIntegral k/lenR
             in e1 i | k <- [0,1..lenI-1]] )

  -> concatMap
       (\k -> let i = -2*pi*fromIntegral k/lenR in e1 i)
       [0,1..lenI-1]

-> idft (f (concatMap
             (\k -> let i = -2*pi*fromIntegral k/lenR in e1 i)
             [0,1..lenI-1]
           )
        )

-> let lenI2 = length fdft
       lenR2 = fromIntegral lenI2
   in [let i2 = 2*pi* fromIntegral n2 / lenR2
       in sum [(fdft !! k) * exp (0 :+ i2 * fromIntegral k) | k <- [0..lenI2-1]]
        | n2 <- [0..lenI2-1]]

-> let lenI2 = length fdft
       lenR2 = fromIntegral lenI2
   in [let i2 = 2*pi* fromIntegral n2 / lenR2
       in sum [(fdft !! k) * exp (0 :+ i2 * fromIntegral k) | k <- [0..lenI2-1]]
        | n2 <- [0..lenI2-1]]
-}

fftEx :: UISF () ()
fftEx =
  proc _ -> do
    f <- withDisplay $ hSlider (1,2000) 440 -< ()
    d <- clockedSFToUISF (0.1) fftDataSig -< f
    let (ds,ts) = unzip d
    _ <- histogram (makeLayout (Fixed {fixedSize = 500})
                               (Fixed {fixedSize = 150}))
                                -< listToMaybe (catMaybes (snd $ unzip ds))
    _ <- realtimeGraph (makeLayout (Fixed {fixedSize = 500})
                                   (Fixed {fixedSize = 150}))
                                   1 Black -< (zip (fst $ unzip ds) ts)
    m  <- withDisplay $ hSlider (1,2000) 20 -< ()
    d1 <- clockedSFToUISF (0.1) fftDataSig -< m
    let (ds1,ts1) = unzip d1
    _  <- histogram (makeLayout (Fixed {fixedSize = 500})
                               (Fixed {fixedSize = 150}))
                                -< listToMaybe (catMaybes (snd $ unzip ds1))
    _ <- realtimeGraph (makeLayout (Fixed {fixedSize = 500})
                                   (Fixed {fixedSize = 150}))
                                   1 Black -< (zip (fst $ unzip ds1) ts1)
    d3 <- clockedSFToUISF (0.1) fftDataSig -< (f+m)
    let (ds3,ts3) = unzip d3
    _  <- histogram (makeLayout (Fixed {fixedSize = 500})
                               (Fixed {fixedSize = 150}))
                                -< listToMaybe (catMaybes (snd $ unzip ds3))
    outA -< ()
    where
      fftDataSig :: SigFun CtrRate Double (Double, SEvent [Double])
      fftDataSig = proc f -> do
        s <- osc (tableSinesN 4096 [1]) 0 -< f
        fftData <- fftA 128 256 -< s
        outA -< (s,fftData)

t0 = runMUI (defaultMUIParams
               {uiSize       = (500,900),
                uiBackground = rgbE 93 10 13,
                uiTitle      = "FFT Test"})
            fftEx

clockedSFToUISF :: forall a b c . (NFData b,Clock c) =>
                   DeltaT -> SigFun c a b -> UISF a [(b,Time)]
clockedSFToUISF buffer ~(ArrowP sf) = let r = rate (undefined :: c)
  in asyncV r buffer (toAutomaton sf)

toAutomaton :: forall a b . SF a b -> Automaton (->) a b
toAutomaton ~(SF f) = Automaton $ \a -> let (b,sf) = f a in (b,toAutomaton sf)

{- Exercise 20.5: Modify the program in 20.5 (fftEx) in the following ways:
 -
 - 1. Add a second slider, and use it to control the frequency of a second
 - oscillator.©
 - 2. Let s1 and s2 be the names of the signals whose frequencies are controlled
 - by the first and second sliders, respectively. Instead of displaying the FFT
 - of just s1, try a variety of combinations of s1 and s2, such as
 - s1 + s2
 - s1 - s2
 - s1 * s2
 - 1/s1 + 1/s2
 - s1/s2
 - Comment on the results.
 - 3. Use s2 to control the frequency of s1 (as was done with vibrato in Ch.19).
 - Plot the fft of s1 and comment on the result.
 - 4. Instead of using osc to generate a pure sine wave, try using other
 - oscillators and/or table generators to create more complex tones, and plot
 - their FFTs. Comment on the results. -}

----------
-- Util --
----------
printComplexL :: [Complex Double] -> IO ()
printComplexL xs =
  let f (i, rl :+ im) =
        do putStr (spaces (3 - length (show i)))
           putStr (show i ++ ":    ("          )
           putStr (niceNum rl ++ ", "         )
           putStr (niceNum im ++ ") \n"        )
   in mapM_ f (zip [0..length xs - 1] xs)

niceNum :: Double -> String
niceNum d =
  let d' = fromIntegral (round (1e10*d))/1e10
      (dec,fra) = break (== '.') (show d')
      (fra',exp) = break (== 'e') fra
  in spaces (3 - length dec) ++ dec ++ take 11 fra'
     ++ exp ++ spaces (12 - length fra' - length exp)

spaces :: Int -> String
spaces n = take n (repeat ' ')


{- |
Module      :  Numeric.GSL.Fitting
Copyright   :  (c) Alberto Ruiz 2010
License     :  GPL

Maintainer  :  Alberto Ruiz (aruiz at um dot es)
Stability   :  provisional
Portability :  uses ffi

Nonlinear Least-Squares Fitting

<http://www.gnu.org/software/gsl/manual/html_node/Nonlinear-Least_002dSquares-Fitting.html>

The example program in the GSL manual (see examples/fitting.hs):

@dat = [
 ([0.0],([6.0133918608118675],0.1)),
 ([1.0],([5.5153769909966535],0.1)),
 ([2.0],([5.261094606015287],0.1)),
 ...
 ([39.0],([1.0619821710802808],0.1))]

expModel [a,lambda,b] [t] = [a * exp (-lambda * t) + b]

expModelDer [a,lambda,b] [t] = [[exp (-lambda * t), -t * a * exp(-lambda*t) , 1]]

(sol,path) = fitModelScaled 1E-4 1E-4 20 (expModel, expModelDer) dat [1,0,0]

\> path
(6><5)
 [ 1.0,  76.45780563978782, 1.6465931240727802, 1.8147715267618197e-2, 0.6465931240727797
 , 2.0, 37.683816318260355,  2.858760367632973,  8.092094813253975e-2, 1.4479636296208662
 , 3.0,    9.5807893736187,  4.948995119561291,   0.11942927999921617, 1.0945766509238248
 , 4.0,  5.630494933603935,  5.021755718065913,   0.10287787128056883, 1.0338835440862608
 , 5.0,  5.443976278682909,  5.045204331329302,   0.10405523433131504,  1.019416067207375
 , 6.0, 5.4439736648994685,  5.045357818922331,   0.10404905846029407, 1.0192487112786812 ]
\> sol
[(5.045357818922331,6.027976702418132e-2),
(0.10404905846029407,3.157045047172834e-3),
(1.0192487112786812,3.782067731353722e-2)]@

-}
-----------------------------------------------------------------------------

module Numeric.GSL.Fitting (
    -- * Levenberg-Marquardt
    nlFitting, FittingMethod(..),
    -- * Utilities
    fitModelScaled, fitModel
) where

import Data.Packed.Internal
import Numeric.LinearAlgebra
import Numeric.GSL.Internal

import Foreign.Ptr(FunPtr, freeHaskellFunPtr)
import Foreign.C.Types
import System.IO.Unsafe(unsafePerformIO)

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

data FittingMethod = LevenbergMarquardtScaled -- ^ Interface to gsl_multifit_fdfsolver_lmsder. This is a robust and efficient version of the Levenberg-Marquardt algorithm as implemented in the scaled lmder routine in minpack. Minpack was written by Jorge J. More, Burton S. Garbow and Kenneth E. Hillstrom.
                   | LevenbergMarquardt -- ^ This is an unscaled version of the lmder algorithm. The elements of the diagonal scaling matrix D are set to 1. This algorithm may be useful in circumstances where the scaled version of lmder converges too slowly, or the function is already scaled appropriately.
        deriving (Enum,Eq,Show,Bounded)

-- | Nonlinear multidimensional least-squares fitting.
nlFitting :: FittingMethod
      -> Double                     -- ^ absolute tolerance
      -> Double                     -- ^ relative tolerance
      -> Int                        -- ^ maximum number of iterations allowed
      -> (Vector Double -> Vector Double)     -- ^ function to be minimized
      -> (Vector Double -> Matrix Double)   -- ^ Jacobian
      -> Vector Double                   -- ^ starting point
      -> (Vector Double, Matrix Double)  -- ^ solution vector and optimization path

nlFitting method epsabs epsrel maxit fun jac xinit = nlFitGen (fi (fromEnum method)) fun jac xinit epsabs epsrel maxit

nlFitGen m f jac xiv epsabs epsrel maxit = unsafePerformIO $ do
    let p   = dim xiv
        n   = dim (f xiv)
    fp <- mkVecVecfun (aux_vTov (checkdim1 n p . f))
    jp <- mkVecMatfun (aux_vTom (checkdim2 n p . jac))
    rawpath <- createMatrix RowMajor maxit (2+p)
    app2 (c_nlfit m fp jp epsabs epsrel (fi maxit) (fi n)) vec xiv mat rawpath "c_nlfit"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toRows $ dropRows (it-1) path
    freeHaskellFunPtr fp
    freeHaskellFunPtr jp
    return (subVector 2 p sol, path)

foreign import ccall safe "nlfit"
    c_nlfit:: CInt -> FunPtr TVV -> FunPtr TVM -> Double -> Double -> CInt -> CInt -> TVM

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

checkdim1 n _p v
    | dim v == n = v
    | otherwise = error $ "Error: "++ show n
                        ++ " components expected in the result of the function supplied to nlFitting"

checkdim2 n p m
    | rows m == n && cols m == p = m
    | otherwise = error $ "Error: "++ show n ++ "x" ++ show p
                        ++ " Jacobian expected in nlFitting"

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

err (model,deriv) dat vsol = zip sol errs where
    sol = toList vsol
    c = max 1 (chi/sqrt (fromIntegral dof))
    dof = length dat - (rows cov)
    chi = norm2 (fromList $ cost (resMs model) dat sol)
    js = fromLists $ jacobian (resDs deriv) dat sol
    cov = inv $ trans js <> js
    errs = toList $ scalar c * sqrt (takeDiag cov)

-- | Higher level interface to 'nlFitting' 'LevenbergMarquardtScaled'. The optimization function and
-- Jacobian are automatically built from a model f vs x = y and its derivatives, and a list of
-- instances (x, (y,sigma)) to be fitted.

fitModelScaled
         :: Double -- ^ absolute tolerance
         -> Double -- ^ relative tolerance
         -> Int    -- ^ maximum number of iterations allowed
         -> ([Double] -> x -> [Double], [Double] -> x -> [[Double]]) -- ^ (model, derivatives)
         -> [(x, ([Double], Double))] -- ^ instances
         -> [Double] -- ^ starting point
         -> ([(Double, Double)], Matrix Double) -- ^ (solution, error) and optimization path
fitModelScaled epsabs epsrel maxit (model,deriv) dt xin = (err (model,deriv) dt sol, path) where
    (sol,path) = nlFitting LevenbergMarquardtScaled epsabs epsrel maxit
                (fromList . cost (resMs model) dt . toList)
                (fromLists . jacobian (resDs deriv) dt . toList)
                (fromList xin)

-- | Higher level interface to 'nlFitting' 'LevenbergMarquardt'. The optimization function and
-- Jacobian are automatically built from a model f vs x = y and its derivatives, and a list of
-- instances (x,y) to be fitted.

fitModel :: Double -- ^ absolute tolerance
         -> Double -- ^ relative tolerance
         -> Int    -- ^ maximum number of iterations allowed
         -> ([Double] -> x -> [Double], [Double] -> x -> [[Double]]) -- ^ (model, derivatives)
         -> [(x, [Double])]  -- ^ instances
         -> [Double] -- ^ starting point
         -> ([Double], Matrix Double) -- ^ solution and optimization path
fitModel epsabs epsrel maxit (model,deriv) dt xin = (toList sol, path) where
    (sol,path) = nlFitting LevenbergMarquardt epsabs epsrel maxit
                (fromList . cost (resM model) dt . toList)
                (fromLists . jacobian (resD deriv) dt . toList)
                (fromList xin)

cost model ds vs = concatMap (model vs) ds

jacobian modelDer ds vs = concatMap (modelDer vs) ds

-- | Model-to-residual for association pairs with sigma, to be used with 'fitModel'.
resMs :: ([Double] -> x -> [Double]) -> [Double] -> (x, ([Double], Double)) -> [Double]
resMs m v = \(x,(ys,s)) -> zipWith (g s) (m v x) ys where g s a b = (a-b)/s

-- | Associated derivative for 'resMs'.
resDs :: ([Double] -> x -> [[Double]]) -> [Double] -> (x, ([Double], Double)) -> [[Double]]
resDs m v = \(x,(_,s)) -> map (map (/s)) (m v x)

-- | Model-to-residual for association pairs, to be used with 'fitModel'. It is equivalent
-- to 'resMs' with all sigmas = 1.
resM :: ([Double] -> x -> [Double]) -> [Double] -> (x, [Double]) -> [Double]
resM m v = \(x,ys) -> zipWith g (m v x) ys where g a b = a-b

-- | Associated derivative for 'resM'.
resD :: ([Double] -> x -> [[Double]]) -> [Double] -> (x, [Double]) -> [[Double]]
resD m v = \(x,_) -> m v x


{-# LANGUAGE ViewPatterns #-}

module Statistics.Covariance.Shrinkage.MMSE (
    mmseCov
  , shrinkCov
  , sampleCov
  , naiveCov
  , pinCoef
  , lwShrinkageCoef
  , rblwShrinkCovCoef
  , rblwShrinkageCoef
  , oracleShrinkCovCoef
  , oracleShrinkageCoef
  ) where

import Numeric.LinearAlgebra hiding ((<>))
import Numeric.Sum (kbn, sumVector)
import Statistics.Sample (mean)
import qualified Data.Vector.Unboxed as VU

mmseCov :: (Matrix Double -> Double) -> Matrix Double -> Herm Double
mmseCov shrink x = shrinkCov s p f
  where
    s = sampleCov x
    f = naiveCov x
    p = shrink x

shrinkCov :: Herm Double -> Double -> Herm Double -> Herm Double
shrinkCov sample coef target = scale (1 - coef) sample `add` scale coef target

sampleCov :: Matrix Double -> Herm Double
sampleCov = snd . meanCov

avgVarCov :: Herm Double -> Herm Double
avgVarCov (unSym -> s) = trustSym $ scale s' (ident p)
  where
    p = cols s
    s' = mean $ takeDiag s

naiveCov :: Matrix Double -> Herm Double
naiveCov = avgVarCov . sampleCov

pinCoef :: Double -> Double
pinCoef = max 0 . min 1

lwShrinkageCoef :: Matrix Double -> Double
lwShrinkageCoef x = numer / denom
  where
    n = rows x
    p = cols x
    s = unSym $ sampleCov x
    trs2 = trace $ s <> s
    tr2s = sq $ trace s
    frob ix = let c = x ? [ix] in sq . norm_Frob $ unSym (mTm c) - s
    numer = sumVector kbn (VU.generate n frob) / fromIntegral (sq n)
    denom = trs2 - tr2s / fromIntegral p

rblwShrinkCovCoef :: Int -> Herm Double -> Double
rblwShrinkCovCoef (fromIntegral -> n) (unSym -> s) = numer / denom
  where
    p = cols s
    trs2 = trace $ s <> s
    tr2s = sq $ trace s
    numer = (n - 2) / n * trs2 + tr2s
    denom = (n + 2) * (trs2 - tr2s / fromIntegral p)

rblwShrinkageCoef :: Matrix Double -> Double
rblwShrinkageCoef x = rblwShrinkCovCoef n (sampleCov x)
  where n = rows x

data OASIter = OASIter Double (Herm Double)

data OAS = OAS Int Int (Herm Double) (Herm Double)

oracleStep :: OAS -> Herm Double -> OASIter
oracleStep (OAS n p s f) sj = OASIter pj' sj'
  where
    n' = fromIntegral n
    p' = fromIntegral p
    trss = trace $ unSym sj <> unSym s
    tr2s = sq . trace $ unSym sj
    numer = (1 - 2 / p') * trss  + tr2s
    denom = (n' + 1 - 2 / p') * trss + (1 - n' / p') * tr2s
    pj' = numer / denom
    sj' = scale (1 - pj') s `add` scale pj' f

oracleIter :: Double -> OAS -> Herm Double -> Double
oracleIter tol oas = go 1 . oracleStep oas
  where
    go p (OASIter p' sj')
      | abs (p - p') < tol = p'
      | otherwise = go p' $ oracleStep oas sj'

oracleShrinkCovCoef :: Int -> Herm Double -> Double
oracleShrinkCovCoef n s = oracleIter 1e-6 (OAS n p s f) s
  where
    p = cols $ unSym s
    f = avgVarCov s

oracleShrinkageCoef :: Matrix Double -> Double
oracleShrinkageCoef x = oracleShrinkCovCoef n (sampleCov x)
  where n = rows x

sq :: Num a => a -> a
sq a = a * a

trace :: Matrix Double -> Double
trace = sumVector kbn . takeDiag


module Haskov where

import System.Random
import Data.Tuple (swap)
import Data.List (foldl')
import Data.Sequence (Seq, (|>))
import qualified Data.Sequence as Seqq
import Data.Map.Strict (Map, keys, foldlWithKey', foldrWithKey', mapWithKey)
import qualified Data.Map.Strict as Map
import Numeric.LinearAlgebra (scale, luSolve, luPacked)
--import qualified Numeric.LinearAlgebra as Lin
import Numeric.LinearAlgebra.Data (Matrix, matrix, col, cols, ident, tr)
import qualified Numeric.LinearAlgebra.Data as Dat
import Numeric.LinearAlgebra.HMatrix (sumElements, cmap)
import qualified Numeric.LinearAlgebra.HMatrix as Hma

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

-- Markov Type --
type IndexMap a = Map a Int
type ReverseIndexMap a = Map Int a
type HMatrixMap = Map (Int, Int) Double

data Markov a = Markov {  imap :: IndexMap a
                        , hmap :: HMatrixMap }

instance (Show a, Ord a) => Show (Markov a) where
    show haskov = show $ toList haskov

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

-- Query --
-- Safe query
lookup :: (Ord a) => a -> a -> Markov a -> Maybe Double
lookup i j (Markov imap hmap)
    | (Map.member i imap) && (Map.member j imap) = Just (hmap Map.! ((imap Map.!i), (imap Map.! j)))
    | otherwise                                  = Nothing

-- Returns number of states
size :: (Ord a) => Markov a -> Int
size (Markov imap hmap) = Map.size imap

-- Returns true if the chain is empty
null :: (Ord a) => Markov a -> Bool
null (Markov imap hmap) = Map.null imap

-- Returns true if input is a state of the chain
member :: (Ord a) => a -> Markov a -> Bool
member i (Markov imap hmap)
    | (Map.member i imap) = True
    | otherwise = False

-- Returns true if input is not a state of the chain
notMember :: (Ord a) => a -> Markov a -> Bool
notMember i haskov = not (member i haskov)

-- States of the Markov chain
states :: (Ord a) => Markov a -> [a]
states (Markov imap hmap) = keys imap

statesI :: (Ord a) => Markov a -> IndexMap a
statesI (Markov imap hmap) = imap

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

-- Construction --
empty :: Markov a
empty = Markov (Map.empty) (Map.empty)

insert :: (Ord a) => a -> a -> Double -> Markov a -> Markov a
insert i j d (Markov imap hmap)
    | hasI && (not hasJ) = Markov (Map.insert j s imap) (Map.insert (imap Map.! i, s) d hmap)
    | (not hasI) && hasJ = Markov (Map.insert i s imap) (Map.insert (s, imap Map.! j) d hmap)
    | (not hasI) && (not hasJ) && (i /= j) = Markov (Map.insert j (s+1) (Map.insert i s imap)) (Map.insert (s, (s+1)) d hmap)
    | (not hasI) && (not hasJ) && (i == j) = Markov (Map.insert i s imap) (Map.insert (s, s) d hmap)
    | otherwise          = Markov imap (Map.insert (imap Map.! i, imap Map.! j) d hmap)
    where
        hasI = Map.member i imap
        hasJ = Map.member j imap
        s = Map.size imap

insertWith :: (Ord a) => (Double -> Double -> Double) -> a -> a -> Double -> Markov a -> Markov a
insertWith f i j d (Markov imap hmap)
    | hasI && (not hasJ) = Markov (Map.insert j s imap) (Map.insert (imap Map.! i, s) d hmap)
    | (not hasI) && hasJ = Markov (Map.insert i s imap) (Map.insert (s, imap Map.! j) d hmap)
    | (not hasI) && (not hasJ) && (i /= j) = Markov (Map.insert j (s+1) (Map.insert i s imap)) (Map.insert (s, (s+1)) d hmap)
    | (not hasI) && (not hasJ) && (i == j) = Markov (Map.insert i s imap) (Map.insert (s, s) d hmap)
    | otherwise          = Markov imap (Map.insertWith f (imap Map.! i, imap Map.! j) d hmap)
    where
        hasI = Map.member i imap
        hasJ = Map.member j imap
        s = Map.size imap

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

-- Lists --
fromList :: (Ord a) => [((a, a), Double)] -> Markov a
fromList list =
    let imap = foldl' toimap Map.empty list
        hmap = foldl' (tohmap imap) Map.empty list
    in Markov imap hmap

toList :: (Ord a) => Markov a -> [((a, a), Double)]
toList (Markov imap hmap) =
    let rIMap = reverseIMap imap
    in foldrWithKey' (combineImapHmap rIMap) [] hmap

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

-- Maps --
fromMap :: (Ord a) => Map (a, a) Double -> Markov a
fromMap mapp =
    let imap = foldlWithKey' toimap' Map.empty mapp
        hmap = foldlWithKey' (tohmap' imap) Map.empty mapp
    in Markov imap hmap

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

-- Matrices --
hmatrix :: (Ord a) => Markov a -> Matrix Double
hmatrix (Markov imap hmap) =
    let m = Map.size imap
        zeros = matrix m (replicate (m*m) 0)
        hlist = Map.toList hmap
    in  Dat.accum zeros (+) hlist

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

-- Chains --

walkFrom :: (Ord a, Show a) => a -> Int -> Markov a -> StdGen -> IO [a]
walkFrom start n haskov gen = do
    let mat = hmatrix haskov
        imap = statesI haskov
        rimap = transMap imap
        nextRow = mat Dat.? [(imap Map.! start)]
    return (start : (steps rimap mat nextRow n gen))

walk :: (Ord a, Show a) => Int -> Markov a -> StdGen -> IO [a]
walk n haskov gen = do
    let mat = hmatrix haskov
        ss = tr . steady $ mat
        imap = statesI haskov
        rimap = transMap imap
    return (steps rimap mat ss n gen)

steps :: (Ord a, Show a) => ReverseIndexMap a -> Matrix Double -> Matrix Double -> Int -> StdGen -> [a]
steps _ _ _ 0 _ = []
steps rimap mat row n gen
    | (sumElements row) == 0 = []
    | otherwise = choice : steps rimap mat newRow (n-1) (snd rand)
    where
        rand = random gen
        index = randomStep row (fst rand) 0.0 0
        choice = rimap Map.! index
        newRow = mat Dat.? [index]

randomStep :: Matrix Double -> Double -> Double -> Int -> Int
randomStep row rand total j
    | newTotal < rand = randomStep row rand newTotal newJ
    | otherwise       = j
    where
        newTotal = total + (row Dat.! 0 Dat.! j)
        newJ = (j + 1) `mod` (cols row)

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

-- Steady State --
steadyState :: (Ord a) =>  Markov a -> [(a, Double)]
steadyState haskov = zip (states haskov) (concat . Dat.toLists $ steady (hmatrix haskov))

steady :: Matrix Double -> Matrix Double
steady mat =
    let m = Dat.rows mat -- Number of states
        q = tr $ ident m - mat -- Transpose (I-matrix size m - hmatrix)
        e = col ((replicate (m-1) 0) ++ [1.0]) -- col of zeros size m, last element is 1.0
        x = luSolve (luPacked q) e -- Solve linear system of Q and e (Qx = e)
    in  scale (1 / (sumElements x)) x  -- Divide x by sum of elements

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

-- Normalize --
normalize :: (Ord a) => Markov a -> Markov a
normalize (Markov imap hmap) =
    let ss = foldlWithKey' sumSeq Seqq.empty hmap
    in Markov imap (mapWithKey (calcNorm ss) hmap)

calcNorm :: Seq Double -> (Int, Int) -> Double -> Double
calcNorm ss (i, j) n = n / (Seqq.index ss i)

sumSeq :: Seq Double -> (Int, Int) -> Double -> Seq Double
sumSeq s (i, j) n =
    if Seqq.length s < (i+1)
    then s |> n
    else Seqq.adjust (+n) i s

norm :: Matrix Double -> Matrix Double
norm row = cmap (/ (sumElements row)) row

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

---- Mapping --
--map :: (Ord a) => (Double -> Double) -> Markov a -> Markov a
--map f (Markov imap hmap) = Markov imap (fmap f hmap)

--mapWithState :: (Ord a) => (a -> Double -> Double) -> Markov a -> Markov a
--mapWithState f (Markov imap hmap) =
    --let rIMap = reverseIMap imap

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

-- Delete/Update --
--delete :: (Ord a) => a -> a -> Markov a -> Markov a
--delete i j (Markov imap hmap) =

adjust :: (Ord a) => (Double -> Double) -> a -> a -> Markov a -> Markov a
adjust f i j (Markov imap hmap) =
    if Map.member i imap && Map.member j imap
    then Markov imap (Map.adjust f (imap Map.! i, imap Map.! j) hmap)
    else Markov imap hmap

-- Helper Functions --
-- for fromList
toimap :: (Ord a) => IndexMap a -> ((a, a), Double) -> IndexMap a
toimap acc ((i, j), d)
    | i /= j && Map.notMember i acc && Map.notMember j acc = Map.insert j (Map.size acc+1) (Map.insert i (Map.size acc) acc)
    | Map.notMember i acc = Map.insert i (Map.size acc) acc
    | Map.notMember j acc = Map.insert j (Map.size acc) acc
    | otherwise = acc

-- For fromMap
toimap' :: (Ord a) => IndexMap a -> (a, a) -> Double -> IndexMap a
toimap' acc (i, j) d
    | i /= j && Map.notMember i acc && Map.notMember j acc = Map.insert j (Map.size acc+1) (Map.insert i (Map.size acc) acc)
    | Map.notMember i acc = Map.insert i (Map.size acc) acc
    | Map.notMember j acc = Map.insert j (Map.size acc) acc
    | otherwise = acc

-- For fromList
tohmap :: (Ord a) => IndexMap a -> HMatrixMap -> ((a, a), Double) -> HMatrixMap
tohmap imap acc ((i, j), n) = Map.insert (imap Map.! i, imap Map.! j) n acc

-- For fromMap
tohmap' :: (Ord a) => IndexMap a -> HMatrixMap -> (a, a) -> Double -> HMatrixMap
tohmap' imap acc (i, j) n = Map.insert (imap Map.! i, imap Map.! j) n acc

combineImapHmap :: (Ord a) => Map Int a -> (Int, Int) -> Double -> [((a, a), Double)] -> [((a, a), Double)]
combineImapHmap rIMap (i, j) d acc = ((rIMap Map.! i, rIMap Map.! j), d) : acc

reverseIMap :: (Ord a) => Map a Int -> Map Int a
reverseIMap imap = foldrWithKey' reverseIMap' Map.empty imap
    where reverseIMap' state i m = Map.insert i state m

-- swap map key and index
transMap :: (Ord k, Ord a) => Map k a -> Map a k
transMap = Map.fromList . map swap . Map.toList

-- Machine Epsilon --

machineE :: Double
machineE = until ((== 1) . (+1)) (/2) 1


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_HADDOCK show-extensions #-}

module NQS.CG
  ( Operator
  , NQS.CG.cg
  , test
  ) where

import Data.Complex
import Control.Monad.Primitive

import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV

import Foreign.Storable

import Lens.Micro

import NQS.Internal.Types
import NQS.Internal.BLAS

type Operator m el = -- forall m. PrimMonad m
     MDenseVector 'Direct (PrimState m) el -- ^ Input
  -> MDenseVector 'Direct (PrimState m) el -- ^ Output
  -> m ()

class ToComplex a where
  promoteToComplex :: RealOf a -> a

instance {-# OVERLAPPABLE #-} (a ~ RealOf a) => ToComplex a where
  promoteToComplex = id

instance {-# OVERLAPS #-} (Num a, a ~ RealOf (Complex a)) => ToComplex (Complex a) where
  promoteToComplex = (:+ 0)

cg :: forall v el m.
     ( PrimMonad m, v ~ MDenseVector 'Direct
     , DOTC v el el, AXPY v el, COPY v el
     , NRM2 v el (RealOf el), SCAL (RealOf el) v el
     , Storable el, Fractional el, ToComplex el, RealFloat (RealOf el)
     )
  => Int -- ^ Max number of iterations
  -> RealOf el -- ^ Tolerance
  -> Operator m el -- ^ A
  -> MDenseVector 'Direct (PrimState m) el -- ^ b
  -> MDenseVector 'Direct (PrimState m) el -- ^ Initial guess x
  -> m (Int, RealOf el)
cg !maxIter !tol operator !b !x0 = do
  r <- newTempVector n
  p <- newTempVector n
  q <- newTempVector n
  bNorm <- norm b
  -- This will probably never ever happen, but we should check anyway
  if bNorm == 0
    then fill x0 0 >> return (0, 0)
    else do
      -- r := b - max * x0
      copy b r
      operator x0 q
      axpy (-1) q r
      -- Are we done already?
      norm r >>= \rNorm -> if rNorm < tol * tol * bNorm
        then return $! (0, sqrt (rNorm / bNorm))
        else do
          -- Start the loop
          copy r p
          go bNorm x0 r p q rNorm 0
  where
    n = b ^. dim
    newTempVector size = MDenseVector size 1 <$> newVectorAligned size 64
    norm x = (\y -> y * y) <$> nrm2 x
    go !bNorm !x !r !p !q !ρ !i
      | i >= maxIter = nrm2 r >>= \rNorm -> return (i, rNorm / bNorm)
      | otherwise = do
        operator p q
        α <- (promoteToComplex ρ /) <$> dotc p q
        axpy   α  p x
        axpy (-α) q r
        norm r >>= \rNorm -> if rNorm < threshold
          then return (i, sqrt (rNorm / bNorm))
          else do copy r q
                  let ρ' = rNorm
                      β  = ρ' / ρ
                  scal β p
                  axpy 1 q p
                  go bNorm x r p q ρ' (i + 1)
      where threshold :: RealOf el
            threshold = tol * tol * bNorm

test :: IO ()
test = do
  let -- operator :: Operator Float
      operator x y = do
        aBuff <- newVectorAligned 9 64 :: IO (MV.MVector (PrimState IO) Float)
        V.copy aBuff $ V.fromList [1, 2, 3, 4, 5, 6, 7, 8, 9]
        let a = MDenseMatrix 3 3 3 aBuff :: MDenseMatrix 'Row (PrimState IO) Float
        gemv NoTranspose 1.0 a x 0.0 y
  bBuff <- V.unsafeThaw $ V.fromList [2, 2, 2 :: Float]
  xBuff <- V.unsafeThaw $ V.fromList [0, 0, 0 :: Float]
  let x = MDenseVector 3 1 xBuff :: MDenseVector 'Direct (PrimState IO) Float
  let b = MDenseVector 3 1 bBuff :: MDenseVector 'Direct (PrimState IO) Float
  print =<< (NQS.CG.cg 30 (1.0E-8 :: Float) operator b x)


{-# OPTIONS_GHC -fplugin Numeric.Units.Dimensional.DK.Solver #-}

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeOperators #-}

module Controls
where

import Numeric.Units.Dimensional.DK.Prelude
import Numeric.LinearAlgebra.Dimensional.DK.Shapes
import qualified Prelude as P

data DimMat (s :: Shape) (e :: *) = DimMat

(<>) :: DimMat s1 e -> DimMat s2 e -> DimMat (ShapeProduct s1 s2) e
_ <> _ = DimMat

add :: DimMat s e -> DimMat s e -> DimMat s e
add _ _ = DimMat

data ContinuousLiSystem (iv :: Dimension) (xs :: [Dimension]) (ys :: [Dimension]) (us :: [Dimension]) e = ContinuousLiSystem
                                                                                {
                                                                                a'' :: DimMat (DivideVectorLists (MapDiv iv xs) xs) e,
                                                                                b'' :: DimMat (DivideVectorLists (MapDiv iv xs) us) e,
                                                                                c'' :: DimMat (DivideVectorLists ys xs) e,
                                                                                d'' :: DimMat (DivideVectorLists ys us) e
                                                                                }

type ContinuousLtiSystem = ContinuousLiSystem DTime

type ExampleSystem = ContinuousLtiSystem
    '[DLength, DVelocity, DPlaneAngle, DAngularVelocity]
    '[DLength, DPlaneAngle]
    '[DForce]
    Double

evaluate :: ContinuousLiSystem iv (x ': xs) (y ': ys) (u ': us) e -> DimMat (VectorShape x xs) e -> DimMat (VectorShape u us) e -> (DimMat (VectorShape (x / iv) (MapDiv iv xs)) e,  DimMat (VectorShape y ys) e)
evaluate sys x u = let
                      a = a'' sys
                      b = b'' sys
                      c = c'' sys
                      d = d'' sys
                      xDot = (a <> x) `add` (b <> u)
                      y = (c <> x) `add` (d <> u)
                    in (xDot, y)

{- Example from http://ctms.engin.umich.edu/CTMS/index.php?example=InvertedPendulum&section=ControlStateSpace -}

-- not really stated on that page, but if you go back a couple of pages in their derivation
-- you can see that the type of u is a 1x1 matrix whose sole element is a force

massOfCart = (0.5 :: Double) *~ (kilo gram)
massOfPendulum = (0.2 :: Double) *~ (kilo gram)
coefficientOfFrictionForCart = (0.1 :: Double) *~ (newton / (meter / second))
lengthToPendulumCenterOfMass = (0.3 :: Double) *~ meter
massMomentOfInertiaOfPendulum = (0.006 :: Double) *~ (kilo gram * meter^pos2)
g = (9.8 :: Double) *~ (meter / second^pos2)

p = massMomentOfInertiaOfPendulum*(massOfCart+massOfPendulum)+(massOfCart*massOfPendulum*lengthToPendulumCenterOfMass*lengthToPendulumCenterOfMass)

a22 = negate (massMomentOfInertiaOfPendulum+massOfPendulum * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass) * coefficientOfFrictionForCart / p
a23 = (massOfPendulum * massOfPendulum * g * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass) / p
a42 = negate (massOfPendulum * lengthToPendulumCenterOfMass * coefficientOfFrictionForCart) / p
a43 = massOfPendulum * g * lengthToPendulumCenterOfMass*(massOfCart + massOfPendulum)/p

b21 = (massMomentOfInertiaOfPendulum + (massOfPendulum * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass)) / p
b41 = massOfPendulum * lengthToPendulumCenterOfMass / p

{-
pendulum = ContinuousLiSystem {
           a'' = [matD| _0, _1, _0, _0;
                        _0, a22, a23, _0;
                        _0, _0, _0, _1;
                        _0, a42, a43, _0 |],
           b'' = [matD| _0;
                        b21;
                        _0;
                        b41 |],
           c'' = [matD| _1, _0, _0, _0;
                        _0, _0, _1, _0 |],
           d'' = zeroes
         } :: ExampleSystem
-}

pendulum = ContinuousLiSystem {
           a'' = DimMat,
           b'' = DimMat,
           c'' = DimMat,
           d'' = DimMat
         } :: ExampleSystem


module Statistics.Information.Continuous.TotalCorrelation where

import Data.Matrix
import Statistics.Information.Continuous.Entropy
import Statistics.Information.Continuous.MutualInfo
import Statistics.Information.Utils.Matrix
import Statistics.Information.Utils.Random
import System.Random

tc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
tc g k base xs = sum hxs - hx where
  hxs = [entropy gcol k base col | (col, gcol) <- cols]
  hx = entropy g2 k base xs
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

ctc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
       Double
ctc g k base xs ys = sum hxs - hx where
  hxs = [centropy gcol k base col ys | (col, gcol) <- cols]
  hx = centropy g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

corex_tcs :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
             Double
corex_tcs g k base xs ys = tc g1 k base xs - ctc g2 k base xs ys where
  (g1, g2) = split g

corex_mis :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
             Double
corex_mis g k base xs ys = sum mixs - mi_all where
  mixs = [mi gcol k base col ys | (col, gcol) <- cols]
  mi_all = mi g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

residual :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
residual g k base xs = sum $ [hlessxi i gi | (i, gi) <- igs] where
  hlessxi i gi = let xi = (colVector $ getCol i xs) in
    centropy gi k base xi $ residualMatrix xs i
  igs = [1..(ncols xs)] `zip` (splitN (ncols xs) g)

dtc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
dtc g k base xs = entropy g1 k base xs - residual g2 k base xs where
  (g1, g2) = split g


-- MandelBrot implementation in Haskell
import Data.Complex
import System.Environment
import Control.Parallel.Strategies
import Criterion.Measurement
--import Graphics.UI.SDL

--initSDL :: IO Surface
--initSDL = do
--    SDL.init [SDL.InitEverything]
--    SDL.setVideoMode (floor imageWidth) (floor imageHeight) 32 []
--    SDL.setCaption "Mandelbrot" "Mandelbrot"-   screen <- getVideoSurface
--    white <- SDL.mapRGB (SDL.surfaceGetPixelFormat screen) 255 255 255
--    fillRect screen Nothing white >> SDL.flip screen
--    return screen

mandelbrot :: Num a => a -> Int -> a
mandelbrot a n = iterate (\z -> z^2 + a) 0 !! n

mbrot_seq :: Int -> Int -> Int -> [[Bool]]
mbrot_seq iter w h = [[magnitude (mandelbrot (x :+ y) iter) < 2
                       | x <- [-2, (-2 + incW) ..   2]]
                       | y <- [ 2,  (2 - incH) ..  -2]]
 where incW = 4.0 / (fromIntegral w)
       incH = 4.0 / (fromIntegral h)

mbrot_par :: Int -> Int -> Int  -> [[Bool]]
mbrot_par iter w h = (mbrot_seq iter w h) `using` parList rdeepseq

draw_image :: Int -> Int -> [[Bool]] -> IO()
draw_image x y xs = print "to be implemented"

main = do
    initializeTime
    start_time <- getTime
    args <- getArgs
    if (length args < 4)
        then print "Expected \"par/seq\"iterations, image width, and image height"
        else do
            let nIter = read (args !! 1) :: Int
            let nWidth = read (args !! 2) :: Int
            let nHeight = read (args !! 3) :: Int
            case (args !! 0) of
                "seq" -> print $ mbrot_seq nIter nWidth nHeight
                "par" -> print $ mbrot_par nIter nWidth nHeight
                otherwise -> print "expected first argument to be `par` or `seq`"
            end_time <- getTime
            print $ "Time elapsed : " ++ show (floor $ 10000 * (end_time - start_time)) ++ "ms"


{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
{-# LANGUAGE StrictData       #-}
module FourierPinwheel.Array where

import qualified Data.Array.Accelerate                  as A
import qualified Data.Array.Accelerate.LLVM.PTX as A
import           Data.Array.Repa                        as R
import           Data.Complex
import           Data.List                              as L
import           Data.Vector.Generic                    as VG
import           Data.Vector.Storable                   as VS
import           DFT.Plan
import           Filter.Utils
import           Graphics.Rendering.Chart.Backend.Cairo
import           Graphics.Rendering.Chart.Easy
import           Image.IO
import           Text.Printf
import           Utils.Array
import           Utils.BLAS
import           Utils.Parallel
import FourierMethod.FourierSeries2D

data FPArray vector = FPArray
  { getFPArrayNumXFreq     :: Int
  , getFPArrayNumYFreq     :: Int
  , getFPArrayNumRFreq     :: Int
  , getFPArrayNumThetaFreq :: Int
  , getFPArrayNumRhoFreq   :: Int
  , getFPArrayNumPhiFreq   :: Int
  , getFPArray             :: [vector]
  }

{-# INLINE toMatrixAcc #-}
toMatrixAcc ::
     (VG.Vector vector e, A.Elt e) => FPArray (vector e) -> A.Acc (A.Array A.DIM2 e)
toMatrixAcc (FPArray numXFreq numYFreq numRFreq numThetaFreq _ _ arr) =
  A.use .
  A.fromList (A.Z A.:. (numRFreq * numThetaFreq) A.:. (numXFreq * numYFreq)) .
  VG.toList . VG.concat $
  arr

{-# INLINE parMapFPArray #-}
parMapFPArray ::
     (NFData vector) => (vector -> vector) -> FPArray vector -> FPArray vector
parMapFPArray f (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs) =
  FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq .
  parMap rdeepseq f $
  vecs

{-# INLINE parZipWithFPArray #-}
parZipWithFPArray ::
     (NFData vector)
  => (vector -> vector -> vector)
  -> FPArray vector
  -> FPArray vector
  -> FPArray vector
parZipWithFPArray f (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs1) =
  FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq .
  parZipWith rdeepseq f vecs1 . getFPArray

plotFPArray ::
     DFTPlan
  -> FilePath
  -> FPArray (VS.Vector (Complex Double))
  -> IO (R.Array R.U R.DIM4 (Complex Double))
plotFPArray plan filePath arr = do
  let planIDBackward =
        DFTPlanID
          IDFT1DG
          [ getFPArrayNumThetaFreq arr
          , getFPArrayNumXFreq arr
          , getFPArrayNumYFreq arr
          ]
          [1, 2]
  vecsR2 <- dftExecuteBatchP plan planIDBackward . getFPArray $ arr
  arrR2 <-
    R.computeUnboxedP .
    makeFilter2DInverse .
    R.fromUnboxed
      (R.Z R.:. getFPArrayNumRFreq arr R.:. getFPArrayNumThetaFreq arr R.:.
       getFPArrayNumXFreq arr R.:.
       getFPArrayNumYFreq arr) .
    VS.convert . VS.concat $
    vecsR2
  arr1 <-
    R.sumP . R.sumS . R.map (\x -> magnitude x Prelude.^ 2) . rotate4D2 $ arrR2
  plotImageRepa filePath .
    ImageRepa 8 .
    R.fromUnboxed
      (R.Z R.:. (1 :: Int) R.:. getFPArrayNumXFreq arr R.:.
       getFPArrayNumYFreq arr) .
    -- VG.map (\x -> log (x + 1)) .
    R.toUnboxed $ arr1
  return arrR2

plotFPArrayAcc ::
     (A.Elt (Complex Double))
  => [A.PTX]
  -> FilePath
  -> Int
  -> Double
  -> Double
  -> Int
  -> FPArray (VS.Vector (Complex Double))
  -> IO (R.Array R.U R.DIM4 (Complex Double))
plotFPArrayAcc ptxs filePath numPoints delta periodR2 numBatch arr@(FPArray numXFreq numYFreq numRFreq numThetaFreq _ _ vecs) = do
  let rows = numRFreq * numThetaFreq
      cols = numXFreq * numYFreq
      arrMat =
        A.transpose .
        A.use .
        A.fromList (A.Z A.:. rows A.:. cols) .
        R.toList .
        makeFilter2DInverse .
        fromUnboxed (Z :. rows :. numXFreq :. numXFreq) . VG.convert . VG.concat $
        vecs
  arrR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numXFreq
      numPoints
      rows
      periodR2
      delta
      numBatch
      arrMat
  (sumP . R.map (\x -> magnitude x ** 2) . rotate3D $ arrR2) >>=
    plotImageRepa filePath .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) . toUnboxed
  return .
    computeS .
    R.reshape (Z :. numRFreq :. numThetaFreq :. numPoints :. numPoints) $
    arrR2

plotFPArrayFreqency ::
     FilePath
  -> FPArray (VS.Vector (Complex Double))
  -> IO ()
plotFPArrayFreqency filePath arr = do
  arr1 <-
    R.sumP .
    R.sumS .
    R.map (\x -> magnitude x Prelude.^ 2) .
    rotate4D2 .
    makeFilter2DInverse .
    R.fromUnboxed
      (R.Z R.:. getFPArrayNumRFreq arr R.:. getFPArrayNumThetaFreq arr R.:.
       getFPArrayNumXFreq arr R.:.
       getFPArrayNumYFreq arr) .
    VS.convert . VS.concat . getFPArray $
    arr
  plotImageRepa filePath .
    ImageRepa 8 .
    R.fromUnboxed
      (R.Z R.:. (1 :: Int) R.:. getFPArrayNumXFreq arr R.:.
       getFPArrayNumYFreq arr) .
    -- VG.map (\x -> x^2) .
    R.toUnboxed $
    arr1

plotRThetaDist ::
     FilePath
  -> FilePath
  -> Int
  -> Int
  -> Int
  -> Double
  -> (Int, Int)
  -> R.Array R.U R.DIM4 (Complex Double)
  -> IO ()
plotRThetaDist filePathTheta filePathR numPoints numTheta numR periodEnv (x, y) arr = do
  let (Z :. numRFreq :. numThetaFreq :. _ :. _) = extent arr
      center = div numPoints 2
      centerR = div numR 2
      centerTheta = div numTheta 2
      centerRFreq = div numRFreq 2
      centerThetaFreq = div numThetaFreq 2
      vecRTheta =
        VG.convert . R.toUnboxed . R.computeUnboxedS . R.slice arr $
        (R.Z R.:. R.All R.:. R.All R.:. (x + center) R.:. (y + center))
      deltaTheta = 2 * pi / Prelude.fromIntegral numTheta
      deltaLogR = log periodEnv / Prelude.fromIntegral numR
      matrix =
        VG.convert .
        R.toUnboxed .
        R.computeUnboxedS .
        R.fromFunction (R.Z R.:. numR R.:. numTheta :. numRFreq :. numThetaFreq) $ \(Z :. r :. t :. rf :. tf) ->
          cis
            (fromIntegral ((r - centerR) * (rf - centerRFreq)) * 2 * pi /
             log periodEnv *
             deltaLogR +
             fromIntegral ((t - centerTheta) * (tf - centerThetaFreq)) *
             deltaTheta)
      m = numR * numTheta
      n = 1
      k = numRFreq * numThetaFreq
  vec <- VS.map magnitude <$> gemmBLAS m n k matrix vecRTheta
  let outputArr = fromUnboxed (Z :. numR :. numTheta) . VG.convert $ vec
      outputR = R.toList . sumS $ outputArr
      outputTheta =
        R.toList .
        sumS .
        R.backpermute (Z :. numTheta :. numR) (\(Z :. i :. j) -> (Z :. j :. i)) $
        outputArr
  toFile def filePathTheta $ do
    layout_title .= printf "Theta (%d , %d)" x y
    plot
      (line
         ""
         [ L.zip
             [ fromIntegral (i - centerTheta) * deltaTheta
             | i <- [0 .. numTheta - 1]
             ]
             outputTheta
         ])
  toFile def filePathR $ do
    layout_title .= printf "LogR (%d , %d)" x y
    plot
      (line
         ""
         [ L.zip
             [fromIntegral (i - centerR) * deltaLogR | i <- [0 .. numR - 1]]
             outputR
         ])


{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Model.Neural (
  -- * Fully connected
  -- ** Feed-forward
    FC, pattern FC, FCp, fcBias, fcWeights
  -- *** With activation function
  , FCA, pattern FCA, _fcaActivation
  -- ** Recurrent
  , FCR, pattern FCR, FCRp, fcrBias, fcrInputWeights, fcrStateWeights
  -- *** With activation function
  , FCRA, pattern FCRA, _fcraStore, _fcraActivation
  ) where

import           Backprop.Learn.Model.Combinator
import           Backprop.Learn.Model.Regression
import           Backprop.Learn.Model.State
import           Data.Tuple
import           GHC.TypeNats
import           Lens.Micro
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import qualified Numeric.LinearAlgebra.Static          as H

-- | Fully connected feed-forward layer with bias.  Parameterized by its
-- initialization distribution.
--
-- Note that this has no activation function; to use as a model with
-- activation function, chain it with an activation function using 'RMap',
-- ':.~', etc.; see 'FCA' for a convenient type synonym and constructor.
--
-- Without any activation function, this is essentially a multivariate
-- linear regression.
--
-- With the logistic function as an activation function, this is
-- essentially multivariate logistic regression. (See 'Logistic')
type FC i o = LinReg i o

pattern FC :: FC i o
pattern FC = LinReg

-- | Convenient synonym for an 'FC' post-composed with a simple
-- parameterless activation function.
type FCA i o = RMap (R o) (R o) (FC i o)

-- | Construct an 'FCA' using a generating function and activation
-- function.
--
-- Some common ones include 'logistic' and @'vmap' 'reLU'@.
pattern FCA
    :: (forall s. Reifies s W => BVar s (R o) -> BVar s (R o))          -- ^ '_fcaActivation'
    -> FCA i o
pattern FCA { _fcaActivation } = RM _fcaActivation FC

type FCp = LRp

fcWeights :: Lens (FCp i o) (FCp i' o) (L o i) (L o i')
fcWeights = lrBeta

fcBias :: Lens' (FCp i o) (R o)
fcBias = lrAlpha

-- | Fully connected recurrent layer with bias.
--
-- Parameterized by its initialization distributions, and also the function
-- to compute the new state from previous input.
--
-- @
-- instance 'Learn' ('R' i) (R o) ('FCR' h i o) where
--     type 'LParamMaybe' (FCR h i o) = ''Just' ('FCRp' h i o)
--     type 'LStateMaybe' (FCR h i o) = 'Just (R h)
-- @
type FCR h i o = Recurrent (R (i + h)) (R i) (R h) (R o) (FC (i + h) o)

-- | Construct an 'FCR'
pattern FCR
    :: (KnownNat h, KnownNat i)
    => (forall s. Reifies s W => BVar s (R o) -> BVar s (R h))          -- ^ '_fcrSTore'
    -> FCR h i o
pattern FCR { _fcrStore } <-
    Rec { _recLoop  = _fcrStore
        , _recLearn = FC
        }
  where
    FCR s = Rec { _recSplit = H.split
                , _recJoin  = (H.#)
                , _recLoop  = s
                , _recLearn = FC
                }
{-# COMPLETE FCR #-}

-- | Convenient synonym for an 'FCR' post-composed with a simple
-- parameterless activation function.
type FCRA h i o = RMap (R o) (R o) (FCR h i o)

-- | Construct an 'FCRA' using a generating function and activation
-- function.
--
-- Some common ones include 'logistic' and @'vmap' 'reLU'@.
pattern FCRA
    :: (KnownNat h, KnownNat i)
    => (forall s. Reifies s W => BVar s (R o) -> BVar s (R h))          -- ^ '_fcraStore'
    -> (forall s. Reifies s W => BVar s (R o) -> BVar s (R o))          -- ^ '_fcraActivation'
    -> FCRA h i o
pattern FCRA { _fcraStore, _fcraActivation }
    = RM _fcraActivation (FCR _fcraStore)

type FCRp h i o = FCp (i + h) o

lensIso :: (s -> (a, x)) -> ((b, x) -> t) -> Lens s t a b
lensIso f g h x = g <$> _1 h (f x)

fcrInputWeights
    :: (KnownNat h, KnownNat i, KnownNat i', KnownNat o)
    => Lens (FCRp h i o) (FCRp h i' o) (L o i) (L o i')
fcrInputWeights = fcWeights . lensIso H.splitCols (uncurry (H.|||))

fcrStateWeights
    :: (KnownNat h, KnownNat h', KnownNat i, KnownNat o)
    => Lens (FCRp h i o) (FCRp h' i o) (L o h) (L o h')
fcrStateWeights = fcWeights . lensIso (swap . H.splitCols) (uncurry (H.|||) . swap)

fcrBias :: Lens' (FCRp h i o) (R o)
fcrBias = fcBias


-- file: Astro/PostNewtonian.hs
-- Post-Newtonian approximations for two-body problem

module Astro.PostNewtonian
	where

import Astro.Misc
import Numeric.LinearAlgebra

import Data.List (foldl')

-- Calculate Post-Newtonian acceleration functions for non-spin
-- dependent terms
mkNonSpinPNs :: Double->Double->Double->Double->
	[(Vector Double->Vector Double->Vector Double)]
mkNonSpinPNs gc c m1 m2 = 	[ a1
				, a2
				, a2_5
				]
			where
	-- Keplerian contribution
	--a0 :: Vector Double -> Vector Double -> Vector Double
	--a0 rv vv = scale (-gm / r^3) rv where r = norm rv

	a1 :: Vector Double -> Vector Double -> Vector Double
	a1 rv vv = {-# SCC "a1" #-}
		scale (-gmperc2 / r^2) $
			scale nbrac nv + scale vbrac vv
		where
			nbrac = -2*(2+eta)*gm/r + (1+3*eta)*v^2 - 3/2*eta*rd^2
			vbrac = -2*(2-eta)*rd
			r	= norm rv
			v	= norm vv
			nv	= unitV rv
			rd	= (rv <.> vv)/r

	a2 :: Vector Double -> Vector Double -> Vector Double
	a2 rv vv = {-# SCC "a2" #-}
		scale (- gmperc2 / (c^2 * r^2)) $
			scale (nbrac/r) rv + scale ((-rd/2) * vbrac) vv
		where
			nbrac = 3/4*(12+29*eta)*(gm/r)^2 + eta*(3-4*eta)*v^4
				+ 15/8*eta*(1-3*eta)*rd^4
				- 3/2*eta*(3-4*eta)*v^2*rd^2
				- 1/2*eta*(13-4*eta)*(gm/r)*v^2
				- (2+25*eta+2*eta^2)*(gm/r)*rd^2
			vbrac = eta*(15+4*eta)*v^2 - (eta+41*eta+8*eta^2)*(gm/r)
				- 3*eta*(3+2*eta)*rd^2
			r	= norm rv
			v	= norm vv
			rd	= (rv <.> vv)/r

	a2_5 :: Vector Double -> Vector Double -> Vector Double
	a2_5 rv vv = {-# SCC "a2_5" #-}
		scale (8/15*(gmperc2)^2*eta/(c*r^2)) $
			scale (nbrac*rd/r) rv + scale (-vbrac) vv
		where
			nbrac 	= 9*v^2 + 17*gm/r
			vbrac 	= 3*v^2 + 9*gm/r
			r	= norm rv
			v	= norm vv
			rd	= (rv <.> vv)/r
	gm 	= gc * m
	gmperc2 = gm / c ^2
	m 	= m1+m2
	eta 	= m1*m2/m^2

-- Calculate Post-Newtonian acceleration functions for spin
-- dependent terms
--mkSpinPNs :: Double->Double->Double->Double->Double->Double->
--	[(Vector Double->Vector Double->Vector Double->Vector Double)]
mkSpinPNs gc c m1 m2 chi q = 	[ spinOrbit
				, quadMonoInteraction
				] where
	--spinOrbit :: Vector Double -> Vector Double -> Vector Double
	--		-> Vector Double
	spinOrbit rv vv s = {-# SCC "SO" #-}
		scale (gm^2/(r^3*c^3)*(1+sqrt (1-4*eta))/4*chi) $
				scale nbrac nv + scale (nsbrac) (cross nv s')
					+ scale (-vsbrac) (cross vv s')
		where
			nbrac = 12 * s' <.> (cross nv vv)
			nsbrac = (9 + 3 * sqrt (1-4*eta))*rd
			vsbrac = 7 + sqrt(1-4*eta)
			r	= norm rv
			v	= norm vv
			nv	= unitV rv
			rd	= (rv <.> vv)/r
			s'	= spinCylToSpher s

	--quadMonoInteraction :: Vector Double -> Vector Double -> Vector Double
	--			-> Vector Double
	quadMonoInteraction rv vv s = {-# SCC "Q" #-}
		-- FIXME: Minus or no minus? Depends on whether
		-- coordinates fixed on m1 or m2, see OJ287 paper, minus
		-- in paper.
		scale (-q*chi^2*3*gc^3*m1^2*m/(2*c^4*r^4)) $
			scale nbrac nv + scale (-2 * nv <.> s') s'
		where
			nbrac = 5 * (nv <.> s')^2 - 1
			r	= norm rv
			v	= norm vv
			nv	= unitV rv
			s'	= spinCylToSpher s
	gm 	= gc * m
	m 	= m1+m2
	eta 	= m1*m2/m^2

-- Time derivative of spin unit vector, in spherical coordinates
spinDt :: Double->Double->Double->Double->Double->Double->
	Vector Double->Vector Double->Vector Double->Vector Double
spinDt gc c m1 m2 chi q rv vv s = cross omega s
	where
	omega = scale (gm*eta/(2*c^2*r^2)*(7+sqrt (1-4*eta))
				/(1+sqrt (1-4*eta))) (cross nv vv)
	r = norm rv
	nv = scale (1/r) rv
	gm 	= gc * m
	m 	= m1+m2
	eta 	= m1*m2/m^2

-- Time derivatives of spin unit vector components in cylindrical
-- coordinates, with x as the axisymmery axis
spinDtCyl :: Double->Double->Double->Double->Double->Double->
	Vector Double->Vector Double->[(Double,Double)->Double]
spinDtCyl gc c m1 m2 chi q rv vv = [dtheta, dxi]
	where
	dtheta (th,xi) =
		-xiPerRho*(dhs@>1*cos th + dhs@>2*sin th) + dhs@>0
		where
		xiPerRho = if abs xi >= 1 then 0 else xi/sqrt (1-xi^2)
	dxi (th,xi) =
		rho * (dhs@>1*sin th - dhs@>2 * cos th)
		where
		rho = if abs xi >= 1 then 0 else sqrt $ 1 - xi^2
	r = norm rv
	nv = scale (1/r) rv
	dhs = omega
	omega = scale (gm*eta/(2*c^2*r^2)*(7+sqrt (1-4*eta))
				/(1+sqrt (1-4*eta))) (cross nv vv)
	gm 	= gc * m
	m 	= m1+m2
	eta 	= m1*m2/m^2

-- Convert spin from cylindrical to cartesian representation
-- Axisymmetry axis is x-axis
spinCylToSpher :: (Double,Double)->Vector Double
spinCylToSpher (theta,xi) = --3|>[rho * cos theta, rho*sin theta, xi]
	3|>[xi, rho * cos theta, rho*sin theta]
	where rho = if abs xi >= 1 then 0 else sqrt $ 1 - xi^2

-- Post-Newtonian acceleration functions for highest supported degree
nonSpinPNAcc gc c m1 m2 rv vv =
	sum [f rv vv | f <- mkNonSpinPNs gc c m1 m2]
spinPNAcc gc c m1 m2 chi q rv vv s =
	sum [f rv vv s | f <- mkSpinPNs gc c m1 m2 chi q]

totalPNAcc gc c m1 m2 chi q rv vv s = ((f1 rv vv) + (f2 rv vv s))
	where (f1, f2) = (nonSpinPNAcc gc c m1 m2,
				spinPNAcc gc c m1 m2 chi q)

-- Explicitly calculated total acceleration
totalPNAccExp gc c m1 m2 chi q rv vv s =
	a1 + spinOrbit + quadMono + a2 + a2_5
	where
	a1 = {-# SCC "a1" #-}
		scale (-gmperc2 / r^2) $
			scale nbrac nv + scale vbrac vv
		where
			nbrac = -2*(2+eta)*gm/r + (1+3*eta)*v^2 - 3/2*eta*rd^2
			vbrac = -2*(2-eta)*rd

	a2 = {-# SCC "a2" #-}
		scale (- gmperc2 / (c^2 * r^2)) $
			scale nbrac nv + scale ((-rd/2) * vbrac) vv
		where
			nbrac = 3/4*(12+29*eta)*(gm/r)^2 + eta*(3-4*eta)*v^4
				+ 15/8*eta*(1-3*eta)*rd^4
				- 3/2*eta*(3-4*eta)*v^2*rd^2
				- 1/2*eta*(13-4*eta)*(gm/r)*v^2
				- (2+25*eta+2*eta^2)*(gm/r)*rd^2
			vbrac = eta*(15+4*eta)*v^2 - (eta+41*eta+8*eta^2)*(gm/r)
				- 3*eta*(3+2*eta)*rd^2

	a2_5 = {-# SCC "a2_5" #-}
		scale (8/15*(gmperc2)^2*eta/(c*r^2)) $
			scale (nbrac*rd) nv + scale (-vbrac) vv
		where
			nbrac 	= 9*v^2 + 17*gm/r
			vbrac 	= 3*v^2 + 9*gm/r

	spinOrbit = {-# SCC "SO" #-}
		scale (gm^2/(r^3*c^3)*(1+sqrt (1-4*eta))/4*chi) $
				scale nbrac nv + scale (nsbrac) (cross nv s')
					+ scale (-vsbrac) (cross vv s')
		where
			nbrac = 12 * s' <.> (cross nv vv)
			nsbrac = (9 + 3 * sqrt (1-4*eta))*rd
			vsbrac = 7 + sqrt(1-4*eta)

	quadMono = {-# SCC "Q" #-}
		scale (-q*chi^2*3*gc^3*m1^2*m/(2*c^4*r^4)) $
			scale nbrac nv + scale (-2 * nv <.> s') s'
		where
			nbrac = 5 * (nv <.> s')^2 - 1
	r	= norm rv
	v	= norm vv
	s'	= spinCylToSpher s
	nv	= scale (1/r) rv
	rd	= (rv <.> vv)/r
	gm 	= gc * m
	gmperc2 = gm / c ^2
	m 	= m1+m2
	eta 	= m1*m2/m^2


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass   #-}
{-# LANGUAGE DeriveGeneric    #-}
{-# LANGUAGE FlexibleContexts #-}
module FokkerPlanck.FourierSeries
  ( fourierMellin
  , sampleScale
  -- , computeFourierCoefficients
  , computeHarmonicsArray
  , computeHarmonicsArraySparse
  -- , computeHarmonicsArrayGPU
  , getHarmonics
  , computeThetaRHarmonics
  , computeFourierSeriesThetaR
  -- , computeFourierSeriesR2
  , normalizeFreqArr
  -- , normalizeFreqArr'
  , plotThetaDimension
  , computeFourierSeriesOfLogPolarHarmonicsArray
  , computeRectangularInverseHarmonics
  ) where

import           Array.UnboxedArray                     as UA
import           Control.DeepSeq
import           Data.Array.IArray                      as IA
import           Data.Array.Repa                        as R
import           Data.Binary
import           Data.Complex
import           Data.DList                             as DL
import           Data.List                              as L
import           Data.Vector.Generic                    as VG
import           Data.Vector.Unboxed                    as VU
import           FokkerPlanck.BrownianMotion
import           FokkerPlanck.Histogram
import           GHC.Generics                           (Generic)
import           Text.Printf
import           Utils.Array
import           Utils.Parallel
import           Utils.Time
import           Graphics.Rendering.Chart.Backend.Cairo
import           Graphics.Rendering.Chart.Easy
import           Image.IO
import           System.FilePath
import           Utils.SimpsonRule
import Numeric.LinearAlgebra as NL
import Numeric.LinearAlgebra.Data as NL
import            Utils.Distribution
import Pinwheel.Base

-- {-# INLINE fourierMellin' #-}
-- fourierMellin' :: Int -> Int -> (Double, Double) -> Complex Double
-- fourierMellin' !angularFreq !radialFreq (!x, !y) =
--   if x == 0 && y == 0
--     then 0
--     else exp $ (-1 * log rho) :+ (tf * atan2 y x + rf * (log rho))

{-# INLINE sampleScale #-}
sampleScale ::
     Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Particle]
  -> Complex Double
sampleScale !phiFreq !rhoFreq !thetaFreq !rFreq !halfLogPeriod =
  L.sum .
  L.map
    (\(Particle newPhi newRho theta newR _) ->
       if newRho == 0 -- || newR == 0
         then 0
         else (cos (phiFreq * newPhi + thetaFreq * (theta - newPhi)) :+ 0) *
              (cis $
               (-pi) * (rhoFreq * (newRho) + rFreq * (newR - newRho)) /
               halfLogPeriod))

-- {-# INLINE computeFourierCoefficients #-}
-- computeFourierCoefficients ::
--      [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> [DList Particle]
--   -> Histogram (Complex Double)
-- computeFourierCoefficients !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !deltaLogRho !xs =
--   Histogram
--     [L.length phiFreqs, L.length rhoFreqs, L.length thetaFreqs, L.length rFreqs]
--     (L.length ys) .
--   toUnboxedVector .
--   UA.accum
--     (+)
--     0
--     ( (1, 1, 1, 1)
--     , ( L.length rFreqs
--       , L.length thetaFreqs
--       , L.length rhoFreqs
--       , L.length phiFreqs)) .
--   L.concat .
--   parMap
--     rdeepseq
--     (\(particle@(Particle phi rho theta r)) ->
--        let !n = floor $ (log r + halfLogPeriod) / deltaLogRho
--            !samples =
--              -- L.map moveParticle  $
--              particle :
--              [ -- (Particle
--              --      phi
--              --      rho
--              --      theta
--              --      (exp $ (fromIntegral i * deltaLogRho - halfLogPeriod)))
--              -- | i <- [0 .. n]
--              ]
--        in L.map
--             (\((rFreq, i), (thetaFreq, j), (rhoFreq, k), (phiFreq, l)) ->
--                ( (i, j, k, l)
--                , -- (deltaLogRho :+ 0) *
--                  1/ (16 * pi^2 * halfLogPeriod :+ 0) *
--                  sampleScale
--                    phiFreq
--                    rhoFreq
--                    thetaFreq
--                    rFreq
--                    halfLogPeriod
--                    samples))
--             freqs) $
--   ys
--   where
--     !ys = DL.toList . DL.concat $ xs
--     !freqs =
--       [ (rFreq', thetaFreq', rhoFreq', phiFreq')
--       | rFreq' <- L.zip rFreqs [1 ..]
--       , thetaFreq' <- L.zip thetaFreqs [1 ..]
--       , rhoFreq' <- L.zip rhoFreqs [1 ..]
--       , phiFreq' <- L.zip phiFreqs [1 ..]
--       ]

-- {-# INLINE normalizeFreqArr #-}
-- normalizeFreqArr ::
--      Double
--   -> [Double]
--   -> [Double]
--   -> R.Array U DIM4 (Complex Double)
--   -> R.Array U DIM4 (Complex Double)
-- normalizeFreqArr !std !phiFreqs !rhoFreqs arr =
--   computeUnboxedS .
--   R.traverse3
--     arr
--     (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
--     (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
--     (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
--     fArr idx *
--     ((exp $
--       (-1) *
--       ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2 + (fRho (Z :. rho)) ^ 2 +
--        (fRho (Z :. r)) ^ 2) /
--       2 /
--       (std ^ 2)) :+
--      0)

{-# INLINE normalizeFreqArr #-}
normalizeFreqArr ::
     Double
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> R.Array U DIM4 (Complex Double)
normalizeFreqArr !stdA !stdR !phiFreqs !rhoFreqs arr
  | stdA == 0 && stdR == 0 = arr
  | stdA == 0 =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
      fArr idx *
      (exp $
       ((-1) * ((fRho (Z :. rho)) ^ 2 + (fRho (Z :. r)) ^ 2) / (2 * stdR ^ 2)) :+
       0)
  | stdR == 0 =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. _ :. theta :. rho :. phi) ->
      fArr idx *
      ((exp $
        (-1) * ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2) /
        (2 * stdA ^ 2)) :+
       0)
  | otherwise =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
      fArr idx *
      ((exp $
        (-1) * ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2) /
        (2 * stdA ^ 2) -
        ((fRho (Z :. rho)) ^ 2 + (fRho (Z :. r)) ^ 2) / (2 * stdR ^ 2)) :+
       0)

-- {-# INLINE computeHarmonicsArray #-}
computeHarmonicsArray ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> IA.Array (Int, Int) (vector (Complex Double))
computeHarmonicsArray !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff =
  let !centerRow = div numRows 2
      !centerCol = div numCols 2
      rangeFunc1 xs ys = [(L.head xs - L.last ys) .. (L.last xs - L.head ys)]
      rangeFunc2 xs ys =
        (round (L.head xs - L.last ys), round (L.last xs - L.head ys))
      (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
      (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
      !xs =
        parMap
          rdeepseq
          (\(!rf, !tf) ->
             let !vec =
                   VG.convert .
                   toUnboxed . computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
                     let !x = fromIntegral (c - centerCol) * deltaCol
                         !y = fromIntegral (r - centerRow) * deltaRow
                         !rho = 0 + (sqrt $ x ^ 2 + y ^ 2)
                         !rho2 =
                           fromIntegral $
                           (c - centerCol) ^ 2 + (r - centerRow) ^ 2
                     in if rho2 > cutoff ^ 2 || (rho <= 0) -- || pi * rho < (abs tf) -- || log ((rho + 1) / rho) > (1 / (2 * rf))
                          then 0
                               -- (x :+ y) ** (tf :+ 0) *
                               -- ((x ^ 2 + y ^ 2) :+ 0) **
                               -- (((-tf - 0.5) :+ rf) / 2)
                          else ((rho :+ 0) ** ((-1) :+ rf)) *
                               (cis (tf * atan2 y x))
             in ((round rf, round tf), vec))
          [ (rf, tf)
          | rf <- rangeFunc1 rhoFreqs rFreqs
          , tf <- rangeFunc1 phiFreqs thetaFreqs
          ]
  in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

-- {-# INLINE computeHarmonicsArraySparse #-}
computeHarmonicsArraySparse ::
     Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> Double
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
computeHarmonicsArraySparse !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff !envelopeSigma =
  let !centerRow = div numRows 2
      !centerCol = div numCols 2
      rangeFunc1 xs ys =
        [(round $ L.head xs - L.last ys) .. (round $ L.last xs - L.head ys)]
      rangeFunc2 xs ys =
        ((round $ (L.head xs - L.last ys)), (round $ (L.last xs - L.head ys)))
      (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
      (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
      !xs =
        parMap
          rseq
          (\(!rf, !tf) ->
             let !arr =
                   computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
                     let !x = fromIntegral (c - centerCol) * deltaCol
                         !y = fromIntegral (r - centerRow) * deltaRow
                         !rho = ((sqrt $ x ^ 2 + y ^ 2))
                         !rho2 =
                           fromIntegral $
                           (c - centerCol) ^ 2 + (r - centerRow) ^ 2
                     in if (rho <= 0) || rho2 > cutoff ^ 2
                          then 0
                          else (x :+ y) ** (fromIntegral tf :+ 0) *
                               (((x ^ 2 + y ^ 2) :+ 0) **
                                (((-(fromIntegral tf) - envelopeSigma) :+ (2 * pi / halfLogPeriod * fromIntegral rf)) /
                                 2))
             in deepSeqArray arr ((rf, tf), arr))
          [ (rf, tf)
          | rf <- rangeFunc1 rhoFreqs rFreqs
          , tf <- rangeFunc1 phiFreqs thetaFreqs
          ]
  in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

-- {-# INLINE computeHarmonicsArrayGPU #-}
-- computeHarmonicsArrayGPU ::
--      (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
--   => Int
--   -> Double
--   -> Int
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> IA.Array (Int, Int) (vector (Complex Double))
-- computeHarmonicsArrayGPU !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff =
--   let !centerRow = div numRows 2
--       !centerCol = div numCols 2
--       rangeFunc1 xs ys =
--         [round (L.head xs - L.last ys) .. round (L.last xs - L.head ys)]
--       rangeFunc2 xs ys =
--         (round (L.head xs - L.last ys), round (L.last xs - L.head ys))
--       (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
--       (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
--       !xs =
--         parMap
--           rdeepseq
--           (\(!rf, !tf) ->
--              let !vec =
--                    VG.convert .
--                    toUnboxed . computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
--                      let !x = fromIntegral (c - centerCol) * deltaCol
--                          !y = fromIntegral (r - centerRow) * deltaRow
--                          !r2 = x ^ 2 + y ^ 2
--                      in if (x == 0 && y == 0) || r2 > cutoff ^ 2
--                           then 0
--                           else cis $
--                                fromIntegral tf * atan2 y x +
--                                (pi * (fromIntegral rf * 0.5 * log r2) /
--                                halfLogPeriod - pi)
--              in ((rf, tf), vec))
--           [ (rf, tf)
--           | rf <-  (rangeFunc1 rhoFreqs rFreqs)
--           , tf <- L.reverse (rangeFunc1 phiFreqs thetaFreqs)
--           ]
--   in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

{-# INLINE getHarmonics #-}
getHarmonics ::
     (VG.Vector vector e)
  => IA.Array (Int, Int) (vector e)
  -> Double
  -> Double
  -> Double
  -> Double
  -> vector e
getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq =
  harmonicsArray IA.! (round (rhoFreq - rFreq), round $ (phiFreq - thetaFreq))

-- {-# INLINE computeFourierSeriesR2 #-}
-- computeFourierSeriesR2 ::
--      Int
--   -> Double
--   -> Int
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> IA.Array (Int, Int) (VU.Vector (Complex Double))
--   -> R.Array U DIM4 (Complex Double)
--   -> [VU.Vector (Complex Double)]
-- computeFourierSeriesR2 numRows deltaRow numCols deltaCol phiFreqs rhoFreqs thetaFreqs rFreqs halfLogPeriod cutoff harmonicsArray arr =
--   let (Z :. (!numRFreq) :. (!numThetaFreq) :. _ :. _) = extent arr
--       !initVec = VU.replicate (VU.length (harmonicsArray IA.! (0, 0))) 0
--   in parMap
--        rdeepseq
--        (\((!r, !rFreq), (!theta, !thetaFreq)) ->
--           L.foldl'
--             (\(!vec) ((!rho, !rhoFreq), (!phi, !phiFreq)) ->
--                VU.zipWith
--                  (+)
--                  vec
--                  (VU.map
--                     (* (arr R.! (Z :. r :. theta :. rho :. phi)))
--                     (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq)))
--             initVec $
--           (,) <$> (L.zip [0 ..] rhoFreqs) <*> (L.zip [0 ..] phiFreqs)) $
--      (,) <$> (L.zip [0 ..] rFreqs) <*> (L.zip [0 ..] thetaFreqs)

-- {-# INLINE computeThetaRHarmonics #-}
computeThetaRHarmonics ::
     Int -> Int -> [Double] -> [Double] -> Double -> [[(Complex Double)]]
computeThetaRHarmonics !numOrientation !numScale !thetaFreqs !rFreqs !halfLogPeriod =
  let !deltaTheta = 2 * pi / fromIntegral numOrientation
      !deltaScale = 2 * halfLogPeriod / fromIntegral numScale
      !numRFreq = L.length rFreqs
      !numThetaFreq = L.length thetaFreqs
  in if numScale == 1
       then parMap
              rdeepseq
              (\o ->
                 L.map
                   (\freq -> cis $ freq * (fromIntegral o * deltaTheta))
                   thetaFreqs) $
            [0 .. numOrientation - 1]
       else parMap
              rdeepseq
              (\(!o, !s) ->
                 R.toList .
                 R.traverse2
                   (fromListUnboxed (Z :. numRFreq) rFreqs)
                   (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
                   (\_ _ -> (Z :. numRFreq :. numThetaFreq)) $ \fRFreq fThetaFreq idx@(Z :. rFreq :. thetaFreq) ->
                   -- exp $
                   -- (-0.5 * (fromIntegral s * deltaScale - halfLogPeriod)) :+
                   -- ((fRFreq (Z :. rFreq)) *
                   --  (fromIntegral s * deltaScale - halfLogPeriod) +
                   --  fThetaFreq (Z :. thetaFreq) * fromIntegral o * deltaTheta)
                   cis
                     ((fRFreq (Z :. rFreq)) *
                      (pi * (fromIntegral s * deltaScale - halfLogPeriod) /
                       halfLogPeriod) +
                      fThetaFreq (Z :. thetaFreq) * fromIntegral o * deltaTheta)
               ) $
            (,) <$> [0 .. numOrientation - 1] <*> [0 .. numScale - 1]

{-# INLINE computeFourierSeriesThetaR #-}
computeFourierSeriesThetaR ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => [[(Complex Double)]]
  -> [vector (Complex Double)]
  -> [vector (Complex Double)]
computeFourierSeriesThetaR !harmonics !vecs =
  let !initVec = VG.replicate (VG.length . L.head $ vecs) 0
  in parMap
       rdeepseq
       (L.foldl'
          (\vec0 (!vec1, !v) -> VG.zipWith (+) vec0 . VG.map (* v) $ vec1)
          initVec .
        L.zip vecs) $
     harmonics

plotThetaDimension ::
     (R.Source s Double)
  => FilePath
  -> String
  -> (Int, Int)
  -> R.Array s DIM3 Double
  -> IO ()
plotThetaDimension folderPath prefix (x', y') inputArr = do
  let centerCol = div cols 2
      centerRow = div rows 2
      x = x' + centerCol
      y = y' + centerRow
      (Z :. numOrientation :. cols :. rows) = extent inputArr
      centerArr =
        fromFunction (Z :. (1 :: Int) :. cols :. rows) $ \(Z :. _ :. i :. j) ->
          if i == x && j == y
            then 1 :: Double
            else 0
  plotImageRepa (folderPath </> printf "%s(%d,%d)_center.png" prefix x' y') .
    ImageRepa 8 . computeS $
    centerArr
  let ys' = R.toList . R.slice inputArr $ (Z :. R.All :. x :. y)
      !maxY = L.maximum ys'
      -- ys = L.map (/maxY) ys'
      ys = ys'
      deltaTheta = (360 ::Double) / fromIntegral numOrientation
      xs = [fromIntegral x * deltaTheta | x <- [0 .. numOrientation - 1]]
  toFile def (folderPath </> printf "%s(%d,%d)_theta.png" prefix x' y') $ do
    layout_title .= printf "%s(%d,%d)" prefix x' y'
    layout_x_axis . laxis_generate .= scaledAxis def (0, 359)
    layout_y_axis . laxis_generate .= scaledAxis def (0, maxY)
    plot (line "" [L.zip xs ys])

-- The codes below compute Fourier series in R2

computeFourierSeriesOfLogPolarHarmonicsArray ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> IA.Array (Int, Int) (vector (Complex Double))
computeFourierSeriesOfLogPolarHarmonicsArray !radius !delta !r2Freq !phiFreq !rhoFreq !thetaFreq !rFreq !halfLogPeriod =
  let ((!angularFreqLB, !angularFreqUB), angularFreqs) =
        freqRange phiFreq thetaFreq
      ((!radialFreqLB, !radialFreqUB), radialFreqs) =
        freqRange phiFreq thetaFreq
      !numPoints' = round $ (2 * radius + 1) / delta
      !numPoints =
        if odd numPoints'
          then numPoints'
          else numPoints' - 1
      !center = div numPoints 2
      !period = 2 * radius + 1
      !periodConstant = delta * 2 * pi / period
      !numR2Freq = 2 * r2Freq + 1
      !std = fromIntegral $ div r2Freq 2
      !simpsonWeights =
        toUnboxed . computeS $
        (computeWeightArrFromListOfShape [numPoints, numPoints] :: R.Array D DIM2 (Complex Double))
      !weightedRectangularHarmonics =
        parMap
          rdeepseq
          (\(freq1, freq2) ->
             VU.zipWith (*) simpsonWeights .
             toUnboxed . computeS . fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
               cis
                 (-periodConstant *
                   fromIntegral (freq1 * (i - center) + freq2 * (j - center))))
          [ (freq1, freq2)
          | freq1 <- [-r2Freq .. r2Freq]
          , freq2 <- [-r2Freq .. r2Freq]
          ]
      !cartesianGrid =
        toUnboxed . computeS . fromFunction (Z :. numPoints :. numPoints) $ \(Z :. c :. r) ->
          let !x = fromIntegral (c - center) * delta
              !y = fromIntegral (r - center) * delta
          in (x, y)
      !simpsonNorm = (delta / 3) ^ 2 :+ 0
      !xs =
        parMap
          rdeepseq
          (\(!radialFreq, !angularFreq) ->
             let !logPolarHarmonics =
                   VU.map (fourierMellin 0.5 angularFreq radialFreq) cartesianGrid
                 coefficients =
                   fromListUnboxed (Z :. numR2Freq :. numR2Freq) .
                   L.map (VU.sum . VU.zipWith (*) logPolarHarmonics) $
                   weightedRectangularHarmonics
                 !coefficientsGaussian =
                   VG.convert .
                   toUnboxed . computeS . R.traverse coefficients id $ \f idx@(Z :. xFreq :. yFreq) ->
                     simpsonNorm * (f idx) *
                     (gaussian2D
                        (fromIntegral $ xFreq - r2Freq)
                        (fromIntegral $ yFreq - r2Freq)
                        std :+
                      0)
             in ((radialFreq, angularFreq), coefficientsGaussian))
          [ (radialFreq, angularFreq)
          | radialFreq <- radialFreqs
          , angularFreq <- angularFreqs
          ]
  in IA.array ((radialFreqLB, angularFreqLB), (radialFreqUB, angularFreqUB)) xs

computeRectangularInverseHarmonics ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> [vector (Complex Double)]
computeRectangularInverseHarmonics !numRows !numCols !delta !radius !maxFreq =
  let !period = 2 * radius + 1
      !periodConstant = delta * 2 * pi / period
      !centerRow = div numRows 2
      !centerCol = div numCols 2
  in parMap
       rdeepseq
       (\(col, row) ->
          VG.fromList
            [ cis
              (periodConstant *
               fromIntegral
                 (freq1 * (row - centerRow) + freq2 * (col - centerCol)))
            | freq1 <- [-maxFreq .. maxFreq]
            , freq2 <- [-maxFreq .. maxFreq]
            ])
       [(col, row) | col <- [0 .. numCols - 1], row <- [0 .. numRows - 1]]

-- Utilities

{-# INLINE freqRange #-}
freqRange :: Int -> Int -> ((Int, Int), [Int])
freqRange n m =
  let !x = n + m
  in ((-x, x), [-x .. x])