{-# 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
   (#!) = (!)

instance (Shape.Indexed size) => Indexed (MatrixShape.Hermitian size) where
   arr#!ij =
      if Shape.inBounds (UArray.shape arr) ij
         then arr UArray.! ij
         else conjugate $ arr ! swap ij

instance
   (MatrixShape.Content lo, MatrixShape.TriDiag diag, MatrixShape.Content up,
    Shape.Indexed size) =>
      Indexed (MatrixShape.Triangular lo diag up size) where
   arr#!ij =
      let sh = UArray.shape arr
      in if Shape.inBounds sh ij
            then arr UArray.! ij
            else
               MatrixShape.caseDiagUpLoSym (MatrixShape.triangularUplo sh)
                  (checkedZero "Diagonal" sh ij)
                  (checkedZero "UpperTriangular" sh ij)
                  (checkedZero "LowerTriangular" sh ij)
                  (arr ! swap 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
   arr#!ij =
      let boxIx = uncurry MatrixShape.InsideBox
          sh = UArray.shape arr
      in if Shape.inBounds sh $ boxIx ij
            then arr UArray.! boxIx ij
            else checkedZero "Banded" sh ij

instance
   (Unary.Natural off, Shape.Indexed size) =>
      Indexed (MatrixShape.BandedHermitian off size) where
   arr#!ij =
      let boxIx = uncurry MatrixShape.InsideBox
          sh = UArray.shape arr
      in if' (Shape.inBounds sh $ boxIx ij)
            (arr UArray.! boxIx ij) $
         if' (Shape.inBounds sh $ boxIx $ swap ij)
            (conjugate $ arr UArray.! boxIx (swap ij))
         (checkedZero "BandedHermitian" sh 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 name sh ij =
   if Shape.inBounds (Box.height sh, Box.width sh) ij
      then zero
      else error $ "Matrix.Indexed." ++ name ++ ": index out of range"