{-# LANGUAGE TypeFamilies #-}
module Test.Indexed where

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

import qualified Numeric.LAPACK.Matrix.Shape.Box as Box
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix ((#!), (#>))

import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable (Array)

import qualified Data.Traversable as Trav
import Data.Maybe.HT (toMaybe)

import qualified Test.QuickCheck as QC


genMatrixIndexGen ::
   (Class.Floating a) =>
   (array -> [ix]) ->
   Gen.Matrix a Int Int array ->
   Gen.Matrix a Int Int (Maybe ix, array)
genMatrixIndexGen indices gen =
   flip Gen.mapGen gen $ \_maxElem m -> do
      let set = indices m
      ij <- Trav.mapM QC.elements $ toMaybe (not $ null set) set
      return (ij,m)

genMatrixIndex ::
   (Shape.Indexed shape, Class.Floating a) =>
   Gen.Matrix a Int Int (Array shape a) ->
   Gen.Matrix a Int Int (Maybe (Shape.Index shape), Array shape a)
genMatrixIndex = genMatrixIndexGen (Shape.indices . Array.shape)

unitDot ::
   (Matrix.Indexed shape, Matrix.MultiplyRight shape,
    Box.HeightOf shape ~ height, Shape.Indexed height, Eq height,
    Box.WidthOf shape ~ width, Shape.Indexed width,
    Class.Floating a, Eq a) =>
   (Maybe (Shape.Index height, Shape.Index width), Array shape a) -> QC.Property
unitDot (mij,m) =
   maybeProperty $
   flip fmap mij $ \(i,j) ->
      m#!(i,j) ==
      Vector.dot
         (Vector.unit (Matrix.height m) i)
         (m #> Vector.unit (Matrix.width m) j)