module CommonTypes
  (LispVal(..),
   LispError(..),
   ThrowsError,
   trapError,
   extractValue,
   Env,
   IOThrowsError,
   nullEnv,
   liftThrows,
   runIOThrows)
  where

import qualified Data.Vector as V
import Data.Ratio
import Data.Complex
import Control.Monad.Except
import Text.ParserCombinators.Parsec
import Data.IORef
import System.IO

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

instance Show LispVal where show = showVal

showVal :: LispVal -> String
showVal (String string) = "\"" ++ string ++ "\""
showVal (Atom name) = name
showVal (Number num) = show num
showVal (Bool value) = if value then "#t" else "#f"
showVal (List contents) = "(" ++ unwordsList contents ++ ")"
showVal (DottedList head tail) = "(" ++ unwordsList head ++ " . " ++ showVal tail ++ ")"
showVal (Vector vector) = "#(" ++ (unwordsList $ V.toList vector) ++ ")"
showVal (Float value) = show value
showVal (Ratio value) = show (numerator value) ++ "/" ++ show (denominator value)
showVal (Complex value) = show (realPart value) ++ "+" ++ show (imagPart value) ++ "i"
showVal (Character char) | char == '\n' = "#\\newline"
                         | char == ' '  = "#\\space"
                         | otherwise    = "#\\" ++ [char]
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 primitve>"

unwordsList = unwords . map showVal

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

instance Show LispError where show = showError

type ThrowsError = Either LispError

showError :: LispError -> String
showError (UnboundVar message varname)  = message ++ ": " ++ varname
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func)    = message ++ ": " ++ show func
showError (NumArgs expected found)      = "Expected " ++ show expected
                                       ++ " args; found values " ++ unwordsList found
showError (TypeMismatch expected found) = "Invalid type: expected " ++ expected
                                       ++ ", found " ++ show found
showError (Parser parseErr)             = "Parse error at " ++ show parseErr

trapError action = catchError action (return . show)

extractValue :: ThrowsError a -> a
extractValue (Right val) = val

type Env = IORef [(String, IORef LispVal)]
type IOThrowsError = ExceptT LispError IO

nullEnv :: IO Env
nullEnv = newIORef []

liftThrows :: ThrowsError a -> IOThrowsError a
liftThrows (Left err) = throwError err
liftThrows (Right val) = return val

runIOThrows :: IOThrowsError String -> IO String
runIOThrows action = runExceptT (trapError action) >>= return . extractValue


