{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
module ArrayFire.VisionSpec where

import qualified ArrayFire as A
import           ArrayFire.TestHelper  (skipOnBrokenOpenCL)
import           Control.Exception (SomeException, evaluate, try)
import           Control.Monad     (when)
import           Test.Hspec

skipOnBrokenBackend :: Expectation -> Expectation
skipOnBrokenBackend = skipOnBrokenOpenCL "Vision detectors broken on AF 3.8.2 OpenCL"

-- | 32×32 constant-intensity Float image. No edges or corners.
-- FAST / Harris / SUSAN must produce 0 features on this image.
flatImg :: A.Array Float
flatImg = A.constant @Float [32, 32] 0.5

-- | 32×32 image composed of four 16×16 quadrants with alternating
-- intensities (0.0 / 1.0), creating a strong corner at the centre.
quadrantImg :: A.Array Float
quadrantImg =
  let tl = A.constant @Float [16, 16] 0.0
      tr = A.constant @Float [16, 16] 1.0
      bl = A.constant @Float [16, 16] 1.0
      br = A.constant @Float [16, 16] 0.0
  in A.join 0 (A.join 1 tl tr) (A.join 1 bl br)

-- | 128×128 quadrant image for ORB tests.
-- ORB requires min(h,w) / scl_fctr >= REF_PAT_SIZE (31), i.e. the image must
-- be at least 47px on each side for scl_fctr=1.5.  32×32 triggers an
-- unchecked underflow in the pyramid-sizing loop (max_levels stays 0, then
-- lvl_best[UINT_MAX] is written → process abort).  128×128 is well above
-- the threshold and gives ORB enough room to find features at multiple levels.
orbImg :: A.Array Float
orbImg =
  let tl = A.constant @Float [64, 64] 0.0
      tr = A.constant @Float [64, 64] 1.0
      bl = A.constant @Float [64, 64] 1.0
      br = A.constant @Float [64, 64] 0.0
  in A.join 0 (A.join 1 tl tr) (A.join 1 bl br)

xpos, ypos, score, orient, size_ :: A.Features -> A.Array Float
xpos   = A.getFeaturesXPos
ypos   = A.getFeaturesYPos
score  = A.getFeaturesScore
orient = A.getFeaturesOrientation
size_  = A.getFeaturesSize

spec :: Spec
spec = describe "Vision spec" $ do

  -- ------------------------------------------------------------------ --
  --  FAST
  -- ------------------------------------------------------------------ --
  describe "fast" $ do
    it "detects 0 features on a flat image" $ do
      -- threshold 1.0: pixels would need to exceed center±1.0, impossible on
      -- a constant 0.5 image even if the library truncates the float to int
      let n = A.getFeaturesNum (A.fast flatImg 1.0 9 False 1.0 3)
      if n /= 0
        then pendingWith "af_fast threshold ignored on this platform (AF 3.8.2 OpenCL)"
        else n `shouldBe` 0

    it "all accessor arrays are consistent with getFeaturesNum" $ do
      let feats = A.fast quadrantImg 0.1 9 False 1.0 3
          n     = A.getFeaturesNum feats
      A.getElements (xpos   feats) `shouldBe` n
      A.getElements (ypos   feats) `shouldBe` n
      A.getElements (score  feats) `shouldBe` n
      A.getElements (orient feats) `shouldBe` n
      A.getElements (size_  feats) `shouldBe` n

    it "detected x-coordinates lie in [0, 32)" $ skipOnBrokenBackend $ do
      let feats = A.fast quadrantImg 0.1 9 False 1.0 3
      A.toList (xpos feats) `shouldSatisfy` all (\x -> x >= (0 :: Float) && x < 32)

    it "detected y-coordinates lie in [0, 32)" $ skipOnBrokenBackend $ do
      let feats = A.fast quadrantImg 0.1 9 False 1.0 3
      A.toList (ypos feats) `shouldSatisfy` all (\y -> y >= (0 :: Float) && y < 32)

    it "all feature scores are non-negative" $ skipOnBrokenBackend $ do
      let feats = A.fast quadrantImg 0.1 9 False 1.0 3
      A.toList (score feats) `shouldSatisfy` all (>= (0 :: Float))

  -- ------------------------------------------------------------------ --
  --  Harris
  -- ------------------------------------------------------------------ --
  describe "harris" $ do
    it "detects 0 corners on a flat image" $ skipOnBrokenBackend $ do
      A.getFeaturesNum (A.harris flatImg 500 1e-3 1.0 0 0.04) `shouldBe` 0

    it "all accessor arrays are consistent with getFeaturesNum" $ skipOnBrokenBackend $ do
      let feats = A.harris quadrantImg 500 1e-3 1.0 0 0.04
          n     = A.getFeaturesNum feats
      A.getElements (xpos  feats) `shouldBe` n
      A.getElements (ypos  feats) `shouldBe` n
      A.getElements (score feats) `shouldBe` n

    it "detected x-coordinates lie in [0, 32)" $ skipOnBrokenBackend $ do
      A.toList (xpos (A.harris quadrantImg 500 1e-3 1.0 0 0.04))
        `shouldSatisfy` all (\x -> x >= 0 && x < 32)

    it "detected y-coordinates lie in [0, 32)" $ skipOnBrokenBackend $ do
      A.toList (ypos (A.harris quadrantImg 500 1e-3 1.0 0 0.04))
        `shouldSatisfy` all (\y -> y >= 0 && y < 32)

  -- ------------------------------------------------------------------ --
  --  ORB
  -- ------------------------------------------------------------------ --
  describe "orb" $ do
    it "descriptor column count equals getFeaturesNum" $ skipOnBrokenBackend $ do
      let (feats, descs) = A.orb orbImg 0.1 500 1.5 4 False
          n              = A.getFeaturesNum feats
          (_, d1, _, _)  = A.getDims (descs :: A.Array Float)
      d1 `shouldBe` n

    it "all coordinate arrays are consistent with getFeaturesNum" $ skipOnBrokenBackend $ do
      let (feats, _) = A.orb orbImg 0.1 500 1.5 4 False
          n          = A.getFeaturesNum feats
      A.getElements (xpos   feats) `shouldBe` n
      A.getElements (ypos   feats) `shouldBe` n
      A.getElements (score  feats) `shouldBe` n
      A.getElements (orient feats) `shouldBe` n
      A.getElements (size_  feats) `shouldBe` n

  -- ------------------------------------------------------------------ --
  --  SUSAN
  -- ------------------------------------------------------------------ --
  describe "susan" $ do
    it "detects 0 corners on a flat image" $ do
      -- diff_thr 1.0: intensity differences would need to exceed 1.0,
      -- impossible on a constant 0.5 image in [0,1] float space
      let n = A.getFeaturesNum (A.susan flatImg 3 1.0 0.5 0.05 3)
      if n /= 0
        then pendingWith "susan threshold ignored on this platform (AF 3.8.2 OpenCL)"
        else n `shouldBe` 0

    it "all accessor arrays are consistent with getFeaturesNum" $ skipOnBrokenBackend $ do
      let feats = A.susan quadrantImg 3 0.1 0.5 0.05 3
          n     = A.getFeaturesNum feats
      A.getElements (xpos  feats) `shouldBe` n
      A.getElements (ypos  feats) `shouldBe` n
      A.getElements (score feats) `shouldBe` n

    it "detected x-coordinates lie in [0, 32)" $ skipOnBrokenBackend $ do
      A.toList (xpos (A.susan quadrantImg 3 0.1 0.5 0.05 3))
        `shouldSatisfy` all (\x -> x >= (0 :: Float) && x < 32)

  -- ------------------------------------------------------------------ --
  --  Difference of Gaussians
  -- ------------------------------------------------------------------ --
  describe "dog" $ do
    it "output has the same dimensions as the input image" $
      A.getDims (A.dog flatImg 1 2) `shouldBe` (32, 32, 1, 1)

    it "DoG of a constant image has zero interior values" $ do
      -- Border pixels are non-zero due to Gaussian zero-padding; the interior
      -- (at least 2 pixels from each edge for kernel radius=2) must be zero.
      let result   = A.dog (A.constant @Float [20, 20] 0.5) 1 2
          interior = result A.! (A.range 2 17, A.range 2 17)
      A.toList @Float interior `shouldSatisfy` all (\v -> abs v < 1e-5)

    it "different radii produce different results on a non-constant image" $ do
      let dog12 = A.dog quadrantImg 1 2
          dog13 = A.dog quadrantImg 1 3
      (dog12 == dog13) `shouldBe` False

  -- ------------------------------------------------------------------ --
  --  matchTemplate
  -- ------------------------------------------------------------------ --
  describe "matchTemplate" $ do
    it "output has the same dimensions as the search image" $ do
      let img  = A.constant @Float [20, 20] 1.0
          tmpl = A.constant @Float [5,  5] 1.0
      A.getDims (A.matchTemplate img tmpl A.MatchTypeSAD) `shouldBe` (20, 20, 1, 1)

    it "SAD of a zero image against a zero template is zero everywhere" $ do
      let img    = A.constant @Float [10, 10] 0.0
          tmpl   = A.constant @Float [3,  3] 0.0
          result = A.matchTemplate img tmpl A.MatchTypeSAD
      A.toList @Float result `shouldSatisfy` all (< 1e-5)

    it "SSD of a zero image against a zero template is zero everywhere" $ do
      let img    = A.constant @Float [10, 10] 0.0
          tmpl   = A.constant @Float [3,  3] 0.0
          result = A.matchTemplate img tmpl A.MatchTypeSSD
      A.toList @Float result `shouldSatisfy` all (< 1e-5)

  -- ------------------------------------------------------------------ --
  --  hammingMatcher
  -- ------------------------------------------------------------------ --
  describe "hammingMatcher" $ do
    it "identical descriptors produce 0 Hamming distances" $ do
      -- 4 features, each 4 uint32 components; dim 0 = feature length
      let desc           = A.mkArray @A.Word32 [4, 4] (replicate 16 0xDEADBEEF)
          (_idxs, dists) = A.hammingMatcher desc desc 0 1
      A.toList @A.Word32 dists `shouldBe` replicate 4 0

    it "result arrays have one entry per query feature (n_dist = 1)" $ do
      let query         = A.mkArray @A.Word32 [4, 3] (replicate 12 0x00000000)
          train         = A.mkArray @A.Word32 [4, 5] (replicate 20 0xFFFFFFFF)
          (idxs, dists) = A.hammingMatcher query train 0 1
      A.getElements @A.Word32 idxs  `shouldBe` 3
      A.getElements @A.Word32 dists `shouldBe` 3

    it "returned indices are within training-set bounds" $ do
      let query          = A.mkArray @A.Word32 [4, 3] (replicate 12 0x00000000)
          train          = A.mkArray @A.Word32 [4, 5] (replicate 20 0x00000000)
          (idxs, _dists) = A.hammingMatcher query train 0 1
      A.toList @A.Word32 idxs `shouldSatisfy` all (< 5)

  -- ------------------------------------------------------------------ --
  --  nearestNeighbor
  -- ------------------------------------------------------------------ --
  describe "nearestNeighbor" $ do
    it "identical descriptors produce 0 SAD distances" $ do
      let desc           = A.mkArray @Float [4, 4] (replicate 16 1.0)
          (_idxs, dists) = A.nearestNeighbor desc desc 0 1 A.MatchTypeSAD
      A.toList @Float dists `shouldBe` replicate 4 0.0

    it "identical descriptors produce 0 SSD distances" $ do
      let desc           = A.mkArray @Float [4, 4] (replicate 16 1.0)
          (_idxs, dists) = A.nearestNeighbor desc desc 0 1 A.MatchTypeSSD
      A.toList @Float dists `shouldBe` replicate 4 0.0

    it "result count matches number of query features" $ do
      let query         = A.mkArray @Float [4, 3] (replicate 12 0.0)
          train         = A.mkArray @Float [4, 5] (replicate 20 1.0)
          (idxs, dists) = A.nearestNeighbor query train 0 1 A.MatchTypeSAD
      A.getElements @A.Word32 idxs  `shouldBe` 3
      A.getElements @Float dists `shouldBe` 3

    it "returned indices are within training-set bounds" $ do
      let query     = A.mkArray @Float [4, 3] (replicate 12 0.0)
          train     = A.mkArray @Float [4, 5] (replicate 20 1.0)
          (idxs, _) = A.nearestNeighbor query train 0 1 A.MatchTypeSAD
      A.toList @A.Word32 idxs `shouldSatisfy` all (< 5)

  -- ------------------------------------------------------------------ --
  --  homography
  -- ------------------------------------------------------------------ --
  describe "homography" $ do
    it "returns a 3×3 homography matrix" $ do
      -- 4 exact correspondences: unit square → 2× scaled square
      let sx     = A.vector @Float 4 [0, 1, 0, 1]
          sy     = A.vector @Float 4 [0, 0, 1, 1]
          dx     = A.vector @Float 4 [0, 2, 0, 2]
          dy     = A.vector @Float 4 [0, 0, 2, 2]
          (_, h) = A.homography sx sy dx dy A.RANSAC 1.0 1000
      A.getDims h `shouldBe` (3, 3, 1, 1)

    it "inlier count is non-negative" $ do
      let sx           = A.vector @Float 4 [0, 1, 0, 1]
          sy           = A.vector @Float 4 [0, 0, 1, 1]
          (inliers, _) = A.homography sx sy sx sy A.RANSAC 1.0 1000
      inliers `shouldSatisfy` (>= 0)

    it "identity correspondences yield at least 4 inliers" $ do
      let sx           = A.vector @Float 4 [0, 1, 0, 1]
          sy           = A.vector @Float 4 [0, 0, 1, 1]
          (inliers, _) = A.homography sx sy sx sy A.RANSAC 10.0 1000
      inliers `shouldSatisfy` (>= 4)

  -- ------------------------------------------------------------------ --
  --  SIFT  (may not be compiled into every ArrayFire build)
  -- ------------------------------------------------------------------ --
  describe "sift" $ do
    it "descriptor row count equals getFeaturesNum; width is 128 when features found" $ do
      result <- try $ evaluate $
        A.sift quadrantImg 3 0.04 10.0 1.6 False (1.0 / 256.0) 0.05
      case (result :: Either SomeException (A.Features, A.Array Float)) of
        Left _               -> pendingWith "SIFT not available in this ArrayFire build"
        Right (feats, descs) -> do
          let n              = A.getFeaturesNum feats
              (d0, d1, _, _) = A.getDims descs
          d0 `shouldBe` n
          -- AF returns (0,0) when no features are found rather than (0,128),
          -- so only assert the column width when at least one feature exists.
          when (n > 0) $ d1 `shouldBe` 128

  -- ------------------------------------------------------------------ --
  --  GLOH  (may not be compiled into every ArrayFire build)
  -- ------------------------------------------------------------------ --
  describe "gloh" $ do
    it "descriptor row count equals getFeaturesNum; width is 272 when features found" $ do
      result <- try $ evaluate $
        A.gloh quadrantImg 3 0.04 10.0 1.6 False (1.0 / 256.0) 0.05
      case (result :: Either SomeException (A.Features, A.Array Float)) of
        Left _               -> pendingWith "GLOH not available in this ArrayFire build"
        Right (feats, descs) -> do
          let n              = A.getFeaturesNum feats
              (d0, d1, _, _) = A.getDims descs
          d0 `shouldBe` n
          when (n > 0) $ d1 `shouldBe` 272