package com.tzld.piaoquan.recommend.server.util;

import com.tzld.piaoquan.recommend.server.model.Video;
import com.tzld.piaoquan.recommend.server.service.funnel.FunnelContext;
import com.tzld.piaoquan.recommend.server.service.funnel.RecallVideoEntry;
import com.tzld.piaoquan.recommend.server.service.funnel.SelectKind;
import com.tzld.piaoquan.recommend.server.service.rank.RankParam;
import com.tzld.piaoquan.recommend.server.service.recall.RecallResult;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;

class RecallUtilsTest {

    @Test
    void countDistinctCandidatesDeduplicatesAndExcludesPreviouslySelectedVideos() {
        RankParam param = rankParam(
                Arrays.asList(
                        recallData("p1", video(1, 10), video(2, 9)),
                        recallData("p2", video(2, 9), video(3, 8)),
                        recallData("ignored", video(4, 7))
                ),
                Collections.emptyList()
        );

        int candidates = RecallUtils.countDistinctCandidates(
                param, Collections.singleton(1L), new HashSet<>(Arrays.asList("p1", "p2")));

        assertEquals(2, candidates);
    }

    @Test
    void coarseRankOnlyAttributesFilteredCandidatesThatWereSelected() {
        RankParam param = rankParam(
                recallData("p1", video(1, 10), video(3, 1)),
                entries("p1",
                        entry(1, true),
                        entry(2, false),
                        entry(3, true))
        );

        RecallUtils.extractAllAndTruncateByCoarseRank(
                1, param, new HashSet<>(), new ArrayList<>(),
                scores(1, 10, 3, 1), Collections.singleton("p1"));

        List<RecallVideoEntry> entries = param.getFunnelContext()
                .getStages123RecallByStrategy().get("p1");
        assertEquals(SelectKind.SELF, entries.get(0).getSelect());
        assertNull(entries.get(1).getSelect(), "filtered video must not enter step 3");
        assertNull(entries.get(2).getSelect(), "ordinary coarse-rank loser must not enter step 3");
    }

    @Test
    void coarseRankMarksDuplicatePushFromAsOther() {
        RankParam param = rankParam(
                Arrays.asList(
                        recallData("p1", video(1, 10)),
                        recallData("p2", video(1, 10))
                ),
                Arrays.asList(
                        entries("p1", entry(1, true)),
                        entries("p2", entry(1, true))
                )
        );

        RecallUtils.extractAllAndTruncateByCoarseRank(
                1, param, new HashSet<>(), new ArrayList<>(),
                scores(1, 10), new HashSet<>(Arrays.asList("p1", "p2")));

        assertEquals(SelectKind.SELF, entry(param, "p1", 0).getSelect());
        assertEquals(SelectKind.OTHER, entry(param, "p2", 0).getSelect());
    }

    @Test
    void coarseRankMarksCandidateSelectedByPreviousPoolAsOther() {
        RankParam param = rankParam(
                Arrays.asList(
                        recallData("personal", video(1, 10)),
                        recallData("nonPersonal", video(1, 10), video(2, 9))
                ),
                Arrays.asList(
                        entries("personal", entry(1, true)),
                        entries("nonPersonal", entry(1, true), entry(2, true))
                )
        );
        Set<Long> selectedIds = new HashSet<>();
        List<Video> result = new ArrayList<>();

        RecallUtils.extractAllAndTruncateByCoarseRank(
                1, param, selectedIds, result, scores(1, 10, 2, 9),
                Collections.singleton("personal"));
        RecallUtils.extractAllAndTruncateByCoarseRank(
                1, param, selectedIds, result, scores(1, 10, 2, 9),
                Collections.singleton("nonPersonal"));

        assertEquals(SelectKind.OTHER, entry(param, "nonPersonal", 0).getSelect());
        assertEquals(SelectKind.SELF, entry(param, "nonPersonal", 1).getSelect());
    }

    private static RankParam rankParam(RecallResult.RecallData data,
                                       Map.Entry<String, List<RecallVideoEntry>> entries) {
        return rankParam(Collections.singletonList(data), Collections.singletonList(entries));
    }

    private static RankParam rankParam(List<RecallResult.RecallData> data,
                                       List<Map.Entry<String, List<RecallVideoEntry>>> entries) {
        FunnelContext context = new FunnelContext();
        for (Map.Entry<String, List<RecallVideoEntry>> item : entries) {
            context.getStages123RecallByStrategy().put(item.getKey(), item.getValue());
        }
        RankParam param = new RankParam();
        param.setRecallResult(new RecallResult(data));
        param.setFunnelContext(context);
        return param;
    }

    private static RecallResult.RecallData recallData(String pushFrom, Video... videos) {
        return new RecallResult.RecallData(pushFrom, Arrays.asList(videos));
    }

    private static Map.Entry<String, List<RecallVideoEntry>> entries(
            String pushFrom, RecallVideoEntry... entries) {
        return new java.util.AbstractMap.SimpleEntry<>(pushFrom, Arrays.asList(entries));
    }

    private static RecallVideoEntry entry(long videoId, boolean filteredIn) {
        RecallVideoEntry entry = new RecallVideoEntry(videoId, 0, 0);
        entry.setFilteredIn(filteredIn);
        entry.setIndexNewAfterFilter(filteredIn ? 0 : -1);
        return entry;
    }

    private static RecallVideoEntry entry(RankParam param, String pushFrom, int index) {
        return param.getFunnelContext().getStages123RecallByStrategy().get(pushFrom).get(index);
    }

    private static Video video(long videoId, double score) {
        Video video = new Video();
        video.setVideoId(videoId);
        video.setRovScore(score);
        return video;
    }

    private static Map<Long, Double> scores(double... values) {
        Map<Long, Double> scores = new HashMap<>();
        for (int i = 0; i < values.length; i += 2) {
            scores.put((long) values[i], values[i + 1]);
        }
        return scores;
    }
}
