{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict              #-}
module Pinwheel.FourierSeries2D where

import           Control.Monad                  as M
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Resource
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX as A
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Complex
import           Data.Conduit                   as C
import           Data.Conduit.List              as CL
import           Data.List                      as L
import           Data.Vector.Generic            as VG
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.FourierSeries2D
import           Math.Gamma
import           Pinwheel.Base
import           Pinwheel.List
import           Utils.BLAS
import           Utils.Distribution
import           Utils.List
import           Utils.Parallel                 hiding ((.|))
import           Utils.SimpsonRule
import           Utils.Time
import DFT.Plan
import Filter.Utils
import Control.Concurrent.Async
import Numeric.GSL.Special.Bessel

pinwheelFourierSeries ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => [Int]
  -> [PTX]
  -> Int
  -> Int
  -> e
  -> e
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> Int
  -> Int
  -> IO (IA.Array (Int, Int) (vector (Complex e)))
pinwheelFourierSeries deviceIDs ptxs numR2Freqs numPoints delta periodR2 phiFreq rhoFreq thetaFreq rFreq sigma numBatchR2 numBatchPinwheelFreqs = do
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- pinwheelFreqs rhoFreq rFreq
        , angularFreq <- pinwheelFreqs phiFreq thetaFreq
        ]
      centerR2Freq = div numR2Freqs 2
  pinwheels <-
    computeUnboxedP .
    R.traverse
      (fromListUnboxed (Z :. (L.length idxs)) idxs)
      (\(Z :. freqs) -> (Z :. freqs :. numR2Freqs :. numR2Freqs)) $ \f (Z :. pFreq :. xFreq :. yFreq) ->
      fourierMellin
        sigma
        (snd $ f (Z :. pFreq))
        (fst $ f (Z :. pFreq))
        ( delta * fromIntegral (xFreq - centerR2Freq)
        , delta * fromIntegral (yFreq - centerR2Freq))
  -- pinwheels <-
  --   M.mapM
  --     (\idx ->
  --        fmap
  --          (CuMat (numR2Freqs ^ 2) (L.length idx) .
  --           CuVecHost . VU.convert . toUnboxed) .
  --        computeP .
  --        R.traverse
  --          (fromListUnboxed (Z :. (L.length idx)) idx)
  --          (\(Z :. freqs) -> (Z :. numR2Freqs :. numR2Freqs :. freqs)) $ \f (Z :. xFreq :. yFreq :. pFreq) ->
  --          fourierMellin
  --            sigma
  --            (snd $ f (Z :. pFreq))
  --            (fst $ f (Z :. pFreq))
  --            ( fromIntegral (xFreq - centerR2Freq)
  --            , fromIntegral (yFreq - centerR2Freq))) .
  --   divideListN numBatchPinwheelFreqs $
  --   idxs
  -- printCurrentTime "Compute Fourier Series... "
  -- arr <-
  --   computeFourierSeriesR2Stream
  --     deviceIDs
  --     ptxs
  --     numR2Freqs
  --     numPoints
  --     periodR2
  --     delta
  --     numBatchR2
  --     pinwheels
  -- printCurrentTime "Compute Fourier Series Done"
  let (radialLB, radialUB) = pinwheelFreqsBound rhoFreq rFreq
      (angularLB, angularUB) = pinwheelFreqsBound phiFreq thetaFreq
  return .
    listArray ((radialLB, angularLB), (radialUB, angularUB)) .
    parMap
      rdeepseq
      (\i ->
         VG.convert . toUnboxed . computeS . R.slice pinwheels $  -- arr $
         (Z :. i :. All :. All)) $
    [0 .. (L.length idxs - 1)]

-- Stream
{-# INLINE source #-}
source :: Int -> Int -> Int -> Int -> ConduitT () (Int, Int) (ResourceT IO) ()
source phiFreq rhoFreq thetaFreq rFreq =
  CL.sourceList
    [ (radialFreq, angularFreq)
    | radialFreq <- pinwheelFreqs rhoFreq rFreq
    , angularFreq <- pinwheelFreqs phiFreq thetaFreq
    ]

conduit ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     , VG.Vector vector (Complex e)
     )
  => [Int]
  -> Int
  -> Int
  -> e
  -> Int
  -> e
  -> [[CuMat (Complex e)]]
  -> ConduitT (Int, Int) [vector (Complex e)] (ResourceT IO) ()
conduit deviceIDs numR2Freqs numPoints periodR2 batchSize sigma inverseR2Harmonics = do
  liftIO $ printCurrentTime "conduit"
  idx <- CL.take batchSize
  unless
    (L.null idx)
    (do let centerR2Freq = div numR2Freqs 2
        pinwheels <-
          fmap
            (CuMat (L.length idx) (numR2Freqs ^ 2) .
             CuVecHost . VU.convert . toUnboxed) .
          liftIO .
          computeP .
          R.traverse
            (fromListUnboxed (Z :. (L.length idx)) idx)
            (\(Z :. freqs) -> (Z :. freqs :. numR2Freqs :. numR2Freqs)) $ \f (Z :. pFreq :. xFreq :. yFreq) ->
            fourierMellin
              sigma
              (snd $ f (Z :. pFreq))
              (fst $ f (Z :. pFreq))
              ( fromIntegral (xFreq - centerR2Freq)
              , fromIntegral (yFreq - centerR2Freq))
        arr <-
          liftIO $
          computeFourierSeriesR2
            deviceIDs
            numR2Freqs
            numPoints
            periodR2
            inverseR2Harmonics
            [pinwheels]
        yield .
          L.map
            (\i ->
               VG.convert . toUnboxed . computeS . R.slice arr $
               (Z :. i :. All :. All)) $
          [0 .. (L.length idx - 1)])

pinwheelFourierSeriesStream ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     , VG.Vector vector (Complex e)
     )
  => [Int]
  -> [PTX]
  -> Int
  -> Int
  -> e
  -> e
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> Int
  -> Int
  -> Int
  -> IO (IA.Array (Int, Int) (vector (Complex e)))
pinwheelFourierSeriesStream deviceIDs ptxs numR2Freqs numPoints delta periodR2 phiFreq rhoFreq thetaFreq rFreq sigma numBatchR2 batchSize numBatchPinwheelFreqs = do
  inverseR2Harmonics <-
    createInverseHarmonicMatriesGPU
      ptxs
      numBatchR2
      numPoints
      numR2Freqs
      periodR2
      delta
  let (radialLB, radialUB) = pinwheelFreqsBound rhoFreq rFreq
      (angularLB, angularUB) = pinwheelFreqsBound rhoFreq rFreq
  fmap (listArray ((radialLB, angularLB), (radialUB, radialUB)) . L.concat) .
    runConduitRes $
    source phiFreq rhoFreq thetaFreq rFreq .|
    conduit
      deviceIDs
      numR2Freqs
      numPoints
      periodR2
      batchSize
      sigma
      inverseR2Harmonics .|
    CL.consume

-- 2D spatial domain to frequency domain:
-- transform: cis (-freq * x)
-- inverse transform: cis (freq * x)

-- Analytical solution type 1
-- The pinwheel harmonics are cis (-freq * x)
{-# INLINE analyticalFourierSeriesFunc1 #-}
analyticalFourierSeriesFunc1 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e))
  => Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> e
  -> Complex e
analyticalFourierSeriesFunc1 angularFreq radialFreq sigma periodR2 periodEnv phi rho =
  let radialConst = 2 * pi / log periodEnv
   in pi * ((0 :+ 1) ^ abs angularFreq) *
      cis (fromIntegral (-angularFreq) * phi) *
      ((periodR2 / (pi * rho) :+ 0) **
       ((2 + sigma) :+ (radialConst * fromIntegral (-radialFreq)))) *
      gamma
        (((2 + fromIntegral (abs angularFreq) + sigma) :+
          (radialConst * fromIntegral (-radialFreq))) /
         2) /
      gamma
        (((fromIntegral (abs angularFreq) - sigma) :+
          (radialConst * fromIntegral radialFreq)) /
         2)

