--[[----------------------------------------------------------------------------

WFMLEmbeddingsTest.lua

--------------------------------------------------------------------------------

ADOBE CONFIDENTIAL
------------------
Copyright 2025 Adobe
All Rights Reserved.

NOTICE: All information contained herein is, and remain the property of
Adobe and its suppliers, if any. The intellectual and
technical concepts contained herein are proprietary to Adobe Systems
Incorporated and its suppliers and may be covered by U.S. and Foreign
Patents, patents in process, and are protected by trade secret or copyright
law.  Dissemination of this information or reproduction of this material is
strictly forbidden unless prior written permission is obtained from
Adobe.

------------------------------------------------------------------------------]]

local WFFileUtils = require 'WFFileUtils'
local KSIdGenerator = require 'KSIdGenerator'
local KSPathUtils = require "KSPathUtils"
local WFPlatformInfo = require "WFPlatformInfo"
local WFStruct = require "WFStruct"
local WFMLImp = require "WFMLImp"
local WFRxTestCase = require "WFRxTestCase"

--------------------------------------------------------------------------------

if WFPlatformInfo.isAndroid then
	return {}
end

--------------------------------------------------------------------------------

local WFMLEmbeddingsTest = {}

--------------------------------------------------------------------------------

function WFMLEmbeddingsTest.setUp()

	local documentDir = KSPathUtils.getStandardFilePath( 'carouselDocuments' )
	local baseName = KSPathUtils.child( documentDir, 'embeddings' )

	WFFileUtils.delete( baseName .. '.idx' )
	WFFileUtils.delete( baseName .. '.dat' )

end

WFMLEmbeddingsTest.tearDown = WFMLEmbeddingsTest.setUp

--------------------------------------------------------------------------------
--@mark Standard Tests
--------------------------------------------------------------------------------

