{-# LANGUAGE TypeFamilies #-}
module Numeric.LAPACK.Matrix.Indexed where

import qualified Numeric.LAPACK.Matrix.Plain.Indexed as ArrIndexed
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Type as Type
import qualified Numeric.LAPACK.Permutation as Perm
import Numeric.LAPACK.Matrix.Type (Matrix)
import Numeric.LAPACK.Scalar (one, zero)
import Numeric.LAPACK.Permutation.Private (Permutation(Permutation))

import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Storable.Unchecked as UArray
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable ((!))


infixl 9 #!

class (Type.Box typ) => Indexed typ where
   (#!) ::
      (Class.Floating a) =>
      Matrix typ a ->
      (Shape.Index (Type.HeightOf typ), Shape.Index (Type.WidthOf typ)) -> a

instance (ArrIndexed.Indexed sh) => Indexed (ArrMatrix.Array sh) where
   ArrMatrix.Array arr #! :: Matrix (Array sh) a
-> (Index (HeightOf (Array sh)), Index (WidthOf (Array sh))) -> a
#! (Index (HeightOf (Array sh)), Index (WidthOf (Array sh)))
ij = Array sh a
arr Array sh a -> (Index (HeightOf sh), Index (WidthOf sh)) -> a
forall sh a.
(Indexed sh, Floating a) =>
Array sh a -> (Index (HeightOf sh), Index (WidthOf sh)) -> a
ArrIndexed.#! (Index (HeightOf sh), Index (WidthOf sh))
(Index (HeightOf (Array sh)), Index (WidthOf (Array sh)))
ij

instance (Shape.Indexed size) => Indexed (Type.Scale size) where
   Type.Scale sh a #! :: Matrix (Scale size) a
-> (Index (HeightOf (Scale size)), Index (WidthOf (Scale size)))
-> a
#! (Index (HeightOf (Scale size))
i,Index (WidthOf (Scale size))
j) =
      if size -> Index size -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset size
sh Index size
Index (HeightOf (Scale size))
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== size -> Index size -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset size
sh Index size
Index (WidthOf (Scale size))
j then a
a else a
forall a. Floating a => a
zero

instance (Shape.Indexed size) => Indexed (Permutation size) where
   Type.Permutation (Permutation perm) #! :: Matrix (Permutation size) a
-> (Index (HeightOf (Permutation size)),
    Index (WidthOf (Permutation size)))
-> a
#! (Index (HeightOf (Permutation size))
i,Index (WidthOf (Permutation size))
j) =
      let psh :: Shape size
psh@(Perm.Shape size
sh) = Array (Shape size) (Element size) -> Shape size
forall sh a. Array sh a -> sh
UArray.shape Array (Shape size) (Element size)
perm
          reindex :: Index size -> Element size
reindex = Shape size -> Int -> Index (Shape size)
forall sh. InvIndexed sh => sh -> Int -> Index sh
Shape.indexFromOffset Shape size
psh (Int -> Element size)
-> (Index size -> Int) -> Index size -> Element size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. size -> Index size -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset size
sh
      in if Array (Shape size) (Element size)
perm Array (Shape size) (Element size)
-> Index (Shape size) -> Element size
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
! Index size -> Element size
reindex Index size
Index (HeightOf (Permutation size))
i Element size -> Element size -> Bool
forall a. Eq a => a -> a -> Bool
== Index size -> Element size
reindex Index size
Index (WidthOf (Permutation size))
j then a
forall a. Floating a => a
one else a
forall a. Floating a => a
zero