{-# LANGUAGE ConstraintKinds, PolyKinds, TypeFamilies #-}
{-# LANGUAGE FlexibleInstances, UndecidableInstances #-}

module CFunctor where

import Generics.SOP.Constraint
import Numeric.LinearAlgebra hiding (C)

class CFunctor f where
  type C f :: * -> Constraint
  cfmap :: (C f a, C f b) => (a -> b) -> f a -> f b

{-
instance {-# OVERLAPPABLE #-} Functor f => CFunctor f where
  type C f = Top
  cfmap = fmap
-}

instance CFunctor Vector where
  type C Vector = And Element (Container Vector)
  cfmap = cmap


{-# language BangPatterns #-}
{-# language DataKinds #-}
{-# language FlexibleContexts #-}
{-# language KindSignatures #-}
{-# language LambdaCase #-}
{-# language PolyKinds #-}
{-# language StandaloneDeriving #-}
{-# language TypeFamilies #-}
{-# language TypeFamilyDependencies #-}
{-# language UndecidableInstances #-}

module Qprob where

import Control.Monad (join)
import Data.Coerce (coerce)
import Data.Kind (Type)
import Linear.Epsilon (nearZero)
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Complex as Complex

data Setting = Classical | Quantum

--type family Setting' s = (s' :: Type) | s' -> s

class V (s :: Setting) where
  type Setting' s = (s' :: Type) | s' -> s
  oneS :: Setting' s
  plusS :: Setting' s -> Setting' s -> Setting' s
  zeroS :: Setting' s
  timesS :: Setting' s -> Setting' s -> Setting' s
  negateS :: Setting' s -> Setting' s
  nearZeroS :: Setting' s -> Bool

instance V 'Classical where
  type Setting' 'Classical = Double
  oneS = one; plusS = plus; zeroS = zero; timesS = times; negateS = negate; nearZeroS = nearZero;
  {-# inline oneS #-}
  {-# inline plusS #-}
  {-# inline zeroS #-}
  {-# inline timesS #-}
  {-# inline negateS #-}
  {-# inline nearZeroS #-}

instance V 'Quantum where
  type Setting' 'Quantum = Complex Double
  oneS = one; plusS = plus; zeroS = zero; timesS = times; negateS = negate; nearZeroS = nearZero;
  {-# inline oneS #-}
  {-# inline plusS #-}
  {-# inline zeroS #-}
  {-# inline timesS #-}
  {-# inline negateS #-}
  {-# inline nearZeroS #-}

data Space (s :: Setting) a = Space !a !(Setting' s)

deriving instance (Show a, Show (Setting' s)) => Show (Space s a)
deriving instance (Eq a, Eq (Setting' s)) => Eq (Space s a)

mapSpace :: (Setting' s -> Setting' s) -> Space s a -> Space s a
mapSpace f (Space a s) = Space a (f s)

-- | @W s a@ is a vector space whose basis elements are labelled
--   by objects of type @a@ and where the coefficients are of type @Settings'' s@.
--
--   This is very similar to standard probability monads except that we
--   allow probabilities to be types other than 'Double'.
newtype W (s :: Setting) a = W { runW :: [Space s a] }

deriving instance (Show a, Show (Setting' s)) => Show (W s a)
deriving instance (Eq a, Eq (Setting' s)) => Eq (W s a)

-- | Transform the probabilities inside of 'W'.
mapW :: (Setting' s -> Setting' s) -> W s a -> W s a
mapW f (W l) = W (List.map (mapSpace f) l)

instance Semigroup (W s a) where
  W x <> W y = W (x <> y)
  {-# inline (<>) #-}

instance Monoid (W s a) where
  mempty = W mempty
  {-# inline mempty #-}

instance Functor (W s) where
  fmap f w = fmapW f w
  {-# inline fmap #-}

fmapW :: (a -> b) -> W s a -> W s b
fmapW f (W l) = W (List.map (\(Space a p) -> Space (f a) p) l)
{-# NOINLINE [1] fmapW #-}

{-# RULES "fmapW/coerce" fmapW coerce = coerce #-}

instance V s => Applicative (W s) where
  pure x = W [Space x oneS]
  {-# inline pure #-}
  W fs <*> W xs = W
    ( do Space f a <- fs
         Space x b <- xs
         pure (Space (f x) (timesS a b))
    )

infixl 7 .*
(.*) :: V s => Setting' s -> W s a -> W s a
a .* b = mapW (a `timesS`) b

type P a = W 'Classical a
type Q a = W 'Quantum   a

star :: Q a -> Q a
star = mapW Complex.conjugate

kron :: V s => W s a -> W s c -> W s (a,c)
kron (W x) (W y) = W
  ( do Space a r1 <- x
       Space c r2 <- y
       pure (Space (a,c) (r1 `timesS` r2))
  )

collect :: (V s, Ord a) => W s a -> W s a
collect (W l) = W $ toList . fromListWith plusS $ l
  where
    toList = Map.foldrWithKey (\k x xs -> (Space k x):xs) []
    fromListWith f xs = fromListWithKey (\_ x y -> f x y) xs
    fromListWithKey f xs =
      let ins t (Space a s) = Map.insertWithKey f a s t
      in List.foldl' ins mempty xs

instance V s => Monad (W s) where
  l >>= f = W $ List.concatMap (\(Space (W d) p) -> List.map (\(Space x q) -> (Space x (p `timesS` q))) d) (runW $ fmap f l)

-- | When we come to observe the state of a quantum system, the quantum state becomes an ordinary probablistic one.
observe :: Ord a => Q a -> P a
observe = W . List.map (\(Space a w) -> Space a (Complex.magnitude (w `times` w))) . runW . collect

rotate :: Double -> Bool -> Q Bool
rotate theta = \case
  True ->
    let theta' = theta :+ 0
    in subtract (cos (theta' / 2) .* pure True) (sin (theta' / 2) .* pure False)
  False ->
    let theta' = theta :+ 0
    in (cos (theta' / 2) .* pure False) <> (sin (theta' / 2) .* pure True)

repeat :: Int -> (a -> a) -> (a -> a)
repeat 0 _ = id
repeat n f = repeat (n - 1) f . f

repeatM :: Monad m => Int -> (a -> m a) -> m a -> m a
repeatM n f = repeat n (>>= f)

snot :: Bool -> Q Bool
snot = rotate (pi / 2)

snot1 :: Int -> P Bool
snot1 n = pure True & repeatM n snot & observe

subtract :: V s => W s a -> W s a -> W s a
subtract a b = a <> ((negateS oneS) .* b)

-- | Quantum Zeno effect.
--
-- A watched pot never boils.
zeno1 :: Int -> P Bool
zeno1 n = pure True & repeatM n (rotate (pi / fromIntegral n)) & collect & observe

ffor :: Functor f => f a -> (a -> b) -> f b
ffor = flip fmap

liftBind :: (Functor f, Monad g) => f (g a) -> (a -> g b) -> f (g b)
liftBind fga f = fmap (>>= f) fga

zeno2 :: Int -> P Bool
zeno2 n = pure True & repeat n
  (\x -> x `ffor` pure `liftBind` rotate (pi / fromIntegral n) `ffor` observe & join
  ) & collect

type MixedState a = P (Q a)

data Experimenter = Experimenter
  { experimenterMemory :: [Bool]
  , experimenterState :: !Bool
  }
  deriving (Eq, Ord)

zeno3 :: Int -> P Bool
zeno3 n = pure (Experimenter [] True) & repeatM n
  (\(Experimenter m s) -> do
    s' <- rotate (pi / fromIntegral n) s
    pure $ Experimenter (s:m) s'
  ) & observe & fmap experimenterState & collect

trimZero :: V s => W s a -> W s a
trimZero = W . List.filter (\(Space _ v) -> not $ nearZeroS v) . runW

simplify :: Ord a => Q a -> Q a
simplify = trimZero . collect


{-# LANGUAGE BangPatterns, DeriveFunctor, DeriveFoldable, DeriveTraversable, DeriveDataTypeable #-}
-- | A scene graph is used to represent any N-dimensional scene. Since it is
--   polymorphic over its contents, we can use it to represent a variety of
--   things, from a physics world to a render tree.
module Game.SceneGraph ( -- * Basic Data Types
                         SceneGraph(..)
                       , TransformView(..)
                       -- * Views
                       , viewTrans
                       , matRep
                       , Quaternion
                       -- * Transformations
                       , translate
                       , scaleGraph
                       , rotate2D
                       , rotate3D
                       , rotateQuat
                       , rotate3DX
                       , rotate3DY
                       , rotate3DZ
                       , rotateEuler
                       , applyPoint
                       , applyVector
                       -- * Graph Visitors
                       , optimize
                       , sceneMap
                       ) where

import Prelewd hiding (toList, join, (<>))

import Impure

import Control.DeepSeq
import Control.Monad.Par
import Data.Function.Pointless
import Data.Typeable
import Storage.List (init)
import Numeric.LinearAlgebra
import Text.Show

-- | Represents an N-dimensional scene graph, with objects at the leaves, and
--   transformations applied in a hierarchical fashion.
--
--   This is an instance of Functor, Foldable, and Traversable, so feel free to
--   walk it like you would any recursive data structure.
data SceneGraph a = Object a
                  | Transform !Transformation !(SceneGraph a)
                  | Branch [SceneGraph a]
    deriving (Show, Functor, Foldable, Typeable)

instance NFData a => NFData (SceneGraph a) where
    rnf (Transform t g) = rnf t `seq` rnf g
    rnf (Object x)      = rnf x
    rnf (Branch xs)     = rnf xs

-- | We represent the basic transformations (translate, scale, rotate) as their
--   "easy" representation, as well as the more complicated matrix. This allows
--   us to optimize the scene graph with the semantic information of chained
--   transformations instead of just blindly multiplying matricies.
--
--   For example, instead of multiplying two (rather complicated) quaternion
--   matricies to chain rotation, we just perform quaternion multiplication.
--
--   By keeping the matricies around too, we can always fall back on matrix
--   multiplication, and we can also avoid recomputing them.
--
--   Oh yeah, and all the matricies use homogenous coordinates, so they're
--   (N+1)x(N+1) in an n-dimensional scene. The more you know.
data Transformation = Translate !(Vector Double)
                                 (Matrix Double)
                    | Scale     !(Vector Double)
                                 (Matrix Double)
                    | Rotate2D  !Double -- θ - radians, clockwise.
                                 (Matrix Double) -- affine representation
                    | AffineTransform !(Matrix Double) -- fall back generic transformations
                                                       -- also used for rotations.
    deriving (Show, Typeable)

-- | A data type used for deconstructing transformations and seeing what they
--   "really are".
data TransformView = TranslateView !(Vector Double)
                   | ScaleView !(Vector Double)
                   | Rotate2DView !Double -- ^ θ - radians, clockwise.
                   | AffineTransformView !(Matrix Double) --  ^ Uses homogenous coordinates.

-- | Use this function to deconstruct a transformation into what it represents.
--
--   We use a more complicated internal data type for representing
--   transformations, and this view can be used to simplify it. The GHC extension
--   "ViewPatterns" will probably be of help here.
--
--   > -- f returns a scaled (by a factor of 2) translation vector.
--   > f :: Transformation -> Vector Double
--   > f (viewTrans -> Translate v) = v*2
viewTrans :: Transformation -> TransformView
viewTrans (Translate v _)     = TranslateView v
viewTrans (Scale v _)         = ScaleView v
viewTrans (Rotate2D θ _)      = Rotate2DView θ
viewTrans (AffineTransform m) = AffineTransformView m

-- | Retrieves the matrix view of a transformation. This will be an (N+1)x(N+1)
--   matrix, where N is the dimensions in the scene.
matRep :: Transformation -> Matrix Double
matRep (Translate _ m)     = m
matRep (Scale _ m)         = m
matRep (Rotate2D _ m)      = m
matRep (AffineTransform m) = m

seq_ :: a -> ()
seq_ = (`seq` ())

instance NFData Transformation where
    rnf (Translate _ m)     = seq_ m
    rnf (Scale     _ m)     = seq_ m
    rnf (Rotate2D  _ m)     = seq_ m
    rnf (AffineTransform _) = ()

-- | Just a 4-element vector, specifying a 3-dimensional rotation. See the
--   wikipedia page for more information on Quaternions.
type Quaternion = Vector Double

-- | Represents a 2-dimensional, clockwise rotation by θ radians. This is
--   probably easier to use than building a rotation matrix manually.
rotate2D :: Double -- ^ θ.
         -> Transformation
rotate2D !θ = Rotate2D θ $ rotate2DMat θ

rotate2DMat :: Double
            -> Matrix Double
rotate2DMat !θ = (3><3) [  c, s, 0
                        , -s, c, 0
                        , 0,  0, 1 ]
    where
        c = cos θ
        s = sin θ

-- | Rotates by a given quaternion.
--
--   Note: This algorithm is untested, and disagrees with wikipedia right now.
--         I'll look into it one day.
rotateQuat :: Quaternion -> Transformation
rotateQuat q = AffineTransform $ (4><4) [ 1-2.0*(yy-zz), 2.0*(xy-zw)    , 2.0*(xz+yw), 0
                                        , 2.0*(xy+zw)  , 1.0-2.0*(xx-zz), 2.0*(yz-xw), 0
                                        , 2.0*(xz-yw)  , 2*(xw+yz)      , 1-2*(xx-yy), 0
                                        , 0            , 0              , 0          , 1 ]
    where
        [ w, x, y, z ] = toList q
        xx = x*x
        yy = y*y
        zz = z*z
        zw = z*w
        xy = x*y
        xz = x*z
        xw = x*w
        yw = y*w
        yz = y*z

-- | Rotates by a given angle around an arbitrary 3-dimensional axis.
rotate3D :: Double -- ^ The amount (in radians) to rotate by.
         -> Vector Double -- ^ The 3-dimensional axis around which we are rotating.
         -> Transformation
-- https://en.wikipedia.org/wiki/Rotation_matrix#Rotation_matrix_from_axis_and_angle
rotate3D !θ !axis = AffineTransform $ (4><4) [ c+x*xoc , x*yoc-zs, x*zoc+ys, 0
                                             , y*xoc+zs, c+y*yoc , y*zoc-xs, 0
                                             , z*xoc-ys, z*yoc+xs, c+z*zoc , 0
                                             ,     0   ,    0    ,    0    , 1 ]
    where
        c = cos θ
        oc = 1-c
        s = sin θ
        [x, y, z] = toList axis
        xs = x*s
        ys = y*s
        zs = z*s
        xoc = x*oc
        yoc = y*oc
        zoc = z*oc

-- | Rotates θ radians clockwise around the X-axis.
rotate3DX :: Double -- θ
          -> Transformation
rotate3DX !θ = AffineTransform $ (4><4) [ 1,  0, 0, 0
                                        , 0,  c, s, 0
                                        , 0, -s, c, 0
                                        , 0,  0, 0, 1 ]
    where
        c = cos θ
        s = sin θ

-- | Rotates θ radians clockwise around the Y axis.
rotate3DY :: Double -- θ
          -> Transformation
rotate3DY !θ = AffineTransform $ (4><4) [  c, 0, -s, 0
                                        ,  0, 1,  0, 0
                                        , -s, 0,  c, 0
                                        ,  0, 0,  0, 1 ]
    where
        c = cos θ
        s = sin θ

-- | Rotates θ radians clockwise around the Y axis.
rotate3DZ :: Double -- θ
          -> Transformation
rotate3DZ !θ = AffineTransform $ (4><4) [  c,  s, 0, 0
                                        , -s,  c, 0, 0
                                        ,  0,  0, 1, 0
                                        ,  0,  0, 0, 1 ]
    where
        c = cos θ
        s = sin θ

-- | Rotates around a yaw, pitch, and roll - all in one transformation.
rotateEuler :: Double -- X-axis rotation.
            -> Double -- Y-axis rotation.
            -> Double -- Z-axis rotation.
            -> Transformation
rotateEuler x y z = let (AffineTransform x') = rotate3DX x
                        (AffineTransform y') = rotate3DY y
                        (AffineTransform z') = rotate3DZ z
                     in AffineTransform $ z' * y' * x'

-- | Adds a number to the end of a vector. Tends to be useful when constructing
--   homogenous coordinates.
vappend :: Element a => a -> Vector a -> Vector a
vappend x = join . flip (:) [ constant x 1 ]

-- | Applies the given modelview matrix to the given vector.
applyVector :: Matrix Double -> Vector Double -> Vector Double
applyVector m v = subVector 0 (dim v) $ m <> vappend 1.0 v

-- | Applied the given modelview matrix to the given point.
applyPoint :: Matrix Double -> Vector Double -> Vector Double
applyPoint m p = subVector 0 (dim p) $ m <> vappend 0.0 p

scaleMat :: Vector Double -> Matrix Double
scaleMat = diag . vappend 1.0

-- | Performs an n-dimensional scaling of the scene graph. Bigger numbers mean
--   bigger objects. 1.0 means no scaling.
scaleGraph :: Vector Double -- ^ The N-dimensional vector of values to scale by.
           -> Transformation
scaleGraph v = Scale v $ scaleMat v

translateMat :: Vector Double -> Matrix Double
translateMat v = fromColumns $ ((<?> error "Empty vector") . init . toColumns . ident $ dim v + 1) `mappend` [vappend 1.0 v]

-- | Translates an n-dimensional scene graph by the given amount. Bigger values
--   mean bigger objects. 1.0 means no scaling.
translate :: Vector Double -- ^ The N-dimensional vector of values to translate by.
          -> Transformation
translate v = Translate v $ translateMat v

-- | Cleans up a scene graph, and does its best to avoid as many
--   multiplications and indirections as possible.
--
--   This will walk the whole graph, but if you're doing so anyhow, it'd be a
--   great way to avoid a lot of extra processing due to inefficient structure.
optimize :: SceneGraph a -> SceneGraph a
optimize x@(Object _) = x
optimize (Branch xs)  = optBranch xs
    where
        optBranch :: [SceneGraph a] -> SceneGraph a
        optBranch []  = Branch [] -- lolwut
        optBranch [x] = optimize x
        optBranch xs' = Branch $ foldr f [] xs'
            where
                -- merges sub-branches up into their parent.
                f y ys = case optimize y of
                            Branch zs -> zs `mappend` ys
                            g         -> g:ys
optimize (Transform t g) = go
    where
        go = let subGraph = optimize g
              in case subGraph of
                Transform t' subsubGraph -> case (t, t') of
                    (Translate v _, Translate v' _) -> Transform (combineTrans v v') subsubGraph
                    (Scale v _, Scale v' _)         -> Transform (combineScale v v') subsubGraph
                    (Rotate2D θ _, Rotate2D τ _)    -> Transform (combineRot θ τ) subsubGraph
                    (_, _)                          -> Transform t subGraph
                _ -> Transform t subGraph
        combineTrans = translate .: (+)
        combineScale = scaleGraph .: (+)
        combineRot   = rotate2D .: (+)

-- | Maps a function taking an object and its modelview matrix, returning a new
--   scene graph of a possibly different type.
--
--   Please optimize the scene graph before calling this function for optimal
--   efficiency.
--
--   To transform a point by the given matrix, use 'homPoint'. To
--   transform a vector, use 'homVector'.
--
--   Don't do the multiplication yourself. You'll probably get it wrong.
sceneMap :: NFData b
         => Int -- ^ The number of dimensions in the scene.
         -> (a -> Matrix Double -> b) -> SceneGraph a -> Par (SceneGraph b)
sceneMap n f = sMap (ident n)
    where
        -- Give a stack of matricies to multiply and a scene graph, returns the
        -- transformed (by f) scene graph.
        sMap m (Object x)  = return . Object $ f x m
        sMap m (Transform t child) = Transform t <$> sMap (matRep t * m) child
        sMap m (Branch xs) = Branch <$> parMapM (sMap m) xs
{-# INLINE sceneMap #-}


--
-- Benchmarks for LUSolve.
--

module LUSolveBenchmark (
    benchmarks
    ) where

import           Criterion                     (Benchmark, bench, bgroup, env,
                                                nf)

import           Numeric.LinearAlgebra.LUSolve (luFactor)

import qualified Data.Matrix.Generic           as M
import qualified Data.Vector.Unboxed           as V
import           System.Random

bundle :: Int -> [ a ] -> [[ a ]]
bundle _ [] = []
bundle n xs = take n xs : bundle n (drop n xs)

type Mat = M.Matrix V.Vector Double

runLUFactor :: Mat -> Mat
runLUFactor = (\(x, _, _) -> x) . luFactor

setupEnv :: IO (Mat, Mat, Mat)
setupEnv = do
  let mVals = randoms (mkStdGen 1) -- not a top level CAF, so will be GC'd promptly
      randomSquareMatrices :: Int -> [ Mat ]
      randomSquareMatrices n = Prelude.map (\vs -> M.fromLists (bundle n vs)) (bundle (n * n) mVals)
      m100:_ = randomSquareMatrices 100
      m500:_ = randomSquareMatrices 500
      m1000:_= randomSquareMatrices 1000
  return (m100, m500, m1000)

benchmarks :: Benchmark
benchmarks = env setupEnv $ \ ~(m100, m500, m1000) ->
  bgroup "luFactor"
    [ bench "100 x 100"   $ nf runLUFactor m100
    , bench "500 x 500"   $ nf runLUFactor m500
    , bench "1000 x 1000" $ nf runLUFactor m1000
    ]


{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

module Roc (testRoc) where

import Data.Csv ((.:), FromNamedRecord(..), decodeByName)
import Codec.Compression.Lzma (decompress)
import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck
import Statistics.Classification (ClassificationScore(..))
import Statistics.Classification.ROC (ROC(..), roc)
import System.Directory (doesFileExist)
import System.FilePath ((</>), (<.>))
import qualified Data.ByteString.Lazy as LB
import qualified Data.Vector as V

checkAUC :: V.Vector (ClassificationScore Bool) -> Bool
checkAUC v = 0 <= auc && auc <= 1
  where auc = rocAUC $ roc v

newtype CScore = CScore { unCS :: ClassificationScore Bool }

instance FromNamedRecord CScore where
  parseNamedRecord o = fmap CScore $ ClassificationScore <$> fmap (== ("up" :: String)) (o .: "label") <*> o .: "prob_up"

testExamples :: TestTree
testExamples = testGroup "examples" $ map (\p -> testCase p (ex p)) validp
  where
    validp = filter (/= "ri_czc-2") $ pth <$> inst <*> [1..5 :: Int]
    inst = ["b_dce", "j_dce", "ri_czc"]
    pth prod n = prod ++ "-" ++ show n
    ex b =
      let p = "data" </> b <.> "csv" <.> "xz"
      in doesFileExist p >>= \case
      True -> decodeByName . decompress <$> LB.readFile p >>= \case
        Left err -> assertFailure err
        Right (_, cs) -> assertBool "auc in range" $ checkAUC (V.map unCS cs)
      False -> assertBool "no file" True

propAucRange :: [ClassificationScore Bool] -> Property
propAucRange l = length l >= 1 ==> checkAUC (V.fromList l)

testRoc :: TestTree
testRoc = testGroup "roc" [
    testProperty "auc range" propAucRange
  , testExamples
  ]


{-# LANGUAGE Arrows #-}

module Control.SF.AuxFunctions (
    Event, edge, quantize,
    presentFFT, fftA,
    toMSF, toRealTimeMSF,
    Control, buffer
) where

import Prelude hiding (init)
import Control.Arrow
import Control.CCA.Types

import Numeric.FFT (fft)
import Data.Complex
import Data.Map (Map)
import qualified Data.Map as Map

-- | Alternative for working with Math.FFT instead of Numeric.FFT
--import qualified Math.FFT as FFT
--import Data.Array.IArray
--import Data.Array.CArray
--myFFT n lst = elems $ (FFT.dft) (listArray (0, n-1) lst)

-- For use with SF Conversions
import Control.Monad.Fix
import Control.SF.SF
import Control.SF.MSF

import Control.Concurrent.MonadIO
import Data.IORef.MonadIO
import qualified Data.Sequence as S
import Data.Foldable (toList)
import Control.DeepSeq

--------------------------------------
-- Generic Types and Functions
--------------------------------------

type Event a = Maybe a

edge :: ArrowInit a => a Bool Bool
edge = proc b -> do
    prev <- init False -< b
    returnA -< prev && not b

-- | Scrutinizes n samples at a time, updating after k new values from a signal function
quantize :: ArrowInit a => Int -> Int -> a b (Event [b])
quantize n k = proc d -> do
    rec (ds,c) <- init ([],0) -< (take n (d:ds), c+1)
    returnA -< if c >= n && c `mod` k == 0 then Just ds else Nothing

--------------------------------------
-- Fast Fourier Transform
--------------------------------------

-- | Converts the vector result of a dft into a map from frequency to magnitude.
--   One common use is:
--      fftA >>> arr (fmap $ presentFFT clockRate)
presentFFT :: Double -> [Double] -> Map Double Double
presentFFT clockRate a = Map.fromList $ map mkAssoc (zip [0..(length a)] a) where
    mkAssoc (i,c) = (freq, c) where
        samplesPerPeriod = fromIntegral (length a)
        freq = fromIntegral i * (clockRate / samplesPerPeriod)

-- | Given a quantization frequency (the number of samples between each
--   successive FFT calculation) and a fundamental period, this will decompose
--   the input signal into its constituent frequencies.
--   NOTE: The fundamental period must be a power of two!
fftA :: ArrowInit a => Int -> Int -> a Double (Event [Double])
fftA qf fp = proc d -> do
    carray <- quantize fp qf -< d :+ 0
    returnA -< fmap (map magnitude . take (fp `div` 2) . fft) carray

--------------------------------------
-- Signal Function Conversions
--------------------------------------

-- | The following two functions are for lifting SFs to MSFs.  The first
--   one is a quick and dirty solution, and the second one appropriately
--   converts a simulated time SF into a real time one.
toMSF :: Monad m => SF a b -> MSF m a b
toMSF (SF sf) = MSF h
    where
      h a = return (b, toMSF sf')
        where (b, sf') = sf a

-- | The clockrate is the simulated rate of the input signal function.
--   The buffer is the number of time steps the given signal function is
--   allowed to get ahead of real time.  Thus, the real amount of time
--   that it can get ahead is the buffer divided by the clockrate seconds.
--   The threadHandler is a where the ThreadId of the forked thread is sent.
--
--   The output signal function takes and returns values in real time.
--   The input must be paired with time, and the return values are the
--   list of bs generated in the given time step and a boolean that is
--   true when time is synced and false when the simulation is running
--   slower than real time.  Note that the returned list will be long
--   if the clockrate is much faster than real time and potentially
--   empty if it's slower.
toRealTimeMSF :: (Monad m, MonadIO m, MonadFix m, NFData b) =>
                 Double -> Int -> (ThreadId -> m ()) -> SF a b
              -> MSF m (a, Double) ([b], Bool)
toRealTimeMSF clockrate buffer threadHandler sf = proc (a, t) -> do
    t' <- init 0 -< t
    rec f <- init 0 -< f'
        let tseg = f + t-t'
            (n,f') = properFraction $ tseg * clockrate
    MSF initFun -< (a, n)
  where
      initFun (a, n) = do
        inp <- newIORef a
        out <- newIORef S.empty
        nvar <- newEmptyMVar
        tid <- liftIO $ forkIO $ worker inp out nvar 1 sf
        threadHandler tid
        h inp out nvar (a, n)
      h inp out nvar (a, n) = do
        writeIORef inp a
        if (n > 0)
          then do
            tryPutMVar nvar n
            b <- atomicModifyIORef out (swap . S.splitAt n)
            return ((toList b, S.length b == n), MSF (h inp out nvar))
          else return (([], True),MSF (h inp out nvar))
      worker inp out nvar n (SF sf) = do
        a <- readIORef inp
        let (b, sf') = sf a
        s <- deepseq b $ atomicModifyIORef out (\s -> (s S.|> b, s))
        n' <- if S.length s >= n+buffer then takeMVar nvar else return n
        worker inp out nvar n' sf'
      swap (a,b) = (b,a)

--------------------------------------
-- A Buffering Arrow (beta)
--------------------------------------

data Control = Play | Pause | Record | Dump | Jump Integer
             | ClearEarlier | ClearLater | ClearAll
-- buffer takes a control signal as well as an input stream of data and
-- returns values appropriate to the control signal:
--
buffer :: ArrowInit a => a (Control, b) [b]
buffer = proc (c,x) -> do
    rec s <- init ([],[]) -< s'
        let (ret, s') = case c of
                Play -> next s
                Pause -> ([], s)
                Record -> ([], cons x s)
                Dump -> (llToList s, s)
                Jump n -> ([], nav n s)
                ClearEarlier -> (fst s, ([], snd s))
                ClearLater -> (snd s, (fst s, []))
                ClearAll -> (llToList s, ([],[]))
    returnA -< ret

-- Helper stuff for buffer
-- =======================
type LensList a = ([a],[a])
cons :: a -> LensList a -> LensList a
cons a (e,l) = (a:e, l)

prev :: LensList a -> ([a], LensList a)
prev l@([],_) = ([], l)
prev (a:e, l) = ([a], (e, a:l))

next :: LensList a -> ([a], LensList a)
next l@(_,[]) = ([], l)
next (e, a:l) = ([a], (a:e, l))

nav :: Integer -> LensList a -> LensList a
nav 0 l = l
nav n l | n > 0 = nav (n-1) (snd $ next l)
nav n l | n < 0 = nav (n+1) (snd $ prev l)

llToList :: LensList a -> [a]
llToList (e, l) = e ++ l


{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# OPTIONS_HADDOCK show-extensions #-}

-- |
-- Module      : VectorSpace
-- Description : Particle Swarm Optimisation
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental

module PSO.VectorSpace
  ( Scalable(..)
  , VectorSpace(..)
  ) where

import           Data.Complex        (Complex (..))
import           Foreign.Storable
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Generic as GV

class Scalable a v where
  scale :: a -> v -> v

instance (Num a) => Scalable a a where
  scale = (*)

instance (Num a) => Scalable a (Complex a) where
  scale c (x :+ y) = (c * x) :+ (c * y)

instance (Num a, Storable a) => Scalable a (V.Vector a) where
  scale c = V.map (c*)

instance (RealFloat a, Storable a) => Scalable a (V.Vector (Complex a)) where
  scale c = scale (c :+ 0)

class (Num a, Num v, Scalable a v) => VectorSpace v a

instance (Num a, Num v, Scalable a v) => VectorSpace v a


{-# LANGUAGE ScopedTypeVariables, BangPatterns, RecordWildCards #-}

{-| Functions used internally by the SDR.Filter module. Most of these are
    not actually used but exist for benchmarking purposes to determine the
    fastest filter implementation.
-}
module SDR.FilterInternal where

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

import qualified Data.Vector.Generic               as VG
import qualified Data.Vector.Generic.Mutable       as VGM
import qualified Data.Vector.Storable              as VS
import qualified Data.Vector.Storable.Mutable      as VSM
import qualified Data.Vector.Fusion.Bundle         as VFB

import           SDR.VectorUtils
import           SDR.Util

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

{-# INLINE filterImperative1 #-}
filterImperative1 :: (PrimMonad m, Functor m, Num a, Mult a b, VG.Vector v a, VG.Vector v b, VGM.MVector vm a) => v b -> Int -> v a -> vm (PrimState m) a -> m ()
filterImperative1 coeffs num inBuf outBuf = go 0
    where
    go offset
        | offset < num = do
            let res = dotProd offset
            VGM.unsafeWrite outBuf offset res
            go $ offset + 1
        | otherwise    = return ()
    dotProd offset = VG.sum $ VG.zipWith mult (VG.unsafeDrop offset inBuf) coeffs

{-# INLINE filterImperative2 #-}
filterImperative2 :: (PrimMonad m, Functor m, Num a, Mult a b, VG.Vector v a, VG.Vector v b, VGM.MVector vm a) => v b -> Int -> v a -> vm (PrimState m) a -> m ()
filterImperative2 coeffs num inBuf outBuf = go 0
    where
    go offset
        | offset < num = do
            let res = dotProd (VG.unsafeDrop offset inBuf)
            VGM.unsafeWrite outBuf offset res
            go $ offset + 1
        | otherwise    = return ()
    dotProd buf = go 0 0
        where
        go !accum j
            | j < VG.length coeffs = go (VG.unsafeIndex buf j `mult` VG.unsafeIndex coeffs j  + accum) (j + 1)
            | otherwise            = accum

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

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

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

foreign import ccall unsafe "filterRR"
    filterRR_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCRR :: FilterRR
filterCRR = filterFFIR filterRR_c

foreign import ccall unsafe "filterRC"
    filterRC_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCRC :: FilterRC
filterCRC = filterFFIC filterRC_c

foreign import ccall unsafe "filterSSERR"
    filterSSERR_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCSSERR :: FilterRR
filterCSSERR = filterFFIR filterSSERR_c

foreign import ccall unsafe "filterSSESymmetricRR"
    filterSSESymmetricRR_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCSSESymmetricRR :: FilterRR
filterCSSESymmetricRR = filterFFIR filterSSESymmetricRR_c

foreign import ccall unsafe "filterSSERC"
    filterSSERC_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCSSERC :: FilterRC
filterCSSERC = filterFFIC filterSSERC_c

foreign import ccall unsafe "filterSSERC2"
    filterSSERC2_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCSSERC2 :: FilterRC
filterCSSERC2 = filterFFIC filterSSERC2_c

foreign import ccall unsafe "filterAVXRR"
    filterAVXRR_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCAVXRR :: FilterRR
filterCAVXRR = filterFFIR filterAVXRR_c

foreign import ccall unsafe "filterAVXSymmetricRR"
    filterAVXSymmetricRR_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCAVXSymmetricRR :: FilterRR
filterCAVXSymmetricRR = filterFFIR filterAVXSymmetricRR_c

foreign import ccall unsafe "filterAVXRC"
    filterAVXRC_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCAVXRC :: FilterRC
filterCAVXRC = filterFFIC filterAVXRC_c

foreign import ccall unsafe "filterAVXRC2"
    filterAVXRC2_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCAVXRC2 :: FilterRC
filterCAVXRC2 = filterFFIC filterAVXRC2_c

foreign import ccall unsafe "filterSSESymmetricRC"
    filterSSESymmetricRC_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCSSESymmetricRC :: FilterRC
filterCSSESymmetricRC = filterFFIC filterSSESymmetricRC_c

foreign import ccall unsafe "filterAVXSymmetricRC"
    filterAVXSymmetricRC_c :: CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

filterCAVXSymmetricRC :: FilterRC
filterCAVXSymmetricRC = filterFFIC filterAVXSymmetricRC_c

-- Decimation

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

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

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

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

foreign import ccall unsafe "decimateRR"
    decimateCRR_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCRR :: DecimateRR
decimateCRR = decimateFFIR decimateCRR_c

foreign import ccall unsafe "decimateRC"
    decimateCRC_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCRC :: DecimateRC
decimateCRC = decimateFFIC decimateCRC_c

foreign import ccall unsafe "decimateSSERR"
    decimateSSERR_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCSSERR :: DecimateRR
decimateCSSERR = decimateFFIR decimateSSERR_c

foreign import ccall unsafe "decimateSSERC"
    decimateSSERC_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCSSERC :: DecimateRC
decimateCSSERC = decimateFFIC decimateSSERC_c

foreign import ccall unsafe "decimateSSERC2"
    decimateSSERC2_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCSSERC2 :: DecimateRC
decimateCSSERC2 = decimateFFIC decimateSSERC2_c

foreign import ccall unsafe "decimateAVXRR"
    decimateAVXRR_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCAVXRR :: DecimateRR
decimateCAVXRR = decimateFFIR decimateAVXRR_c

foreign import ccall unsafe "decimateAVXRC"
    decimateAVXRC_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCAVXRC :: DecimateRC
decimateCAVXRC = decimateFFIC decimateAVXRC_c

foreign import ccall unsafe "decimateAVXRC2"
    decimateAVXRC2_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCAVXRC2 :: DecimateRC
decimateCAVXRC2 = decimateFFIC decimateAVXRC2_c

foreign import ccall unsafe "decimateSSESymmetricRR"
    decimateSSESymmetricRR_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCSSESymmetricRR :: DecimateRR
decimateCSSESymmetricRR = decimateFFIR decimateSSESymmetricRR_c

foreign import ccall unsafe "decimateAVXSymmetricRR"
    decimateAVXSymmetricRR_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCAVXSymmetricRR :: DecimateRR
decimateCAVXSymmetricRR = decimateFFIR decimateAVXSymmetricRR_c

foreign import ccall unsafe "decimateSSESymmetricRC"
    decimateSSESymmetricRC_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCSSESymmetricRC :: DecimateRC
decimateCSSESymmetricRC = decimateFFIC decimateSSESymmetricRC_c

foreign import ccall unsafe "decimateAVXSymmetricRC"
    decimateAVXSymmetricRC_c :: CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

decimateCAVXSymmetricRC :: DecimateRC
decimateCAVXSymmetricRC = decimateFFIC decimateAVXSymmetricRC_c

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

foreign import ccall unsafe "resampleRR"
    resample_c :: CInt -> CInt -> CInt -> CInt -> CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

resampleCRR :: Int -> Int -> Int -> Int -> VS.Vector Float -> VS.Vector Float -> VS.MVector RealWorld Float -> IO ()
resampleCRR num interpolation decimation offset coeffs inBuf outBuf =
    VS.unsafeWith (coerce coeffs) $ \cPtr ->
        VS.unsafeWith (coerce inBuf) $ \iPtr ->
            VSM.unsafeWith (coerce outBuf) $ \oPtr ->
                resample_c (fromIntegral num) (fromIntegral $ VG.length coeffs) (fromIntegral interpolation) (fromIntegral decimation) (fromIntegral offset) cPtr iPtr oPtr

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

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

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

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

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

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

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

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

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

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

resampleFFIC :: (Ptr CFloat -> Ptr CFloat -> IO CInt) -> VS.Vector (Complex Float) -> VSM.MVector RealWorld (Complex Float) -> IO Int
resampleFFIC func inBuf outBuf = liftM fromIntegral $
    VS.unsafeWith (coerce inBuf) $ \iPtr ->
        VSM.unsafeWith (coerce outBuf) $ \oPtr ->
            func iPtr oPtr

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

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

type ResampleRR = Int -> Int -> [Float] -> IO (Int -> Int -> VS.Vector Float -> VS.MVector RealWorld Float -> IO Int)

foreign import ccall unsafe "resample2RR"
    resample2_c :: CInt -> CInt -> CInt -> CInt -> Ptr CInt -> Ptr (Ptr CFloat) -> Ptr CFloat -> Ptr CFloat -> IO CInt

resampleCRR2 :: ResampleRR
resampleCRR2 = mkResampler resample2_c 1

foreign import ccall unsafe "resampleSSERR"
    resampleCSSERR_c :: CInt -> CInt -> CInt -> CInt -> Ptr CInt -> Ptr (Ptr CFloat) -> Ptr CFloat -> Ptr CFloat -> IO CInt

resampleCSSERR :: ResampleRR
resampleCSSERR = mkResampler resampleCSSERR_c 4

foreign import ccall unsafe "resampleAVXRR"
    resampleAVXRR_c :: CInt -> CInt -> CInt -> CInt -> Ptr CInt -> Ptr (Ptr CFloat) -> Ptr CFloat -> Ptr CFloat -> IO CInt

resampleCAVXRR :: ResampleRR
resampleCAVXRR = mkResampler resampleAVXRR_c 8

type ResampleRC = Int -> Int -> [Float] -> IO (Int -> Int -> VS.Vector (Complex Float) -> VS.MVector RealWorld (Complex Float) -> IO Int)

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

foreign import ccall unsafe "resample2RC"
    resample2RC_c :: CInt -> CInt -> CInt -> CInt -> Ptr CInt -> Ptr (Ptr CFloat) -> Ptr CFloat -> Ptr CFloat -> IO CInt

resampleCRC :: ResampleRC
resampleCRC = mkResamplerC resample2RC_c 1

foreign import ccall unsafe "resampleSSERC"
    resampleCSSERC_c :: CInt -> CInt -> CInt -> CInt -> Ptr CInt -> Ptr (Ptr CFloat) -> Ptr CFloat -> Ptr CFloat -> IO CInt

resampleCSSERC :: ResampleRC
resampleCSSERC = mkResamplerC resampleCSSERC_c 4

foreign import ccall unsafe "resampleAVXRC"
    resampleAVXRC_c :: CInt -> CInt -> CInt -> CInt -> Ptr CInt -> Ptr (Ptr CFloat) -> Ptr CFloat -> Ptr CFloat -> IO CInt

resampleCAVXRC :: ResampleRC
resampleCAVXRC = mkResamplerC resampleAVXRC_c 8

{-
 - Cross buffer
-}

{-# INLINE decimateCrossHighLevel #-}
decimateCrossHighLevel :: (PrimMonad m, Functor m, Num a, Mult a b, VG.Vector v a, VG.Vector v b, VGM.MVector vm a) => Int -> v b -> Int -> v a -> v a -> vm (PrimState m) a -> m ()
decimateCrossHighLevel factor coeffs num lastBuf nextBuf outBuf = fill x outBuf
    where
    x = VFB.map dotProd (VFB.iterateN num (+ factor) 0)
    dotProd i = VG.sum $ VG.zipWith mult (VG.unsafeDrop i lastBuf VG.++ nextBuf) coeffs

{-# INLINE filterCrossHighLevel #-}
filterCrossHighLevel :: (PrimMonad m, Functor m, Num a, Mult a b, VG.Vector v a, VG.Vector v b, VGM.MVector vm a) => v b -> Int -> v a -> v a -> vm (PrimState m) a -> m ()
filterCrossHighLevel coeffs num lastBuf nextBuf outBuf = fill (VFB.generate num dotProd) outBuf
    where
    dotProd i = VG.sum $ VG.zipWith mult (VG.unsafeDrop i lastBuf VG.++ nextBuf) coeffs

{-# INLINE resampleCrossHighLevel #-}
resampleCrossHighLevel :: (PrimMonad m, Num a, Mult a b, VG.Vector v a, VG.Vector v b, VGM.MVector vm a) => Int -> Int -> v b -> Int -> Int -> v a -> v a -> vm (PrimState m) a -> m Int
resampleCrossHighLevel interpolation decimation coeffs filterOffset count lastBuf nextBuf outBuf = fill 0 filterOffset 0
    where
    fill i filterOffset inputOffset
        | i < count = do
            let dp = dotProd filterOffset inputOffset
            VGM.unsafeWrite outBuf i dp
            let (q, r)        = divMod (decimation - filterOffset - 1) interpolation
                inputOffset'  = inputOffset + q + 1
                filterOffset' = interpolation - 1 - r
            filterOffset' `seq` inputOffset' `seq` fill (i + 1) filterOffset' inputOffset'
        | otherwise = return filterOffset
    dotProd filterOffset i = VG.sum $ VG.zipWith mult (VG.unsafeDrop i lastBuf VG.++ nextBuf) (stride interpolation (VG.unsafeDrop filterOffset coeffs))

foreign import ccall unsafe "dcBlocker"
    c_dcBlocker :: CInt -> CFloat -> CFloat -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()

dcBlocker :: Int -> Float -> Float -> VS.Vector Float -> VS.MVector RealWorld Float -> IO (Float, Float)
dcBlocker num lastSample lastOutput inBuf outBuf =
    alloca $ \fsp ->
        alloca $ \fop ->
            VS.unsafeWith (coerce inBuf) $ \iPtr ->
                VSM.unsafeWith (coerce outBuf) $ \oPtr -> do
                    c_dcBlocker (fromIntegral num) (realToFrac lastSample) (realToFrac lastOutput) fsp fop iPtr oPtr
                    r1 <- peek fsp
                    r2 <- peek fop
                    return (realToFrac r1, realToFrac r2)


module LispParser where

import Control.Monad (liftM)
import Control.Monad.Error.Class (throwError)
import Data.Complex (Complex (..))
import Data.Ratio ((%))
import LispVal
import Numeric (readFloat, readHex, readOct)
import Text.ParserCombinators.Parsec hiding (spaces)

readOrThrow :: Parser a -> String -> ThrowsError a
readOrThrow parser input = case parse parser "lisp" input of
    Left err  -> throwError $ Parser err
    Right val -> return val

readExpr :: String -> ThrowsError LispVal
readExpr = readOrThrow parseExpr

readExprList :: String -> ThrowsError [LispVal]
readExprList = readOrThrow (endBy parseExpr spaces)

parseLisp :: SourceName -> String -> Either ParseError LispVal
parseLisp = parse parseExpr

parseExpr :: Parser LispVal
parseExpr =
  parseAtom
    <|> parseString
    <|> parseCharacter
    <|> try parseFloat
    <|> try parseRational
    <|> try parseComplex
    <|> parseNumber
    <|> parseQuote
    <|> parseQuasiquote
    <|> do
      _ <- char '('
      x <- try parseList <|> parseDottedList
      _ <- char ')'
      return x

parseList :: Parser LispVal
parseList = liftM List $ sepBy parseExpr spaces

parseDottedList :: Parser LispVal
parseDottedList = do
  x <- endBy parseExpr spaces
  xs <- char '.' >> spaces >> parseExpr
  return $ DottedList x xs

parseQuote :: Parser LispVal
parseQuote = do
  _ <- char '\''
  x <- parseExpr
  return $ List [Atom "quote", x]

parseQuasiquote :: Parser LispVal
parseQuasiquote = do
  _ <- char ','
  x <- parseExpr
  return $ List [Atom "quasiquote", x]

parseAtom :: Parser LispVal
parseAtom = do
  first <- letter <|> symbol
  rest <- many (letter <|> digit <|> symbol)
  let atom = first : rest
  return $ case atom of
    "True" -> Bool True
    "False" -> Bool False
    _ -> Atom atom

parseNumber :: Parser LispVal
parseNumber = parseDecimal <|> parseHex <|> parseOct <|> parseDecimalNotation

parseDecimal :: Parser LispVal
parseDecimal = many1 digit >>= (return . Number . read)

parseDecimalNotation :: Parser LispVal
parseDecimalNotation = do
  _ <- try $ string "#d"
  x <- many1 digit
  (return . Number . read) x

parseHex :: Parser LispVal
parseHex = do
  _ <- try $ string "#x"
  x <- many1 hexDigit
  return $ Number (hexToDigit x)

hexToDigit :: (Eq a, Num a) => String -> a
hexToDigit x = fst $ head (readHex x)

parseOct :: Parser LispVal
parseOct = do
  _ <- try $ string "#o"
  x <- many1 octDigit
  return $ Number (octToDigit x)

octToDigit :: (Eq a, Num a) => String -> a
octToDigit x = fst $ head (readOct x)

parseString :: Parser LispVal
parseString = do
  _ <- char '"'
  x <- many $ noneOf "\"\\" <|> escapedChars
  _ <- char '"'
  return $ String x

parseFloat :: Parser LispVal
parseFloat = do
  whole <- many1 digit
  _ <- char '.'
  frac <- many1 digit
  return . Float . fst . head . readFloat $ whole ++ "." ++ frac

parseRational :: Parser LispVal
parseRational = do
  numerator <- many1 digit
  _ <- char '/'
  denominator <- many1 digit
  return $ Rational (read numerator % read denominator)

parseComplex :: Parser LispVal
parseComplex = do
  real <- many1 digit
  _ <- char '+'
  imaginary <- many1 digit
  _ <- char 'i'
  return $ Complex (read real :+ read imaginary)

parseCharacter :: Parser LispVal
parseCharacter = try (string "\\#") >>= (return . String)

escapedChars :: Parser Char
escapedChars = do
  _ <- char '\\'
  x <- oneOf "\\\"\'nrt"
  return $ case x of
    '\\' -> x
    '"' -> x
    '\'' -> x
    'n' -> '\n'
    't' -> '\t'
    'r' -> '\r'
    _ -> x

symbol :: Parser Char
symbol = oneOf "!$%&|*+-/:<=>?@^_~"

spaces :: Parser ()
spaces = skipMany1 space


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Signal
-- Copyright   :  (c) Alexander Vivian Hugh McPhail 2010, 2014, 2015, 2016
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  uses FFI
--
-- Signal processing functions
--
-----------------------------------------------------------------------------

module Numeric.Signal (
                       S.Convolvable(..),
                       S.Filterable(),
                       -- * Filtering
                       hamming,
                       pwelch,
                       fir,standard_fir,broadband_fir,
                       freqzF,freqzN,
                       filter,broadband_filter,
                       -- * Analytic Signal
                       analytic_signal,analytic_power,analytic_phase,
                       unwrap,
                       -- * Statistics
                       cross_covariance,cross_correlation,cross_spectrum,
                       auto_covariance,auto_correlation,
                       -- * Preprocessing
                       detrend,
                       resize,
                       downsample,
                       deriv,
                       cumulative_sum
                ) where

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

import qualified Numeric.Signal.Internal as S

import Numeric.GSL.Fitting.Linear

import Data.Complex
import Foreign.Storable()

--import Data.Function.Unicode

import qualified Data.List as L

--import Data.Packed.Vector
--import Data.Packed(Container(..))
import Numeric.LinearAlgebra

import qualified Data.Vector.Generic as GV

import qualified Numeric.GSL.Fourier as F

import Prelude hiding(filter)

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

-- | filters the signal
filter :: (S.Filterable a)
         => Vector a   -- ^ zero coefficients
       -> Vector a     -- ^ pole coefficients
       -> Int   -- ^ sampling rate
       -> Vector a     -- ^ input signal
       -> Vector a     -- ^ output signal
filter b a s v = let len = size v
                     w = min s len
                     start = (negate . fromList . reverse . toList . subVector 0 w) v
                     finish = (negate . fromList . reverse . toList . subVector (len-w) w) v
                     v' = vjoin [start,v,finish]
                 in subVector s len $ S.filter_ b a v'

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

-- | Welch (1967) power spectrum density using periodogram/FFT method
pwelch :: Int            -- ^ sampling rate
       -> Int            -- ^ window size
       -> Vector Double  -- ^ input signal
       -> (Vector Double,Vector Double)  -- ^ (frequency index,power density)
pwelch s w v = let w' = max s w -- make window at least sampling rate
                   r  = S.pwelch w' v
                   sd = (fromIntegral s)/2
                   -- scale for sampling rate
                   r' = scale (recip sd) r
                   f  = linspace ((w `div` 2) + 1) (0,sd)
               in (f,r')

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

-- | a broadband FIR
broadband_fir :: (S.Filterable a, Double ~ DoubleOf a, Convert (Complex a)) =>
                Int           -- ^ sampling rate
              -> (Int,Int)     -- ^ (lower,upper) frequency cutoff
              -> Vector a -- ^ filter coefficients
broadband_fir s (l,h) = let o = 501
                            ny = (fromIntegral s) / 2.0
                            fl = (fromIntegral l) / ny
                            fh = (fromIntegral h) / ny
                            f = [0, fl*0.95, fl, fh, fh*1.05, 1]
                            m = [0,0,1,1,0,0]
                            be = zip f m
                        in standard_fir o be

-- | a broadband filter
broadband_filter :: (S.Filterable a, Double ~ DoubleOf a)
                   => Int        -- ^ sampling rate
                 -> (Int,Int)    -- ^ (lower,upper) frequency cutoff
                 -> Vector a            -- ^ input signal
                 -> Vector a            -- ^ output signal
broadband_filter s f v = let b = S.fromDouble $ broadband_fir s f
                         in filter b (scalar 1.0) s v

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

-- | standard FIR filter
-- |   FIR filter with grid a power of 2 greater than the order, ramp = grid/16, hamming window
standard_fir :: (S.Filterable a, Double ~ DoubleOf a, Convert (Complex a)) =>
               Int -> [(a,a)] -> Vector a
standard_fir o be = let grid  = calc_grid o
                        trans_ = grid `div` 16
                    in fir o be grid trans_ $ S.hamming_ (o+1)

calc_grid :: Int -> Int
calc_grid o = let next_power = ceiling (((log $ fromIntegral o) :: Double) / (log 2.0)) :: Int
              in floor $ 2.0 ** ((fromIntegral next_power) :: Double)

-- | produce an FIR filter
fir :: (S.Filterable a
      , Convert (Complex a), Double ~ DoubleOf a) =>
      Int               -- ^ order (one less than the length of the filter)
    -> [(a,a)] -- ^ band edge frequency, nondecreasing, [0, f1, ..., f(n-1), 1]
                        -- ^ band edge magnitude
    -> Int               -- ^ grid spacing
    -> Int               -- ^ transition width
    -> Vector a     -- ^ smoothing window (size is order + 1)
    -> Vector a     -- ^ the filter coefficients
fir o be gn tn w = let mid = o `div` 2
                       (f,m) = unzip be
                       f' = diff (((fromIntegral gn))/((fromIntegral tn))/2.0) f
                       m' = interpolate f m f'
                       grid = interpolate f' m' $ map (\x -> (fromIntegral x)/(fromIntegral gn)) [0..(gn-1)]
                       grid' = map (\x -> x :+ 0) grid
                       b = S.fromDouble $ fst $ fromComplex $ F.ifft $ double $ fromList $ grid' ++ (reverse (drop 1 grid'))
                       b' = vjoin [subVector ((size b)-mid-1) (mid+1) b, subVector 1 (mid+1) b]
                   in b' * w

floor_zero x
    | x < 0.0   = 0.0
    | otherwise = x

ceil_one x
    | x > 1.0   = 1.0
    | otherwise = x

diff :: S.Filterable a => a -> [a] -> [a]
diff _ []  = []
diff _ [x] = [x]
diff inc (x1:x2:xs)
     | x1 == x2     = (floor_zero $ x1-inc):x1:(ceil_one $ x1+inc):(diff inc (L.filter (/= x2) xs))
     | otherwise    = x1:(diff inc (x2:xs))

interpolate :: S.Filterable a => [a] -> [a] -> [a] -> [a]
interpolate _ _ []      = []
interpolate x y (xp:xs) = if xp == 1.0
                             then ((interpolate'' ((length x)-1) x y xp):(interpolate x y xs))
                             else ((interpolate' x y xp):(interpolate x y xs))

interpolate' :: S.Filterable a => [a] -> [a] -> a -> a
interpolate' x y xp = let Just j = L.findIndex (> xp) x
                      in (interpolate'' j x y xp)

interpolate'' :: S.Filterable a => Int -> [a] -> [a] -> a -> a
interpolate'' j x y xp = let x0 = x !! (j-1)
                             y0 = y !! (j-1)
                             x1 = x !! j
                             y1 = y !! j
                         in y0 + (xp - x0) * ((y1 - y0)/(x1-x0))

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

-- | determine the frequency response of a filter, given a vector of frequencies
freqzF :: (S.Filterable a, Double ~ DoubleOf a, S.Filterable (DoubleOf a)) =>
         Vector a     -- ^ zero coefficients
       -> Vector a       -- ^ pole coefficients
       -> Int     -- ^ sampling rate
       -> Vector a       -- ^ frequencies
       -> Vector a       -- ^ frequency response
freqzF b a s f = S.freqz b a ((2*pi/(fromIntegral s)) * f)

-- | determine the frequency response of a filter, given a number of points and sampling rate
freqzN :: (S.Filterable a, Double ~ DoubleOf a) =>
         Vector a     -- ^ zero coefficients
       -> Vector a       -- ^ pole coefficients
       -> Int     -- ^ sampling rate
       -> Int     -- ^ number of points
       -> (Vector a,Vector a)   -- ^ (frequencies,response)
freqzN b a s n = let w' = linspace n (0,((fromIntegral n)-1)/(fromIntegral (2*n)))
                     r = S.freqz b a ((2*pi)*w')
                     in ((fromIntegral s)*w',r)

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

-- | an analytic signal is the original signal with Hilbert-transformed signal as imaginary component
analytic_signal :: Vector Double -> Vector (Complex Double)
analytic_signal = S.hilbert

-- | the power (amplitude^2 = v * (conj c)) of an analytic signal
analytic_power :: S.Filterable a => Vector (Complex Double) -> Vector a
analytic_power = S.complex_power_

-- | the phase of an analytic signal
analytic_phase :: (S.Filterable a) =>
                 Vector (Complex a) -> Vector a
analytic_phase = (uncurry arctan2) . fromComplex

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

-- | remove a linear trend from data
detrend :: Int             -- ^ window size
        -> Vector Double   -- ^ data to be detrended
        -> Vector Double   -- ^ detrended data
detrend w v = let windows = size v `div` w
                  re = size v - (windows * w)
                  re' = if re == 0 then [] else [re]
                  ws = takesV ((replicate windows w) ++ re') v
                  ds = map detrend' ws
                  windows' = (size v - (w `div` 2)) `div` w
                  ws' = takesV (((w `div` 2):(replicate windows' w)) ++ [size v - (w `div` 2) - (windows' * w)]) v
                  ds' = map detrend' ws'
              in (vjoin ds + vjoin ds') / 2
    where detrend' x = let ln = size x
                           t = linspace ln (1.0,fromIntegral ln)
                           (c0,c1,_,_,_,_) = linear t x
                       in x - (scale c1 t + scalar c0)

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

-- | resize the vector to length n by resampling
resize :: S.Filterable a => Int -> Vector a -> Vector a
resize n v = S.downsample_ (size v `div` n) v

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

-- | cross covariance of two signals
--     the cross correlation is computed by dividing the result
--     by the product of the two standard deviations
cross_covariance :: S.Filterable a =>
                   Int           -- ^ maximum delay
                 -> Vector a -- ^ time series
                 -> Vector a -- ^ time series
                 -> (a,a,Vector a) -- ^ (sd_x,sd_y,cov_xy)
cross_covariance = S.cross_covariance_

-- | cross correlation of two signals
cross_correlation :: S.Filterable a =>
                   Int           -- ^ maximum delay
                 -> Vector a -- ^ time series
                 -> Vector a -- ^ time series
                 -> Vector a -- ^ result
cross_correlation l x y = let (sx,sy,r) = S.cross_covariance_ l x y
                          in GV.map(/ (sx*sy)) r

-- | compute the cross spectrum
cross_spectrum :: (S.Filterable a, Double ~ DoubleOf a) =>
                 Int                      -- ^ maximum delay
               -> Vector a                 -- ^ time series
               -> Vector a                 -- ^ time series
               -> Vector (Complex Double)  -- ^ result
cross_spectrum l x y = (\(_,_,c) -> F.fft (complex $ double c)) (cross_covariance l x y)

-- | auto covariance of two signals
--     the auto correlation is computed by dividing the result
--     by the variance
auto_covariance :: S.Filterable a =>
                   Int           -- ^ maximum delay
                 -> Vector a -- ^ time series
                 -> (a,Vector a) -- ^ (var,cov_xx)
auto_covariance l v = let (sd,_,r) = cross_covariance l v v
                      in (sd*sd,r)

-- | auto correlation of two signals
auto_correlation :: S.Filterable a =>
                   Int           -- ^ maximum delay
                 -> Vector a -- ^ time series
                 -> Vector a -- ^ result
auto_correlation l v = let (var,r) = auto_covariance l v
                          in GV.map(/ var) r

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

-- | coefficients of a Hamming window
hamming :: S.Filterable a =>
          Int           -- ^ length
        -> Vector a -- ^ the Hamming coeffficents
hamming = S.hamming_

-- | resample, take one sample every n samples in the original
downsample :: S.Filterable a => Int -> Vector a -> Vector a
downsample = S.downsample_

-- | the difference between consecutive elements of a vector
deriv :: S.Filterable a => Vector a -> Vector a
deriv = S.deriv_

-- | cumulative sum of a series
cumulative_sum :: S.Filterable a =>
                 Vector a
               -> Vector a
cumulative_sum = S.cumulative_sum_

-- | unwrap the phase of signal (input expected to be within (-pi,pi))
unwrap :: S.Filterable a => Vector a -> Vector a
unwrap = S.unwrap_

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