content: {-# LANGUAGE OverloadedLists #-}
  
  module Test.ArraySpec where
  
  import Data.ByteString qualified as BS
  import Data.Massiv.Array as M
  import Data.Word (Word8)
  import GHC.Int
  import Skeletest
  import System.ByteOrder
  import Telescope.Data.Array
  import Telescope.Data.Axes
  
  
  spec :: Spec
  spec = do
    describe "indices" testIndices
    describe "decode" $ do
      -- describe "vector" testDecode
      describe "array" testDecodeArray
      describe "image" testDecodeImage
  
    describe "encode" $ do
      describe "array/image" testEncode
  
  
  testIndices :: Spec
  testIndices = do
    describe "axesIndex" $ do
      it "Ix1" $ do
        ix <- axesIndex $ toRowMajor $ Axes [10]
        ix `shouldBe` Ix1 10
  
      it "Ix2 in reverse order" $ do
        ix <- axesIndex $ toRowMajor $ Axes [3, 2]
        ix `shouldBe` 2 :. 3
  
      it "Ix3 in reverse order" $ do
        ix <- axesIndex $ toRowMajor $ Axes [3, 2, 1]
        ix `shouldBe` 1 :> 2 :. 3
  
    describe "indexAxes" $ do
      -- indexAxes produces row major
      it "Ix1" $ do
        indexAxes @Ix1 2 `shouldBe` Axes [2]
  
      it "Ix2 Row Major" $ do
        indexAxes (2 :. 3) `shouldBe` Axes [2, 3]
  
      it "Ix3 Row Major" $ do
        indexAxes (1 :> 2 :. 3) `shouldBe` Axes [1, 2, 3]
  
      it "Ix3 Column Major" $ do
        sizeAxes (Sz (1 :> 2 :. 3)) `shouldBe` Axes [3, 2, 1]
  
  
  -- generate a series of Word8
  genInput :: Word8 -> BS.ByteString
  genInput start = BS.pack [start .. start + 100]
  
  
  testDecodeArray :: Spec
  testDecodeArray = do
    let input = genInput 0
  
    it "should parse int8 vector" $ do
      let v = decodeVector @Int8 Par BigEndian input
      computeAs P (M.take 6 v) `shouldBe` [0, 1, 2, 3, 4, 5]
  
    it "should decode Ix1 from vector" $ do
      let v = decodeVector @Int8 Par BigEndian input
      a <- fromVector @Ix1 (Axes [3]) $ M.take 3 v
      computeAs P a `shouldBe` [0, 1, 2]
  
    it "should decode Ix2" $ do
      a <- decodeArray @Ix2 @Int8 [2, 3] $ BS.take 6 input
      computeAs P a `shouldBe` [[0, 1, 2], [3, 4, 5]]
  
    it "should decode Ix3" $ do
      a <- decodeArray @Ix3 @Int8 [2, 2, 2] $ BS.take 8 input
      computeAs P a `shouldBe` [[[0, 1], [2, 3]], [[4, 5], [6, 7]]]
  
  
  testDecodeImage :: Spec
  testDecodeImage = do
    let input = genInput 0
    it "should decode image" $ do
      a <- decodeArray @Ix2 @Int8 [2, 3] $ BS.take 6 input
      computeAs P a `shouldBe` [[0, 1, 2], [3, 4, 5]]
  
  
  testEncode :: Spec
  testEncode = do
    let array = delay @Ix2 @P $ M.fromLists' Seq [[1, 2, 3], [4, 5, 6]] :: Array D Ix2 Int8
  
    it "should encode to 6 items" $ do
      BS.length (encodeArray array) `shouldBe` 6
  
    it "int8 should round-trip to original array" $ do
      let enc = encodeArray array :: BS.ByteString
      arr <- decodeArray @Ix2 @Int8 [2, 3] enc
      arr `shouldBe` array
  
    let arrayFloat = delay @Ix2 @P $ M.fromLists' Seq [[1, 2, 3], [4, 5, 6]] :: Array D Ix2 Float
    it "float should round-trip to original array" $ do
      let enc = encodeArray arrayFloat :: BS.ByteString
      arr <- decodeArray [2, 3] enc
      arr `shouldBe` arrayFloat
  
    let arrayInt = delay @Ix2 @P $ M.fromLists' Seq [[1, 2, 3], [4, 5, 6]] :: Array D Ix2 Int
    it "int should round-trip to original array" $ do
      let enc = encodeArray arrayInt :: BS.ByteString
      arr <- decodeArray [2, 3] enc
      arr `shouldBe` arrayInt
  
    it "should encode naxes in correct order" $ do
      size array `shouldBe` Sz (2 :. 3)
      sizeAxes (size array) `shouldBe` Axes [3, 2]
  