{-# LANGUAGE ExistentialQuantification #-}
module Test.Banded.Utility where

import qualified Test.Generator as Gen
import Test.Utility (genArray)

import qualified Numeric.LAPACK.Matrix.Banded as Banded
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import Numeric.LAPACK.Matrix.Shape (UnaryProxy)
import Numeric.LAPACK.Matrix (ShapeInt)

import qualified Numeric.Netlib.Class as Class

import qualified Type.Data.Num.Unary.Proof as Proof
import qualified Type.Data.Num.Unary as Unary
import Type.Data.Num.Unary (unary)
import Type.Base.Proxy (Proxy(Proxy))

import qualified Data.Array.Comfort.Shape as Shape

import Foreign.Storable (Storable)

import Data.Tuple.HT (mapPair)

import qualified Test.QuickCheck as QC


-- cf. MatrixShape.Private
natFromProxy :: (Unary.Natural n) => UnaryProxy n -> Proof.Nat n
natFromProxy Proxy = Proof.Nat

offDiagonals ::
   Banded.Banded sub super vert horiz height width a ->
   (UnaryProxy sub, UnaryProxy super)
offDiagonals = MatrixShape.bandedOffDiagonals . ArrMatrix.shape

offDiagonalNats ::
   (Unary.Natural sub, Unary.Natural super) =>
   Banded.Banded sub super vert horiz height width a ->
   (Proof.Nat sub, Proof.Nat super)
offDiagonalNats = mapPair (natFromProxy, natFromProxy) . offDiagonals


shapeBandedFromFull ::
   (MatrixShape.UnaryProxy sub, MatrixShape.UnaryProxy super) ->
   MatrixShape.Full vert horiz height width ->
   MatrixShape.Banded sub super vert horiz height width
shapeBandedFromFull klu (MatrixShape.Full order extent) =
   MatrixShape.Banded klu order extent


data Square size a =
   forall sub super.
   (Unary.Natural sub, Unary.Natural super) =>
   Square (Banded.Square sub super size a)

instance
   (Show size, Show a, Shape.C size, Storable a) =>
      Show (Square size a) where
   showsPrec p (Square a) = showsPrec p a

genSquare :: (Class.Floating a) => Gen.MatrixInt a (Square ShapeInt a)
genSquare = flip Gen.mapGenDim Gen.squareShape $ \maxElem maxDim shape -> do
   kl <- QC.choose (0, toInteger maxDim)
   ku <- QC.choose (0, toInteger maxDim)
   Unary.reifyNatural kl $ \sub ->
      Unary.reifyNatural ku $ \super ->
      fmap Square $
         genArray maxElem $ shapeBandedFromFull (unary sub, unary super) shape