import Control.Monad (replicateM)
import System.Random.MWC (create, GenIO)
import System.Random.MWC.Distributions (beta, bernoulli, dirichlet)
import Numeric.LinearAlgebra.HMatrix
       (Matrix, matrix, Vector, vector, konst, fromRows, (!), toList, fromList, size)
import qualified Data.Vector.Storable as V

import Utils (multinomial)
import VB (lQ, calculateEQZ, calculateQZ, calculateEQmu, calculateQmu, expectationBeta, em)

generateSampleMatrix :: Double  -- ^ alpha parameter for Beta distribution
                     -> Double  -- ^ beta parameter for Beta distribution
                     -> Vector Double  -- ^ alpha parameter for Dirichlet distr
                     -> Int  -- ^ number of samples
                     -> Int  -- ^ vector size
                     -> Int  -- ^ Number of mixins
                     -> IO (Matrix Double, Vector Double)
generateSampleMatrix a b alpha n d k = do
    gen <- create
    pivT <- dirichlet (toList alpha) gen
    let piv = fromList pivT
    muK <-
        replicateM
            k
            (replicateM
                 d
                 (beta a b gen) >>=
             \m ->
                  return (vector m))
    let mu = fromRows muK
    genData <-
        replicateM
            n
            (generateVectors piv mu gen)
    return (fromRows genData, piv)
    where generateVectors :: Vector Double
                          -> Matrix Double
                          -> GenIO
                          -> IO (Vector Double)
          generateVectors piv mu gen = do
              zn <- multinomial piv gen
              let Just j = V.findIndex (== 1) zn
              xn <-
                  mapM
                      (\i ->
                            bernoulli
                                ((mu ! j) !
                                 i)
                                gen >>=
                            \b ->
                                 if b
                                     then return 1.0
                                     else return 0.0)
                      [0 .. d - 1]
              return (vector xn)

-- | L(q) function must increase
testLQ = do
  let k = 6
      n = 1000
      d = 3
      a = 0.6
      b = 1
      alpha = konst (1 / 6) k
  (mX, piV) <- generateSampleMatrix a b alpha n d k
  let test_k = 10
      test_alpha = konst 0.1 test_k
      test_a = matrix d [0.5 | _ <- [1..test_k*d]]
      test_b = matrix d [0.5 | _ <- [1..test_k*d]]
  gen <- create
  pivT <- dirichlet (toList test_alpha) gen
  let piv = fromList pivT
  let qZ = fromRows (replicate n piv)
  let eQz = calculateEQZ qZ
  let qMu = (test_a, test_b)
  let eQmu = calculateEQmu qMu
  print (lQ mX eQz eQmu)
  let qZ' = calculateQZ mX eQmu eQz
      eQz' = calculateEQZ qZ'
      qMu' = calculateQmu mX eQz' qMu
      eQmu' = calculateEQmu qMu'
  print (lQ mX eQz' eQmu')
  let qZ'' = calculateQZ mX eQmu' eQz'
      eQz'' = calculateEQZ qZ''
      qMu'' = calculateQmu mX eQz'' qMu'
      eQmu'' = calculateEQmu qMu''
  print (lQ mX eQz'' eQmu'')

-- | EM first step. Can't imagine.
testEM = do
  let k = 6
      n = 1000
      d = 3
      a = 0.6
      b = 1
      alpha = konst (1 / 6) k
  (mX, piV) <- generateSampleMatrix a b alpha n d k
  let test_k = 10
      test_alpha = konst 0.1 test_k
      test_a = 0.5
      test_b = 0.5
  gen <- create
  pivT <- dirichlet (toList test_alpha) gen
  let piv = fromList pivT
  print "Original pi vector"
  print piV
  print "Calculated pi vector"
  print (em mX piv test_alpha test_a test_b)

main :: IO ()
main = do
  putStrLn "Test suite not yet implemented"
  let k = 3
      n = 10
      d = 3
      a = 1
      b = 1
      alpha = konst 0.01 k
  d <- generateSampleMatrix a b alpha n d k
  print d
  putStrLn "Test L(q). Must increase."
  testLQ
  putStrLn "Test EM work or not :)"
  testEM
  return ()


{-# LANGUAGE CPP
           , TypeOperators
           , KindSignatures
           , TypeApplications
           , ScopedTypeVariables
           , AllowAmbiguousTypes
           , TypeFamilies
           , FlexibleContexts
           , DeriveFunctor
           , DeriveFoldable
           , DeriveTraversable
           , TemplateHaskell
           , RankNTypes
           , UndecidableInstances
           -- , DeriveAnyClass
           -- , GeneralizedNewtypeDeriving
           -- , DatatypeContexts
  #-}
{-# OPTIONS_GHC -Wall #-}

module Main where

import Control.Applicative
import Control.Monad (unless)
import Data.Complex
import Data.Traversable
import System.Exit (exitFailure)
import Test.QuickCheck (choose, vectorOf)  -- , elements, collect)
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.All (quickCheckAll)

dft :: [Complex Double] -> [Complex Double]
dft xs = [ sum [ x * exp((0.0 :+ (-1.0)) * 2 * pi / lenXs * fromIntegral(k * n))
          | (x, n) <- zip xs [0..]  ]
          | k <- [0..(length xs - 1)]  ]
 where lenXs = fromIntegral $ length xs

type Unop a = a -> a

dft' :: forall f a. ( Applicative f, Traversable f, Sized f
                    , RealFloat a) => Unop (f (Complex a))
dft' xs = omegas (size @f) `cross` xs

cross :: (Functor n, Foldable m, Applicative m, Num a) =>
  n (m a) -> m a -> n a  -- matrix * vector
mat `cross` vec = dot vec <$> mat

dot :: (Foldable f, Applicative f, Num a) =>
  f a -> f a -> a        -- dot product
u `dot` v = sum (liftA2 (*) u v)

newtype (g :. f) a = O (g (f a))
  deriving (Eq, Functor, Foldable, Traversable)

instance (Show a, Show (f a), Show (g String), Functor g) => Show ((g :. f) a) where
  show (O x) = show $ fmap show x

instance (Applicative f, Applicative g) => Applicative (g :. f) where
  pure = O . pure . pure
  O h <*> (O x) = O $ fmap (<*>) h <*> x

class Sized (f :: * -> *) where
    size :: Integer

instance (Sized f, Sized g) => Sized (g :. f) where
    size = size @f * size @g

class FFT f where
  type Reverse f :: * -> *
  fft :: RealFloat a => f (Complex a) -> Reverse f (Complex a)

-- Constraints filled in, by responding to compiler errors incrementally.
instance (Traversable g, Traversable (Reverse g),
          Applicative (Reverse g),
          FFT g,
          Traversable f,
          Applicative f, Applicative (Reverse f),
          FFT f, Sized f, Sized (Reverse g) ) =>
  FFT (g :. f) where
    type Reverse (g :. f) = Reverse f :. Reverse g
    fft = O . fft' . transpose . twiddle . fft' . unO
fft' :: ( Traversable g
        , Applicative (Reverse g)
        , FFT g
        , Traversable f
        , Applicative f
        , RealFloat a ) =>
  g (f (Complex a)) -> Reverse g (f (Complex a))
fft' = transpose . fmap fft . transpose

twiddle :: forall g f a. ( Applicative g, Traversable g, Sized g
                         , Applicative f, Traversable f, Sized f
                         , RealFloat a ) =>
  g (f (Complex a)) -> g (f (Complex a))
twiddle = (liftA2 . liftA2) (*) (omegas (size @(g :. f)))

omegas :: ( Applicative g, Traversable g
          , Applicative f, Traversable f
          , RealFloat a) =>
  Integer -> g (f (Complex a))
-- 'i' was unrecognized.
-- omegas n = powers <$> powers (exp (-i * 2 * pi / fromIntegral n))
omegas n = powers <$> powers (cis (2 * pi / fromIntegral n))

powers :: ( Applicative f, Traversable f
          , Fractional a) => a -> f a
-- I don't have time to digest Conal's "Generic Parallel Scan" talk, right now.
-- powers = fst . lscanAla Product . pure
-- But, I do need to accomodate all applicative functors, as opposed to just lists.
-- So, I'll try the Traversable approach.
powers w = fmap (/ w) . snd . mapAccumL (\x y -> (x * y, x * y)) 1 $ pure w

unO :: (g :. f) a -> g (f a)
unO (O x) = x

-- Copied from Conal's generic-fft library.
transpose :: (Traversable g, Applicative f) => g (f a) -> f (g a)
transpose = sequenceA

-- Quad
newtype Quad a = Quad ((a,a),(a,a))
  deriving (Show, Functor, Foldable, Traversable)
instance Applicative Quad where
  pure x = Quad ((x,x),(x,x))
  Quad ((g,h),(u,v)) <*> Quad ((w,x),(y,z)) = Quad ((g w, h x),(u y, v z))
instance Sized Quad where
  size = 4

-- Pair
data  Pair a = a :# a
  deriving ( Show, Eq, Functor, Foldable, Traversable )
instance Applicative Pair where
  pure x = x :# x
  g :# h <*> (x :# y) = g x :# h y
instance Sized Pair where
  size = 2