function WFMLEmbeddingsTest.rxTestEraseFind( rxTest )

	local documentDir = KSPathUtils.getStandardFilePath( 'carouselDocuments' )
	local baseName = KSPathUtils.child( documentDir, 'embeddings' )

	local embeddings = WFMLImp.WFMLEmbeddings( baseName, 2 )

	local id1 = KSIdGenerator.uuidInLowercaseHex()
	local v1 = WFStruct.pack( '2f', 1.0, 1.0 )

	embeddings:insert( id1, v1 )
	rxTest:expectEqual( true, embeddings:hasKey( id1 ) )

	local id2 = KSIdGenerator.uuidInLowercaseHex()
	local v2 = WFStruct.pack( '2f', 0.0, 0.0 )

	embeddings:insert( id2, v2 )
	rxTest:expectEqual( true, embeddings:hasKey( id2 ) )

	local id3 = KSIdGenerator.uuidInLowercaseHex()
	local v3 = WFStruct.pack( '2f', 1.0, 0.9 )

	embeddings:insert( id3, v3 )
	rxTest:expectEqual( true, embeddings:hasKey( id3 ) )

	local id4 = KSIdGenerator.uuidInLowercaseHex()
	local v4 = WFStruct.pack( '2f', 0.0, 1.0 )

	embeddings:insert( id4, v4 )
	rxTest:expectEqual( true, embeddings:hasKey( id4 ) )

	embeddings:erase( id1 )
	rxTest:expectEqual( false, embeddings:hasKey( id1 ) )

	local vs = WFStruct.pack( '2f', 0.9, 0.9 )

	local matches = embeddings:findClosest( vs, 2, 0.0 )

	rxTest:expectEqual( 2, #matches )
	rxTest:expectEqual( id3, matches[1].id )

end

--------------------------------------------------------------------------------

function WFMLEmbeddingsTest.rxTestInsertFind( rxTest )

	local documentDir = KSPathUtils.getStandardFilePath( 'carouselDocuments' )
	local baseName = KSPathUtils.child( documentDir, 'embeddings' )

	local embeddings = WFMLImp.WFMLEmbeddings( baseName, 2 )

	rxTest:expectEqual( 2, embeddings.vectorSize )
	rxTest:expectEqual( baseName, embeddings.baseName )

	local id1 = KSIdGenerator.uuidInLowercaseHex()
	local v1 = WFStruct.pack( '2f', 1.0, 1.0 )

	embeddings:insert( id1, v1 )
	rxTest:expectEqual( true, embeddings:hasKey( id1 ) )

	local id2 = KSIdGenerator.uuidInLowercaseHex()
	local v2 = WFStruct.pack( '2f', 0.0, 0.0 )

	embeddings:insert( id2, v2 )
	rxTest:expectEqual( true, embeddings:hasKey( id2 ) )

	local id3 = KSIdGenerator.uuidInLowercaseHex()
	local v3 = WFStruct.pack( '2f', 1.0, 0.0 )

	embeddings:insert( id3, v3 )
	rxTest:expectEqual( true, embeddings:hasKey( id3 ) )

	local id4 = KSIdGenerator.uuidInLowercaseHex()
	local v4 = WFStruct.pack( '2f', 0.0, 1.0 )

	embeddings:insert( id4, v4 )
	rxTest:expectEqual( true, embeddings:hasKey( id4 ) )

	local vs = WFStruct.pack( '2f', 0.9, 0.9 )

	local matches = embeddings:findClosest( vs, 2, 0.0 )

	rxTest:expectEqual( 2, #matches )
	rxTest:expectEqual( id1, matches[1].id )

end

--------------------------------------------------------------------------------

function WFMLEmbeddingsTest.rxTestInsertOffload( rxTest )

	local documentDir = KSPathUtils.getStandardFilePath( 'carouselDocuments' )
	local baseName = KSPathUtils.child( documentDir, 'embeddings' )

	local embeddings = WFMLImp.WFMLEmbeddings( baseName, 2 )

	rxTest:expectEqual( 2, embeddings.vectorSize )
	rxTest:expectEqual( baseName, embeddings.baseName )

	local id1 = KSIdGenerator.uuidInLowercaseHex()
	local v1 = WFStruct.pack( '2f', 1.0, 1.0 )

	embeddings:insert( id1, v1 )
	rxTest:expectEqual( true, embeddings:hasKey( id1 ) )

	local id2 = KSIdGenerator.uuidInLowercaseHex()
	local v2 = WFStruct.pack( '2f', 0.0, 0.0 )

	embeddings:insert( id2, v2 )
	rxTest:expectEqual( true, embeddings:hasKey( id2 ) )

	local id3 = KSIdGenerator.uuidInLowercaseHex()
	local v3 = WFStruct.pack( '2f', 1.0, 0.0 )

	embeddings:insert( id3, v3 )
	rxTest:expectEqual( true, embeddings:hasKey( id3 ) )

	embeddings:offloadIndex()

	local id4 = KSIdGenerator.uuidInLowercaseHex()
	local v4 = WFStruct.pack( '2f', 0.0, 1.0 )

	embeddings:insert( id4, v4 )
	rxTest:expectEqual( true, embeddings:hasKey( id4 ) )

	local vs = WFStruct.pack( '2f', 0.9, 0.9 )

	embeddings:offloadIndex()

	local matches = embeddings:findClosest( vs, 2, 0.0 )

	rxTest:expectEqual( 2, #matches )
	rxTest:expectEqual( id1, matches[1].id )

end

--------------------------------------------------------------------------------
--@mark Performance Tests
--------------------------------------------------------------------------------

local function generateRandomVector( size )

	local vector = {}

	-- Random value between -1 and 1
	for i = 1, size do
		vector[ i ] = math.random() * 2 - 1
	end

	-- Normalize it
	local sum = 0
	for _, v in ipairs( vector ) do
		sum = sum + v * v
	end
	sum = math.sqrt( sum )

	if sum > 0 then
		for i, v in ipairs( vector ) do
			vector[ i ] = v / sum
		end
	end

	return vector

end

--------------------------------------------------------------------------------

local function testSearchPerformance( rxTest, numVectors, embeddingSize )

	rxTest:setTimeLimit( 1000 )

	local documentDir = KSPathUtils.getStandardFilePath( 'carouselDocuments' )
	local baseName = KSPathUtils.child( documentDir, 'embeddings' )

	local embeddings = WFMLImp.WFMLEmbeddings( baseName, embeddingSize )

	-- Generate numVectors with a vector size of embeddingSize
	for i = 1, numVectors do
		local vectorData = generateRandomVector( embeddingSize )
		local id = KSIdGenerator.uuidInLowercaseHex()
		local packedData = WFStruct.pack( string.format( '%df', embeddingSize ), unpack( vectorData ) )
		embeddings:insert( id, packedData )
	end

	-- Generate a random search vector
	local searchVector = generateRandomVector( embeddingSize )
	local packedSearchVector = WFStruct.pack( string.format( '%df', embeddingSize ), unpack( searchVector ) )

	-- Test search performance
	rxTest:measurePerformance( function()
		local neighbors = math.floor( numVectors / 100 )
		local results = embeddings:findClosest( packedSearchVector, neighbors, 0.0 )
		rxTest:expectEqual( neighbors, #results )
	end )

	embeddings:offloadIndex()

	-- Test search performance after offloading index
	rxTest:measurePerformance( function()
		local neighbors = math.floor( numVectors / 100 )
		local results = embeddings:findClosest( packedSearchVector, neighbors, 0.0 )
		rxTest:expectEqual( neighbors, #results )
	end )

end

--------------------------------------------------------------------------------

local vectorSizesToTest = {
	{ 1000, 512 },
	{ 10000, 512 }
}

for _, vectorSize in ipairs( vectorSizesToTest ) do

	local numVectors = vectorSize[ 1 ]
	local embeddingSize = vectorSize[ 2 ]

	local testName = string.format( "rxTestSearch_%d_%d_testPerformance", numVectors, embeddingSize )

	WFMLEmbeddingsTest[ testName ] = function( rxTest )

		testSearchPerformance( rxTest, numVectors, embeddingSize )

	end

end

--------------------------------------------------------------------------------

return WFRxTestCase.wrapSuite( 'WFMLEmbeddingsTest', WFMLEmbeddingsTest )
