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

import qualified Numeric.LAPACK.Matrix.Array.Private as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Shape.Omni as Omni
import qualified Numeric.LAPACK.Matrix.Layout.Private as Layout
import qualified Numeric.LAPACK.Matrix.Extent as Extent
import Numeric.LAPACK.Matrix.Array.Private (ArrayMatrix)
import Numeric.LAPACK.Scalar (conjugate, zero)

import qualified Numeric.Netlib.Class as Class

import qualified Type.Data.Num.Unary as Unary

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

import Foreign.Storable (Storable)

import Data.Maybe.HT (toMaybe)
import Data.Maybe (fromMaybe)
import Data.Tuple.HT (swap)


infixl 9 #!

(#!) ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Shape.Indexed height, Shape.Indexed width, Class.Floating a) =>
   ArrayMatrix pack property lower upper meas vert horiz height width a ->
   (Shape.Index height, Shape.Index width) -> a
ArrayMatrix
  pack property lower upper meas vert horiz height width a
a#! :: ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> (Index height, Index width) -> a
#!(Index height, Index width)
ij =
   let shape :: Omni pack property lower upper meas vert horiz height width
shape = ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> Omni pack property lower upper meas vert horiz height width
forall pack property lower upper meas vert horiz height width a.
ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> Omni pack property lower upper meas vert horiz height width
ArrMatrix.shape ArrayMatrix
  pack property lower upper meas vert horiz height width a
a
   in case Omni pack property lower upper meas vert horiz height width
shape of
         Omni.Full Full meas vert horiz height width
fullShape ->
            Full meas vert horiz height width
-> Array
     (Omni pack property lower upper meas vert horiz height width) a
-> Array (Full meas vert horiz height width) a
forall sh1 sh0 a. sh1 -> Array sh0 a -> Array sh1 a
UArray.reshape Full meas vert horiz height width
fullShape (ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> Array
     (Omni pack property lower upper meas vert horiz height width) a
forall pack property lower upper meas vert horiz height width a.
ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> OmniArray
     pack property lower upper meas vert horiz height width a
ArrMatrix.unwrap ArrayMatrix
  pack property lower upper meas vert horiz height width a
a) Array (Full meas vert horiz height width) a
-> Index (Full meas vert horiz height width) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
! (Index height, Index width)
Index (Full meas vert horiz height width)
ij
         Omni.UpperTriangular UpperTriangular height
_ ->
            a
-> ArrayMatrix
     pack property lower upper meas vert horiz height width a
-> (Index height, Index height)
-> a
forall pack prop lower upper meas vert horiz height width shape ix
       a.
(ToPlain pack prop lower upper meas vert horiz height width,
 Plain pack prop lower upper meas vert horiz height width ~ shape,
 Indexed shape, Index shape ~ ix, Storable a) =>
a
-> ArrayMatrix pack prop lower upper meas vert horiz height width a
-> ix
-> a
accessAlt (String
-> Omni pack property lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
forall meas vert horiz height width a pack property lower upper.
(Measure meas, C vert, C horiz, Indexed height, Indexed width,
 Floating a) =>
String
-> Omni pack property lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
checkedZero String
"UpperTriangular" Omni pack property lower upper meas vert horiz height width
shape (Index height, Index width)
ij) ArrayMatrix
  pack property lower upper meas vert horiz height width a
a (Index height, Index height)
(Index height, Index width)
ij
         Omni.LowerTriangular LowerTriangular height
_ ->
            a
-> ArrayMatrix
     pack property lower upper meas vert horiz height width a
-> (Index height, Index height)
-> a
forall pack prop lower upper meas vert horiz height width shape ix
       a.
(ToPlain pack prop lower upper meas vert horiz height width,
 Plain pack prop lower upper meas vert horiz height width ~ shape,
 Indexed shape, Index shape ~ ix, Storable a) =>
a
-> ArrayMatrix pack prop lower upper meas vert horiz height width a
-> ix
-> a
accessAlt (String
-> Omni pack property lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
forall meas vert horiz height width a pack property lower upper.
(Measure meas, C vert, C horiz, Indexed height, Indexed width,
 Floating a) =>
String
-> Omni pack property lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
checkedZero String
"LowerTriangular" Omni pack property lower upper meas vert horiz height width
shape (Index height, Index width)
ij) ArrayMatrix
  pack property lower upper meas vert horiz height width a
a (Index height, Index height)
(Index height, Index width)
ij
         Omni.Symmetric Symmetric height
_ ->
            a
-> ArrayMatrix
     pack property lower upper meas vert horiz height width a
-> (Index height, Index height)
-> a
forall pack prop lower upper meas vert horiz height width shape ix
       a.
(ToPlain pack prop lower upper meas vert horiz height width,
 Plain pack prop lower upper meas vert horiz height width ~ shape,
 Indexed shape, Index shape ~ ix, Storable a) =>
a
-> ArrayMatrix pack prop lower upper meas vert horiz height width a
-> ix
-> a
accessAlt (ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> PlainArray
     pack property lower upper meas vert horiz height width a
forall pack property lower upper meas vert horiz height width a.
ToPlain pack property lower upper meas vert horiz height width =>
ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> PlainArray
     pack property lower upper meas vert horiz height width a
ArrMatrix.toVector ArrayMatrix
  pack property lower upper meas vert horiz height width a
a Array (Symmetric height) a -> Index (Symmetric height) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
! (Index height, Index height) -> (Index height, Index height)
forall a b. (a, b) -> (b, a)
swap (Index height, Index height)
(Index height, Index width)
ij) ArrayMatrix
  pack property lower upper meas vert horiz height width a
a (Index height, Index height)
(Index height, Index width)
ij
         Omni.Hermitian Hermitian height
_ ->
            a
-> ArrayMatrix
     pack property lower upper meas vert horiz height width a
-> (Index height, Index height)
-> a
forall pack prop lower upper meas vert horiz height width shape ix
       a.
(ToPlain pack prop lower upper meas vert horiz height width,
 Plain pack prop lower upper meas vert horiz height width ~ shape,
 Indexed shape, Index shape ~ ix, Storable a) =>
a
-> ArrayMatrix pack prop lower upper meas vert horiz height width a
-> ix
-> a
accessAlt (a -> a
forall a. Floating a => a -> a
conjugate (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> PlainArray
     pack property lower upper meas vert horiz height width a
forall pack property lower upper meas vert horiz height width a.
ToPlain pack property lower upper meas vert horiz height width =>
ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> PlainArray
     pack property lower upper meas vert horiz height width a
ArrMatrix.toVector ArrayMatrix
  pack property lower upper meas vert horiz height width a
a Array (Hermitian height) a -> Index (Hermitian height) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
! (Index height, Index height) -> (Index height, Index height)
forall a b. (a, b) -> (b, a)
swap (Index height, Index height)
(Index height, Index width)
ij) ArrayMatrix
  pack property lower upper meas vert horiz height width a
a (Index height, Index height)
(Index height, Index width)
ij
         Omni.Banded Banded sub super meas vert horiz height width
_ -> ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> (Index height, Index width) -> a
forall pack prop lower upper meas vert horiz height width shape sub
       super a.
(ToPlain pack prop lower upper meas vert horiz height width,
 Plain pack prop lower upper meas vert horiz height width ~ shape,
 Banded sub super meas vert horiz height width ~ shape, Natural sub,
 Natural super, Measure meas, C vert, C horiz, Indexed height,
 Indexed width, Floating a) =>
ArrayMatrix pack prop lower upper meas vert horiz height width a
-> (Index height, Index width) -> a
accessBanded ArrayMatrix
  pack property lower upper meas vert horiz height width a
a (Index height, Index width)
ij
         Omni.UnitBandedTriangular BandedSquare sub super height
_ -> ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> (Index height, Index width) -> a
forall pack prop lower upper meas vert horiz height width shape sub
       super a.
(ToPlain pack prop lower upper meas vert horiz height width,
 Plain pack prop lower upper meas vert horiz height width ~ shape,
 Banded sub super meas vert horiz height width ~ shape, Natural sub,
 Natural super, Measure meas, C vert, C horiz, Indexed height,
 Indexed width, Floating a) =>
ArrayMatrix pack prop lower upper meas vert horiz height width a
-> (Index height, Index width) -> a
accessBanded ArrayMatrix
  pack property lower upper meas vert horiz height width a
a (Index height, Index width)
ij
         Omni.BandedHermitian BandedHermitian offDiag height
_ ->
            a
-> ArrayMatrix
     pack property lower upper meas vert horiz height width a
-> BandedIndex (Index height) (Index height)
-> a
forall pack prop lower upper meas vert horiz height width shape ix
       a.
(ToPlain pack prop lower upper meas vert horiz height width,
 Plain pack prop lower upper meas vert horiz height width ~ shape,
 Indexed shape, Index shape ~ ix, Storable a) =>
a
-> ArrayMatrix pack prop lower upper meas vert horiz height width a
-> ix
-> a
accessAlt
               (a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String
-> Omni pack property lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
forall meas vert horiz height width a pack property lower upper.
(Measure meas, C vert, C horiz, Indexed height, Indexed width,
 Floating a) =>
String
-> Omni pack property lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
checkedZero String
"BandedHermitian" Omni pack property lower upper meas vert horiz height width
shape (Index height, Index width)
ij) a -> a
forall a. Floating a => a -> a
conjugate (Maybe a -> a) -> Maybe a -> a
forall a b. (a -> b) -> a -> b
$
                Array (BandedHermitian offDiag height) a
-> Index (BandedHermitian offDiag height) -> Maybe a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> Maybe a
accessMaybe (ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> PlainArray
     pack property lower upper meas vert horiz height width a
forall pack property lower upper meas vert horiz height width a.
ToPlain pack property lower upper meas vert horiz height width =>
ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> PlainArray
     pack property lower upper meas vert horiz height width a
ArrMatrix.toVector ArrayMatrix
  pack property lower upper meas vert horiz height width a
a) (Index (BandedHermitian offDiag height) -> Maybe a)
-> Index (BandedHermitian offDiag height) -> Maybe a
forall a b. (a -> b) -> a -> b
$ (Index height, Index height)
-> BandedIndex (Index height) (Index height)
forall row column. (row, column) -> BandedIndex row column
boxIx ((Index height, Index height)
 -> BandedIndex (Index height) (Index height))
-> (Index height, Index height)
-> BandedIndex (Index height) (Index height)
forall a b. (a -> b) -> a -> b
$ (Index height, Index height) -> (Index height, Index height)
forall a b. (a, b) -> (b, a)
swap (Index height, Index height)
(Index height, Index width)
ij)
               ArrayMatrix
  pack property lower upper meas vert horiz height width a
a ((Index height, Index height)
-> BandedIndex (Index height) (Index height)
forall row column. (row, column) -> BandedIndex row column
boxIx (Index height, Index height)
(Index height, Index width)
ij)

accessBanded ::
   (Omni.ToPlain pack prop lower upper meas vert horiz height width) =>
   (Omni.Plain pack prop lower upper meas vert horiz height width ~ shape) =>
   (Layout.Banded sub super meas vert horiz height width ~ shape) =>
   (Unary.Natural sub, Unary.Natural super) =>
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Shape.Indexed height, Shape.Indexed width, Class.Floating a) =>
   ArrayMatrix pack prop lower upper meas vert horiz height width a ->
   (Shape.Index height, Shape.Index width) -> a
