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

import qualified Numeric.LAPACK.Matrix.Shape.Private as MatrixShape
import qualified Numeric.LAPACK.Matrix.Shape.Box as Box
import qualified Numeric.LAPACK.Matrix.Extent as Extent
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 (Array, (!))

import Data.Tuple.HT (swap)
import Data.Bool.HT (if')


infixl 9 #!

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

instance
   (Extent.C vert, Extent.C horiz, Shape.Indexed height, Shape.Indexed width) =>
      Indexed (MatrixShape.Full vert horiz height width) where
   #! :: Array (Full vert horiz height width) a
-> (Index (HeightOf (Full vert horiz height width)),
    Index (WidthOf (Full vert horiz height width)))
-> a
(#!) = (!)

instance (Shape.Indexed size) => Indexed (MatrixShape.Hermitian size) where
   Array (Hermitian size) a
arr#! :: Array (Hermitian size) a
-> (Index (HeightOf (Hermitian size)),
    Index (WidthOf (Hermitian size)))
-> a
#!(Index (HeightOf (Hermitian size)),
 Index (WidthOf (Hermitian size)))
ij =
      if Hermitian size -> Index (Hermitian size) -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds (Array (Hermitian size) a -> Hermitian size
forall sh a. Array sh a -> sh
UArray.shape Array (Hermitian size) a
arr) (Index (HeightOf (Hermitian size)),
 Index (WidthOf (Hermitian size)))
Index (Hermitian size)
ij
         then Array (Hermitian size) a
arr Array (Hermitian size) a -> Index (Hermitian size) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
UArray.! (Index (HeightOf (Hermitian size)),
 Index (WidthOf (Hermitian size)))
Index (Hermitian size)
ij
         else a -> a
forall a. Floating a => a -> a
conjugate (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ Array (Hermitian size) a
arr Array (Hermitian size) a -> Index (Hermitian size) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
! (Index size, Index size) -> (Index size, Index size)
forall a b. (a, b) -> (b, a)
swap (Index size, Index size)
(Index (HeightOf (Hermitian size)),
 Index (WidthOf (Hermitian size)))
ij

instance
   (MatrixShape.Content lo, MatrixShape.TriDiag diag, MatrixShape.Content up,
    Shape.Indexed size) =>
      Indexed (MatrixShape.Triangular lo diag up size) where
   Array (Triangular lo diag up size) a
arr#! :: Array (Triangular lo diag up size) a
-> (Index (HeightOf (Triangular lo diag up size)),
    Index (WidthOf (Triangular lo diag up size)))
-> a
#!(Index (HeightOf (Triangular lo diag up size)),
 Index (WidthOf (Triangular lo diag up size)))
ij =
      let sh :: Triangular lo diag up size
sh = Array (Triangular lo diag up size) a -> Triangular lo diag up size
forall sh a. Array sh a -> sh
UArray.shape Array (Triangular lo diag up size) a
arr
      in if Triangular lo diag up size
-> Index (Triangular lo diag up size) -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds Triangular lo diag up size
sh (Index (HeightOf (Triangular lo diag up size)),
 Index (WidthOf (Triangular lo diag up size)))
Index (Triangular lo diag up size)
ij
            then Array (Triangular lo diag up size) a
arr Array (Triangular lo diag up size) a
-> Index (Triangular lo diag up size) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
UArray.! (Index (HeightOf (Triangular lo diag up size)),
 Index (WidthOf (Triangular lo diag up size)))
Index (Triangular lo diag up size)
ij
            else
               (lo, up) -> a -> a -> a -> a -> a
forall lo up a.
(Content lo, Content up) =>
(lo, up) -> a -> a -> a -> a -> a
MatrixShape.caseDiagUpLoSym (Triangular lo diag up size -> (lo, up)
forall lo diag up size. Triangular lo diag up size -> (lo, up)
MatrixShape.triangularUplo Triangular lo diag up size
sh)
                  (String
-> Triangular lo diag up size -> (Index size, Index size) -> a
forall shape a height width.
(Box shape, Floating a, HeightOf shape ~ height, Indexed height,
 WidthOf shape ~ width, Indexed width) =>
String -> shape -> (Index height, Index width) -> a
checkedZero String
"Diagonal" Triangular lo diag up size
sh (Index size, Index size)
(Index (HeightOf (Triangular lo diag up size)),
 Index (WidthOf (Triangular lo diag up size)))
ij)
                  (String
-> Triangular lo diag up size -> (Index size, Index size) -> a
forall shape a height width.
(Box shape, Floating a, HeightOf shape ~ height, Indexed height,
 WidthOf shape ~ width, Indexed width) =>
String -> shape -> (Index height, Index width) -> a
checkedZero String
"UpperTriangular" Triangular lo diag up size
sh (Index size, Index size)
(Index (HeightOf (Triangular lo diag up size)),
 Index (WidthOf (Triangular lo diag up size)))
ij)
                  (String
-> Triangular lo diag up size -> (Index size, Index size) -> a
forall shape a height width.
(Box shape, Floating a, HeightOf shape ~ height, Indexed height,
 WidthOf shape ~ width, Indexed width) =>
String -> shape -> (Index height, Index width) -> a
checkedZero String
"LowerTriangular" Triangular lo diag up size
sh (Index size, Index size)
(Index (HeightOf (Triangular lo diag up size)),
 Index (WidthOf (Triangular lo diag up size)))
ij)
                  (Array (Triangular lo diag up size) a
arr Array (Triangular lo diag up size) a
-> Index (Triangular lo diag up size) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
! (Index size, Index size) -> (Index size, Index size)
forall a b. (a, b) -> (b, a)
swap (Index size, Index size)
(Index (HeightOf (Triangular lo diag up size)),
 Index (WidthOf (Triangular lo diag up size)))
ij)

instance
   (Unary.Natural sub, Unary.Natural super,
    Extent.C vert, Extent.C horiz, Shape.Indexed height, Shape.Indexed width) =>
      Indexed (MatrixShape.Banded sub super vert horiz height width) where
   Array (Banded sub super vert horiz height width) a
arr#! :: Array (Banded sub super vert horiz height width) a
-> (Index (HeightOf (Banded sub super vert horiz height width)),
    Index (WidthOf (Banded sub super vert horiz height width)))
-> a
#!(Index (HeightOf (Banded sub super vert horiz height width)),
 Index (WidthOf (Banded sub super vert horiz height width)))
