content: {-# LANGUAGE OverloadedStrings #-}
  
  module Test.Langchain.VectorStore.Core (tests) where
  
  import Data.Either (fromRight, isRight)
  import Data.Int (Int64)
  import Data.Map (empty)
  import qualified Data.Map.Strict as Map
  import Test.Tasty
  import Test.Tasty.HUnit
  
  import Data.Maybe (fromMaybe, listToMaybe)
  import Langchain.DocumentLoader.Core (Document (..))
  import Langchain.Embeddings.Core
  import Langchain.VectorStore.Core
  import Langchain.VectorStore.InMemory
  
  data MockEmbeddings = MockEmbeddings
    deriving (Show, Eq)
  
  instance Embeddings MockEmbeddings where
    embedQuery _ "World" = pure $ Right [1.0, 0.1, 0.1]
    embedQuery _ "Meet you" = pure $ Right [0.1, 0.1, 1.0]
    embedQuery _ "Both" = pure $ Right [0.5, 0.5, 0.5]
    embedQuery _ _ = pure $ Right [0.0, 0.0, 0.0]
  
    embedDocuments _ docs = pure $ Right $ map determineEmbedding docs
      where
        determineEmbedding doc
          | doc == Document "Hello World" empty = [1.0, 0.1, 0.1]
          | doc == Document "Nice to meet you" empty = [0.1, 0.1, 1.0]
          | doc == Document "Something completely different" empty = [0.3, 0.3, 0.3]
          | otherwise = [0.0, 0.0, 0.0]
  
  createTestDocs :: [Document]
  createTestDocs =
    [ Document "Hello World" empty
    , Document "Nice to meet you" empty
    ]
  
  utilityTests :: TestTree
  utilityTests =
    testGroup
      "Utility Functions Tests"
      [ testCase "dotProduct should compute correct dot product" $ do
          dotProduct [1.0, 2.0, 3.0] [4.0, 5.0, 6.0] @?= 32.0
      , testCase "norm should compute correct Euclidean norm" $ do
          norm [3.0, 4.0] @?= 5.0
      , testCase "cosineSimilarity should compute correct similarity" $ do
          assertBool
            "Expected near same similarity"
            (cosineSimilarity [1.0, 2.0, 3.0] [1.0, 2.0, 3.0] >= 0.999999)
  
          let similarity = cosineSimilarity [1.0, 0.0, 0.0] [0.0, 1.0, 0.0]
          assertBool "Expected near 0" (abs similarity < 0.000001)
  
          let oppSimilarity = cosineSimilarity [1.0, 2.0, 3.0] [-1.0, -2.0, -3.0]
          assertBool "Expected near -1" (abs (oppSimilarity + 1.0) < 0.000001)
      ]
  
  inMemoryTests :: TestTree
  inMemoryTests =
    testGroup
      "InMemory VectorStore Tests"
      [ testCase "emptyInMemoryVectorStore should create empty store" $ do
          let model = MockEmbeddings
              vs = emptyInMemoryVectorStore model
          Map.size (store vs) @?= 0
          embeddingModel vs @?= model
      , testCase "fromDocuments should create store with documents" $ do
          let model = MockEmbeddings
              docs = createTestDocs
          result <- fromDocuments model docs
          assertBool "Expected Right result" (isRight result)
          let vs = fromRight (emptyInMemoryVectorStore model) result
          Map.size (store vs) @?= 2
      , testCase "addDocuments should add documents to store" $ do
          let model = MockEmbeddings
              vs = emptyInMemoryVectorStore model
              docs = createTestDocs
          result <- addDocuments vs docs
          assertBool "Expected Right result" (isRight result)
          let updatedVs = fromRight vs result
          Map.size (store updatedVs) @?= 2
  
          let newDoc = Document "Something completely different" empty
          result2 <- addDocuments updatedVs [newDoc]
          assertBool "Expected Right result" (isRight result2)
          let finalVs = fromRight updatedVs result2
          Map.size (store finalVs) @?= 3
      , testCase "delete should remove documents from store" $ do
          let model = MockEmbeddings
              vs = emptyInMemoryVectorStore model
              docs = createTestDocs
          result <- addDocuments vs docs
          let updatedVs = fromRight vs result
  
          deleteResult <- delete updatedVs [1]
          assertBool "Expected Right result" (isRight deleteResult)
          let afterDeleteVs = fromRight updatedVs deleteResult
          Map.size (store afterDeleteVs) @?= 1
          Map.member (1 :: Int64) (store afterDeleteVs) @?= False
          Map.member (2 :: Int64) (store afterDeleteVs) @?= True
      , testCase "similaritySearch should find similar documents" $ do
          let model = MockEmbeddings
              vs = emptyInMemoryVectorStore model
              docs = createTestDocs
          result <- addDocuments vs docs
          let updatedVs = fromRight vs result
  
          -- Search for "World" - should return "Hello World"
          searchResult1 <- similaritySearch updatedVs "World" 1
          assertBool "Expected Right result" (isRight searchResult1)
          let docs1 = fromRight [] searchResult1
          length docs1 @?= 1
          fromMaybe (Document "" empty) (listToMaybe docs1) @?= Document "Hello World" empty
  
          -- Search for "Meet you" - should return "Nice to meet you"
          searchResult2 <- similaritySearch updatedVs "Meet you" 1
          assertBool "Expected Right result" (isRight searchResult2)
          let docs2 = fromRight [] searchResult2
          length docs2 @?= 1
          fromMaybe (Document "" empty) (listToMaybe docs2) @?= Document "Nice to meet you" empty
  
          -- Search for both documents
          searchResult3 <- similaritySearch updatedVs "Both" 2
          assertBool "Expected Right result" (isRight searchResult3)
          let docs3 = fromRight [] searchResult3
          length docs3 @?= 2
      , testCase "similaritySearchByVector should find similar documents" $ do
          let model = MockEmbeddings
              vs = emptyInMemoryVectorStore model
              docs = createTestDocs
          result <- addDocuments vs docs
          let updatedVs = fromRight vs result
  
          -- Search with vector similar to "Hello World"
          searchResult1 <- similaritySearchByVector updatedVs [1.0, 0.1, 0.1] 1
          assertBool "Expected Right result" (isRight searchResult1)
          let docs1 = fromRight [] searchResult1
          length docs1 @?= 1
          fromMaybe (Document "" empty) (listToMaybe docs1) @?= Document "Hello World" empty
  
          -- Search with vector similar to "Nice to meet you"
          searchResult2 <- similaritySearchByVector updatedVs [0.1, 0.1, 1.0] 1
          assertBool "Expected Right result" (isRight searchResult2)
          let docs2 = fromRight [] searchResult2
          length docs2 @?= 1
          fromMaybe (Document "" empty) (listToMaybe docs2) @?= Document "Nice to meet you" empty
      ]
  
  tests :: TestTree
  tests =
    testGroup
      "Langchain.VectorStore Tests"
      [ utilityTests
      , inMemoryTests
      ]
  