instance FFT Pair where
  type Reverse Pair = Pair
  fft (x :# y) = (x + y) :# (x - y)

-- Tree of Pairs
data Tree a = Leaf (Pair a)
            | Branch (Pair (Tree a))
  deriving (Show, Eq, Functor, Foldable, Traversable)

-- Tree building utility.
mkTree :: Integer -> [a] -> Tree a
mkTree n xs | n == 0    = Leaf (head xs :# head (tail xs))
            | otherwise = Branch (mkTree (n-1) (evens xs) :# mkTree (n-1) (odds xs))

-- Pair of Pair building utility.
mkPofP :: [a] -> Pair (Pair a)
-- mkPofP []        = error "Empty list provided!"
-- mkPofP [w]       = error "Wrong number of elements given!"
-- mkPofP [w,x]     = error "Wrong number of elements given!"
-- mkPofP [w,x,y]   = error "Wrong number of elements given!"
-- mkPofP [w,x,y,z] = O ((w :# x) :# (y :# z))
-- mkPofP xs        = O (unO (mkPofP (evens xs)) :# (unO (mkPofP (odds xs))))
mkPofP [w,x,y,z] = (w :# y) :# (x :# z)
-- mkPofP xs        = mkPofP (evens xs) :# mkPofP (odds xs)
mkPofP _         = error "Wrong number of elements given!"

evens :: [a] -> [a]
evens []     = []
evens (x:xs) = x : odds xs

odds :: [a] -> [a]
odds []     = []
odds (_:xs) = evens xs

newtype Id a = Id {unId :: a}
  deriving ( Show, Eq, Functor, Foldable, Traversable )

instance Applicative Id where
  pure = Id
  Id g <*> (Id x) = Id (g x)

instance Sized Id where
  size = 1

instance FFT Id where
  type Reverse Id = Id
  fft (Id x) = Id x

toList :: Foldable f => f a -> [a]
toList = foldr (:) []

assertion1 :: Bool
assertion1 = dft (map (:+ 0.0) [1.0, 0.0, 0.0, 0.0]) == [1.0 :+ 0.0, 1.0 :+ 0.0, 1.0 :+ 0.0, 1.0 :+ 0.0]

assertion2 :: Bool
assertion2 = dft (map (:+ 0.0) [1.0, 0.0, 0.0, 0.0]) == toList (dft' (fmap (:+ 0.0) (Quad ((1.0, 0.0), (0.0, 0.0)))))

assertion3 :: Bool
assertion3 = fft (fmap (:+ (0 :: Double)) (O ((1 :# 0) :# (0 :# 0)))) == dft' (fmap (:+ (0::Double)) (O ((1 :# 0) :# (0 :# 0))))

exeMain :: IO ()
exeMain = do
  print (powers (cis (2 * pi / 4)) :: Quad (Complex Float))
  print assertion1
  print assertion2
  print assertion3

-- QuickCheck types & propositions
newtype FFTTestVal  f g = FFTTestVal {
  getVal :: (g :. f) (Complex Double)
} -- deriving (Show)

instance (Show (f (Complex Double)), Show (g String), Functor g) =>
  Show (FFTTestVal f g) where
    show = show . getVal

instance (Applicative f, Applicative g) => Arbitrary (FFTTestVal f g) where
    arbitrary = do
        depth  <- choose (2, 5)
        -- rvals <- vectorOf (2^depth) (choose (-1.0, 1.0))
        -- ivals <- vectorOf (2^depth) (choose (-1.0, 1.0))
        -- let values = zipWith (:+) rvals ivals
        -- return $ FFTTestVal (fmap (:+ 0.0) (mkP depth))
        -- return $ FFTTestVal (tree2FuncComp $ mkTree depth values)
        return $ FFTTestVal $ O (pure (pure (depth :+ (0.0::Double))))

-- prop_fail = False  -- Test that QuickCheck is, actually, running.

prop_fft_test testVal = fft val == dft' val
  where types = testVal :: FFTTestVal Pair Pair
        val   = getVal testVal

-- Entry point for unit tests.
return []
runTests = $quickCheckAll

testMain :: IO ()
testMain = do
    -- allPass <- $quickCheckAll -- Run QuickCheck on all prop_ functions
    allPass <- runTests
    unless allPass exitFailure
    putStrLn "Success!"

-- This is a clunky, but portable, way to use the same Main module file
-- for both an application and for unit tests.
-- MAIN_FUNCTION is preprocessor macro set to exeMain or testMain.
-- That way we can use the same file for both an application and for tests.
-- #ifndef MAIN_FUNCTION
-- #define MAIN_FUNCTION exeMain
-- #endif
main :: IO ()
-- main = MAIN_FUNCTION
main = testMain


{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, UndecidableInstances, GADTs, TypeOperators, FlexibleContexts #-}

module Physics.Quantum.Algebra where

import Numeric.Field.Class          (Field)
import Numeric.Algebra.Class        (LeftModule((.*)), Multiplicative((*)), Monoidal(zero))
import Numeric.Additive.Class       (Additive((+)))
import Numeric.Algebra.Unital       (Unital(one))

import Prelude hiding ((+), (*), Complex)

import qualified Data.Map as M
import qualified Data.List as L
import Data.Maybe (catMaybes)

import Data.Complex                 (Complex, conjugate)

import qualified Data.Array.Repa as R
import qualified Data.Array.Repa.Eval as R
import qualified Data.Array.Repa.Repr.Unboxed as R
import qualified Data.Array.Repa.Repr.Vector as R

import Data.Functor.Identity (runIdentity)

import qualified Data.Eigen.SparseMatrix as ES
import qualified Data.Eigen.Matrix as E

-- a vector space m over field k
class (Field k, Additive m) => VectorSpace k m | m -> k where
    (*^) :: k -> m -> m

{-
-- lhs type ‘m’ does not determine rhs type ‘k’
instance (Field k, LeftModule k m) => VectorSpace k m where
    (*^) = (.*)
-}

class VectorSpace k m => LieAlgebra k m | m -> k where
    lieBr :: m -> m -> m

class (VectorSpace k a, VectorSpace k b) => DualVectors k a b | a b -> k where
    dual :: a -> b

class VectorSpace k m => InnerProductSpace k m | m -> k where
    (<.>) :: m -> m -> k

-- Endo-morphism of vector space
-- Data structure for operator may be different for different use cases
-- A Map be suitable for sparse operator, matrix for discreet vector indices, and functions for continuous vector indices.
class VectorSpace k v => SimpleOperator k v o | o -> k, o -> v where
    op :: o -> v -> v

-- Once we have simple operator, we can extract the map information from that
-- and get addition and multiplication operations
data OperatorMap v where
    OperatorMap :: VectorSpace k v => (v -> v) -> OperatorMap v

unOperatorMap :: OperatorMap v -> (v -> v)
unOperatorMap (OperatorMap m) = m

extractOperatorMap :: SimpleOperator k v o => o -> OperatorMap v
extractOperatorMap = OperatorMap . op

instance Additive v => Additive (OperatorMap v) where
    (+) (OperatorMap operator1) (OperatorMap operator2) = OperatorMap $ \vec -> (operator1 vec) + (operator2 vec)

instance  Multiplicative (OperatorMap o) where
    (*) (OperatorMap operator1) (OperatorMap operator2) = OperatorMap $ \vec -> operator1 (operator2 vec)

class (DualVectors k cv v, SimpleOperator k v o) => Operator k v cv o | o -> k where
    (|><|) :: v -> cv -> o

(<|>) :: (InnerProductSpace k v, DualVectors k cv v) => cv -> v -> k
(<|>) covec vec = (dual covec) <.> vec

(|>) :: SimpleOperator k v o => o -> v -> v
(|>) = op

(<|) :: (InnerProductSpace k v, DualVectors k cv v, Operator k v cv o) =>  cv -> o -> (v -> k)
(<|) covector operator  = \vector -> covector <|> (operator |> vector)

class IsComplex a where
    complexConjugate :: a -> a

instance Num a => IsComplex (Complex a) where
    complexConjugate = conjugate

----------------------------------
-- Constructions of vector spaces

-- m is a set of vector indices indices
-- May or may not represent orthogonal vectors, write instance of InnerProductSpace separately
-- kets are represented as maps from indices to coefficients
-- sparse kets have finite number of non zero entries
data SparseKet k m where
    SparseKet :: (Field k, Ord m) => M.Map m k -> SparseKet k m

unSparseKet :: SparseKet k m -> M.Map m k
unSparseKet (SparseKet ketMap) = ketMap

data SparseBra k m where
     SparseBra :: (Field k, Ord m) => M.Map m k -> SparseBra k m

unSparseBra :: SparseBra k m -> M.Map m k
unSparseBra (SparseBra braMap) = braMap

instance (Show k, Show m) => Show (SparseKet k m) where
    show ket@(SparseKet map)  = show map

instance (Show k, Show m) => Show (SparseBra k m) where
    show ket@(SparseBra map)  = show map

instance (Ord m, Field k) => Additive (SparseKet k m) where
    (+) (SparseKet a) (SparseKet b) = SparseKet $ M.unionWith (+) a b

instance (Ord m, Field k) => Additive (SparseBra k m) where
    (+) (SparseBra a) (SparseBra b) = SparseBra $ M.unionWith (+) a b

instance (Ord m, Field k) => VectorSpace k (SparseKet k m) where
    (*^) s (SparseKet v) = SparseKet $ M.map (s *) v

instance (Field k, Ord m) => VectorSpace k (SparseBra k m) where
    (*^) s (SparseBra v) = SparseBra $ M.map (s *) v

instance (Field k, IsComplex k, Ord m) => DualVectors k (SparseKet k m) (SparseBra k m) where
    dual (SparseKet a) = SparseBra $ M.map complexConjugate a

instance (Field k, IsComplex k, Ord m) => DualVectors k (SparseBra k m) (SparseKet k m) where
    dual (SparseBra a) = SparseKet $ M.map complexConjugate a

-- very generic mappings
data InfKet k m = InfKet { unContKet :: (m -> k) }
data InfBra k m = InfBra { unContBra :: (m -> k) }

instance Field k => Additive (InfKet k m) where
    (+) (InfKet a) (InfKet b) = InfKet $ \x -> a x + b x

instance Field k => Additive (InfBra k m) where
    (+) (InfBra a) (InfBra b) = InfBra $ \x -> a x + b x

instance Field k => VectorSpace k (InfKet k m) where
    (*^) s (InfKet v) = InfKet $ \x -> s * (v x)

instance Field k => VectorSpace k (InfBra k m) where
    (*^) s (InfBra v) = InfBra $ \x -> s * (v x)

instance (Field k, IsComplex k) => DualVectors k (InfKet k m) (InfBra k m) where
    dual (InfKet a) = InfBra $ \x -> complexConjugate $ a x

instance (Field k, IsComplex k) => DualVectors k (InfBra k m) (InfKet k m) where
    dual (InfBra a) = InfKet $ \x -> complexConjugate $ a x

-- a vector of coefficients at lattice points
data LatticeKet k m = LatticeKet { unLatticeKet :: R.Array R.D R.DIM1 k }
data LatticeBra k m = LatticeBra { unLatticeBra :: R.Array R.D R.DIM1 k }

instance (Show k, Show m) => Show (LatticeKet k m) where
    -- Delayed representation doesn't have show instance
    show (LatticeKet vec) = show $! compute vec
        where compute :: Show k => R.Array R.D R.DIM1 k -> R.Array R.V R.DIM1 k
              compute = R.computeS

instance (Show k, Show m) => Show (LatticeBra k m) where
    -- Delayed representation doesn't have show instance
    show (LatticeBra vec) = show $! compute vec
        where compute :: Show k => R.Array R.D R.DIM1 k -> R.Array R.V R.DIM1 k
              compute = R.computeS

instance Field k => Additive (LatticeKet k m) where
   (+) (LatticeKet a) (LatticeKet b) = LatticeKet $! R.zipWith (+) a b

instance Field k => Additive (LatticeBra k m) where
   (+) (LatticeBra a) (LatticeBra b) = LatticeBra $! R.zipWith (+) a b

instance Field k => VectorSpace k (LatticeKet k m) where
    (*^) s (LatticeKet v) = LatticeKet $! R.map (s *) v

instance Field k => VectorSpace k (LatticeBra k m) where
    (*^) s (LatticeBra v) = LatticeBra $! R.map (s *) v

instance (Field k, IsComplex k) => DualVectors k (LatticeKet k m) (LatticeBra k m) where
    dual (LatticeKet a) = LatticeBra $! R.map complexConjugate a

instance (Field k, IsComplex k) => DualVectors k (LatticeBra k m) (LatticeKet k m) where
    dual (LatticeBra a) = LatticeKet $! R.map complexConjugate a

-- n x 1 matrix
data SparseMatKet k ck m where
    SparseMatKet :: E.Elem k ck => ES.SparseMatrix k ck -> SparseMatKet k ck m

unSparseMatKet :: SparseMatKet k ck m -> ES.SparseMatrix k ck
unSparseMatKet (SparseMatKet mat) = mat

data SparseMatBra k ck m where
    SparseMatBra :: E.Elem k ck => ES.SparseMatrix k ck -> SparseMatBra k ck m

unSparseMatBra :: SparseMatBra k ck m -> ES.SparseMatrix k ck
unSparseMatBra (SparseMatBra mat) = mat

instance (Show k) => Show (SparseMatKet k ck m) where
    show (SparseMatKet mat) = show mat

instance (Show k) => Show (SparseMatBra k ck m) where
    show (SparseMatBra mat) = show mat

instance Additive (SparseMatKet k ck m) where
    (+) (SparseMatKet a) (SparseMatKet b) = SparseMatKet $! ES.add a b

instance Additive (SparseMatBra k ck m) where
    (+) (SparseMatBra a) (SparseMatBra b) = SparseMatBra $! ES.add a b

instance Field k => VectorSpace k (SparseMatKet k ck m) where
    (*^) s (SparseMatKet v) = SparseMatKet $! ES.scale s v

instance Field k => VectorSpace k (SparseMatBra k ck m) where
    (*^) s (SparseMatBra v) = SparseMatBra $! ES.scale s v

-------------------------------------
-- Constructions of operators

-- Operator formed from outerproduct of two vectors
data OuterOperator k v cv where
    OuterOperator :: (InnerProductSpace k v, DualVectors k cv v) => [(v, cv)] -> OuterOperator k v cv

unOuterOperator :: OuterOperator k v cv -> [(v, cv)]
unOuterOperator (OuterOperator outerProducts) = outerProducts

instance ( InnerProductSpace k v
         , DualVectors k cv v
         )
         => SimpleOperator k v (OuterOperator k v cv) where
    op (OuterOperator outerProducts) vec = L.foldr1 (+) $ map (\(ovec,ocovec) -> (ocovec <|> vec) *^ ovec) $ outerProducts

instance ( InnerProductSpace k v
         , DualVectors k cv v
         )
         => Operator k v cv (OuterOperator k v cv) where
    (|><|) vec covec = OuterOperator [(vec, covec)]

instance Additive (OuterOperator k v cv) where
    (+) (OuterOperator o1) (OuterOperator o2) = OuterOperator $ o1 L.++ o2

-- a sparse operator represented as a map from (i,j) to coefficient of |i><j|
-- Note that (|i><j|) |k> = <j|k> |i> and therefore the operation requires inner product to be defined
data SparseOperator k m = SparseOperator { unSparseOperator :: M.Map (m, m) k }

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => SimpleOperator k (SparseKet k m) (SparseOperator k m) where
    op (SparseOperator opMap) ket = SparseKet $! M.fromListWith (+) $!
        map (\(veci, coveci) -> ( veci
                                , (SparseBra $! M.fromListWith (+) $! catMaybes $! [fmap ((,) coveci) $! M.lookup (veci, coveci) opMap]) <|> ket
                                )
            ) $!
            (M.keys opMap)

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => Operator k (SparseKet k m) (SparseBra k m) (SparseOperator k m) where
    (|><|) (SparseKet vecMap) (SparseBra covecMap) = SparseOperator $! M.fromListWith (+) $!
        catMaybes [(,) (i,j) <$> (fmap (*) (M.lookup i vecMap) <*> (M.lookup j covecMap)) | i <- M.keys vecMap, j <- M.keys covecMap ]

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => Additive (SparseOperator k m) where
    (+) (SparseOperator a) (SparseOperator b) = SparseOperator $! M.unionWith (+) a b

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => Multiplicative (SparseOperator k m) where
     (*) (SparseOperator a) (SparseOperator b) = SparseOperator $! M.fromListWith (+) $! multiplyMaps a b []
         where
             -- chaoose one index i and sum (over j and k) |j><i| |i><k|
             multiplyMaps :: ( Field k, Ord m,
                               InnerProductSpace k (SparseKet k m),
                               DualVectors k (SparseBra k m) (SparseKet k m)
                             )
                             => M.Map (m,m) k -> M.Map (m,m) k -> [((m,m), k)] -> [((m,m), k)]
             multiplyMaps mapA mapB oldResult = case null mapA || null mapB of
                 True -> oldResult
                 False ->
                   -- toList uses list fusion so this should not go through entire map
                   let oneElement = L.take 1 $ M.toList $ mapA
                   in case oneElement of
                          [] -> oldResult
                          elementToUse : _ -> let indexToUse = (snd $! fst elementToUse)

                                                  -- <i|i>
                                                  indexKet = SparseKet $! M.singleton indexToUse one
                                                  indexSquaredNorm =  indexKet <.> indexKet

                                                  -- take all keys from mapA which have this index in snd of key
                                                  (mapAToUse, mapAToPass) = M.partitionWithKey (\key _ -> snd key == indexToUse) mapA
                                                  -- take all keys from mapA which have this index in snd of key
                                                  (mapBToUse, mapBToPass) = M.partitionWithKey (\key _ -> fst key == indexToUse) mapB

                                                  -- we can discard second index from mapAToUse and first index from mapBToUse
                                                  -- we also know that the preserved indices will occur uniquely in index
                                                  -- so we can just take the cross product
                                                  listAToUse = M.toList $! mapAToUse
                                                  listBToUse = M.toList $! mapBToUse
                                                  newResult = [((a,b), va * vb * indexSquaredNorm) | ((a,_), va) <- listAToUse,
                                                                                ((_,b), vb) <- listBToUse
                                                                                ]

                                              in multiplyMaps mapAToPass mapBToPass $! newResult ++ oldResult

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => LeftModule k (SparseOperator k m) where
    (.*) scalar (SparseOperator b) = SparseOperator $! M.map (scalar *) b

-- a more general operator represented as a map from j to Map i_n s_n where operator takes |i_n> to s_n |j>
-- this construct may not suffice when there are uncountable such i_n
data InfOperator k m = InfOperator { unInfOperator :: m -> M.Map m k }

instance (Field k) => SimpleOperator k (InfKet k m) (InfOperator k m) where
    op operator = \(InfKet coeffMap) -> InfKet $ \index -> foldr1 (+) $
        map ( \(invIndex, factor) -> factor * coeffMap invIndex ) $
            M.toList $ unInfOperator operator $ index

instance (Field k, Ord m) => Additive (InfOperator k m) where
    (+) (InfOperator op1) (InfOperator op2) = InfOperator $ \index -> M.unionWith (+) (op1 index) (op2 index)

instance (Field k, Ord m) => Multiplicative (InfOperator k m) where
    (*) (InfOperator op1) (InfOperator op2) = InfOperator $ \index ->
            foldr1 (M.unionWith (+)) $
                map (\(i,s) -> M.map (s *) $ op2 i) $
                    M.toList $ op1 index

-- A matrix to transfom LatticeKet vector
data LatticeOperator k m where
    LatticeOperator :: (Field k, R.Elt k, R.Unbox k, Num k) => R.Array R.D R.DIM2 k -> LatticeOperator k m

unLatticeOperator :: LatticeOperator k m -> R.Array R.D R.DIM2 k
unLatticeOperator (LatticeOperator mat) = mat

-- Num k is required because of sumP in method definition
instance Field k => SimpleOperator k (LatticeKet k m) (LatticeOperator k m) where
    op (LatticeOperator matrix) (LatticeKet vector) = LatticeKet $! R.delay $! runIdentity $! mvMultP matrix vector

instance Additive (LatticeOperator k m) where
    (+) (LatticeOperator op1) (LatticeOperator op2) = LatticeOperator $! R.zipWith (+) op1 op2

-- Num k is required because of sumP in method definition
instance Multiplicative (LatticeOperator k m) where
    (*) (LatticeOperator op1) (LatticeOperator op2) = LatticeOperator $! R.delay $! runIdentity $! mmMultP op1 op2

instance Field k => LeftModule k (LatticeOperator k m) where
    (.*) scalar (LatticeOperator matrix) = LatticeOperator $! R.map (scalar *) matrix

-- Taken from Repa wiki, extended it to Additive and Multiplicative structures
mmMultP :: (Monad m, R.Source r k, Additive k, Multiplicative k, R.Elt k, R.Unbox k, Num k)
        => R.Array r R.DIM2 k
        -> R.Array r R.DIM2 k
        -> m (R.Array R.U R.DIM2 k)
mmMultP a b = R.sumP (R.zipWith (*) aRepl bRepl)
    where
      -- Note: transpose transposes the lowest two dimensions
      t = R.transpose b
      aRepl = R.extend (R.Z R.:.R.All R.:.colsB R.:.R.All) a
      bRepl = R.extend (R.Z R.:.rowsA R.:.R.All R.:.R.All) t

      (R.Z R.:.colsA R.:.rowsA) = R.extent a
      (R.Z R.:.colsB R.:.rowsB) = R.extent b

-- Multiply a martix with a vector
-- based on the above matrix matrix multiplication
mvMultP :: (Monad m, R.Source r k, Additive k, Multiplicative k, R.Elt k, R.Unbox k, Num k)
        => R.Array r R.DIM2 k
        -> R.Array r R.DIM1 k
        -> m (R.Array R.U R.DIM1 k)
mvMultP a b = R.sumP (R.zipWith (*) aRepl bRepl)
    where
      aRepl = a
      bRepl = R.extend (R.Z R.:.rowsA R.:.R.All) b
      (R.Z R.:.rowsA R.:.colsA) = R.extent a

-- A matrix to transfom LatticeKet vector
data SparseMatOperator k ck m where
    SparseMatOperator :: E.Elem k ck => ES.SparseMatrix k ck -> SparseMatOperator k ck m

unSparseMatOperator :: SparseMatOperator k ck m -> ES.SparseMatrix k ck
unSparseMatOperator (SparseMatOperator mat) = mat

-- Num k is required because of sumP in method definition
instance Field k => SimpleOperator k (SparseMatKet k ck m) (SparseMatOperator k ck m) where
    op (SparseMatOperator matrix) (SparseMatKet vector) = SparseMatKet $! ES.mul matrix vector

instance Field k => Additive (SparseMatOperator k ck m) where
    (+) (SparseMatOperator op1) (SparseMatOperator op2) = SparseMatOperator $! ES.add op1 op2

instance Field k => Multiplicative (SparseMatOperator k ck m) where
    (*) (SparseMatOperator op1) (SparseMatOperator op2) = SparseMatOperator $! ES.mul op1 op2

instance Field k => LeftModule k (SparseMatOperator k ck m) where
    (.*) scalar (SparseMatOperator matrix) = SparseMatOperator $! ES.scale scalar matrix


-- Filename: fft.hs
-- Created by: Daniel Winograd-Cort
-- Created on: unknown
-- Last Modified by: Daniel Winograd-Cort
-- Last Modified on: 12/12/2013

-- -- DESCRIPTION --
-- This code was inspired by Euterpea.  It uses UISF to present a GUI that
-- shows the sum of two waves (whose frequencies are specified by the user)
-- as well as the Fast Fourier Transform of that sum.
--
-- This module requires the array and pure-fft packages.

{-# LANGUAGE Arrows #-}
module FRP.UISF.Examples.FFT where
import FRP.UISF
import Numeric.FFT (fft)
import Data.Complex
import Data.Map (Map)
import Data.Maybe (listToMaybe, catMaybes)
import qualified Data.Map as Map

import Data.Array.Unboxed

-- | Alternative for working with Math.FFT instead of Numeric.FFT
--import qualified Math.FFT as FFT
--import Data.Array.IArray
--import Data.Array.CArray
--myFFT n lst = elems $ (FFT.dft) (listArray (0, n-1) lst)

--------------------------------------
-- Sine wave oscillators
--------------------------------------

-- Table definition
type Table = UArray Int Double

-- Sine table generator. Takes an integer representing the number of samples to generate
-- and a list of relative intensities for the overtones of the sine wave.

tableSinesN :: Int -> [Double] -> Table
tableSinesN size amps =
    let wave x   = sum (zipWith (*) [sin (2*pi*x*n) | n <- [1..]] amps)
        delta    = 1 / fromIntegral size
        waveform = take size $ map wave [0,delta..]
        divisor  = (maximum . map abs) waveform
     in listArray (0,size) (map (/divisor) waveform)

-- Two example sine tables.

tab1, tab2 :: Table
tab1 = tableSinesN 4096 [1] -- Basic sine wave
tab2 = tableSinesN 4096 [1.0,0.5,0.33]

-- Table-driven oscillator

osc :: ArrowCircuit a => Table -> Double -> a Double Double
osc table sr = proc freq -> do
    rec
      let delta = 1 / sr * freq
          phase = if next > 1 then frac next else next
      next <- delay 0 -< frac (phase + delta)
    returnA -< ((table!).(`mod` size).round.(*rate)) phase
  where (_,size)    = bounds table
        rate        = fromIntegral size
        frac x = if x > 1 then x - fromIntegral (truncate x) else x

--------------------------------------
-- Fast Fourier Transform
--------------------------------------

-- | Returns n samples of type b from the input stream at a time,
--   updating after k samples.  This function is good for chunking
--   data and is a critical component to fftA
quantize :: ArrowCircuit a => Int -> Int -> a b (SEvent [b])
quantize n k = proc d -> do
    rec (ds,c) <- delay ([],0) -< (take n (d:ds), c+1)
    returnA -< if c >= n && c `mod` k == 0 then Just ds else Nothing

-- | Converts the vector result of a dft into a map from frequency to magnitude.
--   One common use is:
--      fftA >>> arr (fmap $ presentFFT clockRate)
presentFFT :: Double -> [Double] -> Map Double Double
presentFFT clockRate a = Map.fromList $ zipWith (curry mkAssoc) [0..] a where
    mkAssoc (i,c) = (freq, c) where
        samplesPerPeriod = fromIntegral (length a)
        freq = fromIntegral i * (clockRate / samplesPerPeriod)

-- | Given a quantization frequency (the number of samples between each
--   successive FFT calculation) and a fundamental period, this will decompose
--   the input signal into its constituent frequencies.
--   NOTE: The fundamental period must be a power of two!
fftA :: ArrowCircuit a => Int -> Int -> a Double (SEvent [Double])
fftA qf fp = proc d -> do
    carray <- quantize fp qf -< d :+ 0
    returnA -< fmap (map magnitude . take (fp `div` 2) . fft) carray

--------------------------------------
-- UISF Example
--------------------------------------

-- This example shows off the histogram and realtimeGraph widgets by
-- summing two sin waves and displaying them.  Additionally, it makes
-- use of two horizontal sliders.
-- This example also shows off asyncUISFV and how to take a SigFun,
-- of the type used to create sound, and convert it to a UISF.
fftEx :: UISF () ()
fftEx = proc _ -> do
    f1 <- hSlider (1, 2000) 440 -< ()
    _ <- leftRight (label "Freq 1: " >>> display) -< f1
    f2 <- hSlider (1, 2000) 440 -< ()
    _ <- leftRight (label "Freq 2: " >>> display) -< f2
    d <- asyncVT sr 0.1 myAutomaton -< (f1, f2)
    let fft = listToMaybe $ catMaybes $ map (snd . fst) d
        s = map (\((s, _), t) -> (s,t)) d
    _ <- histogram (makeLayout (Stretchy 10) (Fixed 150)) -< fft
    _ <- realtimeGraph (makeLayout (Stretchy 10) (Fixed 150)) 2 Black -< s
    returnA -< ()
  where
    sr = 1000 -- signal rate
    myAutomaton = proc (f1, f2) -> do
        s1 <- osc tab1 sr -< f1
        s2 <- osc tab2 sr -< f2
        let s = (s1 + s2)/2
        fftData <- fftA 100 256 -< s
        returnA -< (s, fftData)

-- This test is run separately from the others.
main :: IO ()
main = runUI (defaultUIParams {uiSize=(500, 600), uiTitle="FFT Example"}) fftEx


{-# LANGUAGE NoMonomorphismRestriction, ImpredicativeTypes ,RankNTypes, TypeOperators #-}
import Kalman
import qualified Data.Vector as V
import Vectorization.Instances
import Local
import Linear.V1
import Data.Distributive
import qualified Data.Foldable as F
import Linear.V2
import Data.Complex
import CartesianProduct
import System.Random.MWC.Distributions
import Vectorization
import Foreign.Storable
import Foreign.Cholesky
import Space.Class
import Control.Applicative
import Control.Monad
import System.Random.MWC
import Graphics.Gnuplot.Simple
import Statistics.Autocorrelation
import Statistics.Sample
import Data.Functor.Compose
import Foreign.FStorable
import Foreign.FStorable

type Model = V1 :|: V1 :|: V1 :|: V1

predSin :: V1 Double -> Model Double -> Model Double
predSin  dt s@(theta :|: omega :|: amp :|: accel)   = s |+|  ((omega*dt )  :|: 0 :|: (dt*accel ) :|: 0 )

projSin :: Model Double -> V1 Double
projSin (theta :|: omega :|: amp :|: accel) = amp * (fmap sin theta )

nx :: Model (Model Double)
nx = diag $ fmap (*1e-3)((1/3) :|: 1 :|: 0.1 :|: 0.1)

getMean :: Functor f => (forall b . g b -> V1 b) -> f (g b , g (g b)) -> f b
getMean s = fmap ( unV1 . s . fst )

getCov :: (Num b ,Functor f )=> (forall b .  g b -> V1 b) -> f (g b , g (g b)) -> f b
getCov s = fmap ( abs .unV1 . s .  unV1 . s . snd )

sigmaBound :: Num b => (forall b . g b -> V1 b) -> [(g b , g (g b))] -> [[b]]
sigmaBound s f =  [zipWith (+) (getMean s f) (getCov s f), getMean s  f,zipWith (-) (getMean s f) (getCov s f)]

sigmas :: Num b => (forall b . g b -> V1 b) -> [[(g b , g (g b))]] -> [[b]]
sigmas s = concat . map (sigmaBound s)

allSigmas :: Num b => [forall b . g b -> V1 b] -> [[(g b , g (g b))]] -> [[[b]]]
allSigmas s d =  map (flip sigmas $ d) s

unV1 (V1 x) = x
execute = do
   seed <- create
   let n = 500
   v <- replicateM n (normal 0 0.25 seed)
   let noisedData = testeSet v
       cleanData = testeSet (replicate n 0)
   let ny :: (V1 (V1 Double))
       ny = 1
       xo :: Model Double
       xo = 0.1 :|: 0.1 :|: 1 :|: 1e-3
       alpha = 1e-2
       iteration  r = full (predSin 1) nx projSin ny r
       backwardIteration = smoothing (predSin 1 ) nx
       forwards =   scanl (flip iteration)  (xo,nx) (fmap V1 $ testeSet v)
       backwards =  reverse $ scanl1 (flip backwardIteration)  $ reverse forwards
       project = fmap (unV1 . projSin . fst )
       getAmp = fmap ( unV1 . amp . fst )
       getAmpCov = fmap ( unV1 .  amp . unV1 . amp . snd)
       get x = fmap ( unV1 . x . fst )
   -- print projections
   plotLists [] [testeSet (replicate n 0) ,testeSet v,project backwards  ]
   -- print all parameters
   mapM_ (plotLists []) (allSigmas [amp ,accel,omega] [backwards])
   let ampCov = V.fromList noiseSignalBackward
       noiseSignalBackward = zipWith (-)  cleanData  (project backwards)
       noiseSignalForward = zipWith (-)  cleanData (project forwards)
   -- Print residuals
   plotLists [] [noiseSignalForward , noiseSignalBackward ,v]
   print $ mean ampCov
   print $ stdDev ampCov
   print $ mean $ V.fromList $  zipWith (-) (getAmp forwards) (getAmp backwards)
   print $ stdDev $ V.fromList $  zipWith (-) (getAmp forwards) (getAmp backwards)

amp :: Model a -> V1 a
amp ( x :|: _ :|: amp :|: _ ) = amp
accel ( x :|: _ :|: amp :|: accel ) = accel
omega ( x :|: omega  :|: amp :|: accel ) = omega
theta ( theta :|: _ :|: amp :|: accel ) = theta

matI :: Num a => Int -> [[a]]
matI n = [ [fromIntegral $ fromEnum $ i == j | i <- [1..n]] | j <- [1..n]]

diag = liftIdent
liftIdent ::(R f,Num a,Storable a,Vectorize f,Functor f,Applicative f)=> f a -> f ( f a)
liftIdent x = fmap (liftA2 (*) x) (fromLists $ matI n )
    where n = dimM x

testeSet v =  zipWith (+)  (map (function.fromIntegral) [1..length v]) v where
  function x =  (1+x/500) * sin(16 * x/500 * pi)


{-# LANGUAGE UnicodeSyntax #-}
{-# LANGUAGE BangPatterns #-}

{-|
Module      : Statistics.FastBayes.Linear
Description : Bayesian linear regression via maximum marginal likelihood.
Copyright   : (c) Melinae, 2014
                  Chad Scherrer, 2014
License     : MIT
Maintainer  : chad.scherrer@gmail.com
Stability   : experimental
Portability : POSIX

This module gives an implementation of Bayesian linear regression, with the scale of the prior chosen by marginal likelihood.

The inputs for a Bayesian linear model are identical to those of a classical linear model, except that in addition to a design matrix and response, we must also specify a prior distribution on the weights and the noise. This leaves us with an open question of how these should be specified.

In his book /Pattern Recognition and Machine Learning/, Christopher Bishop provides details for an approach that simplifies the situation significantly, and allows for much faster inference. The structure of the linear model allows us to integrate the posterior over the weights, resulting in the /marginal likelihood/, expressed as a function of the prior precision and noise precision. This, in turn, can be easily optimized.
-}

module Statistics.FastBayes.Linear
  ( test
  , Fit
  , fit0
  , fit1
  , priorPrecision
  , noisePrecision
  , effectiveNumParameters
  , logEvidence
  , mapWeights
  , intercept
  , hessian
  , features
  , Predict
  , predict
  , studentize
  , studentizeM
  , inputs
  , fitted
  , fittedSD
  )
  where

import qualified Data.Vector.Storable as V
import Data.Packed.Matrix
import Numeric.LinearAlgebra

data Fit a = Fit
  { priorPrecision         :: Double             -- ^The precision (inverse variance) of the prior distribution, determined by maximizing the marginal likelihood
  , noisePrecision         :: Double             -- ^The precision (inverse variance) of the noise
  , effectiveNumParameters :: Double             -- ^The effective number of parameters in the model
  , logEvidence            :: Double             -- ^The log of the evidence, which is useful for model comparison (different features, same response)
  , mapWeights             :: Vector Double      -- ^The MAP (maximum a posteriori) values for the parameter weights
  , intercept              :: Double             -- ^The y-intercept
  , hessian                :: Matrix Double      -- ^The Hessian (matrix of second derivatives) for the posterior distribution
  , features               :: a → Vector Double  -- ^Extract features from a data point
  }

-- instance Show (Fit a) where

-- |@fit0 lim x y@ fits a Bayesian linear model to a design matrix @x@ and response vector @y@. This is an iterative algorithm, resulting in a sequence (list) of (α,β) values. Here α is the prior precision, and β is the noise precision. The @lim@ function passed in is used to specify how the limit of this sequence should be computed.
fit0 ::
  ([(Double, Double)] → (Double, Double)) -- ^How to take the limit of the (α,β) sequence. A simple approach is, /e.g./, @fit (!! 1000) x y@
  → Bool                                  -- ^Whether to studentize (translate and re-scale) the features
  → (a → Vector Double)                   -- ^The features
  → [a]                                   -- ^The input data
  → Vector Double                         -- ^The response vector
  → Fit a
fit0 lim student f xs y = Fit α β γ logEv m 0 h f'
  where
  (x, f') | not student = (fromRows $ map f xs,      f)
          | student     = (                 x0, fs . f)
            where
            (x0, fs) = studentizeM . fromRows . map f $ xs
  n = rows x
  p = cols x
  α0 = 1.0
  β0 = 1.0

  -- A vector of the eigenvalues of xtx
  (_,sqrtEigs,_) = compactSVD x
  eigs = V.map square sqrtEigs

  xtx = trans x <> x
  xty = trans x <> y
  getHessian a b = diag (V.replicate p a) + scale b xtx

  m = scale β $ h <\> xty

  go :: Double → Double → [(Double, Double)]
  go a0 b0 = (a0, b0) : go a b
    where
    h0 = getHessian a0 b0
    m0 = scale b0 $ h0 <\> xty
    c = V.sum $ V.map (\x' → x' / (a0 + x')) eigs
    a = c / (m0 <.> m0)
    b = recip $ (normSq $ y - x <> m0) / (fromIntegral n - c)

  γ = V.sum $ V.map (\x' → x' / (α + x')) eigs

  h = getHessian α β
  logEv = 0.5 *
    ( fromIntegral p * log α
    + fromIntegral n * log β
    - (β * normSq (y - x <> m) + α * (m <.> m))
    - logDetH
    - fromIntegral n * log (2*pi)
    )
    where
    (_,(logDetH, _)) = invlndet h

  (α, β) = lim $ go α0 β0

fit1 ::
  ([(Double, Double)] → (Double, Double)) -- ^How to take the limit of the (α,β) sequence. A simple approach is, /e.g./, @fit (!! 1000) x y@
  → Bool                                  -- ^Whether to studentize (translate and re-scale) the features
  → (a → Vector Double)                   -- ^The features
  → [a]                                   -- ^The input data
  → Vector Double                         -- ^The response vector
  → Fit a
fit1 lim student f xs y = f0 {intercept = μ}
  where
  μ = mean y
  n = V.length y
  f0 = fit0 lim student f xs (y - konst μ n)

square :: Double → Double
square x = x * x

normSq :: Vector Double → Double
normSq x = x <.> x

mean :: Vector Double -> Double
mean v = V.sum v / fromIntegral (V.length v)

studentizeM :: Matrix Double -> (Matrix Double, Vector Double → Vector Double)
studentizeM m = (x, f)
  where
  (xs, fs) = unzip . map studentize . toColumns $ m
  x = fromColumns xs
  f = fromList . zipWith ($) fs . toList

studentize :: Vector Double -> (Vector Double, Double → Double)
studentize v = (scale (1/σ) v0, f)
  where
  n   = V.length v
  μ   = mean v
  v0  = v - konst μ n
  σ   = sqrt $ normSq v0 / fromIntegral n
  f x = (x - μ) / σ

test = myFit 11
  where
  x = [1,1.1..pi]
  y = fromList $ map sin x
  f n x = fromList . map (x **) $ [1..n]
  myFit n = fit1 (!!100) True (f n) x y

data Predict a = Predict
  { inputs          :: [a]
  , fitted          :: [Double]
  , fittedSD        :: [Double]
  }

predict :: Fit a -> [a] -> Predict a
predict myFit xs = Predict xs ys (toList . V.map sqrt $ konst (1/β) n + takeDiag cov)
  where
  n = rows x
  α  = priorPrecision myFit
  β  = noisePrecision myFit
  w  = mapWeights     myFit
  y0 = intercept      myFit
  h  = hessian        myFit
  f  = features       myFit
  x  = fromRows $ map f xs
  cov = x <> (h <\> trans x)
  y = x <> w + konst y0 n
  ys = toList y


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Binomial
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The binomial distribution.  This is the discrete probability
-- distribution of the number of successes in a sequence of /n/
-- independent yes\/no experiments, each of which yields success with
-- probability /p/.

module Statistics.Distribution.Binomial
    (
      BinomialDistribution
    -- * Constructors
    , binomial
    , binomialE
    -- * Accessors
    , bdTrials
    , bdProbability
    ) where

import Control.Applicative
import Data.Aeson            (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary           (Binary(..))
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.SpecFunctions           (choose,logChoose,incompleteBeta,log1p)
import Numeric.MathFunctions.Constants (m_epsilon)

import qualified Statistics.Distribution as D
import qualified Statistics.Distribution.Poisson.Internal as I
import Statistics.Internal

-- | The binomial distribution.
data BinomialDistribution = BD {
      bdTrials      :: {-# UNPACK #-} !Int
    -- ^ Number of trials.
    , bdProbability :: {-# UNPACK #-} !Double
    -- ^ Probability.
    } deriving (Eq, Typeable, Data, Generic)

instance Show BinomialDistribution where
  showsPrec i (BD n p) = defaultShow2 "binomial" n p i
instance Read BinomialDistribution where
  readPrec = defaultReadPrecM2 "binomial" binomialE

instance ToJSON BinomialDistribution
instance FromJSON BinomialDistribution where
  parseJSON (Object v) = do
    n <- v .: "bdTrials"
    p <- v .: "bdProbability"
    maybe (fail $ errMsg n p) return $ binomialE n p
  parseJSON _ = empty

instance Binary BinomialDistribution where
  put (BD x y) = put x >> put y
  get = do
    n <- get
    p <- get
    maybe (fail $ errMsg n p) return $ binomialE n p

instance D.Distribution BinomialDistribution where
    cumulative = cumulative

instance D.DiscreteDistr BinomialDistribution where
    probability    = probability
    logProbability = logProbability

instance D.Mean BinomialDistribution where
    mean = mean

instance D.Variance BinomialDistribution where
    variance = variance

instance D.MaybeMean BinomialDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance BinomialDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy BinomialDistribution where
  entropy (BD n p)
    | n == 0 = 0
    | n <= 100 = directEntropy (BD n p)
    | otherwise = I.poissonEntropy (fromIntegral n * p)

instance D.MaybeEntropy BinomialDistribution where
  maybeEntropy = Just . D.entropy

-- This could be slow for big n
probability :: BinomialDistribution -> Int -> Double
probability (BD n p) k
  | k < 0 || k > n = 0
  | n == 0         = 1
    -- choose could overflow Double for n >= 1030 so we switch to
    -- log-domain to calculate probability
  | n < 1000       = choose n k * p^k * (1-p)^(n-k)
  | otherwise      = exp $ logChoose n k + log p * k' + log1p (-p) * nk'
  where
    k'  = fromIntegral k
    nk' = fromIntegral $ n - k

logProbability :: BinomialDistribution -> Int -> Double
logProbability (BD n p) k
  | k < 0 || k > n          = (-1)/0
  | n == 0                  = 0
  | otherwise               = logChoose n k + log p * k' + log1p (-p) * nk'
  where
    k'  = fromIntegral   k
    nk' = fromIntegral $ n - k

-- Summation from different sides required to reduce roundoff errors
cumulative :: BinomialDistribution -> Double -> Double
cumulative (BD n p) x
  | isNaN x      = error "Statistics.Distribution.Binomial.cumulative: NaN input"
  | isInfinite x = if x > 0 then 1 else 0
  | k <  0       = 0
  | k >= n       = 1
  | otherwise    = incompleteBeta (fromIntegral (n-k)) (fromIntegral (k+1)) (1 - p)
  where
    k = floor x

mean :: BinomialDistribution -> Double
mean (BD n p) = fromIntegral n * p

variance :: BinomialDistribution -> Double
variance (BD n p) = fromIntegral n * p * (1 - p)

directEntropy :: BinomialDistribution -> Double
directEntropy d@(BD n _) =
  negate . sum $
  takeWhile (< negate m_epsilon) $
  dropWhile (not . (< negate m_epsilon)) $
  [ let x = probability d k in x * log x | k <- [0..n]]

-- | Construct binomial distribution. Number of trials must be
--   non-negative and probability must be in [0,1] range
binomial :: Int                 -- ^ Number of trials.
         -> Double              -- ^ Probability.
         -> BinomialDistribution
binomial n p = maybe (error $ errMsg n p) id $ binomialE n p

-- | Construct binomial distribution. Number of trials must be
--   non-negative and probability must be in [0,1] range
binomialE :: Int                 -- ^ Number of trials.
          -> Double              -- ^ Probability.
          -> Maybe BinomialDistribution
binomialE n p
  | n < 0            = Nothing
  | p >= 0 || p <= 1 = Just (BD n p)
  | otherwise        = Nothing

errMsg :: Int -> Double -> String
errMsg n p
  = "Statistics.Distribution.Binomial.binomial: n=" ++ show n
  ++ " p=" ++ show p ++ "but n>=0 and p in [0,1]"


{-# LANGUAGE OverloadedLists#-}
{-# LANGUAGE RankNTypes#-}
{-# LANGUAGE ScopedTypeVariables #-}
--{-# LANGUAGE Strict#-}
module Lib where

import Control.Monad.ST
import Debug.Trace
import qualified Numeric.AD as AD
import Numeric.AD.Internal.Reverse
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import qualified Data.Vector.Storable as VS
import System.Random
import Debug.Trace
import Text.Printf
import Control.Monad
import Control.Concurrent
import Data.List (permutations)

someFunc :: IO ()
someFunc = putStrLn "someFunc"

data Neuron = Neuron {
                     w  :: Vector Double,
                     f  :: Double -> Double }

data MNeuron = MNeuron {
                       mw :: Matrix Double,
                       mf :: Double -> Double }

instance Show Neuron where
  show = show . w

instance Show MNeuron where
  show = show . mw

applyMNeuron :: MNeuron -> Vector Double -> Vector Double
applyMNeuron (MNeuron wagi funkcja) = VS.map funkcja . (wagi #>) . VS.cons 1

applyNeuron :: Neuron -> Vector Double -> Double
applyNeuron (Neuron w f) = f . (w <.>) . VS.cons 1

acFunc :: Double -> Double
acFunc x | x >= 0 = 1
          | otherwise = 0

notN = Neuron [1,-1.5] acFunc

andN = Neuron [-2, 1.5, 1.5] acFunc

nandN = Neuron [2, -1.5, -1.5] acFunc

orN = Neuron [-1, 1.5, 1.5] acFunc

-- 2. Picture classification

trainSet :: [(Vector Double, Double)]
trainSet = [([ 0, 0, 0, 0, 0
              , 0, 1, 1, 0, 0
              , 0, 0, 1, 0, 0
              , 0, 0, 1, 0, 0
              , 0, 0, 1, 0, 0], 1)

           , ([ 0, 0, 1, 1, 0
              , 0, 0, 0, 1, 0
              , 0, 0, 0, 1, 0
              , 0, 0, 0, 0, 0
              , 0, 0, 0, 0, 0], 1)

           , ([ 0, 0, 0, 0, 0
              , 1, 1, 0, 0, 0
              , 0, 1, 0, 0, 0
              , 0, 1, 0, 0, 0
              , 0, 1, 0, 0, 0], 1)

           , ([ 0, 0, 0, 0, 0
              , 0, 1, 1, 1, 0
              , 0, 1, 0, 1, 0
              , 0, 1, 1, 1, 0
              , 0, 0, 0, 0, 0], 0)

           , ([ 0, 0, 0, 0, 0
              , 0, 0, 0, 0, 0
              , 1, 1, 1, 0, 0
              , 1, 0, 1, 0, 0
              , 1, 1, 1, 0, 0], 0)]

test :: [Vector Double]
test = [[ 0, 1, 1, 0, 0
        , 0, 0, 1, 0, 0
        , 0, 0, 1, 0, 0
        , 0, 0, 0, 0, 0
        , 0, 0, 0, 0, 0
        ]
       ,[ 1, 1, 1, 0, 0
        , 1, 0, 1, 0, 0
        , 1, 0, 1, 0, 0
        , 1, 1, 1, 0, 0
        , 0, 0, 0, 0, 0
        ]
       ,[ 0, 0, 1, 0, 0
        , 0, 0, 1, 0, 0
        , 0, 0, 1, 0, 0
        , 0, 0, 1, 0, 0
        , 0, 0, 0, 0, 0
        ]
       ,[ 0, 0, 0, 0, 0
        , 0, 0, 0, 0, 0
        , 1, 1, 1, 0, 0
        , 1, 0, 1, 0, 0
        , 1, 1, 1, 0, 0
        ]]

train st perc trSet = go st 0 perc (cycle trSet) (length trSet)
  where
    go st iters n@(Neuron w f) ((inp,z):xs) c =
      if c <=0 then
        (iters, n)
      else
        go st (iters+1) (Neuron newW f) xs newC
      where
        y = applyNeuron n inp
        newC = if z == y then c - 1 else 5
        newW = w + VS.map (\x -> x * st * (z - y)) (VS.cons 1 inp)

perceptron :: Neuron
perceptron = Neuron (VS.fromList . replicate 26 $ 1) acFunc

-- 3. Picture recognition

prPrint :: Vector Double -> String
prPrint = morph (\x (acc, count) -> ((if x <= 0 then " " else "■") ++ (if count `mod` 5 == 0 then "\n" else []) ++ acc, count+1)) []
  where
    morph f a = fst . VS.foldr f (a,0)

readPic :: String -> Vector Double
readPic = VS.fromList . foldr (\x acc -> case x of
                                  ' ' -> (-1):acc
                                  '*' -> 1:acc
                                  '\n'-> acc) []
pics :: [String]
pics = [ "     \n *** \n * * \n *** \n     "
       , "     \n **  \n  *  \n  *  \n     "]

sgn :: Double -> Double
sgn x
     | x >= 0 = 1
     | x < 0 = -1

picTaker :: MNeuron
picTaker = MNeuron weights sgn
  where
    [v1, v2] = map readPic pics
    w1       = cmap (/25) (v1 `outer` v1)
    w2       = cmap (/25) (v2 `outer` v2)
    weights = w1 + w2

apMNeu2 (MNeuron weights f) = VS.map f . (weights #>)

pics2 :: [String]
pics2 = [ " *** \n * * \n * * \n *** \n     "
        , "  *  \n  *  \n  *  \n  *  \n  *  "]

-- 4. Gradient Descent
gradStep :: Floating a => a
gradStep = 0.1

eps :: Floating a => a
eps = 0.000001

gradDesc :: (Floating a, Ord a, Show a) => (forall b. (Floating b, Ord b, Show b) => [b] -> b) -> [a] -> ([a], a)
gradDesc f x = if diff < eps then (newX, f newX) else gradDesc f newX
  where gr   = AD.grad f x
        newX = zipWith (-) x . map (gradStep *) $ gr
        diff = abs . foldr (+) 0 . zipWith (-) x $ newX

funkcja1 :: Num a => [a] -> a
funkcja1 [x,y,z] = (2 * x * x) + (2 * y * y) + (z * z) - (2 * x * y) - (2 * y * z) - (2 * x) + 3

dfunkcja1 :: Num a => [a] -> [a]
dfunkcja1 [x,y,z] = [(4*x) - (2*y) - 2, 4*y - 2*x - 2*z, 2*z - 2*y]

funkcja2 :: Num a => [a] -> a
funkcja2 [x,y] = (3 * x * x * x * x) + (4 * x * x* x) - (12 * x * x) + (12 * y * y) - (27 * y)

dfunkcja2 :: Num a => [a] -> [a]
dfunkcja2 [x,y] = [12* x* x* x + 12* x* x - 24* x, 24 * y - 24]

-- 5. Backpropagation

beta :: Double
beta = 1.25

krok :: Double
krok = 1

sig :: Double -> Double
sig = (1/).(1+).exp.negate

dsig :: Double -> Double
dsig x = sig x * (1 - sig x)

trainSetXOR :: [(Vector Double, Double)]
trainSetXOR =
  [([0,0],0)
  ,([1,0],1)
  ,([0,1],1)
  ,([1,1],0)]

fstLayer :: IO (Neuron, Neuron)
fstLayer = do
  randVec1 <- VS.replicateM 3 randomIO
  randVec2 <- VS.replicateM 3 randomIO
  return (Neuron randVec1 sig, Neuron randVec2 sig)

outLayer :: IO Neuron
outLayer = do
  randVec <- VS.replicateM 3 randomIO
  return (Neuron randVec sig)

beginning :: ((Neuron, Neuron), Neuron)
beginning = ((Neuron [2,0,1] sig, Neuron [2,0,1] sig), Neuron [2,0,1] sig)

applyXor :: ((Neuron,Neuron), Neuron) -> Vector Double -> Double
applyXor ((h1, h2), o) u = applyNeuron o out
  where
    x1 = applyNeuron h1 u
    x2 = applyNeuron h2 u
    out= [x1, x2]

infixr 3 #*
(#*) :: Double -> Vector Double -> Vector Double
x #* v = VS.map (*x) v

zipWith4 :: (a -> b -> c -> d -> e) -> [a] -> [b] -> [c] -> [d] -> [e]
zipWith4 f x y z = zipWith ($) (zipWith3 f x y z)

trainXOR perc trSet = go 0 perc
  where
    go iters n@((n1@(Neuron w1 f1), n2@(Neuron w2 f2)), n3@(Neuron s f3)) =
      if eps > (maximum . VS.toList .  VS.map abs $ difference) then
        (iters, newN)
      else
--        (if iters `mod` 10000 == 0 then trace (show iters ++ " iterations. Gradients:\n" ++ unlines [show gradS,show gradW1,show gradW2] ++ "\nDifference:\n" ++ show difference ++ "\n") else id ) $
        go (iters+1) newN
      where
        difference = (w1 VS.++ w2 VS.++s) - (newW1 VS.++ newW2 VS.++ newS)
        newN =  ((Neuron newW1 f1, Neuron newW2 f2), Neuron newS f3)
        u   = map (VS.cons 1 . fst) trSet
        x1s = map (applyNeuron n1 . fst) trSet
        x2s = map (applyNeuron n2 . fst) trSet
        x :: [Vector Double]
        x   = zipWith (\x1 x2 -> [1,x1,x2]) x1s x2s
        y   = map (applyXor n . fst) trSet
        z   = map snd trSet
        grs x y z = dsig (x<.>s) * (y - z) #* x
        gradsS= zipWith3 grs x y z
        gradS = foldr (+) [0,0,0] gradsS
        newS = s - (VS.map (*gradStep) gradS)
        [_,s1,s2] = s
        grw w sx x y z u= (y-z) * sx * dsig (x<.>s) * dsig (w<.>u) #* u
        gradsW w s = zipWith4 (grw w s) x y z u
        gradW w s = foldr (+) [0,0,0] $ gradsW w s
        gradW1 = gradW w1 s1
        gradW2 = gradW w2 s2
        newW1 = w1 - (VS.map (*gradStep) $ gradW w1 s1)
        newW2 = w2 - (VS.map (*gradStep) $ gradW w2 s2)

prettyPrintVals :: [Double] -> String
prettyPrintVals = unwords . map (printf "%.4f")

-- Zad. 6 Encoder - Decoder

autoEncTrain = map readPic [ "     \n **  \n  *  \n  *  \n  *  "
                            , "  ** \n   * \n   * \n     \n     "
                            , "     \n**   \n *   \n *   \n *   "]

type EncDec = (MNeuron, MNeuron)

makeDec (_, MNeuron w _) = MNeuron w acFunc

zeroM inp out = (out><inp) $ replicate (out*inp) 0

--startEncDec :: IO EncDec
startEncDec = do --(MNeuron (zeroM 26 16) sig, MNeuron (zeroM 17 25) sig) --do
  m1 <- map (\x -> x - 0.5) <$> replicateM (25*26) (randomIO :: IO Double)
  m2 <- map (\x -> x - 0.5) <$> replicateM (25*17) (randomIO :: IO Double)
  return (MNeuron ((16><26) m1) sig, MNeuron ((25><17) m2) sig)

decode (_,a) = applyMNeuron a

encode (a,_) = applyMNeuron a

ende a = applyMNeuron (makeDec a) . encode a

alfa = 0.8

trainED trSet = trainHelp 0
  where trainHelp old ed@(MNeuron w f1, MNeuron w' f2) =
          if abs (old - rmse) < 0.000001 then ed
          else trace ("RMSE = "++ show rmse) $ trainHelp rmse (MNeuron nW f1, MNeuron nW' f2)
          where nW            = w - cmap (*alfa) gradW
                nW'           = w'- cmap (*alfa) gradW'
                diffDecode :: Vector Double -> Vector Double
                diffDecode    = cmap dsig . (#>) w' . VS.cons 1
                diffEncode    = cmap dsig . (#>) w  . VS.cons 1
                gradDec x y x'= let dx = (x' - x) * diffDecode y
                                in dx `outer` VS.cons 1 y
                gradEnc x y x'= let dx = (x' - x) * diffDecode y
                                    dy = VS.tail (tr w' #> dx) * diffEncode x
                                in dy `outer` VS.cons 1 x
                triples       = map (\x -> (x, encode ed x,decode ed . encode ed $ x )) trSet
                gradsW'       = map (\(x,y,x') -> gradDec x y x') triples
                gradsW        = map (\(x,y,x') -> gradEnc x y x') triples
                rmse          = sqrt . VS.foldr (+) 0 . foldr (\(x,_,x') acc -> (x-x')^2 +acc ) (fromList . replicate 25 $ 0) $ triples
                gradW'        = foldr (+) (zeroM 17 25) gradsW'
                gradW         = foldr (+) (zeroM 26 16) gradsW

-- Zad 7. Sieci Hopfielda

readPic2 :: String -> Vector Double
readPic2 = VS.fromList . foldr (\x acc -> case x of
                                  ' ' -> 0:acc
                                  '*' -> 1:acc
                                  '\n'-> acc) []

jedynka :: Vector Double
jedynka = readPic2 "     \n **  \n  *  \n  *  \n  *  "

wagi :: Matrix Double
wagi = cmap (*2) $ (xs `outer` xs) * (cmap ((+1).negate) $ ident 25)
  where xs = cmap (flip (-) 0.5) jedynka

bias :: Vector Double
bias = cmap (pred.(/2)) (wagi #> fromList (replicate 25 1))

iteration :: Vector Double -> Vector Double
iteration = (flip (-) bias).(wagi #>) >>= VS.zipWith (\u x -> if u > 0 then 1 else if u < 0 then 0 else x)

inp :: Vector Double
inp = readPic2 " *   \n   * \n* *  \n    *\n *   "

process = process . iteration . (prPrint >>= trace)

randInp :: IO (Vector Double)
randInp = VS.fromList . map (\x -> if x > 0.5 then 1 else 0) <$> replicateM 25 (randomIO :: IO Double)

zeroV :: Vector Double
zeroV = readPic2 " *** \n * * \n * * \n * * \n *** "

wagi2 :: Matrix Double
wagi2 = cmap (*2) (((cx `outer` cx) + (dx `outer` dx)) * (cmap ((+1).negate) $ ident 25))
  where dx = cmap (flip (-) 0.5) zeroV
        cx = cmap (flip (-) 0.5) jedynka

bias2 :: Vector Double
bias2 = cmap (pred.pred.(/2)) (wagi2 #> fromList (replicate 25 1))

iterate2 = (flip (-) bias2).(wagi2 #>) >>= VS.zipWith (\u x -> if u > 0 then 1 else if u < 0 then 0 else x)

inp2 :: Vector Double
inp2 = readPic2 "  *  \n * * \n     \n* *  \n    *"

-- Boltzmann

prob :: Double -> Double -> Double
prob t x = (1/).(1+).exp.negate $ (x/t)

boltzStep :: Double -> Vector Double -> IO (Vector Double)
boltzStep temp = VS.mapM (\x -> randomIO >>= \y -> if y <= prob temp x then return 1 else return 0).(flip (-) bias).(wagi#>)

--replicateIO :: (a -> IO a) -> a -> IO [a]
--replicateIO f a = [ a : rest | rest <- replicateIO f nA , nA <- f a ]

annealing :: Double -> Vector Double -> IO ()
annealing = aSt 0

aSt :: Double -> Double -> Vector Double -> IO ()
aSt x inT input = do
  nv <- boltzStep (tann inT x) input
  threadDelay 300000
  putStrLn $ "Step number: " ++ show x
  putStrLn $ "Temperature: " ++ show (tann inT x)
  putStrLn . prPrint $ input
  aSt (x+1) inT nv

bSt :: Double -> Vector Double -> IO ()
bSt temp input = do
  nv <- boltzStep temp input
  threadDelay 300000
  putStrLn . prPrint $ input
  bSt temp nv

tann :: Double -> Double -> Double
tann t = (t/).(1+).log.(1+)

-- TSP

randomPermutation :: Int -> IO [Int]
randomPermutation n = do
  let ps= permutations [0..n-1]
      x = length ps
  r <- randomRIO (0,x-1) :: IO Int
  return $ ps !! r

odleglosci1 = cmap abs (row [1..10] - col [1..10]) - (col (8:replicate 9 0) * row (replicate 9 0 ++ [1])) - (col (replicate 9 0++[8]) * row (1 : replicate 9 0))


{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module      : Data.Array.Accelerate.Math.Qtfd
-- Copyright   : [2017] Rinat Stryungis
-- License     : BSD3
--
-- Maintainer  : Rinat Stryungis <lazybonesxp@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Computation of quadratic time-frequency distributions using the accelerate-fft library.
--
-- This module uses the accelerate-fft library. And the base implementation of fft
-- uses a naive divide-and-conquer fft implementation
-- whose absolute performance is appalling. It also requires that you know on
-- the Haskell side the size of the data being transformed, and that this is
-- a power-of-two in each dimension.
--
-- For performance, compile accelerate-fft against the foreign library bindings (using any
-- number of '-fllvm-ptx', and '-fllvm-cpu' for the accelerate-llvm-ptx, and
-- accelerate-llvm-native backends, respectively), which have none of the above
-- restrictions.
-- Both of this flags are enabled by default.

module Data.Array.Accelerate.Math.Qtfd where

import Data.Array.Accelerate as A
import qualified Data.Array.Accelerate.Data.Complex as ADC
import qualified Data.Array.Accelerate.Math.BornJordan as BJ
import qualified Data.Array.Accelerate.Math.ChoiWilliams as CW
import qualified Data.Array.Accelerate.Math.FFT as AMF
import qualified Data.Array.Accelerate.Math.ModifiedB as EMB
import qualified Data.Array.Accelerate.Math.PseudoWigner as P
import Data.Array.Accelerate.Math.Wigner'

import Data.Complex
import Data.List
import Math.Gamma
-- | Wigner-ville distribution. It takes 1D array of complex floating numbers and returns 2D array of real numbers.
--  Columns of result array represents time and rows - frequency. Frequency range is from 0 to n/4, where n is a sampling frequency.

wignerVille :: (A.RealFloat e, A.IsFloating e, A.FromIntegral Int e, Elt e, Elt (ADC.Complex e), AMF.Numeric e)
  => Acc (Array DIM1 (ADC.Complex e))  -- ^ Data array
  -> Acc (Array DIM2 e)
wignerVille arr =
  let times = A.enumFromN (A.index1 leng) 0 :: Acc (Array DIM1 Int)
      leng = A.length arr
      taumx = taumaxs times
      lims = limits taumx
  in A.map ADC.real . A.transpose . AMF.fft AMF.Forward $ createMatrix arr taumx lims

wignerVilleNew :: (A.RealFloat e, A.IsFloating e, A.FromIntegral Int e, Elt e, Elt (ADC.Complex e), AMF.Numeric e)
  => Acc (Array DIM1 (ADC.Complex e))  -- ^ Data array
  -> Acc (Array DIM2 e)
wignerVilleNew arr =
  let times = A.enumFromN (A.index1 leng) 0 :: Acc (Array DIM1 Int)
      leng = A.length arr
      taumx = taumaxs times
      lims = limits taumx
  in A.map ADC.real . A.transpose . AMF.fft AMF.Forward $ createMatrix arr taumx lims

-- | Pseudo Wigner-ville distribution.
-- It takes 1D array of complex floating numbers, window and returns 2D array of real numbers.
-- Columns of result array represents time and rows - frequency. Frequency range is from 0 to n/4, where n is a sampling frequency.

pWignerVille :: (A.RealFloat e, A.IsFloating e, A.FromIntegral Int e, Elt e, Elt (ADC.Complex e), AMF.Numeric e)
  => Acc (Array DIM1 e)                -- ^ Smoothing window. Length of it must be odd.
  -> Acc (Array DIM1 (ADC.Complex e))  -- ^ Data array
  -> Acc (Array DIM2 e)
pWignerVille window arr =
  let times = A.enumFromN (A.index1 leng) 0 :: Acc (Array DIM1 Int)
      leng = A.length arr
      taumx = P.taumaxs times window
      lims = P.limits taumx
  in A.map ADC.real $ A.transpose $ AMF.fft AMF.Forward $ P.createMatrix arr window taumx lims

-- | Choi-Williams distribution. It takes 1D array of complex floating numbers,
-- and returns 2D array of real numbers.
-- Columns of result array represents time and rows - frequency. Frequency range is from 0 to n/4, where n is a sampling frequency.

choiWilliams :: (A.RealFloat e, A.IsFloating e, A.FromIntegral Int e, Elt e, Elt (ADC.Complex e), AMF.Numeric e)
  => A.Exp e                           -- ^ sigma
  -> Maybe (Acc (Array DIM1 e), Acc (Array DIM1 e))
  -> A.Exp Int                         -- ^ Smoothing window over mu, must be odd and symmetrical
  -> A.Exp Int                         -- ^ Smoothing window over tau, must be odd and symmetrical
  -> Bool                              -- ^ If normalise
  -> Acc (Array DIM1 (ADC.Complex e))  -- ^ Data array
  -> Acc (Array DIM2 e)
choiWilliams sigma mWindowArrays uWindow nWindow normalise arr =
  let times = A.enumFromN (A.index1 leng) 0 :: Acc (Array DIM1 Int)
      leng = A.length arr
      taumx = taumaxs times
      lims = limits taumx
  in A.transpose $ A.map ((*4) . ADC.real) $
    AMF.fft AMF.Forward $
      CW.summedOverMu arr taumx lims sigma mWindowArrays uWindow nWindow normalise

bornJordan :: (A.RealFloat e, A.IsFloating e, A.FromIntegral Int e, Elt e, Elt (ADC.Complex e), AMF.Numeric e)
  => A.Exp e                           -- ^ sigma
  -> Maybe (Acc (Array DIM1 e), Acc (Array DIM1 e))
  -> A.Exp Int                         -- ^ Smoothing window over mu, must be odd and symmetrical
  -> A.Exp Int                         -- ^ Smoothing window over tau, must be odd and symmetrical
  -> Acc (Array DIM1 (ADC.Complex e))  -- ^ Data array
  -> Acc (Array DIM2 e)
bornJordan sigma mWindowArrays uWindow nWindow arr =
  let times = A.enumFromN (A.index1 leng) 0 :: Acc (Array DIM1 Int)
      leng = A.length arr
      taumx = taumaxs times
      lims = limits taumx
  in A.transpose $ A.map ((*2) . ADC.real) $
    AMF.fft AMF.Forward $
      BJ.summedOverMu arr taumx lims sigma mWindowArrays uWindow nWindow

eModifiedB :: (A.RealFloat e, A.IsFloating e, A.FromIntegral Int e, Elt e, Elt (ADC.Complex e), AMF.Numeric e)
  => A.Exp e                           -- ^ alpha
  -> (Acc (Array DIM1 e))              -- ^ gammas
  -> Maybe (Acc (Array DIM1 e), Acc (Array DIM1 e))
  -> A.Exp Int
  -> A.Exp Int
  -> Bool
  -> Acc (Array DIM1 (ADC.Complex e))  -- ^ Data array
  -> Acc (Array DIM2 e)
eModifiedB alpha gammas mWindowArrays uWindow nWindow normalise arr =
  let times = A.enumFromN (A.index1 leng) 0 :: Acc (Array DIM1 Int)
      leng = A.length arr
      taumx = taumaxs times
      lims = limits taumx
  in A.transpose $ A.map (*2) $ A.map ADC.real $ AMF.fft AMF.Forward $ A.transpose $ EMB.summedOverMu arr taumx lims alpha gammas mWindowArrays uWindow nWindow normalise

makeGammaArray2 :: Double -> Int -> [Double]
makeGammaArray2 beta wLength =
  let dwLength = Prelude.fromIntegral wLength :: Double
      v = [(-0.5), ((-0.5) + 1.0/dwLength)..0.5]
      s = Prelude.map (\x -> beta :+ pi*x) v
      f = Prelude.map (\x -> abs $ ((magnitude $ gamma x) Prelude.^ 2)/((gamma beta) Prelude.^ 2)) s
      n = Prelude.floor $ (Prelude.fromIntegral $ Prelude.length f)/2.0
      frst = Data.List.take (n-1) f
      secnd  = Data.List.drop n f
  in secnd Prelude.++ frst


{-# LANGUAGE FlexibleContexts #-}
module STC.PowerMethodNormalization where

import           Data.Array.Repa      as R
import           Data.Array.Unboxed   as AU
import           Data.Complex
import           Data.List            as L
import           Data.Vector.Storable as VS
import           DFT.Plan
import           Filter.Utils
import           Types
import           Utils.Array
import           Utils.Parallel

data PowerMethodNormalizationOption
  = PowerMethodPatchNorm (VS.Vector (Complex Double))
  | PowerMethodGlobal
  | PowerMethodConnection [[(Int, Int)]]
  | PowerMethodConnectionMax [[(Int, Int)]]

-- {-# INLINE makePatchNormFilter #-}
-- makePatchNormFilter ::
--      DFTPlan
--   -> Int
--   -> Int
--   -> Bool
--   -> Int
--   -> IO (DFTPlan, PowerMethodNormalizationOption)
-- makePatchNormFilter plan cols rows True n = do
--   let arr =
--         makeFilter2D . fromFunction (Z :. cols :. rows) $ \(Z :. i :. j) ->
--           if (sqrt . fromIntegral $ (i - div cols 2) ^ 2 + (j - div rows 2) ^ 2) <=
--              fromIntegral n
--             then 1
--             else 0
--   lock <- getFFTWLock
--   (p, vec) <-
--     dft1dGPlan lock plan [cols, rows] [0, 1] . VS.convert . toUnboxed . computeS $
--     arr
--   (newPlan, _) <- idft1dGPlan lock p [cols, rows] [0, 1] $ vec
--   return (newPlan, PowerMethodPatchNorm vec)
-- makePatchNormFilter plan _ _ False _ = return (plan, PowerMethodGlobal)

-- {-# INLINE patchNorm #-}
-- patchNorm ::
--      (R.Source s (Complex Double))
--   => DFTPlan
--   -> VS.Vector (Complex Double)
--   -> R.Array s DIM4 (Complex Double)
--   -> IO (R.Array U DIM4 (Complex Double))
-- patchNorm plan filterF arr = do
--   let (Z :. numThetaFreq :. numScaleFreq :. cols :. rows) = extent arr
--   arrSum <- (R.sumP . rotate4D . rotate4D . R.map magnitude $ arr) >>= R.sumP
--   arrSumF <-
--     dftExecute plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
--     VS.convert . toUnboxed . computeS . R.map (:+ 0) $
--     arrSum
--   arrPatchSum <-
--     fmap (fromUnboxed (Z :. cols :. rows) . VS.convert) .
--     dftExecute plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) $
--     VS.zipWith (*) arrSumF filterF
--   computeP . R.traverse2 arr arrPatchSum const $ \f fNorm idx@(Z :. _ :. _ :. i :. j) ->
--     if fNorm (Z :. i :. j) == 0
--       then 0
--       else f idx / fNorm (Z :. i :. j)

-- {-# INLINE powerMethodNormalization #-}
-- powerMethodNormalization ::
--      (R.Source s (Complex Double))
--   => DFTPlan
--   -> PowerMethodNormalizationOption
--   -> R.Array s DIM4 (Complex Double)
--   -> IO (R.Array U DIM4 (Complex Double))
-- powerMethodNormalization plan (PowerMethodPatchNorm filterF) arr =
--   patchNorm plan filterF arr
-- powerMethodNormalization _ PowerMethodGlobal arr = do
--   s <- R.sumAllP . rotate4D . rotate4D . R.map magnitude $ arr
--   return . computeS . R.map (/ (s :+ 0)) $ arr

{-# INLINE powerMethodNormalization #-}
powerMethodNormalization ::
     (R.Source s (Complex Double))
  => PowerMethodNormalizationOption
  -> R.Array s DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
powerMethodNormalization PowerMethodGlobal arr = do
  -- s <- R.sumAllP . rotate4D . rotate4D . R.map (\x -> (magnitude x) ^ 2) $ arr
  s <- R.foldAllP max 0 . rotate4D . rotate4D . R.map (\x -> magnitude x) $ arr
  return . computeS . R.map (/ (s :+ 0)) $ arr
powerMethodNormalization (PowerMethodConnection xss) arr = do
  let (Z :. _ :. _ :. cols :. rows) = extent arr
      arrNorm =
        fromListUnboxed (Z :. cols :. rows) . AU.elems $
        (AU.accumArray (+) 0 ((0, 0), (cols - 1, rows - 1)) .
         L.concat .
         parMap
           rdeepseq
           (\xs ->
              let s =
                    L.maximum .
                    L.map
                      (\(i, j) ->
                         R.foldAllS max 0 .
                         R.slice (R.map (\x -> magnitude x) arr) $
                         (Z :. All :. All :. i :. j)) $
                    xs
               in L.map
                    (\x ->
                       ( x
                       , if s == 0
                           then 0
                           else 1 / s))
                    xs) $
         xss :: UArray (Int, Int) Double)
  return . computeS . R.traverse2 arr arrNorm const $ \f fNorm idx@(Z :. _ :. _ :. i :. j) ->
    f idx * ((fNorm (Z :. i :. j)) :+ 0)
powerMethodNormalization (PowerMethodConnectionMax xss) arr = do
  let (Z :. numThetaFreq :. numScaleFreq :. cols :. rows) = extent arr
      arrMax =
        fromListUnboxed (Z :. cols :. rows) . AU.elems $
        (AU.accumArray (+) 0 ((0, 0), (cols - 1, rows - 1)) .
         L.concat .
         parMap
           rdeepseq
           (\xs ->
              let s =
                    sqrt .
                    L.maximum .
                    L.map
                      (\(i, j) ->
                         R.sumAllS .
                         R.slice (R.map (\x -> (magnitude x) ^ 2) arr) $
                         (Z :. All :. All :. i :. j)) $
                    xs
               in L.map (\x -> (x, s)) xs) $
         xss :: UArray (Int, Int) Double)
      arrSum =
        fromListUnboxed (Z :. cols :. rows) . AU.elems $
        (AU.accumArray (+) 0 ((0, 0), (cols - 1, rows - 1)) .
         L.concat .
         parMap
           rdeepseq
           (\xs ->
              L.map
                (\(i, j) ->
                   ( (i, j)
                   , sqrt .
                     R.sumAllS . R.slice (R.map (\x -> (magnitude x) ^ 2) arr) $
                     (Z :. All :. All :. i :. j))) $
              xs) $
         xss :: UArray (Int, Int) Double)
  return . computeS . R.traverse3 arr arrMax arrSum (\s1 s2 s3 -> s1) $ \f fMax fSum idx@(Z :. k :. l :. i :. j) ->
    if fSum (Z :. i :. j) > 0.9 * fMax (Z :. i :. j) && k == div numThetaFreq 2 && l == div numScaleFreq 2
      then 1 -- (f idx) / (fSum (Z :. i :. j) :+ 0)
      else 0


{-# LANGUAGE RecordWildCards, BangPatterns #-}

module NeuralNetwork where

import qualified  Statistics.Matrix as M
import            Control.Monad.Random

-- hyperparameters

data NeuralNetwork =
     NeuralNetwork { inputLayerSize  :: Int
                   , outputLayerSize :: Int
                   , hiddenLayerSize :: Int
                   , weight1Matrix   :: M.Matrix
                   , weight2Matrix   :: M.Matrix
                   }

initNetwork :: MonadRandom r => Int -> Int -> Int -> r NeuralNetwork
initNetwork inputSize outputSize hiddenSize = do
    rs  <- getRandomRs (0.0, 1.0)
    rs' <- getRandomRs (0.0, 1.0)
    return $ NeuralNetwork { inputLayerSize  = inputSize
                           , outputLayerSize = outputSize
                           , hiddenLayerSize = hiddenSize
                           , weight1Matrix   = M.fromList inputSize  hiddenSize $ take (inputSize  * hiddenSize) rs
                           , weight2Matrix   = M.fromList hiddenSize outputSize $ take (outputSize * hiddenSize) rs'
                           }

forwardRs :: NeuralNetwork -> M.Matrix -> (M.Matrix, M.Matrix, M.Matrix, M.Matrix)
forwardRs NeuralNetwork{..} matrix =
    let !z2 = matrix * weight1Matrix
        !a2 = M.map sigmoid z2
        !z3 = a2 * weight2Matrix
    in  (M.map sigmoid z3, z2, a2, z3)

forward :: NeuralNetwork -> M.Matrix -> M.Matrix
forward nn matrix =
    let (r, _, _, _) = forwardRs nn matrix
    in r

sigmoid :: Double -> Double
sigmoid z = 1 / (1 + exp (-z))

-- | Derivative of sigmoid function
sigmoidPrime :: Double -> Double
sigmoidPrime z = ez / ((1 + ez) ** 2)
    where ez = exp (-z)

costFunction :: NeuralNetwork -> M.Matrix -> M.Matrix -> Double
costFunction nn@NeuralNetwork{..} x y =
    let yHat = nn `forward` x
        err  = sum . map (**2) . M.toList $ y - yHat
    in  0.5 * err

costFunctionPrime :: NeuralNetwork -> M.Matrix -> M.Matrix -> (M.Matrix, M.Matrix)
costFunctionPrime nn@NeuralNetwork{..} x y =
    let (yHat, z2, a2, z3) = nn `forwardRs` x
        !delta3            =  (- (y - yHat)) `pointMultiply` (M.map sigmoidPrime z3)
        !dJdW2             = (M.transpose a2) * delta3

        !delta2            = delta3 * (M.transpose weight2Matrix) * (M.map sigmoidPrime z2)
        !dJdW1             = M.transpose x * delta2

    in (dJdW1, dJdW2)

pointMultiply :: M.Matrix -> M.Matrix -> M.Matrix
pointMultiply m1 m2 =
    let (r, c) = (M.rows m1, M.cols m1)
        (l1, l2) = (M.toList m1, M.toList m2)
    in M.fromList r c (foldr (\(a,b) xs -> (a * b) : xs) [] (zip l1 l2))

instance Num M.Matrix where
    m1 + m2  = let (r, c) = (M.rows m1, M.cols m1)
                   (l1, l2) = (M.toList m1, M.toList m2)
               in M.fromList r c (foldr (\(a,b) xs -> (a + b) : xs) [] (zip l1 l2))

    negate m      = let (r, c) = (M.rows m, M.cols m)
                        l      = M.toList m
                    in M.fromList r c (map negate l)

    m1 * m2  = m1 `M.multiply` m2

    abs m         = undefined
    signum m      = undefined
    fromInteger m = undefined


{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}

-- {-# OPTIONS_GHC -O -ddump-rule-firings #-}

-- | Provides the class `ParamSet` which is used to represent the set of
-- parameters of a particular model.  The goal of SGD is then to find the
-- parameter values which minimize a given objective function.

module Numeric.SGD.ParamSet
  (
  -- * Class
    ParamSet(..)
  -- * Generics
  , GPMap
  , GAdd
  , GSub
  , GDiv
  , GMul
  , GNorm2
  ) where

import           GHC.Generics
import           GHC.TypeNats (KnownNat)

import           Prelude hiding (div)

import qualified Data.Map.Strict as M

import qualified Numeric.LinearAlgebra.Static as LA

-- | Class of types that can be treated as parameter sets.  It provides basic
-- element-wise operations (addition, multiplication, mapping) which are
-- required to perform stochastic gradient descent.  Many of the operations
-- (`add`, `mul`, `sub`, `div`, etc.) have the same interpretation and follow
-- the same laws (e.g. associativity) as the corresponding operations in `Num`
-- and `Fractional`.
--
-- `zero` takes a parameter set as argument and "zero out"'s all its elements
-- (as in the backprop library).  This allows instances for `Maybe`, `M.Map`,
-- etc., where the structure of the parameter set is dynamic.  This leads to
-- the following property:
--
--     @add (zero x) x = x@
--
-- However, `zero` does not have to obey @(add (zero x) y = y)@.
--
-- A `ParamSet` can be also seen as a (structured) vector, hence `pmap` and
-- `norm_2`.  The latter is not strictly necessary to perform SGD, but it is
-- useful to control the training process.
--
-- `pmap` should obey the following law:
--
--     @pmap id x = x@
--
-- If you leave the body of an instance declaration blank, GHC Generics will be
-- used to derive instances if the type has a single constructor and each field
-- is an instance of `ParamSet`.
class ParamSet a where
  -- | Element-wise mapping
  pmap :: (Double -> Double) -> a -> a

  -- | Zero-out all elements
  zero :: a -> a
  zero = pmap (const 0.0)

--   -- | Element-wise negation
--   neg :: a -> a
--   neg = pmap (\x -> -x)

  -- | Element-wise addition
  add :: a -> a -> a
  -- | Elementi-wise substruction
  sub :: a -> a -> a

  -- | Element-wise multiplication
  mul :: a -> a -> a
  -- | Element-wise division
  div :: a -> a -> a

  -- | L2 norm
  norm_2 :: a -> Double

--   default zero :: (Generic a, GZero (Rep a)) => a -> a
--   zero = genericZero
--   {-# INLINE zero #-}

  default pmap
    :: (Generic a, GPMap (Rep a))
    => (Double -> Double) -> a -> a
  pmap = genericPMap
  {-# INLINE pmap #-}

  default add :: (Generic a, GAdd (Rep a)) => a -> a -> a
  add = genericAdd
  {-# INLINE add #-}

  default sub :: (Generic a, GSub (Rep a)) => a -> a -> a
  sub = genericSub
  {-# INLINE sub #-}

  default mul :: (Generic a, GMul (Rep a)) => a -> a -> a
  mul = genericMul
  {-# INLINE mul #-}

  default div :: (Generic a, GDiv (Rep a)) => a -> a -> a
  div = genericDiv
  {-# INLINE div #-}

  default norm_2 :: (Generic a, GNorm2 (Rep a)) => a -> Double
  norm_2 = genericNorm2
  {-# INLINE norm_2 #-}

{-# RULES
"ParamSet pmap/pmap" forall f g p. pmap f (pmap g p) = pmap (f . g) p
  #-}

-- {-# RULES
-- "ParamSet pmap/add/pmap" forall f g p h q.
--   pmap f (add (pmap g p) (pmap h q))
--   = add (pmap (f . g) p) (pmap (f . h) q)
--   #-}

-- -- | 'add' using GHC Generics; works if all fields are instances of
-- -- 'ParamSet', but only for values with single constructors.
-- genericZero :: (Generic a, GZero (Rep a)) => a -> a
-- genericZero x = to $ gzero (from x)
-- {-# INLINE genericZero #-}

-- | 'add' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericAdd :: (Generic a, GAdd (Rep a)) => a -> a -> a
genericAdd x y = to $ gadd (from x) (from y)
{-# INLINE genericAdd #-}

-- | 'sub' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericSub :: (Generic a, GSub (Rep a)) => a -> a -> a
genericSub x y = to $ gsub (from x) (from y)
{-# INLINE genericSub #-}

-- | 'div' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericDiv :: (Generic a, GDiv (Rep a)) => a -> a -> a
genericDiv x y = to $ gdiv (from x) (from y)
{-# INLINE genericDiv #-}

-- | 'mul' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericMul :: (Generic a, GMul (Rep a)) => a -> a -> a
genericMul x y = to $ gmul (from x) (from y)
{-# INLINE genericMul #-}

-- | 'norm_2' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericNorm2 :: (Generic a, GNorm2 (Rep a)) => a -> Double
genericNorm2 x = gnorm_2 (from x)
{-# INLINE genericNorm2 #-}

-- | 'pmap' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericPMap :: (Generic a, GPMap (Rep a)) => (Double -> Double) -> a -> a
genericPMap f x = to $ gpmap f (from x)
{-# INLINE genericPMap #-}

--------------------------------------------------
-- Generics
--
-- Partially borrowed from the backprop library
--------------------------------------------------

-- -- | Helper class for automatically deriving 'add' using GHC Generics.
-- class GZero f where
--     gzero :: f t -> f t
--
-- instance ParamSet p => GZero (K1 i p) where
--     gzero (K1 x) = K1 (zero x)
--     {-# INLINE gzero #-}
--
-- instance (GZero f, GZero g) => GZero (f :*: g) where
--     gzero (x1 :*: y1) = x2 :*: y2
--       where
--         !x2 = gzero x1
--         !y2 = gzero y1
--     {-# INLINE gzero #-}
--
-- instance GZero V1 where
--     gzero = \case {}
--     {-# INLINE gzero #-}
--
-- instance GZero U1 where
--     gzero _ = U1
--     {-# INLINE gzero #-}
--
-- instance GZero f => GZero (M1 i c f) where
--     gzero (M1 x) = M1 (gzero x)
--     {-# INLINE gzero #-}
--
-- -- instance GZero f => GZero (f :.: g) where
-- --     gzero = Comp1 gzero
-- --     {-# INLINE gzero #-}

-- | Helper class for automatically deriving 'add' using GHC Generics.
class GAdd f where
    gadd :: f t -> f t -> f t

instance ParamSet a => GAdd (K1 i a) where
    gadd (K1 x) (K1 y) = K1 (add x y)
    {-# INLINE gadd #-}

instance (GAdd f, GAdd g) => GAdd (f :*: g) where
    gadd (x1 :*: y1) (x2 :*: y2) = x3 :*: y3
      where
        !x3 = gadd x1 x2
        !y3 = gadd y1 y2
    {-# INLINE gadd #-}

instance GAdd V1 where
    gadd = \case {}
    {-# INLINE gadd #-}

instance GAdd U1 where
    gadd _ _ = U1
    {-# INLINE gadd #-}

instance GAdd f => GAdd (M1 i c f) where
    gadd (M1 x) (M1 y) = M1 (gadd x y)
    {-# INLINE gadd #-}

-- instance GAdd f => GAdd (f :.: g) where
--     gadd (Comp1 x) (Comp1 y) = Comp1 (gadd x y)
--     {-# INLINE gadd #-}

-- | Helper class for automatically deriving 'sub' using GHC Generics.
class GSub f where
    gsub :: f t -> f t -> f t

instance ParamSet a => GSub (K1 i a) where
    gsub (K1 x) (K1 y) = K1 (sub x y)
    {-# INLINE gsub #-}

instance (GSub f, GSub g) => GSub (f :*: g) where
    gsub (x1 :*: y1) (x2 :*: y2) = x3 :*: y3
      where
        !x3 = gsub x1 x2
        !y3 = gsub y1 y2
    {-# INLINE gsub #-}

instance GSub V1 where
    gsub = \case {}
    {-# INLINE gsub #-}

instance GSub U1 where
    gsub _ _ = U1
    {-# INLINE gsub #-}

instance GSub f => GSub (M1 i c f) where
    gsub (M1 x) (M1 y) = M1 (gsub x y)
    {-# INLINE gsub #-}

-- instance GSub f => GSub (f :.: g) where
--     gsub (Comp1 x) (Comp1 y) = Comp1 (gsub x y)
--     {-# INLINE gsub #-}

-- | Helper class for automatically deriving 'mul' using GHC Generics.
class GMul f where
    gmul :: f t -> f t -> f t

instance ParamSet a => GMul (K1 i a) where
    gmul (K1 x) (K1 y) = K1 (mul x y)
    {-# INLINE gmul #-}

instance (GMul f, GMul g) => GMul (f :*: g) where
    gmul (x1 :*: y1) (x2 :*: y2) = x3 :*: y3
      where
        !x3 = gmul x1 x2
        !y3 = gmul y1 y2
    {-# INLINE gmul #-}

instance GMul V1 where
    gmul = \case {}
    {-# INLINE gmul #-}

instance GMul U1 where
    gmul _ _ = U1
    {-# INLINE gmul #-}

instance GMul f => GMul (M1 i c f) where
    gmul (M1 x) (M1 y) = M1 (gmul x y)
    {-# INLINE gmul #-}

-- instance GMul f => GMul (f :.: g) where
--     gmul (Comp1 x) (Comp1 y) = Comp1 (gmul x y)
--     {-# INLINE gmul #-}

-- | Helper class for automatically deriving 'div' using GHC Generics.
class GDiv f where
    gdiv :: f t -> f t -> f t

instance ParamSet a => GDiv (K1 i a) where
    gdiv (K1 x) (K1 y) = K1 (div x y)
    {-# INLINE gdiv #-}

instance (GDiv f, GDiv g) => GDiv (f :*: g) where
    gdiv (x1 :*: y1) (x2 :*: y2) = x3 :*: y3
      where
        !x3 = gdiv x1 x2
        !y3 = gdiv y1 y2
    {-# INLINE gdiv #-}

instance GDiv V1 where
    gdiv = \case {}
    {-# INLINE gdiv #-}

instance GDiv U1 where
    gdiv _ _ = U1
    {-# INLINE gdiv #-}

instance GDiv f => GDiv (M1 i c f) where
    gdiv (M1 x) (M1 y) = M1 (gdiv x y)
    {-# INLINE gdiv #-}

-- instance GDiv f => GDiv (f :.: g) where
--     gdiv (Comp1 x) (Comp1 y) = Comp1 (gdiv x y)
--     {-# INLINE gdiv #-}

-- | Helper class for automatically deriving 'norm_2' using GHC Generics.
class GNorm2 f where
    gnorm_2 :: f t -> Double

instance ParamSet a => GNorm2 (K1 i a) where
    gnorm_2 (K1 x) = norm_2 x
    {-# INLINE gnorm_2 #-}

instance (GNorm2 f, GNorm2 g) => GNorm2 (f :*: g) where
    gnorm_2 (x1 :*: y1) =
      sqrt ((x2 ^ (2 :: Int)) + (y2 ^ (2 :: Int)))
      where
        !x2 = gnorm_2 x1
        !y2 = gnorm_2 y1
    {-# INLINE gnorm_2 #-}

instance GNorm2 V1 where
    gnorm_2 = \case {}
    {-# INLINE gnorm_2 #-}

instance GNorm2 U1 where
    gnorm_2 _ = 0
    {-# INLINE gnorm_2 #-}

instance GNorm2 f => GNorm2 (M1 i c f) where
    gnorm_2 (M1 x) = gnorm_2 x
    {-# INLINE gnorm_2 #-}

-- -- TODO: Make sure this makes sense
-- instance GNorm2 f => GNorm2 (f :.: g) where
--     gnorm_2 (Comp1 x) = gnorm_2 x
--     {-# INLINE gnorm_2 #-}

-- | Helper class for automatically deriving 'pmap' using GHC Generics.
class GPMap f where
    gpmap :: (Double -> Double) -> f t -> f t

instance ParamSet a => GPMap (K1 i a) where
    gpmap f (K1 x) = K1 (pmap f x)
    {-# INLINE gpmap #-}

instance (GPMap f, GPMap g) => GPMap (f :*: g) where
    gpmap f (x1 :*: y1) = x2 :*: y2
      where
        !x2 = gpmap f x1
        !y2 = gpmap f y1
    {-# INLINE gpmap #-}

instance GPMap V1 where
    gpmap _ = \case {}
    {-# INLINE gpmap #-}

instance GPMap U1 where
    gpmap _ _ = U1
    {-# INLINE gpmap #-}

instance GPMap f => GPMap (M1 i c f) where
    gpmap f (M1 x) = M1 (gpmap f x)
    {-# INLINE gpmap #-}

-- instance GPMap f => GPMap (f :.: g) where
--     gpmap f (Comp1 x) = Comp1 (gpmap f x)
--     {-# INLINE gpmap #-}

--------------------------------------------------
-- Basic instances
--------------------------------------------------

instance ParamSet Double where
  zero = const 0
  pmap = id
  add = (+)
  sub = (-)
  mul = (*)
  div = (/)
  norm_2 = abs

instance (ParamSet a, ParamSet b) => ParamSet (a, b) where
  pmap f (x, y) = (pmap f x, pmap f y)
  add (x1, y1) (x2, y2) = (x1 `add` x2, y1 `add` y2)
  sub (x1, y1) (x2, y2) = (x1 `sub` x2, y1 `sub` y2)
  mul (x1, y1) (x2, y2) = (x1 `mul` x2, y1 `mul` y2)
  div (x1, y1) (x2, y2) = (x1 `div` x2, y1 `div` y2)
  norm_2 (x, y)
    = sqrt . sum . map ((^(2::Int)))
    $ [norm_2 x, norm_2 y]

instance (KnownNat n) => ParamSet (LA.R n) where
  zero = const 0
  pmap = LA.dvmap
  add = (+)
  sub = (-)
  mul = (*)
  div = (/)
  norm_2 = LA.norm_2

instance (KnownNat n, KnownNat m) => ParamSet (LA.L n m) where
  zero = const 0
  pmap = LA.dmmap
  add = (+)
  sub = (-)
  mul = (*)
  div = (/)
  norm_2 = LA.norm_2

-- | `Nothing` represents a deactivated parameter set component. If `Nothing`
-- is given as an argument to one of the `ParamSet` operations, the result is
-- `Nothing` as well.
--
-- This differs from the corresponding instance in the backprop library, where
-- `Nothing` is equivalent to `Just 0`.  However, the implementation below
-- seems to correspond adequately enough to the notion that a particular
-- component is either active or not in both the parameter set and the
-- gradient, hence it doesn't make sense to combine `Just` with `Nothing`.
instance (ParamSet a) => ParamSet (Maybe a) where
  zero = fmap zero
  pmap = fmap . pmap

  add (Just x) (Just y) = Just (add x y)
  add _ _ = Nothing

  sub (Just x) (Just y) = Just (sub x y)
  sub _ _ = Nothing

  mul (Just x) (Just y) = Just (mul x y)
  mul _ _ = Nothing

  div (Just x) (Just y) = Just (div x y)
  div _ _ = Nothing

  norm_2 = maybe 0 norm_2

-- | A map with different parameter sets (of the same type) assigned to the
-- individual keys.
--
-- When combining two maps with different sets of keys, only their intersection
-- is preserved.
instance (Ord k, ParamSet a) => ParamSet (M.Map k a) where
  zero = fmap zero
  pmap f = fmap (pmap f)
  add = M.intersectionWith add
  sub = M.intersectionWith sub
  mul= M.intersectionWith mul
  div= M.intersectionWith div
  norm_2 = sqrt . sum . map ((^(2::Int)) . norm_2)  . M.elems


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE TypeSynonymInstances #-}
module DSLsofMath.W06 where
import DSLsofMath.FunExp hiding (eval, derive)
import DSLsofMath.W05
import DSLsofMath.Algebra
import Prelude hiding (Num(..),(/),(^),Fractional(..),Floating(..))
import Prelude (abs)
import Data.Complex ()

evalAll :: Transcendental a => FunExp -> [a -> a]
evalAll e = (evalFunExp e) : evalAll (derive e)

help a b as bs = as * (b : bs) + (a : as) * bs

mulStream :: Ring a => Stream a -> Stream a -> Stream a
mulStream []        _         = []
mulStream _         []        = []
mulStream (a : as)  (b : bs)  = (a*b) :  (as * (b : bs) + (a : as) * bs)

type Stream a = [a]
instance Additive a => Additive (Stream a) where
  zero  = repeat zero
  (+)   = addStream
instance AddGroup a => AddGroup (Stream a) where
  negate = negStream
instance Ring a => Multiplicative (Stream a) where
  one   = one : zero
  (*)   = mulStream

addStream :: Additive a => Stream a -> Stream a -> Stream a
addStream = zipWithLonger (+)

negStream :: AddGroup a => Stream a -> Stream a
negStream = map negate

type Taylor a = Stream a

toMaclaurin :: Ring a => PowerSeries a -> Taylor a
toMaclaurin (Poly as) = zipWith (*) as factorials

fromMaclaurin :: Field a => Taylor a -> PowerSeries a
fromMaclaurin as = Poly (zipWith (/) as factorials)

factorials :: Ring a => [a]
factorials = factorialsFrom 0 1

factorialsFrom :: Ring a => a -> a -> [a]
factorialsFrom n factn = factn : factorialsFrom (n+1) (factn * (n + 1))

ex3, ex4 :: (Eq a, Field a) => Taylor a
ex3 = toMaclaurin (x^3 + two * x)
ex4 = toMaclaurin sinx

ida a = toMaclaurin (evalP (X :+: Const a))

d f a = take 10 (toMaclaurin (evalP (f (X :+: Const a))))

dP f a = toMaclaurin (f (idx + Poly [a]))

integT :: a -> Taylor a -> Taylor a
integT = (:)

integ  ::  Field a => a -> PowerSeries a -> PowerSeries a
integ  a0 (Poly as)  =  Poly (integL a0 as)

integL :: Field a => a -> [a] -> [a]
integL c cs = c : zipWith (/) cs oneUp

type PS a = PowerSeries a

solve :: Field a => a -> (PS a -> PS a) -> PS a
solve f0 g = f              -- solves |f' = g f|, |f 0 = f0|
  where f = integ f0 (g f)

idx  ::  Field a => PS a
idx  =   solve 0 (\_f -> 1)         -- \(f'(x) = 1\), \(f(0) = 0\)

expx  ::  Field a => PS a
expx  =   solve 1 (\f -> f)         -- \(f'(x) = f(x)\), \(f(0) = 1\)
expf  ::  Field a => a -> a
expf  =   evalPS 100 expx

testExp :: Double
testExp = maximum (map diff [0,0.001..1::Double])
  where diff = abs . (expf - exp)  -- using the function instance for |exp|

testExpUnits :: Double
testExpUnits =  testExp / epsilon

epsilon :: Double  -- one bit of |Double| precision
epsilon = last (takeWhile (\x -> 1 + x /= 1) (iterate (/2) 1))

sinx,  cosx  :: Field a =>  PS a
sinx  =  integ 0 cosx
cosx  =  integ 1 (-sinx)

sinf,  cosf  :: Field a =>  a -> a
sinf  =  evalPS 100 sinx
cosf  =  evalPS 100 cosx

sx,cx::[Double]
sx = 0  :  1      :  neg 0           :  frac (neg 1) 6  :  error "TODO"
cx = 1  :  neg 0  :  frac (neg 1) 2  :  0               :  error "TODO"

instance (Eq a, Transcendental a) => Transcendental (PowerSeries a) where
   pi   =  Poly [pi]
   exp  =  expPS
   sin  =  sinPS
   cos  =  cosPS

expPS, sinPS, cosPS :: (Eq a, Transcendental a) => PS a -> PS a
expPS  as  = integ  (exp  (val as))  (expPS as   * deriv as)
sinPS  as  = integ  (sin  (val as))  (cosPS as   * deriv as)
cosPS  as  = integ  (cos  (val as))  (-sinPS as  * deriv as)

val ::  Additive a => PS a  ->  a
val (Poly (a:_))   =   a
val _              =   zero

evalP :: (Eq r, Transcendental r) => FunExp -> PS r
evalP (Const x)    =  Poly [fromRational (toRational x)]
evalP (e1 :+: e2)  =  evalP e1 + evalP e2
evalP (e1 :*: e2)  =  evalP e1 * evalP e2
evalP X            =  idx
evalP (Negate e)   =  negate  (evalP e)
evalP (Recip e)    =  recip   (evalP e)
evalP (Exp e)      =  exp     (evalP e)
evalP (Sin e)      =  sin     (evalP e)
evalP (Cos e)      =  cos     (evalP e)

evalFunExp  ::  Transcendental a => FunExp -> a -> a
evalFunExp  (Const alpha)  =   const (fromRational (toRational alpha))
evalFunExp  X              =   id
evalFunExp  (e1 :+: e2)    =   evalFunExp e1  +  evalFunExp e2
evalFunExp  (e1 :*: e2)    =   evalFunExp e1  *  evalFunExp e2
evalFunExp  (Exp e)        =   exp     (evalFunExp e)
evalFunExp  (Sin e)        =   sin     (evalFunExp e)
evalFunExp  (Cos e)        =   cos     (evalFunExp e)
evalFunExp  (Recip e)      =   recip   (evalFunExp e)
evalFunExp  (Negate e)     =   negate  (evalFunExp e)

derive     (Const _)      =  Const 0
derive     X              =  Const 1
derive     (e1 :+: e2)    =  derive e1  :+:  derive e2
derive     (e1 :*: e2)    =  (derive e1  :*:  e2)  :+:  (e1  :*:  derive e2)
derive     (Recip e)      =  let re = Recip e in Negate (re:*:re) :*: derive e
derive     (Negate e)     =  Negate (derive e)
derive     (Exp e)        =  Exp e :*: derive e
derive     (Sin e)        =  Cos e :*: derive e
derive     (Cos e)        =  Const (-1) :*: Sin e :*: derive e

instance Additive FunExp where
  (+)   =  (:+:)
  zero  = Const 0

instance AddGroup FunExp where
  negate x  = Const (-1) * x

instance Multiplicative FunExp where
  (*)  =  (:*:)
  one  = Const 1

instance MulGroup FunExp where
  recip = Recip

instance Transcendental FunExp where
  pi   =  Const pi
  exp  =  Exp
  sin  =  Sin
  cos  =  Cos

instance Additive a => Additive (a, a) where
  (f, f') + (g, g')  = (f + g,  f' + g')
  zero               = (zero,   zero)

instance AddGroup a => AddGroup (a, a) where
  negate (f, f')     = (negate f, negate f')

instance Ring a => Multiplicative (a,a) where
  (f, f') * (g, g')  = (f * g,  f' * g + f * g')
  one                = (one,zero)

instance Field a => MulGroup (a, a) where
  (f, f') / (g, g')  = (f / g,  (f' * g - g' * f) / (g * g))

instance Transcendental a => Transcendental (a, a) where
  pi = (pi, zero)
  exp  (f, f')       = (exp f,  (exp f) * f')
  sin  (f, f')       = (sin f,  cos f * f')
  cos  (f, f')       = (cos f,  -(sin f) * f')