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

import qualified Test.Generator as Gen
import Test.Logic (Dim)
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 qualified Numeric.LAPACK.Matrix.Shape.Omni as Omni
import qualified Numeric.LAPACK.Matrix.Layout as Layout
import qualified Numeric.LAPACK.Matrix as Matrix
import Numeric.LAPACK.Matrix.Layout (UnaryProxy)

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 ::
   Matrix.Banded sub super meas vert horiz height width a ->
   (UnaryProxy sub, UnaryProxy super)
offDiagonals = MatrixShape.bandedOffDiagonals . ArrMatrix.shape

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


shapeBandedFromFull ::
   (Unary.Natural sub, Unary.Natural super) =>
   (Layout.UnaryProxy sub, Layout.UnaryProxy super) ->
   MatrixShape.Full meas vert horiz height width ->
   MatrixShape.Banded sub super meas vert horiz height width
shapeBandedFromFull klu (Omni.Full (Layout.Full order extent)) =
   Omni.Banded $ Layout.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 :: (Dim size, Class.Floating a) => Gen.Square size a (Square size 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