content: {- |
  Module      : Langchain.VectorStore.InMemory
  Description : In-memory vector store implementation for LangChain Haskell
  Copyright   : (c) 2025 Tushar Adhatrao
  License     : MIT
  Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
  Stability   : experimental
  
  In-memory vector store implementation following LangChain's patterns, supporting:
  
  - Document storage with embeddings
  - Cosine similarity search
  - Integration with embedding models
  
  Example usage:
  
  @
  -- Create store with Ollama embeddings
  ollamaEmb = OllamaEmbeddings "nomic-embed" Nothing Nothing
  inMem = emptyInMemoryVectorStore ollamaEmb
  
  -- Add documents
  docs = [Document "Hello World" mempty, Document "Haskell is functional" mempty]
  updatedStore <- addDocuments inMem docs
  
  -- Perform similarity search
  results <- similaritySearch updatedStore "functional programming" 1
  -- Right [Document "Haskell is functional"...]
  @
  -}
  module Langchain.VectorStore.InMemory
    ( InMemory (..)
    , fromDocuments
    , emptyInMemoryVectorStore
    , norm
    , dotProduct
    , cosineSimilarity
    ) where
  
  import Data.Bifunctor
  import Data.Int (Int64)
  import Data.List (sortBy)
  import qualified Data.Map.Strict as Map
  import Data.Ord (comparing)
  import Langchain.DocumentLoader.Core (Document)
  import Langchain.Embeddings.Core
  import Langchain.Error (LangchainError)
  import Langchain.VectorStore.Core
  
  {- | Compute dot product of two vectors
  Example:
  
  >>> dotProduct [1,2,3] [4,5,6]
  32.0
  -}
  dotProduct :: [Float] -> [Float] -> Float
  dotProduct a b = sum $ zipWith (*) a b
  
  {- | Calculate Euclidean norm of a vector
  Example:
  
  >>> norm [3,4]
  5.0
  -}
  norm :: [Float] -> Float
  norm a = sqrt $ sum $ map (^ (2 :: Int)) a
  
  {- | Calculate cosine similarity between vectors
  Example:
  
  >>> cosineSimilarity [1,2] [2,4]
  1.0
  -}
  cosineSimilarity :: [Float] -> [Float] -> Float
  cosineSimilarity a b = dotProduct a b / (norm a * norm b)
  
  {- | Create empty in-memory store with embedding model
  Example:
  
  >>> emptyInMemoryVectorStore ollamaEmb
  InMemory {_embeddingModel = ..., _store = empty}
  -}
  emptyInMemoryVectorStore :: Embeddings m => m -> InMemory m
  emptyInMemoryVectorStore model = InMemory model Map.empty
  
  {- | Initialize store from documents using embeddings
  Example:
  
  >>> fromDocuments ollamaEmb [Document "Test" mempty]
  Right (InMemory {_store = ...})
  -}
  fromDocuments :: Embeddings m => m -> [Document] -> IO (Either LangchainError (InMemory m))
  fromDocuments model docs = do
    let vs = emptyInMemoryVectorStore model
    addDocuments vs docs
  
  {- | In-memory vector store implementation
  Stores documents with:
  
  - Embedding model reference
  - Map of document IDs to (Document, embedding) pairs
  -}
  data Embeddings m => InMemory m = InMemory
    { embeddingModel :: m
    , store :: Map.Map Int64 (Document, [Float])
    }
    deriving (Show, Eq)
  
  instance Embeddings m => VectorStore (InMemory m) where
    -- \| Add documents with generated embeddings
    --  Example:
    --
    --  >>> addDocuments inMem [doc1, doc2]
    --  Right (InMemory {_store = ...})
    --
    addDocuments inMem docs = do
      eRes <- embedDocuments (embeddingModel inMem) docs
      case eRes of
        Left err -> pure $ Left err
        Right floats -> do
          let currStore = store inMem
              mbMaxKey = Map.lookupMax currStore
              newStore =
                Map.fromList $
                  zip
                    [(maybe 1 (\x -> fst x + 1) mbMaxKey) ..]
                    (zip docs floats)
              newInMem = inMem {store = Map.union newStore currStore}
          pure $ Right newInMem
  
    -- \| Delete documents by ID
    --  Example:
    --
    --  >>> delete inMem [1, 2]
    --  Right (InMemory {_store = ...})
    --
    delete inMem ids = do
      let currStore = store inMem
          newStore = foldl (flip Map.delete) currStore ids
          newInMem = inMem {store = newStore}
      pure $ Right newInMem
  
    -- \| Text-based similarity search
    --  Example:
    --
    --  >>> similaritySearch inMem "Haskell" 2
    --  Right [Document "Haskell is...", Document "Functional programming..."]
    --
    similaritySearch vs query k = do
      eQueryEmbedding <- embedQuery (embeddingModel vs) query
      case eQueryEmbedding of
        Left err -> return $ Left err
        Right queryVec -> similaritySearchByVector vs queryVec k
  
    -- \| Vector-based similarity search
    --  Uses cosine similarity for ranking
    --
    --  Example:
    --
    --  >>> similaritySearchByVector inMem [0.1, 0.3, ...] 3
    --  Right [mostRelevantDoc, ...]
    --
    similaritySearchByVector vs queryVec k = do
      let similarities =
            map
              (second (cosineSimilarity queryVec) . snd)
              (Map.toList $ store vs)
          sorted = sortBy (comparing (negate . snd)) similarities
          -- Sort in descending order
          topK = take k sorted
      return $ Right $ map fst topK
  
  {-
  ghci> let x = OllamaEmbeddings "nomic-embed-text:latest" Nothing Nothing
  ghci> let inMem = emptyInMemoryVectorStore x
  ghci> eRes <- addDocuments inMem [Document "Hello World" empty, Document "Nice to meet you" empty]
  ghci> let newInMem = fromRight inMem eRes
  ghci> similaritySearch newInMem "World" 1
  Right [Document {pageContent = "Hello World", metadata = fromList []}]
  ghci> similaritySearch newInMem "Meet you" 1
  Right [Document {pageContent = "Nice to meet you", metadata = fromList []}]
  -}
  
  {- $examples
  Test case patterns:
  1. Document addition
     >>> addDocuments inMem [Document "Test" mempty]
     Right (InMemory {_store = ...})
  
  2. Similarity search
     >>> similaritySearch inMem "World" 1
     Right [Document "Hello World"...]
  
  3. Vector-based search
     >>> similaritySearchByVector inMem [0.5, 0.5] 1
     Right [mostSimilarDoc]
  -}
  