module Main(main) where

import System.Environment (getArgs)
import Codec.Picture (generateImage, savePngImage, DynamicImage(ImageRGB8), PixelRGB8(PixelRGB8), pixelAt, colorMap, mixWith)
import Numeric.LinearAlgebra.Data ((|>))
import Raytracer.Camera (Camera(Camera), Point(Point), fire_ray, calculate_ray)
import Raytracer.Geometry (cube, square, dist, Collision(Collision), calc_ray)
import Raytracer.Light (Light(Light), computeLighting)
import Data.Monoid ((<>))
import Data.Maybe (isJust)

import Control.Parallel.Strategies (rdeepseq, parMap)
import Control.Concurrent (getNumCapabilities)

help = putStrLn
  "Ray traces an image from a simple scene\n\
  \Run: width height step cx cy cz dx dy dz\n\
  \Where c_ is camera position, d_ is camera direction\n\
  \width and height the size of the image, and step is how many pixels exist between width and height\n\
  \For example, w=200 h=200 step=1 will be a 200x200 image, 200x200x2 will be 100x100, but have the same frame"

test_camera wres hres = Camera 4 3 wres hres (Just 3)
test_cube = cube (PixelRGB8 255 0 255) (3 |> [2, 0, 0]) (3 |> [0, 2, 0]) (3 |> [0, 0, 2]) (3 |> [0, 0, 0])
test_floor = square (PixelRGB8 255 255 255) (3 |> [8,0,0]) (3 |> [0,0,8]) (3 |> [-2,0,-2])
test_mesh = test_cube <> test_floor

test_lights = [
  Light (3|> [-1, 3, 1]) 6 $ PixelRGB8 255 255 255,
  Light (3|> [1, 4, 1]) 5 $ PixelRGB8 0 0 255,
  Light (3|> [6, 0, 0]) 2 $ PixelRGB8 255 0 0,
  Light (3|> [6, 0, 2]) 2 $ PixelRGB8 0 255 0,
  Light (3|> [6, 0, 4]) 2 $ PixelRGB8 0 0 255
  ]

-- This mixes light colours additively
mixColours = foldr (mixWith add) (PixelRGB8 20 20 20)
  where
  add _ c1 c2 = if maxBound - c1 >= c2 then c1 + c2 else maxBound

-- This mixes a pigment with a light subtractively
reflectLight :: PixelRGB8 -> PixelRGB8 -> PixelRGB8
reflectLight surface light = mixWith subColour surface light
  where
  subColour _ cs cl = if cl >= (maxBound - cs) then cl - (maxBound - cs) else minBound

renderPixel camera x y = computePixel $ fire_ray test_mesh ray
  where
  ray = calculate_ray camera $ Point x y
  lights = computeLighting test_mesh test_lights
  computePixel Nothing = PixelRGB8 0 0 0
  computePixel (Just (Collision d c)) = reflectLight c $ mixColours $ lights $ calc_ray ray d

verticalFlip h func camera x y = func camera x (h - y)

parGenerateImage 1 func w h = generateImage func w h
parGenerateImage n func w h = generateImage combine w h
  where
  -- combine rotates through the parallel images putting lines together
  combine x y = pixelAt (images !! (y `mod` n)) x $ y `div` n
  images = parMap rdeepseq generate [0..n-1]
  generate offset = generateImage (\x y -> func x $ y*n + offset) w $ heightFor offset
  -- heightFor offset exists for when the height is not evenly divisible by n
  -- We need to add a couple rows to the earlier images to make up the difference
  heightFor offset = h `div` n + (if offset < h `mod` n then 1 else 0)

