module SimpleParser.SimpleParser where

import System.IO
import Text.Parsec hiding ( spaces )
import Text.Parsec.String
import Control.Monad
import Numeric
import Data.Complex -- for complex number representation in 2.7
import Data.Ratio -- for rational numbers
import Data.Array -- could use haskell vectors instead, harder implementation
import Data.IORef
import Control.Monad.Error --deprecated but following the book :/

spaces :: Parser ()
spaces = skipMany1 space

type IOThrowsError = ErrorT LispError IO

data LispVal = Atom String
               | List [LispVal]
               | DottedList [LispVal] LispVal
               | Number Integer
               | String String
               | Bool Bool
               | Character Char
               | Float Double
               | Complex (Complex Double)
               | Rational Rational
               | Vector (Array Int LispVal)
               | PrimitiveFunc ([LispVal] -> ThrowsError LispVal)
               | Func { params :: [String], vararg :: Maybe String
                      , body :: [LispVal], closure :: Env}
               | IOFunc ([LispVal] -> IOThrowsError LispVal)
               | Port Handle

data LispError = NumArgs Integer [LispVal]
               | TypeMismatch String LispVal
               | Parser ParseError
               | BadSpecialForm String LispVal
               | NotFunction String String
               | UnboundVar String String
               | Default String

showError :: LispError -> String
showError (NumArgs expected found) = "Expected: " ++ show expected ++
  " args but found: " ++ unwordsList found
showError (TypeMismatch expected found) = "Type mismatch; Got: " ++
  show found ++ " but Expected: " ++ expected
showError (Parser error) = "Parse error at: " ++ show error
showError (UnboundVar message var) = message ++ " : " ++ show var
showError (BadSpecialForm message val) = message ++ " : " ++ show val
showError (NotFunction message func) = message ++ " : " ++ show func

instance Show LispError where show = showError
type ThrowsError = Either LispError
type Env = IORef [(String, IORef LispVal)]

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

--------------------------- Exercise 2.1 ----------------------------------------
parseNumber' :: Parser LispVal
parseNumber' =
  do
    s <- many1 digit
    return (Number . read $ s)

--------------------------- Exercise 2.2 ----------------------------------------
parseNumber'' :: Parser LispVal
parseNumber'' = many1 digit >>= return . Number. read

--------------------------- Exercise 2.3 ----------------------------------------
specials :: Parser Char
specials = do char '\\'
              x <- oneOf "\\\"nrt"
              return $ case x of
                '\\' -> x
                '"'  -> x
                'n'  -> '\n'
                'r'  -> '\r'
                't'  -> '\t'

parseString :: Parser LispVal
parseString = do char '"'
                 x <- many $ specials <|> noneOf "\"\\"
                 char '"'
                 return $ String x
--------------------------- Exercise 2.4 ----------------------------------------
--Scheme defines Octal as #o, decimal as #d, hex as #h
symbol :: Parser Char
symbol = oneOf "!$%&|*+-/:<=>?@^_~"

-- so we can no longer parse bools with a prefixed '#'
parseBool :: Parser LispVal
parseBool =
  do
    char '#' --match a #
    (char 't' >> return (Bool True)) <|> (char 'f' >> return (Bool False))

parseNumber :: Parser LispVal
parseNumber = parseDecimal
  <|> parseSchemeDecimal
  <|> parseHex
  <|> parseOct
  <|> parseBin

parseDecimal :: Parser LispVal
parseDecimal = many1 digit >>= return . Number . read

parseSchemeDecimal:: Parser LispVal
parseSchemeDecimal = do try $ string "#d"
                        x <- many1 digit
                        (return . Number . read) x

parseHex :: Parser LispVal
parseHex = do try $ string "#x"
              x <- many1 hexDigit
              return $ Number (hex2dig x)

parseOct :: Parser LispVal
parseOct = do try $ string "#o"
              x <- many1 octDigit
              return $ Number (oct2dig x)

parseBin :: Parser LispVal
parseBin = do try $ string "#b"
              x <- many1 (oneOf "10")
              return $ Number (bin2dig x)

numToList :: Int -> [Int]
numToList = map (read . (:[])) . show --probably slow, could use div and mod

oct2dig x = fst $ readOct x !! 0
hex2dig x = fst $ readHex x !! 0
bin2dig  = bin2dig' 0
bin2dig' digint "" = digint
bin2dig' digint (x:xs) = let old = 2 * digint + (if x == '0' then 0 else 1) in
                         bin2dig' old xs
--------------------------- Exercise 2.5 ----------------------------------------
parseCharacter :: Parser LispVal
--parseChar
             --try to parse a whole match on newline or space
             --or try to match any character not followed by an alpha numeric
             --return thst as a list
             -- now try to match newline or space return the corresponding value
             --if not a match then return the first character match by previous
             -- do
parseCharacter =
  do
    try $ string "#\\"
    x <- try (string "newline" <|> string "space")
         <|> do {str <- anyChar; notFollowedBy alphaNum; return [str]}
    return . Character $ case x of
      "newline" -> '\n'
      "space" -> ' '
      otherwise -> (x !! 0) --use of unsafe head

--------------------------- Exercise 2.6 ----------------------------------------
-- I keep struggling to find the appropriate informatin to match on in the
-- R5RS, I thought I had to implement exact/nonexact in addition to this...
parseFloat :: Parser LispVal
parseFloat =
  do
    x <- many1 digit --many  digits before a '.'
    char '.' -- match on the ','
    y <- many1 digit --more digitts after
    return $ Float . fst . head . readFloat $ (x ++ y) -- readFloat is a zipper

--------------------------- Exercise 2.7 ----------------------------------------
toDouble :: LispVal -> Double
toDouble (Float f) = realToFrac f
toDouble (Number n) = fromInteger n

parseComplexNumbers :: Parser LispVal
parseComplexNumbers =
  do
    a <- parseFloat <|> parseNumber
    char '+' --complex nums have pattern a+bi
    b <- parseFloat <|> parseNumber
    char 'i'
    return (Complex (toDouble a :+ toDouble b))

parseRationalNumbers :: Parser LispVal
parseRationalNumbers =
  do
    a <- many1 digit-- a rational num is denoted by a / b /= 0
    char '/'
    b <- many1 digit
    return (Rational ((read a) % (read b))) --didn't know that (%) existed

------------------------Begin section on recursive parsers-----------------------
parseList :: Parser LispVal
--sepBy parseExpr spaces seperates a string by spaces on parseExpr criteria
parseList = liftM List $ sepBy parseExpr spaces

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

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

parseExpr :: Parser LispVal
parseExpr = parseAtom
  <|> parseString
  <|> try parseBool
  <|> try parseComplexNumbers
  <|> try parseRationalNumbers
  <|> parseNumber --try's are required because these start with a '#'
  <|> parseCharacter
  <|> parseQuasiQuoted
  <|> parseUnQuote
  <|> parseQuoted
  <|> try parseVector'
  <|> do char '(' --first bracker
         x <- try parseList <|> parseDottedList
         char ')'
         return x
--------------------------- Exercise 4.1 ----------------------------------------
parseQuasiQuoted :: Parser LispVal
parseQuasiQuoted =
  do
    char '`'
    x <- parseExpr
    return $ List [Atom "quasiquote", x]

parseUnQuote :: Parser LispVal
parseUnQuote =
  do
    char ','
    x <- parseExpr
    return $ List [Atom "unquote", x]

--------------------------- Exercise 4.2 ----------------------------------------
parseVector :: Parser LispVal
parseVector =
  do
    values <- sepBy parseExpr spaces
    return $ Vector (listArray (0, (length values) - 1) values)

parseVector' :: Parser LispVal
parseVector' =
  do string "#("
     vector <- parseVector
     char ')'
     return vector

--------------------------- Exercise 4.3 ----------------------------------------
     --skip

showVal :: LispVal -> String
showVal (String contents) = "\"" ++ contents ++ "\""
showVal (Atom name) = name
showVal (Number contents) = show contents
showVal (Float num) = show num
showVal (Complex num) = show num
showVal (Rational num) = show num
showVal (Vector arr) = unlines . map showVal . elems $ arr
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (List values) = "(" ++ unwordsList values ++ ")"
showVal (DottedList h t) = "(" ++ unwordsList h ++ " . " ++ showVal t
  ++ ")"
showVal (Character c) = [c]
showVal (PrimitiveFunc _) = "primitive"
showVal (Func {params = args, vararg = varargs, body = body, closure = env}) =
  "(lambda (" ++ unwords (map show args) ++
  (case varargs of
      Nothing -> ""
      Just arg -> " . " ++ arg) ++ ") ...)"
showVal (Port _) = "<IO port>"
showVal (IOFunc _) = "<IO primitive>"

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

instance Show LispVal where show = showVal


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict, StrictData #-}
{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}

------------------------------------------------
-- |
-- Module    :  ArcGraph.EnhancedState
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Implementing manipulations on enhanced states
--
------------------------------------------------

module ArcGraph.EnhancedState where

import GHC.Generics (Generic)

import Control.DeepSeq
import Control.Parallel.Strategies

import Control.Applicative
import Control.Monad
import Control.Monad.ST

import Data.Bifunctor

import Data.STRef
import Data.Maybe as M
import Data.Foldable
import qualified Data.List as L
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IMap
import qualified Data.Set as Set
import qualified Data.BitArray as BA

import qualified Numeric.LinearAlgebra as LA

import ArcGraph
import ArcGraph.Component
import ArcGraph.State

import Numeric.Matrix.Integral
import Numeric.Algebra.FreeModule as FM
import Numeric.Algebra.Frobenius as Frob
import Numeric.Algebra.Presentation
import Numeric.Algebra.Homology

{-- for debug
import Debug.Trace

traceCond :: Bool -> String -> a -> a
traceCond False _  = id
traceCond True msg = trace msg
--}

-----------------------
-- * Enhanced states --
-----------------------
class (DState ds, Ord e) => Enhancement ds e where
  listEnh :: (Alternative f) => Int -> ArcGraph -> ds -> f e
  diffEnh :: ArcGraph -> ds -> e -> FreeMod Int (ds,e)

listEStates :: (Enhancement ds e) => ArcGraph -> Int -> Int -> [(ds,e)]
listEStates ag i j = concatMap (\s -> (,) s <$> listEnh j ag s) $ listStates ag i

-----------------
-- * Instances --
-----------------
newtype MapEState pc = MEState (Map.Map pc SL2B)
  deriving (Eq, Ord, Show, Generic, NFData)