-- The pinwheel is in the frequency domain, the function computes the Fourier series
{-# INLINE analyticalFourierSeries1 #-}
analyticalFourierSeries1 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierSeries1 numPoints delta angularFreq radialFreq sigma periodR2 periodEnv =
  let center = div numPoints 2
  in fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
       let x = fromIntegral $ i - center
           y = fromIntegral $ j - center
           rho = sqrt $ x ^ 2 + y ^ 2
           phi = atan2 y x
       in if rho == 0
            then 0
            else analyticalFourierSeriesFunc1
                   angularFreq
                   radialFreq
                   sigma
                   periodR2
                   periodEnv
                   phi
                   (rho * delta)

-- The pinwheel is in the spatial domain, the function computes the Fourier coefficients
{-# INLINE analyticalFourierCoefficients1 #-}
analyticalFourierCoefficients1 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierCoefficients1 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv =
  let c = ((-1) ^ (abs angularFreq)) :+ 0
  in R.map (* c) $
     analyticalFourierSeries1 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv

-- Analytical solution type 2
-- The pinwheel harmonics are cis (freq * x)
{-# INLINE analyticalFourierSeriesFunc2 #-}
analyticalFourierSeriesFunc2 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e))
  => Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> e
  -> Complex e
analyticalFourierSeriesFunc2 angularFreq radialFreq sigma periodR2 periodEnv phi rho =
  let radialConst = 2 * pi / (log periodEnv)
  in pi * ((0 :+ 1) ^ (abs angularFreq)) *
     ((cis (fromIntegral angularFreq * phi))) *
     ((periodR2 / (pi * rho) :+ 0) **
      ((2 + sigma) :+ (radialConst * fromIntegral radialFreq))) *
     (gamma $
      ((2 + fromIntegral (abs angularFreq) + sigma) :+
       (radialConst * fromIntegral radialFreq)) /
      2) /
     (gamma $
      ((fromIntegral (abs angularFreq) - sigma) :+
       (radialConst * fromIntegral (-radialFreq))) /
      2)

-- The pinwheel is in the frequency domain, the function computes the Fourier series
{-# INLINE analyticalFourierSeries2 #-}
analyticalFourierSeries2 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierSeries2 numPoints delta angularFreq radialFreq sigma periodR2 periodEnv =
  let center = div numPoints 2
      arr =
        fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
          let x = delta * (fromIntegral $ i - center)
              y = delta * (fromIntegral $ j - center)
              rho = sqrt $ x ^ 2 + y ^ 2
              phi = atan2 y x
          in if rho == 0 -- || (rho > period / 2) -- || phi /= 0
               then 0
               else analyticalFourierSeriesFunc2
                      angularFreq
                      radialFreq
                      sigma
                      periodR2
                      periodEnv
                      phi
                      rho
     -- if angularFreq == 0
     --   then fromFunction (Z :. numPoints :. numPoints) $ \_ -> 0
     --   else
  in arr

-- The pinwheel is in the spatial domain, the function computes the Fourier coefficients
{-# INLINE analyticalFourierCoefficients2 #-}
analyticalFourierCoefficients2 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierCoefficients2 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv =
  let c = ((-1) ^ (abs angularFreq)) / periodR2 :+ 0
  in R.map (* c) $
     analyticalFourierSeries2
       numFreqs
       delta
       angularFreq
       radialFreq
       sigma
       periodR2
       periodEnv

-- Analytical solution type 3
-- The pinwheel harmonics are cis (-freq * x)
-- The envelope is r^2 * r^\alpha, where \alpha \in (-2,-0.5)
{-# INLINE analyticalFourierSeriesFunc3 #-}
analyticalFourierSeriesFunc3 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e))
  => Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> e
  -> Complex e
analyticalFourierSeriesFunc3 angularFreq radialFreq sigma periodR2 periodEnv phi rho =
  let radialConst = 2 * pi / log periodEnv
   in ((fromIntegral angularFreq ^ 2 :+ 0) -
       ((sigma + 2) :+ (radialConst * fromIntegral (-radialFreq))) ^ 2) /
      (8 * pi) *
      ((0 :+ 1) ^ abs angularFreq) *
      cis (fromIntegral (-angularFreq) * phi) /
      (rho ^ 2 :+ 0) *
      ((periodR2 / (pi * rho) :+ 0) **
       ((2 + sigma) :+ (radialConst * fromIntegral (-radialFreq)))) *
      gamma
        (((2 + fromIntegral (abs angularFreq) + sigma) :+
          (radialConst * fromIntegral (-radialFreq))) /
         2) /
      gamma
        (((fromIntegral (abs angularFreq) - sigma) :+
          (radialConst * fromIntegral radialFreq)) /
         2)

{-# INLINE analyticalFourierSeries3 #-}
analyticalFourierSeries3 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierSeries3 numPoints delta angularFreq radialFreq sigma periodR2 periodEnv =
  let center = div numPoints 2
  in fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
       let x = fromIntegral $ i - center
           y = fromIntegral $ j - center
           rho = sqrt $ x ^ 2 + y ^ 2
           phi = atan2 y x
       in if rho == 0
            then 0
            else analyticalFourierSeriesFunc3
                   angularFreq
                   radialFreq
                   sigma
                   periodR2
                   periodEnv
                   phi
                   (rho * delta)

{-# INLINE analyticalFourierCoefficients3 #-}
analyticalFourierCoefficients3 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierCoefficients3 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv =
  let c = ((-1) ^ (abs angularFreq)) :+ 0
  in R.map (* c) $
     analyticalFourierSeries3 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv

pinwheelFourierCoefficientsAnatical ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> IA.Array (Int, Int) (vector (Complex e))
pinwheelFourierCoefficientsAnatical numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- pinwheelFreqs rhoFreq rFreq
        , angularFreq <- pinwheelFreqs phiFreq thetaFreq
        ]
      centerR2Freq = div numR2Freqs 2
      (radialLB, radialUB) = pinwheelFreqsBound rhoFreq rFreq
      (angularLB, angularUB) = pinwheelFreqsBound phiFreq thetaFreq
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert .
             toUnboxed .
             computeS -- .
             -- R.zipWith
             --   (*)
             --   (fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
             --      gaussian2D
             --        (fromIntegral $ i - centerR2Freq)
             --        (fromIntegral $ j - centerR2Freq)
             --        (fromIntegral $ div centerR2Freq 1))
            $
             analyticalFourierCoefficients2
               numR2Freqs
               1
               (angularFreq)
               (radialFreq)
               sigma
               periodR2
               periodEnv)
          idxs
  in listArray ((radialLB, angularLB), (radialUB, angularUB)) pinwheels

pinwheelFourierCoefficientsAnaticalList ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
pinwheelFourierCoefficientsAnaticalList numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- [-rhoFreq .. rhoFreq]
        , angularFreq <- [-phiFreq .. phiFreq]
        ]
      centerR2Freq = div numR2Freqs 2
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert . toUnboxed . computeS $
             analyticalFourierCoefficients2
               numR2Freqs
               1
               (angularFreq)
               (radialFreq)
               sigma
               periodR2
               periodEnv)
          idxs
  in pinwheels

pinwheelFourierCoefficientsAnaticalList1 ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
pinwheelFourierCoefficientsAnaticalList1 numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- [-rhoFreq .. rhoFreq]
        , angularFreq <- [-phiFreq .. phiFreq]
        ]
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert . toUnboxed . computeS $
             analyticalFourierCoefficients2
               numR2Freqs
               1
               (angularFreq)
               (radialFreq)
               sigma
               periodR2
               periodEnv)
          idxs
  in pinwheels

pinwheelFourierCoefficientsAnaticalList2 ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
pinwheelFourierCoefficientsAnaticalList2 numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- [rFreq,(rFreq - 1) .. -rFreq]
        , angularFreq <- [thetaFreq,(thetaFreq - 1) .. -thetaFreq]
        ]
      centerR2Freq = div numR2Freqs 2
      radialConstant = 2 * pi / (log periodEnv)
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert . toUnboxed . computeS $
             fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. x :. y) ->
               let xFreq = fromIntegral $ x - centerR2Freq
                   yFreq = fromIntegral $ y - centerR2Freq
                   rho = sqrt $ xFreq ^ 2 + yFreq ^ 2
                   phi = atan2 yFreq xFreq
               in if rho == 0
                    then 0
                    else ((0 :+ (-1)) ^ 1) *
                         ((periodR2 / (pi * rho) :+ 0) **
                          (0 :+ fromIntegral radialFreq * radialConstant)) *
                         (cis $ phi * fromIntegral angularFreq))
          idxs
  in pinwheels

