text:
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE FlexibleContexts         #-}
{-# LANGUAGE FlexibleInstances        #-}
{-# LANGUAGE BangPatterns             #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Packed.Internal.Matrix
-- Copyright   :  (c) Alberto Ruiz 2007
-- License     :  GPL-style
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  provisional
-- Portability :  portable (uses FFI)
--
-- Internal matrix representation
--
-----------------------------------------------------------------------------
-- #hide

module Data.Packed.Internal.Matrix(
    Matrix(..), rows, cols, cdat, fdat,
    MatrixOrder(..), orderOf,
    createMatrix, mat,
    cmat, fmat,
    toLists, flatten, reshape,
    Element(..),
    trans,
    fromRows, toRows, fromColumns, toColumns,
    matrixFromVector,
    subMatrix,
    liftMatrix, liftMatrix2,
    (@@>), atM',
    saveMatrix,
    singleton,
    size, shSize, conformVs, conformMs, conformVTo, conformMTo
) where

import Data.Packed.Internal.Common
import Data.Packed.Internal.Signatures
import Data.Packed.Internal.Vector

import Foreign.Marshal.Alloc(alloca, free)
import Foreign.Marshal.Array(newArray)
import Foreign.Ptr(Ptr, castPtr)
import Foreign.Storable(Storable, peekElemOff, pokeElemOff, poke, sizeOf)
import Data.Complex(Complex)
import Foreign.C.Types
import Foreign.C.String(newCString)
import System.IO.Unsafe(unsafePerformIO)
import Control.DeepSeq

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

{- Design considerations for the Matrix Type
   -----------------------------------------

- we must easily handle both row major and column major order,
  for bindings to LAPACK and GSL/C

- we'd like to simplify redundant matrix transposes:
   - Some of them arise from the order requirements of some functions
   - some functions (matrix product) admit transposed arguments

- maybe we don't really need this kind of simplification:
   - more complex code
   - some computational overhead
   - only appreciable gain in code with a lot of redundant transpositions
     and cheap matrix computations

- we could carry both the matrix and its (lazily computed) transpose.
  This may save some transpositions, but it is necessary to keep track of the
  data which is actually computed to be used by functions like the matrix product
  which admit both orders.

- but if we need the transposed data and it is not in the structure, we must make
  sure that we touch the same foreignptr that is used in the computation.

- a reasonable solution is using two constructors for a matrix. Transposition just
  "flips" the constructor. Actual data transposition is not done if followed by a
  matrix product or another transpose.

-}

data MatrixOrder = RowMajor | ColumnMajor deriving (Show,Eq)

transOrder RowMajor = ColumnMajor
transOrder ColumnMajor = RowMajor
{- | Matrix representation suitable for GSL and LAPACK computations.

The elements are stored in a continuous memory array.

-}

data Matrix t = Matrix { irows :: {-# UNPACK #-} !Int
                       , icols :: {-# UNPACK #-} !Int
                       , xdat :: {-# UNPACK #-} !(Vector t)
                       , order :: !MatrixOrder }
-- RowMajor: preferred by C, fdat may require a transposition
-- ColumnMajor: preferred by LAPACK, cdat may require a transposition

cdat = xdat
fdat = xdat

rows :: Matrix t -> Int
rows = irows

cols :: Matrix t -> Int
cols = icols

orderOf :: Matrix t -> MatrixOrder
orderOf = order

-- | Matrix transpose.
trans :: Matrix t -> Matrix t
trans Matrix {irows = r, icols = c, xdat = d, order = o } = Matrix { irows = c, icols = r, xdat = d, order = transOrder o}

cmat :: (Element t) => Matrix t -> Matrix t
cmat m@Matrix{order = RowMajor} = m
cmat Matrix {irows = r, icols = c, xdat = d, order = ColumnMajor } = Matrix { irows = r, icols = c, xdat = transdata r d c, order = RowMajor}

fmat :: (Element t) => Matrix t -> Matrix t
fmat m@Matrix{order = ColumnMajor} = m
fmat Matrix {irows = r, icols = c, xdat = d, order = RowMajor } = Matrix { irows = r, icols = c, xdat = transdata c d r, order = ColumnMajor}

-- C-Haskell matrix adapter
-- mat :: Adapt (CInt -> CInt -> Ptr t -> r) (Matrix t) r

mat :: (Storable t) => Matrix t -> (((CInt -> CInt -> Ptr t -> t1) -> t1) -> IO b) -> IO b
mat a f =
    unsafeWith (xdat a) $ \p -> do
        let m g = do
            g (fi (rows a)) (fi (cols a)) p
        f m
-- | Creates a vector by concatenation of rows. If the matrix is ColumnMajor, this operation requires a transpose.
--
-- @\> flatten ('ident' 3)
-- 9 |> [1.0,0.0,0.0,0.0,1.0,0.0,0.0,0.0,1.0]@
flatten :: Element t => Matrix t -> Vector t
flatten = xdat . cmat

type Mt t s = Int -> Int -> Ptr t -> s
-- not yet admitted by my haddock version
-- infixr 6 ::>
-- type t ::> s = Mt t s

-- | the inverse of 'Data.Packed.Matrix.fromLists'
toLists :: (Element t) => Matrix t -> [[t]]
toLists m = splitEvery (cols m) . toList . flatten $ m

-- | Create a matrix from a list of vectors.
-- All vectors must have the same dimension,
-- or dimension 1, which is are automatically expanded.
fromRows :: Element t => [Vector t] -> Matrix t
fromRows vs = case compatdim (map dim vs) of
    Nothing -> error "fromRows applied to [] or to vectors with different sizes"
    Just c  -> reshape c . join . map (adapt c) $ vs
  where
    adapt c v | dim v == c = v
              | otherwise = constantD (v@>0) c

-- | extracts the rows of a matrix as a list of vectors
toRows :: Element t => Matrix t -> [Vector t]
toRows m = toRows' 0 where
    v = flatten m
    r = rows m
    c = cols m
    toRows' k | k == r*c  = []
              | otherwise = subVector k c v : toRows' (k+c)

-- | Creates a matrix from a list of vectors, as columns
fromColumns :: Element t => [Vector t] -> Matrix t
fromColumns m = trans . fromRows $ m

-- | Creates a list of vectors from the columns of a matrix
toColumns :: Element t => Matrix t -> [Vector t]
toColumns m = toRows . trans $ m

-- | Reads a matrix position.
(@@>) :: Storable t => Matrix t -> (Int,Int) -> t
infixl 9 @@>
m@Matrix {irows = r, icols = c} @@> (i,j)
    | safe      = if i<0 || i>=r || j<0 || j>=c
                    then error "matrix indexing out of range"
                    else atM' m i j
    | otherwise = atM' m i j
{-# INLINE (@@>) #-}

--  Unsafe matrix access without range checking
atM' Matrix {icols = c, xdat = v, order = RowMajor} i j = v `at'` (i*c+j)
atM' Matrix {irows = r, xdat = v, order = ColumnMajor} i j = v `at'` (j*r+i)
{-# INLINE atM' #-}

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

matrixFromVector o c v = Matrix { irows = r, icols = c, xdat = v, order = o }
    where (d,m) = dim v `quotRem` c
          r | m==0 = d
            | otherwise = error "matrixFromVector"

-- allocates memory for a new matrix
createMatrix :: (Storable a) => MatrixOrder -> Int -> Int -> IO (Matrix a)
createMatrix ord r c = do
    p <- createVector (r*c)
    return (matrixFromVector ord c p)

{- | Creates a matrix from a vector by grouping the elements in rows with the desired number of columns. (GNU-Octave groups by columns. To do it you can define @reshapeF r = trans . reshape r@
where r is the desired number of rows.)

@\> reshape 4 ('fromList' [1..12])
(3><4)
 [ 1.0,  2.0,  3.0,  4.0
 , 5.0,  6.0,  7.0,  8.0
 , 9.0, 10.0, 11.0, 12.0 ]@

-}
reshape :: Storable t => Int -> Vector t -> Matrix t
reshape c v = matrixFromVector RowMajor c v

singleton x = reshape 1 (fromList [x])

-- | application of a vector function on the flattened matrix elements
liftMatrix :: (Storable a, Storable b) => (Vector a -> Vector b) -> Matrix a -> Matrix b
liftMatrix f Matrix { icols = c, xdat = d, order = o } = matrixFromVector o c (f d)

-- | application of a vector function on the flattened matrices elements
liftMatrix2 :: (Element t, Element a, Element b) => (Vector a -> Vector b -> Vector t) -> Matrix a -> Matrix b -> Matrix t
liftMatrix2 f m1 m2
    | not (compat m1 m2) = error "nonconformant matrices in liftMatrix2"
    | otherwise = case orderOf m1 of
        RowMajor    -> matrixFromVector RowMajor    (cols m1) (f (xdat m1) (flatten m2))
        ColumnMajor -> matrixFromVector ColumnMajor (cols m1) (f (xdat m1) ((xdat.fmat) m2))

compat :: Matrix a -> Matrix b -> Bool
compat m1 m2 = rows m1 == rows m2 && cols m1 == cols m2

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

{- | Supported matrix elements.

    This class provides optimized internal
    operations for selected element types.
    It provides unoptimised defaults for any 'Storable' type,
    so you can create instances simply as:
    @instance Element Foo@.
-}
class (Storable a) => Element a where
    subMatrixD :: (Int,Int) -- ^ (r0,c0) starting position
               -> (Int,Int) -- ^ (rt,ct) dimensions of submatrix
               -> Matrix a -> Matrix a
    subMatrixD = subMatrix'
    transdata :: Int -> Vector a -> Int -> Vector a
    transdata = transdataP -- transdata'
    constantD  :: a -> Int -> Vector a
    constantD = constantP -- constant'

instance Element Float where
    transdata  = transdataAux ctransF
    constantD  = constantAux cconstantF

instance Element Double where
    transdata  = transdataAux ctransR
    constantD  = constantAux cconstantR

instance Element (Complex Float) where
    transdata  = transdataAux ctransQ
    constantD  = constantAux cconstantQ

instance Element (Complex Double) where
    transdata  = transdataAux ctransC
    constantD  = constantAux cconstantC

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

transdata' :: Storable a => Int -> Vector a -> Int -> Vector a
transdata' c1 v c2 =
    if noneed
        then v
        else unsafePerformIO $ do
                w <- createVector (r2*c2)
                unsafeWith v $ \p ->
                    unsafeWith w $ \q -> do
                        let go (-1) _ = return ()
                            go !i (-1) = go (i-1) (c1-1)
                            go !i !j = do x <- peekElemOff p (i*c1+j)
                                          pokeElemOff      q (j*c2+i) x
                                          go i (j-1)
                        go (r1-1) (c1-1)
                return w
  where r1 = dim v `div` c1
        r2 = dim v `div` c2
        noneed = r1 == 1 || c1 == 1

-- {-# SPECIALIZE transdata' :: Int -> Vector Double -> Int ->  Vector Double #-}
-- {-# SPECIALIZE transdata' :: Int -> Vector (Complex Double) -> Int -> Vector (Complex Double) #-}

-- I don't know how to specialize...
-- The above pragmas only seem to work on top level defs
-- Fortunately everything seems to work using the above class

-- C versions, still a little faster:

transdataAux fun c1 d c2 =
    if noneed
        then d
        else unsafePerformIO $ do
            v <- createVector (dim d)
            unsafeWith d $ \pd ->
                unsafeWith v $ \pv ->
                    fun (fi r1) (fi c1) pd (fi r2) (fi c2) pv // check "transdataAux"
            return v
  where r1 = dim d `div` c1
        r2 = dim d `div` c2
        noneed = r1 == 1 || c1 == 1

transdataP :: Storable a => Int -> Vector a -> Int -> Vector a
transdataP c1 d c2 =
    if noneed
       then d
       else unsafePerformIO $ do
          v <- createVector (dim d)
          unsafeWith d $ \pd ->
              unsafeWith v $ \pv ->
                  ctransP (fi r1) (fi c1) (castPtr pd) (fi sz) (fi r2) (fi c2) (castPtr pv) (fi sz) // check "transdataP"
          return v
   where r1 = dim d `div` c1
         r2 = dim d `div` c2
         sz = sizeOf (d @> 0)
         noneed = r1 == 1 || c1 == 1

foreign import ccall unsafe "transF" ctransF :: TFMFM
foreign import ccall unsafe "transR" ctransR :: TMM
foreign import ccall unsafe "transQ" ctransQ :: TQMQM
foreign import ccall unsafe "transC" ctransC :: TCMCM
foreign import ccall unsafe "transP" ctransP :: CInt -> CInt -> Ptr () -> CInt -> CInt -> CInt -> Ptr () -> CInt -> IO CInt

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

constant' v n = unsafePerformIO $ do
    w <- createVector n
    unsafeWith w $ \p -> do
        let go (-1) = return ()
            go !k = pokeElemOff p k v >> go (k-1)
        go (n-1)
    return w

-- C versions

constantAux fun x n = unsafePerformIO $ do
    v <- createVector n
    px <- newArray [x]
    app1 (fun px) vec v "constantAux"
    free px
    return v

constantF :: Float -> Int -> Vector Float
constantF = constantAux cconstantF
foreign import ccall unsafe "constantF" cconstantF :: Ptr Float -> TF

constantR :: Double -> Int -> Vector Double
constantR = constantAux cconstantR
foreign import ccall unsafe "constantR" cconstantR :: Ptr Double -> TV

constantQ :: Complex Float -> Int -> Vector (Complex Float)
constantQ = constantAux cconstantQ
foreign import ccall unsafe "constantQ" cconstantQ :: Ptr (Complex Float) -> TQV

constantC :: Complex Double -> Int -> Vector (Complex Double)
constantC = constantAux cconstantC
foreign import ccall unsafe "constantC" cconstantC :: Ptr (Complex Double) -> TCV

constantP :: Storable a => a -> Int -> Vector a
constantP a n = unsafePerformIO $ do
    let sz = sizeOf a
    v <- createVector n
    unsafeWith v $ \p -> do
       alloca $ \k -> do
                      poke k a
                      cconstantP (castPtr k) (fi n) (castPtr p) (fi sz) // check "constantP"
    return v
foreign import ccall unsafe "constantP" cconstantP :: Ptr () -> CInt -> Ptr () -> CInt -> IO CInt

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

-- | Extracts a submatrix from a matrix.
subMatrix :: Element a
          => (Int,Int) -- ^ (r0,c0) starting position
          -> (Int,Int) -- ^ (rt,ct) dimensions of submatrix
          -> Matrix a -- ^ input matrix
          -> Matrix a -- ^ result
subMatrix (r0,c0) (rt,ct) m
    | 0 <= r0 && 0 < rt && r0+rt <= (rows m) &&
      0 <= c0 && 0 < ct && c0+ct <= (cols m) = subMatrixD (r0,c0) (rt,ct) m
    | otherwise = error $ "wrong subMatrix "++
                          show ((r0,c0),(rt,ct))++" of "++show(rows m)++"x"++ show (cols m)

subMatrix'' (r0,c0) (rt,ct) c v = unsafePerformIO $ do
    w <- createVector (rt*ct)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) _ = return ()
                go !i (-1) = go (i-1) (ct-1)
                go !i !j = do x <- peekElemOff p ((i+r0)*c+j+c0)
                              pokeElemOff      q (i*ct+j) x
                              go i (j-1)
            go (rt-1) (ct-1)
    return w

subMatrix' (r0,c0) (rt,ct) (Matrix { icols = c, xdat = v, order = RowMajor}) = Matrix rt ct (subMatrix'' (r0,c0) (rt,ct) c v) RowMajor
subMatrix' (r0,c0) (rt,ct) m = trans $ subMatrix' (c0,r0) (ct,rt) (trans m)

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

-- | Saves a matrix as 2D ASCII table.
saveMatrix :: FilePath
           -> String     -- ^ format (%f, %g, %e)
           -> Matrix Double
           -> IO ()
saveMatrix filename fmt m = do
    charname <- newCString filename
    charfmt <- newCString fmt
    let o = if orderOf m == RowMajor then 1 else 0
    app1 (matrix_fprintf charname charfmt o) mat m "matrix_fprintf"
    free charname
    free charfmt

foreign import ccall unsafe "matrix_fprintf" matrix_fprintf :: Ptr CChar -> Ptr CChar -> CInt -> TM

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

conformMs ms = map (conformMTo (r,c)) ms
  where
    r = maximum (map rows ms)
    c = maximum (map cols ms)

conformVs vs = map (conformVTo n) vs
  where
    n = maximum (map dim vs)

conformMTo (r,c) m
    | size m == (r,c) = m
    | size m == (1,1) = reshape c (constantD (m@@>(0,0)) (r*c))
    | size m == (r,1) = repCols c m
    | size m == (1,c) = repRows r m
    | otherwise = error $ "matrix " ++ shSize m ++ " cannot be expanded to (" ++ show r ++ "><"++ show c ++")"

conformVTo n v
    | dim v == n = v
    | dim v == 1 = constantD (v@>0) n
    | otherwise = error $ "vector of dim=" ++ show (dim v) ++ " cannot be expanded to dim=" ++ show n

repRows n x = fromRows (replicate n (flatten x))
repCols n x = fromColumns (replicate n (flatten x))

size m = (rows m, cols m)

shSize m = "(" ++ show (rows m) ++"><"++ show (cols m)++")"

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

instance (Storable t, NFData t) => NFData (Matrix t)
  where
    rnf m | d > 0     = rnf (v @> 0)
          | otherwise = ()
      where
        d = dim v
        v = xdat m

max_stars_repo_name: curiousleo/liquidhaskell

max_issues_repo_name: curiousleo/liquidhaskell

max_forks_repo_name: curiousleo/liquidhaskell


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

ext: hs

max_stars_repo_name: danielkroeni/arrayfire-haskell

max_issues_repo_name: danielkroeni/arrayfire-haskell

max_forks_repo_name: danielkroeni/arrayfire-haskell


text:
{-# LANGUAGE GADTs,  TypeFamilies,
 StandaloneDeriving, UndecidableInstances,
ScopedTypeVariables, FlexibleInstances, DataKinds,
FunctionalDependencies, PolyKinds,
TypeOperators, RankNTypes,  MultiParamTypeClasses,
 TypeApplications, FlexibleContexts, AllowAmbiguousTypes #-}
-- AllowAmbiguousTypes, ImpredicativeTypes, InstanceSigs, NoImplicitPrelude,

module Fib where

import Vec
import Data.Complex
import Control.Monad  ((<=<))
import GHC.TypeNats
import Data.Proxy

-- data FibAnyon = Id | Tau
-- I started using DataKinds and it ended up being a problem.

data Tau
--data Id
type Id = ()
data FibTree root leaves where
   TTT :: FibTree Tau l -> FibTree Tau r -> FibTree Tau (l,r)
   ITT :: FibTree Tau l -> FibTree Tau r -> FibTree Id (l,r)
   TIT :: FibTree Id l -> FibTree Tau r -> FibTree Tau (l,r)
   TTI :: FibTree Tau l -> FibTree Id r -> FibTree Tau (l,r)
   III :: FibTree Id l -> FibTree Id r -> FibTree Id (l,r)
   TLeaf :: FibTree Tau Tau
   ILeaf :: FibTree Id Id

-- pretty printing would be hella nice
deriving instance Show (FibTree a b)
deriving instance Eq (FibTree a b)

instance Ord (FibTree a b) where
    compare (ITT l r) (ITT l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'
    compare (ITT _ _) _ = LT
    compare _ (ITT _ _) = GT
    compare (TTI l r) (TTI l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'
    compare (TIT l r) (TIT l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'
    compare (TTT l r) (TTT l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'
    compare (III l r) (III l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'

    compare (TTI _ _) _ = LT
    compare _ (TTI _ _) = GT
    compare (TIT _ _) _ = LT
    compare _ (TIT _ _) = GT
    compare (TTT _ _) _ = LT
    compare _ (TTT _ _) = GT
    compare (III _ _) _ = LT
    compare _ (III _ _) = GT
    compare TLeaf TLeaf = EQ
    compare ILeaf ILeaf = EQ

lmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (b,d) -> Q (FibTree e (c,d)))
lmap f (ITT l r) = fmap (\l' -> ITT l' r) (f l)
lmap f (TTI l r) = fmap (\l' -> TTI l' r) (f l)
lmap f (TIT l r) = fmap (\l' -> TIT l' r) (f l)
lmap f (TTT l r) = fmap (\l' -> TTT l' r) (f l)
lmap f (III l r) = fmap (\l' -> III l' r) (f l)

rmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (d,b) -> Q (FibTree e (d,c)))
rmap f (ITT l r) = fmap (\r' -> ITT l r') (f r)
rmap f (TTI l r) = fmap (\r' -> TTI l r') (f r)
rmap f (TIT l r) = fmap (\r' -> TIT l r') (f r)
rmap f (TTT l r) = fmap (\r' -> TTT l r') (f r)
rmap f (III l r) = fmap (\r' -> III l r') (f r)

braid :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid (ITT l r) = W [(ITT r l,  cis $ 4 * pi / 5)]  -- different scalar factors for trivial and non trivial fusion
braid (TTT l r) = W [(TTT r l,  (cis $ - 3 * pi / 5))]
braid (TTI l r) = pure $ TIT r l-- exchange with trivial means nothing
braid (TIT l r) = pure $ TTI r l
braid (III l r) = pure $ III r l

-- The inverse of braid
braid' :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid' = star . braid

tau :: Complex Double
tau =  ((sqrt 5) - 1) / 2 :+ 0 -- 0.618 :+ 0

fmove :: FibTree a (c,(d,e)) -> Q (FibTree a ((c,d),e))
-- fmove (ITT  a  (TTI b c)) = pure $ ITI ( TTT  a b) c -- pure (auto (auto a b) c) -- no maybe not. The internal one isn't auto
fmove (ITT  a  (TIT b c)) = pure $ ITT ( TTI  a b) c
fmove (ITT  a  (TTT b c)) = pure $ ITT ( TTT  a b) c
fmove (ITT  a  (TTI b c)) = pure $ III ( ITT  a b) c

fmove (TIT  a  (TTT b c)) = pure $ TTT ( TIT  a b) c
fmove (TIT  a  (TTI b c)) = pure $ TTI ( TIT  a b) c
fmove (TIT  a  (TIT b c)) = pure $ TIT ( III  a b) c

-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
-- the nontrivial ones have all tau on the leafs and root
-- internal I
fmove (TTI  a  (III b c)) = pure $ TTI ( TTI  a b) c
fmove (TTI  a  (ITT b c)) = W [(TIT ( ITT  a b) c, tau)         , (TTT ( TTT  a b) c, sqrt tau)]
-- internal T
fmove (TTT  a  (TTT b c)) = W [(TIT ( ITT  a b) c, sqrt tau)  ,   (TTT ( TTT  a b) c, - tau   )]
fmove (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c
fmove (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c

fmove (III  a  (ITT b c)) = pure $ ITT ( TIT  a b) c
fmove (III  a  (III b c)) = pure $ III ( III  a b) c

-- largely just a tranpose of the above case.
fmove' :: FibTree a ((c,d),e) -> Q (FibTree a (c,(d,e)))
fmove' (ITT ( TTI  a b) c) = pure $ (ITT  a  (TIT b c))
fmove' (ITT ( TTT  a b) c) = pure $  (ITT  a  (TTT b c))
fmove' (ITT ( TIT  a b) c) = pure $  (III  a  (ITT b c))

--fmoveq (ITT  a  (TTT b c)) = pure $

fmove' (TTI ( TTT  a b) c) = pure $ (TTT  a  (TTI b c))
fmove' (TTI ( TTI  a b) c) = pure $ (TTI  a  (III b c))
fmove' (TTI ( TIT  a b) c) = pure $ TIT  a  (TTI b c)
--fmoveq (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c

--fmoveq (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c
fmove' (TIT ( ITT  a b) c) = W [(TTI  a  (ITT b c), tau)         , (TTT  a  (TTT b c) , sqrt tau)]
fmove' (TIT ( III  a b) c ) = pure $ TIT  a  (TIT b c)

fmove' (TTT ( TTI  a b) c ) = pure $ TTT  a  (TIT b c)
fmove' (TTT ( TIT  a b) c ) = pure $ TIT  a  (TTT b c)
fmove' (TTT ( TTT  a b) c) = W [(TTI  a  (ITT b c), sqrt tau)  , (TTT  a  (TTT b c),   - tau  )]

fmove' (III ( III  a b) c ) = pure $ III  a  (III b c)
fmove' (III ( ITT  a b) c ) = pure $ ITT  a  (TTI b c)

rightUnit :: FibTree e (a,Id) -> Q (FibTree e a)
rightUnit (TTI t _) = pure t
rightUnit (III t _) = pure t

rightUnit' :: FibTree e a -> Q (FibTree e (a,Id))
rightUnit' t@(TTT _ _) = pure (TTI  t ILeaf)
rightUnit' t@(TTI _ _) = pure (TTI  t ILeaf)
rightUnit' t@(TIT _ _) = pure (TTI  t ILeaf)
rightUnit' t@(III _ _) = pure (III  t ILeaf)
rightUnit' t@(ITT _ _) = pure (III  t ILeaf)
rightUnit' t@(ILeaf) = pure (III t ILeaf)
rightUnit' t@(TLeaf) = pure (TTI t ILeaf)

leftUnit :: FibTree e (Id,a) -> Q (FibTree e a)
leftUnit = rightUnit <=< braid

-- braid vs braid' doesn't matter, but it has a nice symmettry.
leftUnit' :: FibTree e a -> Q (FibTree e (Id,a))
leftUnit' = braid' <=< rightUnit'

dot :: FibTree a (b, c) -> FibTree a' (b, c) -> Q (FibTree a' a)
dot x@(TTI _ _) y@(TTI _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TIT _ _) y@(TIT _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TTT _ _) y@(TTT _ _) | x == y = pure TLeaf
                            | otherwise = mempty
dot x@(III _ _) y@(III _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot x@(ITT _ _) y@(ITT _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot _ _ = mempty

class AllTrees a b where
    allTrees :: [FibTree a b]
instance AllTrees Tau Tau where
    allTrees = [TLeaf]
instance AllTrees Id Tau where
        allTrees = []
instance AllTrees Tau Id where
        allTrees = []
instance AllTrees Id Id where
        allTrees = [ILeaf]
instance (AllTrees Tau a,
         AllTrees Id a,
         AllTrees Tau b,
         AllTrees Id b) => AllTrees Id (a,b) where
        allTrees = (III <$> ia <*> ib) <> (ITT <$> ta <*> tb) where
            ta = allTrees @Tau @a
            ia = allTrees @Id @a
            tb = allTrees @Tau @b
            ib = allTrees @Id @b
instance (AllTrees Tau a,
        AllTrees Id a,
        AllTrees Tau b,
        AllTrees Id b) => AllTrees Tau (a,b) where
       allTrees = (TIT <$> ia <*> tb) <> (TTI <$> ta <*> ib) <> (TTT <$> ta <*> tb) where
            ta = allTrees @Tau @a
            ia = allTrees @Id @a
            tb = allTrees @Tau @b
            ib = allTrees @Id @b
t9 = allTrees @Tau @((Tau,Tau),Tau)

-- Resulting type depends on input
-- I think a typefamily type computation might be necessary?
-- pullLeft and pullRight might not need a type class.
-- pullLeft (TTT l r) = fmove (TTT pl r)
--                         where pl = pullLeft l
{-
type family Assoc a where
   Assoc ((a,b),c) = (a,(b,c))

type family PullLeft a where
   PullLeft ((a,b),c) = Assoc (PullLeft (a, b), c)
   PullLeft a = a

type Test1 = PullLeft (((Int,Int),Int),Int)
pullLeft :: forall a b c. (c ~ PullLeft b) => FibTree a b -> Q (FibTree a c)
pullLeft TLeaf = pure Tleaf
pullLeft ILeaf = pure ILeaf
pullLeft t =  do
            t' <- lmap pullLeft t
            fmove' t'

pullLeft t@(ITT l r) = helper t
pullLeft t@(TTT l r) = helper t
pullLeft t@(TTI l r) = helper t
pullLeft t@(TIT l r) = helper t
pullLeft t@(III l r) = helper t where
                       helper :: (d ~ PullLeft (b,c)) => FibTree a (b,c) -> Q (FibTree a d)
                       helper t = do
                                t' <- lmap pullLeft t
                                fmove' t'

-}

{-
Spiritually this function is this. Can't write it this way unforturnately
-- There are typing problems
-- this isn't even right
-- we need to pattern match on ((),)
pullLeft TLeaf = pure Tleaf
pullLeft ILeaf = pure ILeaf
pullLeft t =  do
                t' <- lmap pullLeft t
                fmove' t'

-- but actually a two level tree desctruct
--treeDestruct()

-}
{-
type family Assoc a where
    Assoc ((a,b),c) = (a,(b,c))

 type family PullLeft' a where
    PullLeft' ((a,b),c) = Assoc (PullLeft' (a, b), c)
    PullLeft' a = a
    -}
-- could seperate the typeclass so that I'm using type families instead.

class PullLeftLeaf a b | a -> b where
  pullLeftLeaf :: FibTree c a -> Q (FibTree c b)
instance PullLeftLeaf (Tau,c) (Tau,c) where
  pullLeftLeaf = pure
instance PullLeftLeaf (Id,c) (Id,c) where
  pullLeftLeaf = pure
instance PullLeftLeaf Tau Tau where
  pullLeftLeaf = pure
instance PullLeftLeaf Id Id where
  pullLeftLeaf = pure
instance (PullLeftLeaf (a,b) (a',b'),
          r ~ (a',(b',c))) => PullLeftLeaf ((a, b),c) r where
  pullLeftLeaf t = do
           t' <- lmap pullLeftLeaf t
           fmove' t'

class PullRightLeaf a b | a -> b where -- | a -> b functional dependency causes errors?
    pullRightLeaf :: FibTree c a -> Q (FibTree c b)
{-
instance (PullRightLeaf a a', r ~ ((b,c),a')) => PullRightLeaf ((b,c),a) r where
pullRightLeaf t = rmap pullRightLeaf t
-}
instance PullRightLeaf Tau Tau where
    pullRightLeaf = pure

instance PullRightLeaf Id Id where
    pullRightLeaf = pure

instance PullRightLeaf (c,Tau) (c,Tau) where
    pullRightLeaf = pure

instance PullRightLeaf (c,Id) (c,Id) where
    pullRightLeaf = pure

instance (PullRightLeaf (a,b) (a',b'), r ~ ((c,a'),b')) => PullRightLeaf (c,(a, b)) r where
    pullRightLeaf t = do
        t' <- rmap pullRightLeaf t
        fmove t'

class RightAssoc a b | a -> b where
  rightAssoc :: FibTree c a -> Q (FibTree c b)
instance RightAssoc Tau Tau where
  rightAssoc = pure
instance RightAssoc Id Id where
  rightAssoc = pure
instance (PullLeftLeaf (a,b) (a',b'),
          RightAssoc b' b'',
          r ~ (a', b'')) => RightAssoc (a,b) r where
  rightAssoc t = do
           t' <- pullLeftLeaf t
           rmap rightAssoc t'
-- usually you'll want to force not the root, but the leaves to be some type
-- hence the ordering b c d a for type application
bmove :: forall b c d a. FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))
bmove t = do
           t'  :: FibTree a ((b,c),d) <- fmove t
           t'' :: FibTree a ((c,b),d) <-  lmap braid t'
           fmove' t''
bmove' :: forall b c d a. FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))
bmove' = fmove' <=< (lmap braid') <=< fmove

-- Therei s a general pattern for digging into
-- n and a tell us b, the subpiece of a
-- c and a and n tell use what a looks like with b replaced = d
--data Nat = Z | S Nat

rmapN :: forall n gte s t a b e. (RMapN n gte s t a b, gte ~ (CmpNat n 0)) => (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
rmapN f t = rmapN' @n @gte f t

class RMapN n gte s t a b | n gte s b -> a t where
    rmapN' :: (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
instance (a ~ s, b ~ t) => RMapN 0 'EQ s t a b where
    rmapN' f t = f t -- rmapN = id
instance (RMapN (n-1) gte r r' a b,
               gte ~ (CmpNat (n-1) 0),
               t ~ (l,r')) => RMapN n 'GT (l,r) t a b where
    rmapN' f t = rmap (rmapN @(n-1) f) t
{-
bmoveN :: forall n a b c d e f. RMapN n a (b, (c, d)) (c, (b, d)) e => FibTree f a -> Q (FibTree f e)
bmoveN t = rmapN @n (bmove @b @c @d) t

-- dotN :: forall n a b c d e f. dot :: FibTree a (b, c) -> FibTree a' (b, c) -> Q (FibTree a' a)
fuseN :: forall n a b c d e f g. RMapN n f (b,(c,d)) (a,d) g => FibTree a (b, c) -> FibTree e f -> Q (FibTree e g)
fuseN q t = rmapN @n f t where
    f :: forall r. FibTree r (b,(c,d)) -> Q (FibTree r (a,d))
    f t' = do
            t'' <- fmove t'
            lmap (dot q) t''
-}

type family Count a where
    Count Tau = 1
    Count Id = 1
    Count (a,b) = (Count a) + (Count b)

type family LeftCount a where
    LeftCount (a,b) = Count a

lcamap ::  forall n s t a b e gte .
           (gte ~ CmpNat (LeftCount s) n,
           LCAMap n gte s t a b)
           => (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
lcamap f t = lcamap' @n @gte f t

class LCAMap n gte s t a b | n gte s b -> t a where
    lcamap' :: (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)

instance (n' ~ (n - Count l), -- We're searching in the right subtree. Subtract the leaf number in the left subtree
            lc ~ (LeftCount r), -- dip one level down to order which way we have to go next
            gte ~ (CmpNat lc n'), -- Do we go left, right or have we arrived in the next layer?
            LCAMap n' gte r r' a b,  -- recursive call
            t ~ (l,r') -- reconstruct total return type from recursive return type. Left tree is unaffected by lcamapping
            ) => LCAMap n 'LT (l,r) t a b where
        lcamap' f x = rmap (lcamap @n' f) x

instance (lc ~ (LeftCount l),
            gte ~ (CmpNat lc n),
            LCAMap n gte l l' a b,
            t ~ (l',r)
            ) => LCAMap n 'GT (l,r) t a b where
        lcamap' f x = lmap (lcamap @n f) x

instance (t ~ b, a ~ s) => LCAMap n 'EQ s t a b where -- Base case
    lcamap' f x = f x

class Twiddle s t a b | s b -> t a where
    twiddle :: (forall r. FibTree r a -> Q (FibTree r b)) -> FibTree e s -> Q (FibTree e t)

instance Twiddle ((l,x),(y,r)) ((l,c),r) (x,y) c where
    twiddle f x = do
            x'  <- fmove x -- (((l',x),y),r')
            x'' <- lmap fmove' x' -- ((l',(x,y)),r')
            lmap (rmap f) x''
instance Twiddle (Tau, (y,r)) (c,r) (Tau, y) c where
    twiddle f x = fmove x >>= lmap f
instance Twiddle (Id, (y,r)) (c,r)  (Id, y) c where
    twiddle f x = fmove x >>= lmap f
instance Twiddle ((l,x), Tau) (l,c) (x,Tau) c where
    twiddle f x = fmove' x >>= rmap f
instance Twiddle ((l,x), Id) (l,c) (x,Id) c where
    twiddle f x = fmove' x >>= rmap f
instance Twiddle (Tau, Tau) c (Tau,Tau) c where
    twiddle f x = f x
instance Twiddle (Id, Id) c (Id,Id)  c where
    twiddle f x = f x
instance Twiddle (Tau, Id) c (Tau,Id)  c where
    twiddle f x = f x
instance Twiddle (Id, Tau) c (Id,Tau) c where
    twiddle f x = f x

nmap :: forall (n :: Nat) s t a b a' b' l l' r r' e gte.
    (gte ~ CmpNat (LeftCount s) n,
    LCAMap n gte s t a' b',
    a' ~ (l,r),
    PullRightLeaf l l',
    PullLeftLeaf r r',
    Twiddle (l',r') b' a b) =>
    (forall r. FibTree r a -> Q (FibTree r b)) -> FibTree e s -> Q (FibTree e t)
nmap f z = lcamap @n @s @t @a' @b' (\x -> do
                                           x'  <- lmap pullRightLeaf x
                                           x'' <- rmap pullLeftLeaf x'
                                           twiddle f x'') z

t1 = nmap @2 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t5 = nmap @2 pure (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf)) >>= nmap @3 pure
t2 = nmap @1 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t4 = nmap @1 braid (TTT TLeaf (TTT TLeaf TLeaf))
t3 = nmap @2 braid (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
t6 = rightAssoc (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
t7 = t6 >>= bmove
t8 = t6 >>= rmapN @0 bmove
-- For the category, we probably don't want the association structure.
-- '[Tau, Id, Tau, Id] typelevel list of particles
--

ttt = TTT TLeaf TLeaf
example = starttree >>=
        nmap @1 braid >>=
        nmap @2 braid >>=
        nmap @1 (dot ttt) >>=
        nmap @1 braid' >>=
        nmap @2 (dot ttt) >>=
        nmap @1 (dot ttt) where
        starttree = pure (TTT (TTT TLeaf
                              (TTT TLeaf
                                   TLeaf))
                          TLeaf
                         )
-- would be nice to use an unknot
-- example2 =
-- I should use neighbormap for this.
-- maybe I should call f fuse and make it global?

-- bmoveN t = rmapN @n (bmove :: forall r. FibTree r (b,(c,d)) -> Q (FibTree r (c,(b,d)))) t
-- bmoveN t = rmapN @n @a @(b, (c, d)) @(c, (b, d)) @e bmove t

{-
bmoveN :: forall n a b c d e f. RMapN n a (b, (c, d)) (c, (b, d)) e => Proxy n -> FibTree f a -> Q (FibTree f e)
bmoveN p t = rmapN p (bmove :: FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))) t
-}
   {- do
            t'  <- fmove t
            t'' <-  lmap braid' t'
            fmove' t''
-}

           --
{-
pullLeftLeaf
pullRightLeaf

-}
{-
class Standardize a b | a -> b where
    standardize :: FibTree c a -> Q (FibTree c b)

instance Standardize a b where

type family Append a b where
    Append (a,b) c = (a, Append b c)
    Append a c = (a, c)

type family Leftize where
    Leftize (a,b) = Append (Leftize a) (Leftize b)
    Leftize a = a

fullrightassoc = standardize
completerightassoc
bmove = fmove braid fmove'

standardize ::  FibTree c a -> Q (FibTree c (Leftize a))
standardize (TLeaf) = TLeaf
standardize (ILeaf) = ILeaf
standardize t = do
                x <- pullLeft t
                rmap standardize x

lcamap :: FibTree a b -> Proxy (a :: Int) -> FibTree a ?
lcamap f n t@(TTT l r) | count l == n  = f t
                       | count l < n   = lcamap f (n - count l) r
                       | otherwise     = lcamap f n l

pullRight ()

treerecursor :: FibTree a b -> (l -> r -> c) -> (leaf -> c) -> c
treerecursor (TTT l r) f g = f l r
treerecursor (TLeaf) = g Tleaf

-}


text:
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ScopedTypeVariables #-}

module ModelSpec where

import Control.Applicative
import Codec.Picture
import Codec.Picture.Types
import qualified Data.Aeson as A
import Data.Complex
import qualified Data.HashMap.Strict as HM
import Data.Proxy
import qualified Data.Text as T
import qualified Data.Text.Lazy as TL
import qualified Data.Vector as V
import GHC.Generics
import GHC.Word
import Generics.SOP.Arbitrary
import Generics.SOP
import Test.QuickCheck
import Test.QuickCheck.Function
import Test.QuickCheck.Arbitrary
import Test.Hspec
import Test.Hspec.QuickCheck
import Model

spec = do
  describe "Encoding and decoding"
    $ modifyMaxSize (const 10)
    $ modifyMaxSuccess (const 50) $ do
      it "val should round-trip aeson" $  property $ \(v :: Val) ->
        A.decode (A.encode v) == Just v

instance Arbitrary Val where
  arbitrary = sized arbVal

arbVal :: Int -> Gen Val
arbVal 0 = oneof [ VDouble  <$> arbitrary
                 , VComplex <$> arbitrary
                 , VText    <$> arbitrary
                 , VImage   <$> arbitrary
                 ]
arbVal n = do
  (Positive m) <- arbitrary
  let n' = n `div` (m + 1)
  oneof [ (VList . V.toList) <$> V.replicateM n (arbVal n')
        , VVec1 <$> V.replicateM n' (arbVal m)
        , fmap VVec2 $ do
           x <- (V.replicateM (div n 4) (arbVal (n `div` 4)))
           y <- (V.replicateM (div n 4) (arbVal (n `div` 4)))
           return (V.zip x y)
        , fmap (VVec3 . V.fromList) $ do
           x <- (vectorOf (div n 2) (arbVal (div n 2)))
           y <- (vectorOf (div n 2) (arbVal (div n 2)))
           z <- (vectorOf (div n 2) (arbVal (div n 2)))
           return (zip3 x y z)
        ]

instance Arbitrary PrimComplex where
  arbitrary = fmap PComplex $ liftA2 (:+) arbitrary arbitrary

instance Arbitrary A.Value where
  arbitrary = oneof
    [ fmap (A.String . T.pack) arbitrary
    , fmap (A.Number . realToFrac) (arbitrary :: Gen Double)
    , return A.Null
    , fmap (A.Object . HM.fromList) (zip <$> arbitrary <*> arbitrary)
    , fmap (A.Array . V.fromList) arbitrary
    ]

instance Arbitrary Expr where
  arbitrary = fmap Expr arbitrary

instance Arbitrary T.Text where
  arbitrary = T.pack <$> arbitrary

genImg :: (Pixel a, Arbitrary a)
       => Proxy a
       -> Int
       -> Int
       -> Gen (Image a)
genImg _ wid hgt = do
  f :: Fun (Int,Int) a <- arbitrary
  let img = generateImage (curry (apply f)) wid hgt
  return img

r16 :: (Word16,Word16)
r16 = (minBound,maxBound)

r8 :: (Word8,Word8)
r8 = (minBound,maxBound)

rF :: (Float, Float)
rF = (0,1)

instance Arbitrary a => Arbitrary (V.Vector a) where
  arbitrary = fmap V.fromList arbitrary

instance Arbitrary PixelYA8 where
  arbitrary = liftA2 PixelYA8 (choose r8) (choose r8)

instance Arbitrary PixelYA16 where
  arbitrary = liftA2 PixelYA16 (choose r16) (choose r16)

instance Arbitrary PixelRGB8 where
  arbitrary = liftA3 PixelRGB8 (choose r8) (choose r8) (choose r8)

instance Arbitrary PixelRGB16 where
  arbitrary = liftA3 PixelRGB16 (choose r16) (choose r16) (choose r16)

instance Arbitrary PixelRGBF where
  arbitrary = liftA3 PixelRGBF (choose rF) (choose rF) (choose rF)

liftA4 :: Applicative m
       => (a -> b -> c -> d -> e) -> m a -> m b -> m c -> m d -> m e
liftA4 f a b c d = f <$> a <*> b <*> c <*> d

instance Arbitrary PixelRGBA8 where
  arbitrary = liftA4 PixelRGBA8 (choose r8) (choose r8) (choose r8) (choose r8)

instance Arbitrary PixelRGBA16 where
  arbitrary = liftA4 PixelRGBA16 (choose r16) (choose r16) (choose r16) (choose r16)

instance Arbitrary PixelCMYK8 where
  arbitrary = liftA4 PixelCMYK8 (choose r8) (choose r8) (choose r8) (choose r8)

instance Arbitrary PixelCMYK16 where
  arbitrary = liftA4 PixelCMYK16 (choose r16) (choose r16) (choose r16) (choose r16)

instance Arbitrary PixelYCbCr8 where
  arbitrary = liftA3 PixelYCbCr8 (choose r8) (choose r8) (choose r8)

instance Arbitrary ModelImage where
  arbitrary = genModelImage

genModelImage :: Gen ModelImage
genModelImage = do
    x <- choose (10 :: Int, 200)
    y <- choose (10 :: Int, 200)
    dynImg <- oneof
      [ -- ImageY8     <$> genImg (Proxy :: Proxy Pixel8)      x y
      -- , ImageY16    <$> genImg (Proxy :: Proxy Pixel16)     x y
      -- , ImageYF     <$> genImg (Proxy :: Proxy PixelF)      x y
      -- , ImageYA8    <$> genImg (Proxy :: Proxy PixelYA8)    x y
      -- , ImageYA16   <$> genImg (Proxy :: Proxy PixelYA16)   x y
        ImageRGB8   <$> genImg (Proxy :: Proxy PixelRGB8)   x y
      , ImageRGB16  <$> genImg (Proxy :: Proxy PixelRGB16)  x y
      -- , ImageRGBF   <$> genImg (Proxy :: Proxy PixelRGBF)   x y
      , ImageRGBA8  <$> genImg (Proxy :: Proxy PixelRGBA8)  x y
      , ImageRGBA16 <$> genImg (Proxy :: Proxy PixelRGBA16) x y
      -- , ImageYCbCr8 <$> genImg (Proxy :: Proxy PixelYCbCr8) x y
      -- , ImageCMYK8  <$> genImg (Proxy :: Proxy PixelCMYK8)  x y
      -- , ImageCMYK16 <$> genImg (Proxy :: Proxy PixelCMYK16) x y
      ]
    return (ModelImage $ promoteImage $ convertRGBA8 dynImg)

max_stars_repo_name: CBMM/CBaaS

max_issues_repo_name: CBMM/CBaaS

max_forks_repo_name: CBMM/CBaaS


text:
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
module Math.HiddenMarkovModel.Distribution (
   State(..),
   Emission, Probability, Trained,
   Info(..), Generate(..), EmissionProb(..), Estimate(..),

   Discrete(..), DiscreteTrained(..),
   Gaussian(..), GaussianTrained(..), gaussian,

   CSV(..), HMMCSV.CSVParser, CSVSymbol(..),
   ) where

import qualified Math.HiddenMarkovModel.CSV as HMMCSV
import Math.HiddenMarkovModel.Utility (randomItemProp, normalizeProb)

import qualified Numeric.LinearAlgebra.HMatrix as HMatrix
import qualified Numeric.LinearAlgebra.Algorithms as Algo
import qualified Numeric.Container as NC
import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import Numeric.Container ((<>))
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)
import Foreign.Storable (Storable)

import qualified System.Random as Rnd

import qualified Text.CSV.Lazy.String as CSV
import Text.Read.HT (maybeRead)
import Text.Printf (printf)

import qualified Control.Monad.Exception.Synchronous as ME
import qualified Control.Monad.Trans.Class as MT
import qualified Control.Monad.Trans.State as MS
import Control.DeepSeq (NFData, rnf)
import Control.Monad (liftM2)

import qualified Data.NonEmpty as NonEmpty
import qualified Data.Foldable as Fold
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Array as Array
import qualified Data.List as List
import Data.Foldable (foldMap)
import Data.Tuple.HT (mapFst)
import Data.Array (Array, Ix, listArray, (!))
import Data.Map (Map)
import Data.Maybe (listToMaybe)

newtype State = State Int
   deriving (Eq, Ord, Show, Read, Ix)

instance Enum State where
   toEnum = State
   fromEnum (State n) = n

instance NFData State where
   rnf (State n) = rnf n

type family Probability distr
type family Emission distr
type family Trained distr

class
   (NC.Container Vector (Probability distr), NC.Product (Probability distr)) =>
      Info distr where
   numberOfStates :: distr -> Int

class
   (NC.Container Vector (Probability distr), NC.Product (Probability distr)) =>
      Generate distr where
   generate ::
      (Rnd.RandomGen g, Probability distr ~ prob, Emission distr ~ emission) =>
      distr -> State -> MS.State g emission

class
   (NC.Container Vector (Probability distr), NC.Product (Probability distr)) =>
      EmissionProb distr where
   {-
   This function could be implemented generically in terms of emissionStateProb
   but that would require an Info constraint.
   -}
   emissionProb :: distr -> Emission distr -> Vector (Probability distr)
   emissionStateProb :: distr -> Emission distr -> State -> Probability distr
   emissionStateProb distr e (State s) = NC.atIndex (emissionProb distr e) s

class
   (EmissionProb (Distribution tdistr),
    Trained (Distribution tdistr) ~ tdistr) =>
      Estimate tdistr where
   type Distribution tdistr
   accumulateEmissions ::
      (Distribution tdistr ~ distr, Probability distr ~ prob) =>
      [[(Emission distr, prob)]] -> tdistr
   -- could as well be in Semigroup class
   combine :: tdistr -> tdistr -> tdistr
   normalize :: (Distribution tdistr ~ distr) => tdistr -> distr

newtype Discrete prob symbol = Discrete (Map symbol (Vector prob))
   deriving (Show)

newtype DiscreteTrained prob symbol = DiscreteTrained (Map symbol (Vector prob))
   deriving (Show)

type instance Probability (Discrete prob symbol) = prob
type instance Emission (Discrete prob symbol) = symbol

type instance Trained (Discrete prob symbol) = DiscreteTrained prob symbol

instance (NFData prob, NFData symbol) => NFData (Discrete prob symbol) where
   rnf (Discrete m) = rnf m

instance
   (NFData prob, NFData symbol) =>
      NFData (DiscreteTrained prob symbol) where
   rnf (DiscreteTrained m) = rnf m

instance
   (NC.Container Vector prob, NC.Product prob, Ord symbol) =>
      Info (Discrete prob symbol) where
   numberOfStates (Discrete m) = Vector.dim $ snd $ Map.findMin m

instance
   (NC.Container Vector prob, NC.Product prob, Ord symbol,
    Ord prob, Rnd.Random prob) =>
      Generate (Discrete prob symbol) where
   generate (Discrete m) (State state) =
      randomItemProp $ Map.toAscList $ fmap (flip NC.atIndex state) m

instance
   (NC.Container Vector prob, NC.Product prob, Ord symbol) =>
      EmissionProb (Discrete prob symbol) where
   emissionProb (Discrete m) =
      mapLookup "emitDiscrete: unknown emission symbol" m

instance
   (NC.Container Vector prob, NC.Product prob, Ord symbol) =>
      Estimate (DiscreteTrained prob symbol) where
   type Distribution (DiscreteTrained prob symbol) = Discrete prob symbol
   accumulateEmissions grouped =
      let set = Set.toAscList $ foldMap (Set.fromList . map fst) grouped
          emi = Map.fromAscList $ zip set [0..]
      in  DiscreteTrained $ Map.fromAscList $ zip set $
          transposeVectorList $
          map
             (NC.accum (NC.konst 0 (length set)) (+) .
              map (mapFst
                 (mapLookup "estimateDiscrete: unknown emission symbol" emi)))
             grouped
   combine (DiscreteTrained distr0) (DiscreteTrained distr1) =
      DiscreteTrained $ Map.unionWith NC.add distr0 distr1
   normalize (DiscreteTrained distr) =
      Discrete $ Map.fromAscList $ zip (Map.keys distr) $
      transposeVectorList $ map normalizeProb $
      transposeVectorList $ Map.elems distr

transposeVectorList :: (Matrix.Element a) => [Vector a] -> [Vector a]
transposeVectorList = Matrix.toRows . Matrix.fromColumns

mapLookup :: (Ord k) => String -> Map.Map k a -> k -> a
mapLookup msg dict x =
   Map.findWithDefault (error msg) x dict

newtype Gaussian a = Gaussian (Array State (Vector a, Matrix a, a))
   deriving (Show)

newtype GaussianTrained a = GaussianTrained (Map State (Vector a, Matrix a, a))
   deriving (Show)

type instance Probability (Gaussian a) = a
type instance Emission (Gaussian a) = Vector a

type instance Trained (Gaussian a) = GaussianTrained a

instance (NFData a, Storable a) => NFData (Gaussian a) where
   rnf (Gaussian params) = rnf params

instance (NFData a, Storable a) => NFData (GaussianTrained a) where
   rnf (GaussianTrained params) = rnf params

instance (Algo.Field a) => Info (Gaussian a) where
   numberOfStates (Gaussian params) = Array.rangeSize $ Array.bounds params

instance (Algo.Field a) => Generate (Gaussian a) where
   generate (Gaussian allParams) state = do
      let (center, covarianceChol, _c) = allParams ! state
      seed <- MS.state Rnd.random
      return $
         NC.add center $
         NC.cmap realToFrac
               (NC.randomVector seed NC.Gaussian (Vector.dim center))
            <> covarianceChol

instance (HMatrix.Numeric a, Algo.Field a) => EmissionProb (Gaussian a) where
   emissionProb (Gaussian allParams) x =
      Vector.fromList $ map (emissionProbGen x) $ Array.elems allParams
   emissionStateProb (Gaussian allParams) x s =
      emissionProbGen x $ allParams ! s

emissionProbGen ::
   (HMatrix.Numeric a, Algo.Field a) =>
   Vector a -> (Vector a, Matrix a, a) -> a
emissionProbGen =
   let cholSolve m x = Matrix.flatten $ Algo.cholSolve m $ Matrix.asColumn x
   in  \x (center, covarianceChol, c) ->
         let x0 = NC.sub x center
         in  c * exp ((-1/2) * NC.dot x0 (cholSolve covarianceChol x0))

instance (HMatrix.Numeric a, Algo.Field a) => Estimate (GaussianTrained a) where
   type Distribution (GaussianTrained a) = Gaussian a
   accumulateEmissions =
      let params xs =
             let center =
                    NonEmpty.foldl1Map NC.add (uncurry $ flip NC.scale) xs
                 covariance =
                    NonEmpty.foldl1Map NC.add (\(x,c) -> NC.scale c $ NC.outer x x) xs
             in  (center, covariance, Fold.sum $ fmap snd xs)
      in  GaussianTrained . fmap params . Map.mapMaybe NonEmpty.fetch .
          Map.fromList . zip [State 0 ..]
   combine (GaussianTrained distr0) (GaussianTrained distr1) =
      let comb (center0, covariance0, weight0)
               (center1, covariance1, weight1) =
             (NC.add center0 center1,
              NC.add covariance0 covariance1,
              weight0 + weight1)
      in  GaussianTrained $ Map.unionWith comb distr0 distr1
   {-
     Sum_i (xi-mi) * (xi-mi)^T
   = Sum_i xi*xi^T + Sum_i mi*mi^T - Sum_i xi*mi^T - Sum_i mi*xi^T
   = Sum_i xi*xi^T - Sum_i mi*mi^T
   = Sum_i xi*xi^T - n * mi*mi^T
   -}
   normalize (GaussianTrained distr) =
      let params (centerSum, covarianceSum, weight) =
             let c = recip weight
                 center = NC.scale c centerSum
             in  (center,
                  NC.sub (NC.scale c covarianceSum) (NC.outer center center))
      in  Gaussian $
          Array.array (fst $ Map.findMin distr, fst $ Map.findMax distr) $
          Map.toList $ fmap (gaussianParameters . params) distr

gaussian ::
   (Algo.Field prob) =>
   [(Vector prob, Matrix prob)] -> Gaussian prob
gaussian =
   consGaussian . map gaussianParameters

gaussianParameters ::
   (Algo.Field prob) =>
   (Vector prob, Matrix prob) -> (Vector prob, Matrix prob, prob)
gaussianParameters (center, covariance) =
   gaussianFromCholesky center $ Algo.chol covariance

consGaussian :: [(Vector a, Matrix a, a)] -> Gaussian a
consGaussian xs =
   Gaussian $ listArray (State 0, State $ length xs - 1) xs

gaussianFromCholesky ::
   (Algo.Field prob) =>
   Vector prob -> Matrix prob -> (Vector prob, Matrix prob, prob)
gaussianFromCholesky center covarianceChol =
   let covarianceSqrtDet = NC.prodElements $ Matrix.takeDiag covarianceChol
   in  (center, covarianceChol,
        recip (sqrt (2*pi) ^ Vector.dim center * covarianceSqrtDet))

class CSV distr where
   toCells :: distr -> [[String]]
   parseCells :: Int -> HMMCSV.CSVParser distr

class (Ord symbol) => CSVSymbol symbol where
   cellFromSymbol :: symbol -> String
   symbolFromCell :: String -> Maybe symbol

instance CSVSymbol Char where
   cellFromSymbol = (:[])
   symbolFromCell = listToMaybe

instance CSVSymbol Int where
   cellFromSymbol = show
   symbolFromCell = maybeRead

instance
   (Algo.Field prob, Show prob, Read prob, CSVSymbol symbol) =>
      CSV (Discrete prob symbol) where
   toCells (Discrete m) =
      map
         (\(symbol, probs) ->
            cellFromSymbol symbol : HMMCSV.cellsFromVector probs) $
      Map.toAscList m
   parseCells n =
      fmap (Discrete . Map.fromList) $
      HMMCSV.manyRowsUntilEnd $ parseSymbolProb n

parseSymbolProb ::
   (Algo.Field prob, Read prob, CSVSymbol symbol) =>
   Int -> CSV.CSVRow -> HMMCSV.CSVParser (symbol, Vector prob)
parseSymbolProb n row =
   case row of
      [] -> MT.lift $ ME.throw "missing symbol"
      c:cs ->
         liftM2 (,)
            (let str = CSV.csvFieldContent c
             in  MT.lift $ ME.fromMaybe (printf "unknown symbol %s" str) $
                 symbolFromCell str)
            (do v <- HMMCSV.parseVectorFields cs
                HMMCSV.assert (n == Vector.dim v)
                   (printf "number of states (%d) and size of probability vector (%d) mismatch"
                      n (Vector.dim v))
                return v)

instance (Algo.Field a, Eq a, Show a, Read a) => CSV (Gaussian a) where
   toCells (Gaussian params) =
      List.intercalate [[]] $
      map
         (\(center, covarianceChol, _) ->
            HMMCSV.cellsFromVector center :
            HMMCSV.cellsFromMatrix covarianceChol) $
      Array.elems params
   parseCells n = do
      gs <- HMMCSV.manySepUntilEnd parseSingleGaussian
      HMMCSV.assert (length gs == n) $
         printf "number of states (%d) and number of Gaussians (%d) mismatch"
            n (length gs)
      return $ consGaussian gs

parseSingleGaussian ::
   (Algo.Field prob, Eq prob, Read prob) =>
   HMMCSV.CSVParser (Vector prob, Matrix prob, prob)
parseSingleGaussian = do
   center <- HMMCSV.parseNonEmptyVectorCells
   covarianceChol <- HMMCSV.parseSquareMatrixCells $ Vector.dim center
   HMMCSV.assert (isUpperTriang covarianceChol) $
      "matrices must be upper triangular"
   return $ gaussianFromCholesky center covarianceChol

{-
Maybe this test is too strict.
It would also be ok, and certainly more intuitive
to use an orthogonal but not normalized matrix.
We could get such a matrix from the eigensystem.
-}
isUpperTriang :: (Algo.Field a, Eq a) => Matrix a -> Bool
isUpperTriang m =
   let cleared = Matrix.mapMatrixWithIndex (\(i,j) x -> if i>j then x else 0) m
   in  NC.minElement cleared == 0 &&
       NC.maxElement cleared == 0

max_stars_repo_name: rybern/hmm-hmatrix

max_issues_repo_name: rybern/hmm-hmatrix

max_forks_repo_name: rybern/hmm-hmatrix


text:
module Mandelbrot where

import Diagrams.Backend.Cairo.CmdLine
import Diagrams.Prelude

import Data.Complex

quadratic :: Complex Double -> Complex Double -> Complex Double
quadratic c z = z * z + c

orbit :: Complex Double -> Complex Double -> [Complex Double]
orbit c = iterate (quadratic c)

criticalOrbit :: Complex Double -> [Complex Double]
criticalOrbit = flip orbit 0

maxIter :: Int
maxIter = 64

magOut :: Double
magOut = 2.0

pixel :: [Complex Double] -> Int
pixel = length . takeWhile (\z -> magnitude z <= magOut) . take maxIter

edge :: Int
edge = 256

side :: Int -> Double -> Double -> [Double]
side n v0 v1 =
    let sv = (v1 - v0) / fromIntegral n
    in  [v0, (v0 + sv) .. v1]

sideX :: [Double]
sideX = side edge (-2) 2

sideY :: [Double]
sideY = side edge (-2) 2

outerProduct :: (a -> b -> c) -> [a] -> [b] -> [[c]]
outerProduct f as bs = map (\a -> map (f a) bs) as

grid :: [[Complex Double]]
grid = outerProduct (flip (:+)) sideY sideX

toSquare :: Int -> Diagram B
toSquare n = square 1 # lw medium # fc black # opacity (sqrt o)
    where
        o = fromIntegral n / fromIntegral maxIter

mandel :: [[Diagram B]]
mandel = map (map (toSquare . pixel . criticalOrbit)) grid

mandelImg :: Diagram B
mandelImg = vcat . map hcat $ mandel

mandelBG :: Diagram B
mandelBG = mandelImg # bgFrame 3 pink

max_stars_repo_name: FayeAlephNil/diagrams-fun

max_issues_repo_name: FayeAlephNil/diagrams-fun

max_forks_repo_name: FayeAlephNil/diagrams-fun