accessBanded :: ArrayMatrix pack prop lower upper meas vert horiz height width a
-> (Index height, Index width) -> a
accessBanded ArrayMatrix pack prop lower upper meas vert horiz height width a
a (Index height, Index width)
ij =
   a
-> ArrayMatrix pack prop lower upper meas vert horiz height width a
-> BandedIndex (Index height) (Index width)
-> a
forall pack prop lower upper meas vert horiz height width shape ix
       a.
(ToPlain pack prop lower upper meas vert horiz height width,
 Plain pack prop lower upper meas vert horiz height width ~ shape,
 Indexed shape, Index shape ~ ix, Storable a) =>
a
-> ArrayMatrix pack prop lower upper meas vert horiz height width a
-> ix
-> a
accessAlt (String
-> Omni pack prop lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
forall meas vert horiz height width a pack property lower upper.
(Measure meas, C vert, C horiz, Indexed height, Indexed width,
 Floating a) =>
String
-> Omni pack property lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
checkedZero String
"Banded" (ArrayMatrix pack prop lower upper meas vert horiz height width a
-> Omni pack prop lower upper meas vert horiz height width
forall pack property lower upper meas vert horiz height width a.
ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> Omni pack property lower upper meas vert horiz height width
ArrMatrix.shape ArrayMatrix pack prop lower upper meas vert horiz height width a
a) (Index height, Index width)
ij) ArrayMatrix pack prop lower upper meas vert horiz height width a
a (BandedIndex (Index height) (Index width) -> a)
-> BandedIndex (Index height) (Index width) -> a
forall a b. (a -> b) -> a -> b
$ (Index height, Index width)
-> BandedIndex (Index height) (Index width)
forall row column. (row, column) -> BandedIndex row column
boxIx (Index height, Index width)
ij