{-# INLINE analyticalFourierSeries2' #-}
analyticalFourierSeries2' ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierSeries2' numPoints _ angularFreq radialFreq sigma periodR2 periodEnv =
  let center = div numPoints 2
      arr =
        fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
          let x' = i - center
              y' = j - center
          in if x' == 0 && y' == 0
               then 0
               else let x =
                          if x' == 0
                            then 0
                            else if x' > 0
                                   then 1 / (fromIntegral $ center + 1 - x')
                                   else (-1) / (fromIntegral $ center + 1 + x')
                        y =
                          if y' == 0
                            then 0
                            else if y' > 0
                                   then 1 / (fromIntegral $ center + 1 - y')
                                   else (-1) / (fromIntegral $ center + 1 + y')
                        rho = sqrt $ x ^ 2 + y ^ 2
                        phi = atan2 y x
                    in analyticalFourierSeriesFunc2
                         angularFreq
                         radialFreq
                         sigma
                         periodR2
                         periodEnv
                         phi
                         rho
  in arr

{-# INLINE analyticalFourierCoefficients2' #-}
analyticalFourierCoefficients2' ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierCoefficients2' numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv =
  let c = (-1) ^ (abs angularFreq) :+ 0
  in R.map (* c) $
     analyticalFourierSeries2'
       numFreqs
       delta
       angularFreq
       radialFreq
       sigma
       periodR2
       periodEnv

pinwheelFourierCoefficientsAnatical' ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> IA.Array (Int, Int) (vector (Complex e))
pinwheelFourierCoefficientsAnatical' numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- pinwheelFreqs rhoFreq rFreq
        , angularFreq <- pinwheelFreqs phiFreq thetaFreq
        ]
      centerR2Freq = div numR2Freqs 2
      (radialLB, radialUB) = pinwheelFreqsBound rhoFreq rFreq
      (angularLB, angularUB) = pinwheelFreqsBound phiFreq thetaFreq
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert .
             toUnboxed .
             computeS -- .
             -- R.zipWith
             --   (*)
             --   (fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
             --      gaussian2D
             --        (fromIntegral $ i - centerR2Freq)
             --        (fromIntegral $ j - centerR2Freq)
             --        (fromIntegral $ div centerR2Freq 2))
            $
             analyticalFourierCoefficients2'
               numR2Freqs
               1
               (angularFreq)
               (radialFreq)
               sigma
               periodR2
               periodEnv)
          idxs
  in listArray ((radialLB, angularLB), (radialUB, angularUB)) pinwheels

{-# INLINE idealHighPassFilter #-}
idealHighPassFilter ::
     (VG.Vector vector (Complex Double))
  => Double
  -> Int
  -> vector (Complex Double)
idealHighPassFilter radius numR2Freq =
  let a = 1 / (radius * 2)
      r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      sinc x =
        if x == 0
          then 1
          else (sin (pi * x)) / (pi * x)
  in VG.fromList
       [ if xFreq == 0 && yFreq == 0
         then (1 - 1 / a ^ 2) :+ 0
         else ((-1) * (sinc (xFreq / a)) * (sinc (yFreq / a)) / (a ^ 2)) :+ 0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE idealHighPassFilter1 #-}
idealHighPassFilter1 :: Int -> VS.Vector (Complex Double)
idealHighPassFilter1 numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ if xFreq == 0 && yFreq == 0
         then 0
         else (-1) :+ 0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE gaussianHighPassFilter #-}
gaussianHighPassFilter :: Double -> Int -> VS.Vector (Complex Double)
gaussianHighPassFilter alpha numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ if xFreq == 0 && yFreq == 0
         then (1 - (pi / alpha)) :+ 0
         else ((-1) * (pi / alpha) *
               exp ((pi ^ 2 * (xFreq ^ 2 + yFreq ^ 2)) / (-alpha))) :+
              0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE gaussianHighPassFilter1 #-}
gaussianHighPassFilter1 :: Double -> Double -> Int -> VS.Vector (Complex Double)
gaussianHighPassFilter1 radius alpha numR2Freq =
  let a = 1 / (radius * 2)
      r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      sinc x =
        if x == 0
          then 1
          else (sin (pi * x)) / (pi * x)
  in VS.fromList
       [ (sinc (xFreq / a) * sinc (yFreq / a) / (a ^ 2) -
          (pi / alpha) * exp ((pi ^ 2 * (xFreq ^ 2 + yFreq ^ 2)) / (-alpha))) :+
       0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE idealLowPassFilter #-}
idealLowPassFilter ::
     (VG.Vector vector (Complex Double))
  => Double
  -> Double
  -> Int
  -> vector (Complex Double)
idealLowPassFilter radius periodR2 numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
   in VG.fromList
        [ let rho = 2 * pi * sqrt (xFreq ^ 2 + yFreq ^ 2) / periodR2
           in if rho == 0
                then 0
                else radius / rho  * bessel_J1 (radius * rho) / periodR2 / (2 * pi)^2 :+ 0
        | xFreq <- r2Freqs
        , yFreq <- r2Freqs
        ]

{-# INLINE gaussianLowPassFilter #-}
gaussianLowPassFilter :: Double -> Int -> VS.Vector (Complex Double)
gaussianLowPassFilter std numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ ((std ^ 2) * exp ((-pi) * (xFreq ^ 2 + yFreq ^ 2) * (std ^ 2))) :+
       0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE laplacianLowPassFilter #-}
laplacianLowPassFilter :: Double -> Int -> VS.Vector (Complex Double)
laplacianLowPassFilter a numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ let s2 = xFreq ^ 2 + yFreq ^ 2
         in if xFreq == 0 && yFreq == 0
              then (1 - a / ((4 * pi * pi * s2 + a ^ 2) ** 1.5)) :+ 0
              else (-a) / ((4 * pi * pi * s2 + a ^ 2) ** 1.5) :+ 0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE laplacianHighPassFilter #-}
laplacianHighPassFilter :: Int -> VS.Vector (Complex Double)
laplacianHighPassFilter numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ if xFreq == 0 && yFreq == 0
         then 8 -- -3.33
         else if xFreq == 0 && abs yFreq == 1
                then -1 -- 0.67
                else if abs xFreq == 1 && yFreq == 0
                       then -1 -- 0.67
                       else if abs xFreq == 1 && abs yFreq == 1
                              then -1 -- 0.17
                              else 0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE convolveFrequency #-}
convolveFrequency ::
     (VG.Vector vector (Complex Double))
  => DFTPlan
  -> Int
  -> vector (Complex Double)
  -> IA.Array (Int, Int) (vector (Complex Double))
  -> IO (IA.Array (Int, Int) (vector (Complex Double)))
convolveFrequency plan numR2Freq filter' arr = do
  let dftPlanID = DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]
      idftPlanID = DFTPlanID IDFT1DG [numR2Freq, numR2Freq] [0, 1]
      vecs = L.map VG.convert . IA.elems $ arr
      filter =
        VU.convert .
        toUnboxed .
        computeUnboxedS .
        makeFilter2D . fromUnboxed (Z :. numR2Freq :. numR2Freq) . VG.convert $
        filter'
  dftF <- dftExecute plan dftPlanID filter
  vecsF <- mapConcurrently (dftExecute plan dftPlanID) vecs
  outputs <-
    mapConcurrently (dftExecute plan idftPlanID . VS.zipWith (*) dftF) vecsF
  return . listArray (bounds arr) . L.map VG.convert $ outputs

{-# INLINE convolveFrequency1 #-}
convolveFrequency1 ::
     (VG.Vector vector (Complex Double))
  => DFTPlan
  -> Int
  -> vector (Complex Double)
  -> IA.Array (Int, Int) (vector (Complex Double))
  -> IO (IA.Array (Int, Int) (vector (Complex Double)))
convolveFrequency1 plan numR2Freq filter' arr = do
  let dftPlanID = DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]
      idftPlanID = DFTPlanID IDFT1DG [numR2Freq, numR2Freq] [0, 1]
      vecs = L.map VG.convert . IA.elems $ arr
      filter =
        VU.convert .
        toUnboxed .
        computeUnboxedS .
        makeFilter2D . fromUnboxed (Z :. numR2Freq :. numR2Freq) . VG.convert $
        filter'
  dftF <- dftExecute plan dftPlanID filter
  vecsF <- mapConcurrently (dftExecute plan dftPlanID) vecs
  outputs <-
    mapConcurrently (dftExecute plan idftPlanID . VS.zipWith (*) dftF) vecsF
  return .
    listArray (bounds arr) . L.map VG.convert . L.zipWith (VG.zipWith (-)) vecs $
    outputs

{-# INLINE centerHollow #-}
centerHollow ::
     (VG.Vector vector (Complex Double))
  => Int
  -> IA.Array (Int, Int) (vector (Complex Double))
  -> IA.Array (Int, Int) (vector (Complex Double))
centerHollow numR2Freq arr = centerHollowVector numR2Freq <$> arr

{-# INLINE centerHollowVector #-}
centerHollowVector ::
     (VG.Vector vector (Complex Double))
  => Int
  -> vector (Complex Double)
  -> vector (Complex Double)
centerHollowVector numR2Freq vec =
  let s = VG.sum vec / (fromIntegral (numR2Freq ^ 2) :+ 0)
   in VG.map (\x -> x - s) vec

{-# INLINE centerHollowArray #-}
centerHollowArray ::
     (R.Source s (Complex e), Unbox e, RealFloat e)
  => Int
  -> R.Array s DIM2 (Complex e)
  -> R.Array D DIM2 (Complex e)
centerHollowArray numR2Freq arr =
  let s = sumAllS arr / (fromIntegral (numR2Freq ^ 2) :+ 0)
   in R.map (\x -> x - s) $ arr

{-# INLINE centerHollowArray' #-}
centerHollowArray' ::
     (R.Source s (Complex e), Unbox e, RealFloat e)
  => Int
  -> R.Array s DIM2 (Complex e)
  -> R.Array D DIM2 (Complex e)
centerHollowArray' numR2Freq arr =
  let s = sumAllS arr
      c = div numR2Freq 2
   in R.traverse arr id $ \f idx@(Z :. i :. j) ->
        if i == c && j == c
          then (-s)
          else f idx

envelopIntegral ::
     Double -> Double -> Double -> Double -> Double -> Int -> Complex Double
envelopIntegral a b delta s period radialFreq =
  let m = round $ (b - a) / delta
      n =
        if odd m
          then m
          else m - 1
      -- weights = VU.fromList $ weightsSimpsonRule n
      vec =
        VU.generate n $ \i ->
          let x = (a + fromIntegral i * delta)
          -- in   ((x :+ 0) ** (s :+ freq)) * ((pi :+ 0) ** (0 :+ freq)) *
          --      (gamma ((0.5 :+ 0) * (1 :+ (-freq)))) /
          --      (gamma ((0.5 :+ 0) * (1 :+ freq)))
          in analyticalFourierSeriesFunc2 0 radialFreq s period period 0 x
  -- in (delta / 3 :+ 0) * (VU.sum . VU.zipWith (*) vec $ weights)
  in VU.sum vec

envelopIntegral2D ::
     Int -> Double -> Double -> Double -> Int -> Complex Double
envelopIntegral2D numPoints delta s period radialFreq =
  let center = div numPoints 2
      -- weights = VU.fromList $ weightsSimpsonRule n
      arr = analyticalFourierSeries2 numPoints delta 0 radialFreq s period period
  in sumAllS arr

printEnvelopIntegral ::
     Double -> Double -> Double -> Double -> Double -> Int -> IO ()
printEnvelopIntegral a b delta s period radialFreq =
  let m = round $ (b - a) / delta
      n =
        if odd m
          then m
          else m - 1
      weights = VU.fromList $ weightsSimpsonRule n
      vec =
        VU.generate n $ \i ->
          let x = a + fromIntegral i * delta
          in analyticalFourierSeriesFunc2 0 radialFreq s period period 0 x
  in print $ VU.zipWith (*) weights vec

{-# INLINE createFrequencyArray #-}
createFrequencyArray ::
     (RealFloat e) => Int -> (e -> e -> Complex e) -> R.Array D DIM2 (Complex e)
createFrequencyArray numR2Freqs f =
  fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
    let x = fromIntegral $ i - div numR2Freqs 2
        y = fromIntegral $ j - div numR2Freqs 2
        rho = sqrt $ x ^ 2 + y ^ 2
        phi = atan2 y x
     in f phi rho

{-# INLINE createFrequencyArrayZeroCenter #-}
createFrequencyArrayZeroCenter ::
     (RealFloat e) => Int -> (e -> e -> Complex e) -> R.Array D DIM2 (Complex e)
createFrequencyArrayZeroCenter numR2Freqs f =
  fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
    let x = fromIntegral $ i - div numR2Freqs 2
        y = fromIntegral $ j - div numR2Freqs 2
        rho = sqrt $ x ^ 2 + y ^ 2
        phi = atan2 y x
     in if rho == 0
          then 0
          else f phi rho


{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE CPP                  #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}
{-|
Module      : Grenade.Core.Shape
Description : Dependently typed shapes of data which are passed between layers of a network
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

-}
module Grenade.Core.Shape (
    S (..)
  , Shape (..)
#if MIN_VERSION_singletons(2,6,0)
  , SShape (..)
#else
  , Sing (..)
#endif

  , randomOfShape
  , fromStorable
  , fromStorableV
  , isVectorShape
  ) where

#if MIN_VERSION_singletons(2,6,0)
import           Data.Kind                    (Type)
import           Data.Singletons.TypeLits     (SNat (..))
#endif

import           Control.DeepSeq              (NFData (..))
import           Data.Maybe                   (fromMaybe)
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Vector.Storable         (Vector)
import qualified Data.Vector.Storable         as V
import           GHC.TypeLits                 hiding (natVal)
import qualified Numeric.LinearAlgebra        as NLA
import           Numeric.LinearAlgebra.Static hiding (zipWithVector)
import qualified Numeric.LinearAlgebra.Static as H
import           System.Random.MWC
import           Unsafe.Coerce                (unsafeCoerce)

import           Debug.Trace

import           Grenade.Types
import           Grenade.Utils.Vector

-- | The current shapes we accept.
--   at the moment this is just one, two, and three dimensional
--   Vectors/Matricies.
--
--   These are only used with DataKinds, as Kind `Shape`, with Types 'D1, 'D2, 'D3.
data Shape
  = D1 Nat
  -- ^ One dimensional vector
  | D2 Nat Nat
  -- ^ Two dimensional matrix. Row, Column.
  | D3 Nat Nat Nat
  -- ^ Three dimensional matrix. Row, Column, Channels.

-- | Concrete data structures for a Shape.
--
--   All shapes are held in contiguous memory.
--   3D is held in a matrix (usually row oriented) which has height depth * rows.
data S (n :: Shape) where
  S1D :: ( KnownNat len )
      => R len
      -> S ('D1 len)

  S2D :: ( KnownNat rows, KnownNat columns )
      => L rows columns
      -> S ('D2 rows columns)

  S3D :: ( KnownNat rows
         , KnownNat columns
         , KnownNat depth
         , KnownNat (rows * depth))
      => L (rows * depth) columns
      -> S ('D3 rows columns depth)

  -- HMatrix instances
  S1DV :: (KnownNat len) -- ^ Always a row-vector
       => !(V.Vector RealNum)
       -> S ('D1 len)

  -- HMatrix instances
  S2DV :: (KnownNat rows, KnownNat columns)
       => !(V.Vector RealNum) -- ^ Vector in row or column major accoring to @order@ in @Grenade.Utils.Conversion@. Use Data.Matrix here?
       -> S ('D2 rows columns)

  -- HMatrix instances
  -- S3DH :: (KnownNat rows, KnownNat columns, KnownNat depth, KnownNat (rows * depth))
  --      => HBLAS.MDenseMatrix RealWorld 'HBLAS.Row Double
  --      -> S ('D3 rows columns depth)

isVectorShape :: S x -> Bool
isVectorShape S1DV{} = True
isVectorShape S2DV{} = True
isVectorShape _      = False

instance Show (S n) where
  show (S1D x) = "S1D " ++ show x
  show (S2D x) = "S2D " ++ show x
  show (S3D x) = "S3D " ++ show x
  show (S1DV x) = "S1DV " ++ show x
  show inp@(S2DV x) = "S2DV" ++ show (sz inp) ++ " " ++ show x
    where
      sz ::
           forall rows cols. (KnownNat rows, KnownNat cols)
        => (S ('D2 rows cols))
        -> (Int, Int)
      sz _ =
        let rows = fromIntegral $ natVal (Proxy :: Proxy rows)
            cols = fromIntegral $ natVal (Proxy :: Proxy cols)
         in (rows, cols)

-- Singleton instances.
--
-- These could probably be derived with template haskell, but this seems
-- clear and makes adding the KnownNat constraints simple.
-- We can also keep our code TH free, which is great.
#if MIN_VERSION_singletons(2,6,0)
-- In singletons 2.6 Sing switched from a data family to a type family.
type instance Sing = SShape

data SShape :: Shape -> Type where
  D1Sing :: Sing a -> SShape ('D1 a)
  D2Sing :: Sing a -> Sing b -> SShape ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> SShape ('D3 a b c)
#else
data instance Sing (n :: Shape) where
  D1Sing :: Sing a -> Sing ('D1 a)
  D2Sing :: Sing a -> Sing b -> Sing ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> Sing ('D3 a b c)
#endif

instance KnownNat a => SingI ('D1 a) where
  sing = D1Sing sing
instance (KnownNat a, KnownNat b) => SingI ('D2 a b) where
  sing = D2Sing sing sing
instance (KnownNat a, KnownNat b, KnownNat c, KnownNat (a * c)) => SingI ('D3 a b c) where
  sing = D3Sing sing sing sing

instance SingI x => Num (S x) where
  (+) = n2 (+)
  (-) = n2 (-)
  (*) = n2 (*)
  abs = n1 abs
  signum = n1 signum
  fromInteger x = nk (fromInteger x)

instance SingI x => Fractional (S x) where
  (/) = n2 (/)
  recip = n1 recip
  fromRational x = nk (fromRational x)

instance SingI x => Floating (S x) where
  pi = nk pi
  exp = n1 exp
  log = n1 log
  sqrt = n1 sqrt
  (**) = n2 (**)
  logBase = n2 logBase
  sin = n1 sin
  cos = n1 cos
  tan = n1 tan
  asin = n1 asin
  acos = n1 acos
  atan = n1 atan
  sinh = n1 sinh
  cosh = n1 cosh
  tanh = n1 tanh
  asinh = n1 asinh
  acosh = n1 acosh
  atanh = n1 atanh

--
-- I haven't made shapes strict, as sometimes they're not needed
-- (the last input gradient back for instance)
--
instance NFData (S x) where
  rnf (S1D x)   = rnf x
  rnf (S2D x)   = rnf x
  rnf (S3D x)   = rnf x
  rnf (S1DV !v) = rnf v
  rnf (S2DV !v) = rnf v
  -- rnf (S3DH !_) = ()

-- | Generate random data of the desired shape
randomOfShape :: forall x . (SingI x) => IO (S x)
randomOfShape = do
  seed :: Int <- withSystemRandom . asGenST $ \gen -> uniform gen
  return $ case (sing :: Sing x) of
    D1Sing SNat ->
        S1D (randomVector seed Uniform * 2 - 1)

    D2Sing SNat SNat ->
        S2D (uniformSample seed (-1) 1)

    D3Sing SNat SNat SNat ->
        S3D (uniformSample seed (-1) 1)

-- | Generate a shape from a Storable Vector.
--
--   Returns Nothing if the vector is of the wrong size.
fromStorable :: forall x. SingI x => Vector RealNum -> Maybe (S x)
fromStorable xs =
  case (sing :: Sing x) of
    D1Sing SNat           -> S1D <$> H.create xs
    D2Sing SNat SNat      -> S2D <$> mkL xs
    D3Sing SNat SNat SNat -> S3D <$> mkL xs
  where
    mkL ::
         forall rows columns. (KnownNat rows, KnownNat columns)
      => Vector RealNum
      -> Maybe (L rows columns)
    mkL v =
      let rows = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns = fromIntegral $ natVal (Proxy :: Proxy columns)
       in if rows * columns == V.length v
            then H.create $ NLA.reshape columns v
            else Nothing

fromStorableV :: forall x. SingI x => Vector RealNum -> S x
fromStorableV v =
  case (sing :: Sing x) of
    D1Sing SNat           -> S1DV v
    D2Sing SNat SNat      -> S2DV v
    D3Sing SNat SNat SNat -> error "unexpected case in fromStorableS"

instance SingI x => Serialize (S x) where
  put i =
    case i of
      S1D x -> put (1 :: Int) >> (putListOf put . NLA.toList . H.extract $ x)
      S2D x -> put (1 :: Int) >> (putListOf put . NLA.toList . NLA.flatten . H.extract $ x)
      S3D x -> put (1 :: Int) >> (putListOf put . NLA.toList . NLA.flatten . H.extract $ x)
      S1DV x -> put (2 :: Int) >> putListOf put (V.toList x)
      S2DV x -> put (2 :: Int) >> putListOf put (V.toList x)
  get = do
    (nr :: Int) <- get
    case nr of
      1 -> do
        Just i <- fromStorable . V.fromList <$> getListOf get
        return i
      2 -> fromStorableV . V.fromList <$> getListOf get
      _ -> error "unexpected case in get in Serialize instance in Shape.hs"

-- Helper function for creating the number instances
n1 :: ( forall a. Floating a => a -> a ) -> S x -> S x
n1 f (S1D x)  = S1D (f x)
n1 f (S2D x)  = S2D (f x)
n1 f (S3D x)  = S3D (f x)
n1 f (S1DV x) = S1DV (mapVector f x)
n1 f (S2DV x) = S2DV (mapVector f x)

-- helper function for creating the number instances
n2 :: ( forall a. Floating a => a -> a -> a ) -> S x -> S x -> S x
n2 f (S1D x) (S1D y)    = S1D (f x y)
n2 f (S2D x) (S2D y)    = S2D (f x y)
n2 f (S3D x) (S3D y)    = S3D (f x y)
n2 f (S1DV x) (S1DV y)  = S1DV (zipWithVector f x y)
n2 f (S2DV x) (S2DV y)  = S2DV (zipWithVector f x y)
n2 f (S1D x) y@S1DV {}  = n2 f (S1DV $ extract x) y
n2 f x@S1DV {} (S1D y)  = n2 f x (S1DV $ unsafeCoerce y)
n2 f x@S2D{} y@(S2DV _) = n2 f (toS2DV x) y
n2 f x@(S2DV _) y@S2D{} = n2 f x (toS2DV y)

toS2DV :: S ('D2 i j) -> S ('D2 i j)
toS2DV (S2D x)  = S2DV $ V.concat $ map H.extract . H.toColumns $ x
toS2DV (S2DV x) = S2DV x

-- Helper function for creating the number instances
nk :: forall x. SingI x => RealNum -> S x
nk x = case (sing :: Sing x) of
  D1Sing SNat ->
    S1D (konst x)

  D2Sing SNat SNat ->
    S2D (konst x)

  D3Sing SNat SNat SNat ->
    S3D (konst x)


{-# LANGUAGE DataKinds #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

import FrequencyResponse

import CLaSH.Prelude
import CLaSH.Signal.MultiSignal
import Data.Complex as C
import qualified Prelude as P
import Data.Maybe
import Control.Applicative
import Graphics.EasyPlot
import Data.Map.Strict as Map

{-
Ciruits and stuff
-}

registerP = prepend

firP coeffs x = dotp coeffs (windowP x)
  where
    dotp as bs = sum (zipWith (*) as bs)

iirP cA cB x = r where
    oB =  firP cB x
    oA =  firP cA (registerP def r)
    r = oB - oA

-- fir coeficients
firCoef :: Fractional a => Vec 201 a
firCoef = -6.86056317e-19 :> 1.08643703e-04 :> -6.04051550e-18 :> -2.03379801e-04 :> -2.01638331e-04 :> 1.96946883e-18 :> 2.90791068e-18 :> -4.62180925e-04 :> -1.00665733e-03 :> -9.14037108e-04 :> 3.86774364e-18 :> 1.01950810e-03 :> 1.25131737e-03 :> 6.39186084e-04 :> -8.00554425e-18 :> 2.62712857e-18 :> 3.77239653e-04 :> 4.17356596e-04 :> -1.21745475e-17 :> -2.64065618e-04 :> -5.23918059e-18 :> 2.99576343e-04 :> -2.73859684e-18 :> -6.08801988e-04 :> -6.23370976e-04 :> 1.63484509e-17 :> -7.03525331e-18 :> -1.52438502e-03 :> -3.36082345e-03 :> -3.07661933e-03 :> -1.30788363e-17 :> 3.45259832e-03 :> 4.23241069e-03 :> 2.15427395e-03 :> 0.00000000e+00 :> 2.70383448e-17 :> 1.24466846e-03 :> 1.36380726e-03 :> -2.92823518e-17 :> -8.44354542e-04 :> -4.24392593e-17 :> 9.35427798e-04 :> 2.64177628e-17 :> -1.85473668e-03 :> -1.87576073e-03 :> 2.24314248e-18 :> -5.19238250e-17 :> -4.42290170e-03 :> -9.63775624e-03 :> -8.72271164e-03 :> 3.66725201e-17 :> 9.57794343e-03 :> 1.16209519e-02 :> 5.85688426e-03 :> -6.95248052e-17 :> 1.32593176e-17 :> 3.29422893e-03 :> 3.58089966e-03 :> 9.08965047e-18 :> -2.18539647e-03 :> -2.12919221e-17 :> 2.39194534e-03 :> -8.31327663e-17 :> -4.69672484e-03 :> -4.73138591e-03 :> 5.79766610e-17 :> -5.28939995e-18 :> -1.10712423e-02 :> -2.40989091e-02 :> -2.18047819e-02 :> -6.11496326e-17 :> 2.39909500e-02 :> 2.91782046e-02 :> 1.47558239e-02 :> 3.11706634e-17 :> 1.33673427e-18 :> 8.44218666e-03 :> 9.25308362e-03 :> -9.62845645e-17 :> -5.76785543e-03 :> -2.84709264e-17 :> 6.49485159e-03 :> -4.19102125e-17 :> -1.32387105e-02 :> -1.36437840e-02 :> 1.16146829e-16 :> -4.98075231e-17 :> -3.49290893e-02 :> -7.90710907e-02 :> -7.48583001e-02 :> -3.01478553e-33 :> 9.23333495e-02 :> 1.20794568e-01 :> 6.66720749e-02 :> -9.64385396e-17 :> 4.65544367e-17 :> 5.76578486e-02 :> 7.99902785e-02 :> -8.32650050e-17 :> -1.34035769e-01 :> 8.00710266e-01 :> -1.34035769e-01 :> -8.32650050e-17 :> 7.99902785e-02 :> 5.76578486e-02 :> 4.65544367e-17 :> -9.64385396e-17 :> 6.66720749e-02 :> 1.20794568e-01 :> 9.23333495e-02 :> -3.01478553e-33 :> -7.48583001e-02 :> -7.90710907e-02 :> -3.49290893e-02 :> -4.98075231e-17 :> 1.16146829e-16 :> -1.36437840e-02 :> -1.32387105e-02 :> -4.19102125e-17 :> 6.49485159e-03 :> -2.84709264e-17 :> -5.76785543e-03 :> -9.62845645e-17 :> 9.25308362e-03 :> 8.44218666e-03 :> 1.33673427e-18 :> 3.11706634e-17 :> 1.47558239e-02 :> 2.91782046e-02 :> 2.39909500e-02 :> -6.11496326e-17 :> -2.18047819e-02 :> -2.40989091e-02 :> -1.10712423e-02 :> -5.28939995e-18 :> 5.79766610e-17 :> -4.73138591e-03 :> -4.69672484e-03 :> -8.31327663e-17 :> 2.39194534e-03 :> -2.12919221e-17 :> -2.18539647e-03 :> 9.08965047e-18 :> 3.58089966e-03 :> 3.29422893e-03 :> 1.32593176e-17 :> -6.95248052e-17 :> 5.85688426e-03 :> 1.16209519e-02 :> 9.57794343e-03 :> 3.66725201e-17 :> -8.72271164e-03 :> -9.63775624e-03 :> -4.42290170e-03 :> -5.19238250e-17 :> 2.24314248e-18 :> -1.87576073e-03 :> -1.85473668e-03 :> 2.64177628e-17 :> 9.35427798e-04 :> -4.24392593e-17 :> -8.44354542e-04 :> -2.92823518e-17 :> 1.36380726e-03 :> 1.24466846e-03 :> 2.70383448e-17 :> 0.00000000e+00 :> 2.15427395e-03 :> 4.23241069e-03 :> 3.45259832e-03 :> -1.30788363e-17 :> -3.07661933e-03 :> -3.36082345e-03 :> -1.52438502e-03 :> -7.03525331e-18 :> 1.63484509e-17 :> -6.23370976e-04 :> -6.08801988e-04 :> -2.73859684e-18 :> 2.99576343e-04 :> -5.23918059e-18 :> -2.64065618e-04 :> -1.21745475e-17 :> 4.17356596e-04 :> 3.77239653e-04 :> 2.62712857e-18 :> -8.00554425e-18 :> 6.39186084e-04 :> 1.25131737e-03 :> 1.01950810e-03 :> 3.86774364e-18 :> -9.14037108e-04 :> -1.00665733e-03 :> -4.62180925e-04 :> 2.90791068e-18 :> 1.96946883e-18 :> -2.01638331e-04 :> -2.03379801e-04 :> -6.04051550e-18 :> 1.08643703e-04 :> -6.86056317e-19 :> Nil

-- iir coeficients
iirCoefA :: Fractional a => Vec 8 a
iirCoefB :: Fractional a => Vec 9 a
iirCoefA = -1.8933239813032006 :> 3.422153467072297 :> -3.907588921469733 :> 3.6401868030176012 :> -2.5428541496659465 :> 1.3470500696022936 :> -0.49424202583584054 :> 0.10443313376764539 :> Nil
iirCoefB = 0.0025798085212212327 :> 0.02063846816976986 :> 0.07223463859419452 :> 0.14446927718838903 :> 0.18058659648548628 :> 0.14446927718838903 :> 0.07223463859419452 :> 0.02063846816976986 :> 0.0025798085212212327 :> Nil

{-
Drawing helpers
-}

simFuncLin f skip ph =  snd $ getSpecLin <$> P.last $ P.take skip $ getResponseLin f ph
simFunc f skip ph =  (Map.! 1) $ getSpectrum <$> P.last $ P.take skip $ getResponse f ph
simFuncAcyclic f ph =  (Map.! 1) $ getSpectrum $ getResponseAcyclic f ph

genData f sk pn = (x,y) where
    x = fmap (\a -> (1.0 * pi * fromInteger a / fromInteger pn)) [0, 1 .. pn]
    y = fmap (simFunc f sk) x

-- how is mapping over result like Complex Double -> Double
-- f is ciruit
-- sk skip first sk values from result, for acyclic circuit this should be more than circuitry delay
-- pn - number of linear points from 0 to 1 (nyquist freq)
plt how f sk pn = plot' [] X11 $ Data2D [Style Lines,Title "plot"] [] (P.zip ((/pi) <$> x) (how <$> y))
  where
    (x,y) = genData f sk pn

genDataLin f sk pn = (x,y) where
    x = fmap (\a -> (1.0 * pi * fromInteger a / fromInteger pn)) [0, 1 .. pn]
    y = fmap (simFuncLin f sk) x

pltLin how f sk pn = plot' [] X11 $ Data2D [Style Lines,Title "plot"] [] (P.zip ((/pi) <$> x) (how <$> y))
  where
    (x,y) = genDataLin f sk pn

genDataAcyclic f pn = (x,y) where
    x = fmap (\a -> (1.0 * pi * fromInteger a / fromInteger pn)) [0, 1 .. pn]
    y = fmap (simFuncAcyclic f) x

pltDataAcyclic how f pn = plot' [] X11 $ Data2D [Style Lines,Title "plot"] [] (P.zip ((/pi) <$> x) (how <$> y))
  where
    (x,y) = genDataAcyclic f pn

magLog x = logBase 10 (magnitude x) * 20
magLogMin m x = max m (magLog x)

integrator i = r where
    r = (registerP 0 r) + i

{-
Examples
-}

-- plots working with linear circuits
plotFirLin = pltLin magnitude f 210 100 where
 f = firP firCoef

-- slower and generic version of samplePlotLin
plotFir = plt magnitude f 210 100 where
 f = firP firCoef

-- plot acyclic ciruit
plotFirAcyclic = pltDataAcyclic magnitude f 1000 where
 f = firP firCoef

plotIir = plt magnitude (iirP iirCoefA iirCoefB) 200 100
plotIirLog = pltLin magLog (iirP iirCoefA iirCoefB) 200 100

plotInteg = plt magnitude integrator 100 100

main = plotIir


import Codec.Picture
import Codec.Picture.Types
import Data.Word
import Data.Vector as V
import Data.Complex

adjust :: Double -> Int -> Int -> Double
adjust s xr x = (fromIntegral x - (fromIntegral xr)/2)*s

cmpmake :: Double -> Int -> Int -> Int -> Int -> Complex Double
cmpmake s xr yr x y = adjust s xr x :+ adjust s yr y

polartolch :: (Double, Double) -> Double -> (Double, Double, Double)
polartolch (mag, ph) high = (linearInter high 60 mag, 90, ph)

lchtolab :: (Double, Double, Double) -> (Double, Double, Double)
lchtolab (l, c, h) = (l, c * (cos h), c * (sin h))

labtoxyz :: (Double, Double, Double) -> (Double, Double, Double)
labtoxyz (l, a, b) = let  eps = 216/24389
                          kap = 24389/27
                          fy  = (l + 16)/116
                          fz  = fy - (b/200)
                          fx  = (a/500) + fy
                          x   = if fx**3 > eps then fx**3 else (116*fx - 16)/kap
                          y   = if l > kap*eps then ((l + 16)/116)**3 else l/kap
                          z   = if fz**3 > eps then fz**3 else (116*fz - 16)/kap
                            in (x, y, z)

--https://www.mathworks.com/help/images/ref/lab2xyz.html

a =   (1.0985,  1,          0.3558)
c =   (0.9807,  1,          1.1822)
e =   (1,       1,          1)
d50 = (0.9642,  1.0,        0.8251)
d55 = (0.9568,  1.0,        0.9214)
d65 = (0.95047, 1.00,       1.08883)
icc = (0.9642,  1,          0.8249)

data RGBgamut = RGBgamut {row1 :: (Double, Double, Double),
                          row2 :: (Double, Double, Double),
                          row3 :: (Double, Double, Double),
                          white :: (Double, Double, Double),
                          gamma :: Double}

sRGBgamut = RGBgamut (3.2404542, -1.5371385, -0.4985314)
                      (-0.9692660, 1.8760108, 0.0415560)
                      (0.0556434, -0.2040259, 1.0572252)
                      d65
                      2.4
adobeRGBgamut = RGBgamut (2.0413690, -0.5649464, -0.3446944)
                          (-0.9692660, 1.8760108, 0.0415560)
                          (0.0134474, -0.1183897, 1.0154096)
                          d65
                          (563/256)

adobeWideRGBgamut = RGBgamut (1.4628067, -0.1840623, -0.2743606)
                              (-0.5217933, 1.4472381, 0.0677227)
                              (0.0349342, -0.0968930, 1.2884099)
                              d50
                              (563/256)

proPhotoRGBgamut = RGBgamut (1.3459433, -0.2556075, -0.0511118)
                              (-0.5445989, 1.5081673, 0.0205351)
                              (0,       0,            1.2118128)
                              d50
                              1.8

pointwise :: (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double)
pointwise (a, b, c) (x, y, z) = (a*x, b*y, c*z)

matrixMult :: (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double)
matrixMult (a, b, c) (d, e, f) (g, h, i) (x, y, z) = (a*x + b*y + c*z, d*z + e*y + f*z, g*x + h*y + i*z)

xyztorgb :: RGBgamut -> (Double, Double, Double) -> (Double, Double, Double)
xyztorgb gam (x, y, z) =  let (xl, yl, zl)  = matrixMult (row1 gam) (row2 gam) (row3 gam) $ pointwise (x, y, z) (white gam)
                              compand c     = c**(1/(gamma gam))
                                in (compand xl, compand yl, compand zl)

linearInter :: Double -> Double -> Double -> Double
linearInter highin highout x = highout/highin * x

expInter :: Double -> Double -> Double -> Double
expInter low high x = 1 - base ** (x) where
  base = 0.5

makeSampleSpace :: Int -> Int -> Double -> Vector (Vector (Double, Double))
makeSampleSpace w h s = generate (2*h + 1) (\y -> generate (2*w + 1) (\x -> (s*(fromIntegral x - fromIntegral w), -s*(fromIntegral y - fromIntegral h))))

inmap :: (a -> b) -> Vector (Vector a) -> Vector (Vector b)
inmap f l = V.map (V.map f) l

findextrema :: Vector (Vector (Double, Double)) -> Double
findextrema lst =
  let lst' = V.concat $ toList lst
      lst'' = V.map (\(x, y) -> x) lst'
          in V.maximum lst''

highest :: Vector Double -> Double
highest vec = if V.length vec == 1 then V.head vec else
            (if V.head vec > (highest $ slice 1 (V.length vec - 1) vec) then V.head vec else (highest $ slice 1 (V.length vec - 1) vec))

urx = 1000
ury = 1000
iscale = 0.003

f :: Complex Double -> Complex Double
f z =  cexp $ z

cexp :: Complex Double -> Complex Double
cexp (a :+ b) = (exp a :+ 0) * (cis b)

newimage :: Vector (Vector PixelRGB16)
newimage = let  sampling = makeSampleSpace urx ury iscale
                polars = inmap (\(x,y) -> polar $ f $ x :+ y) sampling
                high = findextrema polars
                lch = inmap (\(x,y) -> polartolch (x,y) high) polars
                lab = inmap lchtolab lch
                xyz = inmap labtoxyz lab
                rgb = inmap (xyztorgb proPhotoRGBgamut) xyz
                rounded = inmap (\(r, g, b) -> PixelRGB16 (round $ r*65530) (round $ g*65530) (round $ b*65530)) rgb
                  in rounded

main :: IO ()
main = writePng "./blah.png" $ generateImage (\x y -> newimage ! y ! x) (2*urx + 1) (2*ury + 1)


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

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

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Class
import           Backprop.Learn.Model.Function
import           Backprop.Learn.Model.Neural
import           Backprop.Learn.Model.Regression
import           Backprop.Learn.Model.State
import           Control.DeepSeq
import           Data.Typeable
import           GHC.Generics                          (Generic)
import           GHC.TypeNats
import           Lens.Micro
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.OneLiner
import           Numeric.Opto.Ref
import           Numeric.Opto.Update
import qualified Data.Binary                           as Bi
import qualified Numeric.LinearAlgebra.Static          as H

-- | "Base" for 'LSTM'.  An 'LSTM' is just an 'LSTM'' where the second half
-- of the input vector is the previous output.
--
-- This is mostly an implementation detail.  It is recommended that you use
-- 'LSTM' and its constructor, instead of directly using this.
data LSTM' (i :: Nat) (o :: Nat) = LSTM'
  deriving (Typeable)

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

-- | 'LSTM' layer parmateters
data LSTMp (i :: Nat) (o :: Nat) =
    LSTMp { _lstmForget :: !(FCp (i + o) o)
          , _lstmInput  :: !(FCp (i + o) o)
          , _lstmUpdate :: !(FCp (i + o) o)
          , _lstmOutput :: !(FCp (i + o) o)
          }
  deriving (Generic, Typeable, Show)

instance NFData (LSTMp i o)
instance (KnownNat i, KnownNat o) => Additive (LSTMp i o)
instance (KnownNat i, KnownNat o) => Scaling Double (LSTMp i o)
instance (KnownNat i, KnownNat o) => Metric Double (LSTMp i o)
instance (KnownNat i, KnownNat o, Ref m (LSTMp i o) v) => AdditiveInPlace m v (LSTMp i o)
instance (KnownNat i, KnownNat o, Ref m (LSTMp i o) v) => ScalingInPlace m v Double (LSTMp i o)
instance (KnownNat i, KnownNat o) => Bi.Binary (LSTMp i o)
instance (KnownNat i, KnownNat o) => Backprop (LSTMp i o)

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

instance (KnownNat i, KnownNat o) => Num (LSTMp i o) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance (KnownNat i, KnownNat o) => Fractional (LSTMp i o) where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational

instance (KnownNat i, KnownNat o) => Floating (LSTMp i o) where
    pi    = gPi
    sqrt  = gSqrt
    exp   = gExp
    log   = gLog
    sin   = gSin
    cos   = gCos
    asin  = gAsin
    acos  = gAcos
    atan  = gAtan
    sinh  = gSinh
    cosh  = gCosh
    asinh = gAsinh
    acosh = gAcosh
    atanh = gAtanh

lstmForget :: Lens' (LSTMp i o) (FCp (i + o) o)
lstmForget f x = (\y -> x { _lstmForget = y }) <$> f (_lstmForget x)

lstmInput :: Lens' (LSTMp i o) (FCp (i + o) o)
lstmInput f x = (\y -> x { _lstmInput = y }) <$> f (_lstmInput x)

lstmUpdate :: Lens' (LSTMp i o) (FCp (i + o) o)
lstmUpdate f x = (\y -> x { _lstmUpdate = y }) <$> f (_lstmUpdate x)

lstmOutput :: Lens' (LSTMp i o) (FCp (i + o) o)
lstmOutput f x = (\y -> x { _lstmOutput = y }) <$> f (_lstmOutput x)

instance (KnownNat i, KnownNat o, KnownNat io, io ~ (i + o)) => Learn (R io) (R o) (LSTM' i o) where
    type LParamMaybe (LSTM' i o) = 'Just (LSTMp i o)
    type LStateMaybe (LSTM' i o) = 'Just (R o)

    runLearn LSTM' (J_ p) x (J_ s) = (h, J_ s')
      where
        forget = logistic $ runLRp (p ^^. lstmForget) x
        input  = logistic $ runLRp (p ^^. lstmInput ) x
        update = tanh     $ runLRp (p ^^. lstmUpdate) x
        s'     = forget * s + input * update
        o      = logistic $ runLRp (p ^^. lstmOutput) x
        h      = o * tanh s'

-- | Long-term short-term memory layer
--
-- <http://colah.github.io/posts/2015-08-Understanding-LSTMs/>
--
-- @
-- instance 'Learn' ('R' i) (R o) ('LSTM' i o) where
--     type 'LParamMaybe' (LSTM i o) = ''Just' ('LSTMp' i o)
--     type 'LStateMaybe' (LSTM i o) = 'Just ('T2' (R o) (R o))     -- cell state, hidden state
-- @
type LSTM i o = Recurrent (R (i + o)) (R i) (R o) (R o) (LSTM' i o)

-- | Construct an 'LSTM'
pattern LSTM :: (KnownNat i, KnownNat o) => LSTM i o
pattern LSTM <- Rec { _recLearn = LSTM' }
  where
    LSTM = Rec
      { _recSplit = H.split
      , _recJoin  = (H.#)
      , _recLoop  = id
      , _recLearn = LSTM'
      }

-- | "Base" for 'GRU'.  An 'GRU' is just an 'GRU'' where the second half
-- of the input vector is the previous output.
--
-- This is mostly an implementation detail.  It is recommended that you use
-- 'GRU' and its constructor, instead of directly using this.
data GRU' (i :: Nat) (o :: Nat) = GRU'
  deriving (Typeable)

-- | 'GRU' layer parmateters
data GRUp (i :: Nat) (o :: Nat) =
    GRUp { _gruMemory :: !(FCp (i + o) o)
         , _gruUpdate :: !(FCp (i + o) o)
         , _gruOutput :: !(FCp (i + o) o)
         }
  deriving (Generic, Typeable, Show)

instance NFData (GRUp i o)
instance (KnownNat i, KnownNat o) => Additive (GRUp i o)
instance (KnownNat i, KnownNat o) => Scaling Double (GRUp i o)
instance (KnownNat i, KnownNat o) => Metric Double (GRUp i o)
instance (KnownNat i, KnownNat o, Ref m (GRUp i o) v) => AdditiveInPlace m v (GRUp i o)
instance (KnownNat i, KnownNat o, Ref m (GRUp i o) v) => ScalingInPlace m v Double (GRUp i o)
instance (KnownNat i, KnownNat o) => Bi.Binary (GRUp i o)
instance (KnownNat i, KnownNat o) => Backprop (GRUp i o)

instance (KnownNat i, KnownNat o) => Initialize (GRUp i o) where
    initialize d g = GRUp <$> initialize d g
                          <*> initialize d g
                          <*> initialize d g

instance (KnownNat i, KnownNat o) => Num (GRUp i o) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance (KnownNat i, KnownNat o) => Fractional (GRUp i o) where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational

instance (KnownNat i, KnownNat o) => Floating (GRUp i o) where
    pi    = gPi
    sqrt  = gSqrt
    exp   = gExp
    log   = gLog
    sin   = gSin
    cos   = gCos
    asin  = gAsin
    acos  = gAcos
    atan  = gAtan
    sinh  = gSinh
    cosh  = gCosh
    asinh = gAsinh
    acosh = gAcosh
    atanh = gAtanh

gruMemory :: Lens' (GRUp i o) (FCp (i + o) o)
gruMemory f x = (\y -> x { _gruMemory = y }) <$> f (_gruMemory x)

gruUpdate :: Lens' (GRUp i o) (FCp (i + o) o)
gruUpdate f x = (\y -> x { _gruUpdate = y }) <$> f (_gruUpdate x)

gruOutput :: Lens' (GRUp i o) (FCp (i + o) o)
gruOutput f x = (\y -> x { _gruOutput = y }) <$> f (_gruOutput x)

instance (KnownNat i, KnownNat o, KnownNat io, io ~ (i + o)) => Learn (R io) (R o) (GRU' i o) where
    type LParamMaybe (GRU' i o) = 'Just (GRUp i o)
    type LStateMaybe (GRU' i o) = 'Nothing

    runLearn GRU' (J_ p) = stateless $ \x ->
        let z      = logistic $ runLRp (p ^^. gruMemory) x
            r      = logistic $ runLRp (p ^^. gruUpdate) x
            r'     = 1 # r
            h'     = tanh     $ runLRp (p ^^. gruOutput) (r' * x)
        in  (1 - z) * snd (split @i x) + z * h'

-- | Gated Recurrent Unit
--
-- <http://colah.github.io/posts/2015-08-Understanding-LSTMs/>
--
-- @
-- instance 'Learn' ('R' i) (R o) ('GRU' i o) where
--     type 'LParamMaybe' (GRU i o) = ''Just' ('GRUp' i o)
--     type 'LStateMaybe' (GRU i o) = 'Just (R o)
-- @
type GRU i o = Recurrent (R (i + o)) (R i) (R o) (R o) (GRU' i o)

-- | Construct an 'GRU'
pattern GRU :: (KnownNat i, KnownNat o) => GRU i o
pattern GRU <- Rec { _recLearn = GRU' }
  where
    GRU = Rec
      { _recSplit = H.split
      , _recJoin  = (H.#)
      , _recLoop  = id
      , _recLearn = GRU'
      }


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RNN.Data
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Conversion from and to text data
--
-----------------------------------------------------------------------------

module AI.Network.RNN.Data  where

import qualified Data.Text as T
import qualified Data.Map as DM
import qualified Data.Set as DS
import Data.Maybe
import Data.Tuple
import Data.List
import Control.Monad.Random as R
import Control.Monad
import Numeric.LinearAlgebra.HMatrix  as M
import Data.Foldable as F (toList)

import qualified Data.SDR as SDR
import qualified Data.IntSet as I

import AI.Network.RNN.Types
import AI.Network.RNN.Util
-- import Debug.Trace

-- | Transform text into training data
-- the encoding of each Char is equilateral encoding
textToTrainData :: T.Text -> TrainData ((DM.Map Int Char),Matrix Double)
textToTrainData t =
    let (ids,m) = T.foldl' toIDs ([],DM.empty) t
        sz = DM.size m
        em = equilateralEncoding sz
        is = map (toArr em) $ reverse ids
        charMap = DM.fromList $ map swap $ DM.assocs m
    in
        TrainData (M.fromList (replicate (sz-1) 0): init is) is (sz-1) (charMap,em)
    where
        toIDs :: ([Int],DM.Map Char Int) -> Char -> ([Int],DM.Map Char Int)
        toIDs (ids,m) c =
            let mmyid = DM.lookup c m
            in case mmyid of
                Just myid -> (myid:ids,m)
                Nothing   ->
                    let myid = DM.size m
                    in (myid:ids,DM.insert c myid m)
--        toArr :: Int -> Int -> Vector Double
--        toArr sz idx = M.fromList $ replicate idx 0 ++ [1] ++ replicate (sz-idx-1) 0
        toArr em idx = flatten $ em ? [idx]

-- | Decode the data from the equilateral encoding
dataToText :: ((DM.Map Int Char),Matrix Double) -> [Vector Double] -> T.Text
dataToText (m,em)  = T.pack . F.toList . fmap toC
    where
        toC :: Vector Double -> Char
        -- toC ds = fromJust $ DM.lookup (maxIndex ds) m
        toC ds = fromJust $ DM.lookup (equilateralDecoding em ds) m

-- | Generate random text with the probabilities given by the vector (not used)
randDataToText :: RandomGen g => DM.Map Int Char -> [Vector Double] -> Rand g T.Text
randDataToText m  = liftM T.pack . mapM (toC . M.toList)
    where
        toC :: RandomGen g => [Double] -> Rand g Char
        toC ls = do
            let m1 = normalize $ map (\(ix,d)->(fromJust $ DM.lookup ix m,d)) $ zip [0..] ls
                m2 = map (\(a,b)->(a,toRational b)) m1
                -- s  = sum $ map (exp . snd) m1
                -- m3 = map (\(a,b)-> (a,exp(b)/s)) m1
            R.fromList m2
            --return $ fst $ last $ sortBy (comparing snd) m3

-- | Generate text using equilateral encoding and a given network
generate :: (RandomGen g,RNNEval a)
    => ((DM.Map Int Char),Matrix Double) -- ^ Character map
    -> Int  -- ^ Number of characters to generate
    -> Int -- ^ Size of the vectors for each characters
    -> a -- network
    -> Rand g T.Text
generate m nb sz rnn =generate' m nb sz rnn dataToText

-- | Generate text using sparse encoding and a given network
generateS :: (RandomGen g,RNNEval a)
    => SDR.SDRSet Char -- ^ Character map
    -> Int  -- ^ Number of characters to generate
    -> Int -- ^ Size of the vectors for each characters
    -> a -- network
    -> Rand g T.Text
generateS m nb sz rnn =generate' m nb sz rnn dataToTextS

-- | Generate text given a decoding function
-- current implementation is not random!
generate' :: (RandomGen g,RNNEval a) => m -> Int -> Int -> a ->
    (m -> [Vector Double] -> T.Text) -> Rand g T.Text
generate' m nb sz rnn f =  do
    let (_,_,alls) = foldl' go (rnn,M.fromList $ replicate sz 0,[]) [1..nb]
    --randDataToText m alls
    return $ f m alls
    where
        go :: (RNNEval a) => (a,Vector Double,[Vector Double]) -> Int -> (a,Vector Double,[Vector Double])
        go (rnn1,is1,oss) _ =
            let (rnn2,os) = evalStep rnn1 is1
                os2 = norm os
            in (rnn2,os2,oss ++ [os])
        norm :: Vector Double -> Vector Double
        norm os = let
            mx = maximum $ M.toList os
            --in M.assoc (M.size os) 0 [(mx,1)]
            in M.fromList $ map (\a->if a == mx then 1 else 0) $ M.toList os

-- | Train data using sparse encoding: 2 values in the vector are set to one for each character
-- This was a test to try to reduce the size of the network but is not efficient
textToTrainDataS :: T.Text -> TrainData (SDR.SDRSet Char)
textToTrainDataS t =
    let charSet = T.foldl' (flip DS.insert) DS.empty t
        sz  = DS.size charSet
        markers = 1
        total = sz * markers
        sdr = SDR.build (sz*markers) markers charSet
        is = T.foldr (\c l->enc total sdr c : l) [] t
    in TrainData (M.fromList (replicate total 0): init is) is total sdr
    where
        enc :: Int -> SDR.SDRSet Char -> Char -> Vector Double
        enc total sdr c =
            let is=SDR.encode sdr c
            in M.assoc total 0 $ zip (I.toList is) (repeat 1)

--textToTrainDataS t =
--    let (ids,m) = T.foldl' toIDs ([],DM.empty) t
--        sz = sparseSize $ DM.size m
--        spr = DM.fromList $ zip [0..] $ sparse $ DM.size m
--        is = map (toArr spr sz) $ reverse ids
--        charMap = DM.fromList $ map (\(c,i)->(fromJust $ DM.lookup i spr,c)) $ DM.assocs m
--    in
--        TrainData (M.fromList (replicate sz 0): init is) is sz charMap
--    where
--        toIDs :: ([Int],DM.Map Char Int) -> Char -> ([Int],DM.Map Char Int)
--        toIDs (ids,m) c =
--            let mmyid = DM.lookup c m
--            in case mmyid of
--                Just myid -> (myid:ids,m)
--                Nothing   ->
--                    let myid = DM.size m
--                    in (myid:ids,DM.insert c myid m)
--        toArr :: DM.Map Int (Int,Int) -> Int -> Int -> Vector Double
--        toArr m sz idx = let
--            (i1,i2) = fromJust $ DM.lookup idx m
--            in M.fromList $ replicate i1 0 ++ [1] ++ replicate (i2-i1-1) 0 ++ [1] ++ replicate (sz-i2-1) 0

-- | Generate text from sparse encoding data
dataToTextS :: SDR.SDRSet Char -> [Vector Double] -> T.Text
dataToTextS sdr = T.pack . map b
    where b v =
            let is = I.fromList $ map fst $ filter (\(_,c)->c>=0.5) $ zipWith (\a b->(b,a)) (M.toList v) [0..]
            in SDR.best sdr is
--dataToTextS m  = T.pack . F.toList . fmap toC
--    where
--        toC :: Vector Double -> Char
--        toC ds = let
--            ix1 = maxIndex ds
--            ix2 = maxIndex $ accum ds const [(ix1,0)]
--            tpl = if ix1 < ix2
--                    then (ix1,ix2)
--                    else (ix2,ix1)
--            mc = DM.lookup tpl m
--            in fromMaybe (fromJust $ DM.lookup (0,1) m) mc
--         --fromJust $ DM.lookup (maxIndex ds) m


----------------------
-- 2018.1
-- sule
-- mnist neutral network for AI course homework
----------------------
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}

module Main where

import Data.List(foldl')
import qualified Data.Array as A
import Control.Monad
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy.Char8 as L
import Data.Serialize
import qualified Data.Vector.Storable as V
import Grenade
import Numeric.LinearAlgebra (maxIndex)
import qualified Numeric.LinearAlgebra.Static as S
import System.Random
import System.Directory
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Diagrams(toFile)
import qualified Data.Array.IO as MA

-- |type of input and output of the network
type IOShape = (S ('D2 28 28), S ('D1 10))

-- |type of the neutral network
type HandRC
  = Network
  '[ Convolution 1 12 5 5 1 1, Pooling 2 2 2 2, Relu
     , Convolution 12 24 5 5 1 1, Pooling 2 2 2 2
     , Reshape, Relu
     , FullyConnected 384 80, Logit
     , FullyConnected 80 10, Logit]
  '[ 'D2 28 28, 'D3 24 24 12, 'D3 12 12 12, 'D3 12 12 12
     , 'D3 8 8 24, 'D3 4 4 24, 'D1 384, 'D1 384
     , 'D1 80, 'D1 80, 'D1 10, 'D1 10]

-- |parse the input and label from a csv file
parseData :: FilePath -> IO [IOShape]
parseData path = do
  strs <- tail . L.lines <$!> L.readFile path
  return (map proc strs)
  where
    proc l =
      ( fromJust $! fromStorable $! V.fromList $! tail lbs
      , fromJust $! fromStorable $! V.fromList $! conv $! head lbs)
      where
        lbs = map (read . L.unpack) $! L.split ',' l
    fromJust (Just a) = a
    conv x' =
      let x = floor x'
      in replicate x 0 ++ [1] ++ replicate (9 - x) 0

-- |train the network with a pair of input and label
trainOne :: LearningParameters -> HandRC -> IOShape -> HandRC
trainOne rate !net p@(d, _) = uncurry (train rate net) p

-- |run the network, get the predict
run :: HandRC -> S ('D2 28 28) -> S ('D1 10)
run = runNet

-- |helper function for random sort
randomSortT :: [Int] -> IO [Int]
randomSortT [] = return []
randomSortT !ls = do
  a <- randomRIO (0, length ls - 1)
  (\ls' -> (ls !! a) : ls') <$!> (randomSortT $! (take a ls ++ drop (a + 1) ls))

-- |helper function for random sort
concatBy::[Int]->[[a]]->[a]
concatBy [] _ = []
concatBy (x:xs) as = (as!!x)++concatBy xs as

-- |implement of random sort
randomSort :: [Int]->Int -> IO [Int]
randomSort ls step=
  if  lg<= step
    then randomSortT ls
    else do
      let ts = lg `div` step
      xr<-randomSortT [0..ts-1]
      res''<-mapM (randomSortT . get) [1 .. div (length ls) step]
      let res'=concatBy xr res''
      if lg `mod` step==0
        then return res'
        else do
          t<-randomRIO (0,ts)
          res<-randomSort (drop (ts*step) ls) step
          return $! take t res'++ res ++drop t res'
  where
    get x = drop (step * (x - 1)) $! take (step * x) ls
    lg =length ls

-- |train one generation
trainOnce ::
     Int->LearningParameters -> Int -> [IOShape] -> [Int] -> HandRC -> IO HandRC
trainOnce gx lp time dts ls !net = do
  let lgT = length ls
  print $! lgT
  foldM
    (\n i -> do
        let ir = ls !! i
        print $! "training the " ++ show i ++ " th data " ++ show ir
        when (i `mod` gx==0) $! do
          let trained = take i dts
          let testD = drop lgT dts
          print $! "testing " ++ show i ++ " records trained before"
          let trainCR = validateList n trained
          appendFile "trained" $! Prelude.show (i, correctRate trainCR)++"\n"
          print $! "testing " ++ show (length dts-lgT) ++ " records for validate"
          let trainTest = validateList n testD
          appendFile "test" $! Prelude.show (i, correctRate trainTest)++"\n"
          plotFromFile ("trained", "test")
        return $!
         trainOne (lp {learningRate = learningRate lp * 0.7 ^ time}) n (dts !! ir))
    net
    [0 .. lgT - 1]

-- |train the nerwork with the given init rate, generations, time to calculate accuracy, csv path, and proportion to train
-- example: trainT 0.007 15 1000 "train.csv" 0.8
-- means train the "train.csv" with a init rate of 0.007, keep training for 15 generations, per 1000 records calculate accuracy
-- if claim a pre-trained network in the directory, it will be loaded
trainT :: FilePath->Double->Int ->Int-> FilePath -> Double -> IO ()
trainT preNet initRate time gx path trainS = do
  dts <- parseData path
  let lg = length dts - 1
  let lgT = 1 + floor (trainS * fromIntegral lg)
  bl<-doesFileExist preNet
  net <- if bl
          then print"loaded">>load preNet
          else randomNetwork
  let lp = LearningParameters initRate 0.9 0.0005
  net' <-
    foldM
      (\nt t -> do
         lsT <- randomSort [0 .. lgT-1] 5
         lsV <- randomSort [lgT .. lg] 5
         trainOnce gx lp t dts lsT nt)
      net
      [1 .. time]
  save net' "net"

calcConfuseMatrix::IO ()
calcConfuseMatrix = do
  net<-load "net"
  dts<-parseData "train.csv"
  ar<-MA.freeze=<<confuseMatrix net dts::IO (A.Array (Int,Int) Int)
  writeFile "confuseMatrix" $ renderA ar

renderA::A.Array (Int,Int) Int->String
renderA = foldl' (\str ((_,a),e)->if a==0 then str++"\n"++show e else str++" "++show e) [] . A.assocs

-- |calculate the confuseMatrix using the given network and records
confuseMatrix::HandRC->[IOShape]->IO (MA.IOUArray (Int,Int) Int)
confuseMatrix net ls = do
  ar<-MA.newArray ((0,0),(9,9)) 0
  mapM_ (\((dt,lb),i)->do
    print i
    let i=(getIdx lb, getIdx $run net dt)
    x<-MA.readArray ar i
    MA.writeArray ar i $x+1) $zip ls [1..]
  return ar

-- |k-cross validation
kCross::Int->Int->Double->FilePath->IO ()
kCross gen k initRate path = do
  lds<-parseData path
  let ws = length lds `div` k
  mapM_ (cross lds ws) [0..k-1]
  where cross lds ws t= do
        let w=t*ws
        let trainD = take w lds++drop (w+ws) lds
        let valD = drop w $! take (w+ws) lds
        net<-trainList initRate t gen trainD
        let tr=validateList net trainD
        let vr=validateList net valD
        save net ("k-cross-net-"++show t)
        putStrLn $!"Train "++show (correctRate tr)++" Validation "++show (correctRate vr)
        writeFile ("k-cross-accuracy-"++show t) $!"Train "++show (correctRate tr)++" Validation "++show (correctRate vr)

-- |helper function for k-cross
justTrain :: LearningParameters -> Int -> [IOShape] -> HandRC -> IO HandRC
justTrain lp time dts !net = do
  let lgT = length dts
  foldM
    (\n i -> do
        print $! "training the " ++ show i ++ " th data"
        return $!
         trainOne (lp {learningRate = learningRate lp * 0.9 ^ time}) n (dts !! i))
    net
    [0 .. lgT - 1]

-- |train a list of records, helper function for k-cross
trainList :: Double->Int->Int->[IOShape] -> IO HandRC
trainList initRate r gen dts = do
  let lg = length dts
  print lg
  net <- randomNetwork :: IO HandRC
  let lp = LearningParameters initRate 0.9 0.0005
  foldM
      (\nt t -> do
         print $! "generation "++show t
         justTrain lp t dts nt)
      net
      [1 .. gen]

-- |get the correctRate from a list of bool
correctRate :: [Bool] -> Double
correctRate [] = 0
correctRate !bs =
  fromIntegral (length (filter id bs)) / fromIntegral (length bs)

-- |validate a list of records
validateList :: HandRC -> [IOShape] -> [Bool]
validateList _ [] = []
validateList !net !ls = map (validateOne net) ls

-- |validate one records
validateOne :: HandRC -> IOShape -> Bool
validateOne !net (!dt, !lb) = judge (run net dt) lb

validate :: IOShape -> IOShape -> Bool
validate (_, !lb') (_, !lb) = judge lb' lb

judge :: S ('D1 10) -> S ('D1 10) -> Bool
judge !x !y = getIdx x == getIdx y

getIdx :: S ('D1 10) -> Int
getIdx (S1D !a) = maxIndex (S.extract a)

-- |save a network to file
save :: HandRC -> FilePath -> IO ()
save net path = B.writeFile path $! runPut (put net)

load :: FilePath -> IO HandRC
load path = do
  modelData <- B.readFile path
  either fail return $! runGet (get :: Get HandRC) modelData

plotFromFile::(FilePath,FilePath)->IO ()
plotFromFile (t,v)=do
  ts<-map read.lines<$>readFile t::IO [(Int,Double)]
  vs<-map read.lines<$>readFile t::IO [(Int,Double)]
  plotAccuracy (trans ts) "train accuracy" "train-accuracy.svg"
  plotAccuracy (trans vs) "validate accuracy" "validate-accuracy.svg"
  where trans xs = filter (\(a,b)->b/=0) $map (\((a,b),z)->(a+z,b)) $ zip xs $concat [replicate 34 x| x<-[0,34000..]]

plotAccuracy::[(Int,Double)]->String->FilePath->IO ()
plotAccuracy xs am path = toFile def path $ do
    layout_title .= "Accuracy"
    setColors [opaque green]
    plot (line am [xs])

runT::IO ()
runT = do
  putStrLn "Enter the pre-trained neutral network file(if not exist then init randomly)"
  net<-getLine
  putStrLn "Enter the init learning rate"
  initRate<-read<$>getLine
  putStrLn "Enter the generations to train(how many repetition)"
  time<-read<$>getLine
  putStrLn "Enter the duration of validation"
  gx<-read<$>getLine
  putStrLn "Enter the percentage for training(1-p for validation)"
  trainS<-read<$>getLine
  putStrLn "Enter the path of training data(train.csv)"
  path<-getLine
  trainT net initRate time gx path trainS

runK::IO ()
runK = do
  putStrLn "Enter the k for k-cross validation"
  k<-read<$>getLine
  putStrLn "Enter the init learning rate"
  initRate<-read<$>getLine
  putStrLn "Enter the generations to train(how many repetition)"
  time<-read<$>getLine
  putStrLn "Enter the path of training data(train.csv)"
  path<-getLine
  kCross time k initRate path

main = do
  putStrLn "choose one action"
  putStrLn "0 -- run randomly training and validation"
  putStrLn "1 -- run k cross validation"
  x<-read<$>getLine::IO Int
  if x==0
    then runT
    else if x==1
      then runK
      else putStrLn "wrong option">>main