{-# LANGUAGE ScopedTypeVariables #-}
module Main (main) where

import Data.Complex (Complex)
import Data.Int (Int8, Int16, Int32, Int64)
import Data.List.NonEmpty (NonEmpty)
import Data.Proxy (Proxy (..))
import Data.Typeable (Typeable, typeRep)
import Data.Word (Word8, Word16, Word32, Word64)
import Numeric.Natural (Natural)
import Test.QuickCheck (Arbitrary (..), Property, counterexample, label, (===),
                        sized, chooseInt, vectorOf)
import Test.QuickCheck.Instances ()
import Test.Tasty (defaultMain, testGroup, TestTree)
import Test.Tasty.QuickCheck (testProperty)

import qualified Data.List as L

import Data.Discrimination
import Utils

main :: IO ()
main = defaultMain $ testGroup "discrimination"
  [ testGroup "examples"
    [ testGroup "nub"
      [ testProperty "List.nub" $
        let prop :: [Word64] -> Property
            prop xs = L.nub xs === nub xs
        in prop

      , testProperty "ordNub" $
        let prop :: [Word64] -> Property
            prop xs = ordNub xs === nub xs
        in prop

      , testProperty "hashNub" $
        let prop :: [Word64] -> Property
            prop xs = hashNub xs === nub xs
        in prop
      ]

    , testGroup "sort"
      [ testProperty "List.sort" $
        let prop :: [Word64] -> Property
            prop xs = L.sort xs === sort xs
        in prop

      , testProperty "introsort" $
        -- for Word64 unstable sort works too
        let prop :: [Word64] -> Property
            prop xs = introsort xs === sort xs
        in prop

      , testProperty "mergesort" $
        let prop :: [Word64] -> Property
            prop xs = mergesort xs === sort xs
        in prop
      ]
    ]

  , testGroup "Grouping"
    [ testGrouping (Proxy :: Proxy ())
    , testGrouping (Proxy :: Proxy Int)
    , testGrouping (Proxy :: Proxy Int8)
    , testGrouping (Proxy :: Proxy Int16)
    , testGrouping (Proxy :: Proxy Int32)
    , testGrouping (Proxy :: Proxy Int64)
    , testGrouping (Proxy :: Proxy Word)
    , testGrouping (Proxy :: Proxy Word8)
    , testGrouping (Proxy :: Proxy Word16)
    , testGrouping (Proxy :: Proxy Word32)
    , testGrouping (Proxy :: Proxy Word64)
    , testGrouping (Proxy :: Proxy Bool)
    , testGrouping (Proxy :: Proxy Ordering)
    , testGrouping (Proxy :: Proxy (Word8,Word8))
    , testGrouping (Proxy :: Proxy (Word8,Word8,Word8))
    , testGrouping (Proxy :: Proxy (Word8,Word8,Word8,Word8))
    , testGrouping (Proxy :: Proxy Rational)
    , testGrouping (Proxy :: Proxy (Complex Word8))
    , testGrouping (Proxy :: Proxy (Maybe Word8))
    , testGrouping (Proxy :: Proxy (Either Word8 Word8))
    , testGrouping (Proxy :: Proxy Char)
    , testGrouping (Proxy :: Proxy String)
    , testGrouping (Proxy :: Proxy (NonEmpty Int))
    , testGrouping (Proxy :: Proxy Natural)
    , testGrouping (Proxy :: Proxy Integer)

    , testGrouping' listToNatural
    , testGrouping' listToInteger
    ]

  , testGroup "Sorting"
    [ testSorting (Proxy :: Proxy ())
    , testSorting (Proxy :: Proxy Int)
    , testSorting (Proxy :: Proxy Int8)
    , testSorting (Proxy :: Proxy Int16)
    , testSorting (Proxy :: Proxy Int32)
    , testSorting (Proxy :: Proxy Int64)
    , testSorting (Proxy :: Proxy Word)
    , testSorting (Proxy :: Proxy Word8)
    , testSorting (Proxy :: Proxy Word16)
    , testSorting (Proxy :: Proxy Word32)
    , testSorting (Proxy :: Proxy Word64)
    , testSorting (Proxy :: Proxy Bool)
    , testSorting (Proxy :: Proxy Ordering)
    , testSorting (Proxy :: Proxy (Word8,Word8))
    , testSorting (Proxy :: Proxy (Word8,Word8,Word8))
    , testSorting (Proxy :: Proxy (Word8,Word8,Word8,Word8))
    , testSorting (Proxy :: Proxy (Maybe Word8))
    , testSorting (Proxy :: Proxy (Either Word8 Word8))
    , testSorting (Proxy :: Proxy Char)
    , testSorting (Proxy :: Proxy String)
    , testSorting (Proxy :: Proxy (NonEmpty Int))
    , testSorting (Proxy :: Proxy Natural)
    , testSorting (Proxy :: Proxy Integer)

    , testSorting' listToNatural
    , testSorting' listToInteger
    ]
  ]

listToNatural :: SmallList Word64 -> Natural
listToNatural = L.foldl' (\x y -> x * 2 ^ (64 :: Int) + fromIntegral y) 0 . getSmallList

listToInteger :: SmallList Int64 -> Integer
listToInteger = L.foldl' (\x y -> x * 2 ^ (64 :: Int) + fromIntegral y) 0 . getSmallList

newtype SmallList a = SmallList { getSmallList :: [a] } deriving (Eq, Show)

instance Arbitrary a => Arbitrary (SmallList a) where
    arbitrary = sized $ \n -> do
        m <- chooseInt (0, min 10 n)
        SmallList <$> vectorOf m arbitrary

    shrink = fmap SmallList . shrink . getSmallList

testGrouping
  :: forall a. (Grouping a, Typeable a, Arbitrary a, Eq a, Show a)
  => Proxy a
  -> TestTree
testGrouping _ = testGrouping' (id :: a -> a)

testGrouping'
  :: forall a b. (Grouping b, Typeable a, Typeable b, Arbitrary a, Eq b, Show a, Show b)
  => (a -> b)
  -> TestTree
testGrouping' f = testGroup name
    [ testProperty "groupingEq" prop_eq
    , testProperty "nub"        prop_nub
    ]
  where
    tra = typeRep (Proxy :: Proxy a)
    trb = typeRep (Proxy :: Proxy b)
    name = if tra == trb then show tra else show trb ++ " from " ++ show tra

    prop_eq :: a -> a -> Property
    prop_eq x' y' =
        counterexample (show (x,y)) $
        label (show lhs) $
        lhs === groupingEq x y
      where
        x = f x'
        y = f y'
        lhs = x == y

    prop_nub :: [a] -> Property
    prop_nub xs' = L.nub xs === nub xs
      where
        xs = take 100 (map f xs')

testSorting
  :: forall a. (Sorting a, Typeable a, Arbitrary a, Ord a, Show a)
  => Proxy a
  -> TestTree
testSorting _ = testSorting' (id :: a -> a)

testSorting'
  :: forall a b. (Sorting b, Typeable a, Typeable b, Arbitrary a, Ord b, Show a, Show b)
  => (a -> b)
  -> TestTree
testSorting' f = testGroup name
    [ testProperty "sortingCompare" prop_cmp
    , testProperty "sort"           prop_sort
    ]
  where
    tra = typeRep (Proxy :: Proxy a)
    trb = typeRep (Proxy :: Proxy b)
    name = if tra == trb then show tra else show trb ++ " from " ++ show tra

    prop_cmp :: a -> a -> Property
    prop_cmp x' y' =
        counterexample (show (x,y)) $
        label (show lhs) $
        lhs === sortingCompare x y
      where
        x = f x'
        y = f y'
        lhs = compare x y

    prop_sort :: [a] -> Property
    prop_sort xs' = L.sort xs === sort xs
      where
        xs = map f xs'


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE InstanceSigs          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE Strict                #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Grenade.Layers.FullyConnected (
    FullyConnected (..)
  , FullyConnected' (..)
  , randomFullyConnected
  , SpecFullyConnected (..)
  , specFullyConnected
  , fullyConnected
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Primitive        (PrimBase, PrimState)
import           Control.Parallel.Strategies
import           Data.Maybe                     (fromMaybe)
import           Data.Reflection                (reifyNat)
import qualified Data.Vector.Storable           as V
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import           System.Random.MWC              hiding (create)
#if MIN_VERSION_singletons(2,6,0)
import           Data.Singletons.TypeLits       (SNat (..))
#endif
import           Data.List                      (foldl')
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.Prelude.Num    ((%*))
import qualified Numeric.LinearAlgebra          as LA
import           Numeric.LinearAlgebra.Static   hiding (zipWithVector)
import           Text.Printf

import           Control.Monad                  (void)
import           Foreign.Storable               (peekElemOff, pokeElemOff, sizeOf)
import           System.IO.Unsafe               (unsafePerformIO)

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build
import           Grenade.Layers.Internal.BLAS
import           Grenade.Layers.Internal.CUDA
import           Grenade.Layers.Internal.Update
import           Grenade.Types
import           Grenade.Utils.Conversion
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore
import           Grenade.Utils.Vector

import           Debug.Trace

-- | A basic fully connected (or inner product) neural network layer.
data FullyConnected i o = FullyConnected
                        !(FullyConnected' i o)             -- Neuron weights
                        !(ListStore (FullyConnected' i o)) -- momentum store
                        deriving (Generic)

instance NFData (FullyConnected i o) where
  rnf (FullyConnected w store) = rnf w `seq` rnf store

instance Show (FullyConnected i o) where
  show FullyConnected {} = "FullyConnected"

-- | How to store the data .
data FullyConnected' i o
  = FullyConnectedHMatrix
      !(R o)   -- ^ Bias
      !(L o i) -- ^ Activations
  | FullyConnectedBLAS
    !(Int, Int)         -- ^ Input, output
    !(V.Vector RealNum) -- ^ Bias, Temporary vector of same size
    !(V.Vector RealNum) -- ^ Activations
  deriving (Generic)

instance Show (FullyConnected' i o) where
  show FullyConnectedBLAS{}    = "FullyConnectedBLAS"
  show FullyConnectedHMatrix{} = "FullyConnectedHMatrix"

instance NFData (FullyConnected' i o) where
  rnf (FullyConnectedHMatrix b w)    = rnf b `seq` rnf w
  rnf (FullyConnectedBLAS !io !b !w) = rnf io `seq` rnf b `seq` rnf w

instance (KnownNat i, KnownNat o, KnownNat (i * o)) => UpdateLayer (FullyConnected i o) where
  type Gradient (FullyConnected i o) = (FullyConnected' i o)
  type MomentumStore (FullyConnected i o) = ListStore (FullyConnected' i o)
  runUpdate opt@OptSGD {} x@(FullyConnected (FullyConnectedHMatrix oldBias oldActivations) store) (FullyConnectedHMatrix biasGradient activationGradient) =
    let (FullyConnectedHMatrix oldBiasMomentum oldMomentum) = getData opt x store
        VectorResultSGD newBias newBiasMomentum = descendVector opt (VectorValuesSGD oldBias biasGradient oldBiasMomentum)
        MatrixResultSGD newActivations newMomentum = descendMatrix opt (MatrixValuesSGD oldActivations activationGradient oldMomentum)
        newStore = setData opt x store (FullyConnectedHMatrix newBiasMomentum newMomentum)
     in FullyConnected (FullyConnectedHMatrix newBias newActivations) newStore
  runUpdate opt@OptAdam {} x@(FullyConnected (FullyConnectedHMatrix oldBias oldActivations) store) (FullyConnectedHMatrix biasGradient activationGradient) =
    let [FullyConnectedHMatrix oldMBias oldMActivations, FullyConnectedHMatrix oldVBias oldVActivations] = getData opt x store
        VectorResultAdam newBias newMBias newVBias = descendVector opt (VectorValuesAdam (getStep store) oldBias biasGradient oldMBias oldVBias)
        MatrixResultAdam newActivations newMActivations newVActivations =
          descendMatrix opt (MatrixValuesAdam (getStep store) oldActivations activationGradient oldMActivations oldVActivations)
        newStore = setData opt x store [FullyConnectedHMatrix newMBias newMActivations, FullyConnectedHMatrix newVBias newVActivations]
     in FullyConnected (FullyConnectedHMatrix newBias newActivations) newStore
  runUpdate opt@OptSGD {} x@(FullyConnected (FullyConnectedBLAS io@(i,o) oldBiasH oldActivationsH) store) (FullyConnectedBLAS _ biasGradient activationGradient) =
      let oldBias = oldBiasH
          oldActivations = oldActivationsH
          (oldMBias, oldMActivations) = case getData opt x store of -- In the first periods until the store is filled newData is called, which will generate FullyConnectedHMatrix instances!
            FullyConnectedBLAS _ oldMBias' oldMActivations' -> (oldMBias', oldMActivations')
            FullyConnectedHMatrix oldMBias' oldMActivations' -> (extract oldMBias', extractM oldMActivations')
          VectorResultSGDV newBias newMBias               = descendVectorV opt (VectorValuesSGDV oldBias biasGradient oldMBias)
          MatrixResultSGDV newActivations newMActivations = descendMatrixV opt (MatrixValuesSGDV oldActivations activationGradient oldMActivations)
          newStore = setData opt x store (FullyConnectedBLAS io newMBias newMActivations)
      in FullyConnected (FullyConnectedBLAS io newBias newActivations) newStore
    where extractM mat = (\(S2DV vec) -> vec) . fromS2D $ S2D mat
  runUpdate opt@OptAdam {} x@(FullyConnected (FullyConnectedBLAS io@(i,o) oldBiasH oldActivationsH) store) (FullyConnectedBLAS _ biasGradient activationGradient) =
      let oldBias = oldBiasH
          oldActivations = oldActivationsH
          (oldMBias, oldMActivations, oldVBias, oldVActivations) = case getData opt x store of -- In the first periods until the store is filled newData is called, which will generate FullyConnectedHMatrix instances!
            [FullyConnectedBLAS _ oldMBias' oldMActivations', FullyConnectedBLAS _ oldVBias' oldVActivations'] -> (oldMBias', oldMActivations', oldVBias', oldVActivations')
            [FullyConnectedHMatrix oldMBias' oldMActivations', FullyConnectedHMatrix oldVBias' oldVActivations'] -> (extract oldMBias', extractM oldMActivations', extract oldVBias', extractM oldVActivations')
            [FullyConnectedBLAS _ oldMBias' oldMActivations', FullyConnectedHMatrix oldVBias' oldVActivations'] -> (oldMBias', oldMActivations', extract oldVBias', extractM oldVActivations')
            xs -> error $ "unexpected data in ListStore in FullyConnected BLAS implementation: " ++ show xs
          VectorResultAdamV newBias newMBias newVBias                      = descendVectorV opt (VectorValuesAdamV (getStep store) oldBias biasGradient oldMBias oldVBias)
          MatrixResultAdamV newActivations newMActivations newVActivations = descendMatrixV opt (MatrixValuesAdamV (getStep store) oldActivations activationGradient oldMActivations oldVActivations)
          newStore = setData opt x store [FullyConnectedBLAS io newMBias newMActivations, FullyConnectedBLAS io newVBias newVActivations]
      in FullyConnected (FullyConnectedBLAS io newBias newActivations) newStore
    where extractM mat = (\(S2DV vec) -> vec) . fromS2D $ S2D mat
  runUpdate opt (FullyConnected layer _) _ = error $ "Unexpected input in runUpdate in FullyConnected layer. Optimizer" ++ show opt ++ ". Layer: " ++ show layer

instance (KnownNat i, KnownNat o, KnownNat (i * o)) => LayerOptimizerData (FullyConnected i o) (Optimizer 'SGD) where
  type MomentumDataType (FullyConnected i o) (Optimizer 'SGD) = FullyConnected' i o
  getData opt x store = head $ getListStore opt x store
  setData opt x store = setListStore opt x store . return
  newData _ _ = FullyConnectedHMatrix (konst 0) (konst 0)

instance (KnownNat i, KnownNat o, KnownNat (i * o)) => LayerOptimizerData (FullyConnected i o) (Optimizer 'Adam) where
  type MomentumDataType (FullyConnected i o) (Optimizer 'Adam) = FullyConnected' i o
  type MomentumExpOptResult (FullyConnected i o) (Optimizer 'Adam) = [FullyConnected' i o]
  getData = getListStore
  setData = setListStore
  newData _ _ = FullyConnectedHMatrix (konst 0) (konst 0)

instance (KnownNat i, KnownNat o) => FoldableGradient (FullyConnected' i o) where
  mapGradient f (FullyConnectedHMatrix bias activations) = FullyConnectedHMatrix (dvmap f bias) (dmmap f activations)
  mapGradient f (FullyConnectedBLAS io bias activations) = FullyConnectedBLAS io (mapVector f bias) (mapVector f activations)
  squaredSums (FullyConnectedHMatrix bias activations) = [sumV . squareV $ bias, sumM . squareM $ activations]
  squaredSums (FullyConnectedBLAS _ bias activations) = [V.sum . mapVector (^(2::Int)) $ bias, V.sum . mapVector (^(2::Int)) $ activations]

runForward :: forall i o. (KnownNat i, KnownNat o) => FullyConnected i o -> S ('D1 i) -> (Tape (FullyConnected i o) ('D1 i) ('D1 o), S ('D1 o))
runForward (FullyConnected (FullyConnectedHMatrix wB wN) _) (S1D v) = (S1D v, S1D (wB + wN #> v))
runForward (FullyConnected (FullyConnectedBLAS _ wB wN) _) (S1DV v) =
  let inp = unsafeMemCopyVectorFromTo wB (createVectorUnsafe (V.length wB))
      !out' = unsafePerformIO $ matXVec BlasNoTranspose wN v 1 inp
   in force out' `seq` (S1DV v, S1DV out')
runForward lay@(FullyConnected (FullyConnectedHMatrix _ _) _ ) v@S1DV{} = runForward lay (toS1D v)
runForward lay@(FullyConnected FullyConnectedBLAS{} _) v@S1D{} = runForward lay (fromS1D v)

runBackward :: forall i o . (KnownNat i, KnownNat o) => FullyConnected i o -> Tape (FullyConnected i o) ('D1 i) ('D1 o) -> S ('D1 o) -> (Gradient (FullyConnected i o), S ('D1 i))
runBackward (FullyConnected (FullyConnectedHMatrix _ wN) _) (S1D x) (S1D dEdy) =
  let wB' = dEdy
      mm' = dEdy `outer` x
            -- calcluate derivatives for next step
      dWs = tr wN #> dEdy
   in (FullyConnectedHMatrix wB' mm', S1D dWs)
runBackward (FullyConnected (FullyConnectedBLAS io@(i, o) _ wN) _) (S1DV x) (S1DV dEdy) =
  let !mm' = unsafePerformIO $ outerV dEdy x
      dWsInp = createVectorUnsafe i
      !dWs' = unsafePerformIO $ matXVec BlasTranspose wN dEdy 0 dWsInp
   in force mm' `seq` force dWs' `seq` (FullyConnectedBLAS io dEdy mm', S1DV dWs')
runBackward l x dEdy = runBackward l x (toLayerShape x dEdy)

instance (KnownNat i, KnownNat o, KnownNat (i * o)) => Layer (FullyConnected i o) ('D1 i) ('D1 o) where
  type Tape (FullyConnected i o) ('D1 i) ('D1 o) = S ('D1 i)
  runForwards = runForward   -- Do a matrix vector multiplication and return the result.
  runBackwards = runBackward -- Run a backpropogation step for a full connected layer.

instance (KnownNat i, KnownNat o) => Serialize (FullyConnected i o) where
  put (FullyConnected w ms) = put w >> put ms
  get = FullyConnected <$> get <*> get

instance (KnownNat i, KnownNat o) => Serialize (FullyConnected' i o) where
  put (FullyConnectedHMatrix b w) = do
    put (0 :: Int)
    putListOf put . LA.toList . extract $ b
    putListOf put . LA.toList . LA.flatten . extract $ w
  put (FullyConnectedBLAS io b w) = do
    put (1 :: Int)
    put io
    putListOf put . V.toList $ b
    putListOf put . V.toList $ w
  get = do
    (nr :: Int) <- get
    case nr of
      0 -> do
        let f = fromIntegral $ natVal (Proxy :: Proxy i)
        b <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
        k <- maybe (fail "Vector of incorrect size") return . create . LA.reshape f . LA.fromList =<< getListOf get
        return $ FullyConnectedHMatrix b k
      1 -> do
        io <- get
        b <- V.fromList <$> getListOf get
        w <- V.fromList <$> getListOf get
        return $ FullyConnectedBLAS io b w
      _ -> error $ "Unexpected nr in get in Serialize of FullyConnected' " ++ show nr

instance (KnownNat i, KnownNat o, KnownNat (i*o)) => RandomLayer (FullyConnected i o) where
  createRandomWith = randomFullyConnected

randomFullyConnected ::
     forall m i o. (PrimBase m, KnownNat i, KnownNat o, KnownNat (i * o))
  => NetworkInitSettings
  -> Gen (PrimState m)
  -> m (FullyConnected i o)
randomFullyConnected (NetworkInitSettings m HMatrix _) gen = do
  wN <- getRandomMatrix i o m gen
  wB <- getRandomVector i o m gen
  return $!! FullyConnected (FullyConnectedHMatrix wB wN) mkListStore
  where i = natVal (Proxy :: Proxy i)
        o = natVal (Proxy :: Proxy o)
randomFullyConnected (NetworkInitSettings m BLAS _) gen = do

  wB <- getRandomVectorV i o o' m gen
  wN <- getRandomVectorV i o (i' * o') m gen
  return $!! FullyConnected (FullyConnectedBLAS (i', o') wB wN) mkListStore
  where
    i = natVal (Proxy :: Proxy i)
    i' = fromIntegral i
    o = natVal (Proxy :: Proxy o)
    o' = fromIntegral o

-------------------- DynamicNetwork instance --------------------

instance (KnownNat i, KnownNat o) => FromDynamicLayer (FullyConnected i o) where
  fromDynamicLayer _ _ _ = SpecNetLayer $ SpecFullyConnected (natVal (Proxy :: Proxy i)) (natVal (Proxy :: Proxy o))

instance ToDynamicLayer SpecFullyConnected where
  toDynamicLayer wInit gen (SpecFullyConnected nrI nrO) =
    reifyNat nrI $ \(pxInp :: (KnownNat i) => Proxy i) ->
      reifyNat nrO $ \(pxOut :: (KnownNat o') => Proxy o') ->
        case singByProxy pxInp %* singByProxy pxOut of
          SNat -> do
            (layer :: FullyConnected i o') <- randomFullyConnected wInit gen
            return $ SpecLayer layer (sing :: Sing ('D1 i)) (sing :: Sing ('D1 o'))

-- | Make a specification of a fully connected layer (see Grenade.Dynamic.Build for a user-interface to specifications).
specFullyConnected :: Integer -> Integer -> SpecNet
specFullyConnected nrI nrO = SpecNetLayer $ SpecFullyConnected nrI nrO

-- | A Fully-connected layer with input dimensions as given in last output layer and output dimensions specified. 1D only!
fullyConnected :: Integer -> BuildM ()
fullyConnected rows = do
  (inRows, _, _) <- buildRequireLastLayerOut Is1D
  buildAddSpec (SpecNetLayer $ SpecFullyConnected inRows rows)
  buildSetLastLayer (rows, 1, 1)

-------------------- GNum instances --------------------

instance (KnownNat i, KnownNat o) => GNum (FullyConnected i o) where
  s |* FullyConnected w store = FullyConnected (s |* w) (s |* store)
  FullyConnected w1 store1 |+ FullyConnected w2 store2 = FullyConnected (w1 |+ w2) (store1 |+ store2)

instance (KnownNat i, KnownNat o) => GNum (FullyConnected' i o) where
  s |* FullyConnectedHMatrix b w = FullyConnectedHMatrix (dvmap (fromRational s *) b) (dmmap (fromRational s *) w)
  s |* FullyConnectedBLAS io b w = FullyConnectedBLAS io (mapVector (fromRational s *) b) (mapVector (fromRational s *) w)
  FullyConnectedHMatrix b1 w1 |+ FullyConnectedHMatrix b2 w2 = FullyConnectedHMatrix (b1 + b2) (w1 + w2)
  FullyConnectedBLAS io b1 w1 |+ FullyConnectedBLAS _ b2 w2 = FullyConnectedBLAS io (zipWithVector (+) b2 b1) (zipWithVector (+) w2 w1)
  x |+ y = error $ "Cannot add different network types in |+ in FullyConnected: " ++ show (x, y)
  sumG xs@(FullyConnectedBLAS io _ _:_) = FullyConnectedBLAS io bs' ws'
    where
      (bs, ws) = unzip $ map (\(FullyConnectedBLAS _ b w) -> (b, w)) xs
      bs' = sumVectors bs `using` rparWith rdeepseq
      ws' = sumVectors ws `using` rparWith rdeepseq
  sumG xs = foldl1 (|+) xs


module Scheme.Parser (readLisp, readLisps) where
import Control.Monad
import Control.Monad.Error (throwError)

import Data.Char
import Data.Maybe
import Data.Complex
import Data.Ratio
import Data.Vector as V (fromList)

import Text.Parsec hiding (spaces)

import Scheme.Types

readLisp :: FilePath -> String -> ThrowsErrorIO LispVal
readLisp fname inp = case parse (spaces >> lispVal) fname inp of
    Left err  -> throwError $ ParserError err
    Right val -> return val

readLisps :: FilePath -> String -> ThrowsErrorIO [LispVal]
readLisps fname inp = case parse lispParser fname inp of
    Left err  -> throwError $ ParserError err
    Right val -> return val

type Parser a = Parsec String () a

lispParser = spaces >> sepEndBy lispVal spaces

lispVal =
    try lispSequence <|> try num    <|> try lispChar <|> try bool <|>
    try atom         <|> stringLisp <|> quoted

comment = char ';' >> manyTill anyChar (try (char '\n' <|> char '\r'))

ignored = void space <|> void comment

spaces = many ignored

spaces1 :: Parser ()
spaces1 = skipMany1 ignored

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

bool = char '#' >> (string "t" <|> string "f")
                >>= (\s -> return $ case s of "t" -> Bool True
                                              "f" -> Bool False)

atom = do first <- letter <|> symbol
          rest  <- many (letter <|> symbol <|> digit)
          return . Atom $ first:rest

lispChar :: Parser LispVal
lispChar = do
    string "#\\"
    special <- optionMaybe $ choice (map (string . fst) specialCharLits)
    maybe (liftM Character anyChar) (\lit -> return (Character . fromJust $ lookup lit specialCharLits)) special
  where specialCharLits = [("newline", '\n'), ("space", ' '), ("tab", '\t'), ("return", '\r')]

num = do
    let radices = [('x', 16), ('d', 10), ('o', 8), ('t', 3), ('b', 2)]
    maybeBase <- optionMaybe (try (char '#' >> oneOf (map fst radices)))
    maybeSign <- optionMaybe (char '+' <|> char '-')
    let sign :: Num a => a -> a
        sign = fromMaybe id $ maybeSign >>= \c -> if c == '-' then Just negate else Nothing
        base = case maybeBase of
                    Nothing -> 10
                    Just s -> fromMaybe 10 (lookup s radices)
    real <- try (liftM (Double   . sign) (parseFloat base))
        <|> try (liftM (Rational . sign) (parseRational base))
        <|>      liftM (Int      . sign) (parseInteger base)
    maybeImaginary <- optionMaybe (try $ num >>= \n -> char 'i' >> return n)
    return $ maybe (Num real) (\(Num imag) ->
        Num $ Complex $ getAsDouble real :+ getAsDouble imag) maybeImaginary

digitCharacters = ['0'..'9'] ++ ['A'..'Z']

toDigit :: (Num a, Enum a) => Char -> a
toDigit c = case lookup (toUpper c) $ zip digitCharacters [0..] of
    Nothing -> error $ "character not between 0 and 9 or A and Z: '" ++ [c] ++ "'"
    Just d -> d

parseInteger :: Int -> Parser Integer
parseInteger base = liftM round $ parseDigits False base

parseFloat :: Int -> Parser Double
parseFloat base = do big <- parseDigits False base
                     char '.'
                     small <- parseDigits True base
                     return $ big + small

parseRational :: Int -> Parser (Ratio Integer)
parseRational base = do n <- parseInteger base
                        char '/'
                        d <- parseInteger base
                        return $ n % d

parseDigits :: (Floating a, Enum a) => Bool -> Int -> Parser a
parseDigits isDecimal base =
    let allowedDigits = take base digitCharacters
        convertStr numAsStr = snd $ foldr
            (\c (plc, val) -> (plc + 1, val + toDigit c * fromIntegral base ** plc))
            (if isDecimal then (-1) * fromIntegral (length numAsStr) else 0.0, 0.0) numAsStr
    in liftM convertStr $ many1 (satisfy (flip elem allowedDigits . toUpper))

escapables = [('t', '\t'), ('n', '\n'), ('r', '\r'), ('\\', '\\')]

stringLisp :: Parser LispVal
stringLisp = do char '"'
                s <- many (noneOf "\"\\" <|> (try $ do
                        char '\\'
                        esc <- anyChar
                        return $ fromMaybe esc $ lookup esc escapables))
                char '"'
                return $ String s

quoted = do char '\''
            v <- lispVal
            return $ List [Atom "quote", v]

lispSequence = do
    isVector <-  (string "("  >> return False)
             <|> (string "#(" >> return True)
    spaces
    vs <- sepEndBy lispVal spaces1
    maybeDot <- optionMaybe (spaces >> char '.' >> spaces1 >> lispVal)
    spaces
    char ')'
    return $ maybe
        (if isVector then Vector (toVector vs) else List vs)
        (DottedList vs) maybeDot

toVector = V.fromList


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : Bench mark morris
-- | Creator: Xiao Ling
-- | Created: 12/8/2015
-- | see    : https://github.com/snoyberg/conduit
-- | TODO   : move this to benchmark directory
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module Bmorris where

import Data.List
import Data.Random
import Data.Conduit
import Data.List.Split
import qualified Data.Conduit.List as Cl

import Control.Monad.Identity

import Criterion.Main

import Core
import Statistics
import Morris

{-----------------------------------------------------------------------------
  Benchmark  counter vs counter'

    Benchmark results: it's unclear whether counter is faster than counter'
    due to noise:

    benchmarking count vs count'/count 500
      time                 6.527 s    (-1.716 s .. 19.59 s)
                           0.692 R²   (0.283 R² .. 1.000 R²)
      mean                 4.232 s    (2.136 s .. 5.686 s)
      std dev              2.186 s    (0.0 s .. 2.519 s)
      variance introduced by outliers: 74% (severely inflated)

    benchmarking count vs count'/count' 500
      time                 3.315 s    (-6.152 s .. 16.65 s)
                           0.357 R²   (0.239 R² .. 1.000 R²)
      mean                 5.474 s    (2.826 s .. 6.971 s)
      std dev              2.351 s    (0.0 s .. 2.592 s)
      variance introduced by outliers: 74% (severely inflated)

    The variance actually went up as # of counters increased

------------------------------------------------------------------------------}

main :: IO ()
main = defaultMain . return $ bgroup "flat vs nested morris" [
        bench "flat list e=d=0.05"   . nfIO $ morris 0.05 0.05 [1..10000]
      , bench "nested list e=d=0.05" . nfIO $ morris' 0.05 0.05 [1..10000]
  ]

_count cs  = runRVar (Cl.sourceList [1..1000] $$ count cs) StdRandom
_count' cs = runRVar (Cl.sourceList [1..1000] $$ count' cs) StdRandom

(t1,m1) = (10,3)   :: (Int,Int)
(t2,m2) = (100,3)  :: (Int,Int)
(t3,m3) = (8000,3)  :: (Int,Int)

cs' :: Num a => Int -> Int -> [[Counter]]
cs' t m  = replicate m $ replicate t 0

cs :: Num a => Int -> Int -> [Counter]
cs t m = replicate (t*m) 0

-- * conclude: morris 2x faster than morris', not surprisingly
--main :: IO ()
--main = defaultMain . return $ bgroup "morris vs morris'" [
--        bench "morris 10"  . nfIO $ morris  0.5 0.5 [1..10]
--      , bench "morris' 10" . nfIO $ morris' 0.5 0.5 [1..10]

--      , bench "morris 100"  . nfIO $ morris  0.5 0.5 [1..100]
--      , bench "morris' 100" . nfIO $ morris' 0.5 0.5 [1..100]

--      , bench "morris 1000"  . nfIO $ morris  0.5 0.5 [1..1000]
--      , bench "morris' 1000" . nfIO $ morris' 0.5 0.5 [1..1000]

--      , bench "morris 10000"  . nfIO $ morris  0.5 0.5 [1..10000]
--      , bench "morris' 10000" . nfIO $ morris' 0.5 0.5 [1..10000]
--  ]

{-----------------------------------------------------------------------------
  Benchmark  t of m vs m of t
------------------------------------------------------------------------------}

---- * conclusion : counters faster by 1 ms on 0.05 0.05
--main :: IO ()
--main = defaultMain . return $ bgroup "t of m vs m of t tests" [
--        bench "eps = 0.5 d = 0.5 counter " . nfIO $ incrsIO $ counters  0.5 0.5
--      , bench "eps = 0.5 d = 0.5 counter'" . nfIO $ incrsIO $ counters' 0.5 0.5

--      , bench "eps = 0.1 d = 0.1 counter" . nfIO $ incrsIO $ counters   0.1 0.1
--      , bench "eps = 0.1 d = 0.1 counter'" . nfIO $ incrsIO $ counters' 0.1 0.1

--      , bench "eps = 0.05 d = 0.05 counter" . nfIO $ incrsIO $ counters   0.05 0.05
--      , bench "eps = 0.05 d = 0.05 counter'" . nfIO $ incrsIO $ counters' 0.05 0.05
--  ]

--incrsIO :: MonadRandom m => [[Counter]] -> m [[Counter]]
--incrsIO  xxs = runRVar (incrs  xxs) StdRandom

--counters :: Eps -> Delta  -> [[Counter]]
--counters e d = replicate m $ replicate t 0
--  where (m,t) = (m' e d, t' e d)

--counters' :: Eps -> Delta -> [[Counter]]
--counters' e d = replicate t $ replicate m 0
--  where (m,t) = (m' e d, t' e d)

--t' e d = round $ 1/(e^2*d)
--m' e d = round . log $ 1/d

---- * given a list of list of counters toss a coin for each counter and incr
--incrs :: [[Counter]] -> RVar [[Counter]]
--incrs = sequence . fmap (sequence . fmap incr)

---- * Increment a counter `x` with probability 1/2^x
--incr :: Counter -> RVar Counter
--incr x = do
--  h <- toss . coin $ 0.5^(round x)
--  return $ if isHead h then (seq () succ x) else seq () x

{-----------------------------------------------------------------------------
    Benchmark time compleity of mean/median
       see : http://www.serpentine.com/criterion/tutorial.html
------------------------------------------------------------------------------}

-- * observation: 5% difference in runtime using mean   vs not
-- *              5% difference in runtime using median vs not
--main :: IO ()
--main = defaultMain . return $ bgroup "morris mean vs no median" [
--      bench "morris median    100 items eps = 0.1   delta = 0.1 " . nfIO $ morris  0.1 0.1   [1..100]
--    , bench "morris no median 100 items eps = 0.1   delta = 0.1 " . nfIO $ morris' 0.1 0.1   [1..100]
--    , bench "morris    median 100 items eps = 0.05  delta = 0.05" . nfIO $ morris  0.05 0.05 [1..100]
--    , bench "morris no median 100 items eps = 0.05  delta = 0.05" . nfIO $ morris' 0.05 0.05 [1..100]
--  ]

{-----------------------------------------------------------------------------
    II. Approximate Mean
------------------------------------------------------------------------------}

-- * Run Morris alpha on stream inputs `xs`
--morrisA :: [a] -> IO Counter
--morrisA xs = flip runRVar StdRandom $ Cl.sourceList xs $$ alpha

---- * Run Morris beta on stream inputs `xs` for `t` independent trials and average
--morrisB :: Int -> [a] -> IO Counter
--morrisB t =  fmap rmean . replicateM t . morrisA

---- * final morris algorithm
---- * Run on stream inputs `xs` for t independent trials for `t = 1/eps`,
---- * and `m` times in parralell, for `m = 1/(e^2 * d)`
---- * and take the median
---- * TODO: make this actually parralell
--morris :: Eps -> Delta -> [a] -> IO Counter
--morris e d = fmap rmedian . replicateM m . morrisB t
--  where (t,m) = (round $ 1/(e^2*d), round $ 1/d)

------------------------------------------------------------------------------
--    III. Utils
-------------------------------------------------------------------------------

---- * Utils * --

---- * A step in morris Algorithm alpha
--alpha :: Sink a RVar Counter
--alpha = (\x -> 2^(round x) - 1) <$> Cl.foldM (\x _ -> incr x) 0

---- * Increment a counter `x` with probability 1/2^x
--incr :: Counter -> RVar Counter
--incr x = do
--  h <- toss . coin $ 0.5^(round x)
--  return $ if isHead h then (seq () succ x) else seq () x

--rmean, rmedian :: (Floating a, Ord a, RealFrac a) => [a] -> Float
--rmean   = fromIntegral . round . mean
--rmedian = fromIntegral . round . median

{-----------------------------------------------------------------------------
    Depricated
------------------------------------------------------------------------------}

-- * Non fmap f version to test list traversal
-- * marginal difference due to fmap g where g = mean or median
--morrisB' :: Int -> [a] -> IO [Counter]
--morrisB' t = replicateM t . morrisA

--morris' :: Eps -> Delta -> [a] -> IO [Counter]
--morris' e d = replicateM m . morrisB t
  --where (t,m) = (round $ 1/(e^2*d), round $ 1/d)

{-----------------------------------------------------------------------------
  Benchmark  median of means
  naive "imperitive" solution worse in large n
------------------------------------------------------------------------------}

--main :: IO ()
--main = defaultMain . return $ bgroup "count vs count'" [
--      --  bench "median  1000"  $ whnf (medianOfMeans  10) [1..1000]
--      --, bench "median' 1000"  $ whnf (medianOfMeans' 10) [1..1000]

--      --, bench "median 10000"  $ whnf (medianOfMeans  10) [1..10000]
--      --, bench "median' 10000"  $ whnf (medianOfMeans' 10) [1..10000]

--      --, bench "median 100000"  $ whnf (medianOfMeans  10) [1..100000]
--      --, bench "median' 100000"  $ whnf (medianOfMeans' 10) [1..100000]

--        bench "median  0.1"  $ whnf (medianOfMeans  1000000) [1..5000000]
--      , bench "median' 0.1"  $ whnf (medianOfMeans' 1000000) [1..5000000]

--      , bench "median  0.5"  $ whnf (medianOfMeans  8000) [1..24000]
--      , bench "median' 0.5"  $ whnf (medianOfMeans' 8000) [1..24000]

--  ]

---- * this is worse!!
--medianOfMeans' :: Counter -> [Counter] -> Counter
--medianOfMeans' t xs = let (ms,_,_) = foldr (tomean t) ([],0,1) xs in median ms

--tomean :: Counter -> Counter -> ([Counter],Counter,Counter) -> ([Counter],Counter,Counter)
--tomean t x (ms,m,c) | c < t      = (ms,m+x,c+1)
--                    | otherwise  = ((fromIntegral . round $ (m+x)/t):ms,0,1)

--medianOfMeans :: Int -> [Counter] -> Counter
--medianOfMeans t = median . fmap mean' . (chunksOf t)
-- fmap median' .

--ms = medianOfMeans 10 [1..100]
--(ms',_,_) = medianOfMeans' 10 [1..100]


{-# LANGUAGE BangPatterns #-}
--{-# XRankNTypes #-}

module RprGD2 (OptimisationType (Rapid , FromData ,WithDepth)
    , mmlRprGdEvl,mml2DsGD -- ,mmlRprTest
    , mml2DsGen, RprMix2.aBar,mapP --  , stdTest
    --,mmlRprTest
    ) where

import Numeric.LinearAlgebra
import Numeric.GSL.Statistics (stddev,lag1auto)
import qualified Data.ByteString.Lazy.Char8 as L (writeFile, pack, unpack, intercalate,appendFile, append,concat)
import System.Random
import Control.Concurrent
import Control.DeepSeq
import Foreign.Storable
import Control.Parallel (par, pseq)
import qualified Control.Parallel.Strategies as St-- (parMap)
import System.IO.Unsafe
import Data.Maybe (isJust,fromMaybe,maybe,listToMaybe,fromJust)
import Data.Time
import Data.Char (isDigit)
--import Text.Printf
import Data.Ord (comparing)
import Control.Monad (liftM,liftM2)
import Data.List (transpose, foldr1, sort, foldl' ,nub, nubBy,sortBy,  minimumBy, partition)
import qualified Data.Map  as M (Map,empty,insertWith' ,mapWithKey,filterWithKey,toList, fromList)
--import IO
--------------------------------------------------------------------------------------------
import HMDIFPrelude (bsReadChr,bs2Int')
import RprMix2 -- (vLength)
import LINreg2
import ResultTypes
import ProcessFile (writeToFile,log2File,writeOrAppend,myTime)
import LMisc hiding (force)
import ListStats (nCr,normal,binomial ,myTake,-- corr,
                  mean, adjustedMatrix,cdf, meanBy, stDivBy, -- stDiv,
                  adjustedMatrixN, adjustedMatrixPN,stNorm,mean,stDivM)
---------------------------------------------------------------------------------------------
----                 optimisation type
data OptimisationType = Rapid | FromData | WithDepth Int deriving (Show, Eq)-- | WithNum Int
---------------------------------------------------------------------------------------------
---- Paralell definitions
concatPMap ::  (a -> [b]) -> [a] -> [b]
concatPMap f = concat . St.parMap (St.parList St.rseq) f

-- parallel map and parallel concatMap
--mapP  = St.parMap (St.dot St.rseq St.rpar)

mapP' :: (a -> b) -> [a] -> [b]
mapP' f = foldl' (\bs a -> let fa = (f a) `St.using`  St.rseq in
                                fa : bs ) []
--
mapP :: (a -> b) -> [a] -> [b]
mapP  f = St.parMap St.rseq f

concatPMap' f   = foldl'Rnf  (\xs x -> xs ++ [St.using (f x) (St.dot St.rseq St.rpar) ] ) []

--- strict fold vector
-- foldVector :: Storable a => (a -> b -> b) -> b -> Vector a -> b
foldVector' :: (NFData b, Storable a) => (b -> a -> b) -> b -> Vector a -> b
foldVector' f a vs = foldVector (\b g x -> g (seqit $ f x b)) id vs a
    where seqit a = rnf a `pseq` a
--------------------------------------------------------------------------------------------------
myMinimum :: (NFData a, Ord a) => [a] -> a
myMinimum   (x:xs) = foldl'Rnf  min x xs
myMinimum     _    = error "error: myMinimum empty list"
--
--mmlRprGD :: Vector Double -> Double -> Double -> Matrix Double -> Mmodel -> (Mmodel, Double)
{-# NOINLINE mmlRprGD #-}
mmlRprGD g ys n cts nSLP dss (ml, model) = mmlRprGD_aux g ys n cts nSLP dss  model

mmlRprGD_aux :: (RandomGen g) => g
     -> Vector Double
     -> Double
     -> [Int]
     -> Maybe Bool
     -> Matrix Double
     -> Mmodel
     -> IO OModel
mmlRprGD_aux g ys n cts nSLP dss model =  do
    return $ Omod (nAbaMod model taU) oldMsgLn
        where
            ------- compute values from the input model, eg the weighs-----------
            nAbaMod  (MkModel a b c x _ )  k  =  (MkModel a b c x k)
            ----
            fL             = (+ 1) . fromIntegral . length
            aA'            = fL cts
            -------------------------------------------------------
            siG            =  sigD  model
            siGE           =  sigE  model
            taU            =  tau   model
            ---------------------------------------------------------
            xXa            =  invS nSLP (dss <> (aBar model)) ---
            abar           = aBar model

            aA             =  fromIntegral . (\ a -> if a == 2 then 1 else a `div` 2)  $ dim  abar
            errFun  a b    =  (r_n model a b)
            rN             =   zipVectorWith  errFun xXa ys -- (r_n model)
            lnR            =   dim rN   -- get the magnitude of rN
            wData :: Matrix Double
            wData          =  diagRect 0 rN lnR lnR
            wError :: Matrix Double
            wError         =  diagRect 0 (mapVector (1 -) rN) lnR lnR
            ---
            xtWx            =  ((trans dss) <> wData) <> dss
            ------------------------------------------------------------
            ---
            nAlpha         =  foldVector (+) 0 rN
            nAlpA1         =  nAlpha - aA - 1 --
            --
            ysSubXa        =  zipVectorWith (-) ys  xXa -- the data minus the predictions
            ySubWySub ww   =  ysSubXa <.> (ww <> ysSubXa) -- y
            --
            dehAbarSq      = (abar <.> abar)
            aBarSqTau1     =   dehAbarSq / (2 * taU^2)
            ----------------------------------------------------------------------------------------
            ------------------- calculate the message length of the input model --------------------
            ----------------------------------------------------------------------------------------
            !oldMsgLn     =   l1 `par` l2 `par` l3 `par` l4 `pseq`   (l1 + l2 + l3 + l4)
            ----------------------------------log 0.01 *------------------------------------------------------
            restrPr       = ((aA + 1)* log 2)/ 2
            restrPrior     = maybe restrPr (\_ -> restrPr) nSLP
            l1             =   (aA + 2.0) * logg taU + aBarSqTau1 - restrPrior  + 0.5 * ((n-2) * log (2 * pi))
            ----------------------------------- log 2 * (aA + 1)/2 ---------------------------------
            l2              =   nAlpA1 * logg siG + (1/(2 * siG^2)) * (ySubWySub wData)
            l3              =  (n - nAlpha) * logg siGE + (1/(2 * siGE^2)) * (ySubWySub wError)  + 0.5 * logg (2 * nAlpha)
            mlLogs          = 0.5 * log (det xtWx) + 0.5 * logg (2*(n - nAlpha)) +  0.5 * log (pi * (aA + 3))
            l4              =   mlLogs - 1.319
----------------------------------------------------------------------------------------------------
dummyMod       = Omod (MkModel 0 0 0 (fromList [1]) 0) 100000000
vunStdNorm' vs = vunStdNorm (vmean vs) (vstDiv vs) vs
{----------------------------------------------------------------------------------------------------
mmlRprGdEvl: the mmlRPR  evaluation function. Takes:
n  - the range in years of roughness data (which corresponds to the number of data points
      when dealing with simulation data)
ms - the number of chains that are joined together - the maintenance strategy.
     Note that the mml2DS uses the mmlRPR on adjaent sections joined together. Hence this
     parameter says how many sections are to be joined for calls to mmlRprGdEvl. When this
     functions is called outside of the mml2DS, this valus is always 1 (i.e we do not make any
     prior assumptions about maintenances on the chain)
ys  - a lsit of the data values for this chain
pp  - the interval in years between maintenenace interventions
gen - a random generator to initialize the process
nSlopes - determines whether to exclude negative slopes: if true, negative slopes are excluded
           otherwise they are not
rprOnly - toggles the applicaion of the MMLRPR function only. The likelihood of maintenance
          is not applied to each interventions
----------------------------------------------------------------------------------------------------}
mmlRprGdEvl :: (RandomGen g) => Int ->
    Int -> -- error search depth limit
    [Vector Double]  -> -- Either (Vector Double) [Double] -> -- the data in standardized for
    Double ->
    g ->
    Maybe Bool -> -- negative slopes (temporarily used for turning on and off the mixture model)
    Int ->
    IO (Maybe ( (OModel,(Vector Double,[Int]) ) , [Double])) --ms
mmlRprGdEvl n dpt ys pp gen nSlopes mtd = -- mN std
    mmlRprAuxGD n dpt m p (filter ((== mLen) . dim) ys) pp gen nSlopes mtd True -- ms -- mN std
        where
            p     =  1 / pp
            --- don need this claculation. we can filter out irelevant values
            m     =  n `div` mtd
            mLen  = maybe 0 (\_ -> maximum $ map dim ys) (listToMaybe ys)
            ---------------------------------------------------------------------------------------------------
            unStd :: (OModel,(Vector Double,[Int])) -> (OModel,(Vector Double,[Int]))
            unStd  (Omod (MkModel c b e ks d)  ml,(vs,ys)) =  (Omod (MkModel c b e (vunStdNorm' ks) d)  ml,(vunStdNorm' vs,ys))
            unSd ms =
                case ms of
                    Nothing -> Nothing
                    Just (a,b) -> Just (unStd a, b)

------------
mmlRprAuxGD n _ m p ysv pp gen nSlopes mtd toOpt = do
    return minMod
    where
        minMod = liftM applyMixture $
                    maybe Nothing (\_ -> Just $ minimumBy (comparing (fst . fst))  rprs1) (listToMaybe rprs1)
        --
        result  rprss =  (snd (minimumBy (comparing (fst . snd)) rprss), map fst rprss)
        f         =  fromIntegral
        (g1,_)    =  split gen
        abr' ::  Matrix Double -> [(Double, (Vector Double, (Double, Vector Double))) ]
        abr'      = (: []) . minimumBy (comparing fst)
                    . zipWith (\f a -> f a) (map (\ys -> mmlLinearRegLP ys 1 1) setLen)
                    . replicate len -- (length ysv)
            where
                len    = length ysv
                setLen = replicate len $ join ysv
        ----------------------------------------------------------------------------
        initM b s =  initModel gen s b -
        initLz (ml, (prd, (s,ab))) =  (ml, (prd , initM ab s))
        --
        mod :: Matrix Double -> [(Double, (Vector Double, Mmodel))]
        mod  mxx  =  map   initLz $ abr'  mxx
        ----------------------------------------------------------------------------
        nN = f . sum $ map dim ysv -- n
        nNs = map ((\n -> [1 .. n]) . f . dim) ysv
        fLn = (+ 1) . f . length
        msgLm (Omod _ l) = l
        appMML g1 nM ct  mmX (ml,(prd , inMd))  = unsafePerformIO $ mmlRprGD g1 (join  ysv) nM ct nSlopes mmX (ml, inMd)
        iTr (mX, cts) = (map (appMML gen nN cts  mX) (mod mX), (mX, cts))
        cMML  = mapP (getPredictions nSlopes) . mapSnd . iTr
        nps cmm  =   realignedNpieces cmm   mtd
        --------------------------------------------------------------------------------------------
        ----------- applying the mixture model to the discovered intervention, only ----------------
        rprs1  =   [ (kx, cm1) | cm1 <- [0 .. m],  kx <- ( concatPMap (mapSnd . iTr) . (nps cm1)) nNs  ]  `St.using` (St.parList St.rseq)
        applyMixture ( (Omod md mln , (mx , cts)), cm) =
            let   (!nMd , !mxL)  = unsafePerformIO $ nextModelR' gen cts nSlopes mx (join ysv) md
                  ncma    = n
                  nncma   =  ncma `nCr` cm
                  ncr     =   (binomial n cm p) / f nncma -- else 1
                  lncr    =   log ncr
                  newOMod = Omod nMd ((lncr + mln) + mxL)
                  prds    = getPredictions nSlopes (newOMod, (mx, cts))
            in    (prds , [lncr]  )
        --------------   end applying the Mixture Model at the end ----------------------------------
        getPredictions nSlp (mm, (xConfig, cts))  =   (mm , (invS nSlp  prds , cts))
            where
               prds     =   xConfig <> (aBar (oMod mm))
        --- retain the domain of the function
        retDom :: (a -> [b]) -> a -> [(a , b)]
        retDom  f  a  =   map (\x -> (a, x)) ( f a)
        ---
        nAbaMod  (MkModel a b c _ x)  k  =  (MkModel a b c k x)

{----------------------------------------------------------------------------------------------------
  The mml2DsGD functions applies the MMLRPR to alogorithm to a list of data
  points (i.e. the readings or chainages) and after
  sectioning them off mml2DsGD years distance gen [[Double]]
--}
--- returning results
mml2DsGen :: (RandomGen g) => g -> TdsParms ->
    OptimisationType -> -- the optimisation type (Rapid , FromData ,WithDepth)
    SecOrSchm ->
    IO SecOrSchm
mml2DsGen g tprms optType ssm =
    case ssm of
        -- analysis results for sections
        Left  sect -> do
            let (a, b , c ,d) = (sCondition sect, sScannerCode sect, sScannerSorrg sect, sCWay sect)
            let mkSect = Section a b c d (sChains sect)
            liftM (Left . mkSect . omd2Results) $ mml2DsGD g tprms optType [xs | xs <- sChains sect, length (cData xs) > 1]
        -- results for schemes
        Right schm -> do
            nScmRecs <- mapM (mkSchemeRecRes g tprms optType) (schmRecs schm)
            return $ Right $ Scheme (schemeName schm) (schmCondition schm) nScmRecs
    where
        mkSchemeRecRes :: (RandomGen g) => g ->
            TdsParms ->
            OptimisationType ->
            SchemeRec ->
            IO SchemeRec
        mkSchemeRecRes g tp opT src = liftM (mkSchmRec . omd2Results) $ mml2DsGD g tp opT chn
            where
                mkSchmRec = SchemeRec scd srd srnm srg stc sec chn
                --
                scd = scmScannerCode src
                srd = roadClassCode src
                srnm = roadName  src
                srg  = scmScannerSorrg  src
                --- carriage way
                stc = startChain   src
                sec = endChain  src
                chn = scmChains src

--}
-- returning a OMod, etc
mml2DsGD :: (RandomGen g) => g ->
    TdsParms -> -- the parameters for the mmlTds
    OptimisationType -> -- the optimisation type (Rapid , FromData ,WithDepth)
    [Chain] -> --- a list of all the data for each chain in a section or scheme
    IO ([(OModel, ([Double], [Int]))], [Int])
mml2DsGD  gen tdPrs optType chxss = do
    ---------------------------------------------
    scrollLog ("Identifying candidate groups for section with length: " ++ show len) -- (max 3 (len `div` 15))
    !mks1  <-   (joinPairsV dist optType 2 len 25 xss)
    ------
    let total = length mks1
    let msgP1 = "Fitting progression rate and identifying outliers for section length: " ++ show len ++". "
    let msgP2 = show total ++": candidate groups found."
    onlycuts <- liftM (maybe 0 (fromJust . bs2Int') . listToMaybe . head) $ bsReadChr  ',' "./mmlTdsFiles/onlycuts.txt"
    -- print ("negslopes is : " ++ show nSlopes)
    if len > 0 then
        if null mks1 then  do
            let noGrupMsg = "NO Groups for section length: " ++ show len
            logLog ("-------------------"++noGrupMsg)
            return  dummy
        else   if onlycuts == 1 then do
            logLog ("only first cuts printed")
            return dummy
        else do
            --
            scrollLog  (msgP1 ++ msgP2)
            let mks2  = [(a,b,c) | ((a,b),c) <- zip (zip mks1 (rGens gen)) [1 ..]]
            let applyRpr = mapP (mml2DsGD_aux1 total minTrend scrollLog)
            let appMin  = liftM (minimumBy compML'')  . sequence .  applyRpr  -- ($!)
            appMin mks2
    else do
        logLog ("Empty section list sent to MML2DS: " ++ show len ++". limit: "++ show lenLim)
        return dummy
    where
        ---- projections --
        xss                  =  map cData chxss
        (!dpt, !pp , !dist)  =  (searchDept tdPrs, rngInYrs tdPrs, rngInDist tdPrs)
        (nSlopes,lim,minAnl) =  (negSlopes tdPrs, maxChn2Aln tdPrs, minChn2Aln tdPrs)
        minTrend             =  minPts2Trd tdPrs
        loggers              =  logFun tdPrs
        scrollLog            =  maybe (\_ -> return ())  (\_ -> loggers !! 0) (listToMaybe loggers)
        logLog               =  maybe (\_ -> return ())  (\_ -> loggers !! 1) (listToMaybe loggers)
        maxChain             =  (foldl' (\n xs -> max (length xs) n) 0 xss)
        lenLim               =  lim * maxChain
        len                  =  length xss
        minLen               =  minAnl * maxChain
        -------------------------------- end of projections ------------------
        dummy                =  ([(dummyMod , ([0] ,[0]))] ,[])
        pm                   =  1 / dist
        compML''             =  comparing  (sumBy (oModMl . fst). fst)
        f                    =  fromIntegral
        {--------------------------- logging the output from the meddage lengths --------------------
        printCutsInfo :: OptimisationType -> ([(OModel, ([Double], [Int]))], [Int]  ) -> IO ()
        printCutsInfo opTp (ks,cs) = do
            let cutsMsg = "\n cuts: "++ printElm cs
            let mlgMsg  = "; \t message length: " ++ (show $ sumBy (oModMl . fst) ks)
            let name = case opTp of
                            Rapid         ->  "fromData" -- "priorOnly"
                            FromData      ->  "fromData"
                            WithDepth _   ->   "fixedDepth"
            let file = "./SimulationFiles/"++name++"groups.txt"
            -- let header = "\n-----------\t cuts and message lengths ------------\n"
            let output = L.pack (cutsMsg ++ mlgMsg)
            writeOrAppend file output
        ---------------------------------------------------------------------------------------}
        mml2DsGD_aux1 :: RandomGen g =>  Int ->
                           Int ->
                         (String -> IO()) ->
                         ( ([([Vector Double], Int)],[Int]), g, Int)  ->
                         IO ([(OModel, ([Double], [Int]))], [Int])
        mml2DsGD_aux1 total minToTrend write ((pgroups,pgcuts),gen,num)  =  do
            write ("Calculating mmlrpr on group: "++ show num ++ " of " ++ show total)
            let max1  = sumBy snd pgroups -- sum pgcut
            ----
            let  calc = [ (result,pgcuts) |
                            let !rprsCalcs =  unsafePerformIO . sequence $ mapP (calcRpR max1 gen minToTrend) pgroups -- do the mlrpr on each grou
                            , all isJust rprsCalcs -- .  filter isJust
                            , let result   = (map fromJust . filter isJust) rprsCalcs]
            if null calc then do
                write ("Invalid mmlrpr group for group "++ show num)
                return dummy

            else do
                retMinComp calc
                where
                    --
                    retMinComp =  return . minimumBy  compML''
                    --
                    calcRpR :: RandomGen g => Int -> g -> Int ->
                               ([Vector Double], Int) -> IO (Maybe (OModel,([Double],[Int]))) -- , [Double]) -
                    calcRpR ns g mtd (vys, ms) = do
                        !rprs <- mmlRprGdEvl ms dpt vys pp g nSlopes mtd
                        case rprs of
                            Just ( (Omod r1 r2, prs),nn) -> do
                                       let smm = ncr' + r2
                                       let str = ("\nncr' is: "++ show ncr' ++" mLen is: "++ show r2 ++". Their sum is: " ++ show smm)
                                       -- L.appendFile "tempLog1.txt" (L.pack str)
                                       return $ Just (Omod r1 smm , (toList (fst prs), snd prs) ) --
                            Nothing ->   return $ Nothing
                        where
                            ncr'  =   (binomial ns ms pm) / f (ns `nCr` ms)
-----------------------------------------------------------------------------------------------------}
------------
--------------------------------------------------------------------------------------------------
findGrpsFromData :: Double -> -- the distance of intervention points
    Int              ->  -- lenght of the input list
    [[Double]] ->           -- ^ the data to aling
    IO [([([Vector Double], Int)], [Int])]
findGrpsFromData dist inLen   =
    return . findCandidateSet .  uncurry calculateFixedGroups1 . findInterventions_aux
    where
        findCandidateSet xs
            | length xs > 100 = maintainBy' avgRelCorr (Just 45) xs
            | otherwise       = (take 10 .  sortBy (comparing avgRelCorr )) xs
        --
        avgRelCorr    =  (meanBy (relativeCorr . fst)) . fst
        -- find possible interventions from the data
        findInterventions_aux ::  [[Double]] ->    ([[Double]] , [[Int]])
        findInterventions_aux  xys =   (xys ,  (nub . findGrps1 0.7) xys)
            where
                findGrps1 tolerance xs =  [ ys | let stDiv = stddev . fromList
                                                , ys <-  findGrps tolerance 38 mmlCorr xs
                                          ] `St.using` (St.parList St.rseq)
        -- calculate fixed groups
        calculateFixedGroups1 :: [[Double]] ->  [[Int]]  -> [([([Vector Double], Int)], [Int])]
        calculateFixedGroups1  xys  =  map (unzip . foldr mkLst [] . filter (not . null . fst) .  mkIntervals vs )
            where
                vs              = (map fromList xys)
                mkLst (ks,k) ps =  ((ks, sumBy dim ks) , k) :  ps
        -- relative correlaton between adjacent chains
        relativeCorr :: [Vector Double] -> Double
        relativeCorr xs
            | length xs < 2 =  4 -- bias against groups with only one chian
            | otherwise     =  mcD xs
            where
                corrDist ps qs =  abs (nCorrV ps qs - 1)
                mcD            =  meanBy (exp . uncurry corrDist) . zipp

----------------------------------------------------------------------------------------------------
--- an infinite list of random generators
rGens :: RandomGen g =>  g -> [g]
rGens g = g : (g1 : rGens g2)
    where (g1, g2) = split g

------------------------------------------------------------------------------------------------------------------------
joinPairsV :: Double -> -- the distance of maintenence interventions
            OptimisationType -> -- the optimisation type (Rapid , FromData ,WithDepth)
            Int ->              -- ^ the minimum number of chains to aligh
            --Int ->            -- ^ maximum
            Int ->              -- ^ length of the input list
            Int ->              -- ^ the maximum number of chains to align
            [[Double]] ->  -- ^ the data to aling
            IO [([([Vector Double], Int)], [Int])]
            --[([(Vector Double, Int, Double )],[Int])]
            -- (joinPairsV pm optType minLen lenLim 3 len lim) xss
joinPairsV dist optType malg xysLen lim xys -- minL maxL
            | null  xys           =  return []
            | xysLen <  4         =  return joinFew -- ' pm --xys
            --- check the length of the combined list
            | lenData < 10        =  return joinFew
            | otherwise           =
                case optType of
                    Rapid         ->  findGrpsFromData dist xysLen xys
                    FromData      ->  findGrpsFromData dist xysLen xys
                    WithDepth d   ->  calculateFixedGroups vxys [(repeat d)]
            where
                lenData           = sumBy length xys
                --
                vxys              =  map fromList xys
                ------------------
                vs                =  join vxys
                ms                =  dim vs
                joinFew           =  [([([vs],dim vs)],[])]

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

-- given a list of possible cuts and the list of chains, we generate the a list
-- of possible maintenance groups defeind by the cuts
calculateFixedGroups :: [Vector Double] ->  [[Int]]  -> IO [([([Vector Double], Int)], [Int])]
calculateFixedGroups  vs  = return . map (unzip . foldr mkLst [] . filter (not . null . fst) . mkIntervals vs)
    where
       mkLst (ks,k) ps =  ((ks, sumBy dim ks) , k) :  ps
       -- return [ (mkIntervals vs cs , cs)  | cs <- ks]


-- | A tool for measuring cycle time of tasks
--
-- Cycle time is the difference between creation date and completion date.
--
module Main (main) where

import RIO

import Asana.Api
import Asana.Api.Gid (Gid, gidToText)
import Asana.Api.Project (Project(pCreatedAt, pGid, pName), getProjects)
import Asana.App
  (AppM, appExt, loadAppWith, parseBugProjectId, parseYear, runApp)
import qualified Data.Csv as Csv
import Data.Foldable (maximum, minimum)
import Data.List (intercalate, nub)
import qualified Data.Map.Strict as Map
import qualified Data.Vector as V
import RIO.ByteString (writeFile)
import RIO.ByteString.Lazy (toStrict)
import RIO.Text (isPrefixOf, unpack)
import RIO.Time (NominalDiffTime, diffUTCTime, toGregorian, utctDay)
import Statistics.Quantile (median, medianUnbiased)
import Statistics.Sample (kurtosis, mean, skewness, stdDev)
import Statistics.Sample.Histogram (histogram)
import System.IO.Temp (emptySystemTempFile)
import Text.Printf (printf)

data AppExt = AppExt
  { appYear :: Integer
  , appBugProject :: Gid
  }

main :: IO ()
main = do
  app <- loadAppWith $ AppExt <$> parseYear <*> parseBugProjectId
  runApp app $ do
    logDebug "Fetch projects"
    year <- asks $ appYear . appExt
    projects <- filter (approvedProject year) <$> getProjects

    taskGids <- fetchRelevantTaskGids projects

    tasks <- pooledForConcurrentlyN maxRequests taskGids getTask
    logDebug "Write task CSV"
    writeTaskCsv tasks

    let
      cycleTimes = zscoreFilter . V.fromList . fmap realToFrac $ mapMaybe
        mayCycleTime
        tasks

    when (null cycleTimes) $ logWarn "No cycle time to compute"

    logDebug "Write histogram CSV"
    writeHistogramCsv cycleTimes

    logDebug "Display Cycle Time"
    let
      fmt = intercalate
        "\n"
        [ "Cycle Time"
        , "- mean:     %d days"
        , "- mediam:   %d days"
        , "- skewness: %.3f"
        , "- kurtosis: %.3f"
        , "- stdDev:   %d days"
        , "- max:      %d days"
        , "- min:      %d days"
        ]
    logInfo . fromString $ printf
      fmt
      (toDays $ mean cycleTimes)
      (toDays $ median medianUnbiased cycleTimes)
      (skewness cycleTimes)
      (kurtosis cycleTimes)
      (toDays $ stdDev cycleTimes)
      (toDays $ maximum cycleTimes)
      (toDays $ minimum cycleTimes)

zscoreFilter :: Vector Double -> Vector Double
zscoreFilter p =
  let
    av = mean p
    sdev = stdDev p
    zs x = (x - av) / sdev
  in V.filter ((<= 3) . zs) p

toDays :: Double -> Int
toDays = floor . (/ 86400)

approvedProject :: Integer -> Project -> Bool
approvedProject compareYear project = year == compareYear && isPrefixOf
  "Iteration"
  (pName project)
  where (year, _, _) = toGregorian $ utctDay (pCreatedAt project)

mayCycleTime :: Task -> Maybe NominalDiffTime
mayCycleTime task = do
  completedAt <- tCompletedAt task
  -- FIXME: Sometimes this is true
  guard $ completedAt > tCreatedAt task
  Just . diffUTCTime completedAt $ tCreatedAt task

writeHistogramCsv :: Vector Double -> AppM ext ()
writeHistogramCsv cycleTimes = do
  filePath <- liftIO $ emptySystemTempFile ".csv"
  let histogramCsv = uncurry V.zip $ histogram @_ @_ @Double 100 cycleTimes
  writeFile filePath . toStrict . Csv.encode $ V.toList histogramCsv
  logInfo . fromString $ "Histogram written to " <> filePath

writeTaskCsv :: [Task] -> AppM ext ()
writeTaskCsv tasks = do
  filePath <- liftIO $ emptySystemTempFile ".csv"
  let
    toRecord task = Map.fromList
      [ ("name" :: String, unpack $ tName task)
      , ("task ID" :: String, unpack $ taskUrl task)
      , ( "cycle time"
        , maybe "" (show . realToFrac @_ @Double) $ mayCycleTime task
        )
      ]
  writeFile filePath
    . toStrict
    . Csv.encodeByName (V.fromList ["name", "task ID", "cycle time"])
    $ toRecord
    <$> tasks
  logInfo . fromString $ "Tasks written to " <> filePath

fetchRelevantTaskGids :: Traversable t => t Project -> AppM AppExt [Gid]
fetchRelevantTaskGids projects = do
  taskGids <-
    fmap (nub . concat)
    . pooledForConcurrentlyN maxRequests projects
    $ \project -> do
        logDebug . fromString $ "Project tasks: " <> show
          (gidToText $ pGid project)
        fmap nGid <$> getProjectTasks (pGid project) AllTasks
  bugProjectGid <- asks $ appBugProject . appExt
  bugTaskGids <- fmap nGid <$> getProjectTasks bugProjectGid AllTasks

  pure $ filter (`notElem` bugTaskGids) taskGids


{-# LANGUAGE RecordWildCards, FlexibleContexts #-}

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

import Control.Error.Util
import Pipes as P
import qualified 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 Network.Socket (SockAddr(SockAddrInet), PortNumber, HostAddress)

import SDR.FFT
import SDR.Plot
import SDR.ArgUtils
import SDR.Util
import SDR.NetworkStream
import Graphics.DynamicGraph.Waterfall
import Graphics.DynamicGraph.Util

data Options = Options {
    fftSize      :: Maybe Int,
    windowWidth  :: Maybe Int,
    windowHeight :: Maybe Int,
    rows         :: Maybe Int,
    colorMap     :: Maybe [GLfloat],
    ip           :: Maybe HostAddress,
    port         :: Maybe PortNumber
}

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
          <$> optional (option (fmap fromIntegral parseSize) (
                 long "size"
              <> short 's'
              <> metavar "SIZE"
              <> help "FFT bin size. Default is 512."
              ))
          <*> 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'."
              ))
          <*> optional (option auto (
                 long "ip"
              <> short 'i'
              <> metavar "IP"
              <> help "IP address to bind to. Default is 0.0.0.0"
              ))
          <*> optional (option auto (
                 long "port"
              <> short 'p'
              <> metavar "PORT"
              <> help "UDP port to bind to. Default is 0x1234"
              ))

opt :: ParserInfo Options
opt = info (helper <*> optParser) (fullDesc <> progDesc "Draw a waterall plot of the samples received over Ethernet" <> header "Ethernet Waterfall")

{-# INLINE interleavedIQSigned256ToFloat #-}
interleavedIQSigned256ToFloat :: (Num a, Integral a, Num b, Fractional b, VG.Vector v1 a, VG.Vector v2 (Complex b)) => v1 a -> v2 (Complex b)
interleavedIQSigned256ToFloat input = VG.generate (VG.length input `quot` 2) convert
    where
    {-# INLINE convert #-}
    convert idx  = convert' (input `VG.unsafeIndex` (2 * idx)) :+ convert' (input `VG.unsafeIndex` (2 * idx + 1))
    {-# INLINE convert' #-}
    convert' val = fromIntegral val / 256

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

    let fftSize' =  fromMaybe 512 fftSize
        window   =  hanning fftSize' :: VS.Vector Double
    rfFFT        <- lift $ fftw fftSize'
    rfSpectrum   <- plotWaterfall (fromMaybe 1024 windowWidth) (fromMaybe 480 windowHeight) fftSize' (fromMaybe 1000 rows) (fromMaybe jet_mod colorMap)

    dev          <- lift $ udpRecvSocket $ SockAddrInet (fromMaybe 0x1234 port) (fromMaybe 0 ip)

    lift $ runEffect $   udpSource dev (fftSize' * 2)
                     >-> P.map (interleavedIQSigned256ToFloat :: VS.Vector CChar -> 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