ij =
      let 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
MatrixShape.InsideBox
          sh :: Banded sub super vert horiz height width
sh = Array (Banded sub super vert horiz height width) a
-> Banded sub super vert horiz height width
forall sh a. Array sh a -> sh
UArray.shape Array (Banded sub super vert horiz height width) a
arr
      in if Banded sub super vert horiz height width
-> Index (Banded sub super vert horiz height width) -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds Banded sub super vert horiz height width
sh (Index (Banded sub super vert horiz height width) -> Bool)
-> Index (Banded sub super vert horiz height width) -> Bool
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)
(Index (HeightOf (Banded sub super vert horiz height width)),
 Index (WidthOf (Banded sub super vert horiz height width)))
ij
            then Array (Banded sub super vert horiz height width) a
arr Array (Banded sub super vert horiz height width) a
-> Index (Banded sub super vert horiz height width) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
UArray.! (Index height, Index width)
-> BandedIndex (Index height) (Index width)
forall row column. (row, column) -> BandedIndex row column
boxIx (Index height, Index width)
(Index (HeightOf (Banded sub super vert horiz height width)),
 Index (WidthOf (Banded sub super vert horiz height width)))
ij
            else String
-> Banded sub super vert horiz height width
-> (Index height, Index width)
-> a
forall shape a height width.
(Box shape, Floating a, HeightOf shape ~ height, Indexed height,
 WidthOf shape ~ width, Indexed width) =>
String -> shape -> (Index height, Index width) -> a
checkedZero String
"Banded" Banded sub super vert horiz height width
sh (Index height, Index width)
(Index (HeightOf (Banded sub super vert horiz height width)),
 Index (WidthOf (Banded sub super vert horiz height width)))
ij

instance
   (Unary.Natural off, Shape.Indexed size) =>
      Indexed (MatrixShape.BandedHermitian off size) where
   Array (BandedHermitian off size) a
arr#! :: Array (BandedHermitian off size) a
-> (Index (HeightOf (BandedHermitian off size)),
    Index (WidthOf (BandedHermitian off size)))