instance (DState ds, PComponent pc) => Enhancement ds (MapEState pc) where
  listEnh j ag st =
    let comps = getComponents (smoothing ag st)
        deg' = j - degree ag st + L.length comps
        subs = filter (\sub -> 2*L.length sub == deg') $ L.subsequences comps
        mapS sub = Map.fromList $! map (\c -> (c, if elem c sub then SLI else SLX)) comps
    in MEState <$> foldr' (\x xs -> pure (mapS x) <|> xs) empty subs

  diffEnh ag st (MEState mp) =
    FM.sumFM $ (diffState ag st :: V.Vector (Int,ds)) >>= \dstp -> do
      let (sign,dst) = dstp
      return $! ((,) dst . MEState . Map.fromList) FM.@$>% sign FM.@*% Frob.tqftZ (doIntersect ag) (Map.toList mp) (getComponents (smoothing ag dst))

----------------------
-- WORK IN PROGRESS --
----------------------{--
newtype BArray = BArray BA.BitArray
  deriving (Eq, Ord, Show, Generic)

instance NFData BArray where
  rnf (BArray x) = x `seq` ()

data BitEState pc = BEState pc BArray
  deriving (Eq, Show, Generic, NFData)
--}

data ArcGraphE ds e = AGraphE {
  arcGraph :: ArcGraph,
  state :: ds,
  enhancement :: e
  } deriving (Eq,Show,Generic,NFData)

enhancements :: (DState ds, PComponent pc, Alternative f) => Int -> ArcGraph -> ds -> f (Map.Map pc SL2B)
enhancements deg ag st =
  let comps = getComponents (smoothing ag st)
      deg' = deg + L.length comps
      subs = filter (\sub -> 2*L.length sub == deg') $ L.subsequences comps
      mapS sub = Map.fromList $! map (\c -> (c, if elem c sub then SLI else SLX)) comps
  in foldr' (\x xs -> pure (mapS x) <|> xs) empty subs

----------------------------------------------------------
-- The computation of (unnormalized) Khovanov homology
-----------------------------------------------------------
-- | The type to carry the data of Khovanov homologies
data KHData a ds e = KHData {
  subject :: ArcGraph,
  rank :: Int,
  tors :: [a],
  cycleV :: [FreeMod a (ds, e)],
  bndryV :: Maybe [FreeMod a (ds, e)] }
  deriving (Show,Eq,Generic, NFData)

vecToSum :: (Coefficient a, Ord b) => [b] -> Vec a -> FreeMod a b
vecToSum bs v = sumFM $! zipWith (@*@%) (vecToList v) bs

cohomologyToKH :: (Coefficient a, DState ds, Enhancement ds e) => ArcGraph -> [(ds, e)] -> Bool -> Homology a -> KHData a ds e
cohomologyToKH ag basis hasBndry hdata =
  let basisAGE = uncurry (AGraphE ag) <$> basis
  in KHData {
    subject = ag,
    rank = L.length (freeCycs hdata),
    tors = fmap fst (torsions hdata),
    cycleV = fmap (vecToSum basis) (freeCycs hdata ++ fmap snd (torsions hdata)),
    bndryV = if hasBndry
             then Just (fmap (vecToSum basis) (bndries hdata))
             else Nothing }

-- | Compute Khovanov homology for given range of cohomological degrees and a given quantum-degree
computeKhovanov :: (ChainEliminable a, NFData ds, NFData e, Show ds, Show e, Eq ds, Eq e, Enhancement ds e) => ArcGraph -> Int -> [ds] -> Bool -> IntMap (KHData a ds e)
computeKhovanov ag qdeg states hasBndry =
  let numCrs = countCross ag
      hdegs = [0..numCrs]
      slimAG = slimCross ag
      basis i = L.concatMap (\st -> (,) st <$> listEnh qdeg slimAG st) (filter ((==i). degree slimAG) states)
      !basisMap = force (IMap.fromAscList (map (\i -> (i,basis i)) hdegs))
      !diffs = flip (parMap rdeepseq) [0..numCrs-1] $ \i ->
        let sbasis = basisMap IMap.! i
            tbasis = basisMap IMap.! (i+1)
        in force $! present (uncurry (diffEnh slimAG)) sbasis tbasis
  in if numCrs == 0
     then -- The case where there is no crossing point;
       IMap.map (\v -> KHData slimAG (L.length v) [] (fmap (1@*@%) v) Nothing) basisMap
     else -- The case where there is at least one crossing point;
       let !hdata = force (homology diffs)
           !hdataMap = IMap.fromAscList $ filter (\hdti -> not (null (freeCycs (snd hdti)) && null (torsions (snd hdti))) ) $ zip hdegs hdata
       in flip IMap.mapWithKey hdataMap $ \i hdt -> cohomologyToKH slimAG (basisMap IMap.! i) hasBndry hdt


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Engine.Matrix.NewMatrix where

import Foreign.Storable (Storable(..))
import qualified Data.Vector.Storable as V

import GHC.TypeLits
import Numeric.LinearAlgebra.Static

newtype Matrix w h t = Matrix (L w h)
  deriving (Num, Fractional, Floating)

newtype Vector l t = Vector (R l)
  deriving (Num, Fractional, Floating)

instance (Sized t (R l) V.Vector) => Sized t (Vector l t) V.Vector where
    konst = Vector . konst
    unwrap (Vector r) = unwrap r

(<&>) :: (KnownNat w1, KnownNat w2, KnownNat h) =>
    Matrix w1 h t -> Matrix w2 h t -> Matrix (w1+w2) h t
(<&>) (Matrix l1) (Matrix l2) = Matrix $ l1 —— l2


{-# LANGUAGE RecordWildCards #-}

import Control.Monad.Trans.Except
import Data.Word
import Data.Int
import Data.Complex
import Foreign.C.Types
import Data.Maybe
import Control.Monad
import Data.Monoid

import Control.Error.Util
import Pipes as P
import Pipes.Prelude as P
import Options.Applicative
import Data.Vector.Storable as VS hiding ((++))
import Data.Vector.Generic as VG hiding ((++))
import Graphics.Rendering.OpenGL

import SDR.Util as U
import SDR.RTLSDRStream
import SDR.FFT
import SDR.Plot
import SDR.ArgUtils
import Graphics.DynamicGraph.Waterfall
import Graphics.DynamicGraph.Util

data Options = Options {
    frequency    :: Word32,
    sampleRate   :: Word32,
    gain         :: Maybe Int32,
    fftSize      :: Maybe Int,
    windowWidth  :: Maybe Int,
    windowHeight :: Maybe Int,
    rows         :: Maybe Int,
    colorMap     :: Maybe [GLfloat]
}

parseColorMap :: ReadM [GLfloat]
parseColorMap = eitherReader func
    where
    func "jet"     = return jet
    func "jet_mod" = return jet_mod
    func "hot"     = return hot
    func "bw"      = return bw
    func "wb"      = return wb
    func arg       = Left $ "Cannot parse colour map: `" ++ arg ++ "'"

optParser :: Parser Options
optParser = Options
          <$> option (fmap fromIntegral parseSize) (
                 long "frequency"
              <> short 'f'
              <> metavar "FREQUENCY"
              <> help "Frequency to tune to"
              )
          <*> option (fmap fromIntegral parseSize) (
                 long "samplerate"
              <> short 'r'
              <> metavar "RATE"
              <> help "Sample rate"
              )
          <*> optional (option auto (
                 long "gain"
              <> short 'g'
              <> metavar "GAIN"
              <> help "Tuner gain"
              ))
          <*> optional (option (fmap fromIntegral parseSize) (
                 long "size"
              <> short 's'
              <> metavar "SIZE"
              <> help "FFT bin size. Default is 8192."
              ))
          <*> optional (option auto (
                 long "width"
              <> short 'w'
              <> metavar "WIDTH"
              <> help "Window width. Default is 1024."
              ))
          <*> optional (option auto (
                 long "height"
              <> short 'h'
              <> metavar "HEIGHT"
              <> help "Window height. Default is 480."
              ))
          <*> optional (option auto (
                 long "rows"
              <> short 'r'
              <> metavar "ROWS"
              <> help "Number of rows in waterfall. Default is 1000."
              ))
          <*> optional (option parseColorMap (
                 long "colorMap"
              <> short 'm'
              <> metavar "COLORMAP"
              <> help "Waterfall color map. Default is 'jet_mod'."
              ))

opt :: ParserInfo Options
opt = info (helper <*> optParser) (fullDesc <> progDesc "Draw a dynamic waterall plot of the received spectrum using OpenGL" <> header "RTLSDR Waterfall")

doIt Options{..} = do
    res <- lift setupGLFW
    unless res (throwE "Unable to initilize GLFW")

    let fftSize' =  fromMaybe 8192 fftSize
        window   =  hanning fftSize' :: VS.Vector Double
    str          <- sdrStream ((defaultRTLSDRParams frequency sampleRate) {tunerGain = gain}) 1 (fromIntegral $ fftSize' * 2)
    rfFFT        <- lift $ fftw fftSize'
    rfSpectrum   <- plotWaterfall (fromMaybe 1024 windowWidth) (fromMaybe 480 windowHeight) fftSize' (fromMaybe 1000 rows) (fromMaybe jet_mod colorMap)
    --rfSpectrum   <- plotFill (maybe 1024 id windowWidth) (maybe 480 id windowHeight) fftSize' (maybe jet_mod id colorMap)
    --rfSpectrum   <- plotTexture (maybe 1024 id windowWidth) (maybe 480 id windowHeight) fftSize' fftSize'

    lift $ runEffect $   str
                     >-> P.map (interleavedIQUnsigned256ToFloat :: VS.Vector CUChar -> VS.Vector (Complex Double))
                     >-> P.map (VG.zipWith (flip mult) window . VG.zipWith mult (halfBandUp fftSize'))
                     >-> rfFFT
                     >-> P.map (VG.map ((* (32 / fromIntegral fftSize')) . realToFrac . magnitude))
                     >-> rfSpectrum

main = execParser opt >>= exceptT putStrLn return . doIt


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

module TensorOps.Learn.NeuralNet.FeedForward
  ( Network(..)
  , buildNet
  , runNetwork
  , trainNetwork
  , induceNetwork
  , nmap
  , (~*)
  , (*~)
  , liftNet
  , netParams
  , networkGradient
  , genNet
  , ffLayer
  ) where

import           Control.Category
import           Control.DeepSeq
import           Control.Monad.Primitive
import           Data.Kind
import           Data.Singletons
import           Data.Singletons.Prelude        (Sing(..))
import           Data.Type.Conjunction
import           Data.Type.Length
import           Data.Type.Product              as TCP
import           Data.Type.Product.Util         as TCP
import           Data.Type.Sing
import           Prelude hiding                 ((.), id)
import           Statistics.Distribution.Normal
import           System.Random.MWC
import           TensorOps.Learn.NeuralNet
import           TensorOps.NatKind
import           TensorOps.TOp                  as TO
import           TensorOps.Tensor               as TT
import           TensorOps.Types
import           Type.Class.Higher
import           Type.Class.Higher.Util
import           Type.Class.Known
import           Type.Class.Witness
import           Type.Family.List
import           Type.Family.List.Util

data Network :: ([k] -> Type) -> k -> k -> Type where
    N :: { _nsPs    :: !(Sing ps)
         , _nOp     :: !(TOp ('[i] ': ps) '[ '[o] ])
         , _nParams :: !(Prod t ps)
         } -> Network t i o

instance NFData1 t => NFData (Network t i o) where
    rnf = \case
      N _ o p -> o `seq` p `deepseq1` ()
    {-# INLINE rnf #-}

buildNet
    :: SingI ps
    => TOp ('[i] ': ps) '[ '[o] ]
    -> Prod t ps
    -> Network t i o
buildNet = N sing

netParams
    :: Network t i o
    -> (forall ps. SingI ps => Prod t ps -> r)
    -> r
netParams n f = case n of
    N o _ p -> f p \\ o

(~*~)
    :: Network t a b
    -> Network t b c
    -> Network t a c
N sPs1 o1 p1 ~*~ N sPs2 o2 p2 =
    N (sPs1 %:++ sPs2) (o1 *>> o2) (p1 `TCP.append'` p2)
        \\ singLength sPs1
infixr 4 ~*~
{-# INLINE (~*~) #-}

instance Category (Network t) where
    id = N SNil idOp Ø
    (.) = flip (~*~)

(~*) :: TOp '[ '[a] ] '[ '[b] ]
     -> Network t b c
     -> Network t a c
f ~* N sO o p = N sO (f *>> o) p
infixr 4 ~*
{-# INLINE (~*) #-}

(*~) :: Network t a b
     -> TOp '[ '[b] ] '[ '[c] ]
     -> Network t a c
N sO o p *~ f = N sO (o >>> f) p
infixl 5 *~
{-# INLINE (*~) #-}

liftNet
     :: TOp '[ '[i] ] '[ '[o] ]
     -> Network t i o
liftNet o = buildNet o Ø

nmap
     :: SingI o
     => (forall a. RealFloat a => a -> a)
     -> Network t i o
     -> Network t i o
nmap f n = n *~ TO.map f
{-# INLINE nmap #-}

runNetwork
    :: (RealFloat (ElemT t), Tensor t)
    => Network t i o
    -> t '[i]
    -> t '[o]
runNetwork (N _ o p) = head' . runTOp o . (:< p)
{-# INLINE runNetwork #-}

trainNetwork
    :: forall i o t. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> ElemT t
    -> t '[i]
    -> t '[o]
    -> Network t i o
    -> Network t i o
trainNetwork loss r x y = \case
    N s o p ->
      let p' = map1 (\(!(s1 :&: o1 :&: g1)) -> TT.zip stepFunc o1 g1 \\ s1)
             $ zipProd3 (singProd s) p (tail' $ netGrad loss x y s o p)
      in  N s o p'
  where
    stepFunc :: ElemT t -> ElemT t -> ElemT t
    stepFunc !o' !g' = o' - r * g'
    {-# INLINE stepFunc #-}
{-# INLINE trainNetwork #-}

induceNetwork
    :: forall i o t. (Tensor t, RealFloat (ElemT t), SingI i)
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> ElemT t
    -> t '[o]
    -> Network t i o
    -> t '[i]
    -> t '[i]
induceNetwork loss r y = \case
    N s o p -> \x -> TT.zip stepFunc x (head' $ netGrad loss x y s o p)
  where
    stepFunc :: ElemT t -> ElemT t -> ElemT t
    stepFunc o' g' = o' - r * g'
    {-# INLINE stepFunc #-}
{-# INLINE induceNetwork #-}

networkGradient
    :: forall i o t r. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> t '[i]
    -> t '[o]
    -> Network t i o
    -> (forall ps. SingI ps => Prod t ps -> r)
    -> r
networkGradient loss x y = \case
    N s o p -> \f -> f (tail' $ netGrad loss x y s o p) \\ s
{-# INLINE networkGradient #-}

netGrad
    :: forall i o ps t. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> t '[i]
    -> t '[o]
    -> Sing ps
    -> TOp ('[i] ': ps) '[ '[o] ]
    -> Prod t ps
    -> Prod t ('[i] ': ps)
netGrad loss x y s o p = (\\ appendSnoc lO (Proxy @'[o])) $
                         (\\ lO                         ) $
                         takeProd @'[ '[o] ] (LS lO)
                       $ gradTOp o' inp
  where
    lO  :: Length ps
    lO = singLength s
    o'  :: ((ps ++ '[ '[o] ]) ~ (ps >: '[o]), Known Length ps)
        => TOp ('[i] ': ps >: '[o]) '[ '[]]
    o' = o *>> loss
    inp  :: Prod t ('[i] ': ps >: '[o])
    inp = x :< p >: y
{-# INLINE netGrad #-}

ffLayer
    :: forall i o m t. (SingI i, SingI o, PrimMonad m, Tensor t)
    => Gen (PrimState m)
    -> m (Network t i o)
ffLayer g = (\w b -> buildNet ffLayer' (w :< b :< Ø))
          <$> genRand (normalDistr 0 0.5) g
          <*> genRand (normalDistr 0 0.5) g
  where
    ffLayer'
        :: TOp '[ '[i], '[o,i], '[o]] '[ '[o] ]
    ffLayer' = firstOp (TO.swap >>> TO.matVec)
           >>> TO.add
    {-# INLINE ffLayer' #-}
{-# INLINE ffLayer #-}

genNet
    :: forall k o i m (t :: [k] -> Type). (SingI o, SingI i, PrimMonad m, Tensor t)
    => [(Integer, Activation k)]
    -> Activation k
    -> Gen (PrimState m)
    -> m (Network t i o)
genNet xs0 f g = go sing xs0
  where
    go  :: forall (j :: k). ()
        => Sing j
        -> [(Integer, Activation k)]
        -> m (Network t j o)
    go sj = (\\ sj) $ \case
      []        -> (*~ getAct f) <$> ffLayer g
      (x,f'):xs -> withNatKind x $ \sl -> (\\ sl) $ do
        n <- go sl xs
        l <- ffLayer g
        return $ l *~ getAct f' ~*~ n
    {-# INLINE go #-}
{-# INLINE genNet #-}


{-# LANGUAGE ScopedTypeVariables, RankNTypes #-}

module Dna.Graph (
readConnection,readAllConnection,
allPossibleNode,graphFromConnection,
Graph,euclidianCycle,
unbalanceNodes,
canBeBalance, balance, initOverlaGraph
)
where

import Data.List
import Data.Tuple
import Data.Functor
import Data.STRef
import Control.Monad
import Data.Maybe
import Control.Applicative
import Numeric.LinearAlgebra (rows, (!), cols, toLists, (?), (¿), toList)
import Numeric.LinearAlgebra.Data (Matrix, I)
import Numeric.LinearAlgebra.Devel
import Control.Monad.ST

import qualified Data.ByteString.Lazy.Char8 as C
import qualified Data.Set as S

data Graph a = GraphInternal {
    node :: S.Set a,
    connection :: Matrix I,
    addCon :: Maybe (a,a)
}

instance (Ord a, Show a) => Show (Graph a) where
    show graph = fromMaybe matrixShow (do
                    (from,to) <- addCon graph
                    return $ matrixShow ++ "\n With Link: " ++ show from  ++ "->" ++ show to)
        where
            linkCount row col = connection graph ! row ! col
            isLink row col =   linkCount row col >= 1
            matrixShow = intercalate "\n"
                [link | row <- [0 .. rows (connection graph) - 1],
                        any (isLink row) [0 .. cols (connection graph) - 1],
                        link <- [show (S.elemAt row (node graph)) ++ " -> " ++
                                 intercalate "," (concatMap (\col -> replicate (fromIntegral $ linkCount row col) $ show  $ S.elemAt col (node graph))
                                                    (filter (isLink row) [0 .. cols (connection graph) - 1]))]]

initOverlaGraph :: Ord a =>  [a] -> (forall s. STMatrix s I -> (a -> Int) -> ST s ()) -> Graph a
initOverlaGraph nodes f = GraphInternal setNode (runSTMatrix $ do
        m <- newMatrix 0 lenNodes lenNodes
        f m (`S.findIndex` setNode)
        return m
    ) Nothing
    where
        setNode = S.fromList nodes
        lenNodes  = length setNode

readConnection :: (Read a) => String -> (a, [a])
readConnection info = case lex info of
    [(node, left)] -> (read node, readListOfNode (snd $ head $ lex left))

readListOfNode :: (Read a) => String -> [a]
readListOfNode nodes = case lex nodes of
    [("", "")] -> []
    [(node, left)] -> read node : readListOfNode (snd $ head $ lex left)

readAllConnection :: (Read a) => String -> [(a, [a])]
readAllConnection = map readConnection . lines

allPossibleNode :: [(a,[a])] -> [a]
allPossibleNode = concatMap $ uncurry (:)

graphFromConnection ::Ord a =>  [(a,[a])] -> Graph a
graphFromConnection connections = GraphInternal
        nodes
        (runSTMatrix $ do
            m <- newMatrix 0 nodesLen nodesLen
            forM_ connections $ \ (from,connectee) ->
                forM_ connectee $ \ to ->
                    modifyMatrix m (findIndex from) (findIndex to) (1+)
            return m)
        Nothing
    where
        nodes = S.fromList $ allPossibleNode connections
        findIndex = flip S.findIndex nodes
        nodesLen = S.size nodes

euclidianCycle :: Ord a => Graph a -> Maybe [(a,a)]
euclidianCycle g = do
    (GraphInternal nodes m newCon) <- balance g
    conn <- runST $ do
        traverseM <- thawMatrix m
        path <- findEuclidianCycleAt nodes traverseM (Just 0)
        findNext nodes path traverseM
    return $ maybe conn (\(from,to) -> init $ uncurry (++) (swap (break ( (to == ).fst ) conn))) newCon
  where findNext nodes path m = do
            newPath <- findNextEuclidianPath nodes m path
            if isJust newPath
                then findNext nodes newPath m
                else return path

findNextEuclidianPath :: Ord a => S.Set a -> STMatrix s I -> Maybe [(a,a)] -> ST s (Maybe [(a,a)])
findNextEuclidianPath nodes mat path = do
    inMat <- freezeMatrix mat
    let newPath = do
            foundPath <- path
            let (left,right) = break (\(from,to) -> sum (map ti $ toList (inMat ! S.findIndex from nodes)) > 0) foundPath
            if null right
                then Nothing
                else Just (right ++ left)
    let maybeNext = fmap ((`S.findIndex` nodes) . fst . head) newPath
    maybeNextPath <- findEuclidianCycleAt nodes mat maybeNext
    return $ liftM2 (++) maybeNextPath newPath

findEuclidianCycleAt :: Ord a => S.Set a -> STMatrix s I -> Maybe Int -> ST s (Maybe [(a,a)])
findEuclidianCycleAt nodes mat startNode = if isJust startNode
                                                then findNextNod startNode
                                                else return Nothing
    where
        findNextNod n =  do
            next <- findNextEuclidianAt nodes mat (fromJust n)
            let cycleEnd = liftM2 (==) next startNode
            case (cycleEnd, next) of
                (Just True, Just to) -> return (Just [(S.elemAt (fromJust n) nodes,S.elemAt to nodes)])
                (Just False, to) -> do
                    list <- findNextNod to
                    return $ liftM2 (:) (Just (S.elemAt (fromJust n) nodes,S.elemAt (fromJust to) nodes)) list
                _ -> return Nothing

findNextEuclidianAt :: Ord a => S.Set a -> STMatrix s I -> Int -> ST s (Maybe Int)
findNextEuclidianAt nodes mat n = do
    con <- mapM (readMatrix mat n) [0..(S.size nodes - 1)]
    let next = find ((>0).snd) (zip [0..] con)
    when (isJust next) $ modifyMatrix mat n (fst . fromJust $ next) (subtract 1)
    return $ fst <$> next

mapFst :: (a -> b) -> [(a,c)] -> [(b,c)]
mapFst f = map (uncurry $ (,) . f)

canBeBalance :: Ord a => Graph a -> Bool
canBeBalance g = case unbal of
    [(_,-1), (_,1)] -> True
    [(_,1), (_,-1)] -> True
    [] -> True
    _ -> False
  where unbal = unbalanceNodes g

balance :: Ord a => Graph a -> Maybe (Graph a)
balance g@(GraphInternal node connection addCon) = case unbal of
    [(to,-1), (from,1)] -> return (GraphInternal node (fst (mutable (correctingMatrix (S.findIndex from node) (S.findIndex to node)) connection)) (Just (from,to)))
    [(from,1), (to,-1)] -> return (GraphInternal node (fst (mutable (correctingMatrix (S.findIndex from node) (S.findIndex to node)) connection)) (Just (from,to)))
    [] -> return g
    _ -> Nothing
  where
    unbal = unbalanceNodes g
    correctingMatrix f t _ m = modifyMatrix m f t (+1)

unbalanceNodes :: Ord a => Graph a -> [(a,Int)]
unbalanceNodes (GraphInternal nodes m _) =filter ((0 /= ) . snd) (mapFst (`S.elemAt` nodes) (map (\idx -> (idx,eulerianNode idx m)) [0..(S.size nodes - 1)]))

eulerianNode :: Int -> Matrix I -> Int
eulerianNode n m = ti (sum (concat colConnection)) - ti (sum (concat rowConnection))
    where
        colConnection = toLists $ m ¿ [n]
        rowConnection = toLists $ m ? [n]

eulerianNodeST :: Int -> STMatrix s I -> ST s Bool
eulerianNodeST n m = do
    colm <- extractMatrix m AllRows (Col n)
    rowm <- extractMatrix m (Row n) AllCols
    return $ (sum . concat . toLists) colm == (sum . concat . toLists) rowm


module Approx ( bezier2biarc
              ) where

import qualified CubicBezier as B
import qualified BiArc as BA
import qualified Line as L

import Data.Bool (bool)
import Linear
import Data.Complex

import Types

-- Approximate a bezier curve with biarcs (Left) and line segments (Right)
bezier2biarc :: B.CubicBezier
             -> Double
             -> [Either BA.BiArc (V2 Double)]
bezier2biarc mbezier resolution
    -- Edge case: all points on the same line -> it is a line
    | (L.isOnLine (L.fromPoints (B._p2 mbezier) (B._p1 mbezier)) (B._c1 mbezier)) &&
      (L.isOnLine (L.fromPoints (B._p2 mbezier) (B._p1 mbezier)) (B._c2 mbezier))
        = [Right (B._p2 mbezier)]
    -- Edge case: p1 == c1, don't split
    | (B._p1 mbezier) == (B._c1 mbezier)
        = approxOne mbezier
    -- Edge case: p2 == c2, don't split
    | (B._p2 mbezier) == (B._c2 mbezier)
        = approxOne mbezier
    -- Split by the inflexion points (if any)
    | otherwise
        = byInflection (B.realInflectionPoint i1) (B.realInflectionPoint i2)
    where
        (i1, i2) = B.inflectionPoints mbezier

        order a b | b < a = (b, a)
                  | otherwise = (a, b)

        byInflection True False = approxOne b1 ++ approxOne b2
            where
                (b1, b2) = B.bezierSplitAt mbezier (realPart i1)

        byInflection False True = approxOne b1 ++ approxOne b2
            where
                (b1, b2) = B.bezierSplitAt mbezier (realPart i2)

        byInflection True True = approxOne b1 ++ approxOne b2 ++ approxOne b3
            where
                (it1, it2') = order (realPart i1) (realPart i2)

                -- Make the first split and save the first new curve. The second one has to be splitted again
                -- at the recalculated t2 (it is on a new curve)
                it2 = (1 - it1) * it2'

                (b1, toSplit) = B.bezierSplitAt mbezier it1
                (b2, b3) = B.bezierSplitAt toSplit it2

        byInflection False False = approxOne mbezier

        -- TODO: make it tail recursive
        approxOne :: B.CubicBezier -> [Either BA.BiArc (V2 Double)]
        approxOne bezier
            -- Approximate bezier length. if smaller than resolution, do not approximate
            | (distance (B._p1 bezier) (B._c1 bezier)) +
              (distance (B._c1 bezier) (B._c2 bezier)) +
              (distance (B._c2 bezier) (B._p2 bezier)) < resolution
                = [Right (B._p2 bezier)]
            -- Edge case: start- and endpoints are the same
            | (B._p1 bezier) == (B._p2 bezier)
                = splitAndRecur 0.5
            -- Edge case: control lines are parallel
            | (L._m t1) == (L._m t2) || (isNaN (L._m t1) && isNaN (L._m t2))
                = splitAndRecur 0.5
            -- Approximation is not close enough yet, refine
            | BA.isStable biarc && maxDistance > resolution
                = splitAndRecur maxDistanceAt
            -- Desired case: approximation is stable and close enough
            | BA.isStable biarc
                = [Left biarc]
            -- Unstable approximation: split the bezier into half, basically switching to
            -- linear approximation mode
            | otherwise
                = splitAndRecur 0.5

            where
                -- Edge case: P1==C1 or P2==C2
                -- there is no derivative at P1 or P2, use the other control point
                c1 = bool (B._c1 bezier) (B._c2 bezier) ((B._p1 bezier) == (B._c1 bezier))
                c2 = bool (B._c2 bezier) (B._c1 bezier) ((B._p2 bezier) == (B._c2 bezier))

                -- V: Intersection point of tangent lines
                t1 = L.fromPoints (B._p1 bezier) c1
                t2 = L.fromPoints (B._p2 bezier) c2
                v = L.intersection t1 t2

                -- G: incenter point of the triangle (P1, V, P2)
                dP2V = distance (B._p2 bezier) v
                dP1V = distance (B._p1 bezier) v
                dP1P2 = distance (B._p1 bezier) (B._p2 bezier)
                g = (dP2V *^ B._p1 bezier + dP1V *^ B._p2 bezier + dP1P2 *^ v) ^/ (dP2V + dP1V + dP1P2)

                -- Calculate the BiArc
                biarc = BA.create (B._p1 bezier) (B._p1 bezier - c1) (B._p2 bezier) (B._p2 bezier - c2) g

                -- Calculate the error
                -- TODO: we only calculate the distance at 8 points (first and last skipped as
                --       they should be precise), seems a resonable approximation as for now
                parameterStep = 1 / 10

                (maxDistance, maxDistanceAt) = maxDistance' 0 0 parameterStep

                maxDistance' m mt t
                    | t < 1
                        = if' (d > m) (maxDistance' d t nt) (maxDistance' m mt nt)
                    | otherwise
                        = (m, mt)
                    where
                        d = distance (BA.pointAt biarc t) (B.pointAt bezier t)
                        nt = t + parameterStep

                splitAndRecur t = let (b1, b2) = B.bezierSplitAt bezier t
                                   in approxOne b1 ++ approxOne b2


module SchemeParser.Parser where

import Numeric (readInt, readFloat, readHex, readOct)
import Data.Ratio
import Data.Complex
import Data.Char (digitToInt)
import Data.Maybe (listToMaybe, fromJust)

import qualified Data.Map as M

import Data.Functor (($>))

import Text.Parsec hiding (spaces)
import Text.Parsec.Prim

import SchemeParser.Types

-- parser helpers

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

spaces :: LispParser ()
spaces = skipMany1 space

escapedChars :: LispParser Char
escapedChars = do
  char '\\'
  c <- oneOf "\"nrt\\"
  return $ case c of
    'n' -> '\n'
    'r' -> '\r'
    't' -> '\t'
    _   -> c

readBin :: Integral a => String -> Maybe a
readBin = fmap fst . listToMaybe . readInt 2 (`elem` "01") digitToInt

--
-- Lisp value parsers
--

-- primitives

parseAtom :: LispParser LispVal
parseAtom = do
  first <- letter <|> symbol
  rest <- many (letter <|> digit <|> symbol)
  let atom = first : rest
  return $ LAtom atom

parseString :: LispParser LispVal
parseString = do
  char '"'
  x <- many $ escapedChars <|> noneOf "\""
  char '"'
  return $ LString x

parseChar :: LispParser LispVal
parseChar = do
  try (string "#\\")
  c <- try (string "newline" <|> string "space")
       <|> do { x <- anyChar; notFollowedBy alphaNum; return [x]}
  return $ LChar $ case c of
    "newline" -> '\n'
    "space"   -> ' '
    _         -> head c

parseBool :: LispParser LispVal
parseBool =
  try (string "#t" $> LBool True) <|> try (string "#f" $> LBool False)

-- http://www.schemers.org/Documents/Standards/R5RS/HTML/r5rs-Z-H-9.html#%_sec_6.2.4
parseNumberRadix :: LispParser LispVal
parseNumberRadix =
  parseBin <|> parseOct <|> parseDec <|> parseHex
  where
    parseBin = do
      try (string "#b")
      LNumber . fromJust . readBin <$> many1 (oneOf "01")
    parseOct = do
      try (string "#o")
      LNumber . fst . (!! 0) . readOct <$> many1 octDigit
    parseDec = do
      try (string "#d")
      parseNumber
    parseHex = do
      try (string "#x")
      LNumber . fst . (!! 0) . readHex <$> many1 hexDigit

parseNumber :: LispParser LispVal
parseNumber = parseNumberRadix <|> LNumber . read <$> many1 digit

parseFloat :: LispParser LispVal
parseFloat = do
  l <- many1 digit
  char '.'
  r <- many1 digit
  return $ LFloat $ fst $ head $ readFloat (l ++ "." ++ r)

parseRatio :: LispParser LispVal
parseRatio = do
  n <- many1 digit
  char '/'
  d <- many1 digit
  return $ LRatio (read n % read d)

parseComplex :: LispParser LispVal
parseComplex = do
  r <- try parseFloat <|> parseNumber
  char '+'
  i <- try parseFloat <|> parseNumber
  char 'i'
  return $ LComplex (toDouble r :+ toDouble i)

-- lists

parseList :: LispParser LispVal
parseList = LList <$> sepBy parseExpr spaces

parseDottedList :: LispParser LispVal
parseDottedList = do
  hd <- endBy parseExpr spaces
  tl <- char '.' >> spaces >> parseExpr
  return $ LDottedList hd tl

parseQuoted :: LispParser LispVal
parseQuoted = do
  char '\''
  x <- parseExpr
  return $ LList [LAtom "quote", x]

parseQuasiQuoted :: LispParser LispVal
parseQuasiQuoted = do
  char '`'
  x <- parseExpr
  return $ LList [LAtom "quasiquote", x]

parseUnquoted :: LispParser LispVal
parseUnquoted = do
  char ','
  x <- parseExpr
  return $ LList [LAtom "unquote", x]

parseUnquoteSpliced :: LispParser LispVal
parseUnquoteSpliced = do
  try $ string ",@"
  x <- parseExpr
  return $ LList [LAtom "unquote-splice", x]

parseVector :: LispParser LispVal
parseVector = LVector <$> sepBy parseExpr spaces

--
-- helpers
--

toDouble :: LispVal -> Double
toDouble (LFloat f) = realToFrac f
toDouble (LNumber n) = fromIntegral n

--

parseExpr :: LispParser LispVal
parseExpr = parseAtom
            <|> parseString
            <|> parseChar
            <|> try parseFloat
            <|> try parseRatio
            <|> try parseComplex
            <|> parseNumber
            <|> parseBool
            <|> parseQuoted
            <|> parseQuasiQuoted
            <|> parseUnquoteSpliced
            <|> parseUnquoted
            <|> char '(' *> (try parseList <|> parseDottedList) <* char ')'
            <|> try (string "#(" *> parseVector <* char ')')


{-# LANGUAGE ScopedTypeVariables #-}

module FM (FMOptions(..), run, parseOpts) where

import           Control.Monad
import           Control.Monad.State.Lazy (StateT(runStateT), MonadState (get, put))
import           Data.Complex
import qualified Data.Vector.Storable as V
import qualified Filter as F
import           IQ
import           Options.Applicative
import           Pipes
import           Render
import qualified Sound.File.Sndfile as SF
import qualified Sound.File.Sndfile.Buffer.Vector as SV
import           System.IO
import           ZeroStuff

data FMOptions = FMOptions
  { inputFile :: String
  , outputFile :: String
  , deviationRatio :: Double
  , upsample :: Int
  , filterOpts :: F.FilterOptions }

parseOpts :: Parser FMOptions
parseOpts = FMOptions
         <$> strOption
             ( long "in-file"
            <> short 'i'
            <> metavar "INPUT_FILE"
            <> help "Input sound file" )
         <*> strOption
             ( long "out-file"
            <> short 'o'
            <> metavar "OUTPUT_FILE"
            <> help "Output IQ file")
         <*> option auto
             ( long "deviation-ratio"
            <> short 'd'
            <> value 1
            <> showDefault
            <> metavar "DEV_RATIO"
            <> help "The raio of max_devation over max_mod_frequency.")
         <*> option auto
             ( long "upsample-factor"
            <> short 'u'
            <> value 10
            <> showDefault
            <> metavar "UPSAMPLE_FAC"
            <> help "Upsample factor, outSR = inSR * UPSAMPLE_FAC")
         <*> F.parseOpts

-- clamp phase to [-pi, pi]
clampPhase :: Double -> Double
clampPhase phase
  | phase > pi = clampPhase $ phase - (2 * pi)
  | phase < -pi = clampPhase $ phase + (2 * pi)
  | otherwise =  phase

modulateFM :: MonadIO  m => Double -> Pipe Float IQ m ()
modulateFM devRatio = void $ flip runStateT (0 :: Double) $ forever $ do
  sample <- lift await
  phase <- get
  let phase' = clampPhase $ phase + realToFrac sample * devRatio
  lift . yield $ realToFrac (cos phase') :+ realToFrac (sin phase')
  put phase'

yieldEvery :: Monad m => Int -> Pipe a a m ()
yieldEvery n = do
  replicateM_ (n - 1) await
  await >>= yield
  yieldEvery n

run :: FMOptions -> IO ()
run opts = do
  (info, Just (aSamps :: SV.Buffer Float)) <- SF.readFile $ inputFile opts

  let sampsVec = SV.fromBuffer aSamps
      upsampleRate = upsample opts * SF.samplerate info
      filterKernel = F.lowPass upsampleRate $ filterOpts opts

  putStr "Processing samples... "
  hFlush stdout

  withBinaryFile (outputFile opts) WriteMode  $ \f ->
    runEffect $ V.mapM_ yield sampsVec
            >-> yieldEvery (SF.channels info)
            >-> modulateFM (deviationRatio opts)
            >-> zeroStuff (upsample opts)
            >-> F.convolve filterKernel
            >-> cfileSink f

  putStrLn "Done."

  putStrLn $  "CFile sample rate: " ++ show  upsampleRate ++ "Hz"


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
module AI.Singularity.Examples.Mnist where
import GHC.TypeLits
import qualified Numeric.LinearAlgebra.Data as LD
import qualified Numeric.LinearAlgebra.Devel as LDev
import qualified Numeric.LinearAlgebra as LA
import qualified Control.Monad.State as S
import Control.Lens
import Data.Proxy
import Control.Monad.ST
import System.Random
import Control.Monad
import Control.Arrow
import qualified Data.Vector.Storable as VS

import AI.Singularity.Network

getLabels :: String -> IO [Float]
getLabels = map (fromIntegral . fromEnum) . drop 8 . BS.unpack <$< BS.readFile

getImages2 :: String -> IO [[Vector 28]]
getImages2 = (map . map) fromListV . splitByLen 28 . splitByLen 28 . map ((/255). fromIntegral . fromEnum) . drop 16 . BS.unpack <$< BS.readFile

getSet3 :: String -> IO [((Vector 400, Vector 384), Vector 10)]
getSet3 = map (toVec >>> first splitV) <$< getSet

getSet2 :: String -> IO [([Vector 28],Vector 10)]
getSet2 nm = zip <$> getImages2 ("mnist/"++nm++"-images-idx3"++rest) <*> (map wynik <$> getLabels ("mnist/"++nm++"-labels-idx1"++rest))
  where rest          = "-ubyte"
        wynik liczba  = Vec . VS.map (\x -> if x == liczba then 1 else 0) . VS.fromList $ [0..9]

getImages :: String -> IO [VS.Vector Float]
getImages = map VS.fromList . splitByLen (28*28) . map ((/255). fromIntegral . fromEnum) . drop 16 . BS.unpack <$< BS.readFile

-- getImagesP :: FilePath -> Producer PB.ByteString IO a
getImagesC fn = let s = do
                      x <- CB.take 28
                      if BS.length x == 0
                        then return ()
                        else yield (map ((/255) . fromIntegral . fromEnum) . BS.unpack $ x) >> s
                    v = do
                      x <- CL.take 28
                      if null x
                        then return ()
                        else yield x >> v
                in CB.sourceFile fn .| (CB.drop 16 >> s) .| CL.map (fromListV :: [Float] -> Vector 28) .| v

getLinearC fn = let s = do
                      x <- BS.unpack <$> CB.take (28*28)
                      if null x
                        then return ()
                        else do
                              yield . map ((/255). fromIntegral . fromEnum) $ x
                              s
                in CB.sourceFile fn .| (CB.drop 16 >> s) .| CL.map (fromListV :: [Float] -> Vector 784)

fix f = let x = f x in x

-- setConduit :: String -> Source
setConduit nm = zipSources (getImagesC im) (getLabelsC lab)
  where im   = "mnist/"++nm++"-images-idx3"++rest
        lab  = "mnist/"++nm++"-labels-idx1"++rest
        rest = "-ubyte"

setConduit2 nm = zipSources (getLinearC im) (getLabelsC lab)
  where im   = "mnist/"++nm++"-images-idx3"++rest
        lab  = "mnist/"++nm++"-labels-idx1"++rest
        rest = "-ubyte"

getLabelsC fn =
  CB.sourceFile fn .|
  (CB.drop 8 >> fix (\f -> CB.head >>= \case { Nothing -> return (); Just x -> yield x >> f} )) .| CL.map (fromEnum >>> wynik)
    where wynik n = Vec . VS.map (\x -> if x == n then 1 else 0) . VS.fromList $ [0..9] :: Vector 10

  --map (fromIntegral . fromEnum) . drop 8 . BS.unpack

getSet :: String -> IO [(VS.Vector Float, Float)]
getSet nm = zip <$> getImages ("mnist/"++nm++"-images-idx3"++rest) <*> getLabels ("mnist/"++nm++"-labels-idx1"++rest)
  where rest = "-ubyte"

printSet :: (VS.Vector Float, Float) -> IO ()
printSet (v,a) = do
  print a
  mapM_ (\(i,val) -> (if i `mod` 28 == 0 then putStr (printVal val ++ "\n") else putStr (printVal val) )) . zip [1..] . VS.toList $ v

printVal :: Float -> String
printVal x = if x < 0.00001 then "  " else "* "

toVec :: (VS.Vector Float, Float) -> (Vector 784, Vector 10)
toVec (!obraz, !liczba) = (Vec obraz, Vec wynik)
    where wynik = VS.map (\x -> if x == liczba then 1 else 0) . VS.fromList $ [0..9]

printVec :: Vector 784 -> String
printVec (Vec !v) =  concatMap (\(i,val) -> (if i `mod` 28 == 0 then printVal val ++ "\n" else printVal val )) . zip [1..] . VS.toList $ v

testVec :: (VS.Vector Float, Float) -> (Vector 784, Int)
testVec (!obrac, !liczba) = (Vec obrac, round liczba)


-- | Functionality for graphing 1-dimensional vectors.
module System.Console.Ansigraph.Internal.Horizontal (
    displayRV
  , displayCV
  , displayPV
  , renderPV
  , renderRV
  , renderCV
) where

import System.Console.Ansigraph.Internal.Core

import Data.Complex
import Control.Monad.IO.Class (MonadIO)
-- for GHC <= 7.8
import Control.Applicative

---- Graphing Infrastructure  ----

barChars = "█▇▆▅▄▃▂▁ "

-- These values delineate regions for rounding to the nearest 1/8.
barVals :: [Double]
barVals = (+ 0.0625) . (/8) <$> [7,6..0]
     -- = [15/16, 13/16, 11/16, 9/16, 7/16, 5/16, 3/16, 1/16]

{- forward and reverse versions of unicode bar selection
   for positive and negative graph regions respectively -}

bars, barsR :: [(Double,Char)]
bars  = zip barVals barChars

barsR = zip barVals (reverse barChars)

selectBar, selectBarR :: Double -> Char
selectBar x = let l = filter (\p -> fst p < x) bars in
  case l of
       []     -> ' '
       (p:_) -> snd p

selectBarR x = let l = filter (\p -> fst p < x) barsR in
  case l of
       []     -> '█'
       (p:_) -> snd p

-- | Simple vector to String rendering that assumes positive input. Yields String of Unicode chars
--   representing graph bars varying in units of 1/8. The IO 'display' functions are preferable
--   for most use cases.
renderPV :: [Double] -> String
renderPV xs = let mx = maximum (filter (>= 0) $ 0:xs) in
              (selectBar . (/mx)) <$> xs

-- | Simple real vector rendering as a pair of strings. The IO 'display' functions are
--   preferable for most use cases.
renderRV :: [Double] -> (String,String)
renderRV l = let rp = l
                 rm = negate <$> rp
                 mx = maximum $ rp ++ rm
  in (selectBar  . (/mx) <$> rp,
      selectBarR . (/mx) <$> rm)

-- | Simple complex vector rendering as a pair of strings. The IO 'display' functions are
--   preferable for most use cases.
renderCV :: [Complex Double] -> (String,String,String,String)
renderCV l = let rp = realPart <$> l
                 rm = negate   <$> rp
                 ip = imagPart <$> l
                 im = negate   <$> ip
                 mx = maximum $ rp ++ rm ++ ip ++ im
  in (selectBar  . (/mx) <$> rp,
      selectBarR . (/mx) <$> rm,
      selectBar  . (/mx) <$> ip,
      selectBarR . (/mx) <$> im)

-- | ANSI based display for positive real vectors. Primarily invoked via 'graph', 'graphWith',
--   'animate', 'animateWith'.
displayPV :: MonadIO m => GraphSettings -> [Double] -> m ()
displayPV s l = let (rp,_) = renderRV l
                    rcol   = realColors s in colorStrLn rcol rp

-- | ANSI based display for real vectors. Primarily invoked via 'graph', 'graphWith',
--   'animate', 'animateWith'.
displayRV :: MonadIO m => GraphSettings -> [Double] -> m ()
displayRV s l = let (rp,rm) = renderRV l
                    rcol    = realColors s
  in do colorStrLn rcol          rp
        colorStrLn (invert rcol) rm

-- | ANSI based display for complex vectors. Primarily invoked via 'graph', 'graphWith',
--   'animate', 'animateWith'.
displayCV :: MonadIO m => GraphSettings -> [Complex Double] -> m ()
displayCV s l = let (rp,rm,ip,im) = renderCV l
                    (rcol,icol)   = colorSets s
  in do colorStrLn rcol          rp
        colorStrLn (invert rcol) rm
        colorStrLn icol          ip
        colorStrLn (invert icol) im


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ViewPatterns        #-}
module FokkerPlanck.GPUKernel where

import           Data.Array.Accelerate              as A
import           Data.Array.Accelerate.Data.Complex as A
import qualified Data.Complex                       as C
import Debug.Trace

{-# INLINE moveParticle #-}
moveParticle ::
     (A.RealFloat a, A.Num a, Elt a-- , FromIntegral Int a
     )
  => Exp (a, a, a, a)
  -> Exp (a, a, a, a)
moveParticle (unlift -> (phi, rho, theta, r)) =
  let !x = theta - phi
      !cosX = A.cos x
      !newPhi = phi + (A.atan2 (r * A.sin x) (rho + r * cosX))
      !newRho = A.sqrt $ (rho * rho + r * r + 2 * r * rho * cosX)
      -- !xx = A.round $ newRho * A.cos newPhi :: Exp Int
      -- !y = A.round $ newRho * A.sin newPhi :: Exp Int
      -- !newRho' = A.sqrt . A.fromIntegral $ xx * xx + y * y
      -- !newPhi' = A.atan2 (A.fromIntegral y) (A.fromIntegral xx)
  in lift (newPhi, newRho, theta, r)

{-# INLINE normalizationTerm #-}
normalizationTerm ::
     (A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a))
  => Exp a
  -> Exp a
  -> Exp (Complex a)
normalizationTerm halfLogPeriod freq =
  (lift $ 1 A.:+ ((-A.pi) * freq / halfLogPeriod)) /
  (A.exp (lift $ halfLogPeriod A.:+ (-A.pi * freq)) -
   A.exp (lift $ (-halfLogPeriod) A.:+ (A.pi * freq)))

-- {-# INLINE coefficient #-}
-- coefficient ::
--      forall a. (A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a), Prelude.Fractional a)
--   => Exp a
--   -> Exp a
--   -> Exp a
--   -> Exp a
--   -> Exp a
--   -> Exp (a, a, a, a)
--   -> Exp (A.Complex a)
-- coefficient halfLogPeriod rFreq thetaFreq rhoFreq phiFreq particle -- (unlift -> (phi, rho, theta, r))
--  =
--   let (phi, rho, theta, r) = unlift particle :: (Exp a, Exp a, Exp a, Exp a)
--      -- (normalizationTerm halfLogPeriod rhoFreq) *
--      -- (normalizationTerm halfLogPeriod rFreq) *
--      -- (lift (A.cos (phiFreq * phi + thetaFreq * (theta - phi)) A.:+ 0)) *
--                            -- (A.cis $
--                            --  (-2 * A.pi) * (rhoFreq * rho + rFreq * (r - rho)) /
--                            --  (A.exp halfLogPeriod))

--   in (lift $
--       ((A.exp $ (-0.5) * (rho + r)) *
--        (A.cos (phiFreq * phi  + thetaFreq * (theta - phi)))
--       ) A.:+
--       0) *
--      (A.cis $ (-1) * ((rhoFreq * rho + rFreq * (r - rho)
--                       ) -- * (2 * A.pi) / (A.exp halfLogPeriod)
--                       -- + phiFreq * phi  + thetaFreq * (theta - phi)
--                      ))
--      -- (lift $ ((A.cos (phiFreq * phi + thetaFreq * (theta - phi) / 2))) A.:+ 0) *
--      -- (A.cis $ (-2 * A.pi) * (rhoFreq * rho + rFreq * (r - rho)) / (A.exp halfLogPeriod))
--      -- ((lift $ rho A.:+ 0) A.** (lift $ (-0.5) A.:+ (rFreq - rhoFreq))) *
--      -- ((lift $ r A.:+ 0) A.** (lift $ (-0.5) A.:+ (-rFreq)))
--      -- (A.cis $ (-A.pi) * (rhoFreq * rho + rFreq * (r - rho)
--      --                        ) / (halfLogPeriod))
--      -- (A.cis $
--      --  (-A.pi) * (rhoFreq * (A.log rho) + rFreq * (A.log (r / rho))) /
--      --  halfLogPeriod)
--   -- in ((lift $ rho A.:+ 0) A.**
--   --     (lift $ 0 A.:+ (-A.pi) * (rhoFreq - rFreq) / halfLogPeriod)) *
--   --    ((lift $ r A.:+ 0) A.** (lift $ 0 A.:+ (-A.pi) * rFreq / halfLogPeriod)) *
--   --    (A.cis $ (-phiFreq) * phi - thetaFreq * (theta - phi))

coefficient ::
     forall a. (A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a), Prelude.Fractional a)
  => Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp (a, a, a, a)
  -> Exp (A.Complex a)
coefficient halfLogPeriod rFreq thetaFreq rhoFreq phiFreq particle =
  let (phi, rho, theta, r) = unlift particle :: (Exp a, Exp a, Exp a, Exp a)
  in (lift $
      ((A.exp $ (-0.5) * (rho + r)) *
       (A.cos (phiFreq * phi + thetaFreq * (theta - phi)))) A.:+
      0) *
     (A.cis $ (-1) * (rhoFreq * rho + rFreq * (r + rho)))

-- {-# INLINE gpuKernel #-}
gpuKernel ::
     forall a. (A.Eq a, A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a), A.FromIntegral Int a, Prelude.Fractional a)
  => Exp a
  -> Exp a
  -> Exp (A.Complex a)
  -> Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (A.Complex a))
gpuKernel !maxScaleExp !halfLogPeriodExp !deltaLogRhoComplexExp freqArr particles =
  let -- delta = constant 0.01
      movedParticleArr =
        compute .
        A.map
          (\particle ->
             let (phi, rho, theta, r) =
                   unlift particle :: (Exp a, Exp a, Exp a, Exp a)
                 logRho = A.log rho
                   -- (A.fromIntegral $
                   --  (A.round (((A.log rho) + halfLogPeriodExp) / delta) :: Exp Int)) *
                   -- delta -
                   -- halfLogPeriodExp
                 -- newRho =
                 --   (A.fromIntegral $ (A.round (rho / delta) :: Exp Int)) * delta :: Exp a
             in lift (phi, logRho, theta, (A.log r) :: Exp a)) -- .
        -- afst .
        -- A.filter
        --   (\particle ->
        --      let (_, rho, _, _) =
        --            unlift particle :: (Exp a, Exp a, Exp a, Exp a)
        --      in rho A.> (A.constant 0) ) -- .
        -- A.map moveParticle
        $
        particles
  in A.map
       (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq))
          -- (lift $ delta A.:+ 0) /
          -- (lift $ (8 * A.pi * A.pi * halfLogPeriodExp) A.:+ 0) *
         ->
          (sfoldl
             (\s particle ->
                s +
                (coefficient
                   halfLogPeriodExp
                   rFreq
                   thetaFreq
                   rhoFreq
                   phiFreq
                   particle))
             0
             (constant Z)
             movedParticleArr) * deltaLogRhoComplexExp
          -- (lift $
          --  (16 * A.pi * A.pi * halfLogPeriodExp * halfLogPeriodExp) A.:+ 0)
             -- movedParticleArr
        ) $
     freqArr

{-# INLINE convolveKernel #-}
convolveKernel ::
     (A.Num a, A.RealFloat a, A.Elt (Complex a))
  => Acc (Array DIM4 (Complex a))
  -> Acc (Array DIM4 (Complex a))
  -> Acc (Scalar Int)
  -> Acc (Scalar Int)
  -> Acc (Array DIM4 (Complex a))
  -> Acc (Vector (Complex a))
convolveKernel coefficients harmonics thetaIdx rIdx input =
  let (Z :. numRhoFreq :. numPhiFreq :. cols :. rows) =
        unlift . shape $ input :: (Z :. Exp Int :. Exp Int :. Exp Int :. Exp Int)
      coefficientsArr =
        A.replicate (lift (Z :. All :. All :. cols :. rows)) .
        slice coefficients $
        (lift (Z :. (the rIdx) :. (the thetaIdx) :. All :. All))
      harmonicsArr =
        backpermute
          (shape input)
          (\(unlift -> Z :. rho :. phi :. col :. row :: (Z :. Exp Int :. Exp Int :. Exp Int :. Exp Int)) ->
             lift (Z :. (rho + the rIdx) :. (phi + the thetaIdx) :. col :. row))
          harmonics
  in flatten .
     A.sum .
     A.sum .
     backpermute
       (lift (Z :. cols :. rows :. numRhoFreq :. numPhiFreq))
       (\(unlift -> Z :. col :. row :. rho :. phi :: (Z :. Exp Int :. Exp Int :. Exp Int :. Exp Int)) ->
          lift (Z :. rho :. phi :. col :. row)) .
     A.zipWith (*) harmonicsArr . A.zipWith (*) coefficientsArr $
     input

{-# INLINE coefficient' #-}
coefficient' ::
     forall a.
     ( A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , Prelude.Fractional a
     )
  => Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp (a, a, a, a, a)
  -> Exp (A.Complex a)
coefficient' sigma rFreq thetaFreq rhoFreq phiFreq particle =
  let (phi, rho, theta, r, v) =
        unlift particle :: (Exp a, Exp a, Exp a, Exp a, Exp a)
  in (lift $ (v * (A.exp $ (sigma - 1) * (rho + r))) A.:+ 0) *
     (A.cis $
      (-1) *
      (rhoFreq * rho + rFreq * (r - rho) + phiFreq * phi +
       thetaFreq * (theta - phi)))

gpuKernel' ::
     forall a.
     ( A.Eq a
     , A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , A.FromIntegral Int a
     , Prelude.Fractional a
     )
  => Exp a
  -> Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (a, a, a, a, a))
  -> Acc (A.Vector (A.Complex a))
gpuKernel' sigma freqArr xs =
  A.map
    (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq)) ->
       A.sfoldl
         (\s particle ->
            s + (coefficient' sigma rFreq thetaFreq rhoFreq phiFreq particle))
         0
         (constant Z)
         xs)
    freqArr

{-# INLINE coefficient'' #-}
coefficient'' ::
     forall a.
     ( A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , Prelude.Fractional a
     )
  => Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp (a, a, a, a, a)
  -> Exp (A.Complex a)
coefficient'' sigma period rFreq thetaFreq rhoFreq phiFreq particle =
  let (phi, rho, theta, r, v) =
        unlift particle :: (Exp a, Exp a, Exp a, Exp a, Exp a)
   in lift
        ((v * A.exp ((sigma - 1) * rho) *
          A.cos (phiFreq * phi + thetaFreq * (theta - phi))
         ) :+
         0) *
      A.cis ((-1) * (2 * A.pi / period * (rhoFreq * rho + rFreq * (r - rho))))

gpuKernel'' ::
     forall a.
     ( A.Eq a
     , A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , A.FromIntegral Int a
     , Prelude.Fractional a
     )
  => Exp a
  -> Exp a
  -> Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (a, a, a, a, a))
  -> Acc (A.Vector (A.Complex a))
gpuKernel'' sigma period freqArr xs =
  A.map
    (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq)) ->
       A.sfoldl
         (\s particle ->
            s +
            coefficient'' sigma period rFreq thetaFreq rhoFreq phiFreq particle)
         0
         (constant Z)
         xs)
    freqArr

{-# INLINE pinwheelAcc #-}
pinwheelAcc ::
     forall a. (A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a), Prelude.Fractional a)
  => Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp (a, a, a, a, A.Complex a)
  -> Exp (A.Complex a)
pinwheelAcc rFreq thetaFreq rhoFreq phiFreq particle =
  let (phi, rho, theta, r, v) =
        unlift particle :: (Exp a, Exp a, Exp a, Exp a, Exp (A.Complex a))
  in v * (lift $ ((A.exp $ (-0.5) * (rho + r))) A.:+ 0) *
     (A.cis $
      (-1) *
      (rhoFreq * rho + rFreq * (r - rho) + phiFreq * phi +
       thetaFreq * (theta - phi)))

pinwheelCoefficientsAcc ::
     forall a.
     ( A.Eq a
     , A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , A.FromIntegral Int a
     , Prelude.Fractional a
     )
  => Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (a, a, a, a, A.Complex a))
  -> Acc (A.Vector (A.Complex a))
pinwheelCoefficientsAcc freqArr xs =
  A.map
    (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq)) ->
       A.sfoldl
         (\s particle ->
            s + (pinwheelAcc rFreq thetaFreq rhoFreq phiFreq particle))
         0
         (constant Z)
         xs)
    freqArr


{-# LANGUAGE TypeFamilies, GeneralizedNewtypeDeriving, DeriveGeneric #-}

module BayesStack.Models.Topic.SharedTaste
  ( -- * Primitives
    STData(..)
  , ItemSource(..)
  , Node(..), Item(..), Topic(..)
  , NodeItem, setupNodeItems
  , Friendship(..), otherFriend, isFriend, getFriends
    -- * Initialization
  , ModelInit
  , randomInitialize, smartInitialize
    -- * Model
  , STModel(..), ItemUnit
  , model, modelLikelihood
  , STModelState (..), getModelState
  , sortTopics
  , ItemVars(..)
    -- * Quantities of interest
  , friendInfluence
  , theta
  ) where

import Prelude hiding (mapM, sum, product)

import Data.EnumMap (EnumMap)
import qualified Data.EnumMap as EM

import Data.Sequence (Seq)
import qualified Data.Sequence as SQ

import Data.Set (Set)
import qualified Data.Set as S

import qualified Data.EnumSet as ES
import qualified Data.Vector as V

import Data.Traversable
import Data.Foldable
import Data.Monoid
import Data.Function (on)
import Data.List (sortBy)

import Control.Monad (liftM)
import Control.Monad.Trans.State
import Control.Monad.Trans.Class

import Data.Random
import Data.Random.Distribution.Bernoulli
import Data.Random.Sequence
import Data.Number.LogFloat hiding (realToFrac)

import BayesStack.Core
import BayesStack.Categorical
import BayesStack.DirMulti
import BayesStack.TupleEnum
import BayesStack.Models.Topic.Types

import Control.Monad (when)
import Control.Monad.IO.Class

import Statistics.Sample
import Data.Serialize (Serialize)
import GHC.Generics

data ItemSource = Shared | Own deriving (Show, Eq, Generic, Enum, Ord)
instance Serialize ItemSource

data STData = STData { stAlphaGammaShared, stAlphaGammaOwn :: Double
                     , stAlphaOmega :: Double
                     , stAlphaPsi :: Double
                     , stAlphaLambda :: Double
                     , stAlphaPhi :: Double
                     , stNodes :: Set Node
                     , stFriendships :: Set Friendship
                     , stItems :: Set Item
                     , stTopics :: Set Topic
                     , stNodeItems :: EnumMap NodeItem (Node, Item)
                     }
               deriving (Show, Eq, Generic)
instance Serialize STData

data ItemVars = ItemVars { ivS :: !ItemSource
                         , ivF :: !Node
                         , ivT :: !Topic
                         }
                deriving (Show, Eq, Generic)
instance Serialize ItemVars

data STModel = STModel { mData :: STData
                       , mGammas :: SharedEnumMap Node (Multinom ItemSource)
                       , mOmegas :: SharedEnumMap Node (Multinom Topic)
                       , mPsis :: SharedEnumMap Node (Multinom Node)
                       , mLambdas :: SharedEnumMap Friendship (Multinom Topic)
                       , mPhis :: SharedEnumMap Topic (Multinom Item)
                       , mVars :: SharedEnumMap NodeItem ItemVars
                       , mSortedTopics :: SharedEnumMap Item [Topic]
                       }

data STModelState = STModelState { msData :: STData
                                 , msGammas :: EnumMap Node (Multinom ItemSource)
                                 , msOmegas :: EnumMap Node (Multinom Topic)
                                 , msPsis :: EnumMap Node (Multinom Node)
                                 , msLambdas :: EnumMap Friendship (Multinom Topic)
                                 , msPhis :: EnumMap Topic (Multinom Item)
                                 , msVars :: EnumMap NodeItem ItemVars
                                 , msLogLikelihood :: Double
                                 } deriving (Show, Generic)
instance Serialize STModelState

type ModelInit = EnumMap NodeItem ItemVars

randomInitialize' :: STData -> ModelInit -> RVar ModelInit
randomInitialize' d init =
  let unset = EM.keysSet (stNodeItems d) `ES.difference` EM.keysSet init
      topics = S.toList $ stTopics d
      randomInit :: NodeItem -> RVar ModelInit
      randomInit ni = do t <- randomElement topics
                         let (n,_) = EM.findWithDefault (error "Can't find nodeItem") ni (stNodeItems d)
                             friends = getFriends (S.toList $ stFriendships d) n
                         f <- randomElement friends
                         s <- bernoulli $ stAlphaGammaShared d
                         let s' = case s of
                                        True  -> Shared
                                        False -> Own
                         return $ EM.singleton ni $ ItemVars s' f t
  in liftM ((init `mappend`) . mconcat) $ forM (ES.toList unset) randomInit

randomInitialize :: STData -> RVar ModelInit
randomInitialize = (flip randomInitialize') EM.empty

smartInitialize :: STData -> RVar ModelInit
smartInitialize d =
  let STData {stTopics=topics, stNodes=nodes, stItems=items, stNodeItems=nodeItems} = d
      STData {stFriendships=friendships, stNodeItems=nis} = d
      nisInv :: EnumMap (Node,Item) [NodeItem]
      nisInv = EM.fromListWith (++) $ map (\(ni,(n,i))->((n,i),[ni])) $ EM.toList nis

      sharedTopics :: Friendship -> StateT (EnumMap Item Topic, EnumMap Friendship (Set Topic)) RVar ModelInit
      sharedTopics fs@(Friendship (a,b)) =
        let findItems n = S.fromList $ map snd $ filter (\(n',i)->n==n') $ toList nodeItems
            sharedItems = toList $ S.intersection (findItems a) (findItems b)
        in liftM mconcat $ forM sharedItems $ \x -> do
             (itemTopics,_) <- get
             t <- if x `EM.member` itemTopics
                    then return $ EM.findWithDefault (error "Item has no topic") x itemTopics
                    else do (_, friendshipTopics) <- get
                            --let possTopics = EM.findWithDefault topics fs friendshipTopics
                            let possTopics = topics -- FIXME
                            t <- lift $ randomElement $ S.toList possTopics
                            modify $ \(a,b)->(EM.insert x t a, b)
                            return t
             modify $ \(a,b)->(a, EM.insertWith S.union fs (S.singleton t) b)
             return $ EM.fromList $ do ax <- EM.findWithDefault (error "ouch") (a,x) nisInv
                                       return (ax, ItemVars Shared b t)
                                 ++ do bx <- EM.findWithDefault (error "ouch") (b,x) nisInv
                                       return (bx, ItemVars Shared a t)
  in do a <- evalStateT (mapM sharedTopics $ S.toList friendships) (EM.empty, EM.empty)
        randomInitialize' d $ mconcat a

data ItemUnit = ItemUnit { iuModel :: STModel
                         , iuNodeItem :: NodeItem
                         , iuFriends :: Set Node
                         , iuN :: Node
                         , iuVars :: Shared ItemVars
                         , iuX :: Item
                         , iuGamma :: Shared (Multinom ItemSource)
                         , iuOmega :: Shared (Multinom Topic)
                         , iuLambdas :: SharedEnumMap Friendship (Multinom Topic)
                         , iuPhis :: SharedEnumMap Topic (Multinom Item)
                         , iuState :: Shared GibbsUpdateState
                         }

model :: STData -> ModelInit -> ModelMonad (Seq ItemUnit, STModel)
model d init =
  do let STData {stTopics=topics, stNodes=nodes, stItems=items, stNodeItems=nis} = d
         STData {stFriendships=friendships} = d
         friends :: EnumMap Node (Set Node)
         friends = foldMap (\n->EM.singleton n $ S.fromList $ getFriends (S.toList friendships) n) nodes
     gammas <- newSharedEnumMap (S.toList nodes) $ \n ->
       --return $ dirMulti [ (Shared, stAlphaGammaShared d)
       --                  , (Own, stAlphaGammaOwn d) ]
       return $ multinom [ (Shared, stAlphaGammaShared d)
                         , (Own, stAlphaGammaOwn d) ]
     omegas <- newSharedEnumMap (S.toList nodes) $ \n ->
       return $ symDirMulti (stAlphaOmega d) (S.toList topics)
     psis <- newSharedEnumMap (S.toList nodes) $ \n ->
       return $ symDirMulti (stAlphaPsi d) (S.toList nodes)
     lambdas <- newSharedEnumMap (S.toList friendships) $ \n ->
       return $ symDirMulti (stAlphaLambda d) (S.toList topics)
     phis <- newSharedEnumMap (S.toList topics) $ \t ->
       return $ symDirMulti (stAlphaPhi d) (S.toList items)

     ivs <- newSharedEnumMap (EM.keys nis) $ \ni ->
       return $ EM.findWithDefault (error "Incomplete initialization") ni init

     sortedTopics <- newSharedEnumMap (S.toList items) $ \x -> return $ S.toList topics

     let model = STModel { mData = d
                         , mGammas = gammas
                         , mOmegas = omegas
                         , mPsis = psis
                         , mLambdas = lambdas
                         , mPhis = phis
                         , mVars = ivs
                         , mSortedTopics = sortedTopics
                         }

     itemUnits <- forM (EM.toList ivs) $ \(ni,iv) ->
       do state <- newGibbsUpdateState
          let (n,x) = nis EM.! ni
          let unit = ItemUnit { iuModel = model
                              , iuNodeItem = ni
                              , iuFriends = friends EM.! n
                              , iuN = n
                              , iuVars = iv
                              , iuX = x
                              , iuGamma = gammas EM.! n
                              , iuOmega = omegas EM.! n
                              , iuLambdas = EM.filterWithKey (\k _->isFriend n k) lambdas
                              , iuPhis = phis
                              , iuState = state
                              }
          getShared iv >>= guSet unit
          return unit
     return (SQ.fromList itemUnits, model)

sortTopics :: STModel -> ModelMonad ()
sortTopics model =
  forM_ (EM.toList $ mSortedTopics model) $ \(x,topics)->do
    d <- getShared topics
    weights <- forM d $ \t->do phi <- getShared $ mPhis model EM.! t
                               return $ sampleProb phi x -- FIXME
    setShared topics $ map snd $ sortBy (flip (compare `on` fst)) $ zip weights d

modelLikelihood :: STModelState -> Probability
modelLikelihood model =
  product $ map likelihood (EM.elems $ msGammas model)
         ++ map likelihood (EM.elems $ msPhis model)
         ++ map likelihood (EM.elems $ msLambdas model)
         ++ map likelihood (EM.elems $ msOmegas model)
         ++ map likelihood (EM.elems $ msPsis model)

instance GibbsUpdateUnit ItemUnit where
  type GUValue ItemUnit = ItemVars
  guProb unit (ItemVars s f t) =
    do gamma <- getShared $ iuGamma unit
       omega <- getShared $ iuOmega unit
       psi <- getShared $ mPsis (iuModel unit) EM.! iuN unit
       phi <- getShared $ iuPhis unit EM.! t
       lambda <- getShared $ iuLambdas unit EM.! Friendship (iuN unit, f)
       case s of
            Shared -> return $ sampleProb gamma s
                             * sampleProb psi f
                             * sampleProb lambda t
                             * sampleProb phi (iuX unit)
            Own -> return $ sampleProb gamma s
                          * sampleProb omega t
                          * sampleProb phi (iuX unit)

  guDomain unit = return $ (do t <- S.toList $ stTopics $ mData $ iuModel unit
                               f <- S.toList $ iuFriends unit
                               return $ ItemVars Shared f t)
                        ++ (do t <- S.toList $ stTopics $ mData $ iuModel unit
                               let f = head $ S.toList $ iuFriends unit
                               return $ ItemVars Own f t)

  guUnset unit =
    do ItemVars s f t <- getShared $ iuVars unit
       let x = iuX unit
           u = iuN unit
           m = iuModel unit
           gamma = iuGamma unit
           omega = iuOmega unit
           lambda = iuLambdas unit EM.! Friendship (iuN unit, f)
           phi = iuPhis unit EM.! t
       gamma `updateShared` decMultinom s
       case s of
            Shared -> do (mPsis m EM.! u) `updateShared` decMultinom f
                         (mPsis m EM.! f) `updateShared` decMultinom u
                         lambda `updateShared` decMultinom t
            Own -> omega `updateShared` decMultinom t
       phi `updateShared` decMultinom x
       return $ ItemVars s f t

  guSet unit iv@(ItemVars s f t) =
    do iuVars unit `setShared` iv
       let x = iuX unit
           u = iuN unit
           m = iuModel unit
           gamma = iuGamma unit
           omega = iuOmega unit
           lambda = iuLambdas unit EM.! Friendship (iuN unit, f)
           phi = iuPhis unit EM.! t
       gamma `updateShared` incMultinom s
       case s of
            Shared -> do (mPsis m EM.! u) `updateShared` incMultinom f
                         (mPsis m EM.! f) `updateShared` incMultinom u
                         lambda `updateShared` incMultinom t
            Own -> omega `updateShared` incMultinom t
       phi `updateShared` incMultinom x

  guState = iuState

getModelState :: STModel -> ModelMonad STModelState
getModelState model =
  do gammas <- getSharedEnumMap $ mGammas model
     omegas <- getSharedEnumMap $ mOmegas model
     psis <- getSharedEnumMap $ mPsis model
     lambdas <- getSharedEnumMap $ mLambdas model
     phis <- getSharedEnumMap $ mPhis model
     vars <- getSharedEnumMap $ mVars model
     let state = STModelState { msData = mData model
                              , msGammas = gammas
                              , msOmegas = omegas
                              , msPsis = psis
                              , msLambdas = lambdas
                              , msPhis = phis
                              , msVars = vars
                              , msLogLikelihood = logFromLogFloat $ modelLikelihood state
                              }
     return state

-- | The analogue of theta in LDA
theta :: STModelState -> Node -> Topic -> Double
theta state u t =
  sampleProb gamma Own * sampleProb omega t
  + sampleProb gamma Shared
  * sum (map (\f->let lambda = msLambdas state EM.! Friendship (u,f)
                  in sampleProb psi f * sampleProb lambda t
             )
         $ getFriends (S.toList $ stFriendships $ msData state) u
        )
  where psi = msPsis state EM.! u
        gamma = msGammas state EM.! u
        omega = msOmegas state EM.! u

friendInfluence :: STModelState -> Node -> Node -> Double
friendInfluence state u f =
  let lambda = msLambdas state EM.! Friendship(u,f)
      vars = map snd
             $ filter (\(ni,iv)->let (u',x) = stNodeItems (msData state) EM.! ni in u==u')
             $ EM.assocs $ msVars state
      tProbF = map (realToFrac . sampleProb lambda . ivT) vars
  in case tProbF of
       [] -> error "friendInfluence: vars is null"
       otherwise -> geometricMean $ V.fromList tProbF


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
module Math.HiddenMarkovModel (
   T(..), Distr.State, state,
   Discrete, DiscreteTrained,
   Gaussian, GaussianTrained,
   uniform,
   generate,
   probabilitySequence,
   Normalized.logLikelihood,
   Normalized.reveal,
   Normalized.posterior,

   Trained(..),
   trainSupervised,
   Normalized.trainUnsupervised,
   mergeTrained, finishTraining, trainMany,
   deviation,

   toCSV,
   fromCSV,
   ) where

import qualified Math.HiddenMarkovModel.Distribution as Distr
import qualified Math.HiddenMarkovModel.Normalized as Normalized
import qualified Math.HiddenMarkovModel.CSV as HMMCSV
import Math.HiddenMarkovModel.Private
          (T(..), Trained(..), mergeTrained, toCells, parseCSV)
import Math.HiddenMarkovModel.Distribution (State(State))
import Math.HiddenMarkovModel.Utility
          (randomItemProp, normalizeProb, attachOnes)

import qualified Numeric.LinearAlgebra.Algorithms as Algo
import qualified Numeric.Container as NC
import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)

import qualified Text.CSV.Lazy.String as CSV

import qualified System.Random as Rnd

import qualified Control.Monad.Exception.Synchronous as ME
import qualified Control.Monad.Trans.State as MS
import qualified Control.Monad.HT as Monad

import qualified Data.NonEmpty as NonEmpty
import qualified Data.Array as Array
import Data.Traversable (Traversable, mapAccumL)
import Data.Foldable (Foldable)
import Data.Array (accumArray)

state :: Int -> State
state = State

type DiscreteTrained prob symbol = Trained (Distr.DiscreteTrained prob symbol) prob
type Discrete prob symbol = T (Distr.Discrete prob symbol) prob

type GaussianTrained a = Trained (Distr.GaussianTrained a) a
type Gaussian a = T (Distr.Gaussian a) a

{- |
Create a model with uniform probabilities
for initial vector and transition matrix
given a distribution for the emissions.
You can use this as a starting point for 'Normalized.trainUnsupervised'.
-}
uniform ::
   (Distr.Info distr, Distr.Probability distr ~ prob) =>
   distr -> T distr prob
uniform distr =
   let n = Distr.numberOfStates distr
       c = recip $ fromIntegral n
   in  Cons {
          initial = NC.constant c n,
          transition = NC.konst c (n,n),
          distribution = distr
       }

probabilitySequence ::
   (Traversable f, Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> f (State, emission) -> f prob
probabilitySequence hmm =
   snd
   .
   mapAccumL
      (\index (State s, e) ->
         (NC.atIndex (transition hmm) . flip (,) s,
          index s * Distr.emissionStateProb (distribution hmm) e (State s)))
      (NC.atIndex (initial hmm))

generate ::
   (Rnd.RandomGen g, Ord prob, Rnd.Random prob,
    Distr.Generate distr, Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> g -> [emission]
generate hmm =
   MS.evalState $
   flip MS.evalStateT (initial hmm) $
   Monad.repeat $ MS.StateT $ \v0 -> do
      s <- randomItemProp $ zip [0..] (Vector.toList v0)
      x <- Distr.generate (distribution hmm) (State s)
      return (x, takeColumn s $ transition hmm)

takeColumn :: (Matrix.Element a) => Int -> Matrix a -> Vector a
takeColumn n  =  Matrix.flatten . Matrix.extractColumns [n]

{- |
Contribute a manually labeled emission sequence to a HMM training.
-}
trainSupervised ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   Int -> NonEmpty.T [] (State, emission) -> Trained tdistr prob
trainSupervised n xs =
   let getState (State s, _x) = s
   in  Trained {
          trainedInitial = NC.assoc n 0 [(getState (NonEmpty.head xs), 1)],
          trainedTransition =
             Matrix.trans $ NC.accum (NC.konst 0 (n,n)) (+) $
             attachOnes $ NonEmpty.mapAdjacent (,) $ fmap getState xs,
          trainedDistribution =
             Distr.accumulateEmissions $ map attachOnes $ Array.elems $
             accumArray (flip (:)) [] (State 0, State (n-1)) $ NonEmpty.flatten xs
       }

finishTraining ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob) =>
   Trained tdistr prob -> T distr prob
finishTraining hmm =
   Cons {
      initial = normalizeProb $ trainedInitial hmm,
      transition =
         Matrix.fromColumns $ map normalizeProb $
         Matrix.toColumns $ trainedTransition hmm,
      distribution = Distr.normalize $ trainedDistribution hmm
   }

trainMany ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob,
    Foldable f) =>
   (trainingData -> Trained tdistr prob) ->
   NonEmpty.T f trainingData -> T distr prob
trainMany train =
   finishTraining . NonEmpty.foldl1Map mergeTrained train

{- |
Compute maximum deviation between initial and transition probabilities.
You can use this as abort criterion for unsupervised training.
We omit computation of differences between the emission probabilities.
This simplifies matters a lot and
should suffice for defining an abort criterion.
-}
deviation ::
   (Algo.Field prob, Ord prob) => T distr prob -> T distr prob -> prob
deviation hmm0 hmm1 =
   deviationVec (initial hmm0) (initial hmm1)
   `max`
   deviationVec (transition hmm0) (transition hmm1)

deviationVec ::
   (Ord a, NC.Container c a) =>
   c a -> c a -> a
deviationVec x y =
   let d = NC.sub x y
   in  NC.maxElement d `max` negate (NC.minElement d)

toCSV ::
   (Distr.CSV distr, Algo.Field prob, Show prob) =>
   T distr prob -> String
toCSV hmm =
   CSV.ppCSVTable $ snd $ CSV.toCSVTable $ HMMCSV.padTable "" $
   toCells hmm

fromCSV ::
   (Distr.CSV distr, Algo.Field prob, Read prob) =>
   String -> ME.Exceptional String (T distr prob)
fromCSV =
   MS.evalStateT parseCSV . map HMMCSV.fixShortRow . CSV.parseCSV


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE FlexibleContexts      #-}

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Random
import           Control.Monad.Trans.Except

import qualified Data.Attoparsec.Text as A
import           Data.List ( foldl' )
#if ! MIN_VERSION_base(4,13,0)
import           Data.Semigroup ( (<>) )
#endif

import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Vector.Storable as V

import           Numeric.LinearAlgebra ( maxIndex )
import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative

import           Grenade
import           Grenade.Utils.OneHot

-- It's logistic regression!
--
-- This network is used to show how we can embed a Network as a layer in the larger MNIST
-- type.
type FL i o =
  Network
    '[ FullyConnected i o, Logit ]
    '[ 'D1 i, 'D1 o, 'D1 o ]

-- The definition of our convolutional neural network.
-- In the type signature, we have a type level list of shapes which are passed between the layers.
-- One can see that the images we are inputing are two dimensional with 28 * 28 pixels.

-- It's important to keep the type signatures, as there's many layers which can "squeeze" into the gaps
-- between the shapes, so inference can't do it all for us.

-- With the mnist data from Kaggle normalised to doubles between 0 and 1, learning rate of 0.01 and 15 iterations,
-- this network should get down to about a 1.3% error rate.
--
-- /NOTE:/ This model is actually too complex for MNIST, and one should use the type given in the readme instead.
--         This one is just here to demonstrate Inception layers in use.
--
type MNIST =
  Network
    '[ Reshape,
       Concat ('D3 28 28 1) Trivial ('D3 28 28 14) (InceptionMini 28 28 1 5 9),
       Pooling 2 2 2 2, Relu,
       Concat ('D3 14 14 3) (Convolution 15 3 1 1 1 1) ('D3 14 14 15) (InceptionMini 14 14 15 5 10), Crop 1 1 1 1, Pooling 3 3 3 3, Relu,
       Reshape, FL 288 80, FL 80 10 ]
    '[ 'D2 28 28, 'D3 28 28 1,
       'D3 28 28 15, 'D3 14 14 15, 'D3 14 14 15, 'D3 14 14 18,
       'D3 12 12 18, 'D3 4 4 18, 'D3 4 4 18,
       'D1 288, 'D1 80, 'D1 10 ]

randomMnist :: MonadRandom m => m MNIST
randomMnist = randomNetwork

convTest :: Int -> FilePath -> FilePath -> LearningParameters -> ExceptT String IO ()
convTest iterations trainFile validateFile rate = do
  net0         <- lift randomMnist
  trainData    <- readMNIST trainFile
  validateData <- readMNIST validateFile
  lift $ foldM_ (runIteration trainData validateData) net0 [1..iterations]

    where
  trainEach rate' !network (i, o) = train rate' network i o

  runIteration trainRows validateRows net i = do
    let trained' = foldl' (trainEach ( rate { learningRate = learningRate rate * 0.9 ^ i} )) net trainRows
    let res      = fmap (\(rowP,rowL) -> (rowL,) $ runNet trained' rowP) validateRows
    let res'     = fmap (\(S1D label, S1D prediction) -> (maxIndex (SA.extract label), maxIndex (SA.extract prediction))) res
    print trained'
    putStrLn $ "Iteration " ++ show i ++ ": " ++ show (length (filter ((==) <$> fst <*> snd) res')) ++ " of " ++ show (length res')
    return trained'

data MnistOpts = MnistOpts FilePath FilePath Int LearningParameters

mnist' :: Parser MnistOpts
mnist' = MnistOpts <$> argument str (metavar "TRAIN")
                   <*> argument str (metavar "VALIDATE")
                   <*> option auto (long "iterations" <> short 'i' <> value 15)
                   <*> (LearningParameters
                       <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                       <*> option auto (long "momentum" <> value 0.9)
                       <*> option auto (long "l2" <> value 0.0005)
                       )

main :: IO ()
main = do
    MnistOpts mnist vali iter rate <- execParser (info (mnist' <**> helper) idm)
    putStrLn "Training convolutional neural network..."

    res <- runExceptT $ convTest iter mnist vali rate
    case res of
      Right () -> pure ()
      Left err -> putStrLn err

readMNIST :: FilePath -> ExceptT String IO [(S ('D2 28 28), S ('D1 10))]
readMNIST mnist = ExceptT $ do
  mnistdata <- T.readFile mnist
  return $ traverse (A.parseOnly parseMNIST) (T.lines mnistdata)

parseMNIST :: A.Parser (S ('D2 28 28), S ('D1 10))
parseMNIST = do
  Just lab <- oneHot <$> A.decimal
  pixels   <- many (A.char ',' >> A.double)
  image    <- maybe (fail "Parsed row was of an incorrect size") pure (fromStorable . V.fromList $ pixels)
  return (image, lab)