module Main (main) where

import FAISS.Index
import FAISS.IndexFlat
import System.Random
import Text.Printf

-- Helper function to generate a matrix of random floats
-- For a real application, consider using a more robust library like mwc-random.
generateRandomMatrix :: Int -> Int -> IO [[Float]]
generateRandomMatrix n d = do
  gen <- newStdGen
  let randoms' = randomRs (0.0, 1.0) gen
      vectors = take (n * d) randoms'
  return $ chunksOf d vectors
  where
    chunksOf :: Int -> [a] -> [[a]]
    chunksOf _ [] = []
    chunksOf m xs = let (h, t) = splitAt m xs in h : chunksOf m t

main :: IO ()
main = do
  let d = 64      -- dimension of the vectors
      n_db = 2000 -- number of vectors in the database
      n_q = 5     -- number of query vectors
      k = 5       -- number of nearest neighbors to find

  -- Generate some random data for our database and queries
  databaseMatrix <- generateRandomMatrix n_db d
  queryMatrix <- generateRandomMatrix n_q d

  putStrLn $ "Creating an IndexFlatL2 index with dimension " ++ show d

-- Create a new IndexFlatL2 index
  eRes <- indexFlatL2NewWith d
  case eRes of
    Left err -> putStrLn $ "Error creating index: " <> err
    Right faissIndex -> do
      putStrLn "Index created successfully."
      
      -- Check if the index is trained (Flat indexes don't require training)
      isTrained <- indexGetIsTrained faissIndex 
      putStrLn $ "Is trained: " ++ show isTrained

      -- Add the database vectors to the index
      putStrLn $ "Adding " ++ show n_db ++ " vectors to the index..."
      eAdd <- indexAdd faissIndex databaseMatrix 
      case eAdd of
        Left err -> putStrLn $ "Error adding vectors: " <> err
        Right () -> do
          -- Get the total number of vectors in the index
          totalVectors <- indexGetNTotal faissIndex
          putStrLn $ "Total vectors in index: " ++ show totalVectors

          -- Search for the k nearest neighbors for each query vector 
          putStrLn $ "Searching for " ++ show k ++ " nearest neighbors for " ++ show n_q ++ " queries..."
          eSearch <- indexSearch faissIndex queryMatrix k 
          case eSearch of
            Left err -> putStrLn $ "Error during search: " <> err
            Right (distances, labels) -> do
              putStrLn "Search completed successfully. Top 5 query results:"
              let distanceChunks = chunksOf k distances
                  labelChunks = chunksOf k labels
              -- Nicely print the first few query results
              mapM_ (printQueryResult) (zip [0..] (zip distanceChunks labelChunks))

      -- Free the index from memory
      indexFree faissIndex
      putStrLn "Index freed. All good."

-- A small helper to pretty-print search results
printQueryResult :: (Int, ([[Float]], [[Int]])) -> IO ()
printQueryResult (i, (dists, lbls)) = do
    printf "Query %d:\n" i
    mapM_ (\(dist, lbl) -> printf "  Label: %-5d | Distance: %.4f\n" lbl dist) (zip (head lbls) (head dists))