text:
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}

-- This file is the new version for C-code FFT algorithm with Nyquist allready on position 1.

-- | This module imports the C-functions for a Fast Fourier Transform and encapsulates them in
--   list-mapping functions. (At first it was a DFT algoritm which I changed, but I didn'd change
--   all the names of the functions, so it is called dft...)
--   The functionality of this module is used for filtering audio data, i. e. for building
--   audio filters.

module Sound.Hommage.DFTFilter
 (

-- | Filtering is done with an FFT algorithm written in C (in the file dft.c).
--   This module imports the c-functions for fft ('analyseDFT') and inverse fft ('syntheseDFT'),
--   which operate on Arrays with 1024 Double values.
--   The array-operations are embedded in a mechanism that does buffering, overlapping and windowing
--   and works like a black box that maps one input value to one output value in the IO monad
--   ('mkAnalyse', 'mkSynthese', 'mkFilterBuffered').
--   Around these IO's some list mapping functions are wrapped ('dftanalyse', 'dftsynthese', 'dftfilter',
--   'dftfilterBy') for convenient use in a Haskell programm.
--
--   Mapping the Fourier coefficients during the filter process means filtering the data.
--   The result depends on the way the coefficients are mapped resp. modified while mapping.
--   The function to do this has the type of 'CoeffMap'.
--
--   /Description of the filtering process:/
--
--   Filtering is done by mapping the coefficients with 'CoeffMap'.
--   For this the wave data is split into parts of 1024 values, where every part overlaps with the
--   preceeding and succeeding part by 512 values. Via the fourier-transform (analyse) every part
--   is mapped  to its frequency spectrum, i. e. its fourier coefficients.
--
--   The 'CoeffMap' action is applied to every frequency spectrum (see 'CoeffArr') in sequence.
--   It reads the coefficients from the first 'CoeffArr' where the input is stored and writes it
--   (or some data derived from it) to the second one. This is the point where the \'real\' filtering
--   happens. Via the inverse fourier-transform (synthese) the resulting coefficients are mapped
--   to their real signal values. These signal values are multiplied with a half cosinus curve for
--   fading in and out (windowing). Then they are mixed in the way their source data was split:
--   Every part overlaps for 512 values with the last and the next one.

 -- * User-Level Filter Functions
   dftanalyse
 , dftsynthese
 , dftfilter
 , dftfilterBy
 , dftfilterBy'

 -- * CoeffArr
 , CoeffArr
 , readCoeffArr
-- , readCoeffArr'
 , writeCoeffArr
-- , writeCoeffArr'
 , storeCoeff
 , unstoreCoeff
 -- * CoeffMap
 , CoeffMap
 , mkCoeffMap
 , coeffmap
 -- * IO Wrapper
 , mkAnalyse
 , mkSynthese
 , mkFilterBuffered
 -- * Interface to C-code
 , analyseDFT
 , syntheseDFT
 , kurveDFT
 )
 where

import GHC.Weak
import GHC.IO
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.Storable
import Data.Array.Storable
import Data.Complex
import Data.Array.IO
import Data.IORef

import Sound.Hommage.Misc
---------------------------------------------------------------------------------------------------
-- ([Complex Double] -> [Complex Double]) -> IO CoeffMap

-- | (Fast) Fourier Transformation and Inverse with a buffer-mapping action for filtering.
dftfilterBy :: IO (CoeffMap ()) -> [Double] -> [Double]
dftfilterBy cm = drop 1536 . inList (mkFilterBuffered cm) . (++ replicate 512 0.0)

-- | (Fast) Fourier Transformation and Inverse with a buffer-mapping action for filtering.
dftfilterBy' :: s -> IO (s -> CoeffMap s) -> [Double] -> [Double]
dftfilterBy' s cm = drop 1536 . inList (mkFilterBuffered' s cm) . (++ replicate 512 0.0)

-- | (Fast) Fourier Transformation and Inverse with an additional argument with sublists of
--   factors to weight the coefficients. See 'mkCoeffMap' for a decription of that argument.
dftfilter :: [Double] -> [Double] -> [Double]
dftfilter ds = dftfilterBy (mkCoeffMap (replicate 512 0 ++ ds))

-- | Inverse (Fast) Fourier Transformation
dftsynthese :: [Complex Double] -> [Double]
dftsynthese = inList mkSynthese

-- | (Fast) Fourier Transformation.
dftanalyse :: [Double] -> [(Complex Double)]
dftanalyse = inList mkAnalyse

---------------------------------------------------------------------------------------------------
foreign import ccall "cinit" initfou :: Int

foreign import ccall "canalyse" cAnalyse :: Ptr Double -> Ptr Double -> IO ()

foreign import ccall "csynthese" cSynthese :: Ptr Double -> Ptr Double -> IO ()

foreign import ccall "ckurve" cKurve :: Ptr Double -> IO ()
---------------------------------------------------------------------------------------------------
-- | Maps an array with 1024 real values to its frequency spectrum with 512 complex values.
--   The spectrum is stored in the second array with 1024 Doubles which are the real and imaginary
--   parts of the complex values. See 'CoeffArr' for a detailed frequency spectrum
--   desription.
analyseDFT :: StorableArray Int Double -> StorableArray Int Double -> IO ()
analyseDFT wave coeffs = seq initfou $
 withStorableArray wave $ \ptr_w ->
 withStorableArray coeffs $ \ptr_c ->
 cAnalyse ptr_w ptr_c

-- | Maps a frequency spectrum to its signal value. Both arrays have a length of 1024; the first one
--   contains the 512 complex values of the spectrum. The real signal value will be stored in
--   the second array. See 'CoeffArr' for a detailed frequency spectrum
--   desription.
syntheseDFT :: StorableArray Int Double -> StorableArray Int Double -> IO ()
syntheseDFT coeffs wave = seq initfou $
 withStorableArray wave $ \ptr_w ->
 withStorableArray coeffs $ \ptr_c ->
 cSynthese ptr_c ptr_w

-- | Fades a signal which is stored in an array of length 1024 in and out.
--   The signal is multiplied with a cosinus curve with range 0..pi.
kurveDFT :: StorableArray Int Double -> IO ()
kurveDFT wave = seq initfou $
 withStorableArray wave $ \ptr_w ->
 cKurve ptr_w
---------------------------------------------------------------------------------------------------
--
---------------------------------------------------------------------------------------------------
-- | Represents a frequency spectrum.
--   Range is from 0 to 1023 (N=1024).
--   Index 0 is constant value, index 1 is Nyquest frequency (real part of the N\/2 frequency).
--   Index 2 and 3 are the complex value for the basefrequency, index 4 and 5
--   are the complex value for the double basefrequency, 6 and 7 for 3 * basefrequency and so on:
--
-- > arr [0] = real part of zero frequency (imaginary part is allways 0)
-- > arr [1] = real part of Nyquest frequency (N/2) (imaginary part is allways 0)
-- > arr [2] = real part of base frequency
-- > arr [3] = imag part of base frequency
-- > arr [4] = real part of base frequency * 2
-- > arr [5] = imag part of base frequency * 2
-- > arr [6] = real part of base frequency * 3
-- > arr [7] = imag part of base frequency * 3
-- > ...
-- > arr [N-4] = real part of base frequency * (N/2 - 2)
-- > arr [N-3] = imag part of base frequency * (N/2 - 2)
-- > arr [N-2] = real part of base frequency * (N/2 - 1)
-- > arr [N-1] = imag part of base frequency * (N/2 - 1)
--
-- Pseudocode:
--
-- > c (0)   = (arr [0] :+ 0)
-- > c (512) = (arr [1] :+ 0)
--
--  and for all i=[1..511]:
--
-- > c (i) = (arr [i*2] :+ arr [i*2+1])
--
type CoeffArr = StorableArray Int Double

{-
-- | Reads a Complex value of a CoeffArr as described at 'CoeffArr'.
readCoeffArr :: CoeffArr -> Int -> IO (Complex Double)
readCoeffArr arr n = if n == 0 then readArray arr 0 >>= \x -> readArray arr 1023 >>= \y -> return (x :+ y)
                               else let j  = 2*n
                                        j' = j-1
                                    in readArray arr j' >>= \x -> readArray arr j >>= \y -> return (x :+ y)

-- | Writes a Complex value of a CoeffArr as described at 'CoeffArr'.
writeCoeffArr :: CoeffArr -> Int -> Complex Double -> IO ()
writeCoeffArr arr n c = if n == 0 then writeArray arr 0 (realPart c) >> writeArray arr 1023 (imagPart c)
                                  else let j  = 2*n
                                           j' = j-1
                                  in writeArray arr j' (realPart c) >> writeArray arr j (imagPart c)
-}

-- | Reads a Complex value of a CoeffArr (with Nyquest as imaginary part of coeff 0).
readCoeffArr :: CoeffArr -> Int -> IO (Complex Double)
readCoeffArr arr n =
 let j  = 2*n
     j' = j+1
 in readArray arr j >>= \x -> readArray arr j' >>= \y -> return (x :+ y)

-- | Writes a Complex value of a CoeffArr (with Nyquest as imaginary part of coeff 0).
writeCoeffArr :: CoeffArr -> Int -> Complex Double -> IO ()
writeCoeffArr arr n c =
 let j  = 2*n
     j' = j+1
 in writeArray arr j (realPart c) >> writeArray arr j' (imagPart c)

storeCoeff :: IOArray Int (Complex Double) -> CoeffArr -> IO ()
storeCoeff arr brr = do
 for 0 (<512) (+1) $ \i -> let i1 = i * 2
                               i2 = i1 + 1
                           in do (x :+ y) <- readArray arr i
                                 writeArray brr i1 x
                                 writeArray brr i2 y

unstoreCoeff :: CoeffArr -> IOArray Int (Complex Double) -> IO ()
unstoreCoeff brr arr = do
 for 0 (<512) (+1) $ \i -> let i1 = i * 2
                               i2 = i1 + 1
                           in do x <- readArray brr i1
                                 y <- readArray brr i2
                                 writeArray arr i (x :+ y)
---------------------------------------------------------------------------------------------------
-- | An action that maps the fourier coefficients from the first array to the second.
--   Modifying the data while mapping means filtering. See 'CoeffArr' for a description of the arrays.
type CoeffMap a = CoeffArr -> CoeffArr -> IO a

-- | Constructs a Coeffmap.
--   The sublists contain the 512 real values with which the basefrequency and the 511 complex fourier coefficients
--   are multiplied (the Nyquist frequency will be zero).
--   The first value is the factor for the constant coefficient, the second for the base frequency, the next one
--   for the double base freq and so on.
--   If the sublists have more than 512 elements, these elements are thrown away.
--   If it is shorter, the array will be filled up with zeros.
{-
mkCoeffMap :: [[Double]] -> IO (CoeffMap ())
mkCoeffMap cs = do
 r <- newIORef cs
 b <- newArray (0,511) 0.0
 return $ \a1 a2 -> let loop i (x:xs) = if i >= 512 then return () else
                                        writeArray b i x >> loop (i+1) xs
                        loop i []     = for i (<512) (+1) (\i -> writeArray b i 0.0)
                    in do xs <- readIORef r
                          let (xi,xr) | null xs   = ([], [])
                                      | otherwise = (head xs, tail xs)
                          writeIORef r xr
                          loop 0 xi
                          coeffmap b a1 a2
-}

mkCoeffMap :: [Double] -> IO (CoeffMap ())
mkCoeffMap cs = do
 r <- newIORef cs
 b <- newArray (0,511) 0.0
 return $ \a1 a2 -> let loop i (x:xs) = if i >= 512
                                         then writeIORef r (x:xs)
                                         else writeArray b i x >> loop (i+1) xs
                        loop i []     = for i (<512) (+1) (\i -> writeArray b i 0.0)
                                        >> writeIORef r []
                    in do xs <- readIORef r
                          loop 0 xs
                          coeffmap b a1 a2

-- | 'coeffmap' takes an Array with 512 values and uses them for a weigted map of the fourier
--   coefficients. The Nyquist frequency is muliplied with 0.0.
coeffmap :: StorableArray Int Double -> CoeffMap ()
coeffmap filt coeffs coeffs' = do
 for 0 (<512) (+1) (\i -> let j  = 2*i
                              j' = j + 1
                          in
                          readArray filt i >>= \c ->
                          readArray coeffs j  >>= \c1 ->
                          readArray coeffs j' >>= \c2 ->
                          writeArray coeffs' j  (c * c1) >>
                          writeArray coeffs' j' (c * c2) )
 writeArray coeffs' 1 0.0

