{-# LANGUAGE FlexibleContexts #-}
module FokkerPlanck.DomainChange where

import           Data.Array.Repa       as R
import           Data.Complex
import           Data.List             as L
import           FokkerPlanck.Types
import           Numeric.LinearAlgebra as NL
import           Types
import           Utils.Array

{-# INLINE r2z1t0Tor2s1t0 #-}
r2z1t0Tor2s1t0 :: Int -> [Double] -> R2Z1T0Array -> R2S1T0Array
r2z1t0Tor2s1t0 numOrientations freqs arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      (Z :. numFreqs :. numT0Freqs :. xLen :. yLen) = extent arr
      mat1 =
        (numOrientations >< numFreqs) . R.toList $
        R.traverse
          (fromListUnboxed (Z :. numFreqs) freqs)
          (const (Z :. numOrientations :. numFreqs))
          (\f (Z :. i :. j) ->
             exp (0 :+ (-1) * (deltaTheta * fromIntegral i) * f (Z :. j)))
      mat2 = (numFreqs >< (xLen * yLen * numT0Freqs)) . R.toList $ arr
   in fromListUnboxed (Z :. numOrientations :. numT0Freqs :. xLen :. yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2z1Tor2s1 #-}
r2z1Tor2s1 ::
     (Source s (Complex Double))
  => Int
  -> [Double]
  -> R.Array s DIM3 (Complex Double)
  -> R.Array U DIM3 (Complex Double)
r2z1Tor2s1 numOrientations freqs arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      (Z :. numFreqs :. xLen :. yLen) = extent arr
      mat1 =
        (numOrientations >< numFreqs) . R.toList $
        R.traverse
          (fromListUnboxed (Z :. numFreqs) freqs)
          (const (Z :. numOrientations :. numFreqs))
          (\f (Z :. i :. j) ->
             exp (0 :+ (-1) * (deltaTheta * fromIntegral i) * f (Z :. j)))
      mat2 = (numFreqs >< (xLen * yLen)) . R.toList $ arr
   in fromListUnboxed (Z :. numOrientations :. xLen :. yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2s1tor2z1 #-}
r2s1tor2z1 ::
     (Source s (Complex Double))
  => [Double]
  -> R.Array s DIM3 (Complex Double)
  -> R.Array U DIM3 (Complex Double)
r2s1tor2z1 freqs arr =
  let (Z :. numOrientations :. xLen :. yLen) = extent arr
      deltaTheta = 2 * pi / (fromIntegral numOrientations)
      numFreqs = L.length freqs
      mat1 =
        (numFreqs >< numOrientations) . R.toList $
        R.traverse
          (fromListUnboxed (Z :. numFreqs) freqs)
          (const (Z :. numFreqs :. numOrientations))
          (\f (Z :. i :. j) ->
             exp (0 :+ (deltaTheta * fromIntegral j) * f (Z :. i)))
      mat2 = (numOrientations >< (xLen * yLen)) . R.toList $ arr
   in fromListUnboxed (Z :. numFreqs :. xLen :. yLen) . NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2z2t0s0Tor2s1rpt0s0 #-}
r2z2t0s0Tor2s1rpt0s0 ::
     Int -> [Double] -> Int -> [Double] -> R2Z2T0S0Array -> R2S1RPT0S0Array
r2z2t0s0Tor2s1rpt0s0 numOrientations thetaFreqs numScales scaleFreqs arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :. xLen :. yLen) =
        extent arr
      mat1 =
        ((numOrientations * numScales) >< (numThetaFreq * numScaleFreq)) .
        R.toList $
        R.traverse3
          (fromFunction
             (Z :. numOrientations :. numScales :. numThetaFreq :. numScaleFreq)
             (const 0))
          (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
          (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
          (\a _ _ -> a)
          (\f1 f2 f3 (Z :. i :. j :. k :. l) ->
             exp
               (0 :+
                (-1) *
                ((deltaTheta * fromIntegral i) * f2 (Z :. k) +
                 (2 * pi * fromIntegral j / fromIntegral numScales) *
                 f3 (Z :. l))))
      mat2 =
        ((numThetaFreq * numScaleFreq) ><
         (xLen * yLen * numTheta0Freq * numScale0Freq)) .
        R.toList $
        arr
   in fromListUnboxed
        (Z :. numOrientations :. numScales :. numTheta0Freq :. numScale0Freq :.
         xLen :.
         yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2z2Tor2s1rp #-}
r2z2Tor2s1rp ::
     (Source s (Complex Double))
  => Int
  -> [Double]
  -> Int
  -> [Double]
  -> Double
  -> R.Array s DIM4 (Complex Double)
  -> R.Array U DIM4 (Complex Double)
r2z2Tor2s1rp numOrientations thetaFreqs numScales scaleFreqs maxScale arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      deltaScale = log maxScale / (fromIntegral numScales)
      (Z :. numThetaFreq :. numScaleFreq :. xLen :. yLen) = extent arr
      mat1 =
        ((numOrientations * numScales) >< (numThetaFreq * numScaleFreq)) .
        R.toList $
        R.traverse3
          (fromFunction
             (Z :. numOrientations :. numScales :. numThetaFreq :. numScaleFreq)
             (const 0))
          (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
          (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
          (\a _ _ -> a)
          (\f1 f2 f3 (Z :. i :. j :. k :. l) ->
             exp
               (0 :+
                (1) *
                ((deltaTheta * fromIntegral i) * f2 (Z :. k) +
                 (2 * pi * fromIntegral j * deltaScale / (log maxScale)) *
                 f3 (Z :. l))))
      mat2 = ((numThetaFreq * numScaleFreq) >< (xLen * yLen)) . R.toList $ arr
  in fromListUnboxed (Z :. numOrientations :. numScales :. xLen :. yLen) .
     NL.toList . flatten $
     mat1 NL.<> mat2

{-# INLINE r2z2Tor2s1rpP #-}
r2z2Tor2s1rpP ::
     (Source s (Complex Double))
  => Int
  -> [Double]
  -> Int
  -> [Double]
  -> Double
  -> R.Array s DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
r2z2Tor2s1rpP numOrientations thetaFreqs numScales scaleFreqs maxR arr = do
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      deltaScale = log maxR / (fromIntegral numOrientations)
      (Z :. numThetaFreq :. numScaleFreq :. xLen :. yLen) = extent arr
      mat2 = ((numThetaFreq * numScaleFreq) >< (xLen * yLen)) . R.toList $ arr
  mat1 <-
    fmap
      (((numOrientations * numScales) >< (numThetaFreq * numScaleFreq)) .
       R.toList) .
    computeUnboxedP $
    R.traverse3
      (fromFunction
         (Z :. numOrientations :. numScales :. numThetaFreq :. numScaleFreq)
         (const 0))
      (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
      (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
      (\a _ _ -> a)
      (\f1 f2 f3 (Z :. i :. j :. k :. l) ->
         exp
           (0 :+
            (1) *
            ((deltaTheta * fromIntegral i) * f2 (Z :. k) +
             (2 * pi * fromIntegral j * deltaScale / (log maxR)) * f3 (Z :. l))))
  return .
    fromListUnboxed (Z :. numOrientations :. numScales :. xLen :. yLen) .
    NL.toList . flatten $
    mat1 NL.<> mat2

{-# INLINE r2z2t0s0Tor2s1rps1rp #-}
r2z2t0s0Tor2s1rps1rp ::
     Int -> [Double] -> [Double] -> Int -> [Double] -> [Double] -> Double -> R2Z2T0S0Array -> R.Array U DIM6 (Complex Double)
r2z2t0s0Tor2s1rps1rp numOrientations thetaFreqs theta0Freqs numScales scaleFreqs scale0Freqs maxScale arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      deltaScale =
        if maxScale == 0
          then 0
          else (2 * pi / (log maxScale)) / (fromIntegral numScales)
      (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :. xLen :. yLen) =
        extent arr
      mat1 =
        ((numOrientations * numScales * numOrientations * numScales) ><
         (numThetaFreq * numScaleFreq * numTheta0Freq * numScale0Freq)) .
        R.toList $
        R.traverse4
          (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
          (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
          (fromListUnboxed (Z :. numTheta0Freq) theta0Freqs)
          (fromListUnboxed (Z :. numScale0Freq) scale0Freqs)
          (\_ _ _ _ ->
             (Z :. numOrientations :. numScales :. numOrientations :. numScales :.
              numThetaFreq :.
              numScaleFreq :.
              numTheta0Freq :.
              numScale0Freq))
          (\ft fs ft0 fs0 (Z :. o :. s :. o0 :. s0 :. tf :. sf :. tf0 :. sf0) ->
             exp
               (0 :+
                (-1) *
                ((deltaTheta * fromIntegral o) * ft (Z :. tf) +
                 (deltaTheta * fromIntegral o0) * ft0 (Z :. tf0) +
                 (deltaScale * fromIntegral s) * fs (Z :. sf) +
                 (deltaScale * fromIntegral s0) * fs0 (Z :. sf0))))
      mat2 =
        ((numThetaFreq * numScaleFreq * numTheta0Freq * numScale0Freq) ><
         (xLen * yLen)) .
        R.toList $
        arr
   in fromListUnboxed
        (Z :. numOrientations :. numScales :. numOrientations :. numScales :.
         xLen :.
         yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2s1rps1rpTor2z2t0s0 #-}
r2s1rps1rpTor2z2t0s0 ::
     Int
  -> [Double]
  -> [Double]
  -> Int
  -> [Double]
  -> [Double]
  -> Double
  -> R.Array U DIM6 (Complex Double)
  -> R2Z2T0S0Array
r2s1rps1rpTor2z2t0s0 numOrientations thetaFreqs theta0Freqs numScales scaleFreqs scale0Freqs maxScale arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      deltaScale =
        if maxScale == 0
          then 0
          else (2 * pi / (log maxScale)) / (fromIntegral numScales)
      numThetaFreq = L.length thetaFreqs
      numScaleFreq = L.length scaleFreqs
      numTheta0Freq = L.length theta0Freqs
      numScale0Freq = L.length scale0Freqs
      (Z :. _ :. _ :. _ :. _ :. xLen :. yLen) =
        extent arr
      mat1 =
        ((numThetaFreq * numScaleFreq * numTheta0Freq * numScale0Freq) ><
         (numOrientations * numScales * numOrientations * numScales)) .
        R.toList $
        R.traverse4
          (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
          (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
          (fromListUnboxed (Z :. numTheta0Freq) theta0Freqs)
          (fromListUnboxed (Z :. numScale0Freq) scale0Freqs)
          (\_ _ _ _ ->
             (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :.
              numScale0Freq :.
              numOrientations :.
              numScales :.
              numOrientations :.
              numScales))
          (\ft fs ft0 fs0 (Z :. tf :. sf :. tf0 :. sf0 :. o :. s :. o0 :. s0) ->
             exp
               (0 :+ (1) *
                ((deltaTheta * fromIntegral o) * ft (Z :. tf) +
                 (deltaTheta * fromIntegral o0) * ft0 (Z :. tf0) +
                 (deltaScale * fromIntegral s) * fs (Z :. sf) +
                 (deltaScale * fromIntegral s0) * fs0 (Z :. sf0))))
      mat2 =
        ((numOrientations * numScales * numOrientations * numScales) ><
         (xLen * yLen)) .
        R.toList $
        arr
   in fromListUnboxed
        (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :.
         xLen :.
         yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

-- x y t s
{-# INLINE stsTor2s1rp #-}
stsTor2s1rp ::
     Int
  -> Double
  -> [Double]
  -> Int
  -> Double
  -> [Double]
  -> Int
  -> [Double]
  -> Int
  -> [Double]
  -> Double
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
stsTor2s1rp numX xLen xFreqs numY yLen yFreqs numOrientations thetaFreqs numScales scaleFreqs maxScale arr = do
  let deltaX = xLen / fromIntegral numX
      deltaY = yLen / fromIntegral numY
      deltaTheta = 2 * pi / fromIntegral numOrientations
      deltaScale = (log maxScale) / fromIntegral numScales
      centerX = div numX 2
      centerY = div numY 2
      vec = NL.fromList . R.toList $ arr
  transformMat <-
    fmap
      (((numX * numY * numOrientations * numScales) ><
        (L.product . listOfShape . extent $ arr)) .
       R.toList) .
    computeUnboxedP .
    R.traverse4
      (fromListUnboxed (Z :. L.length xFreqs) xFreqs)
      (fromListUnboxed (Z :. L.length yFreqs) yFreqs)
      (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
      (fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs)
      (\(Z :. numXFreq) (Z :. numYFreq) (Z :. numThetaFreq) (Z :. numScaleFreq) ->
         (Z :. numX :. numY :. numOrientations :. numScales :. numXFreq :.
          numYFreq :.
          numThetaFreq :.
          numScaleFreq)) $ \fx fy fo fs (Z :. x :. y :. o :. s :. xf :. yf :. orif :. sf) ->
      exp
        (0 :+
         (fo (Z :. orif) * fromIntegral o * deltaTheta +
          2 * pi *
          (fs (Z :. sf) * fromIntegral s * deltaScale / (log maxScale) +
           fx (Z :. xf) * fromIntegral ((x :: Int) - centerX) * deltaX / xLen +
           fy (Z :. yf) * fromIntegral ((y :: Int) - centerY) * deltaY / yLen)))
  return .
    fromListUnboxed (Z :. numX :. numY :. numOrientations :. numScales) .
    NL.toList $
    transformMat NL.#> vec

{-# INLINE stsTor2 #-}
stsTor2 ::
     Int
  -> Double
  -> [Double]
  -> Int
  -> Double
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM2 (Complex Double))
stsTor2 numX xLen xFreqs numY yLen yFreqs arr = do
  let deltaX = xLen / fromIntegral numX
      deltaY = yLen / fromIntegral numY
      centerX = div numX 2
      centerY = div numY 2
  vec <-
    fmap (NL.fromList . R.toList) .
    sumP . sumS . R.map (\x -> (magnitude x) ^ 2 :+ 0) $
    arr
  transformMat <-
    fmap
      (((numX * numY) >< (L.product . L.drop 2 . listOfShape . extent $ arr)) .
       R.toList) .
    computeUnboxedP .
    R.traverse2
      (fromListUnboxed (Z :. L.length xFreqs) xFreqs)
      (fromListUnboxed (Z :. L.length yFreqs) yFreqs)
      (\(Z :. numXFreq) (Z :. numYFreq) ->
         (Z :. numX :. numY :. numXFreq :. numYFreq)) $ \fx fy (Z :. x :. y :. xf :. yf) ->
      exp
        (0 :+
         (2 * pi *
          (fx (Z :. xf) * fromIntegral (x - centerX) * deltaX / xLen +
           fy (Z :. yf) * fromIntegral (y - centerY) * deltaY / yLen)))
  return . fromListUnboxed (Z :. numX :. numY) . NL.toList $
    transformMat NL.#> vec

{-# INLINE stsTor2' #-}
stsTor2' ::
     Int
  -> Double
  -> [Double]
  -> Int
  -> Double
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM2 Double)
stsTor2' numX xLen xFreqs numY yLen yFreqs arr = do
  let deltaX = xLen / fromIntegral numX
      deltaY = yLen / fromIntegral numY
      centerX = div numX 2
      centerY = div numY 2
      (Z :. numXF :. numYF :. numTF :. numSF) = extent arr
      mat = ((numXF * numYF) >< (numTF * numSF)) (R.toList arr)
  transformMat <-
    fmap
      (((numX * numY) >< (L.product . L.drop 2 . listOfShape . extent $ arr)) .
       R.toList) .
    computeUnboxedP .
    R.traverse2
      (fromListUnboxed (Z :. L.length xFreqs) xFreqs)
      (fromListUnboxed (Z :. L.length yFreqs) yFreqs)
      (\(Z :. numXFreq) (Z :. numYFreq) ->
         (Z :. numX :. numY :. numXFreq :. numYFreq)) $ \fx fy (Z :. x :. y :. xf :. yf) ->
      exp
        (0 :+
         (2 * pi *
          (fx (Z :. xf) * fromIntegral (x - centerX) * deltaX / xLen +
           fy (Z :. yf) * fromIntegral (y - centerY) * deltaY / yLen)))
  return .
    sumS .
    sumS .
    R.map (\x -> (magnitude x) ^ 2) .
    fromListUnboxed (Z :. numX :. numY :. numTF :. numSF) . NL.toList . flatten $
    transformMat NL.<> mat


module ForSyDe.Shallow.Example.Synchronous.Equalizer.Test (
  -- * Test data
  -- ** Test signals
  zeros, ones, increm, modCos,
  bassUp, bassDn, trebleUp, trebleDn,
  overrides, distorsionFlags,
  -- ** Test filter coefficients
  lpCoeff, bpCoeff, hpCoeff,
  lpCoeffD, bpCoeffD, hpCoeffD,
  -- ** Test vectors
  testV1, testV2, testV3,
  -- * Test functions
  -- ** Signal processing blocks
  testDFT, testFFT, testDftAgainstFft, testFIR,
  audioFilterI, audioFilterD,
  -- ** Equalizer modules
  testAudioFilterD,
  testButtonControl,
  testDistortionControl,
  testAnalyzer,
  testEqualizer,
  -- * Test scripts
  testFilterScript,
  testAnalyzerScript,
  testFilterAndAnalyzerScript
  )where

import Data.Complex
import System.IO
import System.FilePath
import ForSyDe.Shallow
import ForSyDe.Shallow.Utility.FIR
import ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioAnalyzer
import ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter
import ForSyDe.Shallow.Example.Synchronous.Equalizer.ButtonControl
import ForSyDe.Shallow.Example.Synchronous.Equalizer.DistortionControl
import ForSyDe.Shallow.Example.Synchronous.Equalizer.Equalizer
import ForSyDe.Shallow.Example.Synchronous.Equalizer.EqualizerTypes

---------------
-- TEST DATA --
---------------

-- | test signal for 'testButtonControl'
bassUp = signal [
  Prst Active, Prst Active, Abst, Abst,
  Prst Active, Prst Active, Abst, Abst,
  Prst Active, Prst Active, Abst, Abst,
  Prst Active, Prst Active, Abst, Abst ]

-- | test signal for 'testButtonControl'
bassDn = signal [
  Abst, Abst, Prst Active, Abst,
  Abst, Abst, Prst Active, Abst,
  Abst, Abst, Prst Active, Abst,
  Abst, Abst, Prst Active, Abst ]

-- | test signal for 'testButtonControl'
trebleUp = signal [
  Abst, Abst, Abst, Abst,
  Abst, Abst, Abst, Prst Active,
  Abst, Abst, Abst, Abst,
  Abst, Abst, Abst, Prst Active ]

-- | test signal for 'testButtonControl'
trebleDn = signal [
  Abst, Abst, Abst, Prst Active,
  Abst, Abst, Abst, Prst Active,
  Abst, Abst, Abst,  Prst Active,
  Abst, Abst, Abst, Prst Active ]

-- | test signal for 'testButtonControl'
overrides = signal [
  Abst, Abst, Abst, Abst,
  Prst Lock, Abst, Abst, Abst,
  Abst, Prst CutBass, Abst, Abst,
  Prst Release,  Abst, Abst, Abst ]

-- | example signal for testing
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.DistorsionControl.distortionControl'
distorsionFlags = signal [
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass ]

-- | Low-pass filter coefficients
lpCoeff = vector [
  0.01392741661548,
  0.01396895728902,
  0.01399870011280,
  0.01401657422649,
  0.01402253700635,
  0.01401657422649,
  0.01399870011280,
  0.01396895728902,
  0.01392741661548 ]

-- | Low-pass filter coefficients
lpCoeffD = vector [
   0.03898579822345,
   0.09739504968381,
   0.15360490491115,
   0.19416166962179,
   0.20893350067585,
   0.19416166962179,
   0.15360490491115,
   0.09739504968381,
   0.03898579822345
   ]

-- | Band-pass filter coefficients
bpCoeff = vector [
  0.06318761339784,
  0.08131651217682,
  0.09562326700432,
  0.10478344432968,
  0.10793629404886,
  0.10478344432968,
  0.09562326700432,
  0.08131651217682,
  0.06318761339784 ]

-- | Band-pass filter coefficients
bpCoeffD = vector [
  -0.07845593083988,
   0.00000000000000,
  -0.30707118658796,
   0.00000000000000,
   0.58268794919522,
   0.00000000000000,
  -0.30707118658796,
   0.00000000000000,
  -0.07845593083988
   ]

-- | High-pass filter coefficients
hpCoeff = vector [
  -0.07883878579454,
  -0.09820015927379,
  -0.11354603917221,
  -0.12339860164118,
   0.87320570334018,
  -0.12339860164118,
  -0.11354603917221,
  -0.09820015927379,
  -0.07883878579454 ]

-- | High-pass filter coefficients
hpCoeffD = vector [
   0.03898579822345,
  -0.09739504968381,
   0.15360490491115,
  -0.19416166962179,
   0.20893350067585,
  -0.19416166962179,
   0.15360490491115,
  -0.09739504968381,
   0.03898579822345
  ]

-- | Infinite signal of zeroes
zeros = infiniteS id 0.0

-- | Finite signal of ones, having its length given as input
ones pts = takeS pts $ infiniteS id 1.0

-- | Incremental signal of 400 samples
increm = takeS 400 (infiniteS (+1/200) 0.0)

-- | Additive sinusoidal signal of 400 samples, with 3 sine components.
modCos = mapSY modCosF increm
  where
    modCosF x = cos10 x + cos50 x + cos90 x
    cos10 x = cos (2*pi*10*x)
    cos50 x = cos (2*pi*50*x)
    cos90 x = cos (2*pi*90*x)

-- | example input vector to test the ForSyDe implementation of FFT and DFT
testV1 = mapV toComplex (vector [1, 2])

-- | example input vector to test the ForSyDe implementation of FFT and DFT
testV2 = mapV toComplex (vector [1, 2,3,4])

-- | example input vector to test the ForSyDe implementation of FFT and DFT
testV3 = mapV toComplex (vector [1, 2, 3, 4, 5, 6, 7, 8])

--------------------
-- TEST FUNCTIONS --
--------------------

-- | tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.ButtonControl.buttonControl'
-- function.
testButtonControl = buttonControl overrides bassUp bassDn trebleUp trebleDn
--testButtonInterface = buttonInterface  bassUp bassDn trebleUp trebleDn

toComplex a = a:+0

-- | tests ForSyDe's 'ForSyDe.Shallow.dft' against 'testV3'
testDFT = dft 8 testV3

-- | tests ForSyDe's 'ForSyDe.Shallow.fft' against 'testV3'
testFFT = fft 8 testV3

-- | tests both 'ForSyDe.Shallow.dft' and 'ForSyDe.Shallow.fft'
testDftAgainstFft = zipWithV (-) testDFT testFFT

-- | tests ForSyDe's 'ForSyDe.Shallow.firSY' with 4 coefficients
-- against an impulse signal.
testFIR = firSY coeff s
  where
    coeff = vector [0.1, -0.2, 0.5, 0.2]
    s = signal [1.0, 0, 0, 0, 0]

pts = 4

-- | tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.DistorsionControl.distortionControl'
-- function against 'distorsionFlags'.
testDistortionControl = distortionControl distorsionFlags

-- | Tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioAnalyzer'
-- against a 'audioFilter'ed dummy signal.
testAnalyzer = audioAnalyzer 2 (audioFilter lpCoeff bpCoeff hpCoeff bass treble audioIn)
  where audioIn = ones (pts * 4)
        bass    = zeros
        treble  = zeros

-- | Instantiates
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioFilter'.
audioFilterI = audioFilter lpCoeff bpCoeff hpCoeff zeros zeros

-- | Instantiates
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioFilter'.
audioFilterD = audioFilter lpCoeffD bpCoeffD hpCoeffD zeros zeros

-- | Tests
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioFilter'
-- against a test signal.
testAudioFilterD = audioFilter lpCoeffD bpCoeffD hpCoeffD zeros zeros s
  where s = signal [0.1,0.05,0.3,0.2,-0.5,0.2,0.1,0.3,0.1,-0.1,0.0,0.2]

-- | Tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.Equalizer.equalizer'
-- function against a dummy signal consisting of all ones.
testEqualizer = equalizer lpCoeff bpCoeff hpCoeff 2 bassUp bassDn trebleUp trebleDn audioIn
  where audioIn = ones (pts * 4)

--testButtonInterface = buttonInterface  bassUp bassDn trebleUp trebleDn

-- fs = 200

--lpCoeffF8 = mapV real2Fixed8 lpCoeffD
--bpCoeffF8 = mapV real2Fixed8 bpCoeffD
--hpCoeffF8 = mapV real2Fixed8 hpCoeffD
--sF8    = mapSY real2Fixed8 s

--lpCoeffF16 = mapV real2Fixed16 lpCoeffD
--bpCoeffF16 = mapV real2Fixed16 bpCoeffD
--hpCoeffF16 = mapV real2Fixed16 hpCoeffD
--sF16    = mapSY real2Fixed16 s

--outF16    = audioFilter lpCoeffF16 bpCoeffF16 hpCoeffF16 (mapSY real2Fixed16 zeros) (mapSY real2Fixed16 zeros) sF16
--outF8    = audioFilter lpCoeffF8 bpCoeffF8 hpCoeffF8 (mapSY real2Fixed8 zeros)(mapSY real2Fixed8 zeros) sF8

--writeAudioOut = writeFile "Test/AudioOut.for") . writeS

-------------
-- SCRIPTS --
-------------

-- | Tests the 'audioFilterD' function, against the contents of a
-- @.dat@ file. Dumps in another @.dat@ file. An example input file is
-- provided in the @files/test-equalizer@ folder in the root path of
-- the project.
testFilterScript inFile = do
  contents <- readFile inFile
  writeFile audioOut (writeS (audioFilterD (readS contents)))
  putStrLn $ "\nDone. Data dumped in " ++ audioOut
  where
    audioOut = takeDirectory inFile ++ "/AudioOutFSD.ext"

-- | Tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioAnalyzer'
-- function, against the contents of a @.dat@ file. Dumps in another
-- @.dat@ file. An example input file is provided in the
-- @files/test-equalizer@ folder in the root path of the project.
testAnalyzerScript inFile = do
  fftInfile  <- openFile inFile ReadMode
  fftOutfile <- openFile outFile WriteMode
  contents   <- hGetContents fftInfile
--    hPutStr fftOutfile (writeS (audioAnalyzer 6 ((readS contents2) :: Signal Double)))
  putStr (show (audioAnalyzer pts ((readS contents) :: Signal Double)))
  hPutStr fftOutfile (writeS (audioAnalyzer pts ((readS contents) :: Signal Double)))
--    hClose fftOutfile
  putStr $ "\nDone. Data dumped in " ++ outFile ++ "\n"
  where
    k = 8
    n = 2 ^ k
    sig = takeS (2*(2^k)) input
    input = 0.1 :- 0.2 :- input
    outFile = takeDirectory inFile ++ "/fftOut.dat"

-- | Chains and tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioFilter' and
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioAnalyzer'
-- functions, against the contents of a @.dat@ file. Dumps in another
-- @.dat@ file. An example input file is provided in the
-- @files/test-equalizer@ folder in the root path of the project.
testFilterAndAnalyzerScript inFile = do
  -- Test AudioFilter
  putStr "\n-->Test AudioFilter \n"
  filterInfile     <- openFile inFile ReadMode
  filterContents   <- hGetContents filterInfile
  putStr (show (audioFilterD (readS filterContents)))
  writeFile audioOut (writeS (audioFilterD (readS filterContents)))
  putStrLn $ "\nDone. Data dumped in " ++ audioOut
-- Test AudioAnalyzer
  putStr "\n--> Test AudioAnalyzer \n"
  analyzerInfile   <- openFile audioOut ReadMode
  analyzerOutfile  <- openFile analyzerOut WriteMode
  analyzerContents <- hGetContents analyzerInfile
  putStr (show (audioAnalyzer pts ((readS analyzerContents) :: Signal Double)))
  hPutStr analyzerOutfile (writeS (audioAnalyzer pts ((readS analyzerContents) :: Signal Double)))
  putStrLn $ "\nDone. Data dumped in " ++ analyzerOut
  --fftInfile <- openFile "audioOut.txt" ReadMode
  --fftOutfile <- openFile "fftOut.txt" WriteMode
  --contents <- hGetContents fftInfile
  --putStr (writeS (((readS contents) :: Signal Double)))
  --hPutStr fftOutfile (writeS (((readS contents) :: Signal Double)))
  where
    audioOut    = takeDirectory inFile ++ "/audioOut.dat"
    analyzerOut = takeDirectory inFile ++ "/analyzerOut.dat"


{-# LANGUAGE FlexibleContexts, UndecidableInstances #-}

module School.Unit.UnitActivation
( UnitActivation(..)
, isApplyFail
) where

import Numeric.LinearAlgebra (Container, Matrix, Vector)

data UnitActivation a =
    BatchActivation (Matrix a)
 |  ApplyFail String deriving (Show)

instance (Container Vector a, Eq a, Num a) => Eq (UnitActivation a) where
  (BatchActivation a1) == (BatchActivation a2) =
     a1 == a2
  (ApplyFail _) == (ApplyFail _) = True
  _ == _ = False

isApplyFail :: UnitActivation a -> Bool
isApplyFail (ApplyFail _) = True
isApplyFail _ = False


{-# language DeriveGeneric, DeriveAnyClass, StrictData #-}
{-# language KindSignatures, DataKinds, ScopedTypeVariables #-}
{-|
Module      : GA
Description : Genetic Algorithm for Poly Regression
Copyright   : (c) Fabricio Olivetti de Franca, 2021
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX
|-}

module GA
  ( generational
  , createRndSolution
  , crossover
  , mutate
  , Solution(..)
  , Poly(..)
  , Population
  , Fitness
  ) where

import System.Random
import Numeric.LinearAlgebra (Vector)

import Random
import Control.Monad
import Control.Monad.State.Strict
import GHC.Generics (Generic)
import Control.DeepSeq
import GHC.TypeLits
import Data.Proxy

-- | Data type representing a solution
newtype Poly (n :: Nat) = Poly { _getPoly :: (Bool, Int) } deriving (Show, Generic, NFData)

instance KnownNat n => RandomAlele (Poly n) where
  randomProb p x = Poly <$> biRandomWith p randomBool (randomInt (1, maxK)) (_getPoly x)
    where maxK = fromIntegral $ natVal (Proxy :: Proxy n)
  randomAlele    = Poly <$> randomPair (1, maxK)
    where maxK = fromIntegral $ natVal (Proxy :: Proxy n)

data Solution a = Sol { _chromo    :: [a]                    -- ^ chromossome representation
                      , _coeffs    :: Maybe (Vector Double)  -- ^ coefficients of the regression
                      , _fitness   :: Maybe Double           -- ^ fitness of an individual
                      } deriving (Show, Generic, NFData)

instance Eq (Solution a) where
  (Sol _ _ f1) == (Sol _ _ f2) = f1 == f2
instance Ord (Solution a) where
  (Sol _ _ f1) <= (Sol _ _ f2) = f1 <= f2

type Population a   = [Solution a]
type Fitness a      = Solution a -> Solution a

nVars :: Solution a -> Int
nVars s = length (_chromo s)

-- | Creates a random solution
createRndSolution :: RandomAlele a => Int -> Rnd (Solution a)
createRndSolution nLen =
  do chromo <- replicateM nLen randomAlele
     return $ Sol chromo Nothing Nothing

-- | Tournament between two solutions
tournament :: Population a -> Rnd (Solution a)
tournament pop = do s1 <- chooseRndFrom pop
                    s2 <- chooseRndFrom pop
                    return (min s1 s2)

-- | Creates a new solution from the combination of two solutions
crossover :: Population a -> Rnd (Solution a)
crossover pop = do s1 <- tournament pop
                   s2 <- tournament pop
                   ix <- randomInt (0, nVars s1)
                   return (cx ix s1 s2)

-- | Performs crossover
cx :: Int -> Solution a -> Solution a -> Solution a
cx ix sol1 sol2 = Sol child Nothing Nothing
  where
    recombine xs ys = take ix xs ++ drop ix ys
    child           = recombine (_chromo sol1) (_chromo sol2)

-- | Mutates a single solution
mutate :: RandomAlele a => Double -> Solution a -> Rnd (Solution a)
mutate prob sol = do chromo <- traverse (randomProb prob) (_chromo sol)
                     return $ Sol chromo Nothing Nothing

-- | Select the next generation
generational :: Population a -> Population a -> Rnd (Population a)
generational _ = return


module Helpers
    ( roundToStr
    , fta
    , applyAttrs
    , degToRad
    , radToDeg
    , vectorToFloats ) where

import Numeric.LinearAlgebra.Data (Vector, R, toList)
import qualified Text.Blaze.Svg11 as S
import Text.Printf

roundToStr :: (PrintfArg a, Floating a) => Int -> a -> String
roundToStr = printf "%0.*f"

-- Convert a float to an SVG attribute, preserving up to 6 decimal places.
fta :: Float -> S.AttributeValue
fta f = S.stringValue (roundToStr 6 f)

-- Apply a list of attributes to an SVG element.
applyAttrs :: [S.Attribute] -> S.Svg ->  S.Svg
applyAttrs attributes element = foldr (flip (S.!)) element attributes

-- | Convert from degrees to radians.
radToDeg :: Float -> Float
radToDeg radians = radians * 180 / pi

-- | Convert from radians to degrees.
degToRad :: Float -> Float
degToRad degrees = degrees * pi / 180

-- | [Hmatrix vectors](https://hackage.haskell.org/package/hmatrix) use doubles, not floats
vectorToFloats :: Vector R -> [Float]
vectorToFloats = map realToFrac . toList