main = do
  args <- getArgs
  if length args /= 9 then
    help
  else do
    numCores <- getNumCapabilities
    let (width:height:step:cx:cy:cz:dx:dy:dz:[]) = args
    let w' = (read width) `div` (read step)
    let h' = (read height) `div` (read step)
    let camera = test_camera w' h' (3 |> [read cx, read cy, read cz]) (3 |> [read dx, read dy, read dz])
    let img = parGenerateImage numCores (verticalFlip h' renderPixel camera) w' h'
    savePngImage "test.png" $ ImageRGB8 img


{-# LANGUAGE OverloadedStrings #-}

module Simulation where

import DigitalAGC
import LutLogTests (sfToDouble, ufToDouble)

import Clash.Prelude (
  UFixed, Unsigned, Signed, SFixed(..), Vec(..), Nat, KnownDomain(..), HiddenClockResetEnable, SNat(..),
  createDomain, vInitBehavior, vName, vSystem, InitBehavior(..), Signal,
  simulate, System,
  shiftR,
  register,
  sf, bundle, df,
  riseEvery,
  fLitR, KnownNat
  )

import Prelude
import Data.Maybe (fromMaybe)
import Data.List (group)

import Lucid
import Lucid.Html5
import Graphics.Plotly hiding (sort)
import qualified Graphics.Plotly.Base as GPB
import Graphics.Plotly.Simple
import Graphics.Plotly.Lucid
import Lens.Micro
import System.Process
import Data.Array
import Data.Foldable

import qualified Data.Complex as DC
import Numeric.Transform.Fourier.FFT
import DSP.Basic (linspace)

import DataFlow.Extra
import Test.QuickCheck

import Statistics.Sample
import qualified Data.Vector.Unboxed as V
--import Data.Fixed

--dec :: Int -> [a] -> [a]
--dec n [] = []
--dec n (x : xs) = x : dec n (drop (n-1) xs)

interp :: Int -> [a] -> [a]
interp n [] = []
interp n (x : xs) = replicate n x ++ interp n xs

chunks :: Int -> [a] -> [[a]]
chunks n [] = []
chunks n xs = take n xs : chunks n (drop n xs)

quantiseS :: forall i f . (KnownNat i, KnownNat f) => SNat i -> SNat f -> Double -> Double
quantiseS _ _ x = sfToDouble (fLitR x :: SFixed i f)

quantiseU :: forall i f . (KnownNat i, KnownNat f) => SNat i -> SNat f -> Double -> Double
quantiseU _ _ x = ufToDouble (fLitR x :: UFixed i f)

d0 = SNat :: SNat 0

iSig   = SNat :: SNat 26
iGain  = SNat :: SNat 10
fGain  = SNat :: SNat 10
iAlpha = SNat :: SNat 1
fAlpha = SNat :: SNat 6
iRef   = SNat :: SNat 3
fRef   = SNat :: SNat 12
iLog   = SNat :: SNat 4
fLog   = SNat :: SNat 22

simAgc :: Double -> Double -> Double -> [(Double, Double)] -> [(Double, Double,Double)]
simAgc window ref alpha iqs =
  let window' = 2**window
      powers blk = map (\(i,q)-> quantiseU iSig d0 . sqrt $ i^2 + q^2) blk
      smooth blk = quantiseU iSig d0 $ (/window') $ sum  blk
      alg state x = quantiseS iLog fLog (
                      (quantiseS iLog fLog state) -
                      ((quantiseS iLog fLog $ logBase 10 (smooth $ powers x)) - (quantiseU iRef fRef ref)) *
                      (quantiseU iAlpha fAlpha alpha))
      gains = scanl (\state x ->
                quantiseS iLog fLog $
                alg state (map (\(i,q)->(
                  quantiseS iSig d0 $ (i*) $ quantiseU iGain fGain $ 10**state,
                  quantiseS iSig d0 $ (q*) $ quantiseU iGain fGain $ 10**state)
                  ) x)
              ) 0
              (init $ chunks (round window') iqs)
  in zipWith (\(i,q) s -> let g = quantiseU iGain fGain $ 10**s
                          in (g, quantiseS iSig d0 $ i*g, quantiseS iSig d0 $ q*g))
     iqs (interp (round window') gains)

-- My floating point simulation shows us that we _should_ be ok with our
-- approach... make some tests to verify the log and exp DF behaviour.

createDomain vSystem{vName="SyncDefined", vInitBehavior=Defined}

simInput = let iqs = map (\(i,q)->(shiftR i 1, shiftR q 1)) $ sinInputComplex 1 0.01
           in  take 3000 iqs ++ map (\(i,q)->(shiftR i 4, shiftR q 4)) iqs

-- from 1.25 to 0.75... our ref is basically divided by 4, so this makes sense for a 100x range
-- we have 0.5 range with 6 frac bits => resolution of 1/64, so 32 sensible settings for our signal height.
-- This is all with alpha = 0.9375 and a window of 7
-- Window doesn't really affect this too much
-- Alpha seems to?

-- we have a 2.6 log word ( really 4.4)

-- a = 0.9375 and window = 128 ~ 1.5k         1.5k
-- a = 0.5                       3k           3.5k
-- a = 0.25                      4.5k         5k
-- a = 0.125                     7k           8k
-- a = 0.0625                                 10k

--   alpha             recovery samples   diff      log10  log diff
--        1; 5         2                  ----      0.3    --------
--      0.5; 7.3       5.3                3.3       0.72   0.42
--     0.25; 12        9                  3.7       0.95   0.23
--    0.125; 20        17                 8         1.23   0.27
--   0.0625; 33        31                 14        1.49   0.26

-- ref = 1.13 for 32k
-- for 1000, log 10
--sim = let ref = 1.05 :: UFixed 2 10
--          window = 7 :: Unsigned 5
--          alpha = 2.0 :: UFixed 1 6
--          fLog = Clash.d6
--          fGain = Clash.d6
--          out_gain = Clash.simulate @System (uncurry (digiAgcMult (pure window) (pure ref) (pure alpha)). unbundle) simInput
--      in map (zip [1..]) [
--                           map (fromIntegral . (\(_,x,_)->x)) out_gain
--                         , map (fromIntegral . (\(_,_,x)->x)) out_gain
--                         , map (ufToDouble   . (\(x,_,_)->x)) out_gain
--                         ]

{-
Just thinking about sensible wordlengths for RFSoC...

Our input is always 16 bits.

==========

This will likely be after decimation, used inside the user's RX digital logic so we can expect deal with sampling rates well below the 4 GSamples/s... say probably a max of 512 MHz.

How many window bits are needed for a period of 1 ms?

1e-3 * 512e6 = 18.9 ... let's just say 19

==========

Fullscale signal after log will be...

after ID; we need (nWindow + nSig) bits (35!!!)
after log10; we need (1+log2 35) bits (9) and fLog fractional bits, let's say 9.fLog
after sub with ref; we need 10.fLog
after mul with alpha we need (10+ia).(fLog+fa)

{ What should alpha's range be to allow a good selection of response times? }

After the antilog, our wordlength increases exponentially! We should be super
careful about growing our wordlengths in the log domain.

==========

nSig = 16
nWindow = 19
-}

-- Experimentally found recovery cycles for every tenth decimal step in alpha
-- alphas = [0.1,0.2,...1.9]
recoveryCycles :: [(Int, Int)]
recoveryCycles = [(1,56),(2,40),(3,31),(4,26),(5,22),(6,19),(7,17),(8,15),(9,14),(10,13),(11,12),(12,11),(13,10),(14,10),(15,9),(16,8),(17,8),(18,7),(19,7)]

getRecoveryCycles :: Double -> Int
getRecoveryCycles = fromMaybe (error "No entry found for that alpha") . flip lookup recoveryCycles. round . (*10)

newtype InGain = InGain Double deriving Show
newtype InStepTime = InStepTime Int deriving Show

instance Arbitrary InStepTime where
  arbitrary = fmap InStepTime $ choose (1000,10000)

instance Arbitrary InGain where
  arbitrary = fmap (InGain . recip . fromIntegral) $ choose (1::Int,63)

steppedInput :: InGain -> InGain -> InStepTime -> [(Signed 16,Signed 16)]
steppedInput (InGain g1) (InGain g2) (InStepTime n) =
  let a = map (\(i,q)->(multD i g1, multD q g1)) . take n $ sinInputComplex 1 0.01
      b = map (\(i,q)->(multD i g2, multD q g2))           $ sinInputComplex 1 0.01
  in a ++ b
  where
  multD s d = fromIntegral . round $ fromIntegral s * d

splitInto n [] = []
splitInto n xs = let (a,b) = splitAt n xs
                 in a : splitInto n b

isSteady :: Int -> Double -> Double -> [Double] -> Bool
isSteady n ref percent = (<=n) . maximum . map length . filter (\a->False == a!!0) . group . map (\x->abs (x-ref)/ref < (percent/100))

simDfLogErr ref alpha xs =
  simulate @System (\x ->
    bundle $ df (dfLogErr (pure ref) (pure alpha)) x (riseEvery (SNat :: SNat 3)) (pure True) )
  (cycle xs) :: [(SFixed 4 22, Bool, Bool)]

{--

1) Can I reproduce the variation in simulation? Maybe generating a repeating pattern with a non-integer multiple of the window size

2) Does this also appear in the floating point simulation, because that doesn't have the delays from cordic stuff.

If no to 1) maybe it's to do with my DataFlow implementation?

I'm seeing spikes for 32 cycles in the output. Why is this?! We should only be seeing 512 sample changes because of the window of 2^9

I've fixed it... Two things now:

  1) Are the log10 and pow10 df units *really* returning consistient values?

  2) We can add a feature to set the error state! We can use that in the python instead of feeding it extra samples. Is there anything else we should flush? Maybe just pad with 32ish samples. Wait. isn't that just the reset functionality that should already exist? I've just mapped the areset pin to axi reg 0, bit 1.

--}

simOutPower g1 g2 n =
  let ref = 4.0 :: UFixed 3 12
      alpha = 1.0 :: UFixed 1 6
      window = 9 :: Unsigned 5
      --inputSig = take (10000 + rec_time*(2^window)) $ steppedInput g1 g2 n
      inputSig = take (6000*5) $ steppedInput g1 g2 n
      rec_time = (2+) . getRecoveryCycles $ ufToDouble alpha
      --out_gain = take 15000 $ simAgc (fromIntegral window) (ufToDouble ref) (ufToDouble alpha) (map (\(i,q)->(fromIntegral i, fromIntegral q)) inputSig)
      --out_pow = map (\(_,i,q)-> sqrt $ (i)**2 + (q)**2) out_gain
      ip x =
               let ip = df (testBufferDF (SNat :: SNat 6000) `seqDF` throttleDF (SNat :: SNat 4) `seqDF` dfAgc (pure window) (pure ref) (pure alpha) (pure True))
                   oR = (riseEvery (SNat :: SNat 3))
                   (y, oV, iR) = ip x (pure True) oR
               in bundle (y, oV, oR) :: Signal System ((UFixed 10 15, (Signed 16, Signed 16)), Bool, Bool)
      --outs = drop 1 . take (10000 + rec_time*(2^window))
      outs = drop 1 . take 6000 . filter (\(_,v,r)->v&&r)
             $ simulate @System ip inputSig
      out_gain = map (\((g,(i,q)), v,r)->(g,i,q)) outs
      out_pow = map (\(_,i,q)-> sqrt $ (fromIntegral i)**2 + (fromIntegral q)**2) out_gain
      out_pow_block = map ((/(2^window)) . sum) $ splitInto (2^window) out_pow
      expected_pow = 10**(ufToDouble ref)
  in (inputSig, out_gain, out_pow, out_pow_block, rec_time, expected_pow)

prop_OutPower :: InGain -> InGain -> InStepTime -> Property
prop_OutPower g1 g2 (InStepTime n) =
  let (_, _, _, out_pow_block, rec_time, expected_pow) = simOutPower g1 g2 (InStepTime n)
  in property . isSteady rec_time expected_pow 1 $ drop (ceiling $ (fromIntegral n) / 2**7) out_pow_block

showTest g1 g2 n =
    renderToFile "/tmp/clash/sim.html" $ doctypehtml_ $ do
    head_ $ do meta_ [charset_ "utf-8"]
               plotlyCDN
               reloadCDN
               styleSheet
    body_ $ do
               toHtml $ plotly "time_iq_in" (traceTime simInputTrace)
                          & layout . title ?~ "Time domain I/Q Input"
               toHtml $ plotly "time_iq_out" (traceTime tData)
                          & layout . title ?~ "Time domain I/Q Output"
               toHtml $ plotly "constl_iq_out" (traceConstl tData)
                          & layout . title ?~ "Constellation I/Q"
                          & layout . width ?~ 600
                          & layout . height ?~ 600
               toHtml $ plotly "time_iq_ctrl" (traceTime ctrlData)
                          & layout . title ?~ "Time domain Control"
  where
  (simInput, sim, out_pow, _, _, _) = simOutPower g1 g2 n
  tData = map (zip [1..]) $ [map (\(_,x,_)->fromIntegral x) sim
                            ,map (\(_,_,x)->fromIntegral x) sim
                            ,map (\(x,_,_)->ufToDouble x)   sim
                            --[map (\(_,x,_)-> x) sim
                            --,map (\(_,_,x)-> x) sim
                            --,map (\(x,_,_)->x)   sim
                            ] :: [[(Double, Double)]]
  ctrlData = [tData !! 2, tData !! 2]
  simInputTrace = map (zip [(1::Double)..] . map (sfToDouble . sf (SNat :: SNat 0)))  [map fst simInput, map snd simInput]

corr :: [(Double, Double)] -> Double
corr = correlation . V.fromList

--prop_corrTest g1 g2 n =
--  let (_, outsig, _, _, recTime, _) = simOutPower g1 g2 n
--      insig = steppedInput (InGain 0.3125) (InGain 0.3125) (InStepTime 0)
--      inI  = map (fromIntegral . fst) insig
--      inQ  = map (fromIntegral . snd) insig
--      outI = map (fromIntegral . (\(_,i,_)->i)) outsig
--      outQ = map (fromIntegral . (\(_,_,q)->q)) outsig
--      correlation = corr $ zip (drop recTime inI) (drop recTime outI)
--  in property $ correlation > 0.9

bl10 n = fromIntegral . ceiling $ logBase 2 (logBase 10 (2**n))
be10 n = fromIntegral . ceiling $ logBase 2 (10 ** (2**n))

bitsCalc :: Double -> Double -> Double -> (Double, Double) -> Double
         -> (Double, (Double, Double), (Double, Double))
bitsCalc window sig fLog (iAlpha, fAlpha) fGain =
  let intgInternal = 2**window + sig
      iLog = bl10 sig
      iGain = be10 iLog - fGain
  in (intgInternal, (iLog,fLog), (iGain, fGain))

sinInput :: Double -> Double -> [Signed 16]
sinInput fs dt = map (fromInteger . round . ((2**15-1) * ) .  sin . (2*pi*fs * )) [0.0, dt..]

sinInputComplex :: Double -> Double -> [(Signed 16, Signed 16)]
sinInputComplex fs dt = zip sins coss
  where sins = map (fromInteger . round . ((2**15-1) * ) .  sin . (2*pi*fs * )) [0.0, dt..]
        coss = map (fromInteger . round . ((2**15-1) * ) .  cos . (2*pi*fs * )) [0.0, dt..]

reloadCDN :: Monad m => HtmlT m ()
reloadCDN = script_ [src_ "http://livejs.com/live.js"] $ toHtml (""::String)

styleSheet :: Monad m => HtmlT m ()
styleSheet = style_ [type_ "text/css" ] $
             toHtml (".svg-container {margin: 0 auto !important;}"::String)

traceTime a =  [trace "I" $ a !! 0
               ,trace "Q" $ a !! 1]
  where trace label points = linePlot points
                               & name ?~ label
                               & (GPB.line ?~ (defLine & lineshape ?~ Hv))

traceConstl points = [scatterPlot $ zip (map snd (points!!0)) (map snd (points!!1))]

startServer = callCommand "mkdir -p /tmp/clash; cd /tmp/clash/; python -m http.server > /dev/null 2>&1 &"

return []
runTests = $quickCheckAll

{- TODO

Include bin files in IP make script

Try find what resolution we need to increase in order to better recover very very low signals.

-}


{-# LANGUAGE ScopedTypeVariables, ConstraintKinds, GADTs #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE UndecidableInstances #-} -- See below
{-# LANGUAGE TupleSections #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}
{-# OPTIONS_GHC -fno-warn-unused-binds   #-}

----------------------------------------------------------------------
-- |
-- Module      :  Generic.FFT
-- Copyright   :  (c) 2013 Tabula, Inc.
--
-- Maintainer  :  conal@tabula.com
-- Stability   :  experimental
--
-- FFT via functor composition. Apply to depth-typed perfect binary leaf trees.
----------------------------------------------------------------------
module Generic.FFT
  ( dft
  , idft
  , dct
  , HasFFT(..)
  ) where

import Prelude hiding (sum)

import Control.Applicative (Applicative(..), liftA2)
import Control.Arrow ((***))
import Data.Foldable (Foldable, sum)
import Data.Functor ((<$>))
import Data.Monoid (Monoid(..), Sum(..), (<>))
import Data.Traversable (Traversable(..), mapAccumL)

import Data.Complex (Complex(..), conjugate, realPart)
import Data.Tuple (swap)

import Data.UniformPair

import TypeUnary.Nat

import qualified Data.FTree.BottomUp as B
import qualified Data.FTree.TopDown as T

{--------------------------------------------------------------------
    Misc
--------------------------------------------------------------------}
type Unop a = a -> a

transpose :: (Traversable g, Applicative f) => g (f a) -> f (g a)
transpose = sequenceA

inTranspose ::
     (Traversable f, Traversable k, Applicative g, Applicative h)
  => (g (f a) -> k (h b))
  -> (f (g a) -> h (k b))
inTranspose = transpose --> transpose

infixr 1 -->

-- | Add pre- and post processing
(-->) :: (a' -> a) -> (b -> b') -> ((a -> b) -> (a' -> b'))
(f --> h) g = h . g . f

type R = Double
type C = Complex R

scanL :: (Traversable f, Monoid a) => (a, f a) -> (f a, a)
scanL (a0, as) = swap (mapAccumL h a0 as)
  where
    h a a' = (a <> a', a)

-- TODO: Replace scanL with my efficient parallel version.
-- Prefix (left) sums
sumsL :: (Traversable f, Num a) => (a, f a) -> (f a, a)
sumsL = (fmap getSum *** getSum) . scanL . (Sum *** fmap Sum)

-- Yield a structure counting from 0 to size-1, together with size
counts ::
     forall f a. (TA f, Num a)
  => (f a, a)
counts = sumsL (0, pure 1)

-- Cross product of structures
cross :: (Functor g, Functor f) => g a -> f b -> g (f (a, b))
as `cross` bs = fmap (\a -> fmap (a, ) bs) as

-- All products of numbers from each structure.
products :: (Functor g, Functor f, Num n) => g n -> f n -> g (f n)
products = (fmap . fmap . fmap . fmap) (uncurry (*)) cross

-- as `products` bs = (fmap.fmap) (uncurry (*)) (as `cross` bs)
-- as `products` bs = fmap (\ a -> fmap (\ b -> a*b) bs) as
-- Dot product of structures. Assumes a trie-like f, so that the Applicative
-- instance combines corresponding elements. Perhaps replace Applicative with a
-- more suitable constraint.
dot :: (Applicative f, Foldable f, Num a) => f a -> f a -> a
u `dot` v = sum (liftA2 (*) u v)

i2pi :: C
i2pi = 0 :+ 2 * pi

-- Principle nth root of unity with negated angle
uroot :: Int -> C
uroot n = exp (-i2pi / fromIntegral n)

{--------------------------------------------------------------------
    Unoptimized discrete Fourier transform (DFT)
--------------------------------------------------------------------}
-- Discrete Fourier transform:
-- $X_k = \sum_{n=0}^{N-1} x_n e^{-i 2\pi k \frac{n}{N}}$ for $k = 0,\ldots,N$.
dft :: TA f => Unop (f C)
dft xs = (xs `dot`) <$> rootses

idft ::
     forall f. TA f
  => Unop (f C)
idft = fmap ((/ fromIntegral n) . conjugate) . dft . fmap conjugate
  where
    n :: Int
    indices :: f Int
    (indices, n) = counts

-- Inefficinet; doesn't exploit symmetry, but a starting point for testing
dct :: TA f => Unop (f R)
dct = fmap realPart . dft . fmap (:+ 0)

-- Powers of 'uroot' needed in the DFT:
-- $e^{\frac{-i 2\pi k n}{N}}$ for $k,n = 0,\ldots,N$:
rootses ::
     forall f. TA f
  => f (f C)
rootses = rootCross tot indices indices
  where
    indices :: f Int
    (indices, tot) = counts

rootCross :: (Functor g, Functor f, Integral n) => Int -> g n -> f n -> g (f C)
rootCross tot = (fmap . fmap . fmap . fmap) (uroot tot ^) products

-- rootCross tot is js = (fmap.fmap) (uroot tot ^) (is `products` js)
{--------------------------------------------------------------------
    FFT
--------------------------------------------------------------------}
-- | FFT computation, parametrized by structure
class HasFFT f f' | f -> f' where
  fft :: f C -> f' C

-- Constraint shorthands
type TA f = (Traversable f, Applicative f)

type TAH f f' = (TA f, TA f', HasFFT f f')

-- Binary butterfly.
instance HasFFT Pair Pair where
  fft (a :# b) = a + b :# a - b

-- Decimation in time (DIT)
instance (TAH f f', IsNat n) => HasFFT (B.T f n) (T.T f' n) where
  fft (B.L a) = T.L a
  fft (B.B t) = T.B (fftC t)

-- Decimation in frequency (DIF)? I'm unsure.
instance (TAH f f', IsNat n) => HasFFT (T.T f n) (B.T f' n) where
  fft (T.L a) = B.L a
  fft (T.B t) = B.B (fftC t)

instance HasFFT [] [] where
  fft = dft

--     Variable s `f, f' occur more often than in the instance head
--       in the constraint: TAH f
--     (Use -XUndecidableInstances to permit this)
--     In the instance declaration for `HasFFT (T f n)'
--
-- This warning vanishes when we spell out TAH. Hm.
-- I'd prefer terser definitions like `fft = T.inT' T.l (T.B . fftC)`, but I
-- haven't found a type for `inT'` that GHC likes.
fftsT :: (Applicative f, Traversable g, HasFFT g g') => g (f C) -> f (g' C)
fftsT = fmap fft . transpose

-- FFT of composed functors
fftC :: (TAH f f', TAH g g') => g (f C) -> f' (g' C)
fftC = transpose . fftsT . twiddle . fftsT

{-

--   fftsT     :: g  (f  C) -> f  (g' C)
--   twiddle   :: f  (g' C) -> f  (g' C)
--   fftsT     :: f  (g' C) -> g' (f' C)
--   transpose :: g' (f' C) -> f' (g' C)

-}
-- Multiply by twiddle factors
twiddle ::
     forall f g. (TA f, TA g)
  => Unop (g (f C))
twiddle = (liftA2 . liftA2) (*) rootses'
  where
    rootses' = rootCross (gTot * fTot) gIndices fIndices
      where
        fIndices :: f Int
        (fIndices, fTot) = counts
        gIndices :: g Int
        (gIndices, gTot) = counts

{--------------------------------------------------------------------
    Experimental variation
--------------------------------------------------------------------}
-- This version is a better fit with
-- <https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm#General_factorizations>
-- FFT of composed functors
fftC' :: (TAH f f', TAH g g') => g (f C) -> f' (g' C)
fftC' = fftsT' . transpose . twiddle . fftsT'
  where
    fftsT' :: (TA h, TAH k k') => k (h C) -> k' (h C)
    fftsT' = (inTranspose . fmap) fft

-- Hm. This definition differs from the previous one, since `twiddle` and
-- `transpose` got swapped. I doubt they're equivalent.
-- Types in the fftC' definition (right to left):
--
--   fftsT'    :: g  (f  C) -> g' (f  C)
--   twiddle   :: g' (f  C) -> g' (f  C)
--   transpose :: g' (f  C) -> f  (g' C)
--   fftsT'    :: f  (g' C) -> g' (g' C)
{--------------------------------------------------------------------
    Tests
--------------------------------------------------------------------}
-- {1, 0, 0, 0, ...} <=DFT=> {1, 1, 1, 1, ...}\
-- {1, -1, 1, -1, 1, -1, ...} <=DFT=> {0, 0, ... , N, 0, 0, ...} (where the 'N' occurs at the N/2 position in the output vector.
-- dft . idft ~= id
-- TODO :: show some example outputs and add a series of Quickcheck tests to
-- validate them.
-- TODO: make B / Pair showable

_p1 :: Pair C
_p1 = 1 :# 0

_t1 :: B.T Pair N3 C
_t1 = B.B (B.B (B.B (B.L (((1 :# 0) :# (0 :# 0)) :# ((0 :# 0) :# (0 :# 0))))))

_t2 :: B.T Pair N4 C
_t2 = pure 1

_t3 :: B.T Pair N3 C
_t3 = B.B (pure (1 :# -1))


{-|
Module      : Data.Meta.Multiarm
Description : Matrix function for graphs
Copyright   : (c) Thodoris Papakonstantinou, 2019
License     : GPL-3
Maintainer  : mail@tpapak.com
Stability   : experimental
Portability : POSIX

-}

module Data.Meta.Matrices
  ( createBMatrix
  , createAMatrix
  , createCMatrix
  )
where

import           Control.Applicative
import           Data.List
import qualified Data.Map.Strict               as Map
import qualified Data.Set                      as Set
import           Data.Maybe
import           Data.Either
import           Data.Meta.Effects
import qualified Numeric.LinearAlgebra         as LA
import qualified Numeric.LinearAlgebra.Devel   as LAD
import qualified Data.Graph.AdjacencyList      as G
import qualified Data.Graph.AdjacencyList.BFS      as BFS

-- | create the edge vertex adjacency matrix given a graph
createBMatrix :: G.Graph -> LA.Matrix Double
createBMatrix g =
  let vs  = G.vertices g
      nvs = length vs
      es  = G.edges g
      nes = length es
      rowfromEdge e =
          let
            (u, v) = G.toTuple e
            r      = LA.fromList $ replicate nvs 0.0 :: LA.Vector Double
            r' =
              LAD.mapVectorWithIndex
                (\i _ ->
                  if i == u - 1 then 1.0 else if i == v - 1 then (-1.0) else 0.0
                )
                r :: LA.Vector Double
          in
            r'
      rows = map rowfromEdge es
      b    = LA.fromRows rows
  in  b

createAMatrix :: G.Graph -> LA.Matrix Double
createAMatrix g =
  let vs  = G.vertices g
      nvs = length vs
      es  = G.edges g
      nes = length es
      rowfromEdge e =
          let
            (u, v) = G.toTuple e
            r      = LA.fromList $ replicate (nvs - 1) 0.0 :: LA.Vector Double
            r' =
              LAD.mapVectorWithIndex
                (\i _ ->
                  if i == u - 2 then 1.0 else if i == v - 2 then (-1.0) else 0.0
                )
                r :: LA.Vector Double
           in r'
      rows = map rowfromEdge es
      a    = LA.fromRows rows
  in  a

-- | create the Vertex Edge adjacency matrix given a graph

-- | create the Vertex Edge adjacency matrix given a graph
createCMatrix :: G.Graph -> LA.Matrix Double
createCMatrix g =
  let vs  = G.vertices g
      nvs = length vs
      es  = G.edges g
      nes = length es
      rowfromVertex v =
        let r = map (\e ->
              let (u, v') = G.toTuple e
                  fillcell x | x==u = (-1.0)
                  fillcell x | x==v' = (1.0)
                  fillcell _ = (0.0)
               in  fillcell v
                  ) es
         in LA.fromList $ r :: LA.Vector Double
      rows = map rowfromVertex vs
      c    = LA.fromRows rows
  in  c

-- | create the Vertex Edge adjacency matrix given a graph
createXMatrix :: G.Graph -> LA.Matrix Double
createXMatrix g =
  let vs  = G.vertices g
      nvs = length vs
      es  = G.edges g
      nes = length es
      firstVertex = head vs
      bfsNet = BFS.bfs g firstVertex
      spanningtree = BFS.spanningTree bfsNet
      rowfromVertex v =
        let r = map (\e ->
              let (u, v') = G.toTuple e
                  fillcell x | x==u = (-1.0)
                  fillcell x | x==v' = (1.0)
                  fillcell _ = (0.0)
               in  fillcell v
                  ) es
         in LA.fromList $ r :: LA.Vector Double
      rows = map rowfromVertex vs
      c    = LA.fromRows rows
  in  c


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-|
Module      : Grenade.Layers.Convolution
Maintainter : Theo Charalambous
Description : Convolution layer with support for padding and bias

Convolutions with and without biases are supported, allowing for faster
implementation when we know there is no bias.
-}

module Grenade.Layers.Convolution where

import           Data.Function                       ((&))
import           Data.Kind                           (Type)
import           Data.List                           (foldl1')
import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits            hiding (natVal)
import           GHC.TypeLits
import           Control.DeepSeq
import           Lens.Micro                          ((^.))

import           Numeric.LinearAlgebra               hiding (R, konst,
                                                      uniformSample)
import qualified Numeric.LinearAlgebra               as LA
import           Numeric.LinearAlgebra.Static        hiding (build, toRows, (&),
                                                      (|||), size)

import           Grenade.Core
import           Grenade.Layers.Internal.Convolution
import           Grenade.Layers.Internal.Update
import           Grenade.Onnx
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore

type OutputShapeIsOkay (input :: Nat) (pad :: Nat) (kernel :: Nat) (strides :: Nat) (output :: Nat)
  = ( strides * (output - 1) <= (input - kernel + pad)
    , (input - kernel + pad) <= (output * strides ) - 1 )

data HasBias = WithBias | WithoutBias

data ConvPadding = NoPadding | SameUpper | SameLower | Padding Nat Nat Nat Nat

data Convolution :: HasBias
                 -> ConvPadding
                 -> Nat -- Number of channels, for the first layer this could be RGB for instance.
                 -> Nat -- Number of filters, this is the number of channels output by the layer.
                 -> Nat -- The number of rows in the kernel filter
                 -> Nat -- The number of column in the kernel filter
                 -> Nat -- The row stride of the convolution filter
                 -> Nat -- The columns stride of the convolution filter
                 -> Type where
  Convolution :: ( KnownNat channels
                 , KnownNat filters
                 , KnownNat kernelRows
                 , KnownNat kernelColumns
                 , KnownNat strideRows
                 , KnownNat strideColumns
                 , KnownNat kernelFlattened
                 , kernelFlattened ~ (kernelRows * kernelColumns * channels))
              => !(L kernelFlattened filters) -- The kernel filter weights
              -> !(ListStore (L kernelFlattened filters)) -- The last kernel update (or momentum)
              -> Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns

  BiasConvolution :: ( KnownNat channels
                     , KnownNat filters
                     , KnownNat kernelRows
                     , KnownNat kernelColumns
                     , KnownNat strideRows
                     , KnownNat strideColumns
                     , KnownNat kernelFlattened
                     , kernelFlattened ~ (kernelRows * kernelColumns * channels))
                  => !(L kernelFlattened filters) -- The kernel filter weights
                  -> !(R filters) -- The bias weights
                  -> !(ListStore (L kernelFlattened filters)) -- The last kernel update (or momentum)
                  -> Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns

data Convolution' :: HasBias
                  -> ConvPadding
                  -> Nat -- Number of channels, for the first layer this could be RGB for instance.
                  -> Nat -- Number of filters, this is the number of channels output by the layer.
                  -> Nat -- The number of rows in the kernel filter
                  -> Nat -- The number of column in the kernel filter
                  -> Nat -- The row stride of the convolution filter
                  -> Nat -- The columns stride of the convolution filter
                  -> Type where
  Convolution' :: ( KnownNat channels
                  , KnownNat filters
                  , KnownNat kernelRows
                  , KnownNat kernelColumns
                  , KnownNat strideRows
                  , KnownNat strideColumns
                  , KnownNat kernelFlattened
                  , kernelFlattened ~ (kernelRows * kernelColumns * channels))
               => !(L kernelFlattened filters) -- The kernel filter weights
               -> Convolution' 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns

  BiasConvolution' :: ( KnownNat channels
                      , KnownNat filters
                      , KnownNat kernelRows
                      , KnownNat kernelColumns
                      , KnownNat strideRows
                      , KnownNat strideColumns
                      , KnownNat kernelFlattened
                      , kernelFlattened ~ (kernelRows * kernelColumns * channels))
                   => !(L kernelFlattened filters) -- The kernel filter weights
                   -> !(R filters) -- The bias weights
                   -> Convolution' 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns

{---------------------------}
{--    Layer instances    --}
{---------------------------}

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat ((kernelRows * kernelColumns) * channels)
         , KnownNat (filters * ((kernelRows * kernelColumns) * channels))
         ) => RandomLayer (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  createRandomWith m gen = do
    wN <- getRandomMatrix i i m gen
    wB <- getRandomVector i i m gen
    return $ BiasConvolution wN wB mkListStore
    where
      i = natVal (Proxy :: Proxy ((kernelRows * kernelColumns) * channels))

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat ((kernelRows * kernelColumns) * channels)
         , KnownNat (filters * ((kernelRows * kernelColumns) * channels))
         ) => RandomLayer (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  createRandomWith m gen = do
    wN <- getRandomMatrix i i m gen
    return $ Convolution wN mkListStore
    where
      i = natVal (Proxy :: Proxy ((kernelRows * kernelColumns) * channels))

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         , kernelFlattened ~ (kernelRows * kernelColumns * channels)
         ) => UpdateLayer (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  type Gradient (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) = (Convolution' 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns)

  type MomentumStore (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns)  = ListStore (L (kernelRows * kernelColumns * channels) filters)

  runUpdate opt@OptSGD{} x@(Convolution oldKernel store) (Convolution' kernelGradient) =
    let  momentum = getData opt x store
         result = descendMatrix opt (MatrixValuesSGD oldKernel kernelGradient momentum)
         newStore = setData opt x store (matrixMomentum result)
    in Convolution (matrixActivations result) newStore
  runUpdate opt@OptAdam{} x@(Convolution oldKernel store) (Convolution' kernelGradient) =
    let (m, v) = toTuple $ getData opt x store
        result = descendMatrix opt (MatrixValuesAdam (getStep store) oldKernel kernelGradient m v)
        newStore = setData opt x store [matrixM result, matrixV result]
    in Convolution (matrixActivations result) newStore
    where toTuple [m ,v] = (m, v)
          toTuple xs = error $ "unexpected input of length " ++ show (length xs) ++ "in toTuple in Convolution.hs"

  reduceGradient grads = Convolution' $ dmmap (/ (fromIntegral $ length grads)) (foldl1' add (map (\(Convolution' x) -> x) grads))

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         , kernelFlattened ~ (kernelRows * kernelColumns * channels)
         ) => UpdateLayer (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  type Gradient (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) = (Convolution' 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns)

  type MomentumStore (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns)  = ListStore (L (kernelRows * kernelColumns * channels) filters)

  runUpdate      = undefined
  reduceGradient = undefined

{-----------------------------------}
{--    No Bias Layer instances    --}
{-----------------------------------}

-- | A three dimensional image (or 2d with many channels) can have
--   an appropriately sized convolution filter run across it.
--   Case without bias vector.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , OutputShapeIsOkay inputRows 0 kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols 0 kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithoutBias 'NoPadding channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) where

  type Tape (Convolution 'WithoutBias 'NoPadding channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (Convolution kernel _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters)
        cs = fromIntegral $ natVal (Proxy :: Proxy channels)

        out  = forwardConv2d ex cs ix iy ek fs kx ky sx sy ox oy 0 0
    in  (S3D input, S3D . fromJust . create $ out)

  runBackwards (Convolution kernel _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ek = extract kernel
        ey = extract dEdy
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters)
        cs = fromIntegral $ natVal (Proxy :: Proxy channels)

        (dx, dw)  = backwardConv2d ex cs ix iy ek fs kx ky sx sy 0 0 0 0 ey ox oy
    in  (Convolution' . fromJust . create $ dw, S3D . fromJust . create $ dx)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , KnownNat padt
         , KnownNat padl
         , KnownNat padb
         , KnownNat padr
         , OutputShapeIsOkay inputRows (padt + padb) kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols (padl + padr) kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithoutBias ('Padding padl padt padr padb) channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) where

  type Tape (Convolution 'WithoutBias ('Padding padl padt padr padb) channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (Convolution kernel _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters)
        cs = fromIntegral $ natVal (Proxy :: Proxy channels)
        pl = fromIntegral $ natVal (Proxy :: Proxy padl)
        pt = fromIntegral $ natVal (Proxy :: Proxy padt)

        out  = fromJust . create $ forwardConv2d ex cs ix iy ek fs kx ky sx sy ox oy pl pt
    in  (S3D input, S3D out)

  runBackwards (Convolution kernel _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ek = extract kernel
        ey = extract dEdy
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows )
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols )
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters   )
        cs = fromIntegral $ natVal (Proxy :: Proxy channels  )
        pl = fromIntegral $ natVal (Proxy :: Proxy padl      )
        pt = fromIntegral $ natVal (Proxy :: Proxy padt      )
        pr = fromIntegral $ natVal (Proxy :: Proxy padr      )
        pb = fromIntegral $ natVal (Proxy :: Proxy padb      )

        (dx, dw) = backwardConv2d ex cs ix iy ek fs kx ky sx sy pl pt pr pb ey ox oy
    in  (Convolution' . fromJust . create $ dw, S3D . fromJust . create $ dx)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat channels
         ) => Layer (Convolution 'WithoutBias 'SameUpper channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 inputRows inputCols filters) where

  type Tape (Convolution 'WithoutBias 'SameUpper channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 inputRows inputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (Convolution kernel _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows )
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols )
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters   )
        cs = fromIntegral $ natVal (Proxy :: Proxy channels  )

        pady = (ix - 1) * sx + kx - ix
        padx = (iy - 1) * sy + ky - iy

        padt = div pady 2
        padl = div padx 2

        out  = fromJust . create $ forwardConv2d ex cs ix iy ek fs kx ky sx sy ix iy padl padt
    in  (S3D input, S3D out)

  runBackwards (Convolution kernel _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ek = extract kernel
        ey = extract dEdy
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows )
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols )
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters   )
        cs = fromIntegral $ natVal (Proxy :: Proxy channels  )

        pady = (ix - 1) * sx + kx - ix
        padx = (iy - 1) * sy + ky - iy

        padl = div padx 2
        padt = div pady 2
        padr = padx - padl
        padb = pady - padt

        (dx, dw) = backwardConv2d ex cs ix iy ek fs kx ky sx sy padl padt padr padb ey ix iy
    in  (Convolution' . fromJust . create $ dw, S3D . fromJust . create $ dx)

-- | A three dimensional image (or 2d with many channels) can have
--   an appropriately sized convolution filter run across it.
--   Case without bias vector.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat channels
         ) => Layer (Convolution 'WithoutBias 'SameLower channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 inputRows inputCols filters) where

  type Tape (Convolution 'WithoutBias 'SameLower channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 inputRows inputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (Convolution kernel _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters)
        cs = fromIntegral $ natVal (Proxy :: Proxy channels)

        pady = (ix - 1) * sx + kx - ix
        padx = (iy - 1) * sy + ky - iy

        padt = pady - div pady 2
        padl = padx - div padx 2
        out  = fromJust . create $ forwardConv2d ex cs ix iy ek fs kx ky sx sy ix iy padl padt

    in  (S3D input, S3D out)

  runBackwards (Convolution kernel _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ek = extract kernel
        ey = extract dEdy
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows )
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols )
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters   )
        cs = fromIntegral $ natVal (Proxy :: Proxy channels  )

        pady = (ix - 1) * sx + kx - ix
        padx = (iy - 1) * sy + ky - iy

        padt = pady - padb
        padl = padx - padr
        padr = div padx 2
        padb = div pady 2

        (dx, dw) = backwardConv2d ex cs ix iy ek fs kx ky sx sy padl padt padr padb ey ix iy
    in  (Convolution' . fromJust . create $ dw, S3D . fromJust . create $ dx)

{--------------------------------}
{--    Bias Layer instances    --}
{--------------------------------}

-- | A three dimensional image (or 2d with many channels) can have
--   an appropriately sized convolution filter run across it.
--   Case with bias vector
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , OutputShapeIsOkay inputRows 0 kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols 0 kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithBias 'NoPadding channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) where

  type Tape (Convolution 'WithBias 'NoPadding channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (BiasConvolution kernel biases _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        eb = extract biases
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters)
        cs = fromIntegral $ natVal (Proxy :: Proxy channels)

        out  = fromJust . create $ forwardBiasConv2d ex cs ix iy eb ek fs kx ky sx sy ox oy 0 0

    in  (S3D input, S3D out)

  runBackwards (BiasConvolution kernel _ _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ek = extract kernel
        ey = extract dEdy
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters)
        cs = fromIntegral $ natVal (Proxy :: Proxy channels)

        (dx', dw', db') = backwardBiasConv2d ex cs ix iy ek fs kx ky sx sy 0 0 0 0 ey ox oy
        dx              = fromJust . create $ dx'
        dw              = fromJust . create $ dw'
        db              = fromJust . create $ db'
    in  (BiasConvolution' dw db, S3D dx)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , KnownNat padt
         , KnownNat padl
         , KnownNat padb
         , KnownNat padr
         , OutputShapeIsOkay inputRows (padt + padb) kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols (padl + padr) kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithBias ('Padding padl padt padr padb) channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) where

  type Tape (Convolution 'WithBias ('Padding padl padt padr padb) channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (BiasConvolution kernel biases _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        eb = extract biases
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows )
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols )
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters   )
        cs = fromIntegral $ natVal (Proxy :: Proxy channels  )
        pl = fromIntegral $ natVal (Proxy :: Proxy padl      )
        pt = fromIntegral $ natVal (Proxy :: Proxy padt      )

        out  = fromJust . create $ forwardBiasConv2d ex cs ix iy eb ek fs kx ky sx sy ox oy pl pt

    in  (S3D input, S3D out)

  runBackwards (BiasConvolution kernel _ _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ek = extract kernel
        ey = extract dEdy
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows )
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols )
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters   )
        cs = fromIntegral $ natVal (Proxy :: Proxy channels  )
        pl = fromIntegral $ natVal (Proxy :: Proxy padl      )
        pt = fromIntegral $ natVal (Proxy :: Proxy padt      )
        pr = fromIntegral $ natVal (Proxy :: Proxy padr      )
        pb = fromIntegral $ natVal (Proxy :: Proxy padb      )

        (dx', dw', db') = backwardBiasConv2d ex cs ix iy ek fs kx ky sx sy pl pt pr pb ey ox oy
        dx              = fromJust . create $ dx'
        dw              = fromJust . create $ dw'
        db              = fromJust . create $ db'
    in  (BiasConvolution' dw db, S3D dx)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat channels
         ) => Layer (Convolution 'WithBias 'SameUpper channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 inputRows inputCols filters) where

  type Tape (Convolution 'WithBias 'SameUpper channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 inputRows inputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (BiasConvolution kernel bias _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        eb = extract bias
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters)
        cs = fromIntegral $ natVal (Proxy :: Proxy channels)

        pady = (ix - 1) * sx + kx - ix
        padx = (iy - 1) * sy + ky - iy

        padt = div pady 2
        padl = div padx 2

        out  = fromJust . create $ forwardBiasConv2d ex cs ix iy eb ek fs kx ky sx sy ix iy padl padt

    in  (S3D input, S3D out)

  runBackwards (BiasConvolution kernel _ _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ek = extract kernel
        ey = extract dEdy
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows )
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols )
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters   )
        cs = fromIntegral $ natVal (Proxy :: Proxy channels  )

        pady = (ix - 1) * sx + kx - ix
        padx = (iy - 1) * sy + ky - iy

        padl = div padx 2
        padt = div pady 2
        padr = padx - padl
        padb = pady - padt

        (dx', dw', db') = backwardBiasConv2d ex cs ix iy ek fs kx ky sx sy padl padt padr padb ey ix iy
        dx              = fromJust . create $ dx'
        dw              = fromJust . create $ dw'
        db              = fromJust . create $ db'
    in  (BiasConvolution' dw db, S3D dx)

-- | A three dimensional image (or 2d with many channels) can have
--   an appropriately sized convolution filter run across it.
--   Case without bias vector.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat channels
         ) => Layer (Convolution 'WithBias 'SameLower channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 inputRows inputCols filters) where

  type Tape (Convolution 'WithBias 'SameLower channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 inputRows inputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (BiasConvolution kernel bias _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        eb = extract bias
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters)
        cs = fromIntegral $ natVal (Proxy :: Proxy channels)

        pady = (ix - 1) * sx + kx - ix
        padx = (iy - 1) * sy + ky - iy

        padt = pady - div pady 2
        padl = padx - div padx 2

        out  = fromJust . create $ forwardBiasConv2d ex cs ix iy eb ek fs kx ky sx sy ix iy padl padt
    in  (S3D input, S3D out)

  runBackwards (BiasConvolution kernel _ _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ek = extract kernel
        ey = extract dEdy
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows )
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols )
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        fs = fromIntegral $ natVal (Proxy :: Proxy filters   )
        cs = fromIntegral $ natVal (Proxy :: Proxy channels  )

        pady = (ix - 1) * sx + kx - ix
        padx = (iy - 1) * sy + ky - iy

        padr = div padx 2
        padb = div pady 2
        padl = padx - padr
        padt = pady - padb

        (dx', dw', db') = backwardBiasConv2d ex cs ix iy ek fs kx ky sx sy padl padt padr padb ey ix iy
        dx              = fromJust . create $ dx'
        dw              = fromJust . create $ dw'
        db              = fromJust . create $ db'
    in  (BiasConvolution' dw db, S3D dx)

{---------------------------------}
{--    Other Layer instances    --}
{---------------------------------}

-- | A two dimentional image may have a convolution filter applied to it.
--   Case without bias vector
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , OutputShapeIsOkay inputRows 0 kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols 0 kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithoutBias 'NoPadding 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) where
  type Tape (Convolution 'WithoutBias 'NoPadding 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         ) => Layer (Convolution 'WithoutBias 'SameUpper 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 inputRows inputCols filters) where
  type Tape (Convolution 'WithoutBias 'SameUpper 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 inputRows inputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         ) => Layer (Convolution 'WithoutBias 'SameLower 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 inputRows inputCols filters) where
  type Tape (Convolution 'WithoutBias 'SameLower 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 inputRows inputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

-- | A two dimentional image may have a convolution filter applied to it.
--   Case with bias vector.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , OutputShapeIsOkay inputRows 0 kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols 0 kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithBias 'NoPadding 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) where
  type Tape (Convolution 'WithBias 'NoPadding 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         ) => Layer (Convolution 'WithBias 'SameUpper 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 inputRows inputCols filters) where
  type Tape (Convolution 'WithBias 'SameUpper 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 inputRows inputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         ) => Layer (Convolution 'WithBias 'SameLower 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 inputRows inputCols filters) where
  type Tape (Convolution 'WithBias 'SameLower 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 inputRows inputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

-- | A two dimensional image may have a convolution filter applied to it producing
--   a two dimensional image if both channels and filters is 1.
--   Case without bias vector.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputCols
         , KnownNat outputRows
         , OutputShapeIsOkay inputRows 0 kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols 0 kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithoutBias 'NoPadding 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) where
  type Tape (Convolution 'WithoutBias 'NoPadding 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D back :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) -> (c', S2D back)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         ) => Layer (Convolution 'WithoutBias 'SameUpper 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 inputRows inputCols) where
  type Tape (Convolution 'WithoutBias 'SameUpper 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 inputRows inputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D back :: S ('D3 inputRows inputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 inputRows inputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) -> (c', S2D back)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         ) => Layer (Convolution 'WithoutBias 'SameLower 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 inputRows inputCols) where
  type Tape (Convolution 'WithoutBias 'SameLower 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 inputRows inputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D back :: S ('D3 inputRows inputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 inputRows inputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) -> (c', S2D back)

-- | A two dimensional image may have a convolution filter applied to it producing
--   a two dimensional image if both channels and filters is 1.
--   Case with bias vector.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputCols
         , KnownNat outputRows
         , OutputShapeIsOkay inputRows 0 kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols 0 kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithBias 'NoPadding 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) where
  type Tape (Convolution 'WithBias 'NoPadding 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D back :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) -> (c', S2D back)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         ) => Layer (Convolution 'WithBias 'SameUpper 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 inputRows inputCols) where
  type Tape (Convolution 'WithBias 'SameUpper 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 inputRows inputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D back :: S ('D3 inputRows inputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 inputRows inputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) -> (c', S2D back)

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         ) => Layer (Convolution 'WithBias 'SameLower 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 inputRows inputCols) where
  type Tape (Convolution 'WithBias 'SameLower 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 inputRows inputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D back :: S ('D3 inputRows inputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 inputRows inputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) -> (c', S2D back)

-- | A three dimensional image can produce a 2D image from a convolution with 1 filter
--   Case without bias vector.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputCols
         , KnownNat outputRows
         , KnownNat channels
         , OutputShapeIsOkay inputRows 0 kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols 0 kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithoutBias 'NoPadding channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) where
  type Tape (Convolution 'WithoutBias 'NoPadding channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D back :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1))

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat channels
         ) => Layer (Convolution 'WithoutBias 'SameUpper channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 inputRows inputCols) where
  type Tape (Convolution 'WithoutBias 'SameUpper channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 inputRows inputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D back :: S ('D3 inputRows inputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 inputRows inputCols 1))

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat channels
         ) => Layer (Convolution 'WithoutBias 'SameLower channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 inputRows inputCols) where
  type Tape (Convolution 'WithoutBias 'SameLower channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 inputRows inputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D back :: S ('D3 inputRows inputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 inputRows inputCols 1))

-- | A three dimensional image can produce a 2D image from a convolution with 1 filter
--   Case with bias vector.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputCols
         , KnownNat outputRows
         , KnownNat channels
         , OutputShapeIsOkay inputRows 0 kernelRows strideRows outputRows
         , OutputShapeIsOkay inputCols 0 kernelCols strideCols outputCols
         ) => Layer (Convolution 'WithBias 'NoPadding channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) where
  type Tape (Convolution 'WithBias 'NoPadding channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D back :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1))

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat channels
         ) => Layer (Convolution 'WithBias 'SameUpper channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 inputRows inputCols) where
  type Tape (Convolution 'WithBias 'SameUpper channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 inputRows inputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D back :: S ('D3 inputRows inputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 inputRows inputCols 1))

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat channels
         ) => Layer (Convolution 'WithBias 'SameLower channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 inputRows inputCols) where
  type Tape (Convolution 'WithBias 'SameLower channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 inputRows inputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D back :: S ('D3 inputRows inputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 inputRows inputCols 1))

{--------------------}
{-- ONNX Instances --}
{--------------------}

instance OnnxOperator (Convolution 'WithBias padding channels filters kernelRows kernelCols strideRows strideCols) where
  onnxOpTypeNames _ = ["Conv"]

instance OnnxOperator (Convolution 'WithoutBias padding channels filters kernelRows kernelCols strideRows strideCols) where
  onnxOpTypeNames _ = ["Conv"]

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat channels
         , KnownNat (kernelRows * kernelCols * channels)
         ) => OnnxLoadable (Convolution 'WithoutBias 'NoPadding channels filters kernelRows kernelCols strideRows strideCols) where
  loadOnnxNode inits node = do
    node `doesNotHaveAttribute` "auto_pad"

    node & hasSupportedDilations
    node & hasSupportedGroup

    (node `hasMatchingShape` "kernel_shape") kernelShape
    (node `hasMatchingShape` "strides"     ) strideShape
    hasCorrectPadding node (Proxy :: Proxy 0) (Proxy :: Proxy 0) (Proxy :: Proxy 0) (Proxy :: Proxy 0)

    case node ^. #input of
      [_, w] -> do
        filterWeights <- tr <$> readInitializerTensorIntoMatrix inits w
        return (Convolution filterWeights mkListStore)
      _ -> onnxIncorrectNumberOfInputs
      where
        kernelShape = [natVal (Proxy :: Proxy kernelRows), natVal (Proxy :: Proxy kernelCols)]
        strideShape = [natVal (Proxy :: Proxy strideRows), natVal (Proxy :: Proxy strideCols)]

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat channels
         , KnownNat (kernelRows * kernelCols * channels)
         ) => OnnxLoadable (Convolution 'WithoutBias 'SameUpper channels filters kernelRows kernelCols strideRows strideCols) where
  loadOnnxNode inits node = do
    node `doesNotHaveAttribute` "auto_pad"

    node & hasSupportedDilations
    node & hasSupportedGroup

    (node `hasMatchingShape` "kernel_shape") kernelShape
    (node `hasMatchingShape` "strides"     ) strideShape

    case node ^. #input of
      [_, w] -> do
        filterWeights <- tr <$> readInitializerTensorIntoMatrix inits w
        return (Convolution filterWeights mkListStore)
      _ -> onnxIncorrectNumberOfInputs
      where
        kernelShape = [natVal (Proxy :: Proxy kernelRows), natVal (Proxy :: Proxy kernelCols)]
        strideShape = [natVal (Proxy :: Proxy strideRows), natVal (Proxy :: Proxy strideCols)]

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat channels
         , KnownNat padl, KnownNat padt, KnownNat padr, KnownNat padb
         ) => OnnxLoadable (Convolution 'WithoutBias ('Padding padl padt padr padb) channels filters kernelRows kernelCols strideRows strideCols) where
  loadOnnxNode inits node = do
    node `doesNotHaveAttribute` "auto_pad"

    node & hasSupportedDilations
    node & hasSupportedGroup

    (node `hasMatchingShape` "kernel_shape") kernelShape
    (node `hasMatchingShape` "strides"     ) strideShape

    hasCorrectPadding node (Proxy :: Proxy padl) (Proxy :: Proxy padr) (Proxy :: Proxy padt) (Proxy :: Proxy padb)

    case node ^. #input of
      [_, w] -> do
        filterWeights <- tr <$> readInitializerTensorIntoMatrix inits w
        return (Convolution filterWeights mkListStore)
      _ -> onnxIncorrectNumberOfInputs
      where
        kernelShape = [natVal (Proxy :: Proxy kernelRows), natVal (Proxy :: Proxy kernelCols)]
        strideShape = [natVal (Proxy :: Proxy strideRows), natVal (Proxy :: Proxy strideCols)]

instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat channels
         , KnownNat (kernelRows * kernelCols * channels)
         ) => OnnxLoadable (Convolution 'WithBias padding channels filters kernelRows kernelCols strideRows strideCols) where
  loadOnnxNode inits node = do
    node & hasSupportedDilations
    node & hasSupportedGroup

    (node `hasMatchingShape` "kernel_shape") kernelShape
    (node `hasMatchingShape` "strides"     ) strideShape

    -- todo: proper checking to see if auto_pad attribute is valid

    case node ^. #input of
      [_, w, b] -> do
        filterWeights <- tr <$> readInitializerTensorIntoMatrix inits w
        filterBias    <- readInitializerVector inits b
        return (BiasConvolution filterWeights filterBias mkListStore)
      _ -> onnxIncorrectNumberOfInputs
      where
        kernelShape = [natVal (Proxy :: Proxy kernelRows), natVal (Proxy :: Proxy kernelCols)]
        strideShape = [natVal (Proxy :: Proxy strideRows), natVal (Proxy :: Proxy strideCols)]

{--------------------}
{-- Misc Instances --}
{--------------------}

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         ) =>
         LayerOptimizerData (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) where
  type MomentumExpOptResult (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) = L (kernelRows * kernelColumns * channels) filters
  type MomentumDataType (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) = L (kernelRows * kernelColumns * channels) filters
  getData opt x store = head $ getListStore opt x store
  setData opt x store = setListStore opt x store . return
  newData _ _ = konst 0

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         ) =>
         LayerOptimizerData (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) where
  type MomentumExpOptResult (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) = L (kernelRows * kernelColumns * channels) filters
  type MomentumDataType (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) = L (kernelRows * kernelColumns * channels) filters
  getData = undefined
  setData = undefined
  newData = undefined

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         ) => LayerOptimizerData (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam) where
  type MomentumExpOptResult (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam)  = [L (kernelRows * kernelColumns * channels) filters]
  type MomentumDataType (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam) = L (kernelRows * kernelColumns * channels) filters
  getData = getListStore
  setData = setListStore
  newData _ _ = konst 0

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => LayerOptimizerData (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam) where
  type MomentumExpOptResult (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam)  = [L (kernelRows * kernelColumns * channels) filters]
  type MomentumDataType (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam) = L (kernelRows * kernelColumns * channels) filters
  getData = undefined
  setData = undefined
  newData = undefined

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns) => FoldableGradient (Convolution' hasBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  mapGradient f (Convolution' kernelGradient) = Convolution' (dmmap f kernelGradient)
  mapGradient f (BiasConvolution' kernelGradient biasGradient) = BiasConvolution' (dmmap f kernelGradient) (dvmap f biasGradient)
  squaredSums (Convolution' kernelGradient) = [sumM . squareM $ kernelGradient]
  squaredSums (BiasConvolution' kernelGradient biasGradient) = [sumM . squareM $ kernelGradient, sumV . squareV $ biasGradient ]

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => Serialize (Convolution 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (Convolution w store) = do
    putListOf put . toList . flatten . extract $ w
    put (fmap (toList . flatten . extract) store)
  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy filters)
      wN    <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
      store <- fmap (fromMaybe (error "Vector of incorrect size") . create . reshape f . LA.fromList)  <$> get
      return $ Convolution wN store

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => Serialize (Convolution 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (BiasConvolution w b store) = do
    putListOf put . toList . flatten . extract $ w
    putListOf put . toList . extract $ b
    put (fmap (toList . flatten . extract) store)
  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy filters)
      wN    <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
      wB    <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
      store <- fmap (fromMaybe (error "Vector of incorrect size") . create . reshape f . LA.fromList)  <$> get
      return $ BiasConvolution wN wB store

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => Serialize (Convolution' 'WithoutBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (Convolution' w) = do
    putListOf put . toList . flatten . extract $ w
  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy filters)
      wN    <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
      return $ Convolution' wN

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => Serialize (Convolution' 'WithBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (BiasConvolution' w b) = do
    putListOf put . toList . flatten . extract $ w
    putListOf put . toList . extract $ b
  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy filters)
      wN    <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
      wB    <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
      return $ BiasConvolution' wN wB

instance NFData (Convolution hasBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  rnf (BiasConvolution a b c) = rnf a `seq` rnf b `seq` rnf c
  rnf (Convolution a b)       = rnf a `seq` rnf b

instance NFData (Convolution' hasBias padding channels filters kernelRows kernelColumns strideRows strideColumns) where
  rnf (BiasConvolution' a b) = rnf a `seq` rnf b
  rnf (Convolution' a)       = rnf a

instance Show (Convolution hasBias padding c f k k' s s') where
  show (BiasConvolution _ _ _) = "Bias Convolution"
  show (Convolution _ _)       = "Convolution"