{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLists       #-}
{-# LANGUAGE ViewPatterns          #-}

-- | Main module
module SatOpt (main) where
--module SatOpt (module SatOpt) where

import           SatOpt.FFT                    as SatOpt
import           SatOpt.Optimize               as SatOpt
import           SatOpt.Render                 as SatOpt
import           SatOpt.Utility                as SatOpt
-- GENERATE: import New.Module as SatOpt

import           Data.Colour                   (blend)
import           Data.Colour.RGBSpace
import           Data.Colour.RGBSpace.HSL
import           Data.Complex
import           Data.List
import           Data.Packed.Matrix            (mapMatrix)
import           Debug.Trace                   (trace)
import           Numeric.LinearAlgebra.HMatrix hiding (disp, matrix, size)
import qualified Numeric.LinearAlgebra.HMatrix as HM

data Field2D a = FMat  (Matrix a)
               | FFT   (Matrix a)
               | FDisp (Matrix a)
               | FFunc (ℝ -> ℝ -> a)

data HSource = HPointSources  [(ℝ, ℝ, ℂ)]
             | HCircleSources ℝ [ℂ]
             | HSum [HSource]

hsrcToField :: HSource -> Field2D ℂ
hsrcToField (HPointSources ps)   = FFunc $ pointSources ps
hsrcToField (HCircleSources r p) = FFunc $ circleOfSources r p
hsrcToField (HSum [h])           = hsrcToField h
hsrcToField (HSum (h:hs))        = fieldSum (hsrcToField h) (hsrcToField (HSum hs))

fieldSum :: Field2D ℂ -> Field2D ℂ -> Field2D ℂ
fieldSum (FMat m1)  (FMat m2)  = FMat (m1 + m2)
fieldSum (FFunc f1) (FFunc f2) = FFunc (\x y -> f1 x y + f2 x y)
fieldSum (FMat m1)  (FFunc f2) = fieldSum (FFunc f2) (FMat m1)
fieldSum (FFunc f1) (FMat m2)  = fieldSum (matField (FFunc f1)) (FMat m2)

data HDisp = HConstDisp ℂ
           | HFuncDisp (ℝ -> ℝ -> ℂ)

hdspToField :: HDisp -> Field2D ℂ
hdspToField (HConstDisp p) = constF p
hdspToField (HFuncDisp f) = FFunc f

data HelmConf = HelmConf { helmsrc :: HSource
                         , helmdsp :: HDisp
                         }

data State a = State { amplitude  :: Field2D ℂ
                     , dispersion :: Field2D ℂ
                     , sources    :: Field2D ℂ
                     , time       :: ℝ
                     , freq       :: ℝ
                     , conf       :: a
                     , configure  :: State a -> State a
                     , concrete   :: Bool
                     }

constF :: Element a => a -> Field2D a
constF x = matField (FFunc (\_ _ -> x))

defState :: State a
defState = State { amplitude  = constF (0 :+ 0)
                 , dispersion = constF (1 :+ 0.1)
                 , sources    = constF (0 :+ 0)
                 , time       = 0
                 , freq       = 1
                 , conf       = undefined
                 , configure  = const defState
                 , concrete   = False
                 }

fromFieldC :: Field2D ℂ -> Matrix ℂ
fromFieldC (FFunc f) = sample2DM f [0, 1 .. size - 1] [0, 1 .. size - 1]
fromFieldC (FMat m)  = m
fromFieldC (FFT m)   = ifft2dM m
fromFieldC (FDisp m) = sqrt $ (recip m) - (square normm)

matField :: Element a => Field2D a -> Field2D a
matField (FFunc f) = FMat $ sample2DM f [0, 1 .. size - 1] [0, 1 .. size - 1]
matField x         = x

fftMat :: Field2D ℂ -> Matrix ℂ
fftMat (FFT m)  = m
fftMat (FMat m) = fft2dM m
fftMat f        = fftMat $ matField f

dispMat :: Field2D ℂ -> Matrix ℂ
dispMat (FDisp m) = m
dispMat (FMat m)  = recip $ square m + square normm
dispMat f         = dispMat $ matField f

makeConcrete :: (State a) -> (State a)
makeConcrete st@(State { concrete = True }) = st
makeConcrete (st@(State { dispersion = d, sources = s, time = t, freq = f }))
  = st { amplitude = ampl
       , dispersion = FDisp disp
       , sources = FFT srcs
       , concrete = True }
  where
    disp = dispMat $ matField d
    srcs = fftMat $ matField s
    ampl = FMat $ timeShift f t $ ifft2dM $ srcs * disp

phaseShift :: Matrix ℝ -> Matrix ℂ -> Matrix ℂ
phaseShift p m = mags * zipMatrixWith (:+) (cos angles) (sin angles)
  where
    angles = p + (mapMatrix realPart m `arctan2` mapMatrix imagPart m)
    mags = m * conj m

timeShift :: ℝ -> ℝ -> Matrix ℂ -> Matrix ℂ
timeShift f t m = case matField $ constF (tau * f * t) of
                   FMat p -> phaseShift p m
                   _      -> error "matField output invalid"

size :: Num a => a
size = 256

tau :: ℝ
tau = 2 * pi

rootsOfUnity :: Int -> [ℂ]
rootsOfUnity i = take i $ iterate (rotate (tau / fromIntegral i)) (1 :+ 0)

pointsOnCircle :: Int -> ℝ -> [(ℝ, ℝ)]
pointsOnCircle i r = map (fromC . scaleC r) $ rootsOfUnity i
  where
    fromC (a :+ b) = (a, b)
    scaleC s c = (s :+ 0) * c

-- | Take a position with respect to the center and move its origin to the top right
center :: Floating a => a -> a
center x = x - (size / 2)

pointSources :: [(ℝ, ℝ, ℂ)] -> ℝ -> ℝ -> ℂ
pointSources ps x y = sum $ map (\(cx, cy, s) -> s * delta cx cy x y) ps

circleOfSources :: ℝ -> [ℂ] -> ℝ -> ℝ -> ℂ
circleOfSources rad cs = pointSources $ zipCs $ pointsOnCircle (length cs) rad
  where
    zipCs = map (\(z, (x, y)) -> (x, y, z)) . zip cs

--const $ FFunc testFunc

-- testFunc :: HelmConf -> Field2D ℂ
-- testFunc (HelmConf (HPointSources ps) _)   = FFunc $ pointSources ps
-- testFunc (HelmConf (HCircleSources r p) _) = FFunc $ circleOfSources r p
  -- where
  --   pol m p = mkPolar m ((pi/180)*p)
-- testFunc = pointSources [ (0.48, 0.48, 0.5 :+ 0)
--                         , (0.52, 0.52, (-0.5) :+ 0)
--                         ]

-- 0.1 [ pol 1 0
--                                , pol 1 60
--                                , pol 1 120
--                                , pol 1 120
--                                , pol 1 60
--                                , pol 1 0
--                                ]

normm :: Matrix ℂ
--im, jm :: Matrix ℂ
--im = sample2DM (\x _ -> x :+ 0) [0, 1 .. size - 1] [0, 1 .. size - 1]
--jm = sample2DM (\_ y -> y :+ 0) [0, 1 .. size - 1] [0, 1 .. size - 1]
normm = sample2DM (\x y -> norm x y :+ 0) [0, 1 .. size - 1] [0, 1 .. size - 1]

eps :: ℝ
eps = 5

delta :: ℝ -> ℝ -> ℝ -> ℝ -> ℂ
delta ((*size) -> cx) ((*size) -> cy) ((*size) -> x) ((*size) -> y)
  | norm (center x - cx) (center y - cy) < size * eps    = 1
  | otherwise                                                       = 0

rendCmp_ :: ℂ -> (ℝ, ℝ, ℝ)
rendCmp_ (polar -> (m, p)) = uncurryRGB (\x y z -> (x, y, z)) $ hsl a 1 m
  where
--    a = if p < 0 then 90 else 0
    a = (p + pi) * (180 / pi)

rendCmp :: Matrix ℂ -> (Matrix ℝ, Matrix ℝ, Matrix ℝ)
rendCmp = matrify . unzip3 . map (unzip3 . map rendCmp_) . toLists
  where
    matrify (x, y, z) = (fromLists x, fromLists y, fromLists z)

rotate :: ℝ -> ℂ -> ℂ
rotate r = uncurry mkPolar . (\(m, p) -> (m, p + r)) . polar

visual :: (ℝ, ℝ, ℝ, ℝ) -> ℂ -> ℂ
visual (xa, xb, ya, yb) (polar -> (m, p))
  = mkPolar (((yb - ya) * (m - xa)/(xb - xa)) + ya) p

visualM :: Matrix ℂ -> Matrix ℂ
visualM cs@(fromComplex -> (r, i)) = mapMatrix (visual (small, big, 0, 1)) cs
  where
    !mags  = flatten $ sqrt $ (r * r) + (i * i)
    !small = minElement mags
    !big   = maxElement mags

updateState :: State a -> State a
--updateState = id
updateState s@State { time = t } = s { time = t + 0.01, concrete = False }

data RenderMode = RendAmp | RendSrc | RendDsp | RendAmpMag | RendAmpPhs

renderState :: RenderMode -> State a -> RGBTrips
renderState r = case r of
  RendAmp    -> rend . concAmp
  RendAmpMag -> rend . concAmpMag
  RendAmpPhs -> rend . concAmpPhs
  RendSrc    -> rend . concSrc
  RendDsp    -> rend . concDsp
  where
  wrap f = fromFieldC . f . makeConcrete
  rend =  rendCmp . visualM
  concAmp    = wrap amplitude
  concAmpMag = mag . fromComplex . concAmp
  concAmpPhs = phs . fromComplex . concAmp
  concSrc    = wrap sources
  concDsp    = wrap dispersion
  mag (x, y) = toComplex $ ((x*x) + (y*y), y - y)
  phs (x, y) = toComplex $ ((y `arctan2` x), y - y)

changeRenderMode :: RenderMode -> Conf (State a) -> Conf (State a)
changeRenderMode r c = c { render = renderState r }

getConf :: Conf (State HelmConf) -> HelmConf
getConf (Conf { state = (State { conf = cn })}) = cn

changeConf :: (HelmConf -> HelmConf) -> Conf (State HelmConf) -> Conf (State HelmConf)
changeConf h c@(Conf { state = s@(State { configure = f, conf = cn })})
  = c { state = configure s' s' }
  where
    s' = s { conf = h cn }

relPhase (polar -> (_, p1)) (polar -> (_, p2)) = abs $ p1 - p2

--getPhase c@(HelmConf { helmsrc = HCircleSources _ [c1, c2] }) = Just $ relPhase c1 c2
--getPhase _ = Nothing
getPhase c@(HelmConf { helmsrc = HPointSources [_, (x, _, _)] }) = Just x
getPhase _ = Nothing

changePhase :: Double -> Conf (State HelmConf) -> Conf (State HelmConf)
changePhase p = changeConf shift
  where
    shift c@(getPhase -> m) = case m of
                               Just ph -> c { helmsrc = testSource $ ph + p }
                               Nothing -> c

keyState :: KMState -> Conf (State HelmConf) -> IO (Conf (State HelmConf))
keyState ""  c = return c
keyState "d" c = return $ changeRenderMode RendDsp c
keyState "s" c = return $ changeRenderMode RendSrc c
keyState "a" c = return $ changeRenderMode RendAmp c
keyState "m" c = return $ changeRenderMode RendAmpMag c
keyState "p" c = return $ changeRenderMode RendAmpPhs c
keyState "w" c = do
  let p = changePhase 0.5 c
  putStrLn $ "Shifted by 0.5 to " ++ show (getPhase $ getConf c)
  return p
keyState "q" c = do
  let p = changePhase (-0.5) c
  putStrLn $ "Shifted by -0.5 to " ++ show (getPhase $ getConf c)
  return p
keyState x   c = return $ trace x c

pol m p = mkPolar m (p * pi / 180)

testSource :: Double -> HSource
testSource p = HPointSources [ (size/2, size/2, (-1) :+ 0)
                             , (p, p, 1 :+ 0)
                             ]
-- testSource p = HSum [ HCircleSources 0.2  [ pol 1 0
--                                           , pol 1 0
--                                           , pol 1 0
--                                           , pol 1 0
--                                           ]
--                     , HCircleSources 0.15 [ pol 1 0
--                                           , pol 1 90
--                                           , pol 1 180
--                                           , pol 1 270
--                                           ]
--                     ]

dspfun :: ℝ -> ℝ -> ℂ
dspfun x y
  | outbox (center x) (center y) = 1.2 :+ 0
  | otherwise                    = 1 :+ 0
  where
    test = 1
    outbox a b
      | a < (-32) = True
      | a > 32    = True
      | b < (-32) = True
      | b > 32    = True
      | otherwise = False

testConf :: Double -> HelmConf
testConf p = HelmConf { helmsrc = testSource p
                      , helmdsp = HFuncDisp dspfun
                      }

-- testConf = HelmConf { helmsrc = HCircleSources 0.1 [ pol 1 0
--                                                    , pol 1 0
--                                                    , pol 1 0
--                                                    , pol 1 0
--                                                    , pol 1 0
--                                                    , pol 1 0 ]
--                     , helmdsp = HConstDisp 0.001
--                     }
--   where
--     pol m p = mkPolar m (p * pi / 180)

testConfigure st@(State { conf = HelmConf { helmdsp = d, helmsrc = s }})
  = st { sources = hsrcToField s
       , dispersion = hdspToField d
       , concrete = False
       }

initState :: State HelmConf
initState = defState { freq = 3
                     , conf = testConf 0
                     , configure = testConfigure
                     }

config :: Conf (State HelmConf)
config = changeConf id
         Conf { keyBinds = keyState
              , state    = initState
              , evolve   = updateState
              , render   = renderState RendAmp
              , canvas   = (size, size)
              }

-- | TODO
main :: IO ()
--main = optimize
main = runDisp config


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

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Transpose
Description : Layer that transposes its input
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}

module Grenade.Layers.Transpose
  (
  -- * Layer Definition
    Transpose (..)

  -- * Helper function
  , initTranspose
  ) where

import           Control.DeepSeq                   (NFData (..))

import           Data.Kind                         (Type)
import           Data.Maybe                        (fromJust)
import           Data.Proxy
import           Data.Serialize
import           GHC.TypeLits

import qualified Numeric.LinearAlgebra             as LA
import           Numeric.LinearAlgebra.Static      (R)
import qualified Numeric.LinearAlgebra.Static      as H

import           Grenade.Core
import           Grenade.Layers.Internal.Transpose
import           Grenade.Onnx

 -- todo: we can probably use a type family to represent this much better
data Transpose :: Nat
               -> Shape
               -> Shape
               -> Type where
  Transpose  :: ( KnownNat dimensions )
          => !(R dimensions)
          -> Transpose dimensions input output

instance Show (Transpose dimensions input output) where
  show (Transpose mat) = "Transpose " ++ show mat

instance UpdateLayer (Transpose dimensions input output) where
  type Gradient (Transpose dimensions input output) = ()
  runUpdate _ x _  = x
  reduceGradient _ = ()

instance ( KnownNat dimensions ) => RandomLayer (Transpose dimensions input output) where
  createRandomWith _ _ = pure initTranspose

-- | Initialize a tranpose layer that is equivalent to the identity function
initTranspose :: forall dimensions input output. ( KnownNat dimensions )
              => Transpose dimensions input output
initTranspose =
  let ds    = fromIntegral $ natVal (Proxy :: Proxy dimensions)
      perms = H.fromList [1..ds] :: R dimensions
  in Transpose perms

instance ( KnownNat dimensions ) => Serialize (Transpose dimensions input output) where
  put (Transpose perms) = putListOf put . LA.toList . H.extract $ perms
  get                   = do
    perms <- maybe (fail "Vector of incorrect size") return . H.create . LA.fromList =<< getListOf get
    return $ Transpose perms

instance ( KnownNat i, KnownNat j, KnownNat k, KnownNat l, KnownNat a, KnownNat b, KnownNat c, KnownNat d )
  => Layer (Transpose 4 ('D4 i j k l) ('D4 a b c d)) ('D4 i j k l) ('D4 a b c d) where

  type Tape (Transpose 4 ('D4 i j k l) ('D4 a b c d)) ('D4 i j k l) ('D4 a b c d) = ()

  runForwards (Transpose perms) (S4D x)
    = let n  = fromIntegral $ natVal (Proxy :: Proxy i)
          c  = fromIntegral $ natVal (Proxy :: Proxy j)
          h  = fromIntegral $ natVal (Proxy :: Proxy k)
          w  = fromIntegral $ natVal (Proxy :: Proxy l)
          x' = H.extract x
          perms' = H.extract perms
          r  = transpose4d [n, c, h, w] perms' x'
      in  ((), S4D . fromJust . H.create $ r)

  runBackwards = undefined

instance OnnxOperator (Transpose dimensions input output) where
  onnxOpTypeNames _ = ["Transpose"]

instance OnnxLoadable (Transpose 4 input output) where
  loadOnnxNode _ node = readIntsAttribute "perm" node >>= formatPerm
    where
      formatPerm :: [Int] -> Either OnnxLoadFailure (Transpose 4 input output)
      formatPerm [_, _, n, c, h, w]
        = let centered = map (\x -> fromIntegral $ x - 2) [n, c, h, w]
              perms :: R 4 = H.fromList centered
          in  return $ Transpose perms
      formatPerm _ = loadFailureReason "Permutation shape incorrect of Transpose"

instance NFData (Transpose dims input output) where
  rnf (Transpose perms) = rnf perms


{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE KindSignatures       #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE QuasiQuotes          #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE RecordWildCards      #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TemplateHaskell      #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeInType           #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

import           Control.Category
import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Data.Foldable
import           Data.Kind
import           Data.List hiding                      ((\\))
import           Data.Maybe
import           Data.Singletons
import           Data.Monoid
import           Data.Singletons.TypeLits
import           Data.Time.Clock
import           Options.Applicative
import           Prelude hiding                        ((.), id)
import           Statistics.Distribution.Uniform
import           System.Random.MWC
import           TensorOps.Backend.BTensor
import           TensorOps.Backend.NTensor
import           TensorOps.Learn.NeuralNet
import           TensorOps.Learn.NeuralNet.FeedForward
import           TensorOps.NatKind
import           TensorOps.Types
import           Text.PrettyPrint.ANSI.Leijen hiding   ((<>),(<$>))
import           Text.Printf
import           Type.Class.Higher.Util
import qualified Data.String.Here                      as H
import qualified TensorOps.Tensor                      as TT

netTest
    :: forall k (t :: [k] -> Type).
     ( Tensor t
     , ElemT t ~ Double
     , NFData1 t
     , NFData (t '[FromNat 1])
     , NFData (t '[FromNat 2])
     )
    => Proxy t
    -> Double
    -> Int
    -> [Integer]
    -> GenIO
    -> IO String
netTest _ rate n hs g = withSingI (sFromNat @k (SNat @1)) $
                        withSingI (sFromNat @k (SNat @2)) $ do
    ((inps,outs),tp) <- time $ do
      inps :: [t '[FromNat 2]] <- replicateM n (genRand (uniformDistr (-1) 1) g)
      let outs :: [t '[FromNat 1]]
          outs = flip map inps $ \v -> TT.konst $
                   if v `inCircle` (TT.konst 0.33, 0.33)
                        || v `inCircle` (TT.konst (-0.33), 0.33)
                     then 1
                     else 0
      evaluate . force $ (inps, outs)
    printf "Generated test points (%s)\n" (show tp)
    net0 :: Network t (FromNat 2) (FromNat 1)
            <- genNet (hs `zip` repeat actLogistic) actLogistic g
    let trained = foldl' trainEach net0 (zip inps outs)
          where
            trainEach :: SingI o
                      => Network t i o
                      -> (t '[i], t '[o])
                      -> Network t i o
            trainEach nt (i, o) = nt `deepseq` trainNetwork squaredError rate i o nt
    (trained', tn) <- time $ return trained
    printf "Network trained (%s)\n" (show tn)
    let outMat = [ [ render . TT.unScalar . join TT.dot . runNetwork trained' $
                       fromJust (TT.fromList [x / 25 - 1,y / 10 - 1])
                   | x <- [0..50] ]
                 | y <- [0..20] ]
        render r | r <= 0.2  = ' '
                 | r <= 0.4  = '.'
                 | r <= 0.6  = '-'
                 | r <= 0.8  = '='
                 | otherwise = '#'
    return $ unlines outMat
  where
    inCircle
        :: SingI n
        => t '[n]
        -> (t '[n], Double)
        -> Bool
    v `inCircle` (o, r) = let d = TT.zip (-) v o
                          in  TT.unScalar (d `TT.dot` d) <= r**2

data Opts = O { oRate    :: Double
              , oSamples :: Int
              , oNetwork :: [Integer]
              , oTests   :: [TestT]
              }

opts :: Parser Opts
opts = O <$> option auto
               ( long "rate" <> short 'r' <> metavar "STEP"
              <> help "Neural network learning rate"
              <> value 1 <> showDefault
               )
         <*> option auto
               ( long "samps" <> short 's' <> metavar "COUNT"
              <> help "Number of samples to train the network on"
              <> value 50000 <> showDefault
               )
         <*> option auto
               ( long "layers" <> short 'l' <> metavar "LIST"
              <> help "List of hidden layer sizes"
              <> value [12,8] <> showDefault
               )
         <*> (nub <$> (some (argument readBackend (metavar "BACKEND")))
               <|> pure [TTBLAS TBHMat]
             )

main :: IO ()
main = withSystemRandom $ \g -> do
    O{..} <- execParser $ info (helper <*> opts)
        ( fullDesc
       <> header "tensor-ops-dots - train neural nets with tensor-ops"
       <> progDescDoc (Just d)
        )

    printf "rate: %f | samps: %d | layers: %s\n" oRate oSamples (show oNetwork)

    forM_ oTests $ \t -> do
      printf "Training %s network ...\n" (ttLong t)
      let tester :: Double -> Int -> [Integer] -> GenIO -> IO String
          tester = case t of
            TTNested TVList -> netTest (Proxy @NTensorL)
            TTNested TVVect -> netTest (Proxy @NTensorV)
            TTBLAS   TBHMat -> netTest (Proxy @(BTensorV (HMat Double)))
      putStrLn =<< tester oRate oSamples oNetwork g
  where
    d :: Doc
    d = string [H.here|
Trains a feed-forward neural network on a 2D classifier using tensor-ops
machinery, with the given backends.  (If none provided, backend defaults to 'b')
|]
     <$$> mempty
     <$$> string "Backends:"
     <$$> vsep [ string $ printf "- %s: %s" (ttShort t) (ttLong t)
               | t <- allTests ]

time
    :: NFData a
    => IO a
    -> IO (a, NominalDiffTime)
time x = do
    t1 <- getCurrentTime
    y  <- evaluate . force =<< x
    t2 <- getCurrentTime
    return (y, t2 `diffUTCTime` t1)

data TestT = TTNested TestV
           | TTBLAS   TestB
           deriving (Show, Eq, Ord)
data TestV = TVList
           | TVVect
           deriving (Show, Eq, Ord)
data TestB = TBHMat
           deriving (Show, Eq, Ord)

allTests :: [TestT]
allTests = ([TTNested] <*> [TVList, TVVect]) ++ [TTBLAS TBHMat]

readBackend :: ReadM TestT
readBackend = eitherReader $ \s -> case s of
    "nl" -> Right $ TTNested TVList
    "nv" -> Right $ TTNested TVVect
    "b"  -> Right $ TTBLAS   TBHMat
    o    -> Left  $ "Unknown backend: " ++ o

ttShort
    :: TestT
    -> String
ttShort = \case
    TTNested v -> 'n' : tvShort v
    TTBLAS   b -> 'b' : tbShort b
  where
    tvShort = \case
      TVList -> "l"
      TVVect -> "v"
    tbShort = \case
      TBHMat -> ""

ttLong
    :: TestT
    -> String
ttLong = \case
    TTNested v -> printf "Nested (%s)" (tvLong v)
    TTBLAS   b -> printf "BLAS (%s)" (tbLong b)
  where
    tvLong = \case
      TVList -> "List"
      TVVect -> "Vector"
    tbLong = \case
      TBHMat -> "HMatrix"


{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances #-}

module Math.Matrix
       ( DimVector (..)
       , DimMatrix (..)
       , GPConstraint
       , withMat
       , toMatrix
       , mulM
       , transposeM
       , mapMM
       , mapDiagonalM
       , invSM
       , substractMeanM
       , (+^^)
       , (-^^)
       , (*^^)
       , cholM
       , foldAllSM
       , sumAllSM
       , trace2SM
       , toDimMatrix
       , detSM
       , toMatrixM'
       , zipWithDim
       , zipWithArray
       , identM
       , delayMatrix
       , linearSolveM
       , vectorLength
       , randomMatrixD
       , matrixRowsNum
       , matrixColsNum
       , pinvSM
       , mulPM
       , eigSHM
       )
       where

import Universum hiding (Vector, transpose, map, natVal,
                         zipWith)

import           GHC.TypeLits hiding (someNatVal)

import           Data.Array.Repa
import           Data.Random.Normal (normals)
import           Data.Vector.Unboxed.Base (Unbox)
import           Data.Vinyl.TypeLevel (AllConstrained)
--import           Numeric.Dimensions
import           Numeric.LinearAlgebra.Repa hiding (Matrix, Vector)
import           System.Random (Random, RandomGen)

import PCA.Types
import PCA.Util hiding (randomMatrixD, zipWithArray)

newtype DimVector r (n :: Nat) a
  = DimVector { runVector :: Vector r a }
  deriving Eq

newtype DimMatrix r (y :: Nat) (x :: Nat) a
  = DimMatrix { getInternal :: Matrix r a}

instance Functor (DimVector D n) where
  fmap f (DimVector vector) = DimVector (smap f vector)

withMat
  :: Matrix D a
  -> (forall x y. (KnownNat x, KnownNat y) => DimMatrix D x y a -> k)
  -> k
withMat m f =
  let (Z :. x :. y) = extent m
  in
  case someNatVal (fromIntegral x) of
    SomeNat (Proxy :: Proxy m) -> case someNatVal (fromIntegral y) of
      SomeNat (Proxy :: Proxy n) -> f (DimMatrix @_ @m @n m)

mulM
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , Numeric a
  , x1 ~ y2
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y1 x2 a
mulM (DimMatrix m1) (DimMatrix m2) = DimMatrix $ delay $ m1 `mulS` m2

transposeM
  :: (KnownNat y, KnownNat x)
  => DimMatrix D y x a
  -> DimMatrix D x y a
transposeM (DimMatrix m) = DimMatrix $ transpose m

mapMM
  ::
  ( KnownNat y
  , KnownNat x
  , Unbox a
  , Unbox b
  )
  => (a -> b)
  -> DimMatrix D y x a
  -> DimMatrix D y x b
mapMM f (DimMatrix m) =  DimMatrix $ map f m

mapDiagonalM
  ::
  ( KnownNat y
  , KnownNat x
  , Unbox a
  )
  => (a -> a)
  -> DimMatrix D y x a
  -> DimMatrix D y x a
mapDiagonalM f (DimMatrix m) = DimMatrix $ mapDiagonal f m

invSM
  ::
  ( KnownNat y
  , KnownNat x
  , Field a
  , Numeric a
  , y ~ x
  )
  => DimMatrix D y x a
  -> DimMatrix D y x a
invSM (DimMatrix m) = DimMatrix $ delay $ invS m

substractMeanM
  ::
  ( KnownNat y
  , KnownNat x
  )
  => DimMatrix D y x Double
  -> DimMatrix D y x Double
substractMeanM (DimMatrix m) = DimMatrix $ substractMean m

infixl 6 +^^, -^^
infixl 7 *^^

(+^^)
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , x1 ~ x2
  , y1 ~ y2
  , Num a
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y2 x2 a
(+^^) (DimMatrix m1) (DimMatrix m2) = DimMatrix $ m1 +^ m2

(-^^)
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , x1 ~ x2
  , y1 ~ y2
  , Num a
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y2 x2 a
(-^^) (DimMatrix m1) (DimMatrix m2) = DimMatrix $ m1 -^ m2

(*^^)
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , x1 ~ x2
  , y1 ~ y2
  , Num a
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y2 x2 a
(*^^) (DimMatrix m1) (DimMatrix m2) = DimMatrix $ m1 *^ m2

cholM
  ::
  ( KnownNat y
  , KnownNat x
  , Field a
  , y ~ x
  )
  => DimMatrix D y x a
  -> DimMatrix D y x a
cholM (DimMatrix m) = DimMatrix $ delay $ chol $ trustSym $ computeS m

sumAllSM
  ::
  ( KnownNat y
  , KnownNat x
  , Num a)
  => DimMatrix D y x a
  -> a
sumAllSM (DimMatrix m) = sumAllS m

foldAllSM
  :: (KnownNat y, KnownNat x)
  => (Double -> Double -> Double)
  -> Double
  -> DimMatrix D y x Double
  -> Double
foldAllSM f initValue (DimMatrix m) = foldAllS f initValue m

detSM
  :: (KnownNat y, KnownNat x)
  => DimMatrix D y x Double
  -> Double
detSM (DimMatrix m) = detS m

trace2SM
  :: (KnownNat x, KnownNat y)
  => DimMatrix D x y Double
  -> Double
trace2SM (DimMatrix m) = trace2S $ computeS m

toDimMatrix
  ::
  ( Source r a
  , KnownNat m
  , KnownNat n
  )
  => DimVector r m a
  -> Int
  -> DimMatrix D m n a
toDimMatrix (DimVector arr) desiredSize =
  DimMatrix (toMatrix arr desiredSize)

toMatrixM'
  :: (Source r a, KnownNat n, KnownNat m)
  => DimVector r n a
  -> DimMatrix D n m a
toMatrixM' (DimVector arr) =
  DimMatrix $ fromFunction (Z :. dimension :. 1) generator
  where
    dimension = size . extent $ arr
    generator (Z :. rows :. _) = linearIndex arr rows

zipWithDim
  ::
  ( Source r1 a
  , Source r2 b
  , KnownNat m
  , KnownNat n
  )
  => (a -> b -> c)
  -> DimMatrix r1 m n a
  -> DimMatrix r2 m n b
  -> DimMatrix D m n c
zipWithDim f (DimMatrix mat1) (DimMatrix mat2) =
  DimMatrix $ zipWith f mat1 mat2

zipWithArray
  ::
  ( KnownNat n
  , KnownNat m
  )
  => (a -> b -> c)
  -> DimVector D n a
  -> DimMatrix D n m b
  -> DimMatrix D n m c
zipWithArray f (DimVector array1) (DimMatrix array2) =
  DimMatrix $ zipWith f (toMatrix array1 n) array2
  where
    (Z :. n :. _) = extent array2

identM
  :: forall m n a.
  ( KnownNat m
  , KnownNat n
  , GPConstraint a
  , m ~ n
  )
  => DimMatrix D m n a
identM =
  let dim = fromEnum $
            natVal @m @Proxy Proxy in DimMatrix $
            identD dim

delayMatrix
  ::
  ( KnownNat m
  , KnownNat n
  , Source r a
  )
  => DimMatrix r m n a
  -> DimMatrix D m n a
delayMatrix (DimMatrix matrix) = DimMatrix (delay matrix)

linearSolveM
  ::
  ( Field a
  , AllConstrained KnownNat '[m, n]
  )
  => DimMatrix D m m a
  -> DimMatrix D m n a
  -> Maybe (DimMatrix D m n a)
linearSolveM (DimMatrix mat1) (DimMatrix mat2) =
  case linearSolveS mat1 mat2 of
    Nothing  -> Nothing
    Just sol -> Just . delayMatrix . DimMatrix $ sol

vectorLength
  :: forall r m a. KnownNat m
  => DimVector r m a
  -> Int
vectorLength _ = fromEnum $ natVal (Proxy @m)

matrixRowsNum
  :: forall r m n a. KnownNat m
  => DimMatrix r m n a
  -> Int
matrixRowsNum _ = fromEnum $ natVal (Proxy @m)

matrixColsNum
  :: forall r m n a. KnownNat n
  => DimMatrix r m n a
  -> Int
matrixColsNum _ = fromEnum $ natVal (Proxy @n)

randomMatrixD
  :: forall a g m n.
  ( RandomGen g
  , Random a
  , Unbox a
  , Floating a
  , KnownNat m
  , KnownNat n
  )
  => g
  -> (Int, Int)
  -> DimMatrix D m n a
randomMatrixD gen (rows, cols) =
  let randomList = take (rows * cols) (normals gen) in
    DimMatrix . delay $ fromListUnboxed (Z :. rows :. cols) randomList

type GPConstraint a =
  ( Field a
  , Random a
  , Unbox a
  , Floating a
  , Eq a
  )
{-
hermToMatrixM
  :: Herm a
  -> DimMatrix D m n a
hermToMatrixM = undefined
-}
pinvSM
  ::
  ( KnownNat y
  , KnownNat x
  , Field a
  , Numeric a
  , y ~ x
  )
  => DimMatrix D y x a
  -> DimMatrix D y x a
pinvSM (DimMatrix m) = DimMatrix . delay $ pinvS m

mulPM
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , Numeric a
  , x1 ~ y2
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y1 x2 a
mulPM (DimMatrix m) (DimMatrix n) =
  DimMatrix . delay . runIdentity $ m `mulP` n

eigSHM
  :: (Field a, Numeric a)
  => Herm a
  -> (DimVector D m Double, DimMatrix D m m a)
eigSHM hermM =
  (DimVector $ delay $ fst out, DimMatrix $ delay $ snd out)
  where
    out = eigSH hermM


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards  #-}

module AI.DemoLayer
( LayerDefinition(..)
, Layer(..)

, Weights
, Biases
, LayerWidth
, Connectivity

--, createLayer
) where

import AI.DemoNeuron

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import Data.Binary
import Data.Binary.Put
import qualified Data.ByteString.Lazy as B
import System.Random

data LayerDefinition a = LayerDefinition { neuronType   :: a
                                         , width        :: LayerWidth
                                         , connectivity :: Connectivity
                                         }

data Layer a = Layer { neuron :: a
                     , weights :: Weights
                     , biases :: Biases
                     } deriving (Show)

type Weights = Matrix Double
type Biases = Vector Double
type LayerWidth = Int
type Connectivity = LayerWidth -> LayerWidth -> Weights

instance (Neuron a) => Binary (Layer a) where
  put Layer{..} = do put weights; put biases
  get = do weights <- get; biases <- get; return Layer{..}

-- createLayer :: (Neuron a, RandomGen g)
--               => RandomTransform -> g -> LayerDefinition a
--               -> LayerDefinition a -> Layer
-- createLayer t g li lj = Layer (neuron lj)


-- Lib.hs

module Lib
    ( symbol
    , readExp
    , spaces
    ) where

import Text.ParserCombinators.Parsec hiding (spaces)
import Numeric (readHex, readOct, readFloat)
import Data.Char (digitToInt)
import qualified Data.Complex as C
import Data.Ratio ((%))

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Complex (C.Complex Double)
             | Integer Integer
             | Real Double
             | Rational Rational
             | String String
             | Char Char
             | Bool Bool

readExp :: String -> String
readExp input =
  case parse parseExpr "lisp" input of
    Left err -> "No match: " ++ show err
    Right _  -> "Found value"

parseExpr :: Parser LispVal
parseExpr = parseAtom
        <|> parseString
        <|> parseNumber
        <|> try parseChar
        <|> parseQuoted
        <|> do _ <- char '('
               x <- try parseList <|> try parseDottedList
               _ <- char ')'
               return x
        <|> parseQuasiQuoted
        <|> parseUnquote

parseAtom :: Parser LispVal
parseAtom =
  do first <- letter <|> symbol
     rest <- many $ letter <|> digit <|> symbol
     let atom = first:rest
     return $ case atom of
                "#t" -> Bool True
                "#f" -> Bool False
                _    -> Atom atom

parseString :: Parser LispVal
parseString =
  do _ <- char '"'
     x <- many $ escapedChars <|> noneOf "\""
     _ <- char '"'
     return $ String x

parseNumber :: Parser LispVal
parseNumber = try parseComplex
          <|> try parseReal
          <|> try parseRational
          <|> try parseInteger

parseInteger :: Parser LispVal
parseInteger = parseDec
           <|> parseDec2
           <|> parseOct
           <|> parseHex
           <|> parseBin

parseOct :: Parser LispVal
parseOct = try (string "#o") >> many1 octDigit >>= return . Integer . octToDig
  where octToDig = fst . head . readOct

parseDec :: Parser LispVal
parseDec = Integer . read <$> many1 digit

parseDec2 :: Parser LispVal
parseDec2 = try (string "#d") >> Integer . read <$> many1 digit

parseHex :: Parser LispVal
parseHex = try (string "#x") >> many1 hexDigit >>= return . Integer . hexToDig
  where hexToDig = fst . head . readHex

parseBin :: Parser LispVal
parseBin =
  try (string "#b") >>
  many1 (oneOf "10") >>=
  return . Integer . binToDig
-- x0 + 2 * (x1 + 2 * (x2 + 2 * (x3 + 2 * x4)))
  where binToDig "" = 0
        binToDig s = foldr f 0 ds
          where ds      = (toInteger . digitToInt) <$> reverse s
                f x acc = x + 2 * acc

parseChar :: Parser LispVal
parseChar = parseChar1 <|> parseChar2

parseChar1 :: Parser LispVal
parseChar1 =
  do _ <- try (string "#\\")
     x <- anyChar >>= \c -> notFollowedBy alphaNum >> return c
     return . Char $ x

parseChar2 :: Parser LispVal
parseChar2 =
  try (string "#\\") >>
  try (string "newline" <|> string "space") >>= \x ->
  return . Char $ case x of
                    "newline" -> '\n'
                    "space"   -> ' '

parseReal :: Parser LispVal
parseReal =
  do x <- many1 digit
     _ <- char '.'
     y <- many1 digit
     return . Real . fst . head . readFloat $ x ++ ['.'] ++ y

parseRational :: Parser LispVal
parseRational =
  do n <- many1 digit
     _ <- char '/'
     d <- many1 digit
     return . Rational $ read d % read n

parseComplex :: Parser LispVal
parseComplex =
  do r <- try $ parseReal <|> parseDec
     _ <- char '+'
     i <- try $ parseReal <|> parseDec
     _ <- char 'i'
     return . Complex $ (toDouble r) C.:+ (toDouble i)
  where toDouble (Real f)    = realToFrac f
        toDouble (Integer f) = fromIntegral f

parseList :: Parser LispVal
parseList = List <$> sepBy parseExpr spaces

parseDottedList :: Parser LispVal
parseDottedList =
  do h <- endBy parseExpr spaces
     t <- char '.' >> spaces >> parseExpr
     return $ DottedList h t

parseQuoted :: Parser LispVal
parseQuoted =
  do _ <- char '\''
     x <- parseExpr
     return $ List [Atom "quote", x]

parseQuasiQuoted :: Parser LispVal
parseQuasiQuoted =
  do _ <- char '`'
     e <- parseExpr
     return $ List [Atom "quasiquote", e]

parseUnquote :: Parser LispVal
parseUnquote =
  do _ <- char ','
     e <- parseExpr
     return $ List [Atom "unquote", e]

symbol :: Parser Char
symbol = oneOf "!#$%&|*+-/:<=>?@^_~"

spaces :: Parser ()
spaces = skipMany1 space

escapedChars :: Parser Char
escapedChars = char '\\' >>
               oneOf "\\\"\n\r\t" >>= \x ->
               return $ case x of
                          '\\' -> x
                          '"'  -> x
                          'n'  -> '\n'
                          'r'  -> '\r'


{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
module STC.InitialDistribution where

import           Array.UnboxedArray       as AU
import           Data.Array.Repa          as R
import           Data.Complex
import           Data.List                as L
import           Data.Vector.Generic      as VG
import           Data.Vector.Storable     as VS
import           DFT.Plan
import           FourierPinwheel
import           Pinwheel.FourierSeries2D
import           STC.Convolution
import           STC.DFTArray
import           STC.Point
import           STC.Utils
import           Utils.Distribution
import           Utils.List
import           Utils.Parallel
import Filter.Utils

{-# INLINE computeInitialDistribution #-}
computeInitialDistribution ::
     Int -> Int -> [Double] -> [Double] -> Double -> [Point] -> DFTArray
computeInitialDistribution rows cols phiFreqs rhoFreqs halfLogPeriod points =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
  in if L.any (\(Point _ _ _ s) -> s <= 0) points
       then error $ "computeInitialDistribution: initial scale <= 0."
       else DFTArray rows cols phiFreqs rhoFreqs .
            parMap
              rdeepseq
              (\(!rFreq, !thetaFreq) ->
                 VG.convert .
                 toUnboxedVector .
                 AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
                 L.map
                   (\(Point x y theta scale) ->
                      ( (round x, round y)
                      , cis $
                        (log scale) * (-rFreq) - (thetaFreq * theta * pi / 180))) $
                 points) $
            (,) <$> rhoFreqs <*> phiFreqs

{-# INLINE computeInitialDistribution' #-}
computeInitialDistribution' ::
     Int -> Int -> [Double] -> [Double] -> Double -> [Point] -> DFTArray
computeInitialDistribution' rows cols phiFreqs rhoFreqs  halfLogPeriod points =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
  in if L.any (\(Point _ _ _ s) -> s <= 0) points
       then error $ "computeInitialDistribution: initial scale <= 0."
       else DFTArray rows cols phiFreqs rhoFreqs .
            parMap
              rdeepseq
              (\(!thetaFreq) ->
                 VG.convert .
                 toUnboxedVector .
                 AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
                 L.map
                   (\(Point x y theta _) ->
                      ((round x, round y), cis $ -(thetaFreq * theta * pi / 180))) $
                 points) $
            phiFreqs

{-# INLINE computeInitialDistributionPowerMethod #-}
computeInitialDistributionPowerMethod ::
     Int -> Int -> [Double] -> [Double] -> Double -> [Point] -> DFTArray
computeInitialDistributionPowerMethod rows cols thetaFreqs rFreqs halfLogPeriod points =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      !zeroVec = VG.replicate (rows * cols) 0
  in if L.any (\(Point _ _ _ s) -> s <= 0) points
       then error $ "computeInitialDistributionPowerMethod: initial scale <= 0."
       else DFTArray rows cols thetaFreqs rFreqs .
            L.map
              (\(!rFreq, !thetaFreq) ->
                 if rFreq == 0 && thetaFreq == 0
                   then VG.convert .
                        toUnboxedVector .
                        AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
                        L.map (\(Point x y theta scale) -> ((round x, round y), 1)) $
                        points
                   else zeroVec) $
            (,) <$> rFreqs <*> thetaFreqs

{-# INLINE computeInitialDistributionPowerMethod' #-}
computeInitialDistributionPowerMethod' ::
     Int -> Int -> [Double] -> [Double] -> [Point] -> DFTArray
computeInitialDistributionPowerMethod' rows cols phiFreqs rhoFreqs points =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      !zeroVec = VG.replicate (rows * cols) 0
  in if L.any (\(Point _ _ _ s) -> s <= 0) points
       then error $ "computeInitialDistributionPowerMethod: initial scale <= 0."
       else DFTArray rows cols phiFreqs rhoFreqs .
            L.map
              (\(!thetaFreq) ->
                 if thetaFreq == 0
                   then VG.convert .
                        toUnboxedVector .
                        AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
                        L.map (\(Point x y theta scale) -> ((round x, round y), 1)) $
                        points
                   else zeroVec) $
            phiFreqs

{-# INLINE computeInitialDistributionPowerMethodSparse #-}
computeInitialDistributionPowerMethodSparse ::
     [Double] -> [Double] -> [Point] -> [R.Array U DIM2 (Complex Double)]
computeInitialDistributionPowerMethodSparse thetaFreqs rFreqs points =
  let !numThetaFreq = L.length thetaFreqs
      !numRFreq = L.length rFreqs
      !thetaCenter = div numThetaFreq 2
      !rCenter = div numRFreq 2
      !initArr =
        computeUnboxedS . fromFunction (Z :. numRFreq :. numThetaFreq) $ \(Z :. i :. j) ->
          if i == rCenter && j == thetaCenter
            then 1
            else 0
  in L.replicate (L.length points) initArr

{-# INLINE computeInitialDistributionPowerMethodSparse' #-}
computeInitialDistributionPowerMethodSparse' ::
     [Double] -> [Double] -> [Point] -> [R.Array U DIM1 (Complex Double)]
computeInitialDistributionPowerMethodSparse' thetaFreqs rFreqs points =
  let !numThetaFreq = L.length thetaFreqs
      !thetaCenter = div numThetaFreq 2
      !initArr =
        computeUnboxedS . fromFunction (Z :. numThetaFreq) $ \(Z :. j) ->
          if j == thetaCenter
            then 1
            else 0
      -- !initArr =
      --   computeUnboxedS .
      --   R.traverse (fromListUnboxed (Z :. numThetaFreq) thetaFreqs) id $ \f (Z :. thetaFreq) ->
      --     cis $ -(fromIntegral thetaFreq * 36 * pi / 180)
  in L.replicate (L.length points) initArr

computeInitialDistributionFull' ::
     Int -> Double -> Int -> Int -> [Point] -> DFTArray
computeInitialDistributionFull' numR2Freq period phiFreq rhoFreq points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
  in DFTArray
       numR2Freq
       numR2Freq
       phiFreqs
       rhoFreqs
       [ VG.fromList
         [ L.foldl'
           (\b (Point x y theta _) ->
              b +
              cis
                (-(angularFreq * theta * pi / 180 +
                   (freqX * x + freqY * y) * 2 * pi / period)))
           0
           points
         | freqY <- r2Freqs
         , freqX <- r2Freqs
         ]
       | angularFreq <- phiFreqs
       ]

computeInitialDistributionPowerMethodPinwheelBasis' ::
     Int -> Double -> Double -> Int -> Int -> [Point] -> DFTArray
computeInitialDistributionPowerMethodPinwheelBasis' numR2Freq sigma period phiFreq rhoFreq points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      zeroVec = VG.replicate (numR2Freq ^ 2) 0
      envelope =
        VG.convert .
        toUnboxed . computeS . fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - div numR2Freq 2) ^ 2 + (j - div numR2Freq 2) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
  in DFTArray
       numR2Freq
       numR2Freq
       phiFreqs
       rhoFreqs
       [ if angularFreq == 0
         then VG.zipWith (*) envelope . VG.fromList $
              [ L.foldl'
                (\b (Point x y _ _) ->
                   b + cis (-(freqX * x + freqY * y) * 2 * pi / period))
                0
                points
              | freqY <- r2Freqs
              , freqX <- r2Freqs
              ]
         else zeroVec
       | angularFreq <- phiFreqs
       ]

computeInitialDistributionPowerMethodPinwheelBasis ::
     Int -> Double -> Double -> Int -> Int -> [Point] -> DFTArray
computeInitialDistributionPowerMethodPinwheelBasis numR2Freq sigma period phiFreq rhoFreq points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      zeroVec = VG.replicate (numR2Freq ^ 2) 0
      envelope =
        VG.convert .
        toUnboxed . computeS . fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - div numR2Freq 2) ^ 2 + (j - div numR2Freq 2) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
      -- envelope =
      --   VG.convert . toUnboxed . computeUnboxedS $
      --   analyticalFourierCoefficients2 numR2Freq 1 0 0 sigma period (period * sqrt 2)
  in DFTArray
       numR2Freq
       numR2Freq
       phiFreqs
       rhoFreqs
       [ if angularFreq == 0 && radialFreq == 0
         then VG.zipWith (*) envelope .
              VG.fromList $
              [ L.foldl'
                (\b (Point x y _ _) ->
                   b + cis (-(freqX * x + freqY * y) * 2 * pi / period))
                0
                points
              | freqY <- r2Freqs
              , freqX <- r2Freqs
              ]
         else zeroVec
       | radialFreq <- rhoFreqs
       , angularFreq <- phiFreqs
       ]

computeInitialDistributionFourierPinwheel ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> [Point]
  -> FPArray (vector (Complex Double))
computeInitialDistributionFourierPinwheel numR2Freqs period periodEnv phiFreq rhoFreq thetaFreq rFreq points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      rFreqs = L.map fromIntegral [-rFreq .. rFreq]
   in FPArray
        numR2Freqs
        numR2Freqs
        (2 * rFreq + 1)
        (2 * thetaFreq + 1)
        (2 * rhoFreq + 1)
        (2 * phiFreq + 1) .
      L.map
        (\radialFreq ->
           VG.concat .
           parMap
             rdeepseq
             (\angularFreq ->
                VG.convert . toUnboxed . computeS . makeFilter2D . fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) $
                   [ L.foldl'
                     (\b (Point x y theta scale) ->
                        b +
                        cis
                          (-(angularFreq * theta * pi / 180 +
                             (freqX * x + freqY * y) * 2 * pi / period +
                             2 * pi / log periodEnv * log scale * radialFreq)))
                     0
                     points
                   | freqY <- r2Freqs
                   , freqX <- r2Freqs
                   ]) $
           thetaFreqs) $
      rFreqs

computeInitialDistributionPowerMethodFourierPinwheel ::
     Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
computeInitialDistributionPowerMethodFourierPinwheel numR2Freqs phiFreq rhoFreq thetaFreq rFreq bias =
  let thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      rFreqs = L.map fromIntegral [-rFreq .. rFreq]
      numThetaFreq = 2 * thetaFreq + 1
      zeroVec = VG.replicate (numThetaFreq * numR2Freqs ^ 2) 0
   in FPArray
        numR2Freqs
        numR2Freqs
        (2 * rFreq + 1)
        numThetaFreq
        (2 * rhoFreq + 1)
        (2 * phiFreq + 1) .
      L.map
        (\radialFreq ->
           if radialFreq == 0
             then bias
             else zeroVec) $
      rFreqs

computeInitialDistributionPowerMethodFourierPinwheelFull ::
     Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
computeInitialDistributionPowerMethodFourierPinwheelFull numR2Freqs phiFreq rhoFreq thetaFreq rFreq bias =
  let numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * rFreq + 1
   in FPArray
        numR2Freqs
        numR2Freqs
        (2 * rFreq + 1)
        numThetaFreq
        (2 * rhoFreq + 1)
        (2 * phiFreq + 1) .
      parMap
        rdeepseq
        (\radialFreq ->
           VG.convert .
           toUnboxed .
           computeS .
           R.slice
             (fromUnboxed
                (Z :. numRFreq :. numThetaFreq :. numR2Freqs :. numR2Freqs) .
              VG.convert $
              bias) $
           (Z :. radialFreq :. All :. All :. All)) $
      [0 .. 2 * rFreq]


{-# LANGUAGE LambdaCase, DoAndIfThenElse, GeneralizedNewtypeDeriving #-}
module Truthiness where

import Data.Text (Text)
import qualified Data.Text as Text
import Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString

import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Monoid (Any(..), All(..))
import System.Exit (ExitCode(..))
import Numeric.Natural (Natural)
import Data.Complex (Complex(..))
import Data.Ratio (Ratio)
import Control.Applicative (Alternative(..), Alternative(..))
import Control.Monad (MonadPlus(..))
import Data.Foldable (Foldable(..))

{-| Pythonic truthiness/falsines. see <https://docs.python.org/2.4/lib/truth.html>.

keep around information (e.g. use @Maybe@, a custom flag type, etc.)
to avoid <https://existentialtype.wordpress.com/2011/03/15/boolean-blindness/ Boolean Blindness>,
then drop that information with 'Boolean', when necessary for convenience.

e.g. 'when'

when 'Foldable', instances should satisfy the following law:

@
'falsy' = 'null' . 'toList'
@

when a 'Monoid', instances should satisfy the following law:

@
'falsy' = ('mempty' ==)
@

when a 'Num', instances should satisfy the following law:

@
'falsy' = (0 ==)
@

non-instances:

* @()@
* @Dual@

-}
class Boolean a where

 falsy :: a -> Bool

 truthy :: a -> Bool
 truthy = not . falsy

-- booleans
instance Boolean Bool         where falsy = (== False)
instance Boolean Any          where falsy = falsy . getAny
instance Boolean All          where falsy = falsy . getAll

-- errors
instance Boolean (Maybe a)    where falsy = (\case Nothing -> True; _ -> False) -- not (== Nothing)
instance Boolean (Either e a) where falsy = (\case Left{} -> True; _ -> False)
-- | (includes 'String's)
instance Boolean [a]          where falsy = (\case [] -> True; _ -> False) -- not (== [])

-- numbers
instance Boolean Int          where falsy = (== 0)
instance Boolean Word         where falsy = (== 0)
instance Boolean Integer      where falsy = (== 0)
instance Boolean Natural      where falsy = (== 0)
instance (Integral a) => Boolean (Ratio a)    where falsy = (== 0)
instance (RealFloat a) => Boolean (Complex a)  where falsy = (== 0)

-- containers
instance Boolean (Set a)      where falsy = Set.null
instance Boolean (Map k v)    where falsy = Map.null

-- strings
instance Boolean Char         where falsy = (== '\NUL')
instance Boolean Text         where falsy = Text.null
instance Boolean ByteString   where falsy = ByteString.null

-- et cetera
instance Boolean ExitCode      where falsy = (\case ExitFailure{} -> True; _ -> False)

instance (Foldable t)       => Boolean (WrappedFoldable t a) where falsy = null . toList
instance (Monoid a, Eq a)   => Boolean (WrappedMonoid a)     where falsy = (mempty ==)
instance (Num a,    Eq a)   => Boolean (WrappedNum a)        where falsy = (0 ==)

newtype WrappedFoldable t a = WrappedFoldable { getWrappedFoldable :: t a } deriving (Foldable)
newtype WrappedMonoid   a   = WrappedMonoid   { getWrappedMonoid   :: a } deriving (Monoid,Eq)
newtype WrappedNum      a   = WrappedNum      { getWrappedNum      :: a } deriving (Num,Eq)

-- control

whenB :: (Applicative m, Boolean b) => b -> m () -> m ()
whenB condition action = ifB condition action nothing
{-# INLINEABLE whenB #-}

whenM :: (Monad m, Boolean b) => m b -> m () -> m ()
whenM condition action = ifM condition action nothing
{-# INLINEABLE whenM #-}

guardB :: (Alternative m, Boolean b) => b -> m ()
guardB condition = unlessB condition empty
{-# INLINEABLE guardB #-}

guardM :: (MonadPlus m, Boolean b) => m b -> m ()
guardM condition = unlessM condition mzero
{-# INLINEABLE guardM #-}

unlessB :: (Applicative m, Boolean b) => b -> m () -> m ()
unlessB condition action = ifB condition nothing action
{-# INLINEABLE unlessB #-}

unlessM :: (Monad m, Boolean b) => m b -> m () -> m ()
unlessM condition action = ifM condition nothing action
{-# INLINEABLE unlessM #-}

boolB :: (Boolean b) => a -> a -> b -> a
boolB x y c = ifB c x y
{-# INLINEABLE boolB #-}

ifB :: (Boolean b) => b -> a -> a -> a
ifB c x y =
 if   truthy c
 then x
 else y
{-# INLINEABLE ifB #-}

ifM :: (Monad m, Boolean b) => m b -> m a -> m a -> m a
ifM condition a b = (\c -> ifB c a b) =<< condition
{-# INLINEABLE ifM #-}

nothing :: (Applicative m) => m ()
nothing = pure()

{-$ alternatives

too explicit: <https://hackage.haskell.org/package/base-4.8.2.0/docs/Control-Monad.html>

fewer instances: <https://hackage.haskell.org/package/cond>

-}


import ImageLoader
import Brain

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

import Control.Exception (SomeException (..))
import System.Random     (mkStdGen, randomR)
import Numeric.LinearAlgebra
import Data.List
import qualified Numeric.LinearAlgebra.Data as V
import Debug.Trace

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

type Samples = [(V.Vector Double, V.Vector Double)]

main :: IO ()
main = loadFiles

loadFiles :: IO ()
loadFiles = do
    ex    <- loadExamples
    notEx <- loadNonExamples
    case (ex, notEx)
      of (Left (SomeException e), _) -> printError e
         (_, Left (SomeException e)) -> printError e
         (Right v1, Right v2)        -> runTraining v1 v2
  where
    printError e = putStrLn $ "Exception: " ++ (show e)

runTraining :: Samples -> Samples -> IO ()
runTraining s1 s2 = do
    net <- constructNeuralNetwork 256 [512, 256, 128, 64] 1 all_
    putStrLn "----"
    printSamples all_ net
    putStrLn "----"
    let net' = trainN net 800
    printSamples all_ net'
  where all_   = s1 ++ s2

printSamples :: Samples -> NetworkWithSamples -> IO ()
printSamples [] _       = return ()
printSamples (x:xs) net = do
    putStrLn $ vid ++ " --> " ++ output ++ " ? " ++ expectedOut
    printSamples xs net
  where in_    = fst x
        expectedOut = show $ snd x
        vid    = show $ (vecId in_)
        output = show $ (activate net in_)

-- Given a list of examples, separates into a list of
-- test data (fst) and training data (snd)
sepExTest :: [a] -> ([a], [a])
sepExTest xs = go xs ([], []) (mkStdGen 10)
  where
        go [] a _            = a
        go (v:vs) (a, b) gen = case randomR (0, 99) gen
                               of (n, gen') -> if n < (20 :: Int)
                                                  then go vs (v:a, b) gen'
                                                  else go vs (a, v:b) gen'

vecId :: V.Vector Double -> Double
vecId v = v <.> V.fromList (replicate 256 (1 :: Double))


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Convolution
Description : Convolution layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

This module provides the Convolution layer, which is critical in many computer vision tasks.

-}
module Grenade.Layers.Convolution (
    Convolution (..)
  , Convolution' (..)
  , randomConvolution
  ) where

import           Control.Monad.Random hiding ( fromList )
import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits

#if MIN_VERSION_base(4,11,0)
import           GHC.TypeLits hiding (natVal)
#else
import           GHC.TypeLits
#endif
#if MIN_VERSION_base(4,9,0)
import           Data.Kind (Type)
#endif

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

import           Grenade.Core
import           Grenade.Layers.Internal.Convolution
import           Grenade.Layers.Internal.Update

-- | A convolution layer for a neural network.
--   This uses the im2col convolution trick popularised by Caffe, which essentially turns the
--   many, many, many, many loop convolution into a single matrix multiplication.
--
--   The convolution layer takes all of the kernels for the convolution, which are flattened
--   and then put into columns in the matrix.
--
--   The kernel size dictates which input and output sizes will "fit". Fitting the equation:
--   `out = (in - kernel) / stride + 1` for both dimensions.
--
--   One probably shouldn't build their own layer, but rather use the randomConvolution function.
data Convolution :: 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
              -> !(L kernelFlattened filters) -- The last kernel update (or momentum)
              -> Convolution channels filters kernelRows kernelColumns strideRows strideColumns

data Convolution' :: 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 gradient
               -> Convolution' channels filters kernelRows kernelColumns strideRows strideColumns

instance Show (Convolution c f k k' s s') where
  show (Convolution a _) = renderConv a
    where
      renderConv mm =
        let m  = extract mm
            ky = fromIntegral $ natVal (Proxy :: Proxy k)
            rs = LA.toColumns m
            ms = map (take ky) $ toLists . reshape ky <$> rs

            render n'  | n' <= 0.2  = ' '
                       | n' <= 0.4  = '.'
                       | n' <= 0.6  = '-'
                       | n' <= 0.8  = '='
                       | otherwise =  '#'

            px = (fmap . fmap . fmap) render ms
        in unlines $ foldl1 (zipWith (\a' b' -> a' ++ "   |   " ++ b')) $ px

randomConvolution :: ( MonadRandom m
                     , KnownNat channels
                     , KnownNat filters
                     , KnownNat kernelRows
                     , KnownNat kernelColumns
                     , KnownNat strideRows
                     , KnownNat strideColumns
                     , KnownNat kernelFlattened
                     , kernelFlattened ~ (kernelRows * kernelColumns * channels))
                  => m (Convolution channels filters kernelRows kernelColumns strideRows strideColumns)
randomConvolution = do
    s     <- getRandom
    let wN = uniformSample s (-1) 1
        mm = konst 0
    return $ Convolution wN mm

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => UpdateLayer (Convolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  type Gradient (Convolution channels filters kernelRows kernelColumns strideRows strideColumns) = (Convolution' channels filters kernelRows kernelColumns strideRows strideColumns)
  runUpdate LearningParameters {..} (Convolution oldKernel oldMomentum) (Convolution' kernelGradient) =
    let (newKernel, newMomentum) = descendMatrix learningRate learningMomentum learningRegulariser oldKernel kernelGradient oldMomentum
    in Convolution newKernel newMomentum

  createRandom = randomConvolution

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => Serialize (Convolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (Convolution w _) = 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
      let mm = konst 0
      return $ Convolution wN mm

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

  type Tape (Convolution 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)
        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)

        c  = vid2col kx ky sx sy ix iy ex
        mt = c LA.<> ek
        r  = col2vid 1 1 1 1 ox oy mt
        rs = fromJust . create $ r
    in  (S3D input, S3D rs)
  runBackwards (Convolution kernel _) (S3D input) (S3D dEdy) =
    let ex = extract input
        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)

        c  = vid2col kx ky sx sy ix iy ex

        eo = extract dEdy
        ek = extract kernel

        vs = vid2col 1 1 1 1 ox oy eo

        kN = fromJust . create $ tr c LA.<> vs

        dW = vs LA.<> tr ek

        xW = col2vid kx ky sx sy ix iy dW
    in  (Convolution' kN, S3D . fromJust . create $ xW)

-- | A two dimentional image may have a convolution filter applied to it
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputCols - 1) * strideCols) ~ (inputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * 1)
         , KnownNat (outputRows * filters)
         ) => Layer (Convolution 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) where
  type Tape (Convolution 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)

-- | A two dimensional image may have a convolution filter applied to it producing
--   a two dimensional image if both channels and filters is 1.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputCols - 1) * strideCols) ~ (inputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * 1)
         , KnownNat (outputRows * 1)
         ) => Layer (Convolution 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) where
  type Tape (Convolution 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)

-- | A three dimensional image can produce a 2D image from a convolution with 1 filter
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputCols - 1) * strideCols) ~ (inputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * channels)
         , KnownNat (outputRows * 1)
         ) => Layer (Convolution channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) where
  type Tape (Convolution 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))


{-# LANGUAGE RankNTypes #-}
module Patterns where

import Prelude hiding (map, zip, iterate)
import qualified Prelude (map, zip, iterate)

import Numeric.SpecFunctions (log2)
import Data.Array

type Vec = Array Int

-- high-level patterns
map :: forall a b. (a -> b) -> [a] -> [b]
map = Prelude.map

reduce :: forall a. (a -> a -> a) -> a -> [a] -> [a]
reduce f z xs = [Prelude.foldl f z xs]

zip :: forall a b. [a] -> [b] -> [(a,b)]
zip = Prelude.zip

split :: forall a. Int -> [a] -> [[a]]
split n xs | len `mod` n == 0   = builder (len `div` n) xs
           | otherwise          = error "Length not divisible by n"
    where len = length xs
          builder 1 xs = [take n xs]
          builder i xs = take n xs : builder (i-1) (drop n xs)

join :: forall a. [[a]] -> [a]
join = foldl (++) []

iterate :: forall a. Int -> (a -> a) -> a -> a
iterate 0 f xs = xs
iterate n f xs = iterate (n-1) f (f xs)

reorder :: forall a. [a] -> [a]
reorder = id

-- low-level patterns
mapWorkgroup    = map
mapLocal        = map
mapGlobal       = map
mapWarp         = map
mapLane         = map
mapSeq          = map

reduceSeq :: forall a b. (a -> b -> a) -> a -> [b] -> [a]
reduceSeq f z xs = [Prelude.foldl f z xs]

reorderStride :: forall a. Int -> [a] -> [a]
reorderStride s = id

toLocal :: forall a b. (a -> b) -> (a -> b)
toLocal = id

toGlobal :: forall a b. (a -> b) -> (a -> b)
toGlobal = id

asVector :: forall a. Int -> [a] -> [Vec a]
asVector n xs | len `mod` n == 0    = builder (len `div` n) xs
              | otherwise           = error "Length not divisible by n"
    where len = length xs
          builder 1 xs = [newVector (take n xs)]
          builder i xs = newVector (take n xs) : builder (i-1) (drop n xs)
          newVector xs = array (1, n) (zip [1..n] xs)

asScalar :: forall a. [Vec a] -> [a]
asScalar = join . map elems

vectorize :: forall a b. Int -> (a -> b) -> (Vec a -> Vec b)
vectorize n f = \ a -> array (bounds a) $ map (\ (i, e) -> (i, f e)) $ assocs a

-- utilities
onPairs :: forall a b c. (a -> b -> c) -> ((a, b) -> c)
onPairs f = \ (a, b) -> f a b

-- benchmarks
scal :: forall a. Num a => a -> [a] -> [a]
scal a = map ((*) a)

asum :: forall a. (Ord a, Num a) => [a] -> [a]
asum xs = reduce (+) 0 (map abs xs)

dot :: forall a. Num a => [a] -> [a] -> [a]
dot xs ys = reduce (+) 0 (map (onPairs (*)) (zip xs ys))

gemv :: forall a. Num a => [[a]] -> [a] -> [a] -> a -> a -> [a]
gemv mss xs ys a b = map (onPairs (+)) (zip zs (scal b ys))
    where zs = map (head . scal a . dot xs) mss

-- vector scale
mult3 :: forall a. Num a => a -> a
mult3 x = x * 3

vectorScale1 :: forall a. Num a => [a] -> [a]
vectorScale1 = map mult3

vectorScale2 :: forall a. Num a => [a] -> [a]
vectorScale2 xs = join . map (map mult3) . split (s `div` 2) $ xs
    where s = length xs

vectorScale3 :: forall a. Num a => [a] -> [a]
vectorScale3 xs = join . map (
        asScalar . map (vectorize 4 mult3) . asVector 4
    ) . split (s `div` 2) $ xs
    where s = length xs

vectorScale4 :: forall a. Num a => [a] -> [a]
vectorScale4 xs = join . mapWorkgroup (
        asScalar . mapLocal (vectorize 4 mult3) . asVector 4
    ) . split (s `div` 2) $ xs
    where s = length xs

-- reduction
vecSum0 = reduce (+) 0

vecSum1 xs = vecSum0 . join . mapWorkgroup (
        join . toGlobal (mapLocal (mapSeq id)) . split 1 .
        iterate (log2 wgSize) (
            join . mapLocal (reduceSeq (+) 0) . split 2
        ) .
        join . toLocal (mapLocal (mapSeq id)) . split 1
    ) . split wgSize $ xs
    where wgSize = 128

vecSum2 xs = vecSum0 . join . mapWorkgroup (
        join . toGlobal (mapLocal (mapSeq id)) . split 1 .
        iterate (log2 wgSize) (
            join . mapLocal (reduceSeq (+) 0) . split 2
        ) .
        join . toLocal (mapLocal (reduceSeq (+) 0)) . split 2
    ) . split (2 * wgSize) $ xs
    where wgSize = 128

vecSum3 xs = vecSum0 . join . mapWorkgroup (
        join . toGlobal (mapLocal (mapSeq id)) . split 1 .
        join . mapWarp (
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2
        ) . split 64 .
        iterate (log2 wgSize - log2 64) (
            join . mapLocal (reduceSeq (+) 0) . split 2
        ) .
        join . toLocal (mapLocal (reduceSeq (+) 0)) . split 2
    ) . split (2 * wgSize) $ xs
    where wgSize = 128

vecSum4 xs = vecSum0 . join . mapWorkgroup (
        join . toGlobal (mapLocal (mapSeq id)) . split 1 .
        join . mapWarp (
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2
        ) . split 64 .
        join . mapLocal (reduceSeq (+) 0) . split 2 .
        join . toLocal (mapLocal (reduceSeq (+) 0)) . split 2
    ) . split (2 * wgSize) $ xs
    where wgSize = 128

vecSum5 xs = vecSum0 . join . mapWorkgroup (
        join . toGlobal (mapLocal (mapSeq id)) . split 1 .
        join . mapWarp (
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2 .
            join . mapLane (reduceSeq (+) 0) . split 2
        ) . split 64 .
        join . mapLocal (reduceSeq (+) 0) . split 2 .
        join . toLocal (mapLocal (reduceSeq (+) 0)) .
            split (blockSize `div` wgSize)
    ) . split blockSize $ xs
    where wgSize = 128
          blockSize = 262144