{-# 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"


{-|
Module      : MachineLearning.Model.Regression
Description : TIR expression data structures
Copyright   : (c) Fabricio Olivetti de Franca, 2022
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX

Fitting functions for the coefficients.
-}
module MachineLearning.Model.Regression
  ( fitTask
  , predictTask
  , evalPenalty
  , applyMeasures
  , nonlinearFit
  , tirToMatrix
  ) where

import Data.Bifunctor
import           Data.SRTree                   (evalFun, Function(..), OptIntPow(..), evalFun, inverseFunc)
import MachineLearning.TIR       (TIR(..),  Individual(..), Dataset, Constraint, assembleTree, replaceConsts)
import qualified Data.Vector           as V
import qualified Data.Vector.Storable  as VS
import Data.Vector           ((!))
import           Data.List                     (nub)
import           Data.Vector.Storable          (Vector, splitAt)
import           Numeric.LinearAlgebra         ((<\>), Matrix)
import           Numeric.GSL.Fitting           (nlFitting, FittingMethod(..))
import           Prelude               hiding  (splitAt)

import qualified Numeric.LinearAlgebra                     as LA

import MachineLearning.TIR (Individual(..))
import MachineLearning.Utils.Config (Task(..), Penalty(..))
import MachineLearning.Model.Measure (Measure(..))

-- * IT specific stuff

createQ :: Vector Double -> LA.Matrix Double -> LA.Matrix Double
createQ ys = LA.fromColumns . map (*ys) . LA.toColumns
{-# INLINE createQ #-}

avg :: Vector Double -> Vector Double
avg ys = LA.fromList [LA.sumElements ys / fromIntegral (LA.size ys)]
{-# INLINE avg #-}

-- | transform a data matrix using a TIR expression. This function returns
-- a tuple with the transformed data of the numerator and denominator, respectivelly.
-- Each column of the transformed data represents one term of the TIR expression.
tirToMatrix :: Dataset Double -> TIR -> (LA.Matrix Double, LA.Matrix Double)
tirToMatrix xss (TIR _ p q) = bimap (LA.fromColumns . (bias:)) LA.fromColumns (p', q')
  where
    bias      = V.head xss
    xss'      = V.tail xss
    sigma2mtx = map (\(_, g, ps) -> evalFun g $ evalPi ps)
    evalPi    = foldr (\(ix, k) acc -> acc * (xss' ! ix ^^ k)) 1
    p'        = sigma2mtx p
    q'        = sigma2mtx q
{-# INLINE tirToMatrix #-}

-- | Fits a linear model using l2-penalty
ridge :: Matrix Double -> Vector Double -> Matrix Double
ridge a b = oA <\> oB
  where
   mu = 0.01

   a' = LA.tr a
   b' = LA.tr $ LA.asColumn b
   oA = (a' <> LA.tr a') + (mu * LA.ident (LA.rows a'))
   oB = a' <> LA.tr b'
{-# INLINE ridge #-}

-- | Predicts a linear model
predict :: Matrix Double -> Vector Double -> Vector Double
predict xs w | LA.cols xs == LA.size w = xs LA.#> w
             | otherwise = error $ "predict: " ++ show (LA.size xs) ++ show (LA.size w)
{-# INLINE predict #-}

-- | Solve the OLS *zss*w = ys*
solveOLS :: Matrix Double -> Vector Double -> Vector Double
solveOLS zss ys = zss <\> ys
{-# INLINE solveOLS #-}

-- | Applies OLS and returns a Solution
-- if the expression is invalid, it returns Infinity as a fitness
--regress :: Matrix Double -> Vector Double -> [Vector Double]
regress :: TIR -> Dataset Double -> Vector Double -> [Vector Double]
regress tir xss ys = [ws]
  where
    (zssP, zssQ) = tirToMatrix xss tir
    ys'          = evalFun (inverseFunc $ _funY tir) ys
    zssQy        = createQ ys' zssQ
    zss          = if LA.cols zssQ >= 1
                       then zssP LA.||| negate zssQy
                       else zssP
    ws             = if LA.cols zss == 1
                       then avg ys'
                       else solveOLS zss ys'
-- regress zss ys = [solveOLS zss ys]
{-# INLINE regress #-}

-- | Applies conjugate gradient for binary classification
--classify :: Matrix Double -> Vector Double -> [Vector Double]
classify :: Int -> TIR -> Dataset Double -> Vector Double -> [Vector Double]
classify niter tir xss ys = [ws]
  where
    ws           = nonlinearFit niter zssP zssQ ys sigmoid dsigmoid theta0
    theta0       = LA.konst 0 (LA.cols zssP + LA.cols zssQ)
    (zssP, zssQ) = tirToMatrix xss tir

-- | Applies conjugate gradient for one-vs-all classification
--classifyMult :: Matrix Double -> Vector Double -> [Vector Double]
classifyMult :: Int -> TIR -> Dataset Double -> Vector Double -> [Vector Double]
classifyMult niter tir xss ys = zipWith minimize yss theta0
  where
    numLabels    = length $ nub $ LA.toList ys
    yss          = map f [0 .. numLabels-1]
    minimize y t = nonlinearFit niter zssP zssQ y sigmoid dsigmoid t
    theta0       = replicate numLabels $ LA.konst 0 (LA.cols zssP + LA.cols zssQ)
    (zssP, zssQ) = tirToMatrix xss tir

    f sample = VS.map (\a -> if round a == sample then 1 else 0) ys

-- | chooses the appropriate fitting function
--fitTask :: Task -> Matrix Double -> Vector Double -> [Vector Double]
fitTask :: Task -> TIR -> Dataset Double -> Vector Double -> [Vector Double]
fitTask Regression             = regress
fitTask (RegressionNL niter)   = regressNL niter
fitTask (Classification niter) = classify niter
fitTask (ClassMult niter)      = classifyMult niter
{-# INLINE fitTask #-}

-- | sigmoid function for classification.
sigmoid :: Floating a => a -> a
sigmoid z = 1 / (1+exp(-z))
{-# INLINE sigmoid #-}

-- | derivative sigmoid function for classification.
dsigmoid :: Floating a => a -> a
dsigmoid z = sigmoid z * (1 - sigmoid z)
{-# INLINE dsigmoid #-}

-- | chooses the appropriate prediction function
predictTask :: Task -> [Vector Double] -> Vector Double
predictTask _ []                   = error "predictTask: empty coefficients matrix"
predictTask Regression yss         = head yss
predictTask (RegressionNL _) yss   = head yss
predictTask (Classification _) yss = sigmoid $ head yss
predictTask (ClassMult _) yss      = LA.vector $ map (fromIntegral . LA.maxIndex) $ LA.toRows $ LA.fromColumns $ map sigmoid yss
{-# INLINE predictTask #-}

-- | evals the penalty function
evalPenalty :: Penalty -> Int -> Double -> Double
evalPenalty NoPenalty _   _   = 0.0
evalPenalty (Len c)   len _   = fromIntegral len * c
evalPenalty (Shape c) _   val = val*c
{-# INLINE evalPenalty #-}

-- | applies a list of performance measures
applyMeasures :: [Measure] -> Vector Double -> Vector Double -> [Double]
applyMeasures measures ys ysHat = map ((`uncurry` (ys, ysHat)) . _fun) measures
{-# INLINE applyMeasures #-}

regressNL :: Int -> TIR -> Dataset Double -> Vector Double -> [Vector Double]
regressNL niter tir xss ys = [ws]
  where
    ws           = nonlinearFit niter zssP zssQ ys f f' theta0
    f            = evalFun $ _funY tir
    f'           = derivative $ _funY tir
    --theta0       = LA.konst 1 (LA.cols zssP + LA.cols zssQ)
    theta0       = head $ regress tir xss ys
    (zssP, zssQ) = tirToMatrix xss tir

-- | Non-linear optimization using Levenberg-Marquardt method.
--nonlinearFit :: Monad m => Vector Double -> Matrix Double -> Matrix Double -> Vector Double -> m (Vector Double)
nonlinearFit :: Int
             -> Matrix Double
             -> Matrix Double
             -> Vector Double
             -> (Vector Double -> Vector Double)
             -> (Vector Double -> Vector Double)
             -> Vector Double
             -> Vector Double
nonlinearFit niter zssP zssQ ys f f' theta0 = fst $ nlFitting LevenbergMarquardtScaled 1e-6 1e-6 niter model' jacob' theta0
  where
    model'       = model f ys zssP zssQ
    jacob'       = jacob f' zssP zssQ

-- | calculates the error given the parameter vector beta
model :: (Vector Double -> Vector Double) -> Vector Double -> Matrix Double -> Matrix Double -> Vector Double -> Vector Double
model f ys zssP zssQ beta
  | LA.cols zssQ == 0 = f ysHat_P - ys
  | otherwise         = f ysHat - ys
  where
    (betaP, betaQ) = splitAt (LA.cols zssP) beta
    ysHat_P        = predict zssP betaP
    ysHat_Q        = if LA.cols zssQ == 0 then 0 else predict zssQ betaQ
    ysHat          = ysHat_P / (1 + ysHat_Q)

-- | calculates the Jacobian given the parameter vector beta. Doesn't support
-- the outer transformation function.
jacob :: (Vector Double -> Vector Double) -> Matrix Double -> Matrix Double -> Vector Double -> Matrix Double
jacob f zssP zssQ beta | LA.cols zssQ == 0 = zssP
                       | otherwise         = LA.fromColumns $ pjac <> qjac
  where
    (betaP, betaQ) = splitAt (LA.cols zssP) beta
    ysHat_P        = predict zssP betaP
    ysHat_Q        = predict zssQ betaQ
    ysHat          = f $ ysHat_P / (1 + ysHat_Q)
    pjac           = [ysHat * c / ysHat_Q | c <- LA.toColumns zssP]
    qjac           = [-ysHat * c * (ysHat_P / (1 + ysHat_Q)^2) | c <- LA.toColumns zssQ]

derivative :: (Eq val, Floating val) => Function -> val -> val
derivative Id      = const 1
derivative Abs     = \x -> x / abs x
derivative Sin     = cos
derivative Cos     = negate.sin
derivative Tan     = recip . (**2.0) . cos
derivative Sinh    = cosh
derivative Cosh    = sinh
derivative Tanh    = (1-) . (**2.0) . tanh
derivative ASin    = recip . sqrt . (1-) . (^2)
derivative ACos    = negate . recip . sqrt . (1-) . (^2)
derivative ATan    = recip . (1+) . (^2)
derivative ASinh   = recip . sqrt . (1+) . (^2)
derivative ACosh   = \x -> 1 / (sqrt (x-1) * sqrt (x+1))
derivative ATanh   = recip . (1-) . (^2)
derivative Sqrt    = recip . (2*) . sqrt
derivative Square  = (2*)
derivative Exp     = exp
derivative Log     = recip
{-# INLINE derivative #-}
-- (w1 * p1 + w2 * p2) / (1 + w3 * p3 + w4 * p4)
-- d/dw1 = p1 / (1 + w3 * p3 + w4 * p4)
-- d/dw2 = p2 / (1 + w3 * p3 + w4 * p4)
-- d/dw3 = -p3 * (w1 * p2 + w2 * p2) / (1 + w3 * p3 + w4 * p4)^2
{-
toEv :: SRTree Int Double -> (V.Vector Double -> Double)
toEv (Var !ix) = (`V.unsafeIndex` ix)
toEv (Const !val) = const val
toEv (Add !l !r) = jn (+) (toEv l) (toEv r)
toEv (Mul !l !r) = jn (*) (toEv l) (toEv r)
toEv (Fun Exp !t) = exp . toEv t
{-# INLINE toEv #-}

jn :: (Double -> Double -> Double) -> (V.Vector Double -> Double) -> (V.Vector Double -> Double) -> (V.Vector Double -> Double)
jn op f g = \x -> op (f x) (g x)
{-# INLINE jn #-}
-}


module PNN where

import Data.Complex
import qualified Data.Vector as V
import Data.Matrix
import Data.Random.Normal
import Data.List (sortBy)
import Data.Number.BigFloat
import Math.NumberTheory.Zeta

type BF = BigFloat Prec50

n = 40 :: Int
k = 1 :: Int

s = 2.1 :: Double
sd = 1 / sqrt(s * (fromIntegral n))
seed = 1

makeComplex (x:y:xs) = (x :+ y) : makeComplex xs

mkBigNormals' :: (Double,Double) -> Int -> [BF]
mkBigNormals' (u,sd) seed = map (fromRational . toRational)
   (mkNormals' (u,sd) seed)

complexes = makeComplex (mkBigNormals' (0,sd) seed)

w_res = fromList (n-k) (n-k) complexes
w_in = fromList (n-k) k $ drop (nrows w_res * ncols w_res) complexes
x_res = V.fromList $ take (n-k) $ drop (nrows w_res * ncols w_res +
                                        nrows w_in * ncols w_in) complexes

-- function to predict
-- vector should have length k
epsilon = 1e-40 :: BF
zetas' = zetas epsilon
zs = [z :+ 0 | z <- zetas']
as c = V.singleton (zs !! (c+2))
{-
as c = let c' = 0.01 * (fromIntegral c) in
  V.fromList [sin c' + 0.1,
              cos (2*c'+1),
              3*(sin (c'+1)) - cos (3*c' + 2)]
-}

-- matrix / vector multiply
(*:) mat vec = getCol 1 (multStd mat (colVector vec))

-- vector addition
(=+=) va vb = V.zipWith (+) va vb

-- vector addition
(=-=) va vb = V.zipWith (-) va vb

fromVectors vs = fromLists (map V.toList vs)

mat_a = fromVectors (zipWith (V.++) (map as [0..])
                     (take n (map snd mat_a_gen)))

mat_a_gen = iterate f (0, x_res) where
  f (i, x_res') = (i+1, (w_in *: (as i)) =+= (w_res *: x_res'))

{-
-- not sure why this doesn't work
w_out =
  let m = mat_a <|> (fromVectors (map as [1..n]))
      Right m' = rref m
  in m' --submatrix 1 (n+1) (n+1) (n+k) m'
-}

solve m = backsub (triangular m)
 where
   sortRows = sortBy (\(r:_) (s:_) -> compare (realPart$abs s) (realPart$abs r))
   triangular = triangular' . sortRows

   downsub (u:us) (v:vs) = zipWith (-) (map ((v / u) *) us) vs

   triangular' [] = []
   triangular' rs | all (0 ==) (map head rs) = error "underdetermined"
   triangular' ((p:xs):rs) =
     let r = 1 : map (/ p) xs
         ss = map (downsub r) rs
     in r :  (triangular' ss)

   backsubOne (u:us) (1:vs) = zipWith (-) us (map (u *) vs)
   backsub' us [] = us
   backsub' us (v:vs) = backsub' ((tail v) : map (flip backsubOne v) us) vs
   backsub rs = reverse (backsub' [] rs)

{-
solve' m =
 let rows = length m
     a = map (take rows) m
     b = map (drop rows) m
 in L.transpose (O.many_kets_solution (L.transpose a) (L.transpose b))
-}

w_out = let m = mat_a <|> (fromVectors (map as [1..n]))
        in fromLists (solve (toLists m))

m = (transpose w_out) <-> (w_in <|> w_res)

guesses = map (V.take k) (iterate (m *:) ((as 0) V.++ x_res))

norm1 mat = maximum [V.sum (V.map (realPart . abs) (getCol c mat)) |
   c <- [1 .. ncols mat]]

number_of_guesses = 10
err = norm1 (fromVectors ((zipWith (=-=) (map as [0..]) (take (n+number_of_guesses) guesses))))

someFunc :: IO ()
someFunc = do
   let a = fromList 3 4 [1,6,2,8,9,3,9,0,1,5,2,7]
   let Right ra = rref a
   print (getDiag ra)
   print (getCol 4 ra)


{- |
Description :  Boundary mutation model Markov process
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

Maintainer  :  dominik.schrempf@gmail.com
Stability   :  unstable
Portability :  non-portable (not tested)

This module defines the transition rate matrix of the boundary mutation model.

-}

module BndModel
  ( BMModel(..)
  , MutModel
  , Heterozygosity
  , createBMM
  , createBMMNormalized
  , scaleTreeToBMM
  , getBMMInfoStr
  ) where

import           BndState
import           Data.Maybe
import           DNAModel              hiding (rateMatrix)
import           Numeric.LinearAlgebra
import           RateMatrix            hiding (State)
import           RTree
import           Tools

-- | The boundary mutation models uses an underlying mutation model.
type MutModel = DNAModel

-- | Define a heterozygosity to make function definitions clearer.
type Heterozygosity = Double

-- | A boundary mutation model is defined by its rate matrix. However, it is
-- convenient to keep the underlying mutation model, the population size the
-- heterozygosity and the stationary distribution at hand.
data BMModel = BMModel { bmmRateMatrix     :: RateMatrix
                       , bmmMutModel       :: MutModel
                       , bmmPopSize        :: PopSize
                       , bmmHeterozygosity :: Heterozygosity
                       , bmmStationaryDist :: StationaryDist }

-- Get the mutation rate from one allele to another.
mutRate :: MutModel -> Allele -> Allele -> Double
mutRate (DNAModel m _) a b = m ! i ! j
  where i = fromEnum a
        j = fromEnum b

-- Get state frequency for specific allele.
stationaryFreq :: StationaryDist -> Allele -> Double
stationaryFreq f a = f ! fromEnum a

-- Calculate the rate of frequency shifts.
moranCoef :: PopSize -> Int -> Double
moranCoef n i = iD * (nD - iD) / nD
  where iD = fromIntegral i
        nD = fromIntegral n

-- The transition rate from one boundary state to another.
rate :: MutModel -> State -> State -> Double
rate m s t
  | not $ connected s t = 0.0
  | otherwise           = rate' s t
  where rate' (Ply n i _ _) _ = moranCoef n i
        rate' (Bnd _ a) (Ply _ _ b c)
          | a == b    = mutRate m a c
          | a == c    = mutRate m a b
          | otherwise = error "Cannot compute rate between states."
        rate' _ _ = error "Cannot compute rate between states."

-- The transition rate from one state (index) to another.
rateById :: MutModel -> PopSize -> Int -> Int -> Double
rateById m n i j = rate m s t
  where s = idToState n i
        t = idToState n j

-- The build function (see below) has a weird way of assigning entries to
-- indices. The indices have to be the same data type as the entries. This is
-- just a helper function that changes the indices from Double to Int.
rateByDouble :: MutModel -> PopSize -> Double -> Double -> Double
rateByDouble m n x y = rateById m n (round x) (round y)

-- The rate matrix of the boundary mutation model. The dimension with four
-- alleles will be 4 + 6*(N-1).
rateMatrix :: MutModel -> PopSize -> RateMatrix
rateMatrix m n = setDiagonal $ build (s,s) (rateByDouble m n)
  where s = stateSpaceSize n

-- Normalize the rate matrix such that on average one event happens per unit
-- time.
normalizedRateMatrix :: MutModel -> PopSize -> RateMatrix
normalizedRateMatrix m n = normalizeRates f' m'
  where m' = rateMatrix m n
        -- f  = dnaModelSpecGetStateFreqVec (dnaModelSpec m)
        f' = stationaryDist m n

-- Calculate the heterozygosity at stationarity.
theta :: MutModel -> Heterozygosity
theta (DNAModel m s) = f <.> rDiagZero #> f
  where f         = dnaModelSpecGetStateFreqVec s
        e         = toExchMatrix m f
        (r, _)    = matrixSeparateSymSkew e
        -- The summation excludes the diagonal (a /= b).
        rDiagZero = matrixSetDiagToZero r

-- The normalization constant of the stationary distribution.
norm :: MutModel -> PopSize -> Double
norm m n = 1.0 + harmonic (n-1) * theta m

-- Get entries of the stationary measure (not normalized) for a boundary model state.
stationaryMeasEntry :: MutModel -> State -> Double
stationaryMeasEntry m s =
  let f   = dnaModelSpecGetStateFreqVec (dnaModelSpec m) in
  case s of
    (Bnd _ a)     -> f ! fromEnum a
    (Ply n i a b) -> piA * mAB / (fromIntegral n - fromIntegral i)
                                      + piB * mBA / fromIntegral i
      where piA = stationaryFreq f a
            mAB = mutRate m a b
            piB = stationaryFreq f b
            mBA = mutRate m b a

stationaryMeasEntryById :: MutModel -> PopSize -> Int -> Double
stationaryMeasEntryById m n i = stationaryMeasEntry m s
  where s = idToState n i

stationaryMeasEntryByDouble :: MutModel -> PopSize -> Double -> Double
stationaryMeasEntryByDouble m n x = stationaryMeasEntryById m n (round x)

-- Get the stationary distribution of a boundary mutation model.
stationaryDist :: MutModel -> PopSize -> StationaryDist
stationaryDist m n = scale (1.0/norm m n) $
  build s (stationaryMeasEntryByDouble m n)
  where s = stateSpaceSize n

-- Normalize the mutation coefficients such that the heterozygosity matches a
-- given level (see Eq. 12.14 in my thesis).
normalizeToTheta :: MutModel -> PopSize -> Heterozygosity -> MutModel
normalizeToTheta mo@(DNAModel m _) n h =
  mo { dnaRateMatrix = scale (h / (t * (1.0 - c * h))) m }
  where
    -- The heterozygosity of the boundary mutation model.
    t = theta mo
    -- The branch length multiplicative factor introduced by the coalescent.
    c = harmonic (n-1)

-- | Create a boundary mutation model using the minimal number of necessary
-- ingredients.
createBMM :: DNAModel -> PopSize -> Heterozygosity -> BMModel
createBMM m n h = BMModel rm m' n h f
  where m' = normalizeToTheta m n h
        rm = normalizedRateMatrix m' n
        f  = stationaryDist m n

-- | Create a boundary mutation model without providing a heterozygosity. This
-- means that the mutation model has to be normalized already.
createBMMNormalized :: DNAModel -> PopSize -> BMModel
createBMMNormalized m n = BMModel rm m n h f
  where rm = normalizedRateMatrix m n
        h  = theta m
        f  = stationaryDist m n

-- | The branch lengths of threes in the boundary mutation model are not
-- measured in average number of substitutions per site but in average number of
-- mutations or frequency shifts per site. The conversion factor is just the
-- square of the population size. This function converts the branch lengths of a
-- tree.
scaleTreeToBMM :: BMModel -> RTree a Double -> RTree a Double
scaleTreeToBMM bmm = fmap (* nSq)
  where nSq = fromIntegral (bmmPopSize bmm) ** 2

-- | Report the boundary mutation model specifications.
getBMMInfoStr :: BMModel
              -> Maybe Double
              -> Maybe [Double]
              -> String
getBMMInfoStr bmm ma mrs = unlines $
  [ "Population size: " ++ show (bmmPopSize bmm)
  , "Heterozygosity: " ++ show (bmmHeterozygosity bmm)
  , "Mutation model:"
  , getDNAModelInfoStr (bmmMutModel bmm)
  , "Gamma rate heterogeneity: " ++ show (isJust ma) ] ++
  gammaShape ++ gammaMeans
  where
    gammaShape = maybe [] (\a -> ["Shape parameter: " ++ show a]) ma
    gammaMeans = maybe [] (\rs -> ["This corresponds to uniformly distributed rates: " ++ show rs]) mrs