boxIx :: (row, column) -> Layout.BandedIndex row column
boxIx :: (row, column) -> BandedIndex row column
boxIx = (row -> column -> BandedIndex row column)
-> (row, column) -> BandedIndex row column
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry row -> column -> BandedIndex row column
forall row column. row -> column -> BandedIndex row column
Layout.InsideBox

accessAlt ::
   (Omni.ToPlain pack prop lower upper meas vert horiz height width) =>
   (Omni.Plain pack prop lower upper meas vert horiz height width ~ shape) =>
   (Shape.Indexed shape, Shape.Index shape ~ ix, Storable a) =>
   a -> ArrayMatrix pack prop lower upper meas vert horiz height width a ->
   ix -> a
accessAlt :: a
-> ArrayMatrix pack prop lower upper meas vert horiz height width a
-> ix
-> a
accessAlt a
alt ArrayMatrix pack prop lower upper meas vert horiz height width a
a = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
alt (Maybe a -> a) -> (ix -> Maybe a) -> ix -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array shape a -> Index shape -> Maybe a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> Maybe a
accessMaybe (ArrayMatrix pack prop lower upper meas vert horiz height width a
-> PlainArray pack prop lower upper meas vert horiz height width a
forall pack property lower upper meas vert horiz height width a.
ToPlain pack property lower upper meas vert horiz height width =>
ArrayMatrix
  pack property lower upper meas vert horiz height width a
-> PlainArray
     pack property lower upper meas vert horiz height width a
ArrMatrix.toVector ArrayMatrix pack prop lower upper meas vert horiz height width a
a)