-> a
#!(Index (HeightOf (BandedHermitian off size)),
 Index (WidthOf (BandedHermitian off size)))
ij =
      let 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
MatrixShape.InsideBox
          sh :: BandedHermitian off size
sh = Array (BandedHermitian off size) a -> BandedHermitian off size
forall sh a. Array sh a -> sh
UArray.shape Array (BandedHermitian off size) a
arr
      in Bool -> a -> a -> a
forall a. Bool -> a -> a -> a
if' (BandedHermitian off size
-> Index (BandedHermitian off size) -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds BandedHermitian off size
sh (Index (BandedHermitian off size) -> Bool)
-> Index (BandedHermitian off size) -> Bool
forall a b. (a -> b) -> a -> b
$ (Index size, Index size) -> BandedIndex (Index size) (Index size)
forall row column. (row, column) -> BandedIndex row column
boxIx (Index size, Index size)
(Index (HeightOf (BandedHermitian off size)),
 Index (WidthOf (BandedHermitian off size)))
ij)
            (Array (BandedHermitian off size) a
arr Array (BandedHermitian off size) a
-> Index (BandedHermitian off size) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
UArray.! (Index size, Index size) -> BandedIndex (Index size) (Index size)
forall row column. (row, column) -> BandedIndex row column
boxIx (Index size, Index size)
(Index (HeightOf (BandedHermitian off size)),
 Index (WidthOf (BandedHermitian off size)))
ij) (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$
         Bool -> a -> a -> a
forall a. Bool -> a -> a -> a
if' (BandedHermitian off size
-> Index (BandedHermitian off size) -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds BandedHermitian off size
sh (Index (BandedHermitian off size) -> Bool)
-> Index (BandedHermitian off size) -> Bool
forall a b. (a -> b) -> a -> b
$ (Index size, Index size) -> BandedIndex (Index size) (Index size)
forall row column. (row, column) -> BandedIndex row column
boxIx ((Index size, Index size) -> BandedIndex (Index size) (Index size))
-> (Index size, Index size)
-> BandedIndex (Index size) (Index size)
forall a b. (a -> b) -> a -> b
$ (Index size, Index size) -> (Index size, Index size)
forall a b. (a, b) -> (b, a)
swap (Index size, Index size)
(Index (HeightOf (BandedHermitian off size)),
 Index (WidthOf (BandedHermitian off size)))
ij)
            (a -> a
forall a. Floating a => a -> a
conjugate (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ Array (BandedHermitian off size) a
arr Array (BandedHermitian off size) a
-> Index (BandedHermitian off size) -> a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> a
UArray.! (Index size, Index size) -> BandedIndex (Index size) (Index size)
forall row column. (row, column) -> BandedIndex row column
boxIx ((Index size, Index size) -> (Index size, Index size)
forall a b. (a, b) -> (b, a)
swap (Index size, Index size)
(Index (HeightOf (BandedHermitian off size)),
 Index (WidthOf (BandedHermitian off size)))
ij))
         (String -> BandedHermitian off size -> (Index size, Index size) -> a
forall shape a height width.
(Box shape, Floating a, HeightOf shape ~ height, Indexed height,
 WidthOf shape ~ width, Indexed width) =>
String -> shape -> (Index height, Index width) -> a
checkedZero String
"BandedHermitian" BandedHermitian off size
sh (Index size, Index size)
(Index (HeightOf (BandedHermitian off size)),
 Index (WidthOf (BandedHermitian off size)))
ij)

checkedZero ::
   (Box.Box shape, Class.Floating a,
    Box.HeightOf shape ~ height, Shape.Indexed height,
    Box.WidthOf shape ~ width, Shape.Indexed width) =>
   String -> shape -> (Shape.Index height, Shape.Index width) -> a
checkedZero :: String -> shape -> (Index height, Index width) -> a
checkedZero String
name shape
sh (Index height, Index width)
ij =
   if (height, width) -> Index (height, width) -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds (shape -> HeightOf shape
forall shape. Box shape => shape -> HeightOf shape
Box.height shape
sh, shape -> WidthOf shape
forall shape. Box shape => shape -> WidthOf shape
Box.width shape
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"