module SpectralDistribution where

import Control.Applicative (Applicative, liftA2, pure, (<*>), )
import Control.DeepSeq (NFData, rnf, )

import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold

import Foreign.Storable.Record as Store
import Foreign.Storable (Storable (..), )

import NumericPrelude.Numeric
import NumericPrelude.Base
import Prelude ()


data T a = Cons {centroid, spread :: a}
   deriving (Eq, Show)

instance (NFData a) => NFData (T a) where
   rnf (Cons c s) = rnf (c,s)

instance Functor T where
   fmap = Trav.fmapDefault

instance Applicative T where
   pure x = Cons x x
   Cons fc fs <*> Cons c s = Cons (fc c) (fs s)

instance Fold.Foldable T where
   foldMap = Trav.foldMapDefault

instance Trav.Traversable T where
   traverse f (Cons c s) = liftA2 Cons (f c) (f s)


store :: Storable a => Store.Dictionary (T a)
store =
   Store.run $
   liftA2 Cons
      (Store.element centroid)
      (Store.element spread)

instance Storable a => Storable (T a) where
   sizeOf = Store.sizeOf store
   alignment = Store.alignment store
   peek = Store.peek store
   poke = Store.poke store


mapSpread :: (a -> a) -> T a -> T a
mapSpread f (Cons c s) = Cons c (f s)


{-# INLINE spectralDistribution1 #-}
spectralDistribution1 :: Float -> Float -> Float -> T Float
spectralDistribution1 d0 d1 d2 =
   let r1 = d1/d0 in Cons r1 (d2/d0 - r1^2)


{-# INLINE signedSqrt #-}
signedSqrt :: Float -> Float
signedSqrt x = signum x * sqrt (abs x)

{-# INLINE spectralDistribution2 #-}
spectralDistribution2 :: Float -> Float -> Float -> T Float
spectralDistribution2 d0 d1 d2 =
   let r1 = d1/d0 in Cons (sqrt r1) (signedSqrt (d2/d0 - r1^2))