content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeFamilies #-}
  
  module Test.Massiv.Array.Ops.TransformSpec (spec) where
  
  import qualified Data.Foldable as F
  import Data.Massiv.Array as A
  import Data.Maybe
  import Data.Sequence as S
  import Test.Massiv.Array.Delayed (stackSlices')
  import Test.Massiv.Core
  import Prelude as P
  
  prop_TransposeOuterInner :: Matrix D Int -> Property
  prop_TransposeOuterInner arr = transposeOuter arr === transpose arr
  
  prop_UpsampleDownsample
    :: forall r ix e
     . (Eq (Array r ix e), Show (Array r ix e), Load r ix e, Manifest r e)
    => ArrTiny r ix e
    -> Stride ix
    -> e
    -> Property
  prop_UpsampleDownsample (ArrTiny arr) stride fill =
    arr === compute (downsample stride (compute @r (upsample fill stride arr)))
  
  prop_ExtractAppend
    :: forall r ix e
     . (Eq (Array r ix e), Show (Array r ix e), Manifest r e, Index ix)
    => DimIx ix
    -> ArrIx r ix e
    -> Property
  prop_ExtractAppend (DimIx dim) (ArrIx arr ix) =
    arr === compute (uncurry (append' dim) $ A.splitAt' dim (getDim' ix dim) arr)
  
  prop_SplitExtract
    :: forall r ix e
     . ( Eq e
       , Show e
       , Eq (Array r ix e)
       , Show (Array r ix e)
       , Source r e
       , Load r ix e
       , Manifest r e
       , Ragged L ix e
       )
    => DimIx ix
    -> ArrIx r ix e
    -> Positive Int
    -> Property
  prop_SplitExtract (DimIx dim) (ArrIx arr ix) (Positive n) =
    ((compute @r <$> splitAt' dim i arr) === (left, compute @r (append' dim center right)))
      .&&. ((compute @r splitLeft, splitRight) === (compute @r (append' dim left center), right))
    where
      i = getDim' ix dim
      k = getDim' (unSz (size arr)) dim
      n' = n `mod` (k - i)
      (left, center, right) = throwEither (splitExtractM dim i (Sz n') arr)
      (splitLeft, splitRight) = splitAt' dim (i + n') arr
  
  prop_ConcatAppend
    :: forall r ix
     . (Eq (Array r ix Int), Show (Array r ix Int), Load r ix Int, Manifest r Int)
    => DimIx ix
    -> Comp
    -> Sz ix
    -> NonEmptyList (Fun ix Int)
    -> Property
  prop_ConcatAppend (DimIx dim) comp sz (NonEmpty fns) =
    foldl1 (\arr -> compute @r . append' dim arr) arrs
      === compute @r (concat' dim arrs)
    where
      arrs = P.zipWith (\f i -> makeArray @r comp sz ((+ i) . apply f)) fns [0 ..]
  
  prop_ConcatMConcatOuterM
    :: forall r ix
     . (Eq (Array r ix Int), Show (Array r ix Int), Load r ix Int, Manifest r Int)
    => Comp
    -> Sz ix
    -> NonEmptyList (Fun ix Int)
    -> Property
  prop_ConcatMConcatOuterM comp sz (NonEmpty fns) =
    property $ do
      as <- compute @r <$> concatM (dimensions sz) arrs
      as' <- compute @r <$> concatOuterM (P.map toLoadArray arrs)
      as `shouldBe` as'
    where
      arrs = P.zipWith (\f i -> makeArray @r comp sz ((+ i) . apply f)) fns [0 ..]
  
  prop_AppendMappend
    :: Array D Ix1 Int -> Array D Ix1 Int -> Property
  prop_AppendMappend arr1 arr2 =
    computeAs P (append' 1 arr1 arr2) === computeAs P (toLoadArray arr1 <> toLoadArray arr2)
  
  prop_ConcatMconcat
    :: [Array D Ix1 Int] -> Property
  prop_ConcatMconcat arrs =
    computeAs P (concat' 1 (A.empty : arrs)) === computeAs P (mconcat (fmap toLoadArray arrs))
  
  prop_ExtractSizeMismatch
    :: (Size r, Load r ix e, NFData (Array r Int e)) => ArrTiny r ix e -> Positive Int -> Property
  prop_ExtractSizeMismatch (ArrTiny arr) (Positive n) =
    assertDeepExceptionIO (SizeElementsMismatchException sz sz' ==) $ resizeM sz' arr
    where
      sz = size arr
      sz' = Sz (totalElem sz + n)
  
  -- FIXME: deal with overlapping instances for slices, see #106
  -- prop_stackInnerSlices ::
  --      forall ix.
  --      ( Index ix
  --      , Index (Lower ix)
  --      , Elt M ix Int ~ Array M (Lower ix) Int
  --      , Elt P ix Int ~ Array M (Lower ix) Int
  --      , Show (Array P ix Int)
  --      )
  --   => ArrNE P ix Int
  --   -> Property
  -- prop_stackInnerSlices (ArrNE arr) =
  --   (arr === compute (throwEither (stackInnerSlicesM (innerSlices arr)))) .&&.
  --   (arr === compute (stackSlices' 1 (innerSlices arr)))
  prop_stackInnerSlicesIx2 :: ArrNE P Ix2 Int -> Property
  prop_stackInnerSlicesIx2 (ArrNE arr) =
    (arr === compute (throwEither (stackInnerSlicesM (innerSlices arr))))
      .&&. (arr === compute (stackSlices' 1 (innerSlices arr)))
  
  prop_stackInnerSlicesIx3 :: ArrNE P Ix3 Int -> Property
  prop_stackInnerSlicesIx3 (ArrNE arr) =
    (arr === compute (throwEither (stackInnerSlicesM (innerSlices arr))))
      .&&. (arr === compute (stackSlices' 1 (innerSlices arr)))
  
  prop_stackInnerSlicesIx4 :: ArrNE P Ix4 Int -> Property
  prop_stackInnerSlicesIx4 (ArrNE arr) =
    (arr === compute (throwEither (stackInnerSlicesM (innerSlices arr))))
      .&&. (arr === compute (stackSlices' 1 (innerSlices arr)))
  
  -- prop_stackOuterSlices ::
  --      forall ix.
  --      ( Index ix
  --      , Index (Lower ix)
  --      , Elt M ix Int ~ Array M (Lower ix) Int
  --      , Elt P ix Int ~ Array M (Lower ix) Int
  --      , Show (Array P ix Int)
  --      )
  --   => ArrNE P ix Int
  --   -> Property
  -- prop_stackOuterSlices (ArrNE arr) =
  --   (arr === compute (throwEither (stackOuterSlicesM (outerSlices arr)))) .&&.
  --   (arr === compute (stackSlices' (dimensions (Proxy :: Proxy ix)) (outerSlices arr)))
  
  prop_stackOuterSlicesIx2 :: ArrNE P Ix2 Int -> Property
  prop_stackOuterSlicesIx2 (ArrNE arr) =
    (arr === compute (throwEither (stackOuterSlicesM (outerSlices arr))))
      .&&. (arr === compute (stackSlices' (dimensions (Proxy :: Proxy Ix2)) (outerSlices arr)))
  
  prop_stackOuterSlicesIx3 :: ArrNE P Ix3 Int -> Property
  prop_stackOuterSlicesIx3 (ArrNE arr) =
    (arr === compute (throwEither (stackOuterSlicesM (outerSlices arr))))
      .&&. (arr === compute (stackSlices' (dimensions (Proxy :: Proxy Ix3)) (outerSlices arr)))
  
  prop_stackOuterSlicesIx4 :: ArrNE P Ix4 Int -> Property
  prop_stackOuterSlicesIx4 (ArrNE arr) =
    (arr === compute (throwEither (stackOuterSlicesM (outerSlices arr))))
      .&&. (arr === compute (stackSlices' (dimensions (Proxy :: Proxy Ix4)) (outerSlices arr)))
  
  prop_ZoomWithGridStrideCompute
    :: forall r ix e
     . ( Eq (Array r ix e)
       , Show (Array r ix e)
       , StrideLoad r ix e
       , Manifest r e
       )
    => Array r ix e
    -> Stride ix
    -> e
    -> Property
  prop_ZoomWithGridStrideCompute arr stride defVal =
    ( computeWithStride @r stride' arr'
        === compute (A.replicate @DL Seq (Sz (liftIndex (+ 1) $ unSz (size arr))) defVal)
    )
      .&&. (computeWithStride @r stride' (extract' (pureIndex 1) sz' arr') === compute arr)
    where
      arr' = compute @r (zoomWithGrid defVal stride arr)
      sz' = Sz (liftIndex (subtract 1) $ unSz (size arr'))
      stride' = Stride (liftIndex (+ 1) $ unStride stride)
  
  prop_ZoomStrideCompute
    :: forall r ix e
     . (Eq (Array r ix e), Show (Array r ix e), StrideLoad r ix e, Manifest r e)
    => Array r ix e
    -> Stride ix
    -> Property
  prop_ZoomStrideCompute arr stride = computeWithStride @r stride arr' === compute arr
    where
      arr' = compute @r (zoom stride arr)
  
  type Transform r ix e =
    ( Show e
    , Eq e
    , Arbitrary e
    , Arbitrary ix
    , Arbitrary (Array r ix e)
    , Typeable e
    , Typeable ix
    , CoArbitrary e
    , CoArbitrary ix
    , Function e
    , Function ix
    , Eq (Array r ix e)
    , Eq (Array r ix Int)
    , Show (Array r ix e)
    , Show (Array r ix Int)
    , NFData (Array r ix e)
    , NFData (Array r Int e)
    , Load r ix e
    , Load r ix Int
    , Ragged L ix e
    , Source r e
    , StrideLoad r ix e
    , Manifest r Int
    , Manifest r e
    )
  
  specTransformR
    :: forall r ix e
     . Transform r ix e
    => Spec
  specTransformR =
    describe ("Transform (" ++ showsArrayType @r @ix @e ")") $ do
      prop "UpsampleDownsample" (prop_UpsampleDownsample @r @ix @e)
      prop "ExtractSizeMismatch" (prop_ExtractSizeMismatch @r @ix @e)
      prop "ExtractAppend" (prop_ExtractAppend @r @ix @e)
      prop "SplitExtract" (prop_SplitExtract @r @ix @e)
      prop "ConcatAppend" (prop_ConcatAppend @r @ix)
      prop "ConcatMConcatOuterM" (prop_ConcatMConcatOuterM @r @ix)
      prop "ZoomStrideCompute" (prop_ZoomStrideCompute @r @ix @e)
      prop "ZoomWithGridStrideCompute" (prop_ZoomWithGridStrideCompute @r @ix @e)
  
  spec :: Spec
  spec = do
    it "transposeOuterInner" $ property prop_TransposeOuterInner
    specTransformR @P @Ix1 @Int
    specTransformR @P @Ix2 @Int
    specTransformR @P @Ix3 @Int
    specTransformR @P @Ix4 @Int
    describe "Monoid" $ do
      it "Ix1" $ property prop_AppendMappend
      it "Ix1" $ property prop_ConcatMconcat
    describe "Sequence" $ do
      it "ConsSnoc" $ property prop_ConsSnoc
      it "UnconsUnsnoc" $ property prop_UnconsUnsnoc
    describe "slice+stack" $ do
      -- prop "Ix2 - Inner" (prop_stackInnerSlices @Ix2)
      -- prop "Ix3 - Inner" (prop_stackInnerSlices @Ix3)
      -- prop "Ix4 - Inner" (prop_stackInnerSlices @Ix4)
      -- prop "Ix2 - Outer" (prop_stackOuterSlices @Ix2)
      -- prop "Ix3 - Outer" (prop_stackOuterSlices @Ix3)
      -- prop "Ix4 - Outer" (prop_stackOuterSlices @Ix4)
      prop "Ix2 - Inner" prop_stackInnerSlicesIx2
      prop "Ix3 - Inner" prop_stackInnerSlicesIx3
      prop "Ix4 - Inner" prop_stackInnerSlicesIx4
      prop "Ix2 - Outer" prop_stackOuterSlicesIx2
      prop "Ix3 - Outer" prop_stackOuterSlicesIx3
      prop "Ix4 - Outer" prop_stackOuterSlicesIx4
  
  prop_UnconsUnsnoc :: Array D Ix1 Int -> Bool -> Property
  prop_UnconsUnsnoc arr unconsFirst =
    preJust $ do
      (arr', u, s) <-
        if unconsFirst
          then do
            (u, au) <- unconsM arr
            (as, s) <- unsnocM au
            pure (as, u, s)
          else do
            (as, s) <- unsnocM arr
            (u, au) <- unconsM as
            pure (au, u, s)
      pure (computeAs U (A.snoc (A.cons u (toLoadArray (computeAs U arr'))) s) === compute arr)
  
  preJust :: Testable prop => Maybe prop -> Property
  preJust m = isJust m ==> fromJust m
  
  prop_ConsSnoc :: Array D Ix1 Int -> [SeqOp Int] -> Property
  prop_ConsSnoc arr ops =
    A.toList (computeAs U (F.foldl' applyArraySeqOp (toLoadArray arr) ops))
      === F.toList (F.foldl' applySequenceSeqOp (S.fromList (A.toList arr)) ops)
  
  data SeqOp e = Cons e | Snoc e deriving (Eq, Show)
  
  instance Arbitrary e => Arbitrary (SeqOp e) where
    arbitrary = do
      e <- arbitrary
      elements [Cons e, Snoc e]
  
  applyArraySeqOp :: Array DL Ix1 e -> SeqOp e -> Array DL Ix1 e
  applyArraySeqOp arr = \case
    Cons x -> A.cons x arr
    Snoc x -> A.snoc arr x
  
  applySequenceSeqOp :: Seq a -> SeqOp a -> Seq a
  applySequenceSeqOp arr = \case
    Cons x -> x <| arr
    Snoc x -> arr |> x
  