import Criterion.Main
import Criterion.Main.Options
import Data.Store
import qualified Data.IntSet as Set
import qualified Data.ByteString as BS
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as UV
import Data.Array.Repa ((:.)(..), Z (..), fromUnboxed, computeUnboxedP, computeUnboxedS)

import MXNet.NN.DataIter.Anchor

main = do
    goodIndices <- BS.readFile "examples/goodIndices.bin" >>= decodeIO :: IO (V.Vector Int)
    gtBoxes     <- BS.readFile "examples/gtBoxes.bin"     >>= decodeIO :: IO (V.Vector (UV.Vector Float))
    anchors     <- BS.readFile "examples/anchors.bin"     >>= decodeIO :: IO (V.Vector (UV.Vector Float))
    goodIndices <- return $ Set.fromList $ V.toList goodIndices   :: IO Set.IntSet
    gtBoxes <- return $ V.map (fromUnboxed (Z:.(5::Int))) gtBoxes
    anchors <- return $ V.map (fromUnboxed (Z:.(4::Int))) anchors

    defaultMain
        [ bench "computeUnboxedP" $ whnfIO $ computeUnboxedP $ overlapMatrix goodIndices gtBoxes anchors
        , bench "computeUnboxedS" $ whnf computeUnboxedS $ overlapMatrix goodIndices gtBoxes anchors
        ]