{-
coeffmap' :: ([Complex Double] -> [Complex Double]) -> IO (CoeffMap ())
coeffmap' f = do
 let toDest arr 1024 _   = return ()
     toDest arr n (x:xs) = writeArray arr n x >> toDest arr (n+1) xs
     toDest arr n []     = writeArray arr n 0 >> toDest arr (n+1) []
 return $ \arr brr -> getElems arr >>= toDest brr 0 . f
-}
---------------------------------------------------------------------------------------------------
-- | Constructs an action that maps wave-data to coefficient-data.
--   Has a delay of 512, i. e. the first 512 elements are zero and the result has (these)
--   512 elemets more than the input.
mkAnalyse :: IO (Maybe Double -> IO (Maybe (Complex Double)))
mkAnalyse = do
 warr <- newArray (0, 1023) 0.0
 carr <- newArray (0, 1023) 0.0
 oarr <- newArray (0, 511) (0.0 :+ 0.0)
-- SIGNALSTREAM reada closea <- openSignalStream sa
 rpos <- newIORef 0
 rcnt <- newIORef Nothing
 let read ma = checkpos >> readIORef rcnt >>= maybe (more ma) rest
     more ma = maybe irest snext ma
     rest k | k <= 0    = return Nothing
            | otherwise =        do writeIORef rcnt $ Just (k-1)
                                    pos <- readIORef rpos
                                    x <- readArray oarr pos
                                    writeIORef rpos (pos + 1)
                                    return $ Just x
     irest =        do writeIORef rcnt (Just 511)
                       pos <- readIORef rpos
                       fillrest pos
                       x <- readArray oarr pos
                       writeIORef rpos (pos + 1)
                       return $ Just x
     snext a =        do pos <- readIORef rpos
                         x <- readArray oarr pos
                         writeArray warr (pos+512) a
                         writeIORef rpos (pos + 1)
                         return $ Just x
     checkpos = readIORef rpos >>= \p -> if p < 512 then return () else do
                writeIORef rpos 0
                analyseDFT warr carr
                unstoreCoeff carr oarr
                for 0 (<512) (+1) (\i -> readArray warr (i+512) >>= writeArray warr i)
     fillrest k = for (512+k) (<1024) (+1) (\i -> writeArray warr i 0.0)
 return read
---------------------------------------------------------------------------------------------------
-- | Constructs an action that maps coefficient-data to wave-data.
--   Has a delay of 512, i. e. the first 512 elements are zero and the result has (these)
--   512 elemets more than the input.
mkSynthese :: IO (Maybe (Complex Double) -> IO (Maybe Double))
mkSynthese = do
 iarr <- newArray (0, 511) (0.0 :+ 0.0)
 carr <- newArray (0, 1023) 0.0
 warr <- newArray (0, 1023) 0.0
 (oarr :: IOArray Int Double) <- newArray (0, 511) 0.0
-- SIGNALSTREAM reada closea <- openSignalStream sa
 rpos <- newIORef 0
 rcnt <- newIORef Nothing
 let read ma = checkpos >> readIORef rcnt >>= maybe (more ma) rest
     more ma = maybe irest snext ma
     rest k | k <= 0    = return Nothing
            | otherwise =        do writeIORef rcnt $ Just (k-1)
                                    pos <- readIORef rpos
                                    x <- readArray oarr pos
                                    writeIORef rpos (pos + 1)
                                    return $ Just x
     irest =        do writeIORef rcnt (Just 511)
                       pos <- readIORef rpos
                       fillrest pos
                       x <- readArray oarr pos
                       writeIORef rpos (pos + 1)
                       return $ Just x
     snext a =        do pos <- readIORef rpos
                         x <- readArray oarr pos
                         writeArray iarr pos a
                         writeIORef rpos (pos + 1)
                         return $ Just x
     checkpos = readIORef rpos >>= \p -> if p < 512 then return () else do
                writeIORef rpos 0
                storeCoeff iarr carr
                for 0 (<512) (+1) (\i -> readArray warr (i+512) >>= writeArray oarr i)
                syntheseDFT carr warr
                kurveDFT warr
                for 0 (<512) (+1) (\i -> readArray warr i >>= \x ->
                                         readArray oarr i >>= \y ->  writeArray oarr i (x+y))
     fillrest k = for k (<512) (+1) (\i -> writeArray iarr i (0.0 :+ 0.0))
 return read
---------------------------------------------------------------------------------------------------
-- | Constructs an action that maps wave-data to wave-data via a Fast Foutrier Transform and inverse,
--   filtered by the given 'CoeffMap'.
--   Has a delay of 1024.
mkFilterBuffered :: IO (CoeffMap ()) -> IO (Maybe Double -> IO (Maybe Double))
mkFilterBuffered mkf = do
 f <- mkf
 warr <- newArray (0, 1023) 0.0
 carr <- newArray (0, 1023) 0.0
 carr' <- newArray (0, 1023) 0.0
 warr' <- newArray (0, 1023) 0.0
 (oarr :: IOArray Int Double) <- newArray (0, 511) 0.0
