{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Test.Massiv.Array.MutableSpec (spec) where

import Data.Massiv.Array as A
import Test.Massiv.Core
import Test.Massiv.Core.Mutable
import Test.Massiv.Array.Delayed
import Test.Massiv.Array.Mutable
import Test.Massiv.Array.Load
import GHC.Exts
import Data.Int

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
  loadSpec @r @Ix1 @e
  loadSpec @r @Ix2 @e
  loadSpec @r @Ix3 @e
  loadSpec @r @Ix4 @e
  --mutableSpec @r @Ix5 @e -- slows down the test suite
  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 = assertExceptionIO (== 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 = assertExceptionIO (== 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 =
              assertExceptionIO
                (\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
  specUnboxedMutableR @S @Int16
  specUnboxedMutableR @P @Int16
  specUnboxedMutableR @U @Int16
  atomicIntSpec @Ix1
  atomicIntSpec @Ix2
  atomicIntSpec @Ix3
  atomicIntSpec @Ix4
  atomicIntSpec @Ix5