{-# LANGUAGE MagicHash, CPP, UnboxedTuples, BangPatterns, FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Packed.Internal.Vector
-- Copyright   :  (c) Alberto Ruiz 2007
-- License     :  GPL-style
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  provisional
-- Portability :  portable (uses FFI)
--
-- Vector implementation
--
-----------------------------------------------------------------------------

module Data.Packed.Internal.Vector (
    Vector, dim,
    fromList, toList, (|>),
    join, (@>), safe, at, at', subVector, takesV,
    mapVector, mapVectorWithIndex, zipVectorWith, unzipVectorWith,
    mapVectorM, mapVectorM_, mapVectorWithIndexM, mapVectorWithIndexM_,
    foldVector, foldVectorG, foldLoop, foldVectorWithIndex,
    createVector, vec,
    asComplex, asReal, float2DoubleV, double2FloatV,
    stepF, stepD, condF, condD,
    conjugateQ, conjugateC,
    fwriteVector, freadVector, fprintfVector, fscanfVector,
    cloneVector,
    unsafeToForeignPtr,
    unsafeFromForeignPtr,
    unsafeWith
) where

import Data.Packed.Internal.Common
import Data.Packed.Internal.Signatures
import Foreign.Marshal.Alloc(free)
import Foreign.Marshal.Array(peekArray, pokeArray, copyArray, advancePtr)
import Foreign.ForeignPtr(ForeignPtr, castForeignPtr)
import Foreign.Ptr(Ptr)
import Foreign.Storable(Storable, peekElemOff, pokeElemOff, sizeOf)
import Foreign.C.String
import Foreign.C.Types
import Data.Complex
import Control.Monad(when)
import System.IO.Unsafe(unsafePerformIO)

#if __GLASGOW_HASKELL__ >= 605
import GHC.ForeignPtr           (mallocPlainForeignPtrBytes)
#else
import Foreign.ForeignPtr       (mallocForeignPtrBytes)
#endif

import GHC.Base
#if __GLASGOW_HASKELL__ < 612
import GHC.IOBase hiding (liftIO)
#endif

import qualified Data.Vector.Storable as Vector
import Data.Vector.Storable(Vector,
                            unsafeToForeignPtr,
                            unsafeFromForeignPtr,
                            unsafeWith)

-- | Number of elements
dim :: (Storable t) => Vector t -> Int
dim = Vector.length

-- C-Haskell vector adapter
-- vec :: Adapt (CInt -> Ptr t -> r) (Vector t) r
vec :: (Storable t) => Vector t -> (((CInt -> Ptr t -> t1) -> t1) -> IO b) -> IO b
vec x f = unsafeWith x $ \p -> do
    let v g = do
        g (fi $ dim x) p
    f v
{-# INLINE vec #-}

-- allocates memory for a new vector
createVector :: Storable a => Int -> IO (Vector a)
createVector n = do
    when (n <= 0) $ error ("trying to createVector of dim "++show n)
    fp <- doMalloc undefined
    return $ unsafeFromForeignPtr fp 0 n
  where
    --
    -- Use the much cheaper Haskell heap allocated storage
    -- for foreign pointer space we control
    --
    doMalloc :: Storable b => b -> IO (ForeignPtr b)
    doMalloc dummy = do
#if __GLASGOW_HASKELL__ >= 605
        mallocPlainForeignPtrBytes (n * sizeOf dummy)
#else
        mallocForeignPtrBytes      (n * sizeOf dummy)
#endif

{- | creates a Vector from a list:

@> fromList [2,3,5,7]
4 |> [2.0,3.0,5.0,7.0]@

-}
fromList :: Storable a => [a] -> Vector a
fromList l = unsafePerformIO $ do
    v <- createVector (length l)
    unsafeWith v $ \ p -> pokeArray p l
    return v

safeRead v = inlinePerformIO . unsafeWith v
{-# INLINE safeRead #-}

inlinePerformIO :: IO a -> a
inlinePerformIO (IO m) = case m realWorld# of (# _, r #) -> r
{-# INLINE inlinePerformIO #-}

{- | extracts the Vector elements to a list

@> toList (linspace 5 (1,10))
[1.0,3.25,5.5,7.75,10.0]@

-}
toList :: Storable a => Vector a -> [a]
toList v = safeRead v $ peekArray (dim v)

{- | An alternative to 'fromList' with explicit dimension. The input
     list is explicitly truncated if it is too long, so it may safely
     be used, for instance, with infinite lists.

     This is the format used in the instances for Show (Vector a).
-}
(|>) :: (Storable a) => Int -> [a] -> Vector a
infixl 9 |>
n |> l = if length l' == n
            then fromList l'
            else error "list too short for |>"
  where l' = take n l

-- | access to Vector elements without range checking
at' :: Storable a => Vector a -> Int -> a
at' v n = safeRead v $ flip peekElemOff n
{-# INLINE at' #-}

--
-- turn off bounds checking with -funsafe at configure time.
-- ghc will optimise away the salways true case at compile time.
--
#if defined(UNSAFE)
safe :: Bool
safe = False
#else
safe = True
#endif

-- | access to Vector elements with range checking.
at :: Storable a => Vector a -> Int -> a
at v n
    | safe      = if n >= 0 && n < dim v
                    then at' v n
                    else error "vector index out of range"
    | otherwise = at' v n
{-# INLINE at #-}

{- | takes a number of consecutive elements from a Vector

@> subVector 2 3 (fromList [1..10])
3 |> [3.0,4.0,5.0]@

-}
subVector :: Storable t => Int       -- ^ index of the starting element
                        -> Int       -- ^ number of elements to extract
                        -> Vector t  -- ^ source
                        -> Vector t  -- ^ result
subVector = Vector.slice

{- | Reads a vector position:

@> fromList [0..9] \@\> 7
7.0@

-}
(@>) :: Storable t => Vector t -> Int -> t
infixl 9 @>
(@>) = at

{- | creates a new Vector by joining a list of Vectors

@> join [fromList [1..5], constant 1 3]
8 |> [1.0,2.0,3.0,4.0,5.0,1.0,1.0,1.0]@

-}
join :: Storable t => [Vector t] -> Vector t
join [] = error "joining zero vectors"
join [v] = v
join as = unsafePerformIO $ do
    let tot = sum (map dim as)
    r <- createVector tot
    unsafeWith r $ \ptr ->
        joiner as tot ptr
    return r
  where joiner [] _ _ = return ()
        joiner (v:cs) _ p = do
            let n = dim v
            unsafeWith v $ \pb -> copyArray p pb n
            joiner cs 0 (advancePtr p n)

{- | Extract consecutive subvectors of the given sizes.

@> takesV [3,4] (linspace 10 (1,10))
[3 |> [1.0,2.0,3.0],4 |> [4.0,5.0,6.0,7.0]]@

-}
takesV :: Storable t => [Int] -> Vector t -> [Vector t]
takesV ms w | sum ms > dim w = error $ "takesV " ++ show ms ++ " on dim = " ++ (show $ dim w)
            | otherwise = go ms w
    where go [] _ = []
          go (n:ns) v = subVector 0 n v
                      : go ns (subVector n (dim v - n) v)

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

-- | transforms a complex vector into a real vector with alternating real and imaginary parts
asReal :: (RealFloat a, Storable a) => Vector (Complex a) -> Vector a
asReal v = unsafeFromForeignPtr (castForeignPtr fp) (2*i) (2*n)
    where (fp,i,n) = unsafeToForeignPtr v

-- | transforms a real vector into a complex vector with alternating real and imaginary parts
asComplex :: (RealFloat a, Storable a) => Vector a -> Vector (Complex a)
asComplex v = unsafeFromForeignPtr (castForeignPtr fp) (i `div` 2) (n `div` 2)
    where (fp,i,n) = unsafeToForeignPtr v

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

float2DoubleV :: Vector Float -> Vector Double
float2DoubleV v = unsafePerformIO $ do
    r <- createVector (dim v)
    app2 c_float2double vec v vec r "float2double"
    return r

double2FloatV :: Vector Double -> Vector Float
double2FloatV v = unsafePerformIO $ do
    r <- createVector (dim v)
    app2 c_double2float vec v vec r "double2float2"
    return r

foreign import ccall unsafe "float2double" c_float2double:: TFV
foreign import ccall unsafe "double2float" c_double2float:: TVF

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

stepF :: Vector Float -> Vector Float
stepF v = unsafePerformIO $ do
    r <- createVector (dim v)
    app2 c_stepF vec v vec r "stepF"
    return r

stepD :: Vector Double -> Vector Double
stepD v = unsafePerformIO $ do
    r <- createVector (dim v)
    app2 c_stepD vec v vec r "stepD"
    return r

foreign import ccall unsafe "stepF" c_stepF :: TFF
foreign import ccall unsafe "stepD" c_stepD :: TVV

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

condF :: Vector Float -> Vector Float -> Vector Float -> Vector Float -> Vector Float -> Vector Float
condF x y l e g = unsafePerformIO $ do
    r <- createVector (dim x)
    app6 c_condF vec x vec y vec l vec e vec g vec r "condF"
    return r

condD :: Vector Double -> Vector Double -> Vector Double -> Vector Double -> Vector Double -> Vector Double
condD x y l e g = unsafePerformIO $ do
    r <- createVector (dim x)
    app6 c_condD vec x vec y vec l vec e vec g vec r "condD"
    return r

foreign import ccall unsafe "condF" c_condF :: CInt -> PF -> CInt -> PF -> CInt -> PF -> TFFF
foreign import ccall unsafe "condD" c_condD :: CInt -> PD -> CInt -> PD -> CInt -> PD -> TVVV

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

conjugateAux fun x = unsafePerformIO $ do
    v <- createVector (dim x)
    app2 fun vec x vec v "conjugateAux"
    return v

conjugateQ :: Vector (Complex Float) -> Vector (Complex Float)
conjugateQ = conjugateAux c_conjugateQ
foreign import ccall unsafe "conjugateQ" c_conjugateQ :: TQVQV

conjugateC :: Vector (Complex Double) -> Vector (Complex Double)
conjugateC = conjugateAux c_conjugateC
foreign import ccall unsafe "conjugateC" c_conjugateC :: TCVCV

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

cloneVector :: Storable t => Vector t -> IO (Vector t)
cloneVector v = do
        let n = dim v
        r <- createVector n
        let f _ s _ d =  copyArray d s n >> return 0
        app2 f vec v vec r "cloneVector"
        return r

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

-- | map on Vectors
mapVector :: (Storable a, Storable b) => (a-> b) -> Vector a -> Vector b
mapVector f v = unsafePerformIO $ do
    w <- createVector (dim v)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) = return ()
                go !k = do x <- peekElemOff p k
                           pokeElemOff      q k (f x)
                           go (k-1)
            go (dim v -1)
    return w
{-# INLINE mapVector #-}

-- | zipWith for Vectors
zipVectorWith :: (Storable a, Storable b, Storable c) => (a-> b -> c) -> Vector a -> Vector b -> Vector c
zipVectorWith f u v = unsafePerformIO $ do
    let n = min (dim u) (dim v)
    w <- createVector n
    unsafeWith u $ \pu ->
        unsafeWith v $ \pv ->
            unsafeWith w $ \pw -> do
                let go (-1) = return ()
                    go !k = do x <- peekElemOff pu k
                               y <- peekElemOff pv k
                               pokeElemOff      pw k (f x y)
                               go (k-1)
                go (n -1)
    return w
{-# INLINE zipVectorWith #-}

-- | unzipWith for Vectors
unzipVectorWith :: (Storable (a,b), Storable c, Storable d)
                   => ((a,b) -> (c,d)) -> Vector (a,b) -> (Vector c,Vector d)
unzipVectorWith f u = unsafePerformIO $ do
      let n = dim u
      v <- createVector n
      w <- createVector n
      unsafeWith u $ \pu ->
          unsafeWith v $ \pv ->
              unsafeWith w $ \pw -> do
                  let go (-1) = return ()
                      go !k   = do z <- peekElemOff pu k
                                   let (x,y) = f z
                                   pokeElemOff      pv k x
                                   pokeElemOff      pw k y
                                   go (k-1)
                  go (n-1)
      return (v,w)
{-# INLINE unzipVectorWith #-}

foldVector :: Storable a => (a -> b -> b) -> b -> Vector a -> b
foldVector f x v = unsafePerformIO $
    unsafeWith v $ \p -> do
        let go (-1) s = return s
            go !k !s = do y <- peekElemOff p k
                          go (k-1::Int) (f y s)
        go (dim v -1) x
{-# INLINE foldVector #-}

-- the zero-indexed index is passed to the folding function
foldVectorWithIndex :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b
foldVectorWithIndex f x v = unsafePerformIO $
    unsafeWith v $ \p -> do
        let go (-1) s = return s
            go !k !s = do y <- peekElemOff p k
                          go (k-1::Int) (f k y s)
        go (dim v -1) x
{-# INLINE foldVectorWithIndex #-}

foldLoop f s0 d = go (d - 1) s0
     where
       go 0 s = f (0::Int) s
       go !j !s = go (j - 1) (f j s)

foldVectorG f s0 v = foldLoop g s0 (dim v)
    where g !k !s = f k (at' v) s
          {-# INLINE g #-} -- Thanks to Ryan Ingram (http://permalink.gmane.org/gmane.comp.lang.haskell.cafe/46479)
{-# INLINE foldVectorG #-}

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

-- | monadic map over Vectors
--    the monad @m@ must be strict
mapVectorM :: (Storable a, Storable b, Monad m) => (a -> m b) -> Vector a -> m (Vector b)
mapVectorM f v = do
    w <- return $! unsafePerformIO $! createVector (dim v)
    mapVectorM' w 0 (dim v -1)
    return w
    where mapVectorM' w' !k !t
              | k == t               = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f x
                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
              | otherwise            = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f x
                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
                                       mapVectorM' w' (k+1) t
{-# INLINE mapVectorM #-}

-- | monadic map over Vectors
mapVectorM_ :: (Storable a, Monad m) => (a -> m ()) -> Vector a -> m ()
mapVectorM_ f v = do
    mapVectorM' 0 (dim v -1)
    where mapVectorM' !k !t
              | k == t            = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    f x
              | otherwise         = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    _ <- f x
                                    mapVectorM' (k+1) t
{-# INLINE mapVectorM_ #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
--    the monad @m@ must be strict
mapVectorWithIndexM :: (Storable a, Storable b, Monad m) => (Int -> a -> m b) -> Vector a -> m (Vector b)
mapVectorWithIndexM f v = do
    w <- return $! unsafePerformIO $! createVector (dim v)
    mapVectorM' w 0 (dim v -1)
    return w
    where mapVectorM' w' !k !t
              | k == t               = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f k x
                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
              | otherwise            = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f k x
                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
                                       mapVectorM' w' (k+1) t
{-# INLINE mapVectorWithIndexM #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
mapVectorWithIndexM_ :: (Storable a, Monad m) => (Int -> a -> m ()) -> Vector a -> m ()
mapVectorWithIndexM_ f v = do
    mapVectorM' 0 (dim v -1)
    where mapVectorM' !k !t
              | k == t            = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    f k x
              | otherwise         = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    _ <- f k x
                                    mapVectorM' (k+1) t
{-# INLINE mapVectorWithIndexM_ #-}

mapVectorWithIndex :: (Storable a, Storable b) => (Int -> a -> b) -> Vector a -> Vector b
--mapVectorWithIndex g = head . mapVectorWithIndexM (\a b -> [g a b])
mapVectorWithIndex f v = unsafePerformIO $ do
    w <- createVector (dim v)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) = return ()
                go !k = do x <- peekElemOff p k
                           pokeElemOff      q k (f k x)
                           go (k-1)
            go (dim v -1)
    return w
{-# INLINE mapVectorWithIndex #-}

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

-- | Loads a vector from an ASCII file (the number of elements must be known in advance).
fscanfVector :: FilePath -> Int -> IO (Vector Double)
fscanfVector filename n = do
    charname <- newCString filename
    res <- createVector n
    app1 (gsl_vector_fscanf charname) vec res "gsl_vector_fscanf"
    free charname
    return res

foreign import ccall unsafe "vector_fscanf" gsl_vector_fscanf:: Ptr CChar -> TV

-- | Saves the elements of a vector, with a given format (%f, %e, %g), to an ASCII file.
fprintfVector :: FilePath -> String -> Vector Double -> IO ()
fprintfVector filename fmt v = do
    charname <- newCString filename
    charfmt <- newCString fmt
    app1 (gsl_vector_fprintf charname charfmt) vec v "gsl_vector_fprintf"
    free charname
    free charfmt

foreign import ccall unsafe "vector_fprintf" gsl_vector_fprintf :: Ptr CChar -> Ptr CChar -> TV

-- | Loads a vector from a binary file (the number of elements must be known in advance).
freadVector :: FilePath -> Int -> IO (Vector Double)
freadVector filename n = do
    charname <- newCString filename
    res <- createVector n
    app1 (gsl_vector_fread charname) vec res "gsl_vector_fread"
    free charname
    return res

foreign import ccall unsafe "vector_fread" gsl_vector_fread:: Ptr CChar -> TV

-- | Saves the elements of a vector to a binary file.
fwriteVector :: FilePath -> Vector Double -> IO ()
fwriteVector filename v = do
    charname <- newCString filename
    app1 (gsl_vector_fwrite charname) vec v "gsl_vector_fwrite"
    free charname

foreign import ccall unsafe "vector_fwrite" gsl_vector_fwrite :: Ptr CChar -> TV


{-# LANGUAGE NoImplicitPrelude, TypeSynonymInstances, RankNTypes, StandaloneDeriving,
ScopedTypeVariables, TypeApplications, GADTs, TypeOperators, DeriveTraversable,
FlexibleInstances, AllowAmbiguousTypes, UndecidableInstances, TypeFamilies, LambdaCase #-}
module Lib
    where
import Numeric.LinearAlgebra
import Control.Category
import Prelude hiding ((.), id)
import Control.Arrow ((***))
someFunc :: IO ()
someFunc = putStrLn "someFunc"

{-
type V a = a -> Double
type D a = a -> Double
D (V a)
type V a  = ContT () Q a -- ContT b Q a = (a ->Q b) -> Q b. Has some way of dealing with a index. Some kind of
SelT -- takes a diag, unsummed?
-- (a -> Q b) -> Q a

newtype Dag a = Dag (a -> Q ())

class Dagger a where
    type Dual :: * -> * -- *?
    dag :: a -> Dual a
instance Dagger (D (Vec a)) where
    type Dual = Vec
    dag (D v) = join $ ContT ()
-- dagger needs to embed a search procedure?

instance Dagger (D (Vec Bool)) where
    type Dual = Vec Bool
    dag (D v) = join $ ContT ()
instance Dagger a b where
    dag :: a -> b
    dag' :: b -> a
instance Dagger k where
    type family Dual :: * -> *
    dag :: k a (Dual a) -- dag'?

-- also a kind of dot product
dag1 :: ((a -> Q ()) -> Q ()) -> (((a -> Q ()) -> Q ()) -> Q ()
dag1 f =

-- it's a kind of dot product.
lift :: f Double -> (a -> Double) -> Double
lift w f = fold (+) $ (mapWithIndex \k v -> (f k) * v) w

-- tensor is a_i b_i, mutiply out without summing
tensor :: f a ->  f a  -> f a
tensor = intersectWith (*)

type V' a = (a -> Q ()) -> Q ()
tensor :: V' a -> V' a -> V' a

dag :: (V' Bool -> Q ()) -> V' Bool
dag f = \g -> True False

dag :: (Bool -> Q ()) -> V' Bool
dag :: (Bool -> Q ()) -> (Bool -> Q ()) -> Q ()
dag f g = (f True) (g True) + (f False) (g False)

dag :: (Bool -> Q ()) -> Q Bool
dag f = (f True) >>= true + (f False) >>= false

dag :: Q Bool -> (Bool -> Q ())
dag xs b = lookup b xs
dag :: Eq a => Q a -> (a -> Q ()) -- maybe Ord
dag xs b = lookup b xs

-- If we mixed Hedges select sat solver with the Q monad?
-- (Bool -> Q Bool) -> Q Bool ? the unsummed trace? diag? Hedgest hing would come out weirder.

true :: () -> Q Bool
true _ = pure True
false :: () -> Q Bool
false _ = pure False

-- the seelect function.
-- (b -> m ()) -> m b

-}

type M = Matrix Double

{-

data BMatrix a b = BMatrix M M M M | Id -- A B C D

instance Category BMatrix where
    id = Id
    Id . x = x
    x . Id = x
    (BMatrix a b c d) . (BMatrix e f g h) = BMatrix <\> where
            q = d + e
            qlu = lupacked q
            qinv = lusolve qlu
            a' = a - (b <> (qinv c))
            b' =  b <> (qinv f)
            c' = g <> (qinv c)
            d' = h - g <> (qinv f)

par :: BMatrix a b -> BMatrix c d -> BMatrix (a,c) (b,d)
par (Matrix a b c d) (BMatrix e f g h) = BMatrix (diagBlock [a,e]) (diagBlock [b,f]) (diagBlock [c,g])  (diagBlock [d,h])
dup = BMatrix a (a,a)
dup = BMatrix (diagBlock [ident,ident]) (ident ||| ident) (ident === ident) (diagBlock [ident,ident])
id = BMatrix ident ident ident ident
avg =
-}
data QuadFun a b = QuadFun {mat :: M, vec :: Vector Double, c :: Double } deriving Show

--data QuadFun a b where
--   QuadFun :: (Bounded a, Bounded b, Enum a, Enum b) => M -> Vector Double -> Double -> QuadFun a b }

-- This is kron product.
instance (Enum a, Enum b, Bounded a) => Enum (a,b) where
  toEnum x =  let (a, b) = quotRem x (fromEnum (maxBound :: a)) in
                (toEnum a   , toEnum b)
  fromEnum (a,b) =  (fromEnum (maxBound :: a)) * (fromEnum a) + (fromEnum b)
  {-
instance (Enum a, Enum b, Bounded a) => Enum (Either a b) where
    toEnum x | x > (fromEnum (maxBound :: a)) = Left (toEnum x)
             | otherwise = Right $ toEnum (x - 1 - (fromEnum (maxBound :: a)))
    fromEnum (Left a) =  fromEnum a
    fromEnum (Right b) =  (fromEnum (maxBound :: a)) + 1 + (fromEnum b)
-}
-- deriving instance (Enum a, Enum b) => Enum (a,b)

count :: forall a. (Enum a, Bounded a) => Int
count = (fromEnum (maxBound @a)) - (fromEnum (minBound @a)) + 1

id ::forall a. (Enum a, Bounded a) => QuadFun a a
id = let s = 2 * (count @ a) in QuadFun (ident s) (konst 0 s) 0

--dup :: QuadFun a (a,a)
--dup  = QuadFun (ident

class (Bounded a, Enum a) => BEnum a where
instance (Bounded a, Enum a) => BEnum a where

t7 :: QuadFun () ()
t7 = QuadFun (konst 4 (2,2)) (konst 7 2) 3

t8 = (mat t7) <\> (vec t7)
t4 = par t7 t7
t9 = (mat t4) <\> (vec t4)

compose :: forall a b c. BEnum b => QuadFun b c -> QuadFun a b -> QuadFun a c
compose (QuadFun m' w' c') (QuadFun m w c) = QuadFun m'' w'' c'' where
   n = count @b
   k = rows m -- assuming square
   l = rows m'
   i = ident n
   q = konst 0 (k-n,n)  === i  -- I underneath zeros
   q' = -i === konst 0 (l-n,n)   -- -i above zeros
   m'' = fromBlocks [[m,    q, 0],
                     [tr q ,0, tr q'],
                     [0,    q', m']]
   w'' = vjoin [w, konst 0 n, w']
   c'' = c + c'

identOut :: forall a b. BEnum b => QuadFun a b -> M
identOut (QuadFun m w c) = konst 0 (k-n,n)  === i  where
    n = count @b
    k = rows m
    i = ident n
identIn :: forall a b. BEnum a => QuadFun a b -> M
identIn (QuadFun m w c) = -i === konst 0 (k-n,n)   where
    n = count @a
    k = rows m
    i = ident n
type a :+: b = Either a b
-- I kind of feel like our sign convention is flipped
par :: forall a b c d. (BEnum a, BEnum b, BEnum c, BEnum d) => QuadFun a c -> QuadFun b d -> QuadFun (a :+: b) (c :+: d)
par x@(QuadFun m w c) y@(QuadFun m' w' c') = QuadFun m'' w'' c'' where
    ia = ident (count @a)
    ib = ident (count @b)
    ia' = identIn x
    ib' = identIn y
    ia't = tr ia'
    ib't = tr ib'
    ic = - ident (count @c)
    id = - ident (count @d)
    ic' = identOut x
    id' = identOut y
    ic't = tr ic'
    id't = tr id'
    n = (count @a) + (count @b)
    n' = (count @c) + (count @d)
    --iab = fromBlocks [[0,0,ia,0], [0,0,0,ib], [ia, 0,0,0], [0,ib,0,0]]
    --iab' = fromBlocks [[tr (identIn x),0],  [0 , tr (identIn y)]]
    --zab = konst 0 (n,n)
    -- should be 2 ins + 2 langrange + 2 matr + 2lagarne + 2 outs = 10x10 block matrix
    m'' = fromBlocks [[0 ,0 , ia, 0, 0, 0, 0 ,0,0,0],
                    [0, 0,  0, ib, 0, 0, 0, 0,0,0 ] ,
                    [ia, 0, 0, 0,  ia't, 0,0,0,0, 0],
                    [0, ib, 0, 0,  0,ib't,0, 0,0 ,0],
                    [0, 0, ia',0,  m, 0 , ic',0,0,0],
                    [0, 0,  0, ib',0, m', 0,id',0,0],
                    [0, 0,  0, 0, ic't,0 ,0,0,ic,0],
                    [0, 0,  0, 0,  0,id't,0,0,0,id],
                    [0, 0,  0, 0,  0, 0,  ic,0,0,0],
                    [0, 0,  0, 0,  0, 0,  0,id,0,0]]
    w'' = vjoin [konst 0 (2*n), w,w', konst 0 (2*n')]
    c'' = c + c'

dup :: forall a. BEnum a => QuadFun a (a,a)
dup = QuadFun m 0 0 where -- 1 input, 2 lagrange mutipliers, and 2 outputs.
    ia = ident (count @a)
    m = fromBlocks [ [0, ia,ia, 0,0],
                     [ia, 0, 0, -ia,0],
                     [0, ia, 0, 0, -ia],
                     [0,-ia, 0, 0, 0 ],
                     [0, 0, -ia, 0,0]]

-- fst...?
-- for snd, we could just leave it alone
-- fuse?
-- swap.
data Void
fuse :: forall a. BEnum a => QuadFun (a,a) Void
fuse = QuadFun m 0 0 where -- 2 inputs, 1 lagrange multiplier.
    ia = ident (count @a)
    m = fromBlocks [[0,0,    ia],
                    [0 ,0 , -ia],
                    [ia, -ia, 0]]

-- The analog for this slicing for convex sets may be to slice the
-- dimensions into in and out dimensions. Then

data Cell a = Cell {phi :: Double, j :: Double, next :: a} deriving (Show, Traversable, Foldable, Functor) -- composition of cell is a statically sized vector. Derive applicative?

instance Applicative Cell where
    pure x = Cell 0 0 x
    (Cell phi1 j1 f) <*> (Cell phi2 j2 x) = Cell (phi1 + phi2) (j1 + j2) (f x) -- ? I doubt this makes any sense.
{-
data Cell a = Cell { phi :: a, j :: a}
type f :+: g = Product f g
gaussK :: Lens ((Cell :+: Cell) a) (Cell a)

gaussK :: -> Lens (Cell :*: f a) -- no Cell is very not this.

type Row = (Cell Cell Cell Cell Cell)

Lens (a, b, other) (b, other)

gaussK :: Lens

data Cell2 f a b = Cell2 {phi :: a , j :: a, next :: f b}
data Cell2 f a = Cell2 {phi :: f Double a , j ::f Double, next :: f a}

-- zipA :: f a -> f b -> f (a,b)
-- zipA x y = (,) <$> x <*> y
-- monProd

fmap2 = fmap . fmap
fmap4 = fmap2 . fmap2
fmap8 = fmap4 . fmap4

zipA x y = (,) <$> x <*> y
zip2 = zipA . (fmap zipA) -- hmm. Maybe we should be using Compose. we're going to need famp2. Compose will get us these instances for free.
zip4 = zip2 . (fmap2 zip2)
zip8 = zip4 . (fmap4 zip4)

parK :: Lens (f b) b -> Lens (g a) a -> Lens (f g a) a
parK :: forall b. Lens (f b) (f' b) -> Lens (g a) (g' a) -> Lens (f g a) (f' (g' a))
parK :: forall b. Lens (f g a) (f' g a) -> Lens (g a) (g' a) -> Lens (f g a) (f' (g' a))
parK = compose (fmap l2) l1

Do the y direction inductively

ydir :: Lens (Cell a, Cell a) (a,a)
ydir = \(Cell phi j x, Cell phi2 j2 y) -> ((x,y), \(x',y') -> (Cell phi j x, Cell phi2 j2 y)      )

ydir :: Lens (Cell a, Cell a) (a,a)
ydir = \(a,b) -> ((Cell phi j x) , f) = (gaussK a) in ((Cell phi j y) , g) = (gaussK b) in

Lens ( (X (), (X (), a) (X (), a)
ydir = \(r1, (r2, z) ->  zip8    )    -- = gaussK r1 in  = gaussK r2 in

Lens (f a) a -> Lens (f a) a -> Lens (f a, f a) a
Lens (Cell a) a -> Lens

Lens x y x' y'

-}

{-
can i just use regular lenses? I need to set in kind of a weird way though.
-}

data Lens a b = Lens (a -> (b, b -> a))
-- newtype Lens a b = Lens forall r. (a -> (b -> (b -> a) -> r) -> r) -- cpsify for speed? van Laarhoven?

-- SLens s a b = SLens (a -> ST s (b, b -> ST s a)) mutable update
-- MLens a b = MLens (a -> m (b , b -> m a)) -- monadic lens
-- KLens a b = KLens (a -> (b, k b a)) ---- this is sort of what Conal wrote.
comp :: Lens b c -> Lens a b -> Lens a c
comp (Lens f) (Lens g) = Lens $ \x -> let (y, g') = g x in
                                      let (z, f') = f y in
                                      (z, g' . f')

-- Gauss Seidel of the standard 1 -2 1 matrix
-- j1 is the effective source incliuding the influence from phi values up in the stack
-- Double Cell construction is rather wonky
-- gaussK :: Vec (S (S n)) Double -> Vec (S n)
-- we could also go for the lagrange mutiplier interface.
-- can also push the rho value needed for stability up and down.

-- -2 phi1 + 1 phi2 = ~j1
-- -1 phi1 + -2 phi2 + ...? = j2

{-
We are moving the lower diagonal to the right hand side as the splitting.
That ends up to claculating an effective j based on previous values.
We then triangular solve the upper diagonal, which we are able to find the new diagonal element in terms of the lower values.

-}
-- it's kind of weird that we mutate j on the way down but restore it on the way up.
-- But we do need a way to access j. (Phi (Phi a), J a) (Phi )

gaussK :: Lens (Cell (Cell a)) (Cell a) -- we need the context Cell
gaussK = Lens $ \case (Cell phi1 j1 (Cell phi2 j2 y)) -> (Cell phi2 (j2 - phi1) y , \case (Cell phi2' j2' z) -> let j1' = j1 - phi2' in  -- moving the triangular upsolve to the right hand side
                                                                                Cell (- j1' / 2) j1 (Cell phi2' j2 z))

-- interface in , internals, interface out. was the previous way of talking about it. But then the internals grow, which is fine
-- compose :: Lens in internal out -> Lens in' internal' out' -> Lens in (internal, out, in', internal') out
-- some ind of 2 category? Enriched? The morphisms have this internal structure.
-- this is more a containing relationship. The larger context can be converted into the smaller context.
--

g2 :: Lens (Cell (Cell (Cell a))) (Cell a)
g2 = gaussK `comp` gaussK
g4 = g2 `comp` g2
g8 = g4 `comp` g4
g16 = g8 `comp` g8
g32 = g16 `comp` g16

capZero :: Lens (Cell ()) () -- not even sure i really need this? does runGauss (f `comp` capZero) ~ runGauss f
capZero = Lens $ \case (Cell phi j _) -> ((), \_ -> Cell (- j / 2) j ())

-- runGauss ::  Lens a b -> a -> a -- removes the open ended nature of the thing.
runGauss ::  Lens a () -> a -> a -- this might make more sense as what you really want. This is some kind of cap operation.
runGauss (Lens f) x = let (y , trisolve)  = f x in trisolve y
startingVal :: Cell (Cell (Cell ()))
startingVal = (Cell 0 0 (Cell 0 1 (Cell 0 0 ())))

iters = iterate (runGauss (capZero `comp` g2)) startingVal
--iters' = iterate (runGauss g2) startingVal -- They are different.

sol = ((3><3) [-2,1,0,
               1,-2,1,
               0,1,-2])  <\> (vector [0,1,0])

-- can do inner block solves also (via an iterative method perhaps)
-- this is reminsecent of a multiscale attack.
parC :: Lens a a' -> Lens b b' -> Lens (a,b) (a',b')
parC (Lens f) (Lens g) = Lens $ \(x,y) -> let (x', f') = f x in
                                          let (y', g') = g y in
                                          ((x',y') , f' *** g')

-- profucntor optics paper
data FunList a b t = Done t | More a (FunList a b (b -> t))
-- contents and fill
-- contents (s -> a^n)
-- fill s -> b^n -> t
-- s -> exists n. (a^n, a^n -> s) is a traversal. Could use a Vec. But we already have a Vec. Cell is a Vec
-- s -> exists n, (Vec n a, Vec n a -> s)
-- maybe we do need an applicative. We sort of need to zip together two rows to start going 2d.

-- huh. A block wise schur is kind of the monoidal ppoduct here. No. monoidal product is pure dsum.
-- actually composition? is kind of what plays the game of block wise stuff.
{-

if we want this to be not the case,
comp :: Lens (a,a') a'

schur :: Lens a a' -> (a -> b) -> (b -> a) -> Lens b b' -> Lens (a,b) (a',b')
schur (Lens a) b c (Lens d) = Lens $ \(x,y) -> let (x', a') = a x in
                                               let (y', d') = d x in
                                               (       )
-}
-- Lens (f (g a)) (g a)

                                               -- we can change this all into a continuation form
--

{-
par :: forall a b c d. (BEnum a, BEnum b, BEnum c, BEnum d) => BMatrix a b ->
                         BMatrix c d -> BMatrix (a,c) (b,d)
par q@(QuadFun m w r) q'@(QuadFun m' w' r') = QuadFun q'' w'' r'' where
     a' = diagBlock [sliceA q, sliceA q']
     b' = diagBlock [sliceB q, sliceB q']
     c' = diagBlock [sliceC q, sliceC q']
     d' = diagBlock [sliceD q, sliceD q']
     q'' = fromBlocks [[a',b'], [c',d']]
     u' = vjoin [sliceU v, sliceU v']
     v' = vjoin [sliceV v, sliceV v']
     w'' = vjoin [u', v']
     r'' = r + r'

sliceA :: forall a b. (BEnum a, BEnum b) => QuadFun a b -> M
sliceA (QuadFun m v c) = subMatrix (0,0) (count @a, count @a)  m
sliceB :: forall a b. (BEnum a, BEnum b) => QuadFun a b -> M
sliceB (QuadFun m v c) = subMatrix (0,count @a) (count @a, count @b) m
sliceC :: forall a b. (BEnum a, BEnum b) => QuadFun a b -> M
sliceC (QuadFun m v c) = subMatrix (count @a,0) (count @b, count @a) m
sliceD :: forall a b. (BEnum a, BEnum b) => QuadFun a b -> M
sliceD (QuadFun m v c) = subMatrix (count @a,count @a) (count @b, count @b) m

sliceU :: forall a b. (BEnum a, BEnum b) => QuadFun a b -> Vector Double
sliceU (QuadFun m v c) = subVector 0 (count @a) v
sliceV :: forall a b. (BEnum a, BEnum b) => QuadFun a b -> Vector Double
sliceV (QuadFun m v c) = subVector (count @a) (count @b) v

n = count @a
m = count @b
q = rows m
p = q - m
a = subMatrix (0,0) (count @a, count @a)  m

a = subMatrix
--Other options
-- Keep a matrix with the end and the beginning are the

-- Keep Constraints too?

-- what in the world
-- stack build --ghc-options /usr/lib/libiconv.dylib

-- size
-- rows
-- cols

compose :: QuadFun b c -> QuadFun a b -> QuadFun a c
compose (QuadFun m w c) (QuadFun m' w' c') = QuadFun m'' w'' c'' where
    n = (count @b)
    corner =  fromBlocks [[ident n, 0],[0, -(ident n)]]
    corner' = fromBlocks [[0,0],[corner,0]]
    m'' = fromBlocks [[m ,corner'], [tr corner', m']]
    w'' = vjoin [w, const 0 n, w']
    c'' = c + c

par (QuadFun m w c) (QuadFun m' w' c') =

type a :+: b = Either a b
data QuadFun x = QuadFun

((a,b,c),
 (d,e,f),
 (g,h,i)) =

 data BMatrix' = BMatrix' M M M M M M M M M
 compose (BMatrix' a' b' c' d' e' f' g' h' i') (BMatrix' a b c d e f g h i) =
    BMatrix' a'' b'' c'' d'' e'' f'' g'' h'' i'' where
    a'' = a
    b'' = b ||| c ||| 0
    c'' = 0
    d'' = d === g === 0
    e'' = fromBlocks [[e,f,0], [h, i + a', b'], [0, d', e']]
    f'' = 0 === c' === f'
    g'' =  0
    h'' = 0 ||| g' ||| h'
    i'' = i'
compose' (BVec u' v' w') (BVec u v w) = BVec u (vjoin [v ,w + u', v']) w'

par :: QuadFun a b -> QuadFun c d -> QuadFun (Either a c) (Either b d)
par (BMatrix' a' b' c' d' e' f' g' h' i') (BMatrix' a b c d e f g h i) =
   BMatrix' a'' b'' c'' d'' e'' f'' g'' h'' i'' where
   a'' = diagBlock [a',a]
   b'' = diagBlock [b',b]
   c'' = diagBlock [c',c]
   d'' = diagBlock [d',d]
   e'' = diagBlock [e',e]
   f'' = diagBlock [f',f]
   g'' = diagBlock [g',g]
   h'' = diagBlock [h',h]
   i'' = diagBlock [i',i]

-- If we're just recording the entire program, we might as well just record it in a DSL
-- rather than directly building the matrix
-- data QuadFunDSL = Par QuadFun
-- FreeCat + (Lit HMatrix)

-- iterative matrix solution proximal method?
--
-- gauss seidel
type VBlock = [Vector Double]
-- (V, V)
-- Maybe schur solve works. But with possible splitting
-- Maybe unsolvable schur
-- could damp a little if the diagonal isn't dominant enough
-- (lusolve a (w - V <#| v', ) -- WHat I'm suggesting here is block Jacobi? No
-- No it isn't. Because we're passing back v' it is guass seidel.
-- in order to do block gauss seidel
-- type L = VBlock -> (VBlock, VBlock -> VBlock)

-- Writing in the monoidal categroy style makes parallelism more apparent.

-- [vjoin blocks] -> ([lesser blocks] , update [lesserblocks] -> [blocks]
-- \x:xs -> (a*x + xs, \b -> b )

-- storage in the lens?
-- In my AD lens, I had all the weights in the input. This was ungainly
-- Maybe a choice monad? Everyone could have access to a global state
-- forall s. (ActualLens s a, a -> (b, db -> da))
-- compsoe (ActualLens s a, a -> (b, db -> da))
-- par
-- ActualLens are functional references. Do they allow us to refiy sharing?
-- (ActualLens acc a,   )

-- alternative a -> State s (b, db -> State ds da)

-- Mutiple directions of composition. Horizontsl and vertical?
-- a 2-category?
-- ActualLens s a -> ActualLens s' s

-- Jules Hedges:
-- Lens sigma a b -> Lens sigm' b c -> Lens (sig,sig') a c
-- Lens :: Sigma -> (sigma,  ) -> Lens a b -- hides sigma, we'll never be able to get at it

-}


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-|
Module      : Grenade.Core.Pad
Description : Padding layer for 2D and 3D images
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Pad (
    Pad (..)
  ) where

import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits
import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Layers.Internal.Pad

import           Numeric.LinearAlgebra (konst, subMatrix, diagBlock)
import           Numeric.LinearAlgebra.Static (extract, create)

-- | A padding layer for a neural network.
--
--   Pads on the X and Y dimension of an image.
data Pad  :: Nat
          -> Nat
          -> Nat
          -> Nat -> * where
  Pad  :: Pad padLeft padTop padRight padBottom

instance Show (Pad padLeft padTop padRight padBottom) where
  show Pad = "Pad"

instance UpdateLayer (Pad l t r b) where
  type Gradient (Pad l t r b) = ()
  runUpdate _ x _ = x
  createRandom = return Pad

instance Serialize (Pad l t r b) where
  put _ = return ()
  get = return Pad

-- | A two dimentional image can be padded.
instance ( KnownNat padLeft
         , KnownNat padTop
         , KnownNat padRight
         , KnownNat padBottom
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , (inputRows + padTop + padBottom) ~ outputRows
         , (inputColumns + padLeft + padRight) ~ outputColumns
         ) => Layer (Pad padLeft padTop padRight padBottom) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns) where
  type Tape (Pad padLeft padTop padRight padBottom) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns)  = ()
  runForwards Pad (S2D input) =
    let padl  = fromIntegral $ natVal (Proxy :: Proxy padLeft)
        padt  = fromIntegral $ natVal (Proxy :: Proxy padTop)
        padr  = fromIntegral $ natVal (Proxy :: Proxy padRight)
        padb  = fromIntegral $ natVal (Proxy :: Proxy padBottom)
        m     = extract input
        r     = diagBlock [konst 0 (padt,padl), m, konst 0 (padb,padr)]
    in  ((), S2D . fromJust . create $ r)
  runBackwards Pad _ (S2D dEdy) =
    let padl  = fromIntegral $ natVal (Proxy :: Proxy padLeft)
        padt  = fromIntegral $ natVal (Proxy :: Proxy padTop)
        nrows = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        ncols = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        m     = extract dEdy
        vs    = subMatrix (padt, padl) (nrows, ncols) m
    in  ((), S2D . fromJust . create $ vs)

-- | A two dimentional image can be padded.
instance ( KnownNat padLeft
         , KnownNat padTop
         , KnownNat padRight
         , KnownNat padBottom
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , KnownNat channels
         , KnownNat (inputRows * channels)
         , KnownNat (outputRows * channels)
         , (inputRows + padTop + padBottom) ~ outputRows
         , (inputColumns + padLeft + padRight) ~ outputColumns
         ) => Layer (Pad padLeft padTop padRight padBottom) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels) where
  type Tape (Pad padLeft padTop padRight padBottom) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels)  = ()
  runForwards Pad (S3D input) =
    let padl  = fromIntegral $ natVal (Proxy :: Proxy padLeft)
        padt  = fromIntegral $ natVal (Proxy :: Proxy padTop)
        padr  = fromIntegral $ natVal (Proxy :: Proxy padRight)
        padb  = fromIntegral $ natVal (Proxy :: Proxy padBottom)
        outr  = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        outc  = fromIntegral $ natVal (Proxy :: Proxy outputColumns)
        inr   = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        inc   = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        ch    = fromIntegral $ natVal (Proxy :: Proxy channels)
        m     = extract input
        padded = pad ch padl padt padr padb inr inc outr outc m
    in  ((), S3D . fromJust . create $ padded)

  runBackwards Pad () (S3D gradient) =
    let padl  = fromIntegral $ natVal (Proxy :: Proxy padLeft)
        padt  = fromIntegral $ natVal (Proxy :: Proxy padTop)
        padr  = fromIntegral $ natVal (Proxy :: Proxy padRight)
        padb  = fromIntegral $ natVal (Proxy :: Proxy padBottom)
        outr  = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        outc  = fromIntegral $ natVal (Proxy :: Proxy outputColumns)
        inr   = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        inc   = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        ch    = fromIntegral $ natVal (Proxy :: Proxy channels)
        m     = extract gradient
        cropped = crop ch padl padt padr padb inr inc outr outc m
    in  ((), S3D . fromJust . create $ cropped)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
module FourierMethod.FourierSeries2D
  ( module FourierMethod.FourierSeries2D
  , module FourierMethod.BlockCudaMatrix
  , module FourierMethod.BlockMatrixAcc
  ) where

import           Control.DeepSeq
import           Control.Monad                               as M
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Resource
import qualified Data.Array.Accelerate                       as A
import           Data.Array.Accelerate.LLVM.PTX              as A
import           Data.Array.Accelerate.Numeric.LinearAlgebra as A
import           Data.Array.Repa                             as R
import           Data.Complex
import           Data.Conduit                                as C
import           Data.Conduit.List                           as CL
import           Data.List                                   as L
import           Data.Vector.Storable                        as VS
import           Data.Vector.Unboxed                         as VU
import           Foreign.CUDA.Driver                         as CUDA
import           FourierMethod.BlockCudaMatrix
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2DAcc
import           Utils.BLAS
import           Utils.Distribution
import           Utils.List
import           Utils.Parallel                              hiding ((.|))
import           Utils.Time

{-# INLINE harmonicMatPTX #-}
harmonicMatPTX ::
     ( Storable e
     , A.Elt (Complex e)
     , A.Elt e
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => Int
  -> e
  -> e
  -> e
  -> PTX
  -> [(Int, Int)]
  -> CuMat (Complex e)
harmonicMatPTX numPoints period delta deltaFreq ptx r2Freqs =
  let rows = L.length r2Freqs
  in CuMat rows (numPoints ^ 2) .
     CuVecHost .
     VS.fromList .
     A.toList .
     run1With ptx (harmonicAcc numPoints period delta deltaFreq) .
     A.fromList (A.Z A.:. (L.length r2Freqs)) $
     r2Freqs

-- create column-major inverse harmonics
{-# INLINE inverseHarmonicMatPTX #-}
inverseHarmonicMatPTX ::
     ( Storable e
     , A.Elt (Complex e)
     , A.Elt e
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => Int
  -> e
  -> e
  -> PTX
  -> [(Int, Int)]
  -> CuMat (Complex e)
inverseHarmonicMatPTX numFreqs period delta ptx r2Positions =
  let cols = L.length r2Positions
  in CuMat (numFreqs ^ 2) cols .
     CuVecHost .
     VS.fromList .
     A.toList .
     run1With ptx (inverseHarmonicAcc numFreqs period delta) .
     A.fromList (A.Z A.:. (L.length r2Positions)) $
     r2Positions

{-# INLINE inverseHarmonicMatPTX1 #-}
inverseHarmonicMatPTX1 ::
     ( Storable e
     , A.Elt (Complex e)
     , A.Elt e
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => Int
  -> e
  -> e
  -> PTX
  -> [(Int, Int)]
  -> CuMat (Complex e)
inverseHarmonicMatPTX1 numFreqs period delta ptx r2Positions =
  let rows = L.length r2Positions
  in CuMat rows (numFreqs ^ 2) .
     CuVecHost .
     VS.fromList .
     A.toList .
     run1With ptx (inverseHarmonicAcc1 numFreqs period delta) .
     A.fromList (A.Z A.:. (L.length r2Positions)) $
     r2Positions

createHarmonicMatriesGPU ::
     ( Storable e
     , A.Elt (Complex e)
     , A.Elt e
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => [PTX]
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> IO [[CuMat (Complex e)]]
createHarmonicMatriesGPU ptxs numBatch numPoints numFreqs period delta deltaFreq = do
  let idxs =
        L.map (divideListN numBatch) . divideListN (L.length ptxs) $
        [ (xFreq, yFreq)
        | xFreq <- getListFromNumber numFreqs
        , yFreq <- getListFromNumber numFreqs
        ]
  let output =
        parZipWith
          rdeepseq
          (\ptx -> L.map (harmonicMatPTX numPoints period delta deltaFreq ptx))
          ptxs $
        idxs
  return output

createInverseHarmonicMatriesGPU ::
     ( Storable e
     , A.Elt (Complex e)
     , A.Elt e
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => [PTX]
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> IO [[CuMat (Complex e)]]
createInverseHarmonicMatriesGPU ptxs numBatch numPoints numFreqs period delta = do
  let idxs =
        L.map (divideListN numBatch) . divideListN (L.length ptxs) $
        [ (x, y)
        | x <- getListFromNumber numPoints
        , y <- getListFromNumber numPoints
        ]
  return .
    parZipWith
      rdeepseq
      (\ptx -> L.map (inverseHarmonicMatPTX numFreqs period delta ptx))
      ptxs $
    idxs

computeFourierCoefficientsR2 ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => [Int]
  -> [PTX]
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> Int
  -> [CuMat (Complex e)]
  -> IO (R.Array U DIM3 (Complex e))
computeFourierCoefficientsR2 deviceIDs ptxs numFreqs numPoints period delta deltaFreq numBatchR2Freqs xs = do
  let simpsonNorm = (delta / 3) ^ 2 :+ 0
      cols = L.foldl' (\s mat -> s + getColsCuMat mat) 0 xs
  harmonics <-
    createHarmonicMatriesGPU
      ptxs
      numBatchR2Freqs
      numPoints
      numFreqs
      period
      delta
      deltaFreq
  coefs <- blockMatrixMultiply2 True deviceIDs harmonics xs
  return .
    fromUnboxed (Z :. cols :. numFreqs :. numFreqs) .
    VU.map (* simpsonNorm) . VS.convert . getHostCuMat $
    coefs

computeFourierSeriesR2 ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => [Int]
  -> Int
  -> Int
  -> e
  -> [[CuMat (Complex e)]]
  -> [CuMat (Complex e)]
  -> IO (R.Array U DIM3 (Complex e))
computeFourierSeriesR2 deviceIDs numFreqs numPoints period inverseHarmonics xs = do
  let rows = L.foldl' (\s mat -> s + getRowsCuMat mat) 0 xs
  series <- blockMatrixMultiply1 False deviceIDs xs inverseHarmonics
  return .
    fromUnboxed (Z :. rows :. numPoints :. numPoints) .
    VS.convert . getHostCuMat $
    series

{-# INLINE applyGaussian #-}
applyGaussian ::
     (R.Source s (Complex e), Unbox e, RealFloat e)
  => e
  -> e
  -> R.Array s DIM3 (Complex e)
  -> R.Array D DIM3 (Complex e)
applyGaussian deltaFreq std arr =
  let (Z :. _ :. numFreq :. _) = extent arr
      freq = div numFreq 2
      gaussianCoefficients =
        computeUnboxedS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
          gaussian2D
            (deltaFreq * (fromIntegral $ xFreq - freq))
            (deltaFreq * (fromIntegral $ yFreq - freq))
            std :+
          0
  in R.traverse2 arr gaussianCoefficients const $ \fC fG idx@(Z :. _ :. i :. j) ->
       fC idx * fG (Z :. i :. j)

-- Stream
{-# INLINE indexSource #-}
indexSource :: Int -> Int -> ConduitT () [(Int, Int)] (ResourceT IO) ()
indexSource numIndex numBatch =
  let xs =
        divideListN numBatch $
        [ (xFreq, yFreq)
        | xFreq <- getListFromNumber numIndex
        , yFreq <- getListFromNumber numIndex
        ]
  in CL.sourceList xs

{-# INLINE fourierCoefficientsConduit #-}
fourierCoefficientsConduit ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     , Show e
     )
  => [Int]
  -> [PTX]
  -> Int
  -> e
  -> e
  -> e
  -> [CuMat (Complex e)]
  -> ConduitT [(Int, Int)] (CuMat (Complex e)) (ResourceT IO) ()
fourierCoefficientsConduit deviceIDs ptxs numPoints period delta deltaFreq xs =
  awaitForever $ \idx' -> do
    let idxs = divideListN (L.length deviceIDs) idx'
        harmonics =
          parMap
            rdeepseq
            (uncurry (harmonicMatPTX numPoints period delta deltaFreq)) $
          L.zip ptxs idxs
    liftIO $ printCurrentTime ""
    coefs <- liftIO $ blockMatrixMultiply3 False deviceIDs harmonics xs
    yield $!! coefs

{-# INLINE fourierCoefficientsSink #-}
fourierCoefficientsSink ::
     (Storable e, Unbox e, RealFloat e, Show e)
  => Int
  -> e
  -> ConduitT (CuMat (Complex e)) Void (ResourceT IO) ((R.Array U DIM3 (Complex e)))
fourierCoefficientsSink numR2Freqs delta = do
  xs <- CL.consume
  let cols = getColsCuMat . L.head $ xs
      simpsonNorm = (delta / 3) ^ 2 :+ 0
  return .
    fromUnboxed (Z :. cols :. numR2Freqs :. numR2Freqs) .
    VU.map (* simpsonNorm) .
    VS.convert . getHostCuMat . transposeCuMat . concatCuMat $
    xs

computeFourierCoefficientsR2Stream ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => [Int]
  -> [PTX]
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> Int
  -> [CuMat (Complex e)]
  -> IO (R.Array U DIM3 (Complex e))
computeFourierCoefficientsR2Stream deviceIDs ptxs numFreqs numPoints period delta deltaFreq numBatch xs =
  runConduitRes $
  indexSource numFreqs numBatch .|
  fourierCoefficientsConduit deviceIDs ptxs numPoints period delta deltaFreq xs .|
  fourierCoefficientsSink numFreqs delta

{-# INLINE fourierSeriesConduit #-}
fourierSeriesConduit ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     , Show e
     )
  => [Int]
  -> [PTX]
  -> Int
  -> e
  -> e
  -> [CuMat (Complex e)]
  -> ConduitT [(Int, Int)] (CuMat (Complex e)) (ResourceT IO) ()
fourierSeriesConduit deviceIDs ptxs numFreqs period delta xs =
  awaitForever $ \idx' -> do
    let idxs = divideListN (L.length deviceIDs) idx'
        harmonics =
          parMap
            rdeepseq
            (uncurry (inverseHarmonicMatPTX1 numFreqs period delta)) $
          L.zip ptxs idxs
    liftIO $ printCurrentTime "fourierSeriesConduit"
    coefs <- liftIO $ blockMatrixMultiply3 False deviceIDs harmonics xs
    yield $!! coefs

{-# INLINE fourierSeriesSink #-}
fourierSeriesSink ::
     (Storable e, Unbox e, RealFloat e, Show e)
  => Int
  -> ConduitT (CuMat (Complex e)) Void (ResourceT IO) ((R.Array U DIM3 (Complex e)))
fourierSeriesSink numPoints = do
  xs <- CL.consume
  let cols = getColsCuMat . L.head $ xs
  return .
    fromUnboxed (Z :. cols :. numPoints :. numPoints) .
    VS.convert . getHostCuMat . transposeCuMat . concatCuMat $
    xs

computeFourierSeriesR2Stream ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     )
  => [Int]
  -> [PTX]
  -> Int
  -> Int
  -> e
  -> e
  -> Int
  -> [CuMat (Complex e)]
  -> IO (R.Array U DIM3 (Complex e))
computeFourierSeriesR2Stream deviceIDs ptxs numFreqs numPoints period delta numBatch xs =
  runConduitRes $
    indexSource numPoints numBatch .|
    fourierSeriesConduit deviceIDs ptxs numFreqs period delta xs .|
    fourierSeriesSink numPoints

{-# INLINE fourierSeriesConduitAcc #-}
fourierSeriesConduitAcc ::
     ( A.Floating e
     , A.Elt (Complex e)
     , A.FromIntegral Int e
     , Numeric (Complex e)
     , Unbox e
     )
  => [PTX]
  -> Int
  -> e
  -> e
  -> A.Acc (Matrix (Complex e))
  -> ConduitT [(Int, Int)] (VU.Vector (Complex e)) (ResourceT IO) ()
fourierSeriesConduitAcc ptxs numFreqs period delta x =
  awaitForever $ \idx' -> do
    let idxs = divideListN (L.length ptxs) idx'
        harmonics =
          L.map
            (\r2Positions ->
               inverseHarmonicAcc2 numFreqs period delta .
               A.fromList (A.Z A.:. L.length r2Positions) $
               r2Positions) $
          idxs
    liftIO $ printCurrentTime "fourierSeriesConduit"
    yield $!! blockMatrixMultiply ptxs harmonics x

{-# INLINE fourierSeriesSinkAcc #-}
fourierSeriesSinkAcc ::
     (Unbox e)
  => Int
  -> Int
  -> ConduitT (VU.Vector (Complex e)) Void (ResourceT IO) (R.Array U DIM3 (Complex e))
fourierSeriesSinkAcc numPoints cols = do
  xs <- CL.consume
  computeP .
    R.backpermute
      (Z :. cols :. numPoints :. numPoints)
      (\(Z :. a :. b :. c) -> Z :. b :. c :. a) .
    fromUnboxed (Z :. numPoints :. numPoints :. cols) . VU.concat $
    xs

computeFourierSeriesR2StreamAcc ::
     ( A.Floating e
     , A.Elt (Complex e)
     , A.FromIntegral Int e
     , Numeric (Complex e)
     , Unbox e
     )
  => [PTX]
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> Int
  -> Acc (Matrix (Complex e))
  -> IO (R.Array U DIM3 (Complex e))
computeFourierSeriesR2StreamAcc ptxs numFreqs numPoints cols period delta numBatch x =
  runConduitRes $
  indexSource numPoints numBatch .|
  fourierSeriesConduitAcc ptxs numFreqs period delta (A.compute x) .|
  fourierSeriesSinkAcc numPoints cols

{-# INLINE fourierSeriesConduitAcc' #-}
fourierSeriesConduitAcc' ::
     ( A.Floating e
     , A.Elt (Complex e)
     , A.FromIntegral Int e
     , Numeric (Complex e)
     , Unbox e
     , Prelude.Num e
     )
  => [PTX]
  -> Int
  -> e
  -> e
  -> A.Acc (Matrix (Complex e))
  -> ConduitT [(Int, Int)] (VU.Vector (Complex e)) (ResourceT IO) ()
fourierSeriesConduitAcc' ptxs numFreqs period delta x =
  awaitForever $ \idx' -> do
    let idxs = divideListN (L.length ptxs) idx'
        harmonics =
          L.map
            (\r2Positions ->
               inverseHarmonicAcc2' numFreqs period delta .
               A.fromList (A.Z A.:. L.length r2Positions) $
               r2Positions) $
          idxs
    liftIO $ printCurrentTime "fourierSeriesConduit"
    yield $ blockMatrixMultiply ptxs harmonics x

computeFourierSeriesR2StreamAcc' ::
     ( A.Floating e
     , A.Elt (Complex e)
     , A.FromIntegral Int e
     , Numeric (Complex e)
     , Unbox e
     , Prelude.Num e
     )
  => [PTX]
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> Int
  -> Acc (Matrix (Complex e))
  -> IO (R.Array U DIM3 (Complex e))
computeFourierSeriesR2StreamAcc' ptxs numFreqs numPoints cols period delta numBatch x =
  runConduitRes $
  indexSource numPoints numBatch .|
  fourierSeriesConduitAcc' ptxs numFreqs period delta (A.compute x) .|
  fourierSeriesSinkAcc numPoints cols


-- 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


{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}

module HVX.Primitives
-- TODO(mh): Currently primitives that should generate implicit constraints are
-- unsupported. The should be supported shortily. (2014-06-04)
  ( apply
  , hadd
  , (+~)
  , hmul
  , (*~)
  , habs
  , neg
  , hlog
  , hexp
  , logsumexp
  , hmax
  , hmin
  , norm
  , berhu
  , huber
  , quadform
  , powBaseP0
  , powBaseP01
  , powBaseP1
  , powBaseP1InfEven
  , powBaseP1InfNotInt
  ) where

import Numeric.LinearAlgebra hiding (i)
import Numeric.LinearAlgebra.LAPACK

import HVX.Internal.DCP
import HVX.Internal.Matrix
import HVX.Internal.Primitives
import HVX.Internal.Util

apply :: (Vex vf, Mon mf, Vex ve, Mon me, Vex (ApplyVex vf mf ve me), Mon (ApplyMon mf me))
  => Fun vf mf -> Expr ve me -> Expr (ApplyVex vf mf ve me) (ApplyMon mf me)
apply = EFun

hadd :: (Vex v1, Mon m1, Vex v2, Mon m2, Vex (AddVex v1 v2), Mon (AddMon m1 m2))
  => Expr v1 m1 -> Expr v2 m2 -> Expr (AddVex v1 v2) (AddMon m1 m2)
hadd = EAdd

infixl 6 +~
(+~) :: (Vex v1, Mon m1, Vex v2, Mon m2, Vex (AddVex v1 v2), Mon (AddMon m1 m2))
  => Expr v1 m1 -> Expr v2 m2 -> Expr (AddVex v1 v2) (AddMon m1 m2)
(+~) = hadd

-- Constructors that enforce DCP constraints.
hmul :: (Vex v1, Mon m1, Vex (ApplyVex Affine Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Expr Affine Const -> Expr v1 m1 -> Expr (ApplyVex Affine Nonmon v1 m1) (ApplyMon Nonmon m1)
hmul (EConst a) e = apply (Mul a) e
hmul _ _ = error "the left argument of a multiply must be a constant"

infixl 7 *~
(*~) :: (Vex v1, Mon m1, Vex (ApplyVex Affine Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Expr Affine Const -> Expr v1 m1 -> Expr (ApplyVex Affine Nonmon v1 m1) (ApplyMon Nonmon m1)
(*~) (EConst a) e = apply (Mul a) e
(*~) _ _ = error "the left argument of a multiply must be a constant"

habs :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
habs = apply Abs

neg :: (Vex v1, Mon m1, Vex (ApplyVex Affine Noninc v1 m1), Mon (ApplyMon Noninc m1))
  => Expr v1 m1 -> Expr (ApplyVex Affine Noninc v1 m1) (ApplyMon Noninc m1)
neg = apply Neg

hlog :: (Vex v1, Mon m1, Vex (ApplyVex Concave Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Concave Nondec v1 m1) (ApplyMon Nondec m1)
hlog = apply Log

hexp :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Convex Nondec v1 m1) (ApplyMon Nondec m1)
hexp = apply Exp

logsumexp :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Convex Nondec v1 m1) (ApplyMon Nondec m1)
logsumexp = apply LogSumExp

hmax :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Convex Nondec v1 m1) (ApplyMon Nondec m1)
hmax = apply Max

hmin :: (Vex v1, Mon m1, Vex (ApplyVex Concave Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Concave Nondec v1 m1) (ApplyMon Nondec m1)
hmin = apply Min

norm :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
norm p
  | p == infinity = error "Internal: Infinity norm should become max . abs."
  | 1 <= p = apply (Norm p)
  | otherwise = error "Internal: Norm only supports p >= 1."

berhu :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
berhu m
  | 0 < m = apply (Berhu m)
  | otherwise = error "Internal: Berhu only supports m >= 0."

huber :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
huber m
  | 0 < m = apply (Huber m)
  | otherwise = error "Internal: Huber only supports m >= 0."

quadform :: (Mon m1, Vex (ApplyVex Convex Nonmon Affine m1), Mon (ApplyMon Nonmon m1))
  => Expr Affine Const -> Expr Affine m1 -> Expr (ApplyVex Convex Nonmon Affine m1) (ApplyMon Nonmon m1)
quadform (EConst a) e
  | rows a == cols a
    && fpequalsMat a (trans a)
    && 0 <= maxElement (eigOnlyS a) = apply (Quadform a) e
  | otherwise = error "Matrices in quadratic forms must be positive semidefinite."
quadform _ _ = error "The matrix sandwitched by the quadratic form must be constant."

powBaseP0 :: (Vex v1, Mon m1, Vex (ApplyVex Affine Const v1 m1), Mon (ApplyMon Const m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Affine Const v1 m1) (ApplyMon Const m1)
powBaseP0 p
  | p `fpequals` 0 = apply PowBaseP0
  | otherwise = error "Internal: PowBaseP0 only supports p == 0."

powBaseP01 :: (Vex v1, Mon m1, Vex (ApplyVex Concave Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Concave Nondec v1 m1) (ApplyMon Nondec m1)
powBaseP01 p
  | 0 < p && p < 1 = apply (PowBaseP01 p)
  | otherwise = error "Internal: PowBaseP01 only supports 0 < p < 1."

powBaseP1 :: (Vex v1, Mon m1, Vex (ApplyVex Affine Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Affine Nondec v1 m1) (ApplyMon Nondec m1)
powBaseP1 p
  | p `fpequals` 1 = apply PowBaseP1
  | otherwise = error "Internal: PowBaseP1 only supports p == 1."

powBaseP1InfEven :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
powBaseP1InfEven p
  | 1 < p && isInteger p && even intP = apply (PowBaseP1InfEven intP)
  | otherwise = error "Internal: PowBaseP1InfEven only supports even p > 1."
  where intP = round p :: Integer

powBaseP1InfNotInt :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nondec v1 m1) (ApplyMon Nondec m1)
powBaseP1InfNotInt p
  | 1 < p && not (isInteger p) = apply (PowBaseP1InfNotInt p)
  | otherwise = error "Internal: PowBaseP1InfNotInt only supports non integral p > 1."