-- SIGNALSTREAM reada closea <- openSignalStream sa
 rpos <- newIORef 0
 rcnt <- newIORef Nothing
 let read ma = checkpos >> readIORef rcnt >>= maybe (more ma) rest
     more ma = maybe irest snext ma
     rest k | k <= 0    = return Nothing
            | otherwise =        do writeIORef rcnt $ Just (k-1)
                                    pos <- readIORef rpos
                                    x <- readArray oarr pos
                                    writeIORef rpos (pos + 1)
                                    return $ Just x
     irest =        do writeIORef rcnt (Just 511)
                       pos <- readIORef rpos
                       fillrest pos
                       x <- readArray oarr pos
                       writeIORef rpos (pos + 1)
                       return $ Just x
     snext a =        do pos <- readIORef rpos
                         x <- readArray oarr pos
                         writeArray warr (pos+512) a
                         writeIORef rpos (pos + 1)
                         return $ Just x
     checkpos = readIORef rpos >>= \p -> if p < 512 then return () else do
                writeIORef rpos 0
                analyseDFT warr carr
                for 0 (<512) (+1) (\i -> readArray warr (i+512) >>= writeArray warr i >>
                                         readArray warr' (i+512) >>= writeArray oarr i)
                f carr carr'
                syntheseDFT carr' warr'
                kurveDFT warr'
                for 0 (<512) (+1) (\i -> readArray warr' i >>= \x ->
                                         readArray oarr i >>= \y ->  writeArray oarr i (x+y))

     fillrest k = for (512+k) (<1024) (+1) (\i -> writeArray warr i 0.0)
 return read
---------------------------------------------------------------------------------------------------

-- | Constructs an action that maps wave-data to wave-data via a Fast Foutrier Transform and inverse,
--   filtered by the given 'CoeffMap'.
--   Has a delay of 1024.
mkFilterBuffered' :: s -> IO (s -> CoeffMap s) -> IO (Maybe Double -> IO (Maybe Double))
mkFilterBuffered' st mkf = do
 stref <- newIORef st
 f <- mkf
 warr <- newArray (0, 1023) 0.0
 carr <- newArray (0, 1023) 0.0
 carr' <- newArray (0, 1023) 0.0
 warr' <- newArray (0, 1023) 0.0
 (oarr :: IOArray Int Double) <- newArray (0, 511) 0.0
-- SIGNALSTREAM reada closea <- openSignalStream sa
 rpos <- newIORef 0
 rcnt <- newIORef Nothing
 let read ma = checkpos >> readIORef rcnt >>= maybe (more ma) rest
     more ma = maybe irest snext ma
     rest k | k <= 0    = return Nothing
            | otherwise =        do writeIORef rcnt $ Just (k-1)
                                    pos <- readIORef rpos
                                    x <- readArray oarr pos
                                    writeIORef rpos (pos + 1)
                                    return $ Just x
     irest =        do writeIORef rcnt (Just 511)
                       pos <- readIORef rpos
                       fillrest pos
                       x <- readArray oarr pos
                       writeIORef rpos (pos + 1)
                       return $ Just x
     snext a =        do pos <- readIORef rpos
                         x <- readArray oarr pos
                         writeArray warr (pos+512) a
                         writeIORef rpos (pos + 1)
                         return $ Just x
     checkpos = readIORef rpos >>= \p -> if p < 512 then return () else do
                writeIORef rpos 0
                analyseDFT warr carr
                for 0 (<512) (+1) (\i -> readArray warr (i+512) >>= writeArray warr i >>
                                         readArray warr' (i+512) >>= writeArray oarr i)
                s <- readIORef stref
                s' <- f s carr carr'
                writeIORef stref s'
                syntheseDFT carr' warr'
                kurveDFT warr'
                for 0 (<512) (+1) (\i -> readArray warr' i >>= \x ->
                                         readArray oarr i >>= \y ->  writeArray oarr i (x+y))

     fillrest k = for (512+k) (<1024) (+1) (\i -> writeArray warr i 0.0)
 return read
---------------------------------------------------------------------------------------------------

{-
-- | Reads a Complex value of a CoeffArr
readCoeffArr :: CoeffArr -> Int -> IO (Complex Double)
readCoeffArr arr n = if n == 0 then readArray arr 0 >>= \x -> readArray arr 1023 >>= \y -> return (x :+ y)
                               else let j  = 2*n
                                        j' = j-1
                                    in readArray arr j' >>= \x -> readArray arr j >>= \y -> return (x :+ y)

-- | Writes a Complex value of a CoeffArr
writeCoeffArr :: CoeffArr -> Int -> Complex Double -> IO ()
writeCoeffArr arr n c = if n == 0 then writeArray arr 0 (realPart c) >> writeArray arr 1023 (imagPart c)
                                  else let j  = 2*n
                                           j' = j-1
                                  in writeArray arr j' (realPart c) >> writeArray arr j (imagPart c)

-- | Copies coefficient data from an IOArray with Complex values (512 elements) to a
--   StorableArray with Doubles (1024 elements)
storeCoeff :: IOArray Int (Complex Double) -> CoeffArr -> IO ()
storeCoeff arr brr = do
 for 1 (<512) (+1) $ \i -> let i2 = i * 2
                               i1 = i2 - 1
                           in do (x :+ y) <- readArray arr i
                                 writeArray brr i1 x
                                 writeArray brr i2 y
 (x :+ y) <- readArray arr 0
 writeArray brr 0 x
 writeArray brr 1023 y

-- | Copies coefficient data from a StorableArray with Doubles (1024 elements) to
--   an IOArray with Complex values (512 elements).
unstoreCoeff :: CoeffArr -> IOArray Int (Complex Double) -> IO ()
unstoreCoeff brr arr = do
 for 1 (<512) (+1) $ \i -> let i2 = i * 2
                               i1 = i2 - 1
                           in do x <- readArray brr i1
                                 y <- readArray brr i2
                                 writeArray arr i (x :+ y)
 x <- readArray brr 0
 y <- readArray brr 1023
 writeArray arr 0 (x :+ y)

coeffmap :: StorableArray Int Double -> CoeffMap ()
coeffmap filt coeffs coeffs' = do
 for 1 (<512) (+1) (\i -> let j' = 2*i
                              j  = j' - 1
                          in
                          readArray filt i >>= \c ->
                          readArray coeffs j  >>= \c1 ->
                          readArray coeffs j' >>= \c2 ->
                          writeArray coeffs' j  (c * c1) >>
                          writeArray coeffs' j' (c * c2) )
 readArray filt 0 >>= \c -> readArray coeffs 0 >>= \c' -> writeArray coeffs' 0 (c * c')
 writeArray coeffs' 1023 0.0
-}

max_stars_repo_name: aische/hommage

max_issues_repo_name: aische/hommage

max_forks_repo_name: aische/hommage


text:
{-|
Module      : ODMatrix.ElementaryDecomposition
Description : This module is responsible for the representation of Elementary operations and the process of elementary decomposition of the path between Origin-Destination matrices.

-}
module ODMatrix.ElementaryDecomposition (
    ElementaryMatrix(..)
  , getChildren
  , applyElementary
  , applyPath
  , applicables
  , elementaryValue
  , opposite
  , isNullElementary
  , ODM
  , Bounds
  ) where

  import Numeric.LinearAlgebra
  --import Data.Tree (Tree(..), unfoldTree)

  import ODMatrix (ODM)

  data ElementaryMatrix =
    NullElementary Int |
    Elementary Int (Int, Int) (Int, Int)
    deriving (Show)

  -- | Size of the elementary matrix
  emSize :: ElementaryMatrix -> Int
  emSize (NullElementary s) = s
  emSize (Elementary s _ _) = s

  -- | Bounds of the capacity of the unit (Lower bound, Upper bound).
  type Bounds = (Double, Double)

  -- Size, visited, s
  type ES = (ElementaryMatrix, ODM)

  applyElementaries :: ODM -> [ElementaryMatrix] -> [ODM]
  applyElementaries s es = map (applyElementary s) es

  -- | Apply a elementary operation to the given ODM.
  applyElementary :: ODM -> ElementaryMatrix -> ODM
  applyElementary a e = a + (toMatrix e)

  -- | Apply a sequence of elementary operations to the given ODM.
  applyPath :: ODM -> [ElementaryMatrix] -> ODM
  applyPath = foldl applyElementary

  -- | Given S, such that A + S = B, this function computes the first elementary permutation matrix E such that S' + E = S.
  getChildren :: ODM                  -- ^ Full path permutation matrix
               -> [ElementaryMatrix]  -- ^ All possible elementary matrices
  getChildren odm = es ++ map opposite es
    where al = toAssocList odm
          s = rows odm
          es = [ Elementary s (r1,c1) (r2,c2) |
                  ((r1,c1),_) <- al,
                  ((r2,c2),_) <- al,
                  r1 > r2, c1 < c2 ]

  opposite :: ElementaryMatrix -> ElementaryMatrix
  opposite (Elementary s (r1,c1) (r2,c2)) = Elementary s (r2,c1) (r1,c2)
  opposite (NullElementary s) = NullElementary s

  -- getChildren s = foldl (\acc p -> acc ++ opposites s p pivots) [] pivots
  --   where -- Filtra solo los potenciales pivotes a los que puede agregarse o quitarse un pasajero.
  --         pivots = [ p | p@((i,j),x) <- toAssocList s, i <= j, x /= 0]

  -- reciprocals :: Pivot -> [Pivot] -> [Pivot]
  -- reciprocals ((i,j),x) pivots =
  --   [ p | p@((h,k),y) <- pivots, h < i, j < k, x/y > 0]

  -- opposites :: ODM -> Pivot -> [Pivot] -> [ElementaryMatrix]
  -- opposites s p@(i,j) pivots =
  --   [ Elementary (rows s) (i,j) (h,k)
  --     | ((h,k),_) <- reciprocals p pivots,
  --       atIndex s (h,j) / x < 0,  -- Opposite signs
  --       atIndex s (i,k) / x < 0 ]

  -- | Filter the elementary matrices that cannot be applied to the odm
  applicables :: Bounds             -- ^ Bounds of the values of the target
              -> ODM      -- ^ Target Origin-Destination Matrix
              -> [ElementaryMatrix] -- ^ List of potential elementary operations
              -> [ElementaryMatrix] -- ^ List of applicable elementary operations
  applicables b s = filter (applicable b s)

  -- | Indicates if an elementary operation is applicable for a given origin destination matrix with the given bounds.
  applicable :: Bounds            -- ^ Bounds of the values of the target
             -> ODM               -- ^ Target Origin-Destination Matrix
             -> ElementaryMatrix  -- ^ Elementary operation to apply
             -> Bool
  applicable _ _ (NullElementary _) = True
  applicable (mn,mx) s (Elementary _ (i,j) (h,k)) =
    p (i,k) > mn && p (h,k) < mx &&
    p (i,j) < mx && p (h,j) > mn
    where p = atIndex s

  -- | Compute the value of a elementary matrix in base of a distance functions between cells.
  elementaryValue :: (Int -> Int -> Double) -- ^ Measure function
                  -> ElementaryMatrix       -- ^ E
                  -> Double                 -- ^ Value of E
  elementaryValue _ (NullElementary _) = 0
  elementaryValue m (Elementary _ (r1,_ ) (r2,_)) = s * 2 * m r1 r2
    where s = fromIntegral . signum $ (r1 - r2)

  -- * Convertions

  -- | Transform the internal representation of a Elementary Matrix in a Matrix Double
  toMatrix :: ElementaryMatrix -> ODM
  toMatrix (NullElementary n) = assoc (n,n) 0 []
  toMatrix (Elementary n idx1@(r1,c1) idx2@(r2,c2)) =
    assoc (n,n) 0 [(idx1,1), (idx2,1), ((r1,c2),-1), ((r2,c1),-1)]

  -- | Transform a matrix to an assoc list.
  -- Recovers only the valid, different to 0, elements of the ODM.
  toAssocList :: ODM -> [((Int, Int), Double)]
  toAssocList s = [ c | c@((i,j),v) <- zip idxs elems, i <= j, v /= 0]
    where n = rows s
          idxs = [ (i,j) | i <- [0..(n-1)], j <- [0..(n-1)] ]
          elems = toList . flatten $ s

  -- | Return True if the elementary matrix is a null elementary matrix.
  isNullElementary :: ElementaryMatrix -> Bool
  isNullElementary (NullElementary s) = True
  isNullElementary _ = False

ext: hs

max_stars_repo_name: renecura/odmatrix

max_issues_repo_name: renecura/odmatrix

max_forks_repo_name: renecura/odmatrix


text:
module SemiGr.SemiGrSpec (main,spec) where

import Test.Hspec
import Test.QuickCheck()
import Data.Complex(Complex((:+)))
import SemiGr.Semi(fromList, toList, LinkedList(..))

main :: IO ()
main = hspec spec

spec :: Spec
spec = do

  describe "Basic" $ do

    it "create linked list from list" $ do
      let a = [1, 2, 3, 4]
        in (fromList a) `shouldBe` LinkedList 1 (LinkedList 2 (LinkedList 3 (LinkedList 4 Empt)))

    it "join two linked list with semigroup opr" $ do
      let { a = fromList [1,2,3]
          ; b = fromList [4,5,6]}
          in (toList $ a <> b) `shouldBe` [1,4,2,5,3,6]


text:
{-# LANGUAGE ApplicativeDo                            #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveFunctor                            #-}
{-# LANGUAGE ExistentialQuantification                #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving               #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

-- |
-- Module      : Data.Bin
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Tools for aggregating numeric values into a set of discrete bins
-- according to some binning specification.
--
-- See 'withBinner' for main usage information, and 'Bin' for the main
-- binned data type, and 'binFreq' for a common usage example.
--

module Data.Bin (
  -- * Specifying the binning
    BinSpec(..), linBS, logBS, gaussBS
  -- ** 'BinView'
  , BinView, binView, linView, logView, gaussView
  -- ** Inspecting 'BinSpec'
  , binSpecIntervals, binSpecCenters
  -- * Creating and manipulating bins
  , Bin, Binner, withBinner, fromFin, bins
  -- ** Inspecting bins
  , binFin, binRange, binMin, binMax, binMid
  , binFinExt, binFinComp
  -- ** Showing bins
  , displayBin, displayBinDouble
  -- *** In-depth inspection
  , Pointed(..), pointed, pElem, binIx, fromIx, expandFin, unexpandFin
  -- * Untyped
  , SomeBin(..), sameBinSpec
  -- * Handy use patterns
  , binFreq
  ) where

import           Control.Monad
import           Data.Finite
import           Data.Foldable
import           Data.Profunctor
import           Data.Proxy
import           Data.Reflection
import           Data.Tagged
import           Data.Type.Equality
import           GHC.TypeNats
import           Numeric.SpecFunctions
import           Text.Printf
import           Unsafe.Coerce
import qualified Data.Map              as M
import qualified Data.Vector.Sized     as SV

-- | A bidirectional "view" to transform the data type before binning.
--
-- See 'linView' for a linear binning, and 'logView' for a logarithmic
-- binning.  You can construct your own custom transformer using 'binView'.
--
-- This type is essentially 'Iso' from the /lens/ library, and any 'Iso''
-- from lens can be used here.  However, it is important that all of these
-- represent /monotonic/ isomorphisms.
type BinView a b = forall p. Profunctor p => p b b -> p a a

-- | Construct a 'BinView' based on "to" and "from" functions
--
-- It is important that the "to" and "from" functions be /inverses/ of each
-- other.  Furthermore, both "to" and "from" should be __monotonic__.
binView
    :: (a -> b)       -- ^ "to"
    -> (b -> a)       -- ^ "from"
    -> BinView a b
binView = dimap

-- | Linear binning
linView :: BinView a a
linView = binView id id

-- | Logarithmic binning (smaller bins at lower levels, larger bins at
-- higher levels).
logView :: Floating a => BinView a a
logView = binView log exp

-- | Binning based on a Gaussian Distribution.  Bins "by standard
-- deviation"; there are more bins the closer to the mean you get, and less
-- bins the farther away.
gaussView
    :: RealFrac a
    => a           -- ^ center / mean
    -> a           -- ^ standard deviation
    -> BinView a Double
gaussView μ σ = binView to from
  where
    to   = erf . realToFrac . (/ σ) . subtract μ
    from = (+ μ) . (* σ) . realToFrac . invErf

view :: BinView a b -> a -> b
view v = runForget (v (Forget id))

review :: BinView a b -> b -> a
review v = unTagged . v . Tagged

-- | Specification of binning.
--
-- A @'BinSpec' n a b@ will bin values of type @a@ into @n@ bins, according
-- to a scaling in type @b@.
--
-- Constructor is meant to be used with type application syntax to indicate
-- @n@, like @'BinSpec' @5 0 10 'linView'@
data BinSpec (n :: Nat) a b =
        BS { bsMin  :: a             -- ^ lower bound of values
           , bsMax  :: a             -- ^ upper bound of values
           , bsView :: BinView a b   -- ^ binning view
           }

-- | Convenient constructor for a 'BinSpec' for a linear scaling.
--
-- Meant to be used with type application syntax:
--
-- @
-- 'linBS' @5 0 10
-- @
linBS
    :: forall n a. ()
    => a                            -- ^ Lower bound
    -> a                            -- ^ Upper bound
    -> BinSpec n a a
linBS mn mx = BS mn mx linView

-- | Convenient constructor for a 'BinSpec' for a logarithmic scaling.
--
-- Meant to be used with type application syntax:
--
-- @
-- 'logBS' @5 0 10
-- @
logBS
    :: forall n a. Floating a
    => a                            -- ^ Lower bound
    -> a                            -- ^ Upper bound
    -> BinSpec n a a
logBS mn mx = BS mn mx logView

-- | Convenient constructor for a 'BinSpec' for a gaussian scaling.  Uses
-- the midpoint as the inferred mean.
--
-- Meant to be used with type application syntax:
--
-- @
-- 'gaussBS' @5 3 0 10
-- @
--
-- indicates that you want 5 bins.
gaussBS
    :: forall n a. RealFrac a
    => a                            -- ^ Standard Deviation
    -> a                            -- ^ Lower bound
    -> a                            -- ^ Upper bound
    -> BinSpec n a Double
gaussBS σ mn mx = BS mn mx (gaussView ((mn + mx)/2) σ)

-- | Data type extending a value with an extra "minimum" and "maximum"
-- value.
data Pointed a = Bot
               | PElem !a
               | Top
  deriving (Show, Eq, Ord, Functor)

-- | @since 0.1.3.0
instance (Enum a, Bounded a) => Enum (Pointed a) where
    toEnum i
      | i == 0 = Bot
      | i == (fromEnum (maxBound :: a) + 2) = Top
      | otherwise = PElem (toEnum (i - 1))
    fromEnum = \case
      Bot     -> 0
      PElem x -> fromEnum x + 1
      Top     -> fromEnum (maxBound :: a) + 2

-- | @since 0.1.3.0
instance Bounded (Pointed a) where
    minBound = Bot
    maxBound = Top

-- | Church-style deconstructor for 'Pointed', analogous to 'maybe',
-- 'either', and 'bool'.
--
-- @since 0.1.1.0
pointed
    :: b                -- ^ return if 'Bot'
    -> (a -> b)         -- ^ apply if 'PElem'
    -> b                -- ^ return if 'Top'
    -> Pointed a
    -> b
pointed bot pelem top = \case
    Bot     -> bot
    PElem x -> pelem x
    Top     -> top

-- | Extract the item from a 'Pointed' if it is neither the extra minimum
-- or maximum.
pElem :: Pointed a -> Maybe a
pElem = pointed Nothing Just Nothing

-- | A @'Bin' s n@ is a single bin index out of @n@ partitions of the
-- original data set, according to a 'BinSpec' represented by @s@.
--
-- All 'Bin's with the same @s@ follow the same 'BinSpec', so you can
-- safely use 'binRange' 'withBinner'.
--
-- It has useful 'Eq' and 'Ord' instances.
--
-- Actually has @n + 2@ partitions, since it also distinguishes values
-- that are outside the 'BinSpec' range.
newtype Bin s n = Bin { _binIx :: Pointed (Finite n) }
  deriving (Eq, Ord, Enum, Bounded)

-- | A more specific version of 'binFin' that indicates whether or not the
-- value was too high or too low for the 'BinSpec' range.
binIx :: Bin s n -> Pointed (Finite n)
binIx = _binIx

-- | Extract, potentially, the 'Bin' index.  Will return 'Nothing' if the
-- original value was outside the 'BinSpec' range.
--
-- See 'binIx' for a more specific version, which indicates if the original
-- value was too high or too low.  Also see 'binFinExt', which extends the
-- range of the 'Finite' to embed lower or higher values.
binFin :: Bin s n -> Maybe (Finite n)
binFin = pElem . binIx

-- | Like 'binFin', but return the true "n + 2" slot number of a 'Bin',
-- where 'minBound' is "below minimum" and 'maxBound' is "above maximum"
--
-- @since 0.1.1.0
binFinExt
    :: KnownNat n
    => Bin s n
    -> Finite (1 + n + 1)
binFinExt = expandFin . binIx

-- | Like 'binFin', but squishes or compresses "below minimum" to "above
-- maximum" bins into the 'Finite', counting them in the same bin as the
-- minimum and maximum bin, respectively.
--
-- @since 0.1.1.0
binFinComp
    :: KnownNat n
    => Bin s n
    -> Finite n
binFinComp = pointed minBound id maxBound . binIx

tick
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> b
tick BS{..} = totRange / fromIntegral (natVal (Proxy @n))
  where
    totRange = view bsView bsMax - view bsView bsMin

packPointed
    :: KnownNat n
    => Integer
    -> Pointed (Finite n)
packPointed n
    | n < 0     = Bot
    | otherwise = maybe Top PElem . packFinite $ n

mkBin_
    :: forall n a b. (KnownNat n, RealFrac b)
    => BinSpec n a b
    -> a
    -> Pointed (Finite n)
mkBin_ bs = packPointed
          . floor
          . (/ tick bs)
          . subtract (scaleIn (bsMin bs))
          . scaleIn
  where
    scaleIn = view (bsView bs)

mkBin
    :: forall n a b s. (KnownNat n, RealFrac b, Reifies s (BinSpec n a b))
    => a
    -> Bin s n
mkBin = Bin . mkBin_ (reflect (Proxy @s))

-- | The type of a "binning function", given by 'withBinner'.  See
-- 'withBinner' for information on how to use.
type Binner s n a = a -> Bin s n

-- | With a 'BinSpec', give a "binning function" that you can use to create
-- bins within a continuation.
--
-- @
-- 'withBinner' myBinSpec $ \toBin ->
--     show (toBin 2.8523)
-- @
--
-- Uses a Rank-N continution to ensure that you can only compare 'Bin's
-- constructed from the same 'BinSpec'/binning function.
withBinner
    :: (KnownNat n, RealFrac b)
    => BinSpec n a b
    -> (forall s. Reifies s (BinSpec n a b) => Binner s n a -> r)
    -> r
withBinner bs f = reify bs $ \(_ :: Proxy s) -> f @s mkBin

-- | Generate a vector of the boundaries deriving the bins from
-- a 'BinSpec'.  Can be useful for debugging.
binSpecIntervals
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> SV.Vector (n + 1) a
binSpecIntervals bs = SV.generate $ \i ->
    case strengthen i of
      Just (fromIntegral->i') -> scaleOut $ i' * t + scaleIn (bsMin bs)
      Nothing                 -> bsMax bs
  where
    t        = tick bs
    scaleIn  = view (bsView bs)
    scaleOut = review (bsView bs)

-- | Generate a vector of the scaled centers of each bin derived from
-- a 'BinSpec'.  Can be useful for debugging.
--
-- Note that for non-linear scales, this centering respects the scale.  For
-- example, for logarithmic scales, this is a "log-aware" center: it gives
-- the midpoint in log-space, not linear-space.
--
-- @since 0.1.3.0
binSpecCenters
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> SV.Vector n a
binSpecCenters bs = SV.generate $ \i ->
    scaleOut $ (fromIntegral i + 0.5) * t + scaleIn (bsMin bs)
  where
    t        = tick bs
    scaleIn  = view (bsView bs)
    scaleOut = review (bsView bs)

binRange_
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> Pointed (Finite n)
    -> (Maybe a, Maybe a)
binRange_ bs = \case
    Bot     -> ( Nothing         , Just (SV.head v))
    PElem i -> ( Just (v `SV.index` weaken i)
               , Just (v `SV.index` shift i )
               )
    Top     -> ( Just (SV.last v), Nothing         )
  where
    v        = binSpecIntervals @n bs

-- | Extract the minimum and maximum of the range indicabed by a given
-- 'Bin'.
--
-- A 'Nothing' value indicates that we are outside of the normal range of
-- the 'BinSpec', so is "unbounded" in that direction.
binRange
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => Bin s n
    -> (Maybe a, Maybe a)
binRange = binRange_ (reflect (Proxy @s)) . binIx

-- | Extract the minimum of the range indicabed by a given 'Bin'.
--
-- A 'Nothing' value means that the original value was below the minimum
-- limit of the 'BinSpec', so is "unbounded" in the lower direction.
binMin
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => Bin s n
    -> Maybe a
binMin = fst . binRange

-- | Extract the maximum of the range indicabed by a given 'Bin'.
--
-- A 'Nothing' value means that the original value was above the maximum
-- limit of the 'BinSpec', so is "unbounded" in the upper direction.
binMax
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => Bin s n
    -> Maybe a
binMax = snd . binRange

-- | Extract the center of the range indicabed by a given 'Bin'.
--
-- A 'Nothing' value means that the original value was either below or above the maximum
-- limit of the 'BinSpec'.
--
-- Note that for non-linear scales, this centering respects the scale.  For
-- example, for logarithmic scales, this is a "log-aware" center: it gives
-- the midpoint in log-space, not linear-space.
--
-- @since 0.1.3.0
binMid
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => Bin s n
    -> Maybe a
binMid = binMid_ (reflect (Proxy @s)) . binIx

binMid_
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> Pointed (Finite n)
    -> Maybe a
binMid_ bs = fmap (v `SV.index`) . pElem
  where
    v = binSpecCenters @n bs

-- | Display the interval maintained by a 'Bin'.
displayBin
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => (a -> String)        -- ^ how to display a value
    -> Bin s n
    -> String
displayBin f b = printf "%s .. %s" mn' mx'
  where
    (mn, mx) = binRange b
    mn' = case mn of
            Nothing -> "(-inf"
            Just m  -> "[" ++ f m
    mx' = case mx of
            Nothing -> "+inf)"
            Just m  -> f m ++ ")"

-- | Display the interval maintained by a 'Bin', if the 'Bin' contains
-- a 'Double'.
displayBinDouble
    :: forall n b s. (KnownNat n, Fractional b, Reifies s (BinSpec n Double b))
    => Int                      -- ^ number of decimal places to round
    -> Bin s n
    -> String
displayBinDouble d = displayBin (printf ("%." ++ show d ++ "f"))

instance (KnownNat n, Show a, Fractional b, Reifies s (BinSpec n a b)) => Show (Bin s n) where
    showsPrec d b = showParen (d > 10) $
      showString "Bin " . showString (displayBin @n show b)

-- | Generate a histogram: given a container of @a@s, generate a frequency
-- map of how often values in a given discrete bin occurred.
--
-- @
-- xs :: [Double]
-- xs = [1..100]
--
-- main :: IO ()
-- main = withBinner (logBS @10 5 50) $ \toBin ->
--     mapM_ (\(b, n) -> putStrLn (displayBinDouble 4 b ++ "\t" ++ show n))
--   . M.toList
--   $ binFreq toBin xs
-- @
--
-- @
-- (-inf .. 5.0000)        4
-- [5.0000 .. 6.2946)      2
-- [6.2946 .. 7.9245)      1
-- [7.9245 .. 9.9763)      2
-- [9.9763 .. 12.5594)     3
-- [12.5594 .. 15.8114)    3
-- [15.8114 .. 19.9054)    4
-- [19.9054 .. 25.0594)    6
-- [25.0594 .. 31.5479)    6
-- [31.5479 .. 39.7164)    8
-- [39.7164 .. 50.0000)    10
-- [50.0000 .. +inf)       51
-- @
binFreq
    :: forall n t a s. Foldable t
    => Binner s n a
    -> t a
    -> M.Map (Bin s n) Int
binFreq toBin = M.unionsWith (+) . map go . toList
  where
    go :: a -> M.Map (Bin s n) Int
    go x = M.singleton (toBin x) 1

-- | Construct a 'Bin' if you know the bin number you want to specify, or
-- if the bin is over or under the maximum.
fromIx :: Pointed (Finite n) -> Bin s n
fromIx = Bin

-- | Construct a 'Bin' if you know the bin number you want to specify.  See
-- 'fromIx' if you want to specify bins that are over or under the maximum,
-- as well.
fromFin :: Finite n -> Bin s n
fromFin = fromIx . PElem

-- | An (ordered) list of all 'Bin's.
--
-- @since 0.1.3.0
bins :: KnownNat n => [Bin s n]
bins = Bin <$> binser []
  where
    binser = (Bot :)
           . foldr ((.) . (:) . PElem) id finites
           . (Top :)

-- | "Expand" a 'Pointed' containing a 'Finite' to a wider-ranged 'Finite'.
-- Used for 'binFinExt'
--
-- @since 0.1.2.0
expandFin :: KnownNat n => Pointed (Finite n) -> Finite (1 + n + 1)
expandFin = pointed minBound (weaken . shift) maxBound

-- | The inverse of 'expandFin': "re-pack" a 'Finite' back into
-- a 'Pointed' containing a narrower-ranged 'Finite'.
--
-- @since 0.1.2.0
unexpandFin :: KnownNat n => Finite (1 + n + 1) -> Pointed (Finite n)
unexpandFin x = case unshift x of
  Nothing -> Bot
  Just y  -> case strengthen y of
    Nothing -> Top
    Just z  -> PElem z

-- | A @'SomeBin' a n@ is @'Bin' s n@, except with the 'BinSpec' s hidden.
-- It's useful for returning out of 'withBinner'.
--
-- It has useful 'Eq' and 'Ord' instances.
--
-- To be able to "unify" two 'Bin's inside a 'SomeBin', use 'sameBinSpec'
-- to verify that the two 'SomeBin's were created with the same 'BinSpec'.
data SomeBin a n = forall s b. (Fractional b, Reifies s (BinSpec n a b))
    => SomeBin { getSomeBin :: Bin s n }

deriving instance (KnownNat n, Show a) => Show (SomeBin a n)

-- | Compares if the ranges match.  Note that this is less performant than
-- comparing the original 'Bin's, or extracting and using 'sameBinSpec'.
instance (KnownNat n, Eq a) => Eq (SomeBin a n) where
    SomeBin x == SomeBin y = binRange x == binRange y

-- | Lexicographical ordering -- compares the lower bound, then the upper
-- bounds.  Note that this is less performant than comparing the original
-- 'Bin's, or extracting and using 'sameBinSpec'
instance (KnownNat n, Ord a) => Ord (SomeBin a n) where
    compare (SomeBin x) (SomeBin y) = compare (binRange x) (binRange y)

-- | Verify that the two reified 'BinSpec' types refer to the same one,
-- allowing you to use functions like '==' and 'compare' on 'Bin's that you
-- get out of a 'SomeBin'.
sameBinSpec
    :: forall s t n a b p. (Reifies s (BinSpec n a b), Reifies t (BinSpec n a b), KnownNat n, Eq a, Fractional b)
    => p s
    -> p t
    -> Maybe (s :~: t)
sameBinSpec _ _ = do
    guard $ binSpecIntervals bs1 == binSpecIntervals bs2
    guard $ binSpecCenters bs1 == binSpecCenters bs2
    pure (unsafeCoerce Refl)
  where
    bs1 = reflect (Proxy @s)
    bs2 = reflect (Proxy @t)

max_stars_repo_name: mstksg/bins

max_issues_repo_name: mstksg/bins

max_forks_repo_name: mstksg/bins


text:
-- Copyright 2016 TensorFlow authors.
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}

module TensorFlow.Gradient
    ( GradientCompatible
    , gradients
    ) where

import Control.Monad (forM, zipWithM)
import Control.Monad.State.Strict (State, evalState, gets, modify)
import Data.ByteString (ByteString)
import Data.Complex (Complex)
import Data.Default (def)
import Data.Int (Int32, Int64)
import Data.Foldable (foldlM)
import Data.List (foldl', sortBy)
import Data.Map.Strict (Map)
import Data.Maybe (fromMaybe, maybeToList, mapMaybe)
import Data.Ord (comparing)
import Data.ProtoLens.TextFormat (showMessage)
import Data.Set (Set)
import Data.Text (Text)
import Data.Tuple (swap)
import Lens.Family2 (Lens', view, (&), (^.), (.~), (%~))
import Lens.Family2.State.Strict (uses)
import Lens.Family2.Stock (at, intAt)
import Lens.Family2.Unchecked (lens, iso)
import Prelude hiding (sum)
import Text.Printf (printf)
import qualified Data.Graph.Inductive.Basic as FGL
import qualified Data.Graph.Inductive.Graph as FGL
import qualified Data.Graph.Inductive.PatriciaTree as FGL
import qualified Data.Graph.Inductive.Query.DFS as FGL
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified Data.Text as Text

import qualified TensorFlow.GenOps.Core as CoreOps
import TensorFlow.Build
    ( MonadBuild
    , Build
    , build
    , renderedNodeDefs
    , opDef
    , opAttr
    , opInputs
    )
import TensorFlow.BuildOp
import TensorFlow.Ops
    ( addN
    , broadcastGradientArgs
    , expandDims
    , fill
    , matMul
    , matMul'
    , reducedShape
    , reluGrad
    , reshape
    , scalar
    , shape
    , softmaxCrossEntropyWithLogits
    , sum
    , scalarize
    , vector
    , zerosLike
    )
import TensorFlow.Output
    ( NodeName(..)
    , Output(..)
    , OutputIx(..)
    , outputIndex
    )
import TensorFlow.Tensor
    ( Tensor(..)
    , Value
    , render
    , expr
    , Rendered
    , tensorNodeName
    , renderedOutput
    , renderValue
    , ToTensor(..)
    )
import TensorFlow.Types (Attribute, OneOf, TensorType, attrLens)
import Proto.Tensorflow.Core.Framework.NodeDef
    (NodeDef, attr, input, op, name)

type GradientCompatible a =
    -- TODO(fmayle): MaxPoolGrad doesn't support Double for some reason.
    (Num a, OneOf '[ Float, Complex Float, Complex Double ] a)

-- TODO(fmayle): Support control flow.
-- TODO(fmayle): Support gate_gradients-like option to avoid race conditions.
-- TODO(fmayle): Do we need to consider control inputs? See _PendingCount in
-- tensorflow/python/ops/gradients.py.
-- TODO(fmayle): Maybe store the gradient functions and numOutputs on the OpDef.

-- | Gradient of @y@ w.r.t. each element of @xs@.
gradients :: forall a v1 t m . ( MonadBuild m
                               , Rendered t
                               , ToTensor t
                               , GradientCompatible a
                               )
          => Tensor v1 a  -- ^ The output of the graph.
          -> [t a]        -- ^ Tensors for which gradients are computed.
          -> m [Tensor Value a]
gradients y xs = build $ do
    -- The gradients are computed using "reverse accumulation", similarly to
    -- what is described here:
    -- https://en.wikipedia.org/wiki/Automatic_differentiation#The_chain_rule.2C_forward_and_reverse_accumulation
    --
    -- The code is summarised as follows:
    --
    -- 1. Create an fgl graph of the relevant nodes (ops) and edges (tensors).
    -- 2. Initialize the gradient of y to 1 (∂y/∂y = 1) and the rest of tensor's
    --    gradients to nothing.
    -- 3. Process the nodes in reverse topological order (i.e. each node comes
    --    after all of its outputs so that the output gradients for a node have
    --    been completely calculated before it is processed):
    --      a. Record the gradient for each of the node's output tensors (∂y/∂w
    --         for each output tensor w).
    --      b. Calculate the gradient of y w.r.t. each of the node's input
    --         tensors using the gradients of the node's output tensors.
    --
    --         Written differently, for each output tensor w and input tensor v:
    --           ∂y/∂w = ...            (calculated in previous steps)
    --           ∂w/∂v = ...            (op specific)
    --           ∂y/∂v = ∂y/∂w * ∂w/∂v  (technically, if tensor v is an input
    --                                   to multiple nodes, then this is only
    --                                   part of ∂y/∂v)
    --
    -- 4. Lookup the recorded gradient for each x in xs.

    y' <- renderValue y
    let yName = tensorNodeName y'
    yOne <- render $ fill (shape y') (scalar 1)
    -- TODO(fmayle): Move this into Build.hs and call it unsafeNodeDefFromName?
    nodeDefLookup :: (NodeName -> NodeDef) <- uses renderedNodeDefs $
        (\f x -> fromMaybe (error $ "no NodeDef found for " ++ show x) (f x))
        . flip Map.lookup
    let (gr, nodeMap) = createGraph yName nodeDefLookup
    -- Set gradient of y to one.
    -- TODO: nicer
    let initPending :: Map.Map FGL.Node (PendingGradients a)
            = Map.empty & (at (nodeMap Map.! yName)
                                . nonEmpty
                                . outputIxAt (outputIndex $ renderedOutput y')
                                . nonEmpty
                                .~ [yOne]
                                )
    -- Calculate the gradients of y w.r.t. each node in the graph.
    gradientMap <- graphGrads gr initPending
    -- Lookup the gradients for each x.
    forM xs $ \x ->
        let Output i xName = renderedOutput x
        in maybe (render $ zerosLike $ toTensor x) return $ do
            n <- nodeMap ^. at xName
            gradientMap ^. at n . nonEmpty . outputIxAt i

outputIxAt :: OutputIx -> Lens' (IntMap.IntMap v) (Maybe v)
outputIxAt = intAt . unOutputIx

-- | Incomplete gradients of a node's outputs.
--
-- The lists represent partial sums. The key is an OutputIx sans newtype.
type PendingGradients a = IntMap.IntMap [Tensor Value a]

-- | Gradients of a node's outputs. The key is an OutputIx sans newtype.
-- TODO: precache the rendering?
type Gradients a = IntMap.IntMap (Tensor Value a)

-- | Graph of TensorFlow operations.
type Graph = FGL.Gr NodeDef EdgeLabel

-- | Data associated with an edge.
--
-- Pair of
--   1. Output index of a tensor from the source node.
--   2. Input index that the tensor connects to on the destination node.
type EdgeLabel = (OutputIx, OutputIx)

-- | State used for calculating gradients.
data GradientsState a = GradientsState
                      { _gradientsPending :: !(Map FGL.Node (PendingGradients a))
                      , _gradientsResult  :: !(Map FGL.Node (Gradients a))
                      }

gradientsPending :: Lens' (GradientsState a) (Map FGL.Node (PendingGradients a))
gradientsPending = lens _gradientsPending (\x y -> x { _gradientsPending = y })

gradientsResult :: Lens' (GradientsState a) (Map FGL.Node (Gradients a))
gradientsResult = lens _gradientsResult (\x y -> x { _gradientsResult = y })

-- TODO(fmayle): Use something like Data.List.Safe.
-- | Safe version of (!!).
safeIndex :: [a] -> Int -> Maybe a
_      `safeIndex` n | n < 0 = Nothing
[]     `safeIndex` _         = Nothing
(x:_)  `safeIndex` 0         = Just x
(_:xs) `safeIndex` n         = xs `safeIndex` (n-1)

-- Copy of http://hackage.haskell.org/package/lens-3.9.0.2/docs/Control-Lens-Iso.html#v%3anon
anon :: a -> (a -> Bool) -> Lens' (Maybe a) a
anon a p = iso (fromMaybe a) go where
  go b | p b       = Nothing
       | otherwise = Just b

non :: Eq a => a -> Lens' (Maybe a) a
non a = anon a (a==)

-- | Lens that defaults Nothing to mempty.
nonEmpty :: (Monoid (t v), Foldable t) => Lens' (Maybe (t v)) (t v)
nonEmpty = anon mempty null

-- TODO: strictness (e.g., foldlM')

-- | Calculate the gradients for every node in a graph.
graphGrads :: forall a. GradientCompatible a
           => Graph
           -> Map FGL.Node (PendingGradients a)
           -- ^ Initial gradients (usually just 1 for the node of interest).
           -> Build (Map FGL.Node (Gradients a))
graphGrads gr initPending = view gradientsResult <$> foldlM go initState nodeOrder
  where
    initState = GradientsState initPending Map.empty
    -- Reverse topological sort.
    -- TODO(fmayle): Filter out nodes that are not successors of any x in xs to
    -- avoid calculating gradients that won't be used.
    nodeOrder = FGL.topsort $ FGL.grev gr
    go :: GradientsState a -> Int -> Build (GradientsState a)
    go state node = do
        -- Aggregate the accumulated gradients for this node.
        outputGrads <-
                sumPendingGradient (state ^. gradientsPending . at node . nonEmpty)
        if null outputGrads
           then pure state
           else do
              let ctx = FGL.context gr node
              inputGrads <- calculateInputGrads ctx outputGrads gr
              -- Calculate the gradients for each of the node's inputs.
              let nextState = state & gradientsResult %~ Map.insert node outputGrads
              pure $ updatePendingGradients ctx inputGrads nextState

-- | Reduce accumulated gradients for each output to one Tensor.
sumPendingGradient :: GradientCompatible a
                   => PendingGradients a -> Build (Gradients a)
sumPendingGradient = sequence . IntMap.mapMaybe f
  where
    f [] = Nothing
    f [x] = Just (pure x)
    f xs = Just (render $ addN xs)

-- | Calculate the gradients of a node's input tensors.
--
-- This is mostly just a wrapper around opGrad.
calculateInputGrads :: forall a. GradientCompatible a
                    => FGL.Context NodeDef EdgeLabel
                    -> Gradients a  -- ^ Output gradients of the node.
                    -> Graph
                    -> Build [Maybe (Tensor Value a)]
calculateInputGrads (inputEdges, _, nodeDef, _) outputGrads gr = do
    fullOutGrads <- fullOutputGrads (numOutputs nodeDef) (nodeDefName nodeDef)
                        outputGrads
    traverse (traverse render) $ opGrad (nodeDef ^. op) nodeDef inputTensors fullOutGrads
  where
    -- Create a tensor from an edge (technically an Output, but it seems less
    -- confusing to refer to it as a tensor here).
    edgeToTensor :: (EdgeLabel, FGL.Node) -> Output
    edgeToTensor ((i, _), n) =
        case FGL.lab gr n of
            Just edgeNodeDef -> Output i (NodeName $ edgeNodeDef ^. name)
            Nothing -> error $ "calculateInputGrads: missing input node for "
                               ++ Text.unpack (nodeDef ^. name)
    -- Input tensors, sorted by input index.
    inputTensors = map edgeToTensor $ sortBy (comparing (snd . fst)) inputEdges

-- | Convert a Map of gradients to a list, with zeros for missing outputs.
fullOutputGrads :: (TensorType a, Num a)
                => OutputIx  -- ^ Number of outputs.
                -> NodeName
                -> Gradients a
                -> Build [Tensor Value a]
fullOutputGrads n o gs =
    mapM (\i -> maybe (render $ zero i) return (gs ^. outputIxAt i)) [0..n-1]
  where
    -- A tensor of zeros with the same shape as the i'th output.
    zero i = zerosLike $ toT (Output i o)

-- | Update the pending gradients of a node's inputs.
updatePendingGradients :: forall a. (TensorType a, Num a)
                       => FGL.Context NodeDef EdgeLabel
                       -> [Maybe (Tensor Value a)]
                       -- ^ Gradient of each input tensor.
                       -> GradientsState a
                       -> GradientsState a
updatePendingGradients (inputEdges, _, nodeDef, _) inputGrads initState =
    foldl' go initState inputEdges
  where
    go :: GradientsState a -> (EdgeLabel, FGL.Node) -> GradientsState a
    go state ((outIndex, OutputIx inIndex), node) =
        case maybeGradient of
            Nothing -> state
            Just g ->
                -- Add to the list of pending gradients for this tensor.
                state & gradientsPending
                      . at node
                      . nonEmpty
                      . outputIxAt outIndex
                      . nonEmpty
                      %~ (g:)
      where
        badSizeErr = error $ printf "updatePendingGradients: bad input index \
                                    \%d for inputGrads of length %d in %s"
                                    inIndex (length inputGrads)
                                    (show (nodeDef ^. name))
        maybeGradient = fromMaybe badSizeErr (safeIndex inputGrads inIndex)

-- | Create a graph that includes a node and its transitive dependencies.
createGraph :: NodeName -> (NodeName -> NodeDef)
            -> (Graph, Map NodeName FGL.Node)
createGraph nodeName nodeDefLookup = (FGL.nmap nodeDefLookup graph, nodeMap)
  where
    -- Parse a tensor name.
    parseTensorName :: Text -> Maybe (NodeName, OutputIx)
    parseTensorName n
        | Text.null n        = error "parseTensorName: empty name"
        | Text.head n == '^' = Nothing  -- Control edge
        | otherwise          =
            let (nm, indexStr) = Text.breakOn ":" n
                index | Text.null indexStr = 0
                      | otherwise = read $ Text.unpack $ Text.tail indexStr
            in Just (NodeName nm, OutputIx index)

    -- Build a map from node name to outward edges.
    --
    -- The state is the set of visited nodes.
    collect :: Maybe (NodeName, OutputIx, OutputIx)
            -> NodeName
            -> State (Set NodeName)
                     (Map NodeName [(NodeName, OutputIx, OutputIx)])
    collect outgoingEdge nm = do
        let nextLookup = Map.singleton nm (maybeToList outgoingEdge)
        seen <- gets (Set.member nm)
        modify (Set.insert nm)
        if seen
            then pure nextLookup
            else do
                let inputs = nodeDefLookup nm ^. input
                    recurse inIndex (parentName, outIndex) =
                        collect (Just (nm, outIndex, inIndex)) parentName
                subEdgeLookups <-
                    zipWithM recurse [0..] $ mapMaybe parseTensorName inputs
                pure $ Map.unionsWith (++) (nextLookup:subEdgeLookups)

    edgeLookup = evalState (collect Nothing nodeName) Set.empty
    -- Associate an ID with each node name.
    nodeMap = Map.fromList $ zip (Map.keys edgeLookup) [0..]
    -- Create the graph.
    graph = FGL.mkGraph (swap <$> Map.toList nodeMap)
                        [ (nodeMap Map.! n, nodeMap Map.! m, (i, j))
                        | (n, edges) <- Map.toList edgeLookup
                        , (m, i, j) <- edges
                        ]

-- | Function to compute the gradient of y w.r.t. each input.
--
-- Let y be an arbitrary tensor
-- and [w_0, ..., w_n] be the output tensors of a node
-- and [v_0, ..., v_n] be the input tensors of the same node.
--
-- Given [∂y/∂w_0, ..., ∂y/∂w_n] and [v_0, ..., v_n], a GradientFunc computes
-- [∂y/∂v_0, ..., ∂y/∂v_n] for a particular op type.
--
-- A Nothing gradient is equivalent to zero (but allows for short circuiting
-- computation when all the gradients for something are Nothing).
type GradientFunc a = NodeDef
                    -> [Output]
                    -- ^ Input tensors.
                    -> [Tensor Value a]
                    -- ^ Gradient of y w.r.t. each output tensor.
                    -> [Maybe (Tensor Build a)]
                    -- ^ Gradient of y w.r.t. each input tensor.

-- TODO(fmayle): Assert the type is correct.
-- | Create a Tensor from an Output.
toT :: Output -> Tensor Build a
toT = Tensor . pure

-- | Wrapper around `TensorFlow.GenOps.Core.slice` that builds vectors from scalars for
-- simple slicing operations.
flatSlice :: forall v1 t . TensorType t
         => Tensor v1 t    -- ^ __input__
         -> Int32          -- ^ __begin__: specifies the offset into the first dimension of
                           -- 'input' to slice from.
         -> Int32          -- ^ __size__: specifies the number of elements of the first dimension
                           -- of 'input' to slice. If size is -1, all remaining elements in the dimension
                           -- are included in the slice (i.e. this is equivalent to setting
                           -- size = input.dim_size(0) - begin).
         -> Tensor Build t -- ^ __output__
flatSlice t begin size = CoreOps.slice t (vector [begin]) (vector [size])

nodeDefName :: NodeDef -> NodeName
nodeDefName = NodeName . view name

-- | Gradient helper for binary component wise operations
-- See https://github.com/tensorflow/tensorflow/blob/e9de087fa7f59c39bbe12ac2c83c5547c83f746c/tensorflow/core/ops/math_grad.cc#L329
gradForBinaryCwise :: ( OneOf '[ Int32, Int64, Float, Double, Complex Float, Complex Double ] t
                      )
                   => (Tensor v1 t, Tensor v1 t)
                   -> (Tensor v1 t, Tensor v1 t)
                   -> [ Maybe (Tensor Build t) ]
gradForBinaryCwise (x, gx) (y, gy) =
    [ Just dx
    , Just dy ]
  where
    dx = reshape (sum gx rx) sx
    dy = reshape (sum gy ry) sy
    sx = shape x
    sy = shape y
    (rx, ry) = broadcastGradientArgs sx sy

-- | The gradient function for an op type.
--
-- These implementations should match their python counterparts in:
-- third_party/tensorflow/python/ops/*_grad.py
opGrad :: forall a . GradientCompatible a => Text -> GradientFunc a

opGrad "Abs" _ [toT -> x] [dz] = [Just $ expr dz * signum x]
opGrad "Neg" _ [_] [dz] = [Just $ negate $ expr dz]
opGrad "Relu" _ [toT -> x] [dz] = [Just $ reluGrad dz x]
opGrad "ReluGrad" _ [_, toT -> x ] [dz] = [Just $ reluGrad dz x, Just $ CoreOps.zerosLike x]

opGrad "Concat" _ _ix [dy]
    -- Concat concatenates input tensors
    --   x1 of shape s1 = [k1, ..., ki_1, ..., kn]
    --   x2 of shape s2 = [k1, ..., ki_2, ..., kn]
    --    .           .     .          .        .
    --    .           .     .          .        .
    --    .           .     .          .        .
    --   xm of shape sm = [k1, ..., ki_m, ..., kn]
    --  along dimension i to an output tensor
    --   y  of shape sy = [k1, ..., k, ..., kn]
    --  where k = sum ki = sum [ki_1,...,ki_m]
    --
    --  The incoming gradient dy from backpropagation is
    --   simply forwarded split across input tensors yielding dx.
    --   Forwarded gradients have shapes s = [s1, ..., sm].
    | m == 1    = Nothing : [Just $ expr dy]
    | otherwise = Nothing : map Just (dx `reshapeZip` s)
  where
    reshapeZip = zipWith reshape
    dx = CoreOps.splitV (fromIntegral m) dy ki _i
    s  :: [Tensor Build Int32]
    s  = map shape x
    x  :: [Tensor Build a]
    x  = map toT $ tail _ix
    -- i: concat dimension. Adjusted modulo n to handle negative indices.
    _i = toT (head _ix) `CoreOps.floorMod` n
    i  = reshape _i $ vector [1 :: Int32]
    -- sizes along concatenated dimension
    ki :: Tensor Build Int32
    ki = CoreOps.concat 0 $ map (\t -> CoreOps.slice t i $ vector [1 :: Int32]) s
    m  = length x
    n  = CoreOps.rank (head x)

opGrad "Square" _ [toT -> x] [dz] =
    -- TODO(fmayle): Handle complex numbers.
    -- TODO(fmayle): The python code makes dz a control dependency of the 2*x
    -- (for performance reasons?). Will need to put these functions in the Build
    -- monad to replicate that.
    [Just $ dz `CoreOps.mul` (2 * x)]

opGrad "Gather" _ [toT -> x, toT -> indices] [dz] =
    -- TODO(fmayle): The python version uses a better performance implementation
    -- when the shape is known without having to run the graph.
    -- TODO(fmayle): We shouldn't convert the result to a dense tensor. Sparse
    -- tensor support will require some thinking.
    [ Just $ CoreOps.unsortedSegmentSum values indices' numRows
    , Nothing
    ]
  where
    -- TODO(gnezdo): Use colocateWith but it requires Build monad.
    denseShape = shape (x :: Tensor Build a)
    numRows = scalarize $ flatSlice denseShape 0 1
    valuesShape = CoreOps.concat 0 [ allDimensions
                                   , flatSlice denseShape 1 (-1)
                                   ]
    values = reshape dz valuesShape
    -- TODO(fmayle): This could be either Int32 or Int64.
    indices' = reshape indices allDimensions :: Tensor Build Int32

opGrad "Max" _ [toT -> x, toT -> indices] [dz] =
    [Just $ indicators `CoreOps.div` numSelected * dz', Nothing]
  where
    sx = shape (x :: Tensor Build a)
    outputShapeKeptDims = reducedShape sx (indices :: Tensor Build Int32)
    y = CoreOps.max x indices
    y' = reshape y outputShapeKeptDims
    dz' = reshape dz outputShapeKeptDims
    indicators = CoreOps.cast $ CoreOps.equal y' x
    numSelected = reshape (sum indicators indices) outputShapeKeptDims

-- Min and Max have identical gradient implementations.
opGrad "Min" u v w = opGrad "Max" u v w

-- Element wise maximum gradient
-- See https://github.com/tensorflow/tensorflow/blob/e9de087fa7f59c39bbe12ac2c83c5547c83f746c/tensorflow/core/ops/math_grad.cc#L473
opGrad "Maximum" _ [toT -> x, toT -> y] [dz] =
    gradForBinaryCwise (x, gx) (y, gy)
  where
    xmask = CoreOps.greaterEqual x y
    gx = CoreOps.select xmask dz (CoreOps.zerosLike dz)
    gy = CoreOps.select (CoreOps.logicalNot xmask) dz (CoreOps.zerosLike dz)

opGrad "Sum" _ [toT -> x, toT -> indices] [dz] =
    [ Just $ CoreOps.tile grad tileScaling, Nothing ]
  where
    -- TODO(gnezdo): Implement the fast-path from math_grad._SumGrad.
    sx = shape (x :: Tensor Build a)
    outputShapeKeptDims = reducedShape sx (indices :: Tensor Build Int32)
    tileScaling = safeShapeDiv sx outputShapeKeptDims
    grad = reshape dz outputShapeKeptDims

opGrad "Mean" u v@[toT -> x, _] w =
    [Just $ dz `CoreOps.div` CoreOps.cast factor, Nothing]
  where
    [Just dz, Nothing] = opGrad "Sum" u v w
    inputShape = shape (x :: Tensor Build a)
    outputShape = shape (dz :: Tensor Build a)
    -- TODO(fmayle): Add fast path when shape is known.
    inputSize = CoreOps.prod inputShape $ rangeOfRank inputShape
    outputSize = CoreOps.prod outputShape $ rangeOfRank outputShape
    factor = safeShapeDiv inputSize outputSize

opGrad "Add" _ [toT -> x, toT -> y] [dz] =
    [ Just $ reshape (sum dz rx) sx
    , Just $ reshape (sum dz ry) sy ]
  where
    sx = shape (x :: Tensor Build a)
    sy = shape (y :: Tensor Build a)
    (rx, ry) = broadcastGradientArgs sx sy

-- Copies the gradients to all inputs
-- Not broadcasting
opGrad "AddN" _ inputs [dz] =
    map ((const . Just . expr) dz) inputs

opGrad "Sub" u v w =
    [Just x, Just (-y)]
  where
    [Just x, Just y] = opGrad "Add" u v w

opGrad "SoftmaxCrossEntropyWithLogits" _ [toT -> x, toT -> y] [dz, _] =
    [ Just $ expandDims dz (-1) * snd (softmaxCrossEntropyWithLogits x y)
    , Nothing ]

opGrad "Mul" _ [toT -> x, toT -> y] [dz] =
    -- TODO(fmayle): Handle complex numbers.
    [ Just $ reshape (sum (dz `CoreOps.mul` y) rx) sx
    , Just $ reshape (sum (x `CoreOps.mul` dz) ry) sy ]
  where
    sx = shape (x :: Tensor Build a)
    sy = shape (y :: Tensor Build a)
    (rx, ry) = broadcastGradientArgs sx sy

opGrad "Div" _ [toT -> x, toT -> y] [dz] =
    -- TODO(fmayle): Handle complex numbers.
    -- TODO(gnezdo): Provide Fractional instance and use '/' instead of div.
    [ Just $ reshape (sum (dz `CoreOps.div` y) rx) sx
    , Just $ reshape (sum (dz `CoreOps.mul` (negate x `CoreOps.div` (y * y)))
                         ry)
                sy
    ]
  where
    sx = shape (x :: Tensor Build a)
    sy = shape (y :: Tensor Build a)
    (rx, ry) = broadcastGradientArgs sx sy

opGrad "MatMul" nodeDef [toT -> x, toT -> y] [dz] =
    let transposeA = lookupAttr nodeDef "transpose_a"
        transposeB = lookupAttr nodeDef "transpose_b"
        transAttrs a b =
            (opAttr "transpose_a" .~ a) . (opAttr "transpose_b" .~ b)
    in case (transposeA, transposeB) of
       (False, False) ->
           [ Just $ matMul' (transAttrs False True) dz y
           , Just $ matMul' (transAttrs True False) x dz]
       (False, True) ->
           [ Just $ matMul dz y
           , Just $ matMul' (transAttrs True False) dz x]
       (True, False) ->
           [ Just $ matMul' (transAttrs False True) y dz
           , Just $ matMul x dz]
       (True, True) ->
           [ Just $ matMul' (transAttrs True True) y dz
           , Just $ matMul' (transAttrs True True) dz x]

opGrad "Transpose" _ [_, toT -> p] [dz] =
    [ Just $ CoreOps.transpose dz
            (CoreOps.invertPermutation p :: Tensor Build Int32)
    , Nothing
    ]

opGrad "Conv2D" nodeDef [toT -> x, toT -> y] [dz] =
    [ Just $ CoreOps.conv2DBackpropInput'
                ((opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "use_cudnn_on_gpu" .~ useCudnnOnGpu)
                    . (opAttr "data_format" .~ dataFormat))
                (shape x) y dz
    , Just $ CoreOps.conv2DBackpropFilter'
                ((opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "use_cudnn_on_gpu" .~ useCudnnOnGpu)
                    . (opAttr "data_format" .~ dataFormat))
                x (shape y) dz
    ]
  where
    strides = lookupAttr nodeDef "strides" :: [Int64]
    padding = lookupAttr nodeDef "padding" :: ByteString
    useCudnnOnGpu = lookupAttr nodeDef "use_cudnn_on_gpu" :: Bool
    dataFormat = lookupAttr nodeDef "data_format" :: ByteString

opGrad "Conv2DBackpropInput" nodeDef [_, toT -> x, toT -> y] [dz] =
    [ Nothing
    , Just $ CoreOps.conv2DBackpropFilter'
                ((opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "use_cudnn_on_gpu" .~ useCudnnOnGpu)
                    . (opAttr "data_format" .~ dataFormat))
                dz (shape x) y
    , Just $ CoreOps.conv2D'
                ((opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "use_cudnn_on_gpu" .~ useCudnnOnGpu)
                    . (opAttr "data_format" .~ dataFormat))
                dz x
    ]
  where
    strides = lookupAttr nodeDef "strides" :: [Int64]
    padding = lookupAttr nodeDef "padding" :: ByteString
    useCudnnOnGpu = lookupAttr nodeDef "use_cudnn_on_gpu" :: Bool
    dataFormat = lookupAttr nodeDef "data_format" :: ByteString

opGrad "MaxPool" nodeDef [toT -> x] [dz] =
    [ Just $ CoreOps.maxPoolGrad'
                ((opAttr "ksize" .~ ksize)
                    . (opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "data_format" .~ dataFormat))
                x output dz
    ]
  where
    output :: Tensor Build a
    output = toT $ Output 0 (nodeDefName nodeDef)
    ksize = lookupAttr nodeDef "ksize" :: [Int64]
    strides = lookupAttr nodeDef "strides" :: [Int64]
    padding = lookupAttr nodeDef "padding" :: ByteString
    dataFormat = lookupAttr nodeDef "data_format" :: ByteString

opGrad "Reshape" _ [toT -> x, _] [dz] =
    [Just $ reshape dz $ shape (x :: Tensor Build a), Nothing]

opGrad "OneHot" _ _ _ = [Nothing, Nothing, Nothing, Nothing]
opGrad "TruncatedNormal" _ _ _ = [Nothing]

opGrad "RefIdentity" _ _ [dz] = [Just $ expr dz]
opGrad "Cast" nodeDef _ [dz] = [Just reverseCast]
  where
    -- TODO(gnezdo): too permissive, python only allows float types as src_type.
    reverseCast =
        pureOp [] $ pure (opDef "Cast"
                 & opAttr "DstT" .~ (lookupAttr nodeDef "SrcT" :: ByteString)
                 & opAttr "SrcT" .~ (lookupAttr nodeDef "DstT" :: ByteString)
                 & opInputs .~ [renderedOutput dz])

opGrad "DynamicStitch" nodeDef inputs [dz] =
    replicate halfLen Nothing ++ valuesGrads
  where
    halfLen =
        let len = length inputs
            half = len `div` 2
        in if 2 * half == len
           then half
           else error ("Uneven input size " ++ show (len, showMessage nodeDef))
    valuesGrads = [ Just $ CoreOps.gather dz (toT idx :: Tensor Build Int32)
                  | idx <- take halfLen inputs
                  ]

opGrad "DynamicPartition" nodeDef [toT -> xs, toT -> indices] dz =
    [ Just reconstructed, Nothing ]
  where
    reconstructed = CoreOps.reshape stitched
                    (CoreOps.shape (xs :: Tensor Build a) :: Tensor Build Int32)
    stitched = CoreOps.dynamicStitch partitionedIndices dz
    partitionedIndices = CoreOps.dynamicPartition np originalIndices indices
    np = lookupAttr nodeDef "num_partitions" :: Int64
    originalIndices =
        CoreOps.reshape (CoreOps.range 0 (CoreOps.size indices) 1) prefixShape
    prefixShape = shapeInt32 indices
    shapeInt32 t = CoreOps.shape t :: Tensor Build Int32

opGrad "Select" _ [toT -> c, toT -> x, _] [dz] =
    [ Nothing
    , Just $ CoreOps.select c dz zeros
    , Just $ CoreOps.select c zeros dz
    ]
  where zeros = CoreOps.zerosLike x

-- TODO(gnezdo): Unlike Python, no control dependency on dz.
opGrad "Log" _ [toT -> x] [dz] = [ Just $ dz `CoreOps.mul` CoreOps.inv x ]
-- TODO(gnezdo): Reuse the output instead of doing another exp,
-- though, it is probably CSE'd away anyway.
opGrad "Exp" _ [toT -> x] [dz] = [ Just $ dz `CoreOps.mul` CoreOps.exp x ]
opGrad "SparseSegmentSum" _ [toT -> x, toT -> y, toT -> t] [dz] =
    [ Just $ CoreOps.unsortedSegmentSum
             (CoreOps.gather dz (t :: Tensor Build Int32))
             (y :: Tensor Build Int32) inputRows
    , Nothing
    , Nothing
    ]
  where inputRows = flatSlice (shape (x :: Tensor Build a)) 0 1

opGrad "LabelClasses" _ _ _ = [Nothing, Nothing]
opGrad "LabelWeights" _ _ _ = [Nothing]
opGrad "Size" _ _ _ = [Nothing]

-- TODO (jcberentsen): Python implementation uses set_shape for
-- static shape inference, which is unsupported.
-- TODO: implement support for static shape inference
opGrad "Tile" _ [toT -> x, toT -> multiples] [dz] =
    [Just inputGrad, Nothing]
  where
    inputGrad = sum reshapedDz axes
    inputShape = shape (x :: Tensor Build a)
    packed = CoreOps.pack [multiples, inputShape]
    perm = vector [1, 0 :: Int32]
    splitShape = CoreOps.reshape (CoreOps.transpose packed perm) allDimensions
    axes = CoreOps.range 0 (CoreOps.size splitShape) (2 :: Tensor Build Int32)
    reshapedDz = CoreOps.reshape dz splitShape

opGrad "ZerosLike" _ _ _ = [Nothing]
opGrad "Fill" _ _ [dz] = [Nothing, Just $ sum dz rx]
  where
    rx = rangeOfRank dz

-- Treat read ops as an identity function on the variable. This allows us to
-- take gradients w.r.t. to the variable handle instead of the result of a read
-- op. If a variable is read multiple times, the gradients will propagate back
-- through each read.
opGrad "ReadVariableOp" _ _ [dz] = [Just $ expr dz]

-- TODO(fmayle): These can go away if we properly prune the graph.
opGrad "Const" _ _ _ = [Nothing, Nothing]
opGrad "Placeholder" _ _ _ = []
opGrad "VarHandleOp" _ _ _ = []
opGrad "Variable" _ _ _ = []

opGrad n nodeDef ins grads =
    error $ "no gradient implemented for " ++
            show (n, length ins, length grads, showMessage nodeDef, ins)

-- | The number of outputs for an op type.
numOutputs :: NodeDef -> OutputIx
numOutputs o =
    case o ^. op of
        "Abs" -> 1
        "Add" -> 1
        "AddN" -> 1
        "Cast" -> 1
        "Const" -> 1
        "Concat" -> 1
        "Conv2D" -> 1
        "Conv2DBackpropInput" -> 1
        "Div" -> 1
        "DynamicStitch" -> 1
        "DynamicPartition" ->
            fromIntegral (lookupAttr o "num_partitions" :: Int64)
        "Exp" -> 1
        "Gather" -> 1
        "LabelClasses" -> 1
        "LabelWeights" -> 1
        "Log" -> 1
        "MatMul" -> 1
        "Max" -> 1
        "Maximum" -> 1
        "MaxPool" -> 1
        "Mean" -> 1
        "Min" -> 1
        "Mul" -> 1
        "Neg" -> 1
        "Placeholder" -> 1
        "OneHot" -> 1
        "ReadVariableOp" -> 1
        "RefIdentity" -> 1
        "Relu" -> 1
        "ReluGrad" -> 1
        "Reshape" -> 1
        "Select" -> 1
        "Size" -> 1
        "SoftmaxCrossEntropyWithLogits" -> 2
        "Square" -> 1
        "SparseSegmentSum" -> 1
        "Sub" -> 1
        "Sum" -> 1
        "Tile" -> 1
        "Transpose" -> 1
        "TruncatedNormal" -> 1
        "VarHandleOp" -> 1
        "Variable" -> 1
        "ZerosLike" -> 1
        "Fill" -> 1
        _ -> error $ "numOutputs not implemented for " ++ show (o ^. op)

-- Divides `x / y` assuming `x, y >= 0`, treating `0 / 0 = 0`
safeShapeDiv :: Tensor v1 Int32 -> Tensor v2 Int32 -> Tensor Build Int32
safeShapeDiv x y = x `CoreOps.div` (CoreOps.maximum y 1)

allDimensions :: Tensor Build Int32
allDimensions = vector [-1 :: Int32]

rangeOfRank :: forall v1 t. TensorType t => Tensor v1 t -> Tensor Build Int32
rangeOfRank x = CoreOps.range 0 (CoreOps.rank x) 1

lookupAttr ::  Attribute a1 => NodeDef -> Text -> a1
lookupAttr nodeDef attrName = nodeDef ^. attr . at attrName . non def . attrLens


text:
{-# LANGUAGE RankNTypes #-}

module Test.Grenade.Layers.Internal.Reference where

import           Grenade.Types
import           Numeric.LinearAlgebra
import qualified Numeric.LinearAlgebra.Static as H
import           GHC.TypeLits

im2col :: Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
im2col nrows ncols srows scols m =
  let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
  in  im2colFit starts nrows ncols m

vid2col :: Int -> Int -> Int -> Int -> Int -> Int -> [Matrix RealNum] -> Matrix RealNum
vid2col nrows ncols srows scols inputrows inputcols ms =
  let starts = fittingStarts inputrows nrows srows inputcols ncols scols
      subs   = fmap (im2colFit starts nrows ncols) ms
  in  foldl1 (|||) subs

im2colFit :: [(Int,Int)] -> Int -> Int -> Matrix RealNum -> Matrix RealNum
im2colFit starts nrows ncols m =
  let imRows = fmap (\start -> flatten $ subMatrix start (nrows, ncols) m) starts
  in  fromRows imRows

col2vid :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> [Matrix RealNum]
col2vid nrows ncols srows scols drows dcols m =
  let starts = fittingStart (cols m) (nrows * ncols) (nrows * ncols)
      r      = rows m
      mats   = fmap (\s -> subMatrix (0,s) (r, nrows * ncols) m) starts
      colSts = fittingStarts drows nrows srows dcols ncols scols
  in  fmap (col2imfit colSts nrows ncols drows dcols) mats

col2im :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2im krows kcols srows scols drows dcols m =
  let rs       = map toList $ toColumns m
      rs'      = zip [0..] rs
      indicies = (\[a,b] -> (a,b)) <$> sequence [[0..(krows-1)], [0..(kcols-1)]]
      accums   = concatMap (\(offset, column) -> zipWith (comb offset) indicies column) rs'
  in accum (konst 0 (drows, dcols)) (+) accums
  where
    comb o (i, j) x = let w = (div (dcols - kcols) scols) + 1
                          (a, b) = divMod o w
                      in  ((i + srows * a, j + scols * b), x)

col2imfit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2imfit starts krows kcols drows dcols m =
  let indicies   = (\[a,b] -> (a,b)) <$> sequence [[0..(krows-1)], [0..(kcols-1)]]
      convs      = fmap (zip indicies . toList) . toRows $ m
      pairs      = zip convs starts
      accums     = concatMap (\(conv',(stx',sty')) -> fmap (\((ix,iy), val) -> ((ix + stx', iy + sty'), val)) conv') pairs
  in  accum (konst 0 (drows, dcols)) (+) accums

poolForward :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
poolForward nrows ncols srows scols outputRows outputCols m =
  let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
  in  poolForwardFit starts nrows ncols outputRows outputCols m

poolForwardList :: Functor f => Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> f (Matrix RealNum) -> f (Matrix RealNum)
poolForwardList nrows ncols srows scols inRows inCols outputRows outputCols ms =
  let starts = fittingStarts inRows nrows srows inCols ncols scols
  in  poolForwardFit starts nrows ncols outputRows outputCols <$> ms

poolForwardFit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
poolForwardFit starts nrows ncols _ outputCols m =
  let els    = fmap (\start -> maxElement $ subMatrix start (nrows, ncols) m) starts
  in  matrix outputCols els

poolBackward :: Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum -> Matrix RealNum
poolBackward krows kcols srows scols inputMatrix gradientMatrix =
  let inRows     = rows inputMatrix
      inCols     = cols inputMatrix
      starts     = fittingStarts inRows krows srows inCols kcols scols
  in  poolBackwardFit starts krows kcols inputMatrix gradientMatrix

poolBackwardList :: Functor f => Int -> Int -> Int -> Int -> Int -> Int -> f (Matrix RealNum, Matrix RealNum) -> f (Matrix RealNum)
poolBackwardList krows kcols srows scols inRows inCols inputMatrices =
  let starts     = fittingStarts inRows krows srows inCols kcols scols
  in  uncurry (poolBackwardFit starts krows kcols) <$> inputMatrices

poolBackwardFit :: [(Int,Int)] -> Int -> Int -> Matrix RealNum -> Matrix RealNum -> Matrix RealNum
poolBackwardFit starts krows kcols inputMatrix gradientMatrix =
  let inRows     = rows inputMatrix
      inCols     = cols inputMatrix
      inds       = fmap (\start -> maxIndex $ subMatrix start (krows, kcols) inputMatrix) starts
      grads      = toList $ flatten gradientMatrix
      grads'     = zip3 starts grads inds
      accums     = fmap (\((stx',sty'),grad,(inx, iny)) -> ((stx' + inx, sty' + iny), grad)) grads'
  in  accum (konst 0 (inRows, inCols)) (+) accums

-- | These functions are not even remotely safe, but it's only called from the statically typed
--   commands, so we should be good ?!?!?
--   Returns the starting sub matrix locations which fit inside the larger matrix for the
--   convolution. Takes into account the stride and kernel size.
fittingStarts :: Int -> Int -> Int -> Int -> Int -> Int -> [(Int,Int)]
fittingStarts nrows kernelrows steprows ncols kernelcols stepcolsh =
  let rs = fittingStart nrows kernelrows steprows
      cs = fittingStart ncols kernelcols stepcolsh
      ls = sequence [rs, cs]
  in  fmap (\[a,b] -> (a,b)) ls

-- | Returns the starting sub vector which fit inside the larger vector for the
--   convolution. Takes into account the stride and kernel size.
fittingStart :: Int -> Int -> Int -> [Int]
fittingStart width kernel steps =
  let go left | left + kernel < width
              = left : go (left + steps)
              | left + kernel == width
              = [left]
              | otherwise
              = []
  in go 0

convBackProp :: Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int
             -> Int -> Int
             -> (Matrix RealNum, Matrix RealNum)
convBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols =
  let fs       = [0..filters-1]
      hs       = [let h_start = h * strideRows in (h, h_start) | h <- [0..outRows - 1]]
      ws       = [let w_start = w * strideCols in (w, w_start) | w <- [0..outCols - 1]]
      fhws     = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      dX_accum = concatMap dxLoop fhws
      dX       = accum (konst 0 (channels * inRows, inCols)) (+) dX_accum
      dW_accum = concatMap dwLoop fhws
      dW       = accum (konst 0 (kernelRows * kernelCols * channels, filters)) (+) dW_accum
  in (dX, dW)
  where
    dxLoop (f, (h, h_start), (w, w_start))
      = [ ((c * inRows + h_start + i, w_start + j), (indexAtdOut f h w) * (indexAtKernel f c i j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]]

    dwLoop (f, (h, h_start), (w, w_start))
      = [ ((c * kernelRows * kernelCols + i * kernelCols + j, f), (indexAtdOut f h w) * (indexAtIn c (h_start + i) (w_start + j))) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]]

    indexAtIn       c x y = input `atIndex` (c * inRows + x, y)
    indexAtdOut     c x y = dout `atIndex` (c * outRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

convForwards :: Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int -> Int
             -> Int -> Int
             -> Int -> Int
             -> Matrix RealNum
convForwards input channels inRows inCols kernel filters kernelRows kernelCols outRows outCols strideRows strideCols =
  let fs     = [0..filters-1]
      hs     = [(h, h * strideRows) | h <- [0..outRows - 1], h * strideRows + kernelRows - 1 < inRows]
      ws     = [(w, w * strideCols) | w <- [0..outCols - 1], w * strideCols + kernelCols - 1 < inCols]
      fhws   = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      accums = map loopIter fhws
  in accum (konst 0 (filters * outRows, outCols)) (+) accums
  where
    loopIter (f, (h, h_start), (w, w_start))
      = ((f * outRows + h, w), sum [ (indexAtKernel f c i j) * (indexAtIn c (h_start + i) (w_start + j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]])

    indexAtIn c x y = input `atIndex` (c * inRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

convForwardsWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> Matrix RealNum
convForwardsWithPadding input channels inRows inCols kernel filters kernelRows kernelCols outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
  in convForwards paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols outRows outCols strideRows strideCols

convBackPropWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> (Matrix RealNum, Matrix RealNum)
convBackPropWithPadding input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
      (dX', dW)   = convBackProp paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      accums'     = [((c * inRows + x, y), dX' `atIndex` (c * padded_r + x + padt, y + padl)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      dX          = accum (konst 0 (channels * inRows, inCols)) (+) accums'
  in  (dX, dW)

biasConvForwards :: Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int -> Int
                 -> Vector RealNum
                 -> Int -> Int
                 -> Int -> Int
                 -> Matrix RealNum
biasConvForwards input channels inRows inCols kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols =
  let fs     = [0..filters-1]
      hs     = [(h, h * strideRows) | h <- [0..outRows - 1], h * strideRows + kernelRows - 1 < inRows]
      ws     = [(w, w * strideCols) | w <- [0..outCols - 1], w * strideCols + kernelCols - 1 < inCols]
      fhws   = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      accums = map loopIter fhws
  in accum (konst 0 (filters * outRows, outCols)) (+) accums
  where
    loopIter (f, (h, h_start), (w, w_start))
      = ((f * outRows + h, w), indexAtBias f + sum [ (indexAtKernel f c i j) * (indexAtIn c (h_start + i) (w_start + j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]])

    indexAtBias i = bias `atIndex` i
    indexAtIn c x y = input `atIndex` (c * inRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

biasConvForwardsWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Vector RealNum
                        -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> Matrix RealNum
biasConvForwardsWithPadding input channels inRows inCols kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
  in biasConvForwards paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols

biasConvBackProp :: Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int
                 -> Int -> Int
                 -> (Matrix RealNum, Matrix RealNum, Vector RealNum)
biasConvBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols =
  let (dX, dW)  = convBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      db_accums = [ (f, sum [ indexAtdOut f i j | i <- [0..outRows-1], j <- [0..outCols-1]]) | f <- [0..filters-1]]
      db        = accum (konst 0 filters) (+) db_accums
  in (dX, dW, db)
  where
    indexAtdOut     c x y = dout `atIndex` (c * outRows + x, y)

biasConvBackPropWithPadding :: Matrix RealNum -> Int -> Int -> Int
                            -> Matrix RealNum -> Int -> Int -> Int
                            -> Matrix RealNum -> Int -> Int
                            -> Int -> Int
                            -> Int -> Int -> Int -> Int
                            -> (Matrix RealNum, Matrix RealNum, Vector RealNum)
biasConvBackPropWithPadding input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols padl padt padr padb =
  let accums        = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r      = inRows + padt + padb
      padded_c      = inCols + padl + padr
      paddedInput   = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
      (dX', dW, db) = biasConvBackProp paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      accums'       = [((c * inRows + x, y), dX' `atIndex` (c * padded_r + x + padt, y + padl)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      dX            = accum (konst 0 (channels * inRows, inCols)) (+) accums'
  in  (dX, dW, db)

naiveFullyConnectedRunForwards :: forall i o. (KnownNat i, KnownNat o)
                               => H.L o i       -- Weights
                               -> H.R o         -- Biases
                               -> H.R i         -- Input
                               -> (H.R i, H.R o)  -- (Tape, Output)
naiveFullyConnectedRunForwards w b i = (i, b + (w H.#> i))

naiveFullyConnectedBackprop :: forall i o. (KnownNat i, KnownNat o)
                            => H.L o i             -- Weights
                            -> (H.R i, H.R o)        -- (Tape, Output)
                            -> (H.L o i, H.R o, H.R i) -- (NablaW, NablaB, derivatives)
naiveFullyConnectedBackprop w (tape, out) = (w', b', d)
  where
    b' = out
    w' = H.outer out tape
    d  = (H.tr w) H.#> out

-- Implementation reference: https://papers.nips.cc/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf
naiveLRNForwards :: RealNum -> RealNum -> RealNum -> Int -> [[[RealNum]]] -> [[[RealNum]]]
naiveLRNForwards a b k n values = [
    [[ g ch ro co | co <- [0..length (values!!ch!!ro) - 1] ] | ro <- [0..length (values!!ch) - 1]] | ch <- [0..cs-1]
  ]
  where
    cs = length values
    f ch ro co = values!!ch!!ro!!co
    g ch ro co = (f ch ro co) / den
      where
        den  = den' ** b
        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]
        summation = sum [ (f j ro co) ** 2 | j <- [lower..upper]]
        den' = k + a * summation

naiveLRNBackwards :: RealNum         -- a
                    -> RealNum       -- b
                    -> RealNum       -- k
                    -> Int          -- n
                    -> [[[RealNum]]] -- inputs
                    -> [[[RealNum]]] -- backpropagated error
                    -> [[[RealNum]]] -- error to propagate further
naiveLRNBackwards a b k n values errs = [
    [[ ng ch ro co | co <- [0..length (values!!ch!!ro) - 1] ] | ro <- [0..length (values!!ch) - 1]] | ch <- [0..cs-1]
  ]
  where
    cs = length values
    f  ch ro co = values!!ch!!ro!!co
    nf ch ro co = errs!!ch!!ro!!co
    c ch ro co = den
      where
        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]
        summation = sum [ (f j ro co) ** 2 | j <- [lower..upper]]
        den = k + a * summation
    ng ch ro co = t1 - t2
      where
        t1 = (c ch ro co) ** (-b) * (nf ch ro co)
        t2 = 2 * b * a * (f ch ro co) * (c ch ro co) ** (-b - 1) * s
        s  = sum [(f q ro co) * (nf q ro co) | q <- [lower..upper]]

        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]

max_stars_repo_name: th-char/grenade

max_issues_repo_name: th-char/grenade

max_forks_repo_name: th-char/grenade