content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  
  module Test.Massiv.Array.MutableSpec where
  
  import Data.Int
  import Data.Massiv.Array as A
  import GHC.Exts
  import Test.Massiv.Array.Delayed
  import Test.Massiv.Array.Load
  import Test.Massiv.Array.Mutable
  import Test.Massiv.Core
  import Test.Massiv.Core.Mutable
  
  type MutableArraySpec r ix e =
    ( Show e
    , Eq e
    , Arbitrary e
    , Arbitrary ix
    , CoArbitrary e
    , Function e
    , Eq (Array r ix e)
    , Show (Array r ix e)
    , Eq (Vector r e)
    , Show (Vector r e)
    , Load r ix e
    , Arbitrary (Array r ix e)
    , Manifest r e
    , Stream r ix e
    )
  
  type MutableSpec r e =
    ( Typeable e
    , IsList (Array r Ix1 e)
    , Item (Array r Ix1 e) ~ e
    , MutableArraySpec r Ix1 e
    , MutableArraySpec r Ix2 e
    , MutableArraySpec r Ix3 e
    , MutableArraySpec r Ix4 e
    , MutableArraySpec r Ix5 e
    )
  
  localMutableSpec :: forall r ix e. MutableArraySpec r ix e => Spec
  localMutableSpec = do
    describe "toStream/toList" $
      it "toStream" $
        property (prop_toStream @r @ix @e)
    describe "Manifest operations" $ do
      it "write" $ property (prop_Write @r @ix @e)
      it "modify" $ property (prop_Modify @r @ix @e)
      it "swap" $ property (prop_Swap @r @ix @e)
  
  specMutableR :: forall r e. MutableSpec r e => Spec
  specMutableR = do
    unsafeMutableSpec @r @Ix1 @e
    unsafeMutableSpec @r @Ix2 @e
    unsafeMutableSpec @r @Ix3 @e
    unsafeMutableSpec @r @Ix4 @e
    unsafeMutableSpec @r @Ix5 @e
    mutableSpec @r @Ix1 @e
    mutableSpec @r @Ix2 @e
    mutableSpec @r @Ix3 @e
    mutableSpec @r @Ix4 @e
    -- mutableSpec @r @Ix5 @e -- slows down the test suite
    loadSpec @r @Ix1 @e
    loadSpec @r @Ix2 @e
    loadSpec @r @Ix3 @e
    loadSpec @r @Ix4 @e
    localMutableSpec @r @Ix1 @e
    localMutableSpec @r @Ix2 @e
    localMutableSpec @r @Ix3 @e
    localMutableSpec @r @Ix4 @e
    localMutableSpec @r @Ix5 @e
    describe "NonFlat" $ do
      specMutableNonFlatR @r @Ix2 @e
      specMutableNonFlatR @r @Ix3 @e
      specMutableNonFlatR @r @Ix4 @e
      specMutableNonFlatR @r @Ix5 @e
    describe "toStream/toList" $
      it "toStreamIsList" $
        property (prop_toStreamIsList @r @e)
  
  specMutableNonFlatR
    :: forall r ix e
     . ( Arbitrary ix
       , Typeable e
       , Arbitrary e
       , Index (Lower ix)
       , Load r ix e
       , Manifest r e
       , Eq (Array r (Lower ix) e)
       , Show (Array r (Lower ix) e)
       , Show (Array r ix e)
       )
    => Spec
  specMutableNonFlatR = do
    describe (showsArrayType @r @ix @e "") $
      prop "outerSliceMArrayM" $
        prop_outerSliceMArrayM @r @ix @e
  
  specUnboxedMutableR :: forall r e. MutableSpec r e => Spec
  specUnboxedMutableR = do
    specMutableR @r @e
    unsafeMutableUnboxedSpec @r @Ix1 @e
    unsafeMutableUnboxedSpec @r @Ix2 @e
    unsafeMutableUnboxedSpec @r @Ix3 @e
    unsafeMutableUnboxedSpec @r @Ix4 @e
    unsafeMutableUnboxedSpec @r @Ix5 @e
  
  prop_Write
    :: forall r ix e
     . (Index ix, Manifest r e, Eq e, Show e)
    => Array r ix e
    -> ix
    -> e
    -> Property
  prop_Write arr ix e =
    monadicIO $
      run $ do
        marr <- thaw arr
        A.read marr ix >>= \case
          Nothing ->
            let withExcept = assertDeepExceptionIO (== IndexOutOfBoundsException (size arr) ix)
             in pure
                  ( withExcept (writeM marr ix e)
                      .&&. (write marr ix e `shouldReturn` False)
                      .&&. (write_ marr ix e `shouldReturn` ())
                  )
          Just olde ->
            pure $
              property $ do
                indexM arr ix `shouldReturn` olde
                A.write marr ix e `shouldReturn` True
                A.read marr ix `shouldReturn` Just e
  
                marr' <- thaw arr
                writeM marr' ix e `shouldReturn` ()
                arr' <- freeze (getComp arr) marr'
                indexM arr' ix `shouldReturn` e
  
                arr'' <- withMArray_ arr (\_ ma -> write_ ma ix e)
                index' arr'' ix `shouldBe` e
  
  prop_Modify
    :: forall r ix e
     . (Index ix, Manifest r e, Eq e, Show e)
    => Array r ix e
    -> Fun e e
    -> ix
    -> Property
  prop_Modify arr f ix =
    monadicIO $
      run $ do
        marr <- thaw arr
        modify marr (pure . apply f) ix >>= \case
          Nothing ->
            let withExcept = assertDeepExceptionIO (== IndexOutOfBoundsException (size arr) ix)
             in pure
                  ( withExcept (void $ indexM arr ix)
                      .&&. withExcept (void $ readM marr ix)
                      .&&. withExcept (void $ modifyM marr (pure . apply f) ix)
                      .&&. withExcept (modifyM_ marr (pure . apply f) ix)
                      .&&. (modify_ marr (pure . apply f) ix `shouldReturn` ())
                  )
          Just e ->
            pure $
              property $ do
                let fM = pure . apply f
                    fe = apply f e
                indexM arr ix `shouldReturn` e
                A.read marr ix `shouldReturn` Just fe
  
                marr' <- thawS arr
                readM marr' ix `shouldReturn` e
                modifyM marr' fM ix `shouldReturn` e
                arr' <- freezeS marr'
                indexM arr' ix `shouldReturn` fe
  
                arr'' <- withMArrayS_ arr (\ma -> modify_ ma fM ix)
                index' arr'' ix `shouldBe` fe
  
  prop_Swap
    :: forall r ix e
     . (Index ix, Manifest r e, Eq e, Show e)
    => Array r ix e
    -> ix
    -> ix
    -> Property
  prop_Swap arr ix1 ix2 =
    monadicIO $
      run $ do
        marr <- thaw arr
        swap marr ix1 ix2 >>= \case
          Nothing ->
            let withExcept =
                  assertDeepExceptionIO
                    ( \case
                        IndexOutOfBoundsException _ _ -> True
                        _ -> False
                    )
             in pure
                  ( withExcept (void $ indexM arr ix1 >> indexM arr ix2)
                      .&&. withExcept (void $ readM marr ix1 >> readM marr ix2)
                      .&&. withExcept (void $ swapM marr ix1 ix2)
                      .&&. withExcept (void $ swapM marr ix2 ix1)
                      .&&. withExcept (swapM_ marr ix1 ix2)
                      .&&. withExcept (swapM_ marr ix2 ix1)
                      .&&. (swap_ marr ix1 ix2 `shouldReturn` ())
                      .&&. (swap_ marr ix2 ix1 `shouldReturn` ())
                  )
          Just (e1, e2) ->
            pure $
              property $ do
                indexM arr ix1 `shouldReturn` e1
                indexM arr ix2 `shouldReturn` e2
                readM marr ix1 `shouldReturn` e2
                readM marr ix2 `shouldReturn` e1
  
                marr' <- thawS arr
                swapM marr' ix1 ix2 `shouldReturn` (e1, e2)
                arr' <- freezeS marr'
                indexM arr' ix1 `shouldReturn` e2
                indexM arr' ix2 `shouldReturn` e1
  
                let arr'' = withMArrayST_ arr (\ma -> swap_ ma ix1 ix2)
                index' arr'' ix1 `shouldBe` e2
                index' arr'' ix2 `shouldBe` e1
  
  prop_outerSliceMArrayM
    :: forall r ix e
     . ( Index ix
       , Index (Lower ix)
       , Manifest r e
       , Eq (Array r (Lower ix) e)
       , Show (Array r (Lower ix) e)
       )
    => ArrNE r ix e
    -> Property
  prop_outerSliceMArrayM (ArrNE arr) =
    forAll genIxInAndOut $ \(iIn, iOut) ->
      propIO $ do
        marr <- thawS arr
        (outerSliceMArrayM marr iIn >>= freezeS) `shouldReturn` arr !> iIn
        outerSliceMArrayM marr iOut `shouldThrow` (== IndexOutOfBoundsException (Sz nOuter) iOut)
    where
      (Sz nOuter, _) = unconsSz $ size arr
      genIxInAndOut = do
        let n = max 0 (nOuter - 1)
        iIn <- chooseInt (0, n)
        iOut <- oneof [chooseInt (minBound, -1), chooseInt (n, maxBound)]
        pure (iIn, iOut)
  
  spec :: Spec
  spec = do
    specMutableR @B @Int16
    specMutableR @BN @Int16
    specMutableR @BL @Int16
    atomicIntSpec @Ix1
    atomicIntSpec @Ix2
    atomicIntSpec @Ix3
    atomicIntSpec @Ix4
    atomicIntSpec @Ix5
  