{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns        #-}
--------------------------------------------------------------------------------
-- |
-- Module      : ArrayFire.Arith
-- Copyright   : David Johnson (c) 2019-2020
-- License     : BSD 3
-- Maintainer  : David Johnson <djohnson.m@gmail.com>
-- Stability   : Experimental
-- Portability : GHC
--
-- Arithmetic functions over 'Array'
--
-- @
-- module Main where
--
-- import qualified ArrayFire as A
--
-- main :: IO ()
-- main = print $ A.scalar \@Int 1 \`A.add\` A.scalar \@Int 1
-- -- 2
-- @
--------------------------------------------------------------------------------
module ArrayFire.Arith where

import Prelude                  (Bool(..), ($), (.), flip, fromEnum, fromIntegral, Real, RealFrac)

import Data.Coerce
import Data.Proxy
import Data.Complex

import ArrayFire.FFI
import ArrayFire.Internal.Arith
import ArrayFire.Internal.Types

-- | Adds two 'Array' objects
--
-- >>> A.scalar @Int 1 `A.add` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--         2
add
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of add
add x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_add arr arr1 arr2 1

-- | Adds two 'Array' objects
--
-- >>> (A.scalar @Int 1 `A.addBatched` A.scalar @Int 1) True
-- ArrayFire Array
-- [1 1 1 1]
--         2
addBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of add
addBatched x y (fromIntegral . fromEnum -> batch) =
  x `op2` y $ \arr arr1 arr2 ->
    af_add arr arr1 arr2 batch

-- | Subtracts two 'Array' objects
--
-- >>> A.scalar @Int 1 `A.sub` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--         0
sub
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of sub
sub x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_sub arr arr1 arr2 1

-- | Subtracts two 'Array' objects
--
-- >>> (A.scalar @Int 1 `subBatched` A.scalar @Int 1) True
-- ArrayFire Array
-- [1 1 1 1]
--         0
subBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of sub
subBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_sub arr arr1 arr2 batch

-- | Multiply two 'Array' objects
--
-- >>> A.scalar @Int 2 `mul` A.scalar @Int 2
-- ArrayFire Array
-- [1 1 1 1]
--         4
mul
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of mul
mul x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_mul arr arr1 arr2 1

-- | Multiply two 'Array' objects
--
--
-- >>> (A.scalar @Int 2 `mulBatched` A.scalar @Int 2) True
-- ArrayFire Array
-- [1 1 1 1]
--         4
mulBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of mul
mulBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_mul arr arr1 arr2 batch

-- | Divide two 'Array' objects
--
-- >>> A.scalar @Int 6 `A.div` A.scalar @Int 3
-- ArrayFire Array
-- [1 1 1 1]
--         2
div
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of div
div x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_div arr arr1 arr2 1

-- | Divide two 'Array' objects
--
-- >>> (A.scalar @Int 6 `A.divBatched` A.scalar @Int 3) True
-- ArrayFire Array
-- [1 1 1 1]
--         2
divBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of div
divBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_div arr arr1 arr2 batch

-- | Test if on 'Array' is less than another 'Array'
--
-- >>> A.scalar @Int 1 `A.lt` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--         0
-- >>> A.scalar @Int 1 < A.scalar @Int 1
-- False
lt
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of less than
lt x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_lt arr arr1 arr2 1

-- | Test if on 'Array' is less than another 'Array'
--
-- >>> (A.scalar @Int 1 `A.ltBatched` A.scalar @Int 1) True
-- ArrayFire Array
-- [1 1 1 1]
--         0
-- >>> A.scalar @Int 1 < A.scalar @Int 1
-- False
ltBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of less than
ltBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_lt arr arr1 arr2 batch

-- | Test if an 'Array' is greater than another 'Array'
--
-- >>> A.scalar @Int 1 `A.gt` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--         0
-- >>> A.scalar @Int 1 > A.scalar @Int 2
-- False
gt
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of gt
gt x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_gt arr arr1 arr2 1

-- | Test if an 'Array' is greater than another 'Array'
--
-- >>> (A.scalar @Int 1 `gtBatched` A.scalar @Int 1) True
-- False
gtBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of gt
gtBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_gt arr arr1 arr2 batch

-- | Test if one 'Array' is less than or equal to another 'Array'
--
-- >>> A.scalar @Int 1 `A.le` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--         1
-- >>> A.scalar @Int 1 <= A.scalar @Int 1
-- False
le
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of less than or equal
le x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_le arr arr1 arr2 1

-- | Test if one 'Array' is less than or equal to another 'Array'
--
-- >>> (A.scalar @Int 1 `A.leBatched` A.scalar @Int 1) True
-- ArrayFire Array
-- [1 1 1 1]
--         1
-- >>> A.scalar @Int 1 <= A.scalar @Int 1
-- True
leBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of less than or equal
leBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_le arr arr1 arr2 batch

-- | Test if one 'Array' is greater than or equal to another 'Array'
--
-- >>> A.scalar @Int 1 `A.ge` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--         1
-- >>> A.scalar @Int 1 >= A.scalar @Int 1
-- True
ge
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of greater than or equal
ge x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_ge arr arr1 arr2 1

-- | Test if one 'Array' is greater than or equal to another 'Array'
--
--
-- >>> (A.scalar @Int 1 `A.geBatched` A.scalar @Int 1) True
--
geBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of greater than or equal
geBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_ge arr arr1 arr2 batch

-- | Test if one 'Array' is equal to another 'Array'
--
-- >>> A.scalar @Int 1 `A.eq` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--         1
--
-- >>> A.scalar @Int 1 == A.scalar @Int 1
-- True
eq
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of equal
eq x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_eq arr arr1 arr2 1

-- | Test if one 'Array' is equal to another 'Array'
--
-- >>> (A.scalar @Int 1 `A.eqBatched` A.scalar @Int 1) True
--
eqBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of equal
eqBatched x y (fromIntegral . fromEnum -> batch) =
  x `op2` y $ \arr arr1 arr2 ->
    af_eq arr arr1 arr2 batch

-- | Test if one 'Array' is not equal to another 'Array'
--
-- >>> A.scalar @Int 1 `A.neq` A.scalar @Int 1
-- ArrayFire Array
--[1 1 1 1]
--         0
-- >>> A.scalar @Int 1 /= A.scalar @Int 1
-- False
neq
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of not equal
neq x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_neq arr arr1 arr2 1

-- | Test if one 'Array' is not equal to another 'Array'
--
-- >>> (A.scalar @Int 1 `A.neqBatched` A.scalar @Int 1) True
-- False
neqBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of not equal
neqBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_neq arr arr1 arr2 batch

-- | Logical 'and' one 'Array' with another
--
-- >>> A.scalar @Int 1 `A.and` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--          1
--
and
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of and
and x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_and arr arr1 arr2 1

-- | Logical 'and' one 'Array' with another
--
-- >>> (A.scalar @Int 1 `andBatched` A.scalar @Int 1) True
-- ArrayFire Array
-- [1 1 1 1]
--         1
andBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of and
andBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_and arr arr1 arr2 batch

-- | Logical 'or' one 'Array' with another
--
-- >>> A.scalar @Int 1 `A.or` A.scalar @Int 1
-- ArrayFire Array
-- [1 1 1 1]
--         1
--
or
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of or
or x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_or arr arr1 arr2 1

-- | Logical 'or' one 'Array' with another
--
--
-- >>> (A.scalar @Int 1 `A.orBatched` A.scalar @Int 1) True
-- ArrayFire Array
-- [1 1 1 1]
--         1
orBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of or
orBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_or arr arr1 arr2 batch

-- | Not the values of an 'Array'
--
-- >>> A.not (A.scalar @Int 1)
-- ArrayFire Array
-- [1 1 1 1]
--         0
not
  :: AFType a
  => Array a
  -- ^ Input 'Array'
  -> Array a
  -- ^ Result of 'not' on an 'Array'
not = flip op1 af_not

-- | Bitwise and the values in one 'Array' against another 'Array'
--
-- >>> A.bitAnd (A.scalar @Int 1) (A.scalar @Int 1)
-- ArrayFire Array
-- [1 1 1 1]
--         1
bitAnd
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of bitwise and
bitAnd x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_bitand arr arr1 arr2 1

-- | Bitwise and the values in one 'Array' against another 'Array'
--
--- >>> A.bitAndBatched (A.scalar @Int 1) (A.scalar @Int 1) True
-- ArrayFire Array
-- [1 1 1 1]
--         1
bitAndBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of bitwise and
bitAndBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_bitand arr arr1 arr2 batch

-- | Bitwise or the values in one 'Array' against another 'Array'
--
-- >>> A.bitOr (A.scalar @Int 1) (A.scalar @Int 1)
-- ArrayFire Array
-- [1 1 1 1]
--         1
bitOr
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of bit or
bitOr x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_bitor arr arr1 arr2 1

-- | Bitwise or the values in one 'Array' against another 'Array'
--
-- >>> A.bitOrBatched (A.scalar @Int 1) (A.scalar @Int 1) False
-- ArrayFire Array
-- [1 1 1 1]
--         1
bitOrBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of bit or
bitOrBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_bitor arr arr1 arr2 batch

-- | Bitwise xor the values in one 'Array' against another 'Array'
--
-- >>> A.bitXor (A.scalar @Int 1) (A.scalar @Int 1)
-- ArrayFire Array
-- [1 1 1 1]
--         0
bitXor
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of bit xor
bitXor x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_bitxor arr arr1 arr2 1

-- | Bitwise xor the values in one 'Array' against another 'Array'
--
-- >>> A.bitXorBatched (A.scalar @Int 1) (A.scalar @Int 1) False
-- ArrayFire Array
-- [1 1 1 1]
--         0
bitXorBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of bit xor
bitXorBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_bitxor arr arr1 arr2 batch

-- | Left bit shift the values in one 'Array' against another 'Array'
--
-- >>> A.bitShiftL (A.scalar @Int 1) (A.scalar @Int 1)
-- ArrayFire Array
-- [1 1 1 1]
--         2
bitShiftL
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of bit shift left
bitShiftL x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_bitshiftl arr arr1 arr2 1

-- | Left bit shift the values in one 'Array' against another 'Array'
--
-- >>> A.bitShiftLBatched (A.scalar @Int 1) (A.scalar @Int 1) False
-- ArrayFire Array
-- [1 1 1 1]
--         2
bitShiftLBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of bit shift left
bitShiftLBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_bitshiftl arr arr1 arr2 batch

-- | Right bit shift the values in one 'Array' against another 'Array'
--
-- >>> A.bitShiftR (A.scalar @Int 1) (A.scalar @Int 1)
-- ArrayFire Array
-- [1 1 1 1]
--         0
bitShiftR
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of bit shift right
bitShiftR x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_bitshiftr arr arr1 arr2 1

-- | Right bit shift the values in one 'Array' against another 'Array'
--
-- >>> A.bitShiftRBatched (A.scalar @Int 1) (A.scalar @Int 1) False
-- ArrayFire Array
-- [1 1 1 1]
--         0
bitShiftRBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of bit shift left
bitShiftRBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_bitshiftr arr arr1 arr2 batch

-- | Cast one 'Array' into another
--
--
-- >>> A.cast (A.scalar @Int 1) :: Array Double
-- ArrayFire Array
-- [1 1 1 1]
--    1.0000
cast
  :: forall a b . (AFType a, AFType b)
  => Array a
  -- ^ Input array to cast
  -> Array b
    -- ^ Result of cast
cast afArr =
  coerce $ afArr `op1` (\x y -> af_cast x y dtyp)
    where
      dtyp = afType (Proxy @ b)

-- | Find the minimum of two 'Array's
--
-- >>> A.minOf (A.scalar @Int 1) (A.scalar @Int 0)
-- ArrayFire Array
-- [1 1 1 1]
--         0
minOf
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of minimum of
minOf x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_minof arr arr1 arr2 1

-- | Find the minimum of two 'Array's
--
-- >>> A.minOfBatched (A.scalar @Int 1) (A.scalar @Int 0) False
-- ArrayFire Array
-- [1 1 1 1]
--         0
minOfBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of minimum of
minOfBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_minof arr arr1 arr2 batch

-- | Find the maximum of two 'Array's
--
-- >>> A.maxOf (A.scalar @Int 1) (A.scalar @Int 0)
-- ArrayFire Array
-- [1 1 1 1]
--         1
maxOf
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of maximum of
maxOf x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_maxof arr arr1 arr2 1

-- | Find the maximum of two 'Array's
--
-- >>> A.maxOfBatched (A.scalar @Int 1) (A.scalar @Int 0) False
-- ArrayFire Array
--[1 1 1 1]
--         1
maxOfBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of maximum of
maxOfBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_maxof arr arr1 arr2 batch

-- | Should take the clamp
--
--
-- >>> clamp (A.scalar @Int 2) (A.scalar @Int 1) (A.scalar @Int 3)
-- ArrayFire Array
-- [1 1 1 1]
--          2
--
clamp
  :: Array a
  -- ^ input
  -> Array a
  -- ^ lower bound
  -> Array a
  -- ^ upper bound
  -> Array a
  -- ^ Result of clamp
clamp a b c =
  op3 a b c $ \arr arr1 arr2 arr3 ->
    af_clamp arr arr1 arr2 arr3 1

-- | Should take the clamp
--
-- >>> (clampBatched (A.scalar @Int 2) (A.scalar @Int 1) (A.scalar @Int 3)) True
-- ArrayFire Array
-- [1 1 1 1]
--          2
clampBatched
  :: Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Third input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of clamp
clampBatched a b c (fromIntegral . fromEnum -> batch) =
  op3 a b c $ \arr arr1 arr2 arr3 ->
    af_clamp arr arr1 arr2 arr3 batch

-- | Find the remainder of two 'Array's
--
-- >>> A.rem (A.vector @Int 10 [1..]) (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--         0          0          0          0          0          0          0          0          0          0
rem
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of remainder
rem x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_rem arr arr1 arr2 1

-- | Find the remainder of two 'Array's
--
--
-- >>> A.remBatched (A.vector @Int 10 [1..])  (vector @Int 10 [1..]) True
--
remBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of remainder
remBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_rem arr arr1 arr2 batch

-- | Take the 'mod' of two 'Array's
--
-- >>> A.mod (A.vector @Int 10 [1..]) (A.vector @Int 10 [1..])
-- ArrayFire Array
--[10 1 1 1]
--         0          0          0          0          0          0          0          0          0          0
mod
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of mod
mod x y = do
  x `op2` y $ \arr arr1 arr2 ->
    af_mod arr arr1 arr2 1

-- | Take the 'mod' of two 'Array's
--
-- >>> A.modBatched (vector @Int 10 [1..]) (vector @Int 10 [1..]) True
-- ArrayFire Array
-- [10 1 1 1]
--         0          0          0          0          0          0          0          0          0          0
modBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of mod
modBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_mod arr arr1 arr2 batch

-- | Take the absolute value of an array
--
-- >>> A.abs (A.scalar @Int (-1))
-- ArrayFire Array
-- [1 1 1 1]
--    1.0000
--
abs
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'abs'
abs = flip op1 af_abs

-- | Find the arg of an array
--
-- >>> A.arg (vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--         0          0          0          0          0          0          0          0          0          0
arg
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'arg'
arg = flip op1 af_arg

-- | Find the sign of two 'Array's
--
-- >>> A.sign (vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--   0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000
sign
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'sign'
sign = flip op1 af_sign

-- | Round the values in an 'Array'
--
-- >>> A.round (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     2.0000     3.0000     4.0000     5.0000     6.0000     7.0000     8.0000     9.0000    10.0000
round
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'round'
round = flip op1 af_round

-- | Truncate the values of an 'Array'
--
-- >>> A.trunc (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     2.0000     3.0000     4.0000     5.0000     6.0000     7.0000     8.0000     9.0000    10.0000
trunc
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'trunc'
trunc = flip op1 af_trunc

-- | Take the floor of all values in an 'Array'
--
-- >>> A.floor (A.vector @Int 10 [10,9..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     2.0000     3.0000     4.0000     5.0000     6.0000     7.0000     8.0000     9.0000    10.0000
floor
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'floor'
floor = flip op1 af_floor

-- | Take the ceil of all values in an 'Array'
--
-- >>> A.ceil (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     2.0000     3.0000     4.0000     5.0000     6.0000     7.0000     8.0000     9.0000    10.0000
ceil
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'ceil'
ceil = flip op1 af_ceil

-- | Take the sin of all values in an 'Array'
--
-- >>> A.sin (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.8415     0.9093     0.1411    -0.7568    -0.9589    -0.2794     0.6570     0.9894     0.4121    -0.5440
sin
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'sin'
sin = flip op1 af_sin

-- | Take the cos of all values in an 'Array'
--
-- >>> A.cos (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.5403    -0.4161    -0.9900    -0.6536     0.2837     0.9602     0.7539    -0.1455    -0.9111    -0.8391
cos
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'cos'
cos = flip op1 af_cos

-- | Take the tan of all values in an 'Array'
--
-- >>> A.tan (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.5574    -2.1850    -0.1425     1.1578    -3.3805    -0.2910     0.8714    -6.7997    -0.4523     0.6484
tan
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'tan'
tan = flip op1 af_tan

-- | Take the asin of all values in an 'Array'
--
-- >>> A.asin (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.5708        nan        nan        nan        nan        nan        nan        nan        nan        nan
--
asin
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'asin'
asin = flip op1 af_asin

-- | Take the acos of all values in an 'Array'
--
-- >>> A.acos (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.0000        nan        nan        nan        nan        nan        nan        nan        nan        nan
acos
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'acos'
acos = flip op1 af_acos

-- | Take the atan of all values in an 'Array'
--
-- >>> A.atan (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.7854     1.1071     1.2490     1.3258     1.3734     1.4056     1.4289     1.4464     1.4601     1.4711
atan
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'atan'
atan = flip op1 af_atan

-- | Take the atan2 of all values in an 'Array'
--
-- >>> A.atan2 (A.vector @Int 10 [1..]) (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.7854     0.7854     0.7854     0.7854     0.7854     0.7854     0.7854     0.7854     0.7854     0.7854
atan2
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of atan2
atan2 x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_atan2 arr arr1 arr2 1

-- | Take the atan2 of all values in an 'Array'
--
-- >>> A.atan2Batched (A.vector @Int 10 [1..]) (A.vector @Int 10 [1..]) True
-- ArrayFire Array
-- [10 1 1 1]
--    0.7854     0.7854     0.7854     0.7854     0.7854     0.7854     0.7854     0.7854     0.7854     0.7854
atan2Batched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of atan2
atan2Batched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_atan2 arr arr1 arr2 batch

-- | Take the cplx2 of all values in an 'Array'
--
-- >>> A.cplx2 (A.vector @Int 10 [1..]) (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--         (1.0000,1.0000)          (2.0000,2.0000)          (3.0000,3.0000)          (4.0000,4.0000)          (5.0000,5.0000)          (6.0000,6.0000)          (7.0000,7.0000)          (8.0000,8.0000)          (9.0000,9.0000)          (10.0000,10.0000)
cplx2
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of cplx2
cplx2 x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_cplx2 arr arr1 arr2 1

-- | Take the cplx2Batched of all values in an 'Array'
--
-- >>> A.cplx2Batched (A..vector @Int 10 [1..]) (A.vector @Int 10 [1..]) True
-- ArrayFire Array
-- [10 1 1 1]
--         (1.0000,1.0000)          (2.0000,2.0000)          (3.0000,3.0000)          (4.0000,4.0000)          (5.0000,5.0000)          (6.0000,6.0000)          (7.0000,7.0000)          (8.0000,8.0000)          (9.0000,9.0000)          (10.0000,10.0000)
cplx2Batched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of cplx2
cplx2Batched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_cplx2 arr arr1 arr2 batch

-- | Execute cplx
--
-- >>> A.cplx (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--         (1.0000,0.0000)          (2.0000,0.0000)          (3.0000,0.0000)          (4.0000,0.0000)          (5.0000,0.0000)          (6.0000,0.0000)          (7.0000,0.0000)          (8.0000,0.0000)          (9.0000,0.0000)          (10.0000,0.0000)
cplx
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'atan'
cplx = flip op1 af_cplx

-- | Execute real
--
-- >>> A.real (A.scalar @(Complex Double) (10 :+ 11)) :: Array Double
-- ArrayFire Array
-- [10 1 1 1]
--    10.0000
real
  :: (AFType a, AFType (Complex b), RealFrac a, RealFrac b)
  => Array (Complex b)
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'real'
real = flip op1d af_real

-- | Execute imag
--
-- >>> A.imag (A.scalar @(Complex Double) (10 :+ 11)) :: Array Double
-- ArrayFire Array
-- [10 1 1 1]
--    11.0000
imag
  :: (AFType a, AFType (Complex b), RealFrac a, RealFrac b)
  => Array (Complex b)
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'imag'
imag = flip op1d af_imag

-- | Execute conjg
--
-- >>> A.conjg (A.vector @Double 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     2.0000     3.0000     4.0000     5.0000     6.0000     7.0000     8.0000     9.0000    10.0000
conjg
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'conjg'
conjg = flip op1 af_conjg

-- | Execute sinh
--
-- >>> A.sinh (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.1752     3.6269    10.0179    27.2899    74.2032   201.7132   548.3161  1490.4788  4051.5419 11013.2329
sinh
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'sinh'
sinh = flip op1 af_sinh

-- | Execute cosh
--
--
-- >>> A.cosh (A.vector @Double 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.5431     3.7622    10.0677    27.3082    74.2099   201.7156   548.3170  1490.4792  4051.5420 11013.2329
cosh
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'cosh'
cosh = flip op1 af_cosh

-- | Execute tanh
--
-- >>> A.tanh (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.7616     0.9640     0.9951     0.9993     0.9999     1.0000     1.0000     1.0000     1.0000     1.0000
tanh
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'tanh'
tanh = flip op1 af_tanh

-- | Execute asinh
--
-- >>> A.asinh (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.8814     1.4436     1.8184     2.0947     2.3124     2.4918     2.6441     2.7765     2.8934     2.9982
asinh
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'tanh'
asinh = flip op1 af_asinh

-- | Execute acosh
--
-- >>> A.acosh (A.vector @Double 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.0000     1.3170     1.7627     2.0634     2.2924     2.4779     2.6339     2.7687     2.8873     2.9932
acosh
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'tanh'
acosh = flip op1 af_acosh

-- | Execute atanh
--
-- >>> A.atanh (A.vector @Double 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--        inf        nan        nan        nan        nan        nan        nan        nan        nan        nan
atanh
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'tanh'
atanh = flip op1 af_atanh

-- | Execute root
--
-- >>> A.root (A.vector @Double 10 [1..]) (A.vector @Double 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     1.4142     1.4422     1.4142     1.3797     1.3480     1.3205     1.2968     1.2765     1.2589
root
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of root
root x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_root arr arr1 arr2 1

-- | Execute rootBatched
--
-- >>> A.rootBatched (vector @Double 10 [1..]) (vector @Double 10 [1..]) True
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     1.4142     1.4422     1.4142     1.3797     1.3480     1.3205     1.2968     1.2765     1.2589
rootBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of root
rootBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_root arr arr1 arr2 batch

-- | Execute pow
--
-- >>> A.pow (A.vector @Int 10 [1..]) 2
-- ArrayFire Array
-- [10 1 1 1]
--         1          4          9         16         25         36         49         64         81        100
pow
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Array a
  -- ^ Result of pow
pow x y =
  x `op2` y $ \arr arr1 arr2 ->
    af_pow arr arr1 arr2 1

-- | Execute powBatched
--
-- >>> A.powBatched (A.vector @Int 10 [1..]) (A.constant @Int [1] 2) True
-- ArrayFire Array
-- [10 1 1 1]
--         1          4          9         16         25         36         49         64         81        100
powBatched
  :: AFType a
  => Array a
  -- ^ First input
  -> Array a
  -- ^ Second input
  -> Bool
  -- ^ Use batch
  -> Array a
  -- ^ Result of powBatched
powBatched x y (fromIntegral . fromEnum -> batch) = do
  x `op2` y $ \arr arr1 arr2 ->
    af_pow arr arr1 arr2 batch

-- | Raise an 'Array' to the second power
--
-- >>> A.pow2 (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--         2          4          8         16         32         64        128        256        512       1024
pow2
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'pow2'
pow2 = flip op1 af_pow2

-- | Execute exp on 'Array'
--
-- >>> A.exp (A.vector @Double 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    2.7183     7.3891    20.0855    54.5982   148.4132   403.4288  1096.6332  2980.9580  8103.0839 22026.4658
exp
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'exp'
exp = flip op1 af_exp

-- | Execute sigmoid on 'Array'
--
-- >>> A.sigmoid (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.7311     0.8808     0.9526     0.9820     0.9933     0.9975     0.9991     0.9997     0.9999     1.0000
sigmoid
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'sigmoid'
sigmoid = flip op1 af_sigmoid

-- | Execute expm1
--
-- >>> A.expm1 (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.7183     6.3891    19.0855    53.5981   147.4132   402.4288  1095.6332  2979.9580  8102.0840 22025.4648
expm1
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'expm1'
expm1 = flip op1 af_expm1

-- | Execute erf
--
-- >>> A.erf (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.8427     0.9953     1.0000     1.0000     1.0000     1.0000     1.0000     1.0000     1.0000     1.0000
erf
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'erf'
erf = flip op1 af_erf

-- | Execute erfc
--
-- >>> A.erfc (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.1573     0.0047     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000     0.0000
erfc
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'erfc'
erfc = flip op1 af_erfc

-- | Execute log
--
-- >>> A.log (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.0000     0.6931     1.0986     1.3863     1.6094     1.7918     1.9459     2.0794     2.1972     2.3026
log
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'log'
log = flip op1 af_log

-- | Execute log1p
--
-- >>> A.log1p (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.6931     1.0986     1.3863     1.6094     1.7918     1.9459     2.0794     2.1972     2.3026     2.3979
log1p
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'log1p'
log1p = flip op1 af_log1p

-- | Execute log10
--
-- >>> A.log10 (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.0000     0.3010     0.4771     0.6021     0.6990     0.7782     0.8451     0.9031     0.9542     1.0000
log10
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'log10'
log10 = flip op1 af_log10

-- | Execute log2
--
-- >>> A.log2 (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.0000     1.0000     1.5850     2.0000     2.3219     2.5850     2.8074     3.0000     3.1699     3.3219
log2
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'log2'
log2 = flip op1 af_log2

-- | Execute sqrt
--
-- >>> A.sqrt (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     1.4142     1.7321     2.0000     2.2361     2.4495     2.6458     2.8284     3.0000     3.1623
sqrt
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'sqrt'
sqrt = flip op1 af_sqrt

-- | Execute cbrt
--
-- >>> A.cbrt (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    1.0000     1.2599     1.4422     1.5874     1.7100     1.8171     1.9129     2.0000     2.0801     2.1544
cbrt
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'cbrt'
cbrt = flip op1 af_cbrt

-- | Execute factorial1
--
-- >>> A.factorial1 (A.vector @Int 10 [1..])
--
factorial
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'factorial'
factorial = flip op1 af_factorial

-- | Execute tgamma
--
--
-- >>> 'tgamma' (vector @Int 10 [1..])
--
tgamma
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'tgamma'
tgamma = flip op1 af_tgamma

-- | Execute lgamma
--
-- >>> A.lgamma (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--    0.0000     0.0000     0.6931     1.7918     3.1781     4.7875     6.5793     8.5252    10.6046    12.8018
lgamma
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'lgamma'
lgamma = flip op1 af_lgamma

-- | Execute isZero
--
-- >>> A.isZero (A.vector @CBool 10 (repeat 0))
-- ArrayFire Array
-- [10 1 1 1]
--         1          1          1          1          1          1          1          1          1          1
isZero
  :: AFType a
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Result of calling 'isZero'
isZero = (`op1` af_iszero)

-- | Execute isInf
--
-- >>> A.isInf (A.vector @Double 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--         0          0          0          0          0          0          0          0          0          0
isInf
  :: (Real a, AFType a)
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ will contain 1's where input is Inf or -Inf, and 0 otherwise.
isInf = (`op1` af_isinf)

-- | Execute isNaN
--
-- >>> A.isNaN $ A.acos (A.vector @Int 10 [1..])
-- ArrayFire Array
-- [10 1 1 1]
--         0          1          1          1          1          1          1          1          1          1
isNaN
  :: forall a. (AFType a, Real a)
  => Array a
  -- ^ Input array
  -> Array a
  -- ^ Will contain 1's where input is NaN, and 0 otherwise.
isNaN = (`op1` af_isnan)


module ForeignHABQT where

import Data.Bifunctor (bimap)
import Data.Complex (Complex(..))
import Data.List (unzip)
import qualified Data.Vector as V
import Foreign
import Foreign.C
import HABQTlib.Data
import HABQTlib.MeasurementProcessing (PurePOVM)
import qualified Numeric.LinearAlgebra as LA

tdm :: DensityMatrix
tdm = DensityMatrix $ LA.real $ (2 LA.>< 2) [1 ..]

type CMatrix = [[CDouble]]

convertDM :: DensityMatrix -> (CMatrix, CMatrix)
convertDM (DensityMatrix dm) =
  let (rp, ip) = LA.fromComplex dm
      conv = (map . map) realToFrac . LA.toLists
   in (conv rp, conv ip)

unmarshallSV :: Dim -> Ptr CDouble -> Ptr CDouble -> IO PureStateVector
unmarshallSV dim rPtr iPtr = do
  rl <- peekArray dim rPtr
  il <- peekArray dim iPtr
  let toHM = (dim LA.>< 1) . map realToFrac
  return . PureStateVector $ LA.toComplex (toHM rl, toHM il)

convertWPSV :: WeighedPureStateVector -> ([CDouble], [CDouble])
convertWPSV (WeighedPureStateVector (w, PureStateVector sv)) =
  let sv' = LA.scale (sqrt w :+ 0) sv
      (rp, ip) = LA.fromComplex sv'
      conv = map realToFrac . LA.toList . head . LA.toColumns
   in (conv rp, conv ip)

convertPOVM :: PurePOVM -> ([CDouble], [CDouble])
convertPOVM povm =
  let svs = V.toList povm
      csvs = map convertWPSV svs
   in bimap concat concat . unzip $ csvs

marshallPOVM :: PurePOVM -> Ptr CDouble -> Ptr CDouble -> IO ()
marshallPOVM povm rPtr iPtr = do
  let (r, i) = convertPOVM povm
  pokeArray rPtr r
  pokeArray iPtr i

marshallDM :: DensityMatrix -> Ptr CDouble -> Ptr CDouble -> IO ()
marshallDM dm rPtr iPtr = do
  let (r, i) = convertDM dm
  pokeArray rPtr $ concat r
  pokeArray iPtr $ concat i


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module IHaskell.Display.Spiral where

import Data.Complex.Cyclotomic (Cyclotomic)
import qualified Data.Text as T
import IHaskell.Display
import Text.LaTeX

import Spiral.Array
import Spiral.Exp
import Spiral.SPL
import Spiral.Util.Pretty.LaTeX

displayLaTeX :: Pretty a => a -> IO Display
displayLaTeX = display . IHaskell.Display.latex . T.unpack . render . math . ppr

instance IHaskellDisplay Cyclotomic where
  display = displayLaTeX

instance IHaskellDisplay (Const a) where
  display = displayLaTeX

instance IHaskellDisplay (Exp a) where
  display = displayLaTeX

instance (Pretty a, IArray r DIM2 a) => IHaskellDisplay (Matrix r a) where
  display = displayLaTeX

instance (Num a, Pretty a) => IHaskellDisplay (SPL a) where
  display = displayLaTeX


module Isumi.Math.FFT.Internal.Bruun2D
  ( fft2dBruun
  ) where

import           Data.Complex
import qualified Data.Vector.Generic           as GV
import           Isumi.Math.FFT.Internal.Bruun (fftBruun, fftBruunC)
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data    (Matrix)

fft2dBruun :: Matrix Double -> Maybe (Matrix (Complex Double))
fft2dBruun m = do
  rows' <- traverse (fftBruun . GV.convert) (LA.toRows m)
  cols' <- traverse (fftBruunC . GV.convert) . LA.toColumns . LA.fromRows
           $ fmap GV.convert rows'
  pure $ LA.fromColumns (fmap GV.convert cols')


{-# 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]"


{-|
Module: MachineLearning.Types
Description: Common Types
Copyright: (c) Alexander Ignatyev, 2016
License: BSD-3
Stability: experimental
Portability: POSIX

Common type definitions used in all modules.
-}

module MachineLearning.Types
(
    R
  , Vector
  , Matrix
)

where

import qualified Numeric.LinearAlgebra.Data as LAD

-- | Scalar Type (Double)
type R = LAD.R

-- | Vector Types (of Doubles)
type Vector = LAD.Vector R

-- | Matrix Types (of Doubles)
type Matrix = LAD.Matrix R


{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE UndecidableInstances   #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Conversion
-- Copyright   :  (c) Alberto Ruiz 2010
-- License     :  BSD3
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--
-- Conversion routines
--
-----------------------------------------------------------------------------

module Internal.Conversion (
    Complexable(..), RealElement,
    module Data.Complex
) where

import           Control.Arrow       ((***))
import           Data.Complex
import           Internal.Matrix
import           Internal.Vector
import           Internal.Vectorized

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

-- | Supported single-double precision type pairs
class (Element s, Element d) => Precision s d | s -> d, d -> s where
    double2FloatG :: Vector d -> Vector s
    float2DoubleG :: Vector s -> Vector d

-- instance Precision Float Double where
--     double2FloatG = double2FloatV
--     float2DoubleG = float2DoubleV

-- instance Precision (Complex Float) (Complex Double) where
--     double2FloatG = asComplex . double2FloatV . asReal
--     float2DoubleG = asComplex . float2DoubleV . asReal

instance Precision I Z where
    double2FloatG = long2intV
    float2DoubleG = int2longV

-- | Supported real types
class (Element t, Element (Complex t), RealFloat t)
    => RealElement t

-- instance RealElement Double
instance RealElement Float

-- | Structures that may contain complex numbers
class Complexable c where
    toComplex'   :: (RealElement e) => (c e, c e) -> c (Complex e)
    fromComplex' :: (RealElement e) => c (Complex e) -> (c e, c e)
    comp'        :: (RealElement e) => c e -> c (Complex e)
    single'      :: Precision a b => c b -> c a
    double'      :: Precision a b => c a -> c b

instance Complexable Vector where
    toComplex' = toComplexV
    fromComplex' = fromComplexV
    comp' v = toComplex' (v,constantD 0 (dim v))
    single' = double2FloatG
    double' = float2DoubleG

-- | creates a complex vector from vectors with real and imaginary parts
toComplexV :: (RealElement a) => (Vector a, Vector a) ->  Vector (Complex a)
toComplexV (r,i) = asComplex $ flatten $ fromColumns [r,i]

-- | the inverse of 'toComplex'
fromComplexV :: (RealElement a) => Vector (Complex a) -> (Vector a, Vector a)
fromComplexV z = (r,i) where
    [r,i] = toColumns $ reshape 2 $ asReal z

instance Complexable Matrix where
    toComplex' = uncurry $ liftMatrix2 $ curry toComplex'
    fromComplex' z = (reshape c *** reshape c) . fromComplex' . flatten $ z
        where c = cols z
    comp' = liftMatrix comp'
    single' = liftMatrix single'
    double' = liftMatrix double'


{-# 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))


module Signal where

import Control.Monad.Fix
import Data.Complex

type Time = Double
type Metric = Complex Double
newtype Signal a = Signal {at :: Metric -> a}

instance Functor Signal where
    fmap f (Signal a) = Signal $ \time -> f (a time)

instance Applicative Signal where
    pure x = Signal $ const x
    sf <*> (Signal a) = Signal $ \time -> (sf `at` time) (a time)

instance Monad Signal where
    return = pure
    sa >>= f = Signal $ \time -> f (sa `at` time) `at` time

instance MonadFix Signal where
    mfix f = Signal $ \time -> fix ((`at` time) . f)

instance Eq a => Eq (Signal a) where
    (Signal x) == (Signal y) = x 0 == y 0

instance Ord a => Ord (Signal a) where
    (Signal x) <= (Signal y) = x 0 <= y 0

instance Num a => Num (Signal a) where
    sx + sy = Signal $ \time -> (sx `at` time) + (sy `at` time)
    sx * sy = Signal $ \time -> (sx `at` time) * (sy `at` time)
    abs sx = Signal $ \time -> abs $ sx `at` time
    fromInteger i = Signal $ const $ fromInteger i
    signum sx = Signal $ \time -> signum $ sx `at` time
    negate sx = Signal $ \time -> negate $ sx `at` time

instance Fractional a => Fractional (Signal a) where
    sx / sy = Signal $ \time -> (sx `at` time) / (sy `at` time)
    fromRational i = Signal $ const $ fromRational i

instance Show a => Show (Signal a) where
    show (Signal a) = show $ a 0


module Wavefront.Converter where

import qualified Types as T
import Wavefront.Types
import Numeric.Vector
import VecUtil

face2Tri :: [Vertex_] -> [TexCoord_] -> [Normal_] -> Face_ -> T.Tri
face2Tri vs tcs ns (Face_ p1 tc1 n1 p2 tc2 n2 p3 tc3 n3) =
  let (t1, b1) = tangentBitangent $ normalToVec $ ns !! n1
      (t2, b2) = tangentBitangent $ normalToVec $ ns !! n2
      (t3, b3) = tangentBitangent $ normalToVec $ ns !! n3
      v1 = T.Vertex (vertexToVec $ vs !! p1) (vec3 1 1 1) (texCoordToVec $ tcs !! tc1) (normalToVec $ ns !! n1) t1 b1
      v2 = T.Vertex (vertexToVec $ vs !! p2) (vec3 1 1 1) (texCoordToVec $ tcs !! tc2) (normalToVec $ ns !! n2) t2 b2
      v3 = T.Vertex (vertexToVec $ vs !! p3) (vec3 1 1 1) (texCoordToVec $ tcs !! tc3) (normalToVec $ ns !! n3) t3 b3
  in  T.Tri v1 v2 v3

obj2Mesh :: [WavefrontLine] -> [T.Tri]
obj2Mesh ls = let verts = map vertexFromLine $ filter isVertex ls
                  texCoords = map texCoordFromLine $ filter isTexCoord ls
                  normals = map normalFromLine $ filter isNormal ls
                  faces = map faceFromLine $ filter isFace ls
              in  map (face2Tri verts texCoords normals) faces


{-# LANGUAGE BangPatterns #-}
-- import System.Environment
import Data.Array.Repa as Repa
import Data.Array.Repa.Eval as Eval
import Data.Array.Repa.IO.BMP as BMP
import Data.Array.Repa.Algorithms.Pixel as Pixel
import Data.Vector.Unboxed.Base as Unboxed

import Control.Monad.Identity

import GHC.Word
import Data.Complex
import Data.List
import Julia
import Rendering
import RealFunctions
import ComplexFunctions

width :: (Num a) => a
width = 600

height :: (Num a) => a
height = 600

pixelRenderer :: (RealFloat a) => Int -> Int -> a
pixelRenderer x y = smooth samples -- avg' (Prelude.map (f) samples)
    where samples = sample 128 (magnitudeLimit 2) dragon start
          start   = afold (remapped x y)
          afold   = (1.5*(cis(pi/(4)))*).(foldPhase (pi/6))
          f z     = cos (atan2 (imagPart z) (realPart z)) --(cos . abs . phase) z
          f1 x z  = 1 + (z*x)
          f2 x z  = 1 - (z*x)

paramRenderer :: (RealFloat a) => a -> Int -> Int -> a
paramRenderer progress x y = smooth samples -- progress between 0 and 100%
    where samples = sample 128 (magnitudeLimit 2) (dragon) start
          start   = ((cis n *).(foldPhase (pi/6))) (remapped x y)
          n       = progress * 2 * pi

remapped :: (RealFloat a) => Int -> Int -> Complex a
remapped x y = (remap (fromIntegral x :+ fromIntegral y) width height)

avg :: (Num a, Fractional a) => [a] -> a
avg l = (sum l) / ((fromIntegral . length) l)

avg' :: (Num a, Fractional a) => [a] -> a
avg' = uncurry (/) . foldl' (\(s,n) x -> (s+x, n+1)) (0, 0)

-- 2D shaped array
shape :: DIM2
shape = Z :. (height :: Int) :. (width :: Int)

canvas :: Array U DIM2 Double
canvas = fromListUnboxed shape (take (fromIntegral (size shape)) (repeat 0))

render :: (Source r a, RealFloat a, Elt a) => Array r DIM2 a -> Array D DIM2 a
render !input = traverse input id (\_ (Z :. y :. x) -> pixelRenderer x y)

-- render a frame with progress between 0% and 100%
renderFrame :: (Source r a, RealFloat a, Elt a) => Array r DIM2 a -> a -> Array D DIM2 a
renderFrame !input progress = traverse input id (\_ (Z :. y :. x) -> paramRenderer progress x y)

convert :: (Shape sh) => Array U sh Double
     -> Array D sh (GHC.Word.Word8, GHC.Word.Word8, GHC.Word.Word8)
convert !array = Repa.map (rgb8OfGreyDouble) array

getMinimum :: (Ord b, Num b, Source r b, Elt b, Unbox b, Shape sh) =>
    Array r sh b -> IO (b)
getMinimum !array = foldAllP (min) 100000 array

getMaximum :: (Ord b, Num b, Source r b, Elt b, Unbox b, Shape sh) =>
    Array r sh b -> IO (b)
getMaximum !array = foldAllP (max) 0 array

correctExpo :: (Source r b, Unbox b, Elt b, Floating b, Shape sh, RealFrac b) =>
    Array r sh b -> b -> b -> IO (Array U sh b)
correctExpo !array !mini !maxi = computeP $ Repa.map ((map' mini maxi 0.0 1.0)) array
--
compute :: IO ()
compute = do
    img <- computeP (render canvas) :: IO (Array U DIM2 Double) -- super intensive
    mini <- getMinimum img
    maxi <- getMaximum img
    corrected <- correctExpo img mini maxi :: IO (Array U DIM2 Double)
    image <- computeP (convert corrected)
    writeImageToBMP "out.bmp" image

computeFrame :: Double -> Int -> IO ()
computeFrame progress frameNum = do
    img <- computeP (renderFrame canvas progress) :: IO (Array U DIM2 Double) -- super intensive
    mini <- getMinimum img
    maxi <- getMaximum img
    corrected <- correctExpo img mini maxi :: IO (Array U DIM2 Double)
    image <- computeP (convert corrected)
    let filename = "gifs/" Data.List.++ (show frameNum) Data.List.++ ".bmp"
    writeImageToBMP filename image

main :: IO ()
main = compute -- foldl (>>) (return ()) (Data.List.zipWith3 ($) (repeat computeFrame) progresses [1..])
    where progresses = takeWhile (<= 1.0) $ iterate (+0.005) 0.0


-- $ ghc --make -O -rtsopts -threaded parallel.hs
-- $ ./parallel 3000 +RTS -N4 -s -A200M

import System.Environment(getArgs)
import Numeric.LinearAlgebra
import Control.Parallel.Strategies
import System.Time

inParallel = parMap rwhnf id

-- matrix product decomposed into p parallel subtasks
parMul p x y = fromBlocks [ inParallel ( map (x <>) ys ) ]
    where [ys] = toBlocksEvery (rows y) (cols y `div` p) y

main = do
    n <- (read . head) `fmap` getArgs
    let m = ident n :: Matrix Double
    time $ print $ maxElement $ takeDiag $ m <> m
    time $ print $ maxElement $ takeDiag $ parMul 2 m m
    time $ print $ maxElement $ takeDiag $ parMul 4 m m
    time $ print $ maxElement $ takeDiag $ parMul 8 m m

time act = do
    t0 <- getClockTime
    act
    t1 <- getClockTime
    print $ tdSec $ normalizeTimeDiff $ diffClockTimes t1 t0


module Numeric.LinearAlgebra.Matrix.Class where

-- import Numeric.LinearAlgebra.Vector

class Functor m => Matrix m where
  mDim     :: m a -> Int
  mElement :: m a -> Int -> Int -> a
  mIndexOf :: (Ord a) => (a -> a -> Bool) -> m a -> (Int, Int)
  mZip     :: (a -> b -> c) -> m a -> m b -> m c
  -- | mFold is foldl1'
  mFold    :: (a -> a -> a) -> m a -> a
  det      :: Num a => m a -> a

{-# INLINE mApply #-}
mApply :: Functor f => f (a -> b) -> a -> f b
mApply f m = fmap ($ m) f

(.+.) :: (Num k, Matrix m) => m k -> m k -> m k
(.+.) = mZip (+)
(.-.) :: (Num k, Matrix m) => m k -> m k -> m k
(.-.) = mZip (-)

{-
(.*.) :: (k ~ Element m, Num k, Matrix m) => m -> m -> m
m .*. n = mIdxMap m $ \i j -> sum [ mElement m i k * mElement n k j | k <- [ 0 .. 3 ] ]
-}
{-
(.*>) :: (k ~ Element m, k ~ Scalar v, Num k, Matrix m, Vector v)
      => m -> v -> v
m .*> v | vDim v == mDim m = flip vIdxMap v $ \k -> sum [ mElement m i k * vElement v k | i <- [ 0 .. mDim m ] ]
        | otherwise        = error "Dimensions do not match"
-}
-- (*.) :: (k ~ Element m, Num k, Matrix m) => m -> k -> m
-- m *. k = mMap (k*) m

{-
transpose :: Matrix m => m -> m
transpose m = flip mIdxMap m $ \i j -> mElement j i m
-}


{-# LANGUAGE DeriveFunctor, DeriveDataTypeable,DeriveGeneric  #-}
module Statistics.Test.Types (
    Test(..)
  , isSignificant
  , TestResult(..)
  , significant
  , PositionTest(..)
  ) where

import Control.DeepSeq  (NFData(..))
import Control.Monad    (liftM3)
import Data.Aeson       (FromJSON, ToJSON)
import Data.Binary      (Binary (..))
import Data.Data (Typeable, Data)
import GHC.Generics

import Statistics.Types (PValue)

-- | Result of hypothesis testing
data TestResult = Significant    -- ^ Null hypothesis should be rejected
                | NotSignificant -- ^ Data is compatible with hypothesis
                  deriving (Eq,Ord,Show,Typeable,Data,Generic)

instance Binary   TestResult where
  get = do
      sig <- get
      if sig then return Significant else return NotSignificant
  put = put . (== Significant)
instance FromJSON TestResult
instance ToJSON   TestResult
instance NFData   TestResult

-- | Result of statistical test.
data Test distr = Test
  { testSignificance :: !(PValue Double)
    -- ^ Probability of getting value of test statistics at least as
    --   extreme as measured.
  , testStatistics   :: !Double
    -- ^ Statistic used for test.
  , testDistribution :: distr
    -- ^ Distribution of test statistics if null hypothesis is correct.
  }
  deriving (Eq,Ord,Show,Typeable,Data,Generic,Functor)

instance (Binary   d) => Binary   (Test d) where
  get = liftM3 Test get get get
  put (Test sign stat distr) = put sign >> put stat >> put distr
instance (FromJSON d) => FromJSON (Test d)
instance (ToJSON   d) => ToJSON   (Test d)
instance (NFData   d) => NFData   (Test d) where
  rnf (Test _ _ a) = rnf a

-- | Check whether test is significant for given p-value.
isSignificant :: PValue Double -> Test d -> TestResult
isSignificant p t
  = significant $ p >= testSignificance t

-- | Test type for test which compare positional (mean,median etc.)
--   information of samples.
data PositionTest
  = SamplesDiffer
    -- ^ Test whether samples differ in position. Null hypothesis is
    --   samples are not different
  | AGreater
    -- ^ Test if first sample (A) is larger than second (B). Null
    --   hypothesis is first sample is not larger than second.
  | BGreater
    -- ^ Test if second sample is larger than first.
  deriving (Eq,Ord,Show,Typeable,Data,Generic)

instance Binary   PositionTest where
  get = do
    i <- get
    case (i :: Int) of
      0 -> return SamplesDiffer
      1 -> return AGreater
      2 -> return BGreater
      _ -> fail "Invalid PositionTest"
  put SamplesDiffer = put (0 :: Int)
  put AGreater      = put (1 :: Int)
  put BGreater      = put (2 :: Int)
instance FromJSON PositionTest
instance ToJSON   PositionTest
instance NFData   PositionTest

-- | significant if parameter is 'True', not significant otherwise
significant :: Bool -> TestResult
significant True  = Significant
significant False = NotSignificant


{-# 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      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}

-- | TODO
module SatOpt.Utility where

import           Data.Packed.Matrix            (buildMatrix, liftMatrix2)
import           Data.Packed.Vector            (dim, zipVectorWith)
import           Numeric.LinearAlgebra.Data
import           Numeric.LinearAlgebra.HMatrix

zipMatrixWith ::(Element a, Element b, Element c)
                => (a -> b -> c)
                -> Matrix a
                -> Matrix b
                -> Matrix c
zipMatrixWith f = liftMatrix2 (zipVectorWith f)

sample2D :: (a -> a -> b) -> [a] -> [a] -> [[b]]
sample2D f xs ys = [[f x y | y <- ys] | x <- xs]

sample2DM :: (Element a, Element b, Indexable (Vector a) a)
             => (a -> a -> b)
             -> Vector a
             -> Vector a
             -> Matrix b
sample2DM f vx vy = buildMatrix (dim vx) (dim vy) (\(i, j) -> f (vx ! i) (vy ! j))

square :: Num a => a -> a
square x = x * x

clamp :: Double -> Double
clamp x
  | x > 1     = 1
  | x < 0     = 0
  | otherwise = x

norm :: Floating a => a -> a -> a
norm x y = sqrt $ square x + square y


module GPLVM.Types
       ( Covariance
       , Distribution
       , EigenVectors
       , InputObservations (..)
       , LatentSpacePoints (..)
       , Matrix (..)
       , Mean (..)
       , ObservedData (..)
       , Vector (..)
       , unLatentSpacePoints
       , unObservedData
       , unInputObs
       ) where

import Universum hiding (Vector, transpose)

import Control.Lens (makeLenses)

import Data.Array.Repa
import Data.Array.Repa.Repr.ForeignPtr
-- import Foreign.Storable
import Numeric.LinearAlgebra.Repa hiding (Matrix, Vector)

type Distribution r a = Matrix r a -> a

type Mean r a = Matrix r a -> a

type Covariance a = Herm a

type Matrix r a = Array r DIM2 a

type Vector r a = Array r DIM1 a

type EigenVectors a = Matrix D a

-- Real matrix, represented as unboxed vector.
-------

newtype LatentSpacePoints a = LatentSpacePoints { _unLatentSpacePoints :: Matrix D a }
    deriving Eq

makeLenses ''LatentSpacePoints

newtype InputObservations a = InputObservations { _unInputObs :: Vector D a }
    deriving Eq

newtype ObservedData a = ObservedData { _unObservedData :: Matrix D a }
    deriving Eq

makeLenses ''ObservedData
makeLenses ''InputObservations


{-# 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