package com.easyagents.rag.retrieval;

import com.easyagents.core.document.Document;
import org.junit.Assert;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class RagScoreNormalizerTest {

    @Test
    public void shouldNormalizeKeywordScoresToZeroAndOneRange() {
        Document first = document(1, 9D, RagRetrievalMetadataKeys.KEYWORD_SCORE);
        Document second = document(2, 0D, RagRetrievalMetadataKeys.KEYWORD_SCORE);

        RagScoreNormalizer.normalize(Arrays.asList(first, second), RetrievalMode.KEYWORD, false);

        Assert.assertEquals(0.9D, first.getScore(), 0.0001D);
        Assert.assertEquals(0D, second.getScore(), 0.0001D);
    }

    @Test
    public void shouldNormalizeHybridFusionScoreByRrfUpperBound() {
        Document document = document(1, 2D / (RrfFusionStrategy.DEFAULT_RRF_K + 1D), RagRetrievalMetadataKeys.FUSION_SCORE);

        RagScoreNormalizer.normalize(Arrays.asList(document), RetrievalMode.HYBRID, false);

        Assert.assertEquals(1D, document.getScore(), 0.0001D);
    }

    @Test
    public void shouldNormalizeRerankScoresByMinMax() {
        List<Document> documents = Arrays.asList(
            document(1, 10D, RagRetrievalMetadataKeys.RERANK_SCORE),
            document(2, 20D, RagRetrievalMetadataKeys.RERANK_SCORE),
            document(3, 30D, RagRetrievalMetadataKeys.RERANK_SCORE)
        );

        RagScoreNormalizer.normalize(documents, RetrievalMode.HYBRID, true);

        Assert.assertEquals(0D, documents.get(0).getScore(), 0.0001D);
        Assert.assertEquals(0.5D, documents.get(1).getScore(), 0.0001D);
        Assert.assertEquals(1D, documents.get(2).getScore(), 0.0001D);
    }

    @Test
    public void shouldFallbackToRankBasedNormalizationWhenRerankScoresAreEqual() {
        List<Document> documents = Arrays.asList(
            document(1, 5D, RagRetrievalMetadataKeys.RERANK_SCORE),
            document(2, 5D, RagRetrievalMetadataKeys.RERANK_SCORE),
            document(3, 5D, RagRetrievalMetadataKeys.RERANK_SCORE)
        );

        RagScoreNormalizer.normalize(documents, RetrievalMode.HYBRID, true);

        Assert.assertEquals(1D, documents.get(0).getScore(), 0.0001D);
        Assert.assertEquals(0.5D, documents.get(1).getScore(), 0.0001D);
        Assert.assertEquals(0D, documents.get(2).getScore(), 0.0001D);
    }

    private Document document(Object id, Double score, String metadataKey) {
        Document document = new Document();
        document.setId(id);
        document.setScore(score);
        document.addMetadata(metadataKey, score);
        return document;
    }
}
