#!/usr/bin/env python3

import random
import unittest

import torch
from reagent.ope.estimators.slate_estimators import (
    DCGSlateMetric,
    NDCGSlateMetric,
    Slate,
    SlateItem,
    SlateItemProbabilities,
    SlateItemValues,
    SlateSlotItemProbabilities,
    SlateSlots,
)


class TestEstimator(unittest.TestCase):
    def setUp(self) -> None:
        random.seed(1234)
        torch.random.manual_seed(1234)
        self._item_relevances = [1.0, 0.5, 2.5, 2.0, 5.0]
        self._slot_item_relevances = [
            [1.0, 0.5, 2.5, 2.0, 5.0],
            [1.5, 1.0, 2.0, 1.0, 4.0],
            [3.0, 2.5, 0.5, 4.0, 2.0],
        ]
        self._item_rewards = [3.0, 8.0, 0.0, 4.0, 1.0]
        self._item_size = len(self._item_relevances)
        self._slate_size = 3
        self._slots = SlateSlots(self._slate_size)

    def test_slate_item_probabilities(self):
        probs = SlateItemProbabilities(self._item_relevances)
        slate = probs.sample_slate(self._slots)
        slate_prob = probs.slate_probability(slate)
        self.assertAlmostEqual(slate_prob, 0.017825312, places=2)
        slot_item_expectations = probs.slot_item_expectations(self._slots)
        slot_rewards = slot_item_expectations.expected_rewards(
            SlateItemValues(self._item_rewards)
        )
        diff = slot_rewards.values - torch.tensor([1.81818, 2.13736, 2.66197])
        self.assertAlmostEqual(diff.sum().item(), 0, places=5)

    def test_slate_slot_item_probabilities(self):
        probs = SlateSlotItemProbabilities(
            [SlateItemValues(vs) for vs in self._slot_item_relevances]
        )
        slate = probs.sample_slate(self._slots)
        slate_prob = probs.slate_probability(slate)
        self.assertAlmostEqual(slate_prob, 0.02139037, places=2)
        slot_item_expectations = probs.slot_item_expectations()
        slot_rewards = slot_item_expectations.expected_rewards(
            SlateItemValues(self._item_rewards)
        )
        diff = slot_rewards.values - torch.tensor([1.81818, 2.51352, 7.36929])
        self.assertAlmostEqual(diff.sum().item(), 0, places=5)

    def test_metrics(self):
        dcg = DCGSlateMetric()
        ndcg = NDCGSlateMetric(SlateItemValues([1.0, 2.5, 2.0, 3.0, 1.5, 0.0]))
        item_rewards = SlateItemValues([2.0, 1.0, 0.0, 3.0, 1.5, 2.5])
        slate = Slate([SlateItem(1), SlateItem(3), SlateItem(2)])
        reward = dcg(slate.slots, slate.slot_values(item_rewards))
        self.assertAlmostEqual(reward, 5.416508275)
        reward = ndcg(slate.slots, slate.slot_values(item_rewards))
        self.assertAlmostEqual(reward, 0.473547669)
        slate = Slate([SlateItem(5), SlateItem(0), SlateItem(4)])
        reward = dcg(slate.slots, slate.slot_values(item_rewards))
        self.assertAlmostEqual(reward, 7.463857073)
        reward = ndcg(slate.slots, slate.slot_values(item_rewards))
        self.assertAlmostEqual(reward, 0.652540703)