accessMaybe ::
   (Shape.Indexed sh, Storable a) =>
   UArray.Array sh a -> Shape.Index sh -> Maybe a
accessMaybe :: Array sh a -> Index sh -> Maybe a
accessMaybe Array sh a
arr Index sh
ij =
   Bool -> a -> Maybe a
forall a. Bool -> a -> Maybe a
toMaybe (sh -> Index sh -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds (Array sh a -> sh
forall sh a. Array sh a -> sh
UArray.shape Array sh a
arr) Index sh
ij) (Array sh a
arr Array sh a -> Index sh -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
UArray.! Index sh
ij)

checkedZero ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Shape.Indexed height, Shape.Indexed width, Class.Floating a) =>
   String ->
   Omni.Omni pack property lower upper meas vert horiz height width ->
   (Shape.Index height, Shape.Index width) -> a
checkedZero :: String
-> Omni pack property lower upper meas vert horiz height width
-> (Index height, Index width)
-> a
checkedZero String
name Omni pack property lower upper meas vert horiz height width
sh (Index height, Index width)
ij =
   if (height, width) -> Index (height, width) -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds (Omni pack property lower upper meas vert horiz height width
-> height
forall meas vert horiz pack property lower upper height width.
(Measure meas, C vert, C horiz) =>
Omni pack property lower upper meas vert horiz height width
-> height
Omni.height Omni pack property lower upper meas vert horiz height width
sh, Omni pack property lower upper meas vert horiz height width
-> width
forall meas vert horiz pack property lower upper height width.
(Measure meas, C vert, C horiz) =>
Omni pack property lower upper meas vert horiz height width
-> width
Omni.width Omni pack property lower upper meas vert horiz height width
sh) (Index height, Index width)
Index (height, width)
ij
      then a
forall a. Floating a => a
zero
      else String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"Matrix.Indexed." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": index out of range"