{-# LANGUAGE AllowAmbiguousTypes #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Main (main) where

import Injection
import Projection

import Data.Complex (Complex ((:+)))
import Data.Dynamic (Dynamic)
import Data.Fixed (Fixed, E6)
import Data.Functor.Const (Const)
import Data.Functor.Identity (Identity)
import Data.List.NonEmpty (NonEmpty ((:|)))
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Monoid (Dual)
import Data.Monoid (Product)
import Data.Monoid (Sum)
import Data.Monoid (All, Any)
import qualified Data.Monoid as Monoid (First, Last)
import Data.Ord (Down (..))
import Data.Ratio (Ratio, (%))
import Data.Semigroup (Max, Min)
import qualified Data.Semigroup as Semigroup (First, Last)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text.Lazy as Lazy (Text)
import Numeric.Natural (Natural)
import Test.Hspec
import Test.QuickCheck hiding (Fixed)
import Test.QuickCheck.Instances ()

main :: IO ()
main = hspec $ do
    describe "instance Injection Integer Integer" $ do
        it "is resolvable" (resolveInjection @Integer @Integer)
    describe "instance Injection String String" $ do
        it "is resolvable" (resolveInjection @String @String)
    describe "instance Injection Text Text" $ do
        it "is resolvable" (resolveInjection @Text @Text)
    describe "instance Injection Lazy.Text Lazy.Text" $ do
        it "is resolvable" (resolveInjection @Lazy.Text @Lazy.Text)
    describe "instance Retraction Integer Dynamic" $ do
        it "is resolvable" (resolveRetraction @Integer @Dynamic)
        it "is the left inverse of inject" (lawLeftInverse @Integer @Dynamic)
    describe "instance Injection (Maybe Integer) [Integer]" $ do
        it "is resolvable" (resolveInjection @(Maybe Integer) @[Integer])
        it "outputs a value with the same length as the input" $ do
            let test = inject @(Maybe Integer) @[Integer]
            length (test Nothing) `shouldBe` 0
            length (test (Just 1)) `shouldBe` 1
        it "is injective" (lawInjective @(Maybe Integer) @[Integer])
    describe "instance Retraction (Maybe Integer) [Integer]" $ do
        it "is resolvable" (resolveRetraction @(Maybe Integer) @[Integer])
        it "is the left inverse of inject" (lawLeftInverse @(Maybe Integer) @[Integer])
    describe "instance Injection Natural Integer" $ do
        it "is resolvable" (resolveInjection @Natural @Integer)
        it "injects the value itself" $ do
            let test = inject @Natural @Integer
            test 0 `shouldBe` 0
            test 1 `shouldBe` 1
            test 2 `shouldBe` 2
        it "is injective" (lawInjective @Natural @Integer)
    describe "instance Retraction Natural Integer" $ do
        it "is resolvable" (resolveRetraction @Natural @Integer)
        it "is the left inverse of inject" (lawLeftInverse @Natural @Integer)
        it "is defined over non-negative integers" $ do
            retract @Natural @Integer 0 `shouldBe` Just 0
            retract @Natural @Integer 1 `shouldBe` Just 1
        it "is not defined over negative integers" $ do
            retract @Natural @Integer (-1) `shouldBe` Nothing
    describe "instance Injection Text String" $ do
        it "is resolvable" (resolveInjection @Text @String)
        it "is injective" (lawInjective @Text @String)
    describe "instance Injection Lazy.Text String" $ do
        it "is resolvable" (resolveInjection @Lazy.Text @String)
        it "is injective" (lawInjective @Lazy.Text @String)
    describe "instance Injection Text Lazy.Text" $ do
        it "is resolvable" (resolveInjection @Text @Lazy.Text)
        it "is injective" (lawInjective @Text @Lazy.Text)
    describe "instance Injection Lazy.Text Text" $ do
        it "is resolvable" (resolveInjection @Lazy.Text @Text)
        it "is injective" (lawInjective @Lazy.Text @Text)
    describe "instance Injection Integer (Fixed a)" $ do
        it "is resolvable" (resolveInjection @Integer @(Fixed E6))
        it "is injective" (lawInjective @Integer @(Fixed E6))
    describe "instance Retraction Integer (Fixed a)" $ do
        it "is resolvable" (resolveRetraction @Integer @(Fixed E6))
        it "is the left inverse of inject" (lawLeftInverse @Integer @(Fixed E6))
        it "is defined over integers" $ do
            retract @Integer @(Fixed E6) 0 `shouldBe` Just 0
            retract @Integer @(Fixed E6) 1 `shouldBe` Just 1
            retract @Integer @(Fixed E6) (-1) `shouldBe` Just (-1)
        it "is not defined over fractions" $ do
            retract @Integer @(Fixed E6) (1 / 2) `shouldBe` Nothing
    describe "instance Injection Integer (Const Integer String)" $ do
        it "is resolvable" (resolveInjection @Integer @(Const Integer String))
        it "is injective" (lawInjective @Integer @(Const Integer String))
    describe "instance Injection (Const Integer String) Integer" $ do
        it "is resolvable" (resolveInjection @(Const Integer String) @Integer)
        it "is injective" (lawInjective @(Const Integer String) @Integer)
    describe "instance Injection Integer (Ratio Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Ratio Integer))
        it "is injective" (lawInjective @Integer @(Ratio Integer))
    describe "instance Retraction Integer (Ratio Integer)" $ do
        it "is resolvable" (resolveRetraction @Integer @(Ratio Integer))
        it "is the left inverse of inject" (lawLeftInverse @Integer @(Ratio Integer))
        it "is defined over integers" $ do
            retract @Integer @(Ratio Integer) (0 % 1) `shouldBe` Just 0
            retract @Integer @(Ratio Integer) (1 % 1) `shouldBe` Just 1
            retract @Integer @(Ratio Integer) (1 % (-1)) `shouldBe` Just (-1)
        it "is not defined over fractions" $ do
            retract @Integer @(Ratio Integer) (1 % 2) `shouldBe` Nothing
    describe "instance Injection Double (Complex Double)" $ do
        it "is resolvable" (resolveInjection @Double @(Complex Double))
        it "is injective" (lawInjective @Double @(Complex Double))
    describe "instance Retraction Double (Complex Double)" $ do
        it "is resolvable" (resolveRetraction @Double @(Complex Double))
        it "is the left inverse of inject" (lawLeftInverse @Double @(Complex Double))
        it "is defined over real numbers" $ do
            retract @Double @(Complex Double) (0 :+ 0) `shouldBe` Just 0
            retract @Double @(Complex Double) (1 :+ 0) `shouldBe` Just 1
            retract @Double @(Complex Double) ((-1) :+ 0) `shouldBe` Just (-1)
        it "is not defined over imaginary and complex numbers" $ do
            retract @Double @(Complex Double) (0 :+ 1) `shouldBe` Nothing
            retract @Double @(Complex Double) (1 :+ 1) `shouldBe` Nothing
    describe "instance Injection Integer (Identity Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Identity Integer))
        it "is injective" (lawInjective @Integer @(Identity Integer))
    describe "instance Injection (Identity Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Identity Integer) @Integer)
        it "is injective" (lawInjective @(Identity Integer) @Integer)
    describe "instance Injection (NonEmpty Integer) [Integer]" $ do
        it "is resolvable" (resolveInjection @(NonEmpty Integer) @[Integer])
        it "is injective" (lawInjective @(NonEmpty Integer) @[Integer])
    describe "instance Retraction (NonEmpty Integer) [Integer]" $ do
        it "is resolvable" (resolveRetraction @(NonEmpty Integer) @[Integer])
        it "is the left inverse of inject" (lawLeftInverse @(NonEmpty Integer) @[Integer])
        it "is not defined on the empty list" $ do
            retract @(NonEmpty Integer) @[Integer] [] `shouldBe` Nothing
        it "is defined over non-empty lists" $ do
            retract @(NonEmpty Integer) @[Integer] [0] `shouldBe` Just (0 :| [])
            retract @(NonEmpty Integer) @[Integer] [0, 1] `shouldBe` Just (0 :| [1])
    describe "instance Injection Integer (Down Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Down Integer))
        it "is injective" (lawInjective @Integer @(Down Integer))
    describe "instance Injection (Down Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Down Integer) @Integer)
        it "is injective" (lawInjective @(Down Integer) @Integer)
    describe "instance Injection Integer (Product Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Product Integer))
        it "is injective" (lawInjective @Integer @(Product Integer))
    describe "instance Injection (Product Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Product Integer) @Integer)
        it "is injective" (lawInjective @(Product Integer) @Integer)
    describe "instance Injection Integer (Sum Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Sum Integer))
        it "is injective" (lawInjective @Integer @(Sum Integer))
    describe "instance Injection (Sum Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Sum Integer) @Integer)
        it "is injective" (lawInjective @(Sum Integer) @Integer)
    describe "instance Injection Integer (Dual Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Dual Integer))
        it "is injective" (lawInjective @Integer @(Dual Integer))
    describe "instance Injection (Dual Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Dual Integer) @Integer)
        it "is injective" (lawInjective @(Dual Integer) @Integer)
    describe "instance Injection Integer (Monoid.Last Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Monoid.Last Integer))
        it "is injective" (lawInjective @Integer @(Monoid.Last Integer))
    describe "instance Retraction Injection (Monoid.Last Integer)" $ do
        it "is resolvable" (resolveRetraction @Integer @(Monoid.Last Integer))
        it "is the left inverse of inject" (lawLeftInverse @Integer @(Monoid.Last Integer))
        it "is not defined over mempty" $ do
            retract @Integer @(Monoid.Last Integer) mempty `shouldBe` Nothing
    describe "instance Injection Integer (Monoid.First Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Monoid.First Integer))
        it "is injective" (lawInjective @Integer @(Monoid.First Integer))
    describe "instance Retraction Injection (Monoid.First Integer)" $ do
        it "is resolvable" (resolveRetraction @Integer @(Monoid.First Integer))
        it "is the left inverse of inject" (lawLeftInverse @Integer @(Monoid.First Integer))
        it "is not defined over mempty" $ do
            retract @Integer @(Monoid.First Integer) mempty `shouldBe` Nothing
    describe "instance Injection Integer (Semigroup.First Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Semigroup.First Integer))
        it "is injective" (lawInjective @Integer @(Semigroup.First Integer))
    describe "instance Injection (Semigroup.First Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Semigroup.First Integer) @Integer)
        it "is injective" (lawInjective @(Semigroup.First Integer) @Integer)
    describe "instance Injection Integer (Semigroup.Last Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Semigroup.Last Integer))
        it "is injective" (lawInjective @Integer @(Semigroup.Last Integer))
    describe "instance Injection (Semigroup.Last Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Semigroup.Last Integer) @Integer)
        it "is injective" (lawInjective @(Semigroup.Last Integer) @Integer)
    describe "instance Injection Integer (Max Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Max Integer))
        it "is injective" (lawInjective @Integer @(Max Integer))
    describe "instance Injection (Max Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Max Integer) @Integer)
        it "is injective" (lawInjective @(Max Integer) @Integer)
    describe "instance Injection Integer (Min Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(Min Integer))
        it "is injective" (lawInjective @Integer @(Min Integer))
    describe "instance Injection (Min Integer) Integer" $ do
        it "is resolvable" (resolveInjection @(Min Integer) @Integer)
        it "is injective" (lawInjective @(Min Integer) @Integer)
    describe "instance Injection Integer (String -> Integer)" $ do
        it "is resolvable" (resolveInjection @Integer @(String -> Integer))
    describe "instance Injection Bool Any" $ do
        it "is resolvable" (resolveInjection @Bool @Any)
        it "is injective" (lawInjective @Bool @Any)
    describe "instance Injection Any Bool" $ do
        it "is resolvable" (resolveInjection @Any @Bool)
        it "is injective" (lawInjective @Any @Bool)
    describe "instance Injection Bool All" $ do
        it "is resolvable" (resolveInjection @Bool @All)
        it "is injective" (lawInjective @Bool @All)
    describe "instance Injection All Bool" $ do
        it "is resolvable" (resolveInjection @All @Bool)
        it "is injective" (lawInjective @All @Bool)
    describe "instance Projection Integer Integer" $ do
        it "is resolvable" (resolveProjection @Integer @Integer)
        it "is surjective" (lawSurjective @Integer @Integer id)
    describe "instance Section Integer Integer" $ do
        it "is resolvable" (resolveSection @Integer @Integer)
        it "is the right inverse of project" (lawRightInverse @Integer @Integer)
    describe "instance Projection [(Integer, Integer)] (Map Integer Integer)" $ do
        it "is resolvable" (resolveProjection @[(Integer, Integer)] @(Map Integer Integer))
        it "is surjective" (lawSurjective @[(Integer, Integer)] @(Map Integer Integer) Map.toList)
    describe "instance Projection [Integer] (Set Integer)" $ do
        it "is resolvable" (resolveProjection @[Integer] @(Set Integer))
        it "is surjective" (lawSurjective @[Integer] @(Set Integer) Set.toList)

resolveInjection :: forall from into. Injection from into => Expectation
resolveInjection = seq (inject @from @into) return ()

resolveRetraction :: forall from into. Retraction from into => Expectation
resolveRetraction = seq (retract @from @into) return ()

resolveProjection :: forall whole part. Projection whole part => Expectation
resolveProjection = seq (project @whole @part) return ()

resolveSection :: forall whole part. Section whole part => Expectation
resolveSection = seq (update @whole @part) return ()

lawInjective
    :: forall from into
    .  Injection from into
    => Arbitrary from
    => (Eq from, Show from)
    => (Eq into, Show into)
    => Property
lawInjective = property $ \from1 from2 ->
    let into1 = inject @from @into from1
        into2 = inject @from @into from2
    in (from1 /= from2) ==> (into1 =/= into2)

lawLeftInverse
    :: forall from into
    .  Retraction from into
    => Arbitrary from
    => (Eq from, Show from)
    => Property
lawLeftInverse = property $ \from ->
    let into = inject @from @into from
        from' = retract @from @into into
    in Just from === from'

lawSurjective
    :: forall whole part
    .  Projection whole part
    => Arbitrary part
    => (Eq part, Show part)
    => (part -> whole)  -- ^ any total right inverse of @project@
    -> Property
lawSurjective rightInverse = property $ \part ->
    -- This test actually proves nothing: that a right inverse of 'project'
    -- merely exists proves that it is surjective. Here, we are just probing the
    -- putative right inverse.
    project (rightInverse part) === part

lawRightInverse
    :: forall whole part
    .  Section whole part
    => Arbitrary whole
    => Show whole
    => Arbitrary part
    => (Eq part, Show part)
    => Property
lawRightInverse = property $ \part whole ->
    project @whole @part (update whole part) === part

instance Arbitrary a => Arbitrary (Down a) where
    arbitrary = Down <$> arbitrary


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeSynonymInstances #-}

-- |
-- Module      :  Spiral.Util.Pretty.Maple
-- Copyright   :  (c) 2017-2020 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Spiral.Util.Pretty.Maple (
    Pretty(..)
  ) where

import Data.Ratio
import Data.Complex (Complex(..))
import Data.Monoid ((<>))
import Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Vector as V
import Text.PrettyPrint.Mainland
import qualified Text.PrettyPrint.Mainland.Class as Pretty

import Spiral.Array
import Spiral.Exp
import Spiral.SPL
import Spiral.Util.Pretty (Assoc(..),
                           Fixity(..),
                           HasFixity(..),
                           addPrec,
                           mulPrec,
                           mulPrec1,
                           infixl_,
                           precOf)

class Pretty a where
    {-# MINIMAL pprPrec | ppr #-}
    ppr     :: a -> Doc
    pprPrec :: Int -> a -> Doc

    pprList     :: [a] -> Doc
    pprPrecList :: Int -> [a] -> Doc

    ppr        = pprPrec 0
    pprPrec _  = ppr

    pprPrecList _  = pprList
    pprList xs     = list (map ppr xs)

instance Pretty a => Pretty [a] where
    ppr     = pprList
    pprPrec = pprPrecList

instance Pretty Bool where
    ppr True  = text "true"
    ppr False = text "false"

instance Pretty Char where
    ppr = char

instance Pretty Int where
    ppr = int

instance Pretty Integer where
    ppr = integer

pprRealFrac :: (RealFrac a, Pretty.Pretty a) => a -> Doc
pprRealFrac x
    | isIntegral = Pretty.ppr (ceiling x :: Integer)
    | otherwise  = Pretty.ppr x
  where
    x' :: Integer
    x' = ceiling x

    isIntegral :: Bool
    isIntegral = fromIntegral x' == x

instance Pretty Float where
    ppr = pprRealFrac

instance Pretty Double where
    ppr = pprRealFrac

instance Pretty a => Pretty (Ratio a)  where
    ppr x = text "Fraction" <> parens (commasep [ppr (numerator x), ppr (denominator x)])

instance Pretty a => Pretty (Set a) where
    ppr xs = enclosesep lbrace rbrace comma (map ppr (Set.toList xs))

instance (Pretty e, IArray r DIM2 e) => Pretty (Matrix r e) where
    ppr m = text "Matrix" <> (parens . ppr . toLists) m

instance Pretty (Const a) where
    pprPrec _ (BoolC x)     = ppr x
    pprPrec _ (IntC x)      = ppr x
    pprPrec _ (IntegerC x)  = ppr x
    pprPrec _ (FloatC x)    = ppr x
    pprPrec _ (DoubleC x)   = ppr x
    pprPrec _ (RationalC x) = ppr x

    pprPrec _ (ComplexC r i)
        | r == 0 && i == 0    = char '0'
        | r == 0 && i == 1    = char 'I'
        | r == 0 && i == (-1) = text "-I"
        | r == 0              = ppr i <> char '*' <> char 'I'
        | i == 0              = ppr r
        | otherwise           = text "Complex" <> pprArgs [ppr r, ppr i]

    pprPrec _ (W _ 0 _) = text "1"
    pprPrec _ (W n 1 _) = text "w_" <> Pretty.ppr n
    pprPrec _ (W n k _) = text "w_" <> Pretty.ppr n <> char '^' <> Pretty.ppr k

    pprPrec p (CycC x) = text (showsPrec p x "")

    pprPrec p (PiC x) = parensIf (p > mulPrec) $
                        Pretty.pprPrec mulPrec1 x <> char '*' <> text "Pi"

    pprPrec p (ModularC x) = Pretty.pprPrec p x

instance Pretty Var where
    ppr = Pretty.ppr

instance Pretty (Exp a) where
    pprPrec p (ConstE c) = pprPrec p c
    pprPrec p (VarE v)   = pprPrec p v

    pprPrec p (UnopE op@Neg e) =
        unop p op e

    pprPrec _ (UnopE op e) =
        Pretty.ppr op <> parens (Pretty.ppr e)

    pprPrec _ (BinopE Quot e1 e2) =
        text "iquo" <> pprArgs [ppr e1, ppr e2]

    pprPrec _ (BinopE Rem e1 e2) =
        text "irem" <> pprArgs [ppr e1, ppr e2]

    pprPrec p (BinopE op e1 e2) =
        infixop p op e1 e2

    pprPrec _ (IdxE ev eis) =
        ppr ev <> (brackets . commasep) [ppr (ix+1) | ix <- eis]

    pprPrec p (ComplexE er ei) =
        parensIf (p > addPrec) $
        pprComplex (er :+ ei)

    pprPrec _ (ReE e) =
        text "Re" <> parens (ppr e)

    pprPrec _ (ImE e) =
        text "Im" <> parens (ppr e)

    pprPrec p (BBinopE op e1 e2) =
        infixop p op e1 e2

    pprPrec _ (IfE e1 e2 e3) =
        text "if" <+> ppr e1 <+>
        text "then" <+> ppr e2 <+>
        text "else" <+> ppr e3 <+>
        text "end if"

instance Pretty Unop where
    ppr Neg    = char '-'
    ppr Abs    = text "abs"
    ppr Signum = text "signum"
    ppr Exp    = text "exp"
    ppr Log    = text "log"
    ppr Sqrt   = text "sqrt"
    ppr Sin    = text "sin"
    ppr Cos    = text "cos"
    ppr Asin   = text "asin"
    ppr Acos   = text "acos"
    ppr Atan   = text "atan"
    ppr Sinh   = text "sinh"
    ppr Cosh   = text "cosh"
    ppr Asinh  = text "asinh"
    ppr Acosh  = text "acosh"
    ppr Atanh  = text "atanh"

instance Pretty Binop where
    ppr Add  = char '+'
    ppr Sub  = char '-'
    ppr Mul  = char '*'
    ppr Quot = text "`quot`"
    ppr Rem  = text "`rem`"
    ppr Div  = text "`div`"
    ppr Mod  = text "`mod`"
    ppr FDiv = char '/'

instance Pretty BBinop where
    ppr Eq = text "="
    ppr Ne = text "<>"
    ppr Lt = text "<"
    ppr Le = text "<="
    ppr Ge = text ">="
    ppr Gt = text ">"

instance (Num e, Pretty e) => Pretty (SPL e) where
    pprPrec _ (I n)      = text "IdentityMatrix" <> parens (ppr n)
    pprPrec _ (Diag xs)  = text "DiagonalMatrix" <> parens (list (map ppr (V.toList xs)))
    pprPrec _ (Kron a b) = text "KroneckerProduct" <> pprArgs [ppr a, ppr b]
    pprPrec _ (DSum a b) = text "DiagonalMatrix" <> parens (list [ppr a, ppr b])
    pprPrec p (Prod a b) = infixop p POp a b
    --pprPrec _ (Prod a b) = text "MatrixMatrixMultiply" <> pprArgs [ppr a, ppr b]
    pprPrec p m          = pprPrec p (toMatrix m)

data MatrixBinop = POp
  deriving (Eq, Ord, Show)

instance HasFixity MatrixBinop where
    fixity POp = infixl_ 7

instance Pretty MatrixBinop where
    ppr POp = char '.'

pprArgs :: [Doc] -> Doc
pprArgs = parens . commasep

pprComplex :: (Eq a, Num a, Pretty a) => Complex a -> Doc
pprComplex (r :+ 0) = ppr r
pprComplex (0 :+ i) = ppr i <> char '*' <> char 'I'
pprComplex (r :+ i) = ppr r <+> char '+' <+> ppr i <> char '*' <> char 'I'

unop :: (Pretty a, Pretty op, HasFixity op)
     => Int -- ^ precedence of context
     -> op  -- ^ operator
     -> a
     -> Doc
unop prec op x =
    parensIf (prec > precOf op) $
    ppr op <> pprPrec (precOf op) x

infixop :: (Pretty a, Pretty b, Pretty op, HasFixity op)
        => Int -- ^ precedence of context
        -> op  -- ^ operator
        -> a   -- ^ left argument
        -> b   -- ^ right argument
        -> Doc
infixop prec op l r =
    parensIf (prec > opPrec) $
    pprPrec leftPrec l <+> ppr op <+/> pprPrec rightPrec r
  where
    leftPrec | opAssoc == RightAssoc = opPrec + 1
             | otherwise             = opPrec

    rightPrec | opAssoc == LeftAssoc = opPrec + 1
              | otherwise            = opPrec

    Fixity opAssoc opPrec = fixity op


{-# LANGUAGE BangPatterns #-}

import Gauge.Main (defaultMain, bench, whnf)

import Data.List (unfoldr)
import Data.Maybe (fromJust)
import Data.Primitive.Contiguous (PrimArray, fromListN)
import Statistics.Array.Types (AscList)
import System.Random (random, mkStdGen)

import qualified Statistics.Array as Stats
import qualified Statistics.Array.Types as Asc

main :: IO ()
main = defaultMain
  [ bench "sort-16" $ whnf (fromJust . Asc.fromArray) input16
  , bench "sort-128" $ whnf (fromJust . Asc.fromArray) input128
  , bench "sort-1024" $ whnf (fromJust . Asc.fromArray) input1024
  , bench "sort-65536" $ whnf (fromJust . Asc.fromArray) input65536
  , bench "list-derivative" $ whnf Stats.listDerivative asc1024
  , bench "median-of-absolute-deviations" $ whnf Stats.mad asc1024
  ]
  where
  input16, input128, input1024, input65536 :: PrimArray Int
  !input16 = fromListN 16 $ take 16 . drop 0 $ inputInf
  !input128 = fromListN 128 $ take 128 . drop 16 $ inputInf
  !input1024 = fromListN 1024 $ take 1024 . drop (16+128) $ inputInf
  !input65536 = fromListN 65536 $ take 65536 . drop (16+128+1024) $ inputInf
  asc1024 :: AscList PrimArray Int
  asc1024 = fromJust $ Asc.fromArray input1024

inputInf :: [Int]
inputInf = unfoldr (Just . random) seed
  where
  seed = mkStdGen 1234567


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# Language QuasiQuotes #-}

module MLP.App (
   runInIO,
   run,
   train,
) where

import MLP.Types
    ( curEpoch,
      curError,
      dataSet,
      epochs,
      hiddenLayers,
      learningRate,
      minErrorRate,
      network,
      numInputs,
      numOutputs,
      params,
      runApp,
      topology,
      testSet,
      trainSet,
      App,
      MonadFileSystem(printText),
      Parameters,
      Pattern(Pattern),
      State(State) )
import MLP.Config (readConf, mkDataSet)
import MLP.Network (AllCon, ArrList, Learn(..), Net, MLP(..), NetInstance, arrLast)
import Numeric.LinearAlgebra.Static (R)
import Control.Lens (view, use, (+=), (^.), (.=))
import Control.Monad (when, forM_, forM)
import qualified Data.Vector.Storable as V (snoc, sum)
import Control.Monad.Writer.Class (tell)
import PyF (fmt)
import qualified Data.Text.Lazy as TL
import Control.Monad.Random.Class ( MonadRandom )
import Data.Singletons ( withSomeSing )
import Data.Singletons.Prelude.Tuple ()
import Data.Singletons.Prelude.List ()
import Fcf ()
import Fcf.Data.List ()
import Data.Singletons.TypeLits ( KnownNat, Nat, SNat(..) )
import Data.Constraint ( Dict(..) )
import Numeric.Natural ()
import Data.Maybe ( fromJust )
import MLP.NatList ( List(List), SNatList(..) )

type NetConstraints (i :: Nat) (hs :: [Nat]) (o :: Nat) =
   (AllCon KnownNat (i ': o ': hs), AllCon KnownNat (Topo (Net i hs o)), Learn (Net i hs o))

netDict :: forall (i :: Nat) (o :: Nat) (hs :: [Nat]). (AllCon KnownNat (i ': o ': hs)) =>
           SNatList ('List hs) ->
           Dict (NetConstraints i hs o)
netDict =
   \case
      SNatNil -> Dict
      SNatCons _ SNatNil -> Dict
      SNatCons _ (SNatCons _ SNatNil) -> Dict
      SNatCons _ (SNatCons _ (SNatCons _ SNatNil)) -> Dict

runInIO :: FilePath -> IO ()
runInIO fl = do
   conf <- readConf fl
   let i  = conf ^. topology.numInputs
       hs = conf ^. topology.hiddenLayers
       o  = conf ^. topology.numOutputs

   withSomeSing i $ \(SNat :: SNat i) ->
      withSomeSing o $ \(SNat :: SNat o) ->
         withSomeSing (List (0:hs)) $
            \case
               (SNatCons _ snatlist) ->
                  case netDict @i @o snatlist of
                     (Dict :: Dict (NetConstraints i hs o)) -> run @i @hs @o @IO conf

run :: forall i hs o m. (NetConstraints i hs o, MonadFileSystem m, MonadRandom m) => Parameters -> m ()
run conf = do
   dt <- mkDataSet @i @o conf
   net <- create @(Net i hs o)

   let st = State dt conf net 0 0
       ((a, st'),log) = runApp train st

   mapM_ printText log

trainOnSet :: forall i o hs. NetConstraints i hs o => App i hs o Double
trainOnSet = do
   trainset <- use (dataSet.trainSet)
   lr <- use (params.learningRate)

   errorRates <- forM trainset $ \(Pattern i o) -> do
      net <- use network

      let outs   = netOut i net
          deltas = mkDeltas outs o net
          actOut = fromJust $ arrLast @o outs
          err = netError @(Net i hs o) actOut o

      network .= newNet lr i deltas outs net

      -- tell [fmt| \n ==> Pattern Error {err:.4f} \n |]
      -- tell [fmt| \n ==> Deltas {deltas:s} \n |]
      -- tell [fmt| \n ==> Expected Output {o:s} \n |]
      -- tell [fmt| \n ==> Actual Output {actOut:s} \n |]
      -- tell [fmt| \n Network Dump \n {net:s} \n|]

      return err

   let totalErr = sum errorRates / fromIntegral (length trainset)

   tell [[fmt| \n ==> Epoch Train Error {totalErr:f} \n |]]

   return totalErr

evaluateOnSet :: forall i o hs. NetConstraints i hs o => App i hs o Double
evaluateOnSet = do
   testset <- use (dataSet.testSet)
   net <- use network

   errorRates <- forM testset $ \(Pattern i o) -> do

      let outs   = netOut i net
          deltas = mkDeltas outs o net
          actOut = fromJust $ arrLast @o outs
          err = netError @(Net i hs o) actOut o

      return err

   let totalErr = sum errorRates / fromIntegral (length testset)
   curError .= totalErr

   tell [[fmt| \n ==> Epoch Test Error {totalErr:f} \n |]]

   return totalErr

loop :: App i hs o () -> App i hs o ()
loop trainAction = do
   epochTotal <- use (params.epochs)
   epochNow <- use curEpoch

   when (epochNow < epochTotal) $ do
      trainAction

      curError <- use curError
      minError <- use (params.minErrorRate)

      when (curError > minError) $ do
         loop trainAction

train :: NetConstraints i hs o => App i hs o ()
train = do
   tell ["Starting Training process...\n"]

   loop $ do
      epoch <- use curEpoch

      tell [[fmt| \n ==> Epoch {epoch:d} \n |]]

      trainOnSet
      evaluateOnSet

      curEpoch += 1


{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE GADTs                    #-}
module Grenade.Layers.Internal.Update (
    descendMatrix
  , descendVector
  , MatrixInputValues (..)
  , MatrixResult (..)
  , VectorInputValues (..)
  , VectorResult (..)
  ) where

import           Data.Maybe                   (fromJust)
import qualified Data.Vector.Storable         as U (unsafeFromForeignPtr0,
                                                    unsafeToForeignPtr0)
import           Foreign                      (mallocForeignPtrArray, withForeignPtr)
import           Foreign.Ptr                  (Ptr)
import           GHC.TypeLits
import           Numeric.LinearAlgebra        (Vector, flatten)
import qualified Numeric.LinearAlgebra.Devel  as U
import           Numeric.LinearAlgebra.Static
import           System.IO.Unsafe             (unsafePerformIO)

import           Grenade.Core.Optimizer
import           Grenade.Types

data MatrixInputValues rows columns
  = MatrixValuesSGD
      !(L rows columns) -- ^ current weights
      !(L rows columns) -- ^ gradients
      !(L rows columns) -- ^ last update (old momentum)
  | MatrixValuesAdam
      !Int   -- ^ Step
      !(L rows columns) -- ^ current weights
      !(L rows columns) -- ^ gradients
      !(L rows columns) -- ^ current m
      !(L rows columns) -- ^ current v

data MatrixResult rows columns
  = MatrixResultSGD
      { matrixActivations :: !(L rows columns) -- ^ new activations (weights)
      , matrixMomentum    :: !(L rows columns) -- ^ new momentum
      }
  | MatrixResultAdam
      { matrixActivations :: !(L rows columns) -- ^ new activations (weights)
      , matrixM           :: !(L rows columns) -- ^ new m
      , matrixV           :: !(L rows columns) -- ^ new v
      }

data VectorInputValues r
  = VectorValuesSGD
      !(R r) -- ^ current weights
      !(R r) -- ^ gradients
      !(R r) -- ^ last update (old momentum)
  | VectorValuesAdam
      !Int   -- ^ Step
      !(R r) -- ^ current weights
      !(R r) -- ^ current gradients
      !(R r) -- ^ current m
      !(R r) -- ^ current v

data VectorResult r
  = VectorResultSGD
      { vectorBias     :: !(R r) -- ^ new activations (bias)
      , vectorMomentum :: !(R r) -- ^ new momentum
      }
  | VectorResultAdam
      { vectorBias :: !(R r) -- ^ new activations (bias)
      , vectorM    :: !(R r) -- ^ new m
      , vectorV    :: !(R r) -- ^ new v
      }

descendMatrix :: (KnownNat rows, KnownNat columns) => Optimizer o -> MatrixInputValues rows columns -> MatrixResult rows columns
descendMatrix (OptSGD rate momentum regulariser) (MatrixValuesSGD weights gradient lastUpdate) =
  let (rows, cols) = size weights
      len          = rows * cols
      -- Most gradients come in in ColumnMajor,
      -- so we'll transpose here before flattening them
      -- into a vector to prevent a copy.
      --
      -- This gives ~15% speed improvement for LSTMs.
      weights'     = flatten . tr . extract $ weights
      gradient'    = flatten . tr . extract $ gradient
      lastUpdate'  = flatten . tr . extract $ lastUpdate
      (vw, vm)     = descendUnsafeSGD len rate momentum regulariser weights' gradient' lastUpdate'

      -- Note that it's ColumnMajor, as we did a transpose before
      -- using the internal vectors.
      mw           = U.matrixFromVector U.ColumnMajor rows cols vw
      mm           = U.matrixFromVector U.ColumnMajor rows cols vm
  in  MatrixResultSGD (fromJust . create $ mw) (fromJust . create $ mm)
descendMatrix (OptAdam alpha beta1 beta2 epsilon lambda) (MatrixValuesAdam step weights gradient m v) =
  let (rows, cols) = size weights
      len          = rows * cols
      -- Most gradients come in in ColumnMajor,
      -- so we'll transpose here before flattening them
      -- into a vector to prevent a copy.
      --
      -- This gives ~15% speed improvement for LSTMs.
      weights'  = flatten . tr . extract $ weights
      gradient' = flatten . tr . extract $ gradient
      m'        = flatten . tr . extract $ m
      v'        = flatten . tr . extract $ v
      (vw, vm, vv)     = descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights' gradient' m' v'

      -- Note that it's ColumnMajor, as we did a transpose before
      -- using the internal vectors.
      mw           = U.matrixFromVector U.ColumnMajor rows cols vw
      mm           = U.matrixFromVector U.ColumnMajor rows cols vm
      mv           = U.matrixFromVector U.ColumnMajor rows cols vv
  in  MatrixResultAdam (fromJust . create $ mw) (fromJust . create $ mm) (fromJust . create $ mv)
descendMatrix opt _ = error $ "optimzer does not match to MatrixInputValues in implementation! Optimizer: " ++ show opt

descendVector :: (KnownNat r) => Optimizer o -> VectorInputValues r -> VectorResult r
descendVector (OptSGD rate momentum regulariser) (VectorValuesSGD weights gradient lastUpdate) =
  let len          = size weights
      weights'     = extract weights
      gradient'    = extract gradient
      lastUpdate'  = extract lastUpdate
      (vw, vm)     = descendUnsafeSGD len rate momentum regulariser weights' gradient' lastUpdate'
  in  VectorResultSGD (fromJust $ create vw) (fromJust $ create vm)
descendVector (OptAdam alpha beta1 beta2 epsilon lambda) (VectorValuesAdam step weights gradient m v) =
  let len       = size weights
      weights'  = extract weights
      gradient' = extract gradient
      m'        = extract m
      v'        = extract v
      (vw, vm, vv)     = descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights' gradient' m' v'
  in  VectorResultAdam (fromJust $ create vw) (fromJust $ create vm) (fromJust $ create vv)
descendVector opt _ = error $ "optimzer does not match to VectorInputValues in implementation! Optimizer: " ++ show opt

descendUnsafeSGD :: Int -> RealNum -> RealNum -> RealNum -> Vector RealNum -> Vector RealNum -> Vector RealNum -> (Vector RealNum, Vector RealNum)
descendUnsafeSGD len rate momentum regulariser weights gradient lastUpdate =
  unsafePerformIO $ do
    outWPtr <- mallocForeignPtrArray len
    outMPtr <- mallocForeignPtrArray len
    let (wPtr, _) = U.unsafeToForeignPtr0 weights
    let (gPtr, _) = U.unsafeToForeignPtr0 gradient
    let (lPtr, _) = U.unsafeToForeignPtr0 lastUpdate

    withForeignPtr wPtr $ \wPtr' ->
      withForeignPtr gPtr $ \gPtr' ->
        withForeignPtr lPtr $ \lPtr' ->
          withForeignPtr outWPtr $ \outWPtr' ->
            withForeignPtr outMPtr $ \outMPtr' ->
              descend_sgd_cpu len rate momentum regulariser wPtr' gPtr' lPtr' outWPtr' outMPtr'

    return (U.unsafeFromForeignPtr0 outWPtr len, U.unsafeFromForeignPtr0 outMPtr len)

descendUnsafeAdam ::
     Int -- Len
  -> Int -- Step
  -> RealNum -- Alpha
  -> RealNum -- Beta1
  -> RealNum -- Beta2
  -> RealNum -- Epsilon
  -> RealNum -- Lambda
  -> Vector RealNum -- Weights
  -> Vector RealNum -- Gradient
  -> Vector RealNum -- M
  -> Vector RealNum -- V
  -> (Vector RealNum, Vector RealNum, Vector RealNum)
descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights gradient m v =
  unsafePerformIO $ do
    outWPtr <- mallocForeignPtrArray len
    outMPtr <- mallocForeignPtrArray len
    outVPtr <- mallocForeignPtrArray len
    let (wPtr, _) = U.unsafeToForeignPtr0 weights
    let (gPtr, _) = U.unsafeToForeignPtr0 gradient
    let (mPtr, _) = U.unsafeToForeignPtr0 m
    let (vPtr, _) = U.unsafeToForeignPtr0 v
    withForeignPtr wPtr $ \wPtr' ->
      withForeignPtr gPtr $ \gPtr' ->
        withForeignPtr mPtr $ \mPtr' ->
          withForeignPtr vPtr $ \vPtr' ->
            withForeignPtr outWPtr $ \outWPtr' ->
              withForeignPtr outMPtr $ \outMPtr' ->
                withForeignPtr outVPtr $ \outVPtr' -> descend_adam_cpu len step alpha beta1 beta2 epsilon lambda wPtr' gPtr' mPtr' vPtr' outWPtr' outMPtr' outVPtr'
    return (U.unsafeFromForeignPtr0 outWPtr len, U.unsafeFromForeignPtr0 outMPtr len, U.unsafeFromForeignPtr0 outVPtr len)

foreign import ccall unsafe
    descend_sgd_cpu
      :: Int -> RealNum -> RealNum -> RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> IO ()

foreign import ccall unsafe
    descend_adam_cpu
      :: Int -> Int -> RealNum -> RealNum -> RealNum -> RealNum -> RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> IO ()


module Statistics.Quantile.Bench where

import Control.Applicative

import Criterion

import Pipes
import qualified Pipes.Prelude as P

import Statistics.Quantile.Types
import Statistics.Quantile.Util

import System.IO

dummy :: Selector IO
dummy = Selector dummy'

dummy' :: Quantile -> Stream IO -> IO Double
dummy' _ (Stream src) = do
  runEffect $ src >-> P.map show >-> P.stdoutLn
  pure 0

benchSelector :: [Int]
              -> String
              -> Quantile
              -> Selector IO
              -> FilePath
              -> Benchmark
benchSelector ns name q (Selector select) fp =
  bgroup name $ benchSelector' <$> ns
  where benchSelector' n = bench' n fp (select q)

benchBufferingSelector :: (Int -> Int)
                       -> [Int]
                       -> String
                       -> Quantile
                       -> (Int -> Int -> Selector IO)
                       -> FilePath
                       -> Benchmark
benchBufferingSelector f ns name q select fp =
  bgroup name $ benchSelector' <$> ns
  where benchSelector' n = bench' n fp (select' n q)

        select' n = unSelector  (select (f n) n)

bench' :: Int -> FilePath -> (Stream IO -> IO Double) -> Benchmark
bench' n fp f = do
  bench (show n) $ nfIO $ do
    fh <- liftIO $ openFile fp ReadMode
    (f . takeStream n $ streamHandle fh)

bufSize :: (Int -> Int)
        -> Int
        -> Int
bufSize f n
  | n < 10000 = n
  | otherwise = f n

sqrtN :: Int -> Int
sqrtN = bufSize (floor . sqrt . fromIntegral)

logN :: Int -> Int
logN = bufSize (floor . logBase 2 . fromIntegral)


module Numerical.HBLAS.Lapack.FFI where
import Foreign.Ptr
import Foreign()
import Foreign.C.Types
import Data.Complex
import Data.Int

{-

stylenote: we will not use the LAPACKE_* operations, only the
LAPACKE_*_work variants that require an explicitly provided work buffer.

This is to ensure that solver routine allocation behavior is transparent

-}

{-
void LAPACK_dgesvx( char* fact, char* trans, lapack_int* n, lapack_int* nrhs,
                    double* a, lapack_int* lda, double* af, lapack_int* ldaf,
                    lapack_int* ipiv, char* equed, double* r, double* c,
                    double* b, lapack_int* ldb, double* x, lapack_int* ldx,
                    double* rcond, double* ferr, double* berr, double* work,
                    lapack_int* iwork, lapack_int *info );

-}

{-
    fortran FFI conventions!
-}

--type Stride_C =

newtype Fact_C = Fact_C CChar
newtype Trans_C = Trans_C CChar
newtype Stride_C = Stride_C Int32
newtype Equilib_C = Equilib_C CChar

type Fun_FFI_GESVX el = Ptr Fact_C  {- fact -}-> Ptr Trans_C {- trans -}
    -> Ptr Int32  {-n -}-> Ptr Int32 {- NRHS -}->
    Ptr el {- a -} -> Ptr Stride_C {- lda -} -> Ptr Double {- af -} -> Ptr Stride_C  {- ldf-}->
    Ptr Int32 -> Ptr Equilib_C {- equed -} -> Ptr el {- r -} -> Ptr el  ->
    Ptr el {- b -} -> Ptr Stride_C {- ld b   -} -> Ptr el {- x -} -> Ptr Stride_C {- ldx -}->
    Ptr el {-rcond -}-> Ptr el {- ferr-} -> Ptr el {-berr-} -> Ptr el {-work-}->
    Ptr Int32 {-iwork -}-> Ptr Int32 {-info  -} -> IO ()

{-

the prefixes mean s=single,d=double,c=complex float,d=complex double

fact will be a 1 character C string
either
    "F", then the inputs af and ipiv already contain the permuted LU factorization
        (act as input rather than result params)
    "E", Matrix input A will be equilibriated if needed, then copied to AF and Factored
    "N", matrix input A will be copied to AF

-}

{-
Xgesvx  is the s -sing

im assuming for now that any real use of *gesvx routines, or any other
n^3 complexity algs from LAPACK, are on inputs typically  n>=15, which means > 1000 flops,
which is > 1µs, and thus ok to
-}

--need to get around to wrapping these, but thats for another day
foreign import ccall  "sgesvx_"  sgesvx :: Fun_FFI_GESVX Float
foreign import ccall  "dgesvx_"  dgesvx :: Fun_FFI_GESVX Double
foreign import ccall  "cgesvx_"  cgesvx :: Fun_FFI_GESVX (Complex Float)
foreign import ccall  "zgesvx_"  zgesvx :: Fun_FFI_GESVX (Complex Double)

--lapack_int ?syev_(  char *jobz, char *uplo, lapack_int *n, ?* a, lapack_int * lda, ?* w );
-- ? is Double or Float

newtype JobTy = JBT CChar
newtype UploTy = UPLT CChar
newtype Info = Info Int32

--basic symmetric eigen value solvers
type SYEV_FUN_FFI elem = Ptr JobTy -> Ptr UploTy -> Ptr Int32  -> Ptr elem -> Ptr Int32 -> Ptr elem -> Ptr Info-> IO ()
foreign import ccall "ssyev_" ssyev_ffi :: SYEV_FUN_FFI Float
foreign import ccall "dsyev_" dsyev_ffi :: SYEV_FUN_FFI Double

{-unsafe versions of lapack routines are meant to ONLY be used for workspace queries-}
foreign import ccall unsafe "ssyev_" ssyev_ffi_unsafe :: SYEV_FUN_FFI Float
foreign import ccall unsafe "dsyev_" dsyev_ffi_unsafe :: SYEV_FUN_FFI Double

--lapack_int LAPACKE_<?>gesv( int matrix_order, lapack_int n, lapack_int nrhs, <datatype>* a, lapack_int lda, lapack_int* ipiv, <datatype>* b, lapack_int ldb );
--call sgesv( n, nrhs, a, lda, ipiv, b, ldb, info )

type GESV_FUN_FFI elem = Ptr Int32  {- n -} -> Ptr Int32 {- nrhs -} -> Ptr elem {- a -}
            -> Ptr Int32 {-  lda -} -> Ptr Int32 {- permutation vector -}
            -> Ptr elem  {- b -} -> Ptr Int32 {- ldb -} -> Ptr Info -> IO ()
-- | basic Linear system solvers. they act inplace
foreign import ccall  "sgesv_"  sgesv_ffi  ::GESV_FUN_FFI Float
foreign import ccall  "dgesv_"  dgesv_ffi :: GESV_FUN_FFI Double
foreign import ccall  "cgesv_"  cgesv_ffi :: GESV_FUN_FFI (Complex Float)
foreign import ccall  "zgesv_"  zgesv_ffi :: GESV_FUN_FFI (Complex Double)

-- / not sure if linear solvers ever are run in < 1 microsecond size instances in practice
foreign import ccall unsafe "sgesv_"  sgesv_ffi_unsafe  ::GESV_FUN_FFI Float
foreign import ccall unsafe "dgesv_"  dgesv_ffi_unsafe :: GESV_FUN_FFI Double
foreign import ccall unsafe "cgesv_"  cgesv_ffi_unsafe :: GESV_FUN_FFI (Complex Float)
foreign import ccall unsafe "zgesv_"  zgesv_ffi_unsafe :: GESV_FUN_FFI (Complex Double)

{- only provide  -}


module Numeric.Digamma (digamma) where

import Numeric.SpecFunctions (digamma)


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module      : Graphics.Color.Algebra.Elevator
-- Copyright   : (c) Alexey Kuleshevich 2018-2020
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Graphics.Color.Algebra.Elevator
  ( Elevator(..)
  , module Data.Word
  , clamp01
  ) where

import Data.Complex
import Data.Int
import Data.Typeable
import Data.Vector.Storable (Storable)
import Data.Vector.Unboxed (Unbox)
import Data.Word
import GHC.Float
import Text.Printf

infixl 7 //

defFieldFormat :: FieldFormat
defFieldFormat = FieldFormat Nothing Nothing Nothing Nothing False "" 'v'

-- | A class with a set of functions that allow for changing precision by shrinking and
-- streatching the values.
class (Show e, Eq e, Num e, Typeable e, Unbox e, Storable e) => Elevator e where
  maxValue :: e

  minValue :: e

  fieldFormat :: e -> FieldFormat
  fieldFormat _ = defFieldFormat

  -- | This is a pretty printer for the value.
  toShowS :: e -> ShowS
  default toShowS :: PrintfArg e => e -> ShowS
  toShowS e = formatArg e (fieldFormat e)

  -- | Values are scaled to @[0, 255]@ range.
  toWord8 :: e -> Word8

  -- | Values are scaled to @[0, 65535]@ range.
  toWord16 :: e -> Word16

  -- | Values are scaled to @[0, 4294967295]@ range.
  toWord32 :: e -> Word32

  -- | Values are scaled to @[0, 18446744073709551615]@ range.
  toWord64 :: e -> Word64

  -- | Values are scaled to @[0.0, 1.0]@ range.
  toRealFloat :: (Elevator a, RealFloat a) => e -> a

  -- | Values are scaled from @[0.0, 1.0]@ range.
  fromRealFloat :: (Elevator a, RealFloat a) => a -> e

  -- | Values are scaled to @[0.0, 1.0]@ range.
  toFloat :: e -> Float
  toFloat = toRealFloat

  -- | Values are scaled to @[0.0, 1.0]@ range.
  toDouble :: e -> Double
  toDouble = toRealFloat

  -- | Values are scaled from @[0.0, 1.0]@ range.
  fromDouble :: Double -> e
  fromDouble = fromRealFloat

  -- | Division that works for integral types as well as floating points. May throw an exception.
  (//) :: e -> e -> e

-- | Lower the precision
dropDown :: forall a b. (Integral a, Bounded a, Integral b, Bounded b) => a -> b
dropDown !e = fromIntegral $ fromIntegral e `div` ((maxBound :: a) `div`
                                                   fromIntegral (maxBound :: b))
{-# INLINE dropDown #-}

-- | Increase the precision
raiseUp :: forall a b. (Integral a, Bounded a, Integral b, Bounded b) => a -> b
raiseUp !e = fromIntegral e * ((maxBound :: b) `div` fromIntegral (maxBound :: a))
{-# INLINE raiseUp #-}

-- | Convert to fractional with value less than or equal to 1.
squashTo1 :: forall a b. (Fractional b, Integral a, Bounded a) => a -> b
squashTo1 !e = fromIntegral e / fromIntegral (maxBound :: a)
{-# INLINE squashTo1 #-}

-- | Convert to integral streaching it's value up to a maximum value.
stretch :: forall a b. (RealFloat a, Integral b, Bounded b) => a -> b
stretch !e = round (fromIntegral (maxBound :: b) * clamp01 e)
{-# INLINE stretch #-}

-- | Clamp a value to @[0, 1]@ range.
clamp01 :: RealFloat a => a -> a
clamp01 !x = min (max 0 x) 1
{-# INLINE clamp01 #-}

float2Word32 :: Float -> Word32
float2Word32 d'
  | d' <= 0 = 0
  | d > 4.294967e9 = maxBound
  | otherwise = round d
  where
    d = maxWord32 * d'
{-# INLINE float2Word32 #-}

-- | Same as:
-- >>> fromIntegral (maxBound :: Word32) :: Float
-- 4.2949673e9
--
maxWord32 :: Float
maxWord32 = F# 4.2949673e9#
{-# INLINE maxWord32 #-}

double2Word64 :: Double -> Word64
double2Word64 d'
  | d' <= 0 = 0
  | d > 1.844674407370955e19 = maxBound
  | otherwise = round d
  where
    d = maxWord64 * d'
{-# INLINE double2Word64 #-}

-- | Differs from `fromIntegral` due to: [GHC #17782](https://gitlab.haskell.org/ghc/ghc/issues/17782)
--
-- >>> fromIntegral (maxBound :: Word64) :: Double
-- 1.844674407370955e19
--
maxWord64 :: Double
maxWord64 = D# 1.8446744073709552e19##
{-# INLINE maxWord64 #-}

{-# RULES
"fromRealFloat :: Double -> Word" fromRealFloat = fromDouble :: Double -> Word
"fromRealFloat :: Double -> Word64" fromRealFloat = fromDouble :: Double -> Word64
"fromRealFloat :: Float -> Word32" fromRealFloat = float2Word32
 #-}

-- | Values between @[0, 255]]@
instance Elevator Word8 where
  maxValue = maxBound
  minValue = minBound
  fieldFormat _ = defFieldFormat {fmtWidth = Just 3, fmtChar = 'd'}
  toWord8 = id
  {-# INLINE toWord8 #-}
  toWord16 = raiseUp
  {-# INLINE toWord16 #-}
  toWord32 = raiseUp
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp
  {-# INLINE toWord64 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  fromDouble = toWord8
  {-# INLINE fromDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = toWord8
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 65535]]@
instance Elevator Word16 where
  maxValue = maxBound
  minValue = minBound
  fieldFormat _ = defFieldFormat { fmtWidth = Just 5, fmtChar = 'd'}
  toWord8 = dropDown
  {-# INLINE toWord8 #-}
  toWord16 = id
  {-# INLINE toWord16 #-}
  toWord32 = raiseUp
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp
  {-# INLINE toWord64 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  fromDouble = toWord16
  {-# INLINE fromDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = toWord16
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 4294967295]@
instance Elevator Word32 where
  maxValue = maxBound
  minValue = minBound
  fieldFormat _ = defFieldFormat { fmtWidth = Just 10, fmtChar = 'd'}
  toWord8 = dropDown
  {-# INLINE toWord8 #-}
  toWord16 = dropDown
  {-# INLINE toWord16 #-}
  toWord32 = id
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp
  {-# INLINE toWord64 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  fromDouble = toWord32
  {-# INLINE fromDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = toWord32
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 18446744073709551615]@
instance Elevator Word64 where
  maxValue = maxBound
  minValue = minBound
  fieldFormat _ = defFieldFormat { fmtWidth = Just 20, fmtChar = 'd'}
  toWord8 = dropDown
  {-# INLINE toWord8 #-}
  toWord16 = dropDown
  {-# INLINE toWord16 #-}
  toWord32 = dropDown
  {-# INLINE toWord32 #-}
  toWord64 = id
  {-# INLINE toWord64 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  fromDouble = double2Word64
  {-# INLINE fromDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = toWord64
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 18446744073709551615]@ on 64bit
instance Elevator Word where
  maxValue = maxBound
  minValue = minBound
#if WORD_SIZE_IN_BITS < 64
  fieldFormat _ = defFieldFormat { fmtWidth = Just 10, fmtChar = 'd'}
  toWord64 = dropDown
  {-# INLINE toWord64 #-}
  fromDouble = stretch
  {-# INLINE fromDouble #-}
#else
  fieldFormat _ = defFieldFormat { fmtWidth = Just 20, fmtChar = 'd'}
  toWord64 (W64# w#) = (W# w#)
  {-# INLINE toWord64 #-}
  fromDouble = toWord64 . double2Word64
  {-# INLINE fromDouble #-}
#endif
  toWord8 = dropDown
  {-# INLINE toWord8 #-}
  toWord16 = dropDown
  {-# INLINE toWord16 #-}
  toWord32 = dropDown
  {-# INLINE toWord32 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 127]@
instance Elevator Int8 where
  maxValue = maxBound
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 3, fmtChar = 'd'}
  toWord8 = fromIntegral . max 0
  {-# INLINE toWord8 #-}
  toWord16 = raiseUp . max 0
  {-# INLINE toWord16 #-}
  toWord32 = raiseUp . max 0
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp . max 0
  {-# INLINE toWord64 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 32767]@
instance Elevator Int16 where
  maxValue = maxBound
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 5, fmtChar = 'd'}
  toWord8 = dropDown . max 0
  {-# INLINE toWord8 #-}
  toWord16 = fromIntegral . max 0
  {-# INLINE toWord16 #-}
  toWord32 = raiseUp . max 0
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp . max 0
  {-# INLINE toWord64 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 2147483647]@
instance Elevator Int32 where
  maxValue = maxBound
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 10, fmtChar = 'd'}
  toWord8 = dropDown . max 0
  {-# INLINE toWord8 #-}
  toWord16 = dropDown . max 0
  {-# INLINE toWord16 #-}
  toWord32 = fromIntegral . max 0
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp . max 0
  {-# INLINE toWord64 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 9223372036854775807]@
instance Elevator Int64 where
  maxValue = maxBound
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 19, fmtChar = 'd'}
  toWord8 = dropDown . max 0
  {-# INLINE toWord8 #-}
  toWord16 = dropDown . max 0
  {-# INLINE toWord16 #-}
  toWord32 = dropDown . max 0
  {-# INLINE toWord32 #-}
  toWord64 = fromIntegral . max 0
  {-# INLINE toWord64 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 9223372036854775807]@ on 64bit
instance Elevator Int where
  maxValue = maxBound
  minValue = 0
#if WORD_SIZE_IN_BITS < 64
  fieldFormat _ = defFieldFormat { fmtWidth = Just 10, fmtChar = 'd'}
  toWord64 = dropDown . max 0
  {-# INLINE toWord64 #-}
#else
  fieldFormat _ = defFieldFormat { fmtWidth = Just 19, fmtChar = 'd'}
  toWord64 = fromIntegral . max 0
  {-# INLINE toWord64 #-}
#endif
  toWord8 = dropDown . max 0
  {-# INLINE toWord8 #-}
  toWord16 = dropDown . max 0
  {-# INLINE toWord16 #-}
  toWord32 = dropDown . max 0
  {-# INLINE toWord32 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0.0, 1.0]@
instance Elevator Float where
  maxValue = 1
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 11, fmtPrecision = Just 8, fmtChar = 'f'}
  toWord8 = stretch
  {-# INLINE toWord8 #-}
  toWord16 = stretch
  {-# INLINE toWord16 #-}
  toWord32 = float2Word32
  {-# INLINE toWord32 #-}
  toWord64 = stretch
  {-# INLINE toWord64 #-}
  toFloat = id
  {-# INLINE toFloat #-}
  toDouble = float2Double
  {-# INLINE toDouble #-}
  fromDouble = toFloat
  {-# INLINE fromDouble #-}
  toRealFloat = uncurry encodeFloat . decodeFloat
  {-# INLINE toRealFloat #-}
  fromRealFloat = uncurry encodeFloat . decodeFloat
  {-# INLINE fromRealFloat #-}
  (//) = (/)
  {-# INLINE (//) #-}

-- | Values between @[0.0, 1.0]@
instance Elevator Double where
  maxValue = 1
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 19, fmtPrecision = Just 16, fmtChar = 'f' }
  toWord8 = stretch
  {-# INLINE toWord8 #-}
  toWord16 = stretch
  {-# INLINE toWord16 #-}
  toWord32 = stretch
  {-# INLINE toWord32 #-}
  toWord64 = double2Word64
  {-# INLINE toWord64 #-}
  toFloat = double2Float
  {-# INLINE toFloat #-}
  toDouble = id
  {-# INLINE toDouble #-}
  fromDouble = id
  {-# INLINE fromDouble #-}
  toRealFloat = uncurry encodeFloat . decodeFloat
  {-# INLINE toRealFloat #-}
  fromRealFloat = uncurry encodeFloat . decodeFloat
  {-# INLINE fromRealFloat #-}
  (//) = (/)
  {-# INLINE (//) #-}

{-# RULES
"toRealFloat   :: Double -> Double / Float -> Float" toRealFloat = id
"toRealFloat   :: Double -> Float"                   toRealFloat = double2Float
"toRealFloat   :: Float -> Double"                   toRealFloat = float2Double
"fromRealFloat :: Double -> Double / Float -> Float" fromRealFloat = id
"fromRealFloat :: Double -> Float"                   fromRealFloat = double2Float
"fromRealFloat :: Float -> Double"                   fromRealFloat = float2Double
 #-}

-- | Discards imaginary part and changes precision of real part.
instance (PrintfArg e, Elevator e, RealFloat e) => Elevator (Complex e) where
  maxValue = maxValue :+ maxValue
  minValue = minValue :+ minValue
  toShowS (r :+ i) = toShowS r . formatArg i ((fieldFormat i) {fmtSign = Just SignPlus}) . ('i' :)
  toWord8 = toWord8 . realPart
  {-# INLINE toWord8 #-}
  toWord16 = toWord16 . realPart
  {-# INLINE toWord16 #-}
  toWord32 = toWord32 . realPart
  {-# INLINE toWord32 #-}
  toWord64 = toWord64 . realPart
  {-# INLINE toWord64 #-}
  toFloat = toFloat . realPart
  {-# INLINE toFloat #-}
  toDouble = toDouble . realPart
  {-# INLINE toDouble #-}
  fromDouble = (:+ 0) . fromDouble
  {-# INLINE fromDouble #-}
  toRealFloat = toRealFloat . realPart
  {-# INLINE toRealFloat #-}
  fromRealFloat = (:+ 0) . fromRealFloat
  {-# INLINE fromRealFloat #-}
  (//) = (/)
  {-# INLINE (//) #-}


{-# LANGUAGE DeriveAnyClass  #-}
{-# LANGUAGE DeriveGeneric   #-}
{-# LANGUAGE RecursiveDo     #-}
{-# LANGUAGE TemplateHaskell #-}
module Main where

--import           Circuits.Network
import           Control.Lens
import           Control.Monad.Fix
import qualified Data.Vector                  as V
import qualified Data.Vector.Storable         as VS
import qualified Data.Vector.Storable.Mutable as VSM
import           GHC.Generics
import qualified Numeric.LinearAlgebra        as HMatrix
import qualified Numeric.LinearAlgebra.Data   as HMatrix
import           Text.Printf

import           Circuits.Analysis.DC
import           Circuits.Circuit
import           Circuits.Components.Linear

data TestCircuit = TestCircuit
  { _gnd :: Node
  , _n1  :: Node
  , _vs  :: VoltageSource
  , _r   :: Resistor
  }
  deriving (Show, Generic, SimulateDC, HasVariables)

c1 :: Circuit TestCircuit
c1 = buildCircuit $ do
  gnd <- ground
  n1 <- newNode
  vs <- newBranch
  return TestCircuit
    { _gnd = gnd
    , _n1  = n1
    , _vs  = VoltageSource (Independent 5) vs gnd n1
    , _r   = Resistor 100 gnd n1
    }

data RC = RC
  { _rcGnd :: Node
  , _rcN1  :: Node
  , _rcN2  :: Node
  , _rcVS  :: VoltageSource
  , _rcR   :: Resistor
  , _rcC   :: Capacitor
  }
  deriving (Show, Generic, SimulateDC, HasVariables)
makeLenses ''RC

rc :: Circuit RC
rc = buildCircuit $ do
  gnd <- ground
  n1 <- newNode
  n2 <- newNode
  vs <- newBranch
  return RC
    { _rcGnd = gnd
    , _rcN1  = n1
    , _rcN2  = n2
    , _rcVS  = VoltageSource (Independent 5) vs gnd n2
    , _rcR   = Resistor 100 gnd n1
    , _rcC   = Capacitor 0.01 n1 n2
    }

main :: IO ()
main = do
  let c2 = step SteadyState c1
  print c2
  let rcs = transient (Transient 0 4 0.1) rc
  forMOf_ (traverse . Control.Lens.to (over _2 $ view $ model . rcC . voltageAcross)) rcs $ uncurry $ printf "%.1f ~ %.2f\n"
  --mapM_ (\(t, c) -> printf "t = %.1f\n%s\n\n" t (show c)) rcs
  return ()

data Transient = Transient
  { _transientStartTime :: Double
  , _transientStopTime  :: Double
  , _transientTimeStep  :: Double
  }

transient :: (SimulateDC c, HasVariables c) => Transient -> Circuit c -> [(Double, Circuit c)]
transient (Transient start stop dt) c0 = go c0 start where
  go c t
    | t > stop = []
    | otherwise = (t, c) : go (step (Step t dt) c) (t + dt)

  {-
  sys <- newMSystem 1 1
  stamp sys (VoltageSource (Independent 5) (Var 1) Ground (Var 0))
  stamp sys (Resistor 100 Ground (Var 0) :: Resistor NodeRef NodeRef Double)
  stamp sys (Resistor 200 Ground (Var 0) :: Resistor NodeRef NodeRef Double)
  mv <- HMatrix.reshape 2 <$> VS.unsafeFreeze (_msystemCoeff sys)
  rhs <- HMatrix.asColumn <$> VS.unsafeFreeze (_msystemRhs sys)
  print mv
  print rhs
  print $ HMatrix.linearSolve mv rhs
-}
{-  let simple = empty "GND"
        & part "R1" .~ Just (Resistor 100 "N1" "GND")
        & part "C1" .~ Just (Capacitor 0.01 "N2" "N1")
        & part "Vs" .~ Just (VoltageSource 5 "GND" "N2")

  case compile simple of
    Left err -> putStrLn err
    Right c -> do
      let cs = take 100 $ iterate (step 0.01) c
      mapM_ print cs
-}


{-# LANGUAGE CPP
           , GADTs
           , Rank2Types
           , DataKinds
           , TypeFamilies
           , FlexibleContexts
           , UndecidableInstances
           , LambdaCase
           , MultiParamTypeClasses
           , OverloadedStrings
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs -fsimpl-tick-factor=1000 -fno-warn-orphans #-}
module Language.Hakaru.Runtime.LogFloatPrelude where

#if __GLASGOW_HASKELL__ < 710
import           Data.Functor                    ((<$>))
import           Control.Applicative             (Applicative(..))
#endif
import           Data.Foldable                   as F
import qualified System.Random.MWC               as MWC
import qualified System.Random.MWC.Distributions as MWCD
import           Data.Number.Natural
import           Data.Number.LogFloat            hiding (sum, product)
import qualified Data.Number.LogFloat            as LF
import           Data.STRef
import qualified Data.Vector                     as V
import qualified Data.Vector.Unboxed             as U
import qualified Data.Vector.Generic             as G
import qualified Data.Vector.Generic.Mutable     as M
import           Control.Monad
import           Control.Monad.ST
import           Numeric.SpecFunctions           (logBeta)
import           Prelude                         hiding (init, sum, product, exp, log, (**), pi)
import qualified Prelude                         as P
import           Language.Hakaru.Runtime.CmdLine (Parseable(..), Measure(..), makeMeasure)

-- This Read instance really should be the logfloat package
instance Read LogFloat where
    readsPrec p s = [(logFloat x, r) | (x, r) <- readsPrec p s]

instance Parseable LogFloat where
  parse = return . read

type family MinBoxVec (v1 :: * -> *) (v2 :: * -> *) :: * -> *
type instance MinBoxVec V.Vector v        = V.Vector
type instance MinBoxVec v        V.Vector = V.Vector
type instance MinBoxVec U.Vector U.Vector = U.Vector

type family MayBoxVec a :: * -> *
type instance MayBoxVec ()           = U.Vector
type instance MayBoxVec Int          = U.Vector
type instance MayBoxVec Double       = U.Vector
type instance MayBoxVec LogFloat     = U.Vector
type instance MayBoxVec Bool         = U.Vector
type instance MayBoxVec (U.Vector a) = V.Vector
type instance MayBoxVec (V.Vector a) = V.Vector
type instance MayBoxVec (a,b)        = MinBoxVec (MayBoxVec a) (MayBoxVec b)

newtype instance U.MVector s LogFloat = MV_LogFloat (U.MVector s Double)
newtype instance U.Vector    LogFloat = V_LogFloat  (U.Vector    Double)

instance U.Unbox LogFloat

instance M.MVector U.MVector LogFloat where
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
#if __GLASGOW_HASKELL__ > 710
  {-# INLINE basicInitialize #-}
#endif
  {-# INLINE basicUnsafeReplicate #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  {-# INLINE basicClear #-}
  {-# INLINE basicSet #-}
  {-# INLINE basicUnsafeCopy #-}
  {-# INLINE basicUnsafeGrow #-}
  basicLength (MV_LogFloat v) = M.basicLength v
  basicUnsafeSlice i n (MV_LogFloat v) = MV_LogFloat $ M.basicUnsafeSlice i n v
  basicOverlaps (MV_LogFloat v1) (MV_LogFloat v2) = M.basicOverlaps v1 v2
  basicUnsafeNew n = MV_LogFloat `liftM` M.basicUnsafeNew n
#if __GLASGOW_HASKELL__ > 710
  basicInitialize (MV_LogFloat v) = M.basicInitialize v
#endif
  basicUnsafeReplicate n x = MV_LogFloat `liftM` M.basicUnsafeReplicate n (logFromLogFloat x)
  basicUnsafeRead (MV_LogFloat v) i = logToLogFloat `liftM` M.basicUnsafeRead v i
  basicUnsafeWrite (MV_LogFloat v) i x = M.basicUnsafeWrite v i (logFromLogFloat x)
  basicClear (MV_LogFloat v) = M.basicClear v
  basicSet (MV_LogFloat v) x = M.basicSet v (logFromLogFloat x)
  basicUnsafeCopy (MV_LogFloat v1) (MV_LogFloat v2) = M.basicUnsafeCopy v1 v2
  basicUnsafeMove (MV_LogFloat v1) (MV_LogFloat v2) = M.basicUnsafeMove v1 v2
  basicUnsafeGrow (MV_LogFloat v) n = MV_LogFloat `liftM` M.basicUnsafeGrow v n

instance G.Vector U.Vector LogFloat where
  {-# INLINE basicUnsafeFreeze #-}
  {-# INLINE basicUnsafeThaw #-}
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicUnsafeIndexM #-}
  {-# INLINE elemseq #-}
  basicUnsafeFreeze (MV_LogFloat v) = V_LogFloat `liftM` G.basicUnsafeFreeze v
  basicUnsafeThaw (V_LogFloat v) = MV_LogFloat `liftM` G.basicUnsafeThaw v
  basicLength (V_LogFloat v) = G.basicLength v
  basicUnsafeSlice i n (V_LogFloat v) = V_LogFloat $ G.basicUnsafeSlice i n v
  basicUnsafeIndexM (V_LogFloat v) i
                = logToLogFloat `liftM` G.basicUnsafeIndexM v i
  basicUnsafeCopy (MV_LogFloat mv) (V_LogFloat v)
                = G.basicUnsafeCopy mv v
  elemseq _ x z = G.elemseq (undefined :: U.Vector a) (logFromLogFloat x) z

type Prob = LogFloat

lam :: (a -> b) -> a -> b
lam = id
{-# INLINE lam #-}

app :: (a -> b) -> a -> b
app f x = f x
{-# INLINE app #-}

let_ :: a -> (a -> b) -> b
let_ x f = let x1 = x in f x1
{-# INLINE let_ #-}

ann_ :: a -> b -> b
ann_ _ a = a
{-# INLINE ann_ #-}

exp :: Double -> Prob
exp = logToLogFloat
{-# INLINE exp #-}

log :: Prob -> Double
log = logFromLogFloat
{-# INLINE log #-}

betaFunc :: Prob -> Prob -> Prob
betaFunc a b = exp (logBeta (fromProb a) (fromProb b))

uniform :: Double -> Double -> Measure Double
uniform lo hi = makeMeasure $ MWC.uniformR (lo, hi)
{-# INLINE uniform #-}

normal :: Double -> Prob -> Measure Double
normal mu sd = makeMeasure $ MWCD.normal mu (fromProb sd)
{-# INLINE normal #-}

beta :: Prob -> Prob -> Measure Prob
beta a b = makeMeasure $ \g ->
  unsafeProb <$> MWCD.beta (fromProb a) (fromProb b) g
{-# INLINE beta #-}

gamma :: Prob -> Prob -> Measure Prob
gamma a b = makeMeasure $ \g ->
  unsafeProb <$> MWCD.gamma (fromProb a) (fromProb b) g
{-# INLINE gamma #-}

categorical :: MayBoxVec Prob Prob -> Measure Int
categorical a = makeMeasure $ MWCD.categorical (U.map prep a)
  where prep p = fromLogFloat (p / m)
        m      = G.maximum a
{-# INLINE categorical #-}

plate :: (G.Vector (MayBoxVec a) a) =>
         Int -> (Int -> Measure a) -> Measure (MayBoxVec a a)
plate n f = G.generateM (fromIntegral n) $ \x ->
             f (fromIntegral x)
{-# INLINE plate #-}

bucket :: Int -> Int -> (forall s. Reducer () s a) -> a
bucket b e r = runST
             $ case r of Reducer{init=initR,accum=accumR,done=doneR} -> do
                          s' <- initR ()
                          F.mapM_ (\i -> accumR () i s') [b .. e - 1]
                          doneR s'
{-# INLINE bucket #-}

data Reducer xs s a =
    forall cell.
    Reducer { init  :: xs -> ST s cell
            , accum :: xs -> Int -> cell -> ST s ()
            , done  :: cell -> ST s a
            }

r_fanout :: Reducer xs s a
         -> Reducer xs s b
         -> Reducer xs s (a,b)
r_fanout Reducer{init=initA,accum=accumA,done=doneA}
         Reducer{init=initB,accum=accumB,done=doneB} = Reducer
   { init  = \xs       -> liftM2 (,) (initA xs) (initB xs)
   , accum = \bs i (s1, s2) ->
             accumA bs i s1 >> accumB bs i s2
   , done  = \(s1, s2) -> liftM2 (,) (doneA s1) (doneB s2)
   }
{-# INLINE r_fanout #-}

r_index :: (G.Vector (MayBoxVec a) a)
        => (xs -> Int)
        -> ((Int, xs) -> Int)
        -> Reducer (Int, xs) s a
        -> Reducer xs s (MayBoxVec a a)
r_index n f Reducer{init=initR,accum=accumR,done=doneR} = Reducer
   { init  = \xs -> V.generateM (n xs) (\b -> initR (b, xs))
   , accum = \bs i v ->
             let ov = f (i, bs) in
             accumR (ov,bs) i (v V.! ov)
   , done  = \v -> fmap G.convert (V.mapM doneR v)
   }
{-# INLINE r_index #-}

r_split :: ((Int, xs) -> Bool)
        -> Reducer xs s a
        -> Reducer xs s b
        -> Reducer xs s (a,b)
r_split b Reducer{init=initA,accum=accumA,done=doneA}
          Reducer{init=initB,accum=accumB,done=doneB} = Reducer
   { init  = \xs -> liftM2 (,) (initA xs) (initB xs)
   , accum = \bs i (s1, s2) ->
             if (b (i,bs)) then accumA bs i s1 else accumB bs i s2
   , done  = \(s1, s2) -> liftM2 (,) (doneA s1) (doneB s2)
   }
{-# INLINE r_split #-}

r_add :: Num a => ((Int, xs) -> a) -> Reducer xs s a
r_add e = Reducer
   { init  = \_ -> newSTRef 0
   , accum = \bs i s ->
             modifySTRef' s (+ (e (i,bs)))
   , done  = readSTRef
   }
{-# INLINE r_add #-}

r_nop :: Reducer xs s ()
r_nop = Reducer
   { init  = \_ -> return ()
   , accum = \_ _ _ -> return ()
   , done  = \_ -> return ()
   }
{-# INLINE r_nop #-}

pair :: a -> b -> (a, b)
pair = (,)
{-# INLINE pair #-}

true, false :: Bool
true  = True
false = False

nothing :: Maybe a
nothing = Nothing

just :: a -> Maybe a
just = Just

left :: a -> Either a b
left = Left

right :: b -> Either a b
right = Right

unit :: ()
unit = ()

data Pattern = PVar | PWild
newtype Branch a b =
    Branch { extract :: a -> Maybe b }

ptrue, pfalse :: a -> Branch Bool a
ptrue  b = Branch { extract = extractBool True  b }
pfalse b = Branch { extract = extractBool False b }
{-# INLINE ptrue  #-}
{-# INLINE pfalse #-}

extractBool :: Bool -> a -> Bool -> Maybe a
extractBool b a p | p == b     = Just a
                  | otherwise  = Nothing
{-# INLINE extractBool #-}

pnothing :: b -> Branch (Maybe a) b
pnothing b = Branch { extract = \ma -> case ma of
                                         Nothing -> Just b
                                         Just _  -> Nothing }

pjust :: Pattern -> (a -> b) -> Branch (Maybe a) b
pjust PVar c = Branch { extract = \ma -> case ma of
                                           Nothing -> Nothing
                                           Just x  -> Just (c x) }
pjust _ _ = error "TODO: Runtime.Prelude{pjust}"

pleft :: Pattern -> (a -> c) -> Branch (Either a b) c
pleft PVar f = Branch { extract = \ma -> case ma of
                                           Right _ -> Nothing
                                           Left x -> Just (f x) }
pleft _ _ = error "TODO: Runtime.Prelude{pLeft}"

pright :: Pattern -> (b -> c) -> Branch (Either a b) c
pright PVar f = Branch { extract = \ma -> case ma of
                                            Left _ -> Nothing
                                            Right x -> Just (f x) }
pright _ _ = error "TODO: Runtime.Prelude{pRight}"

ppair :: Pattern -> Pattern -> (x -> y -> b) -> Branch (x,y) b
ppair PVar  PVar c = Branch { extract = (\(x,y) -> Just (c x y)) }
ppair _     _    _ = error "ppair: TODO"

uncase_ :: Maybe a -> a
uncase_ (Just a) = a
uncase_ Nothing  = error "case_: unable to match any branches"
{-# INLINE uncase_ #-}

case_ :: a -> [Branch a b] -> b
case_ e [c1]     = uncase_ (extract c1 e)
case_ e [c1, c2] = uncase_ (extract c1 e `mplus` extract c2 e)
case_ e bs_      = go bs_
  where go []     = error "case_: unable to match any branches"
        go (b:bs) = case extract b e of
                      Just b' -> b'
                      Nothing -> go bs
{-# INLINE case_ #-}

branch :: (c -> Branch a b) -> c -> Branch a b
branch pat body = pat body
{-# INLINE branch #-}

dirac :: a -> Measure a
dirac = return
{-# INLINE dirac #-}

pose :: Prob -> Measure a -> Measure a
pose _ a = a
{-# INLINE pose #-}

superpose :: [(Prob, Measure a)]
          -> Measure a
superpose pms = do
  i <- categorical (G.fromList $ map fst pms)
  snd (pms !! i)
{-# INLINE superpose #-}

reject :: Measure a
reject = Measure $ \_ -> return Nothing

nat_ :: Int -> Int
nat_ = id

int_ :: Int -> Int
int_ = id

unsafeNat :: Int -> Int
unsafeNat = id

nat2prob :: Int -> Prob
nat2prob = fromIntegral

fromInt  :: Int -> Double
fromInt  = fromIntegral

nat2int  :: Int -> Int
nat2int  = id

nat2real :: Int -> Double
nat2real = fromIntegral

fromProb :: Prob -> Double
fromProb = fromLogFloat

unsafeProb :: Double -> Prob
unsafeProb = logFloat

real_ :: Rational -> Double
real_ = fromRational

prob_ :: NonNegativeRational -> Prob
prob_ = fromRational . fromNonNegativeRational

infinity :: Double
infinity = 1/0

abs_ :: Num a => a -> a
abs_ = abs

(**) :: Prob -> Double -> Prob
(**) = pow
{-# INLINE (**) #-}

pi :: Prob
pi = unsafeProb P.pi
{-# INLINE pi #-}

thRootOf :: Int -> Prob -> Prob
thRootOf a b = b ** (recip $ fromIntegral a)
{-# INLINE thRootOf #-}

array
    :: (G.Vector (MayBoxVec a) a)
    => Int
    -> (Int -> a)
    -> MayBoxVec a a
array n f = G.generate (fromIntegral n) (f . fromIntegral)
{-# INLINE array #-}

arrayLit :: (G.Vector (MayBoxVec a) a) => [a] -> MayBoxVec a a
arrayLit = G.fromList
{-# INLINE arrayLit #-}

(!) :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int -> a
a ! b = a G.! (fromIntegral b)
{-# INLINE (!) #-}

size :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int
size v = fromIntegral (G.length v)
{-# INLINE size #-}

reduce
    :: (G.Vector (MayBoxVec a) a)
    => (a -> a -> a)
    -> a
    -> MayBoxVec a a
    -> a
reduce f n v = G.foldr f n v
{-# INLINE reduce #-}

class Num a => Num' a where
    product :: Int -> Int -> (Int -> a) -> a
    product a b f = F.foldl' (\x y -> x * f y) 1 [a .. b-1]
    {-# INLINE product #-}
    summate :: Int -> Int -> (Int -> a) -> a
    summate a b f = F.foldl' (\x y -> x + f y) 0 [a .. b-1]
    {-# INLINE summate #-}

instance Num' Int
instance Num' Double
instance Num' LogFloat where
    product a b f = LF.product (map f [a .. b-1])
    {-# INLINE product #-}
    summate a b f = LF.sum     (map f [a .. b-1])
    {-# INLINE summate #-}

run :: Show a
    => MWC.GenIO
    -> Measure a
    -> IO ()
run g k = unMeasure k g >>= \case
           Just a  -> print a
           Nothing -> return ()

iterateM_
    :: Monad m
    => (a -> m a)
    -> a
    -> m b
iterateM_ f = g
    where g x = f x >>= g

withPrint :: Show a => (a -> IO b) -> a -> IO b
withPrint f x = print x >> f x


{-# LANGUAGE BangPatterns, OverloadedStrings, RecordWildCards, RelaxedPolyRec, ViewPatterns #-}

module Network.HTTP.LoadTest.Report
    (
      reportBasic
    , reportEvents
    , reportFull
    , writeReport
    -- * Other reports
    , csvEvents
    -- * Helper functions
    , buildTime
    ) where

import Control.Monad (forM_)
import Criterion.Analysis (SampleAnalysis(..), OutlierEffect(..),
                           OutlierVariance(..))
import Data.Data (Data)
import Data.Function (on)
import Data.Maybe (fromMaybe)
import Data.Monoid (mappend, mconcat, mempty)
import Data.Text (Text)
import Data.Text.Buildable (build)
import Data.Text.Format (prec, shortest)
import Data.Text.Lazy.Builder (Builder)
import Data.Vector (Vector)
import Network.HTTP.LoadTest.Types (Analysis(..), Basic(..), Event(..),
                                    Summary(..), summEnd)
import Paths_pronk (getDataFileName)
import Prelude hiding (print)
import Statistics.Function (sort)
import Statistics.Resampling.Bootstrap (Estimate(..))
import Statistics.Sample.KernelDensity (kde)
import System.IO (Handle)
import System.IO.Unsafe (unsafePerformIO)
import Text.Hastache (MuType(..))
import Text.Hastache.Context (mkGenericContext)
import qualified Criterion.Report as R
import qualified Data.ByteString.Lazy as L
import qualified Data.HashMap.Strict as H
import qualified Data.List as List
import qualified Data.MeldableHeap as Q
import qualified Data.Text.Format as T
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector as V
import qualified Text.Hastache as H

reportBasic :: Handle -> Analysis Basic -> IO ()
reportBasic h Analysis{..} = do
  let print a b = T.hprint h a b
  print "latency:\n" ()
  print "    mean:    {}\n" [time (mean latency)]
  print "    std dev: {}\n" [time (stdDev latency)]
  print "    99%:     {}\n    99.9%:   {}\n" (time latency99, time latency999)
  print "\nthroughput:  {}\n" [rate throughput]

reportFull :: (IO () -> IO ()) -> Handle -> Analysis SampleAnalysis -> IO ()
reportFull whenLoud h Analysis{..} = do
  let print a b = T.hprint h a b
  print "latency:\n" ()
  print "    mean:    {}\n" [time (estPoint (anMean latency))]
  whenLoud $ do
    print "      lower: {}\n" [time (estLowerBound (anMean latency))]
    print "      upper: {}\n" [time (estUpperBound (anMean latency))]
  print "    std dev: {}\n" [time (estPoint (anStdDev latency))]
  whenLoud $ do
    print "      lower: {}\n" [time (estLowerBound (anStdDev latency))]
    print "      upper: {}\n" [time (estUpperBound (anStdDev latency))]
  effect h (anOutlierVar latency)
  print "    99%:     {}\n    99.9%:   {}\n" (time latency99, time latency999)
  print "\nthroughput:  {}\n" [rate throughput]

time :: Double -> Builder
time = buildTime 4

rate :: Double -> Builder
rate r = prec 4 r `mappend` " req/sec"

buildTime :: Int -> Double -> Builder
buildTime precision t
     | t < 1e-3  = prec precision (t * 1e6) `mappend` " usec"
     | t < 1     = prec precision (t * 1e3) `mappend` " msec"
     | otherwise = prec precision t `mappend` " sec"

effect :: Handle -> OutlierVariance -> IO ()
effect h OutlierVariance{..} =
    case ovEffect of
      Unaffected -> return ()
      _ -> T.hprint h "    estimates {} affected by outliers ({}%)\n"
           (howMuch, T.fixed 1 (ovFraction * 100))
    where howMuch = case ovEffect of
                      Unaffected -> "not" :: Text
                      Slight     -> "slightly"
                      Moderate   -> "moderately"
                      Severe     -> "severely"

reportEvents :: Handle -> Vector Summary -> IO ()
reportEvents h sumv = do
  let evtMap = G.foldl' go H.empty . G.map summEvent $ sumv
      go m e = H.insertWith (+) (classify e) (1::Int) m
      classify Timeout          = 0
      classify HttpResponse{..} = respCode
  T.hprint h "responses:\n" ()
  forM_ (List.sort . H.toList $ evtMap) $ \(e,n) -> do
    let nameOf 0 = "timeout "
        nameOf k = "HTTP " `mappend` build k
    T.hprint h "    {} {}\n" (nameOf e, T.left 7 ' ' n)
  T.hprint h "\n" ()

csvEvents :: Vector Summary -> Builder
csvEvents sums = "start,elapsed,event\n" `mappend` G.foldr go mempty sums
  where
    firstStart = summStart (G.head sums)
    go Summary{..} b = mconcat [
                         shortest $ summStart - firstStart
                       , ","
                       , shortest summElapsed
                       , ","
                       , classify summEvent
                       , "\n"
                       ] `mappend` b
    classify Timeout          = "timeout"
    classify HttpResponse{..} = build respCode

-- | The path to the template and other files used for generating
-- reports.
templateDir :: FilePath
templateDir = unsafePerformIO $ getDataFileName "templates"
{-# NOINLINE templateDir #-}

writeReport :: (Data a) => FilePath -> Handle -> Double -> Analysis a -> IO ()
writeReport template h elapsed a@Analysis{..} = do
  let context "include" = MuLambdaM $
                          R.includeFile [templateDir, R.templateDir]
      context "elapsed"   = MuVariable elapsed
      context "latKdeTimes" = R.vector "x" latKdeTimes
      context "latKdePDF" = R.vector "x" latKdePDF
      context "latKde"    = R.vector2 "time" "pdf" latKdeTimes latKdePDF
      context "latValues" = MuList . map mkGenericContext . G.toList $ lats
      context "thrTimes" = R.vector "x" thrTimes
      context "thrValues" = R.vector "x" thrValues
      context "concTimes" = R.vector "x" . U.fromList $ map fstS conc
      context "concValues" = R.vector "x" . U.fromList $ map sndS conc
      context n = mkGenericContext a n
      (latKdeTimes,latKdePDF) = kde 128 . G.convert . G.map summElapsed $ latValues
      lats = G.map (\s -> s { summStart = summStart s - t }) latValues
          where t = summStart . G.head $ latValues
      (thrTimes,thrValues) = graphThroughput (min (G.length latValues) 50) elapsed latValues
      conc = graphConcurrency lats
  tpl <- R.loadTemplate [".",templateDir] template
  bs <- H.hastacheStr H.defaultConfig tpl context
  L.hPutStr h bs

data T = T (U.Vector Double) {-# UNPACK #-} !Double

-- | Compute a graph of throughput, requests completed per time
-- interval.
graphThroughput :: Int          -- ^ Number of time slices.
                -> Double       -- ^ Amount of time elapsed.
                -> V.Vector Summary -> (U.Vector Double, U.Vector Double)
graphThroughput slices elapsed sumv =
    (G.generate slices $ \i -> fromIntegral i * timeSlice,
     G.unfoldrN slices go (T endv 0))
  where go (T v i) = Just (fromIntegral (G.length a), T b j)
           where (a,b) = G.span (<=t) v
                 t = start + (j * timeSlice)
                 j = i+1
        timeSlice = elapsed / fromIntegral slices
        start = summStart . G.head $ sumv
        endv = G.convert . sort . G.map summEnd $ sumv

data S = S {
      fstS :: {-# UNPACK #-} !Double
    , sndS :: {-# UNPACK #-} !Int
    }

-- | Compute a graph of concurrency.
graphConcurrency :: V.Vector Summary -> [S]
graphConcurrency = scanl1 f . filter ((/=0) . sndS) . map (foldl1 (flip f)) .
                   List.groupBy ((==) `on` fstS) . go Q.empty . G.toList
  where
    f (S _ i) (S t j) = S t (i+j)
    go q es@(Summary{..}:xs)
        | summStart < t = S summStart 1 : go insQ xs
        | otherwise     = S t (-1)      : go delQ es
      where (t,delQ) = fromMaybe (1e300,q) $ Q.extractMin q
            insQ     = Q.insert (summStart+summElapsed) q
    go q _ = drain q
      where drain (Q.extractMin -> Just (t,q')) = S t (-1) : drain q'
            drain _ = []


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ViewPatterns     #-}
{-# LANGUAGE TypeFamilies     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnicodeSyntax #-}
--
-- |
-- Module      : Main
-- Description : Measures execution time of MC loop for 1D Heisenberg
-- Hamiltonian. Try @--help@ to get usage information.
-- Copyright   : (c) Tom Westerhout, 2018
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental
module Main where

import           Prelude                 hiding ( map
                                                , zipWithM
                                                )

import           Debug.Trace
import qualified System.Random.MWC             as MWC
-- import           Control.Lens hiding((<.>))
import           Control.Monad.Reader    hiding ( zipWithM )
import           Control.Monad.Primitive
import qualified Data.List                     as L
import           Data.Complex
import           Data.Semigroup
import           Data.Vector.Storable           ( Vector
                                                , (!)
                                                )
import qualified Data.Vector.Storable          as V
import qualified Data.Vector.Storable.Mutable  as MV
import           System.Exit
import           System.IO               hiding ( hGetLine )
import           System.Environment             ( getArgs
                                                , getProgName
                                                )
import           Foreign.Storable
import           Data.Text                      ( Text )
import           Data.Text.IO                   ( hGetLine )
import qualified Data.Text.IO                  as T

import           Lens.Micro
import           Lens.Micro.Extras

import           PSO.Random
import           PSO.FromPython
import           NQS.Rbm (Rbm, mkRbm)
import           NQS.Rbm.Mutable -- (sampleGradients)
import           NQS.Internal.Hamiltonian
import           NQS.Internal.Types -- (ℂ, ℝ)
import           NQS.Internal.Rbm (unsafeThawRbm)
import           NQS.Internal.LAPACK

fromPyFile :: FilePath -> IO Rbm
fromPyFile name = withFile name ReadMode toRbm
  where toRight :: Either String (Vector ℂ) -> Vector ℂ
        toRight (Right x) = x
        toRight (Left x)  = error x
        toRbm h = do
          hGetLine h
          a <- trace ("a...") $ toRight <$> readVector <$> hGetLine h
          b <- trace ("b...") $ toRight <$> readVector <$> hGetLine h
          s <- hGetLine h
          T.putStrLn s
          let !w = trace ("w...") $ toRight $ readMatrix s
          return $ mkRbm a b w

randomRbm :: Int -> Int -> (ℝ, ℝ) -> (ℝ, ℝ) -> (ℝ, ℝ) -> IO Rbm
randomRbm n m (lowV, highV) (lowH, highH) (lowW, highW) = do
  g <- mkMWCGen (Just 123)
  flip runReaderT g $ do
    visible <- uniformVector n (lowV :+ lowV, highV :+ highV)
    hidden  <- uniformVector m (lowH :+ lowH, highH :+ highH)
    weights <- uniformVector (n * m) (lowW :+ lowW, highW :+ highW)
    return $ trace ("mkRbm...") (mkRbm visible hidden weights)

numberSteps :: (Int, Int, Int) -> Int
numberSteps (low, high, step) = (high - low - 1) `div` step + 1

main = do -- NQS.Internal.LAPACK.test
  let filename = "/home/tom/src/tcm-swarm/cbits/test/input/rbm_6_6_0.in"
  rbm <- unsafeThawRbm =<< fromPyFile filename :: IO (MRbm (PrimState IO))
  -- rbm <- unsafeThawRbm =<< randomRbm 100 200 (-0.1, 0.1) (-0.1, 0.1) (-0.05, 0.05)
  let n = sizeVisible rbm
      config = defaultMCConfig & steps .~ (1000, 21000 * n, n)
                               & magnetisation .~ (Just 0)
      nParams = size rbm
      nSteps = numberSteps $ config ^. steps
  print n
  hamiltonian <- heisenberg (1, 1) (Just 5.0) (zip [0..] ([1..(n - 1)] ++ [0]))
  moments <- MV.new 4
  -- f <- newDenseVector nParams
  -- grad <- newDenseMatrix (nSteps * config ^. runs) nParams
  sampleMoments config hamiltonian rbm moments -- f grad
  e <- MV.read moments 0
  print ["Hello!", show e]


module SimpleParser where

import SimpleDefs
import Control.Monad (liftM)
import Text.ParserCombinators.Parsec hiding (spaces)
import Numeric (readFloat, readHex, readOct)
import Data.Complex (Complex((:+)))
import Data.Char (toLower)
import Data.Ratio (Rational, (%))

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

spaces :: Parser ()
spaces = skipMany1 space

{-
parseChar :: Parser LispVal
parseChar = do string "#\\"
               c <- many1 letter
               return $ Char $ case map toLower c of
                   "space" -> ' '
                   "newline" -> '\n'
                   [x] -> x
-}

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

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

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

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

{-
parseRadixNumber :: Parser LispVal
parseRadixNumber = char '#' >>
                (   parseDecimal
                <|> parseOctal
                <|> parseHexadecimal
                <|> parseBinary
                )

parseDecimal, parseHexadecimal, parseBinary, parseOctal :: Parser LispVal
parseDecimal = do char 'd'
                  number <- many1 digit
                  return . Number . read $ number

parseHexadecimal = do char 'x'
                      number <- many $ digit <|> (oneOf "abcdef")
                      return . Number . fst . head . readHex $ number

parseBinary = do char 'b'
                 number <- many1 . oneOf $ "01"
                 return . Number . readBin $ number

parseOctal = do char 'o'
                number <- many1 . oneOf $ "01234567"
                return . Number . fst . head . readOct $ number

readBin :: String -> Integer
readBin bin = fst $ foldl (\(n, ix) x -> (n + (num x) * (2 ^ ix), ix - 1))
                          (0, (length bin) - 1)
                          bin where
                  num '0' = 0
                  num '1' = 1
-}
parseNumber :: Parser LispVal
parseNumber = parseSimpleNumber -- <|> parseRadixNumber

{-
parseFloat :: Parser LispVal
parseFloat = do before <- many1 digit
                char '.'
                after <- many1 digit
                return . Float . fst . head . readFloat $ before ++ '.' : after

parseRational :: Parser LispVal
parseRational = do before <- fmap read $ many1 digit
                   char '/'
                   after <- fmap read $ many1 digit
                   return . Rational $ (before % after)

parseComplex :: Parser LispVal
parseComplex = do real <- fmap toFloat $ parseFloat <|> parseSimpleNumber
                  char '+'
                  imag <- fmap toFloat $ parseFloat <|> parseSimpleNumber
                  char 'i'
                  return . Complex $ (real :+ imag) where
                      toFloat (Float x) = x
                      toFloat (Number x) = fromInteger x
-}

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

parseList :: Parser LispVal
parseList = liftM List $ sepBy parseExpr spaces

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

parseExpr :: Parser LispVal
parseExpr = parseAtom
        -- <|> parseBool
        -- <|> parseChar
        <|> parseString
        -- <|> parseComplex
        -- <|> parseRational
        -- <|> parseFloat
        <|> parseNumber
        <|> parseQuoted
        <|> do char '('
               x <- (try parseList) <|> parseDottedList
               char ')'
               return x


{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
module Statistics.BBVI.Propagator
  ( DistCell(..)
  , DistCells
  , DistCellss
  , SampleVector
  , SampleDouble
  , Gradient
  , Memory
  , Time
  , mergeGeneric
  , mergeGenerics
  , mergeGenericss
  , dist
  , time
  , defaultDistCell
  )
where

import           Statistics.BBVI.Class
import qualified Data.Vector                   as V
import           Data.Propagator

-- | a sample with a vector type e.g. sample from a dirichlet
type SampleVector = V.Vector Double

-- | a sample with a real type e.g. number sampled from a gaussian
type SampleDouble = Double

type Gradient = V.Vector Double

type Memory = V.Vector Double

-- | number of times a distribution cell has been updated
type Time = Int

-- | distribution cell
data DistCell a
  =
    U !Memory !Gradient -- ^ update to a distribution cell
  | Node !Time !Memory !a -- ^ distribution cell with a distribution
                          -- of type a
  deriving (Show, Eq, Ord, Read)

-- !(Gradient -> DistCell a -> (Memory, Gradient))

-- | helper function for initializing distribution cells
defaultDistCell
  :: DistUtil a
  => a -- ^ initial point for distribution
  -> DistCell a -- ^ distribution cell
defaultDistCell d = Node 1 (V.replicate (nParams d) 0) d

-- | time accessor for a distribution cell
time :: DistCell a -> Time
time (Node t _ _) = t
time (U{}       ) = error "called time on update propnode!"
-- memory (Node _ m _) = m

-- | distribution accessor for a distribution cell
dist :: DistCell a -> a
dist (Node _ _ d) = d
dist (U{}       ) = error "called dist on update propnode!"

-- | single cell representing a vector of distribution cells
type DistCells a = V.Vector (DistCell a)

-- | single cell representing an array (vector of vector) of distribution cells
type DistCellss a = V.Vector (V.Vector (DistCell a))

-- | generic merge for DistCells: useful for customizing
-- "quiesence" thresholds with cellWith
mergeGeneric
  :: DistUtil a
  => Time -- ^ maximum number of updates to a cell
  -> Double -- ^ threshold for information to be considered "new"; if
            -- the l2-norm of the change in the distribution's
            -- parameters is less than this threshold, the cell
            -- remains unchanged
  -> DistCell a -- ^ current cell
  -> DistCell a -- ^ proposed update to cell
  -> Change (DistCell a)
mergeGeneric maxStep delta !x1 !x2 = m x1 x2
 where
  m no@(Node t memory d) (U memUp gradUp)
    | norm gradUp < delta = Change False no
    | t >= maxStep        = Change False no
    | otherwise           = Change True updateNode
   where
    updateNode = Node (t + 1)
                      (V.zipWith (+) memory memUp)
                      (fromParamVector newQ)
      where newQ = V.zipWith (+) (toParamVector d) gradUp
  m (U _ _) _ = Contradiction mempty "Trying to update a gradient"
  -- CAREFUL: below is dangerous if I start doing the ideas i thought
  -- about: changing maxstep and elta node for local optmizations
  m no1@(Node t1 _m1 d1) no2@(Node t2 _m2 d2)
    | t1 >= maxStep
    = Change False no1
    | (t2 > t1)
      && (norm (V.zipWith (-) (toParamVector d1) (toParamVector d2)) >= delta)
    = Change True no2
    | otherwise
    = Change False no1

-- | generalization of 'mergeGeneric' to cells of vectors of distributions
mergeGenerics
  :: DistUtil a
  => Time -- ^ maximum number of updates to the cells
  -> Double -- ^ threshold for information to be considered "new"; if
            -- the l2-norm of change in the distribution's parameters
            -- is less than this threshold for all distributions, the
            -- cell remains unchanged
  -> DistCells a -- ^ current cells
  -> DistCells a -- ^ proposed update to cells
  -> Change (DistCells a)
mergeGenerics m d x1 x2 = V.sequence . V.zipWith (mergeGeneric m d) x1 $ x2

-- | generalization of 'mergeGeneric' to cells of arrays of distributions
mergeGenericss
  :: DistUtil a
  => Time -- ^ maximum number of updates to the cells
  -> Double -- ^ threshold for information to be considered "new"; if
            -- the l2-norm of change in the distribution's parameters
            -- is less than this threshold for all distributions, the
            -- cell remains unchanged
  -> DistCellss a -- ^ current cell
  -> DistCellss a -- ^ proposed update to cell
  -> Change (DistCellss a)
mergeGenericss m d v1 v2 = V.sequence . V.zipWith (mergeGenerics m d) v1 $ v2

instance DistUtil a => Propagated (DistCell a) where
  merge = mergeGeneric 1000000 1e-16

instance DistUtil a => Propagated (DistCells a) where
  merge ns updates = V.sequence $ V.zipWith merge ns updates

instance DistUtil a => Propagated (DistCellss a) where
  merge ns updates = V.sequence $ V.zipWith merge ns updates

norm :: V.Vector Double -> Double
norm = sqrt . V.sum . V.map (^ (2 :: Int))